Repository: run-llama/llama_deploy

Issue:
Add environment variables to deployment definition
At the moment the only way to pass env vars to a running instance of Llama Deploy is through setting them before starting the process, which is not practical in case of a remote installation or docker containers.

We should add the option to declare the environment variables needed by each service like this:
```yaml
services:
  test-workflow:
    name: Test Workflow
    env:
        OPENAI_API_KEY: "secret!"
    source:
      type: git
      name: https://github.com/run-llama/llama_deploy.git
    path: tests/apiserver/data/workflow:my_workflow
```
then the Control plane would propagate the env var when spinning up the service.


Discussion:

Patch:
diff --git a/llama_deploy/apiserver/config_parser.py b/llama_deploy/apiserver/config_parser.py
index 774efa16..ce5ad1e4 100644
--- a/llama_deploy/apiserver/config_parser.py
+++ b/llama_deploy/apiserver/config_parser.py
@@ -58,6 +58,8 @@ class Service(BaseModel):
     path: str | None = None
     host: str | None = None
     port: int | None = None
+    env: dict[str, str] | None = Field(None)
+    env_files: list[str] | None = Field(None, alias="env-files")
     python_dependencies: list[str] | None = Field(None, alias="python-dependencies")
     ts_dependencies: dict[str, str] | None = Field(None, alias="ts-dependencies")
 
diff --git a/llama_deploy/apiserver/deployment.py b/llama_deploy/apiserver/deployment.py
index 80243b07..d12c97ba 100644
--- a/llama_deploy/apiserver/deployment.py
+++ b/llama_deploy/apiserver/deployment.py
@@ -1,11 +1,14 @@
 import asyncio
 import importlib
+import os
 import subprocess
 import sys
+from dotenv import dotenv_values
 from multiprocessing.pool import ThreadPool
 from pathlib import Path
 from typing import Any
 
