Skip to content

Commit

Permalink
Features/jerry/function level env (#46)
Browse files Browse the repository at this point in the history
* completed draft for function level env var

* completed func level env variable + added new tests

* fix issues

* Remove test_workflow_builder_func_env_var.py and fix tests

---------

Co-authored-by: vgsteiger <[email protected]>
  • Loading branch information
jerryyiransun and vGsteiger authored Jan 11, 2024
1 parent 98c24f5 commit efe3aef
Show file tree
Hide file tree
Showing 12 changed files with 236 additions and 25 deletions.
7 changes: 7 additions & 0 deletions docs/design.md
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,12 @@ workflow = MultiXServerlessWorkflow("workflow_name")
}
],
},
environment_variables=[
{
"key": "example_key",
"value": "example_value"
}
],
)
```

Expand All @@ -161,6 +167,7 @@ The dictionary has two keys:
Each provider is a dictionary with two keys:
- `name`: The name of the provider.
- `config`: A dictionary that contains the configuration for the specific provider.
- `environment_variables`: This parameter represents a list of dictionaries, each designed for setting environment variables specifically for a function. Users must adhere to a structured format within each dictionary. This format requires two entries: "key" and "value". The "key" entry should contain the name of the environment variable, serving as an identifier. The "value" entry holds the corresponding value assigned to that variable.

- Within a function, a user can register a call to another function with the following annotation:

Expand Down
2 changes: 1 addition & 1 deletion multi_x_serverless/deployment/client/cli/config_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@


class EnvironmentVariable(BaseModel):
name: str = Field(..., title="The name of the environment variable")
key: str = Field(..., title="The name of the environment variable")
value: str = Field(..., title="The value of the environment variable")


Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
workflow_name: "{{ workflow_name }}"
environment_variables:
- name: "ENV_VAR_1"
- key: "ENV_VAR_1"
value: "value_1"
iam_policy_file: "iam_policy.json"
home_regions:
Expand Down
1 change: 1 addition & 0 deletions multi_x_serverless/deployment/client/cli/template/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
}
],
},
environment_variables=[{"key": "example_key", "value": "example_value"}],
)
def first_function(event: dict[str, Any]) -> dict[str, Any]:
payload = {
Expand Down
10 changes: 7 additions & 3 deletions multi_x_serverless/deployment/client/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,13 +42,17 @@ def python_version(self) -> str:
return "python3.11"

@property
def environment_variables(self) -> dict[str, Any]:
def environment_variables(self) -> dict[str, str]:
list_of_env_variables: list[dict] = self._lookup("environment_variables")
if list_of_env_variables is None:
return {}
env_variables: dict[str, Any] = {}
env_variables: dict[str, str] = {}
for env_variable in list_of_env_variables:
env_variables[env_variable["name"]] = env_variable["value"]
if not isinstance(env_variable["value"], str):
raise RuntimeError("Environment variable value need to be a str")
if not isinstance(env_variable["key"], str):
raise RuntimeError("Environment variable key need to be a str")
env_variables[env_variable["key"]] = env_variable["value"]
return env_variables

@property
Expand Down
20 changes: 18 additions & 2 deletions multi_x_serverless/deployment/client/deploy/workflow_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,11 +38,14 @@ def build_workflow(self, config: Config) -> Workflow: # pylint: disable=too-man
else:
providers = config.regions_and_providers["providers"]
self._verify_providers(providers)

merged_env_vars = self.merge_environment_variables(
function.environment_variables, config.environment_variables
)
resources.append(
Function(
name=function_deployment_name,
# TODO (#22): Add function specific environment variables
environment_variables=config.environment_variables,
environment_variables=merged_env_vars,
runtime=config.python_version,
handler=function.handler,
role=function_role,
Expand Down Expand Up @@ -158,3 +161,16 @@ def get_function_role(self, config: Config, function_name: str) -> IAMRole:
filename = os.path.join(config.project_dir, ".multi-x-serverless", "iam_policy.yml")

return IAMRole(role_name=role_name, policy=filename)

def merge_environment_variables(
self, function_env_vars: Optional[list[dict[str, str]]], config_env_vars: dict[str, str]
) -> dict[str, str]:
if not function_env_vars:
return config_env_vars

merged_env_vars: dict[str, str] = dict(config_env_vars)
# overwrite config env vars with function env vars if duplicate
for env_var in function_env_vars:
merged_env_vars[env_var["key"]] = env_var["value"]

return merged_env_vars
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,14 @@ def __init__(
name: str,
entry_point: bool,
regions_and_providers: dict,
environment_variables: list[dict[str, str]],
):
self.function_callable = function_callable
self.name = name
self.entry_point = entry_point
self.handler = function_callable.__name__
self.regions_and_providers = regions_and_providers if len(regions_and_providers) > 0 else None
self.environment_variables = environment_variables if len(environment_variables) > 0 else None
self.validate_function_name()

def validate_function_name(self) -> None:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -311,6 +311,7 @@ def register_function(
name: str,
entry_point: bool,
regions_and_providers: dict,
environment_variables: list[dict[str, str]],
) -> None:
"""
Register a function as a serverless function.
Expand All @@ -320,15 +321,15 @@ def register_function(
At this point we only need to register the function with the wrapper, the actual deployment will be done
later by the deployment manager.
"""
wrapper = MultiXServerlessFunction(function, name, entry_point, regions_and_providers)
wrapper = MultiXServerlessFunction(function, name, entry_point, regions_and_providers, environment_variables)
self.functions[function.__name__] = wrapper

# TODO (#22): Add function specific environment variables
def serverless_function(
self,
name: Optional[str] = None,
entry_point: bool = False,
regions_and_providers: Optional[dict] = None,
environment_variables: Optional[list[dict[str, str]]] = None,
) -> Callable[..., Any]:
"""
Decorator to register a function as a Lambda function.
Expand Down Expand Up @@ -366,6 +367,21 @@ def serverless_function(
if regions_and_providers is None:
regions_and_providers = {}

if environment_variables is None:
environment_variables = []
else:
if not isinstance(environment_variables, list):
raise RuntimeError("environment_variables must be a list of dicts")
for env_variable in environment_variables:
if not isinstance(env_variable, dict):
raise RuntimeError("environment_variables must be a list of dicts")
if "key" not in env_variable or "value" not in env_variable:
raise RuntimeError("environment_variables must be a list of dicts with keys 'key' and 'value'")
if not isinstance(env_variable["key"], str):
raise RuntimeError("environment_variables must be a list of dicts with 'key' as a string")
if not isinstance(env_variable["value"], str):
raise RuntimeError("environment_variables must be a list of dicts with 'value' as a string")

def _register_handler(func: Callable[..., Any]) -> Callable[..., Any]:
handler_name = name if name is not None else func.__name__

Expand All @@ -392,7 +408,7 @@ def wrapper(*args, **kwargs): # type: ignore # pylint: disable=unused-argument
wrapper.routing_decision = {} # type: ignore
wrapper.entry_point = entry_point # type: ignore
wrapper.original_function = func # type: ignore
self.register_function(func, handler_name, entry_point, regions_and_providers)
self.register_function(func, handler_name, entry_point, regions_and_providers, environment_variables)
return wrapper

return _register_handler
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,13 @@ def test_build_workflow_multiple_entry_points(self):
function1.name = "function1"
function1.handler = "function1"
function1.regions_and_providers = {}
function1.environment_variables = {}
function2 = Mock(spec=MultiXServerlessFunction)
function2.entry_point = True
function2.name = "function2"
function2.handler = "function1"
function2.regions_and_providers = {"providers": []}
function2.environment_variables = {}
self.config.workflow_app.functions = {"function1": function1, "function2": function2}
with self.assertRaisesRegex(RuntimeError, "Multiple entry points defined"):
self.builder.build_workflow(self.config)
Expand All @@ -45,13 +47,15 @@ def test_build_workflow_merge_case_self_cycle(self):
function1.name = "function1"
function1.handler = "function1"
function1.regions_and_providers = {}
function1.environment_variables = []
function1.is_waiting_for_predecessors = Mock(return_value=False)

function2 = Mock(spec=MultiXServerlessFunction)
function2.entry_point = False
function2.name = "function2"
function2.handler = "function2"
function2.regions_and_providers = {"providers": []}
function2.environment_variables = []
function2.is_waiting_for_predecessors = Mock(return_value=True) # This is a merge function

# Mock the workflow app to return the successors of function1
Expand All @@ -74,13 +78,15 @@ def test_build_workflow_merge_case_multiple_incoming(self):
function1.name = "function1"
function1.handler = "function1"
function1.regions_and_providers = {}
function1.environment_variables = []
function1.is_waiting_for_predecessors = Mock(return_value=False)

function2 = Mock(spec=MultiXServerlessFunction)
function2.entry_point = False
function2.name = "function2"
function2.handler = "function2"
function2.regions_and_providers = {"providers": []}
function2.environment_variables = []
function2.is_waiting_for_predecessors = Mock(return_value=True) # This is a merge function

# Mock the workflow app to return the successors of function1
Expand Down Expand Up @@ -140,6 +146,7 @@ def test_build_workflow_self_call(self):
function1.name = "function1"
function1.handler = "function1"
function1.regions_and_providers = {}
function1.environment_variables = []
function1.is_waiting_for_predecessors = Mock(return_value=False)

# Mock the workflow app to return the successors of function1
Expand All @@ -162,13 +169,15 @@ def test_build_workflow_merge_working(self):
function1.name = "function1"
function1.handler = "function1"
function1.regions_and_providers = {}
function1.environment_variables = []
function1.is_waiting_for_predecessors = Mock(return_value=False)

function2 = Mock(spec=MultiXServerlessFunction)
function2.entry_point = False
function2.name = "function2"
function2.handler = "function2"
function2.regions_and_providers = {"providers": []}
function2.environment_variables = []
function2.is_waiting_for_predecessors = Mock(return_value=True) # This is a merge function

# Mock the workflow app to return the successors of function1
Expand All @@ -189,13 +198,15 @@ def test_build_workflow_cycle_in_function_calls(self):
function1.name = "function1"
function1.handler = "function1"
function1.regions_and_providers = {}
function1.environment_variables = []
function1.is_waiting_for_predecessors = Mock(return_value=False)

function2 = Mock(spec=MultiXServerlessFunction)
function2.entry_point = False
function2.name = "function2"
function2.handler = "function2"
function2.regions_and_providers = {"providers": []}
function2.environment_variables = []
function2.is_waiting_for_predecessors = Mock(return_value=False) # This is a merge function

# Mock the workflow app to return the successors of function1
Expand All @@ -220,13 +231,15 @@ def test_build_workflow_merge_cycle(self):
function1.name = "function1"
function1.handler = "function1"
function1.regions_and_providers = {}
function1.environment_variables = []
function1.is_waiting_for_predecessors = Mock(return_value=False)

function2 = Mock(spec=MultiXServerlessFunction)
function2.entry_point = False
function2.name = "function2"
function2.handler = "function2"
function2.regions_and_providers = {"providers": []}
function2.environment_variables = []
function2.is_waiting_for_predecessors = Mock(return_value=True)

self.config.workflow_app.get_successors = Mock(side_effect=[[function2], [function2], []])
Expand All @@ -251,6 +264,75 @@ def test_get_function_role_without_policy_file(self, mock_join):
self.assertEqual(role.name, "test_function-role")
self.assertEqual(role.policy, "/path/to/default_policy")

def test_build_func_environment_variables(self):
# function 1 (empty function level environment variables)
function1 = Mock(spec=MultiXServerlessFunction)
function1.entry_point = True
function1.name = "function1"
function1.handler = "function1"
function1.regions_and_providers = {}
function1.environment_variables = []

# function 2 (no overlap with global environment variables)
function2 = Mock(spec=MultiXServerlessFunction)
function2.entry_point = False
function2.name = "function2"
function2.handler = "function1"
function2.regions_and_providers = {"providers": []}
function2.environment_variables = [{"key": "ENV_3", "value": "function2_env_3"}]

# function 3 (overlap with global environment variables)
function3 = Mock(spec=MultiXServerlessFunction)
function3.entry_point = False
function3.name = "function2"
function3.handler = "function1"
function3.regions_and_providers = {"providers": []}
function3.environment_variables = [{"key": "ENV_1", "value": "function3_env_1"}]

self.builder = WorkflowBuilder()
self.config = Mock(spec=Config)
self.config.workflow_name = "test_workflow"
self.config.workflow_app.functions = {"function1": function1, "function2": function2, "function3": function3}
self.config.environment_variables = {
"ENV_1": "global_env_1",
"ENV_2": "global_env_2",
}
self.config.python_version = "3.8"
self.config.home_regions = []
self.config.project_dir = "/path/to/project"
self.config.iam_policy_file = None
self.config.regions_and_providers = {"providers": []}
self.config.workflow_app.get_successors.return_value = []

workflow = self.builder.build_workflow(self.config)

self.assertEqual(len(workflow._resources), 3)
built_func1 = workflow._resources[0]
built_func2 = workflow._resources[1]
built_func3 = workflow._resources[2]
self.assertEqual(
built_func1.environment_variables,
{
"ENV_1": "global_env_1",
"ENV_2": "global_env_2",
},
)
self.assertEqual(
built_func2.environment_variables,
{
"ENV_1": "global_env_1",
"ENV_2": "global_env_2",
"ENV_3": "function2_env_3",
},
)
self.assertEqual(
built_func3.environment_variables,
{
"ENV_1": "function3_env_1",
"ENV_2": "global_env_2",
},
)


if __name__ == "__main__":
unittest.main()
2 changes: 1 addition & 1 deletion multi_x_serverless/tests/deployment/client/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ def test_python_version(self):
self.assertTrue(self.config.python_version.startswith("python"))

def test_environment_variables(self):
self.config.project_config["environment_variables"] = [{"name": "ENV", "value": "test"}]
self.config.project_config["environment_variables"] = [{"key": "ENV", "value": "test"}]
self.assertEqual(self.config.environment_variables, {"ENV": "test"})

def test_home_regions(self):
Expand Down
Loading

0 comments on commit efe3aef

Please sign in to comment.