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

src/runpod_flash/cli/commands/build.py

@@ -17,20 +17,96 @@
 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
 from .build_utils.lb_handler_generator import LBHandlerGenerator
 from .build_utils.manifest import ManifestBuilder
 from .build_utils.resource_config_generator import generate_all_resource_configs
-from .build_utils.scanner import RuntimeScanner
+from .build_utils.scanner import RuntimeScanner, defines_endpoint
 
 logger = logging.getLogger(__name__)
 
 console = Console()
 
 
+def validate_local_module_imports(files: list[Path], project_dir: Path) -> None:
+    """Fail the build when shipped code imports a local module the ignores dropped.
+
+    Walks the import closure of every shipped ``.py`` file in *files*. A local
+    import that resolves to a file under *project_dir* which is not itself among
+    *files* was excluded by the ignore rules (``.gitignore`` or the built-in
+    defaults). Force-including it would silently override a deliberate exclusion;
+    omitting it would break the worker with ``ModuleNotFoundError``. So the build
+    is refused with an actionable error naming the excluded file and its importer.
+
+    Strictness for *unresolvable* imports (broken relative import, non-UTF-8
+    bytes, syntax error) is scoped to endpoint files: an ``@remote``/``@Endpoint``
+    entry point fails the build loudly via ``LocalModuleResolutionError`` --
+    shipping it would produce a broken tarball -- while an incidental
+    (non-endpoint) file is skipped with a warning, matching pre-existing behavior
+    of shipping such files untouched.
+
+    Raises:
+        LocalModuleResolutionError: an endpoint's import closure cannot be
+            resolved, or any shipped file imports a local module the ignore rules
+            excluded from the build.
+    """
+    project_dir = project_dir.resolve()
+    present = {f.resolve() for f in files}
+    # excluded module file -> the shipped file that imported it
+    excluded: dict[Path, Path] = {}
+
+    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):
+                # run_build() only catches LocalModuleResolutionError to emit a
+                # clean error. Syntax errors already arrive wrapped (see
+                # local_modules._walk) and re-raise directly; a raw
+                # UnicodeDecodeError from a non-UTF-8 dependency file is
+                # normalized here so an endpoint build fails loudly rather than
+                # surfacing a raw traceback.
+                if isinstance(e, LocalModuleResolutionError):
+                    raise
+                raise LocalModuleResolutionError(
+                    f"Cannot resolve local imports for endpoint file {py_file}: {e}"
+                ) from e
+            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:
+                excluded.setdefault(p, py_file.resolve())
+
+    if excluded:
+        listing = "\n".join(
+            f"  {p.relative_to(project_dir)} (imported by "
+            f"{importer.relative_to(project_dir)})"
+            for p, importer in sorted(excluded.items())
+        )
+        raise LocalModuleResolutionError(
+            "Shipped code imports local modules that the build ignore rules "
+            "(.gitignore or built-in defaults) exclude:\n"
+            f"{listing}\n\n"
+            "Shipping them would silently override a deliberate exclusion, and "
+            "omitting them would break the worker with ModuleNotFoundError. Remove "
+            "the matching ignore pattern or stop importing these modules from "
+            "shipped code."
+        )
+
+
 def compute_source_fingerprint(project_dir: Path, files: list[Path]) -> str:
     """Compute a SHA-256 fingerprint of project source files.
 
@@ -261,6 +337,11 @@ def run_build(
 
     spec = load_ignore_patterns(project_dir)
     files = get_file_tree(project_dir, spec)
+    try:
+        validate_local_module_imports(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/cli/commands/build_utils/scanner.py

@@ -50,6 +50,49 @@ def file_to_module_path(file_path: Path, project_root: Path) -> str:
     return str(rel).replace(os.sep, ".").replace("/", ".")
 
 
+# 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.
+
+    This is the AST-only counterpart to the scanner's live discovery
+    (``discover_remote_functions``): the scanner imports modules and inspects
+    stamped ``__remote_config__`` objects, which requires the imports to
+    succeed. This helper is consulted precisely when imports may be broken (the
+    pre-copy local-module validation), so it cannot go through the import path.
+
+    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
+
+
 @dataclass
 class RemoteFunctionMetadata:
     """Metadata about a @remote decorated function or class."""

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,71 @@
+"""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 shutil
+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).resolve()
+    inserted = False
+    try:
+        for rel_path, source in modules.items():
+            # ``modules`` is untrusted request-body input: an absolute path or
+            # ``..`` segments in rel_path would escape ``root`` and let a caller
+            # write anywhere on disk. Fail secure by rejecting any path that does
+            # not resolve to a location strictly under the temp dir.
+            dest = (root / rel_path).resolve()
+            if not dest.is_relative_to(root):
+                raise ValueError(
+                    f"module path escapes materialization dir: {rel_path!r}"
+                )
+            dest.parent.mkdir(parents=True, exist_ok=True)
+            dest.write_text(source, encoding="utf-8")
+
+        sys.path.insert(0, tmpdir)
+        inserted = True
+        yield tmpdir
+    finally:
+        if inserted:
+            try:
+                sys.path.remove(tmpdir)
+            except ValueError:
+                log.warning(
+                    "flash module temp dir %s already removed from sys.path", tmpdir
+                )
+        shutil.rmtree(tmpdir, ignore_errors=True)