feat: SLS-360 import shipped local modules before executing user code

pyproject.toml

@@ -14,7 +14,11 @@ dependencies = [
     "huggingface_hub>=0.32.0",
     "fastapi>=0.115.0",
     "uvicorn[standard]>=0.34.0",
-    "runpod-flash>=1.8.0",
+    # TEMPORARY (SLS-360 coordinated release): this worker uses
+    # runpod_flash.runtime.module_loader, which is not in a published runpod-flash
+    # yet. Pin to the flash branch until runpod/flash#352 is merged and released,
+    # then REVERT to a released constraint (e.g. "runpod-flash>=1.18.0").
+    "runpod-flash @ git+https://github.com/runpod/flash.git@deanquinanola/sls-360-local-module-bundling",
 ]
 
 [dependency-groups]

src/function_executor.py

@@ -6,6 +6,7 @@
 from typing import Dict, Any
 
 from runpod_flash.protos.remote_execution import FunctionRequest, FunctionResponse
+from runpod_flash.runtime.module_loader import materialized_modules
 from serialization_utils import SerializationUtils
 
 
@@ -34,30 +35,36 @@ async def execute(self, request: FunctionRequest) -> FunctionResponse:
             logger.addHandler(log_handler)
 
             try:
-                # Execute function code in namespace
+                # Execute function code in namespace, with any shipped local
+                # modules importable on sys.path for the duration of the exec
+                # AND the subsequent call. Local modules are typically imported
+                # lazily inside the function body (executed at call time, not
+                # at def-time), so the materialized path must stay live until
+                # the function has actually run.
                 namespace: Dict[str, Any] = {}
-                if request.function_code:
-                    exec(request.function_code, namespace)
-
-                if request.function_name not in namespace:
-                    return FunctionResponse(
-                        success=False,
-                        result=f"Function '{request.function_name}' not found in the provided code",
-                    )
-
-                func = namespace[request.function_name]
-
-                # Deserialize arguments
-                args = SerializationUtils.deserialize_args(request.args)
-                kwargs = SerializationUtils.deserialize_kwargs(request.kwargs)
-
-                # Execute the function (handle both sync and async)
-                if inspect.iscoroutinefunction(func):
-                    # Async function - await directly
-                    result = await func(*args, **kwargs)
-                else:
-                    # Sync function - call directly
-                    result = func(*args, **kwargs)
+                with materialized_modules(getattr(request, "modules", {}) or {}):
+                    if request.function_code:
+                        exec(request.function_code, namespace)
+
+                    if request.function_name not in namespace:
+                        return FunctionResponse(
+                            success=False,
+                            result=f"Function '{request.function_name}' not found in the provided code",
+                        )
+
+                    func = namespace[request.function_name]
+
+                    # Deserialize arguments
+                    args = SerializationUtils.deserialize_args(request.args)
+                    kwargs = SerializationUtils.deserialize_kwargs(request.kwargs)
+
+                    # Execute the function (handle both sync and async)
+                    if inspect.iscoroutinefunction(func):
+                        # Async function - await directly
+                        result = await func(*args, **kwargs)
+                    else:
+                        # Sync function - call directly
+                        result = func(*args, **kwargs)
 
             except Exception as e:
                 # Combine output streams

tests/unit/test_function_executor.py

@@ -215,6 +215,49 @@ async def gpu_matrix_multiply(input_data: dict) -> dict:
         assert result["result_shape"] == [500, 500]
 
 
+class TestModuleImports:
+    """Test that shipped local modules are importable during function execution."""
+
+    def setup_method(self):
+        """Setup for each test method."""
+        self.executor = FunctionExecutor()
+
+    async def test_execute_can_import_shipped_local_module(self):
+        """Test that a shipped local module can be imported and used."""
+        request = FunctionRequest(
+            function_name="handler",
+            function_code=(
+                "def handler():\n"
+                "    import flash_shipped_greeting\n"
+                "    return flash_shipped_greeting.x()\n"
+            ),
+            args=[],
+            kwargs={},
+            modules={"flash_shipped_greeting.py": "def x():\n    return 7\n"},
+        )
+
+        response = await self.executor.execute(request)
+
+        assert response.success is True
+        result = cloudpickle.loads(base64.b64decode(response.result))
+        assert result == 7
+
+    async def test_execute_without_modules_still_works(self):
+        """Test that execution without any shipped modules still works."""
+        request = FunctionRequest(
+            function_name="handler",
+            function_code="def handler():\n    return 1\n",
+            args=[],
+            kwargs={},
+        )
+
+        response = await self.executor.execute(request)
+
+        assert response.success is True
+        result = cloudpickle.loads(base64.b64decode(response.result))
+        assert result == 1
+
+
 class TestErrorHandling:
     """Test error handling in function execution."""