feat: bundle local (non-pip) modules for serverless endpoints

src/runpod_flash/cli/commands/build.py

@@ -17,7 +17,9 @@
 from rich.console import Console
 
 from runpod_flash.cli.utils.formatting import print_error, print_warning
+from runpod_flash.core.exceptions import LocalModuleResolutionError
 from runpod_flash.core.resources.constants import MAX_TARBALL_SIZE_MB
+from runpod_flash.stubs.local_modules import resolve_local_modules
 
 from ..utils.ignore import get_file_tree, load_ignore_patterns
 from .build_utils.handler_generator import HandlerGenerator
@@ -30,6 +32,83 @@
 
 console = Console()
 
+# Decorator names that mark a function/class as a Flash endpoint entry point.
+ENDPOINT_DECORATOR_NAMES: frozenset[str] = frozenset({"remote", "Endpoint"})
+
+
+def _defines_endpoint(py_file: Path) -> bool:
+    """Check whether *py_file* defines a function/class decorated as a Flash endpoint.
+
+    Recognizes ``@remote``, ``@remote(...)``, ``@Endpoint(...)``, and attribute
+    forms (e.g. ``@rf.Endpoint``). Parses the file as AST only -- it does not
+    import it, so decorators do not need to be resolvable.
+
+    Returns:
+        True if an endpoint decorator is found, False if none is found or the
+        file cannot be read/parsed.
+    """
+    try:
+        tree = ast.parse(py_file.read_text(encoding="utf-8"))
+    except (OSError, UnicodeDecodeError, SyntaxError):
+        return False
+
+    for node in ast.walk(tree):
+        if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
+            continue
+        for decorator in node.decorator_list:
+            target = decorator.func if isinstance(decorator, ast.Call) else decorator
+            if isinstance(target, ast.Attribute):
+                name = target.attr
+            elif isinstance(target, ast.Name):
+                name = target.id
+            else:
+                continue
+            if name in ENDPOINT_DECORATOR_NAMES:
+                return True
+
+    return False
+
+
+def augment_with_local_modules(files: list[Path], project_dir: Path) -> list[Path]:
+    """Force-include local modules imported by project code but dropped by ignores.
+
+    Walks the import closure of every non-ignored ``.py`` file in *files* and adds
+    any backing local file under *project_dir* that is not already present.
+
+    Strictness is scoped to endpoint files: if a file defines a ``@remote``/
+    ``@Endpoint`` entry point and its import closure can't be resolved (broken
+    relative import, non-UTF-8 bytes, syntax error), the build fails loudly via
+    ``LocalModuleResolutionError`` -- shipping it would produce a broken
+    tarball. Incidental (non-endpoint) files with the same problems are skipped
+    with a warning instead, matching pre-existing behavior of shipping them
+    untouched.
+    """
+    project_dir = project_dir.resolve()
+    present = {f.resolve() for f in files}
+    extra: dict[Path, None] = {}
+
+    for py_file in [f for f in files if f.suffix == ".py"]:
+        try:
+            resolved = resolve_local_modules(
+                py_file.read_text(encoding="utf-8"), py_file, project_dir
+            )
+        except (LocalModuleResolutionError, UnicodeDecodeError, SyntaxError) as e:
+            if _defines_endpoint(py_file):
+                raise
+            print_warning(
+                console, f"Skipping local-module resolution for {py_file}: {e}"
+            )
+            continue
+
+        for warning in resolved.warnings:
+            print_warning(console, warning)
+        for abs_path in resolved.files.values():
+            p = Path(abs_path).resolve()
+            if p not in present:
+                extra[p] = None
+
+    return list(files) + list(extra)
+
 
 def compute_source_fingerprint(project_dir: Path, files: list[Path]) -> str:
     """Compute a SHA-256 fingerprint of project source files.
@@ -261,6 +340,11 @@ def run_build(
 
     spec = load_ignore_patterns(project_dir)
     files = get_file_tree(project_dir, spec)
+    try:
+        files = augment_with_local_modules(files, project_dir)
+    except LocalModuleResolutionError as e:
+        print_error(console, str(e))
+        raise typer.Exit(1)
 
     # Resolved later by ManifestBuilder from resource configs (or the override
     # above). Pip wheel selection re-reads this via _resolve_pip_python_version.

src/runpod_flash/core/exceptions.py

@@ -39,3 +39,21 @@ def _default_message() -> str:
             "\n"
             "Get a key: https://docs.runpod.io/get-started/api-keys"
         )
+
+
+class LocalModuleResolutionError(Exception):
+    """Raised when an imported local module cannot be resolved to a bundle-able file.
+
+    Covers relative imports that do not resolve, package submodules whose file is
+    missing, and local modules that live outside the project root. External names
+    (stdlib / installed pip packages) are NOT errors — they are left to the worker
+    image.
+    """
+
+
+class LocalModulePayloadTooLargeError(Exception):
+    """Raised when inline local-module source exceeds the live-serverless size cap.
+
+    The live path ships module source inside the request payload; past the cap the
+    user should switch to ``flash deploy``.
+    """

src/runpod_flash/protos/remote_execution.py

@@ -43,6 +43,15 @@ class FunctionRequest(BaseModel):
         default=None,
         description="Optional list of system dependencies to install before executing the function",
     )
+    modules: Dict[str, str] = Field(
+        default_factory=dict,
+        description=(
+            "Local (non-pip) module source files to make importable on the worker, "
+            "keyed by POSIX relative path (e.g. 'utils.py', 'helpers/__init__.py'). "
+            "Written to a temp dir on sys.path before the function code is exec'd. "
+            "Additive and backward compatible: absent/empty means today's behavior."
+        ),
+    )
 
     # NEW FIELDS FOR CLASS SUPPORT
     execution_type: str = Field(

src/runpod_flash/runtime/lb_handler.py

@@ -25,6 +25,8 @@
 from fastapi import FastAPI, File, Form, Request
 from pydantic import BaseModel, create_model
 
+from runpod_flash.runtime.module_loader import materialized_modules
+
 logger = logging.getLogger(__name__)
 
 _BODY_METHODS = frozenset({"POST", "PUT", "PATCH", "DELETE"})
@@ -320,10 +322,46 @@ async def execute_remote_function(request: Request) -> Dict[str, Any]:
                         "error": f"Failed to deserialize arguments: {e}",
                     }
 
-                # Execute function in isolated namespace
+                # Execute function in isolated namespace.
+                #
+                # NOTE: shipped local modules must stay importable through the
+                # function CALL, not just through exec(). Endpoint functions
+                # frequently import local siblings inside their body (i.e.
+                # def-now/call-later), so `materialized_modules` has to remain
+                # active until the function has actually run (including
+                # `await` for async functions). Exiting the context manager
+                # right after exec() (and before the call) would drop the temp
+                # dir from sys.path before those in-body imports execute,
+                # causing a spurious ModuleNotFoundError for otherwise-correct
+                # code. The with-block is therefore widened to also cover the
+                # function lookup and the call itself.
                 namespace: Dict[str, Any] = {}
                 try:
-                    exec(function_code, namespace)
+                    with materialized_modules(body.get("modules", {}) or {}):
+                        exec(function_code, namespace)
+
+                        # Get function from namespace
+                        if function_name not in namespace:
+                            return {
+                                "success": False,
+                                "error": f"Function '{function_name}' not found in executed code",
+                            }
+
+                        func = namespace[function_name]
+
+                        # Execute function
+                        try:
+                            result = func(*args, **kwargs)
+
+                            # Handle async functions
+                            if inspect.iscoroutine(result):
+                                result = await result
+                        except Exception as e:
+                            logger.error(f"Function execution failed: {e}")
+                            return {
+                                "success": False,
+                                "error": f"Function execution failed: {e}",
+                            }
                 except SyntaxError as e:
                     logger.error(f"Syntax error in function code: {e}")
                     return {
@@ -337,29 +375,6 @@ async def execute_remote_function(request: Request) -> Dict[str, Any]:
                         "error": f"Error executing function code: {e}",
                     }
 
-                # Get function from namespace
-                if function_name not in namespace:
-                    return {
-                        "success": False,
-                        "error": f"Function '{function_name}' not found in executed code",
-                    }
-
-                func = namespace[function_name]
-
-                # Execute function
-                try:
-                    result = func(*args, **kwargs)
-
-                    # Handle async functions
-                    if inspect.iscoroutine(result):
-                        result = await result
-                except Exception as e:
-                    logger.error(f"Function execution failed: {e}")
-                    return {
-                        "success": False,
-                        "error": f"Function execution failed: {e}",
-                    }
-
                 # Serialize result
                 try:
                     result_b64 = serialize_arg(result)

src/runpod_flash/runtime/module_loader.py

@@ -0,0 +1,61 @@
+"""Materialize inline local-module source onto the worker's import path.
+
+The live-serverless path ships local module files in ``FunctionRequest.modules``.
+Before the worker ``exec``s the function code, those files must exist on disk and
+be importable. This context manager writes them to a temp dir, prepends it to
+``sys.path``, and cleans up afterward.
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+import tempfile
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+
+log = logging.getLogger(__name__)
+
+
+@contextmanager
+def materialized_modules(modules: dict[str, str]) -> Iterator[str | None]:
+    """Write *modules* to a temp dir on ``sys.path`` for the duration of the block.
+
+    Args:
+        modules: POSIX relative path -> module source text. Empty means no-op.
+
+    Yields:
+        The temp dir path added to ``sys.path``, or ``None`` when *modules* is empty.
+
+    Concurrency note: this mutates the process-global ``sys.path``. It assumes one
+    function executes at a time per worker process (the current worker model). If a
+    worker runs multiple handler invocations concurrently (e.g. async concurrency >
+    1), the inserted temp dir is visible to other in-flight invocations and cleanup
+    on exit could remove files mid-import. Isolating sys.path per invocation is a
+    follow-up if concurrent execution is enabled.
+    """
+    if not modules:
+        yield None
+        return
+
+    tmpdir = tempfile.mkdtemp(prefix="flash_modules_")
+    root = Path(tmpdir)
+    for rel_path, source in modules.items():
+        dest = root / rel_path
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        dest.write_text(source, encoding="utf-8")
+
+    sys.path.insert(0, tmpdir)
+    try:
+        yield tmpdir
+    finally:
+        try:
+            sys.path.remove(tmpdir)
+        except ValueError:
+            log.warning(
+                "flash module temp dir %s already removed from sys.path", tmpdir
+            )
+        import shutil
+
+        shutil.rmtree(tmpdir, ignore_errors=True)