+
 from llama_deploy import (
     Client,
     ControlPlaneServer,
@@ -154,11 +157,15 @@ def _load_services(self, config: Config) -> list[WorkflowService]:
             # Install dependencies
             self._install_dependencies(service_config)
 
+            # Set environment variables
+            self._set_environment_variables(service_config, destination)
+
             # Search for a workflow instance in the service path
             pythonpath = (destination / service_config.path).parent.resolve()
             sys.path.append(str(pythonpath))
             module_name, workflow_name = Path(service_config.path).name.split(":")
             module = importlib.import_module(module_name)
+
             workflow = getattr(module, workflow_name)
             workflow_config = WorkflowServiceConfig(
                 host=service_config.host,
@@ -177,6 +184,26 @@ def _load_services(self, config: Config) -> list[WorkflowService]:
 
         return workflow_services
 
+    @staticmethod
+    def _set_environment_variables(
+        service_config: Service, root: Path | None = None
+    ) -> None:
+        """Sets environment variables for the service."""
+        env_vars: dict[str, str | None] = {}
+
+        if service_config.env:
+            env_vars.update(**service_config.env)
+
+        if service_config.env_files:
+            for env_file in service_config.env_files:
+                # use dotenv to parse env_file
+                env_file_path = root / env_file if root else Path(env_file)
+                env_vars.update(**dotenv_values(env_file_path))
+
+        for k, v in env_vars.items():
+            if v:
+                os.environ[k] = v
+
     @staticmethod
     def _install_dependencies(service_config: Service) -> None:
         """Runs `pip install` on the items listed under `python-dependencies` in the service configuration."""
diff --git a/poetry.lock b/poetry.lock
index 0c2b9dd6..2555dd31 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3078,4 +3078,4 @@ solace = ["solace-pubsubplus"]
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.10,<4.0"
-content-hash = "4bd6c9526b331b24bdd7209a26797d91174b86e71d7ef5a072aec8b1b43fcf03"
+content-hash = "51bc5b5556f3cbc9b13ec3b3cdb7bd98ecb7e597a7fe79de33e7c78c2c831488"
diff --git a/pyproject.toml b/pyproject.toml
index 3836ef1f..9c6eeccc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -52,6 +52,7 @@ gitpython = "^3.1.43"
 python-multipart = "^0.0.18"
 typing_extensions = "^4.0.0"
 asgiref = "^3.8.1"
+python-dotenv = "^1.0.1"
 
 [tool.poetry.extras]
 kafka = ["aiokafka", "kafka-python-ng"]


Test patch:
diff --git a/e2e_tests/apiserver/deployments/deployment_env_git.yml b/e2e_tests/apiserver/deployments/deployment_env_git.yml
new file mode 100644
index 00000000..b0f509db
--- /dev/null
+++ b/e2e_tests/apiserver/deployments/deployment_env_git.yml
@@ -0,0 +1,19 @@
+name: EnvironmentVariablesGit
+
+control-plane:
+  port: 8000
+
+default-service: test_env_workflow
+
+services:
+  workflow_git:
+    name: Git Workflow
+    source:
+      type: git
+      name: https://github.com/run-llama/llama_deploy.git
+    env:
+      VAR_1: x  # this gets overwritten because VAR_1 also exists in the provided .env
+      VAR_2: y
+    env-files:
+      - tests/apiserver/data/.env  # relative to source path
+    path: tests/apiserver/data/workflow:env_reader_workflow
diff --git a/e2e_tests/apiserver/deployments/deployment_env_local.yml b/e2e_tests/apiserver/deployments/deployment_env_local.yml
new file mode 100644
index 00000000..104c5934
--- /dev/null
+++ b/e2e_tests/apiserver/deployments/deployment_env_local.yml
@@ -0,0 +1,19 @@
+name: EnvironmentVariablesLocal
+
+control-plane:
+  port: 8000
+
+default-service: test_env_workflow
+
+services:
+  test_env_workflow:
+    name: Workflow
+    source:
+      type: local
+      name: ./e2e_tests/apiserver/deployments/src
+    env:
+      VAR_1: x  # this gets overwritten because VAR_1 also exists in the provided .env
+      VAR_2: y
+    env-files:
+      - .env  # relative to source path
+    path: workflow_env:workflow
diff --git a/e2e_tests/apiserver/deployments/src/.env b/e2e_tests/apiserver/deployments/src/.env
new file mode 100644
index 00000000..25f126b2
--- /dev/null
+++ b/e2e_tests/apiserver/deployments/src/.env
@@ -0,0 +1,2 @@
+VAR_1=z
+API_KEY=123
diff --git a/e2e_tests/apiserver/deployments/src/workflow_env.py b/e2e_tests/apiserver/deployments/src/workflow_env.py
new file mode 100644
index 00000000..362225ad
--- /dev/null
+++ b/e2e_tests/apiserver/deployments/src/workflow_env.py
@@ -0,0 +1,43 @@
+import asyncio
+import os
+
+from llama_index.core.workflow import (
+    Context,
+    StartEvent,
+    StopEvent,
+    Workflow,
+    step,
+)
+
+
+class MyWorkflow(Workflow):
+    @step()
+    async def run_step(self, ctx: Context, ev: StartEvent) -> StopEvent:
+        var_1 = os.environ.get("VAR_1")
+        var_2 = os.environ.get("VAR_2")
+        api_key = os.environ.get("API_KEY")
+        return StopEvent(
+            # result depends on variables read from environment
+            result=(f"var_1: {var_1}, " f"var_2: {var_2}, " f"api_key: {api_key}")
+        )
+
+
+workflow = MyWorkflow()
+
+
+async def main(w: Workflow):
+    h = w.run()
+    print(await h)
+
+
+if __name__ == "__main__":
+    import os
+
+    # set env variables
+    os.environ["VAR_1"] = "x"
+    os.environ["VAR_1"] = "y"
+    os.environ["API_KEY"] = "123"
+
+    w = MyWorkflow()
+
+    asyncio.run(main(w))
diff --git a/e2e_tests/apiserver/test_env_vars_git.py b/e2e_tests/apiserver/test_env_vars_git.py
new file mode 100644
index 00000000..e2616909
--- /dev/null
+++ b/e2e_tests/apiserver/test_env_vars_git.py
@@ -0,0 +1,22 @@
+import asyncio
+from pathlib import Path
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_read_env_vars_git(apiserver, client):
+    here = Path(__file__).parent
+
+    with open(here / "deployments" / "deployment_env_git.yml") as f:
+        await client.apiserver.deployments.create(f)
+        await asyncio.sleep(5)
+
+    session = await client.core.sessions.create()
+
+    # run workflow
+    result = await session.run(
+        "workflow_git", env_vars_to_read=["VAR_1", "VAR_2", "API_KEY"]
+    )
+
+    assert result == "VAR_1: x, VAR_2: y, API_KEY: 123"
diff --git a/e2e_tests/apiserver/test_env_vars_local.py b/e2e_tests/apiserver/test_env_vars_local.py
new file mode 100644
index 00000000..93e14dcd
--- /dev/null
+++ b/e2e_tests/apiserver/test_env_vars_local.py
@@ -0,0 +1,20 @@
+import asyncio
+from pathlib import Path
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_read_env_vars_local(apiserver, client):
+    here = Path(__file__).parent
+
+    with open(here / "deployments" / "deployment_env_local.yml") as f:
+        await client.apiserver.deployments.create(f)
+        await asyncio.sleep(5)
+
+    session = await client.core.sessions.create()
+
+    # run workflow
+    result = await session.run("test_env_workflow")
+
+    assert result == "var_1: z, var_2: y, api_key: 123"
diff --git a/e2e_tests/basic_streaming/test_run_client.py b/e2e_tests/basic_streaming/test_run_client.py
index df5f79e0..b6041e00 100644
--- a/e2e_tests/basic_streaming/test_run_client.py
+++ b/e2e_tests/basic_streaming/test_run_client.py
@@ -5,7 +5,7 @@
 
 @pytest.mark.e2e
 def test_run_client(services):
-    client = Client(timeout=10)
+    client = Client(timeout=20)
 
     # sanity check
     sessions = client.sync.core.sessions.list()
@@ -39,7 +39,7 @@ def test_run_client(services):
 @pytest.mark.e2e
 @pytest.mark.asyncio
 async def test_run_client_async(services):
-    client = Client(timeout=10)
+    client = Client(timeout=20)
 
     # test streaming
     session = await client.core.sessions.create()
diff --git a/tests/apiserver/data/.env b/tests/apiserver/data/.env
new file mode 100644
index 00000000..df512608
--- /dev/null
+++ b/tests/apiserver/data/.env
@@ -0,0 +1,1 @@
+API_KEY=123
diff --git a/tests/apiserver/data/env_variables.yaml b/tests/apiserver/data/env_variables.yaml
new file mode 100644
index 00000000..9e53565c
--- /dev/null
+++ b/tests/apiserver/data/env_variables.yaml
@@ -0,0 +1,20 @@
+name: MyDeployment
+
+control-plane:
+  port: 8000
+
+message-queue:
+  type: simple
+  host: "127.0.0.1"
+  port: 8001
+
+default-service: myworkflow
+
+services:
+  myworkflow:
+    name: My Python Workflow
+    env:
+      VAR_1: x
+      VAR_2: y
+    env-files:
+      - .env
diff --git a/tests/apiserver/data/example.yaml b/tests/apiserver/data/example.yaml
index f84b3003..7e73656a 100644
--- a/tests/apiserver/data/example.yaml
+++ b/tests/apiserver/data/example.yaml
@@ -26,6 +26,11 @@ services:
       # we can also support installing a req file relative to `path`
       # if source is a git repository
       - "requirements.txt"
+    env:
+      VAR_1: x
+      VAR_2: y
+    env-files:
+      - ./.env
 
   another-workflow:
     # A LITS workflow available in a git repo (might be the same)
diff --git a/tests/apiserver/data/workflow/__init__.py b/tests/apiserver/data/workflow/__init__.py
index 234cdc5a..5b5a2e59 100644
--- a/tests/apiserver/data/workflow/__init__.py
+++ b/tests/apiserver/data/workflow/__init__.py
@@ -1,3 +1,4 @@
-from .workflow_test import MyWorkflow
+from .workflow_test import MyWorkflow, _TestEnvWorkflow
 
 my_workflow = MyWorkflow()
+env_reader_workflow = _TestEnvWorkflow()
diff --git a/tests/apiserver/data/workflow/workflow_test.py b/tests/apiserver/data/workflow/workflow_test.py
index 90f7b596..f8343ecd 100644
--- a/tests/apiserver/data/workflow/workflow_test.py
+++ b/tests/apiserver/data/workflow/workflow_test.py
@@ -1,3 +1,4 @@
+import os
 from llama_index.core.workflow import Context, StartEvent, StopEvent, Workflow, step
 
 
@@ -5,3 +6,10 @@ class MyWorkflow(Workflow):
     @step
     def do_something(self, ctx: Context, ev: StartEvent) -> StopEvent:
         return StopEvent(result=f"Received: {ev.data}")
+
+
+class _TestEnvWorkflow(Workflow):
+    @step()
+    async def read_env_vars(self, ctx: Context, ev: StartEvent) -> StopEvent:
+        env_vars = [f"{v}: {os.environ.get(v)}" for v in ev.get("env_vars_to_read")]
+        return StopEvent(result=", ".join(env_vars))
diff --git a/tests/apiserver/test_config_parser.py b/tests/apiserver/test_config_parser.py
index 6cbcf185..fff27c6b 100644
--- a/tests/apiserver/test_config_parser.py
+++ b/tests/apiserver/test_config_parser.py
@@ -21,6 +21,8 @@ def do_assert(config: Config) -> None:
     assert wf_config.port == 1313
     assert wf_config.python_dependencies
     assert len(wf_config.python_dependencies) == 3
+    assert wf_config.env == {"VAR_1": "x", "VAR_2": "y"}
+    assert wf_config.env_files == ["./.env"]
 
     wf_config = config.services["another-workflow"]
     assert wf_config.name == "My LITS Workflow"
diff --git a/tests/apiserver/test_deployment.py b/tests/apiserver/test_deployment.py
index ad4a281f..64e71a12 100644
--- a/tests/apiserver/test_deployment.py
+++ b/tests/apiserver/test_deployment.py
@@ -136,6 +136,21 @@ def test__install_dependencies(data_path: Path) -> None:
         mocked_subp.check_call.assert_not_called()
 
 
+def test__set_environment_variables(data_path: Path) -> None:
+    config = Config.from_yaml(data_path / "env_variables.yaml")
+    service_config = config.services["myworkflow"]
+    with mock.patch("llama_deploy.apiserver.deployment.os.environ") as mocked_osenviron:
+        # Assert the sub process cmd receives the list of dependencies
+        Deployment._set_environment_variables(service_config, root=data_path)
+        mocked_osenviron.__setitem__.assert_has_calls(
+            [
+                mock.call("VAR_1", "x"),
+                mock.call("VAR_2", "y"),
+                mock.call("API_KEY", "123"),
+            ]
+        )
+
+
 def test__install_dependencies_raises(data_path: Path) -> None:
     config = Config.from_yaml(data_path / "python_dependencies.yaml")
     service_config = config.services["myworkflow"]