Merge pull request #56 from fal-ai/batuhan/fea-609-support-agent-environment-definitions-in

feat: Support agent environment definitions in the server

Author: isidentical

src/isolate/server/server.py

@@ -12,7 +12,9 @@
 from grpc import ServicerContext, StatusCode
 
 from isolate.backends import EnvironmentCreationError, IsolateSettings
+from isolate.backends.common import active_python
 from isolate.backends.local import LocalPythonEnvironment
+from isolate.backends.virtualenv import VirtualPythonEnvironment
 from isolate.connections.grpc import AgentError, LocalPythonGRPC
 from isolate.logs import Log, LogLevel, LogSource
 from isolate.server import definitions
@@ -23,6 +25,14 @@
 
 # Number of threads that the gRPC server will use.
 MAX_THREADS = int(os.getenv("MAX_THREADS", 5))
+_AGENT_REQUIREMENTS_TXT = os.getenv("AGENT_REQUIREMENTS_TXT")
+
+if _AGENT_REQUIREMENTS_TXT is not None:
+    with open(_AGENT_REQUIREMENTS_TXT) as stream:
+        AGENT_REQUIREMENTS = stream.read().splitlines()
+else:
+    AGENT_REQUIREMENTS = []
+
 
 # Number of seconds to observe the queue before checking the termination
 # event.
@@ -69,6 +79,18 @@ def Run(
         for environment in environments:
             environment.apply_settings(run_settings)
 
+        primary_environment = environments[0]
+
+        if AGENT_REQUIREMENTS:
+            python_version = getattr(
+                primary_environment, "python_version", active_python()
+            )
+            agent_environ = VirtualPythonEnvironment(
+                requirements=AGENT_REQUIREMENTS,
+                python_version=python_version,
+            )
+            environments.insert(1, agent_environ)
+
         extra_inheritance_paths = []
         if INHERIT_FROM_LOCAL:
             local_environment = LocalPythonEnvironment()

tests/test_server.py

@@ -259,3 +259,33 @@ def test_server_multiple_envs(
     raw_result = run_request(stub, request)
 
     assert from_grpc(raw_result) == "0.6.0 2.8.2"
+
+
+@pytest.mark.parametrize("python_version", ["3.8"])
+def test_agent_requirements_custom_version(
+    stub: definitions.IsolateStub,
+    monkeypatch: Any,
+    python_version: str,
+) -> None:
+    requirements = ["pyjokes==0.6.0"]
+    agent_requirements = ["dill==0.3.5.1", f"{REPO_DIR}[grpc]"]
+    monkeypatch.setattr("isolate.server.server.AGENT_REQUIREMENTS", agent_requirements)
+
+    env_definition = define_environment(
+        "virtualenv",
+        requirements=requirements,
+        python_version=python_version,
+    )
+    request = definitions.BoundFunction(
+        function=to_serialized_object(
+            partial(
+                eval,
+                "__import__('sysconfig').get_python_version(), __import__('pyjokes').__version__",
+            ),
+            method="dill",
+        ),
+        environments=[env_definition],
+    )
+
+    raw_result = run_request(stub, request)
+    assert from_grpc(raw_result) == ("3.8", "0.6.0")