src/runpod_flash/stubs/live_serverless.py

@@ -7,13 +7,17 @@
 import threading
 import cloudpickle
 import logging
+from pathlib import Path
+
+from ..core.exceptions import LocalModulePayloadTooLargeError
 from ..core.resources import LiveServerless
 from ..protos.remote_execution import (
     FunctionRequest,
     FunctionResponse,
     RemoteExecutorStub,
 )
 from ..runtime.serialization import serialize_args, serialize_kwargs
+from .local_modules import MAX_INLINE_MODULE_BYTES, resolve_local_modules
 
 log = logging.getLogger(__name__)
 
@@ -66,6 +70,46 @@ def get_function_source(func):
     return function_source, source_hash
 
 
+def build_modules_map(func_source: str, source_file: str | None) -> dict[str, str]:
+    """Collect the endpoint's local-module source for inline shipping.
+
+    Resolves the transitive local-import closure of *func_source* against the
+    source file's directory and returns {relative_path: source_text}. Enforces the
+    inline size cap.
+
+    Args:
+        func_source: Extracted function source (from ``get_function_source``).
+        source_file: Absolute path of the file the function was defined in, or
+            ``None`` (e.g. functions with no ``__file__``), in which case nothing
+            is bundled.
+
+    Raises:
+        LocalModulePayloadTooLargeError: total inline source exceeds the cap.
+    """
+    if not source_file:
+        return {}
+
+    project_root = Path(source_file).resolve().parent
+    resolved = resolve_local_modules(func_source, source_file, project_root)
+    for warning in resolved.warnings:
+        log.warning(warning)
+
+    modules: dict[str, str] = {}
+    total = len(func_source.encode("utf-8"))
+    for rel_path, abs_path in resolved.files.items():
+        source = Path(abs_path).read_text(encoding="utf-8")
+        total += len(source.encode("utf-8"))
+        modules[rel_path] = source
+
+    if total > MAX_INLINE_MODULE_BYTES:
+        raise LocalModulePayloadTooLargeError(
+            f"Inline module payload is {total} bytes, over the "
+            f"{MAX_INLINE_MODULE_BYTES}-byte live-serverless cap. "
+            f"Use `flash deploy` for endpoints with large local dependencies."
+        )
+    return modules
+
+
 class LiveServerlessStub(RemoteExecutorStub):
     """Adapter class to make Runpod endpoints look like gRPC stubs."""
 
@@ -122,6 +166,9 @@ async def prepare_request(
 
             request["function_code"] = _SERIALIZED_FUNCTION_CACHE[src_hash]
 
+        source_file = original_func.__globals__.get("__file__")
+        request["modules"] = build_modules_map(source, source_file)
+
         # Serialize arguments using cloudpickle
         if args:
             request["args"] = serialize_args(args)

src/runpod_flash/stubs/load_balancer_sls.py

@@ -18,7 +18,7 @@
 )
 from runpod_flash.core.resources.constants import DEFAULT_LB_STUB_TIMEOUT
 
-from .live_serverless import get_function_source
+from .live_serverless import build_modules_map, get_function_source
 
 log = logging.getLogger(__name__)
 
@@ -238,6 +238,9 @@ async def _prepare_request(
             "accelerate_downloads": accelerate_downloads,
         }
 
+        source_file = original_func.__globals__.get("__file__")
+        request["modules"] = build_modules_map(source, source_file)
+
         # Serialize arguments using cloudpickle + base64
         if args:
             request["args"] = serialize_args(args)