diff --git a/src/runpod_flash/cli/commands/build.py b/src/runpod_flash/cli/commands/build.py index 65414a97..48042be5 100644 --- a/src/runpod_flash/cli/commands/build.py +++ b/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,93 @@ 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): + # 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: + 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 +350,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. diff --git a/src/runpod_flash/core/exceptions.py b/src/runpod_flash/core/exceptions.py index 2de78944..b2446ee8 100644 --- a/src/runpod_flash/core/exceptions.py +++ b/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``. + """ diff --git a/src/runpod_flash/protos/remote_execution.py b/src/runpod_flash/protos/remote_execution.py index e86626c7..3d10b7b5 100644 --- a/src/runpod_flash/protos/remote_execution.py +++ b/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( diff --git a/src/runpod_flash/runtime/lb_handler.py b/src/runpod_flash/runtime/lb_handler.py index 22c930b1..86445b4f 100644 --- a/src/runpod_flash/runtime/lb_handler.py +++ b/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) diff --git a/src/runpod_flash/runtime/module_loader.py b/src/runpod_flash/runtime/module_loader.py new file mode 100644 index 00000000..8aea7aa8 --- /dev/null +++ b/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) diff --git a/src/runpod_flash/stubs/live_serverless.py b/src/runpod_flash/stubs/live_serverless.py index fa6b7552..85a6c060 100644 --- a/src/runpod_flash/stubs/live_serverless.py +++ b/src/runpod_flash/stubs/live_serverless.py @@ -7,6 +7,9 @@ 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, @@ -14,6 +17,7 @@ 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) diff --git a/src/runpod_flash/stubs/load_balancer_sls.py b/src/runpod_flash/stubs/load_balancer_sls.py index bf97c2c4..cf81b901 100644 --- a/src/runpod_flash/stubs/load_balancer_sls.py +++ b/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) diff --git a/src/runpod_flash/stubs/local_modules.py b/src/runpod_flash/stubs/local_modules.py new file mode 100644 index 00000000..bd6ef692 --- /dev/null +++ b/src/runpod_flash/stubs/local_modules.py @@ -0,0 +1,257 @@ +"""Discover the transitive closure of local (non-pip, non-stdlib) modules an +endpoint imports, so they can be bundled at deploy time or shipped inline for +live-serverless execution. + +A "local" module is a ``.py`` file or package directory under the project root. +Anything that resolves to the standard library or an installed site-package is +treated as external and left to pip / the worker image — never bundled, never an +error. +""" + +from __future__ import annotations + +import ast +import logging +import sys +from dataclasses import dataclass, field +from pathlib import Path + +from ..core.exceptions import LocalModuleResolutionError + +log = logging.getLogger(__name__) + +# Total inline module payload cap for the live path. Runpod /run payloads are +# limited (~10 MB); stay well under and fail loudly past this. +MAX_INLINE_MODULE_BYTES = 8 * 1024 * 1024 # 8 MiB + + +@dataclass +class ResolvedModules: + """Result of a local-module resolution. + + files: POSIX relative path -> absolute path, for every local file to ship. + warnings: human-readable notes (e.g. unresolvable dynamic imports). + """ + + files: dict[str, str] = field(default_factory=dict) + warnings: list[str] = field(default_factory=list) + + +def _top_level(name: str) -> str: + return name.split(".", 1)[0] + + +def is_stdlib(name: str) -> bool: + """True when *name*'s top-level package is part of the standard library.""" + return _top_level(name) in sys.stdlib_module_names + + +def local_path_for(dotted: str, search_dirs: list[Path]) -> Path | None: + """Return the file backing a dotted module name under any search dir. + + Checks the module file (``a/b.py``) then the package init (``a/b/__init__.py``) + in each search dir in order. Returns the first match, or ``None`` when the name + does not correspond to a local file (i.e. it is external). + """ + rel = Path(*dotted.split(".")) + for base in search_dirs: + module_file = base / rel.with_suffix(".py") + if module_file.is_file(): + return module_file + package_init = base / rel / "__init__.py" + if package_init.is_file(): + return package_init + return None + + +def resolve_local_modules( + source: str, source_file: str | Path, project_root: str | Path +) -> ResolvedModules: + """Resolve the transitive set of local module files imported by *source*. + + Args: + source: Python source to scan (a function body or a whole module). + source_file: File *source* came from; its directory is the base for + sibling and relative-import resolution. + project_root: Only files under this root are considered local; included + paths are made relative to it. For the live path this is the source + file's directory; for the build path it is the Flash project dir. + + Returns: + ResolvedModules with the file map (POSIX relative path -> absolute path) + and any warnings. + + Raises: + LocalModuleResolutionError: a relative import, a local package submodule, + or a local file outside the project root could not be handled. + + Note (live-serverless path): when *project_root* is the endpoint source file's + directory, a module imported by ABSOLUTE name from a PARENT directory (e.g. + `import shared` where `shared.py` lives above the endpoint) is classified as + external and silently omitted — it will raise ModuleNotFoundError on the worker. + Only relative imports fail loudly. Keep an endpoint at or above its local + dependencies, or use `flash deploy` (which uses the project directory as the + root). Tracked as a follow-up. + """ + root = Path(project_root).resolve() + start_dir = Path(source_file).resolve().parent + result = ResolvedModules() + visited: set[Path] = set() + _walk(source, start_dir, root, result, visited, origin=str(source_file)) + return result + + +def _walk( + source: str, + current_dir: Path, + root: Path, + result: ResolvedModules, + visited: set[Path], + origin: str, +) -> None: + try: + tree = ast.parse(source) + except SyntaxError as exc: + raise LocalModuleResolutionError( + f"Could not parse {origin} while resolving local imports: {exc}" + ) from exc + + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + _consider(alias.name, 0, current_dir, root, result, visited, origin) + elif isinstance(node, ast.ImportFrom): + if node.level > 0 and node.module is None: + # `from . import a, b` — each name is a candidate local submodule. + # `from . import *` carries the sentinel name "*", which is not a + # module; fall back to resolving the package __init__ itself. + names = [alias.name for alias in node.names if alias.name != "*"] + if names: + for name in names: + _consider( + name, + node.level, + current_dir, + root, + result, + visited, + origin, + ) + else: + _consider( + "", + node.level, + current_dir, + root, + result, + visited, + origin, + ) + else: + _consider( + node.module or "", + node.level, + current_dir, + root, + result, + visited, + origin, + ) + elif isinstance(node, ast.Call): + _consider_dynamic(node, current_dir, root, result, visited, origin) + + +def _consider( + dotted: str, + level: int, + current_dir: Path, + root: Path, + result: ResolvedModules, + visited: set[Path], + origin: str, +) -> None: + if level == 0: + if not dotted or is_stdlib(dotted): + return + path = local_path_for(dotted, [current_dir, root]) + if path is None: + return # external (pip / worker image); not our concern + else: + base = current_dir + for _ in range(level - 1): + base = base.parent + if dotted: + path = local_path_for(dotted, [base]) + else: + candidate = base / "__init__.py" + path = candidate if candidate.is_file() else None + if path is None: + raise LocalModuleResolutionError( + f"{origin}: relative import (level={level}, module={dotted!r}) " + f"could not be resolved to a local file under {base}" + ) + _include(path, root, result, visited) + + +def _consider_dynamic( + node: ast.Call, + current_dir: Path, + root: Path, + result: ResolvedModules, + visited: set[Path], + origin: str, +) -> None: + func = node.func + is_import_module = isinstance(func, ast.Attribute) and func.attr == "import_module" + is_dunder_import = isinstance(func, ast.Name) and func.id == "__import__" + if not (is_import_module or is_dunder_import): + return + if ( + node.args + and isinstance(node.args[0], ast.Constant) + and isinstance(node.args[0].value, str) + ): + _consider(node.args[0].value, 0, current_dir, root, result, visited, origin) + else: + lineno = getattr(node, "lineno", "?") + result.warnings.append( + f"{origin}:{lineno}: unresolvable dynamic import (non-literal argument); " + f"ensure the module is on the worker image or bundle it explicitly" + ) + + +def _include( + path: Path, root: Path, result: ResolvedModules, visited: set[Path] +) -> None: + path = path.resolve() + if path in visited: + return + try: + rel = path.relative_to(root) + except ValueError as exc: + raise LocalModuleResolutionError( + f"Imported local module {path} lives outside the project root {root} " + f"and cannot be included; move it under the project or use `flash deploy`." + ) from exc + visited.add(path) + result.files[rel.as_posix()] = str(path) + _include_ancestor_inits(path, root, result, visited) + _walk( + path.read_text(encoding="utf-8"), + path.parent, + root, + result, + visited, + origin=str(path), + ) + + +def _include_ancestor_inits( + path: Path, root: Path, result: ResolvedModules, visited: set[Path] +) -> None: + parent = path.parent + while parent != root and root in parent.parents: + init = parent / "__init__.py" + if init.is_file() and init.resolve() not in visited: + _include(init, root, result, visited) + parent = parent.parent diff --git a/tests/unit/_lb_prepare_request_sibling.py b/tests/unit/_lb_prepare_request_sibling.py new file mode 100644 index 00000000..bd61564c --- /dev/null +++ b/tests/unit/_lb_prepare_request_sibling.py @@ -0,0 +1,6 @@ +"""Sibling module used by test_load_balancer_sls_stub.py to exercise local-module +resolution through LoadBalancerSlsStub._prepare_request. Not a test module itself.""" + + +def sibling_value() -> int: + return 42 diff --git a/tests/unit/_live_serverless_prepare_request_sibling.py b/tests/unit/_live_serverless_prepare_request_sibling.py new file mode 100644 index 00000000..972c71b7 --- /dev/null +++ b/tests/unit/_live_serverless_prepare_request_sibling.py @@ -0,0 +1,6 @@ +"""Sibling module used by test_stub_live_serverless.py to exercise local-module +resolution through LiveServerlessStub.prepare_request. Not a test module itself.""" + + +def sibling_value() -> int: + return 99 diff --git a/tests/unit/cli/commands/test_build_local_modules.py b/tests/unit/cli/commands/test_build_local_modules.py new file mode 100644 index 00000000..52ea253f --- /dev/null +++ b/tests/unit/cli/commands/test_build_local_modules.py @@ -0,0 +1,103 @@ +from pathlib import Path + +import pytest + +from runpod_flash.cli.commands.build import ( + _defines_endpoint, + augment_with_local_modules, +) +from runpod_flash.core.exceptions import LocalModuleResolutionError + + +def test_force_includes_module_dropped_by_ignore_rule(tmp_path: Path): + # endpoint imports a module whose name matches the built-in test_*.py ignore rule + (tmp_path / "test_helpers.py").write_text("def h():\n return 1\n") + endpoint = tmp_path / "endpoint.py" + endpoint.write_text( + "import test_helpers\n\ndef handler():\n return test_helpers.h()\n" + ) + + # simulate the ignore filter having dropped test_helpers.py + files = [endpoint] + augmented = augment_with_local_modules(files, tmp_path) + + assert (tmp_path / "test_helpers.py") in augmented + assert endpoint in augmented + + +def test_leaves_files_untouched_when_no_local_imports(tmp_path: Path): + endpoint = tmp_path / "endpoint.py" + endpoint.write_text("import os\n\ndef handler():\n return os.getpid()\n") + files = [endpoint] + assert augment_with_local_modules(files, tmp_path) == [endpoint] + + +def test_endpoint_file_with_unresolvable_import_raises(tmp_path: Path): + endpoint = tmp_path / "endpoint.py" + endpoint.write_text( + "from . import nonexistent_sibling\n\n" + "from runpod_flash import Endpoint\n\n" + "@Endpoint\n" + "def handler():\n" + " return nonexistent_sibling.value\n" + ) + files = [endpoint] + + with pytest.raises(LocalModuleResolutionError): + augment_with_local_modules(files, tmp_path) + + +def test_endpoint_file_with_non_utf8_sibling_raises_resolution_error(tmp_path: Path): + # A valid endpoint importing a sibling with non-UTF-8 bytes triggers a raw + # UnicodeDecodeError during resolution. It must be normalized to + # LocalModuleResolutionError so run_build() reports it cleanly. + (tmp_path / "badsibling.py").write_bytes(b"\xff\xfe not valid utf-8\n") + endpoint = tmp_path / "endpoint.py" + endpoint.write_text( + "import badsibling\n\n" + "from runpod_flash import Endpoint\n\n" + "@Endpoint\n" + "def handler():\n" + " return badsibling.value\n" + ) + files = [endpoint] + + with pytest.raises(LocalModuleResolutionError, match="endpoint.py") as excinfo: + augment_with_local_modules(files, tmp_path) + # the raw decode error is preserved as the cause for debugging + assert isinstance(excinfo.value.__cause__, UnicodeDecodeError) + + +def test_non_endpoint_file_with_unresolvable_import_does_not_raise(tmp_path: Path): + incidental = tmp_path / "incidental.py" + incidental.write_text( + "from . import nonexistent_sibling\n\n" + "def helper():\n" + " return nonexistent_sibling.value\n" + ) + files = [incidental] + + augmented = augment_with_local_modules(files, tmp_path) + + assert augmented == [incidental] + + +def test_defines_endpoint_true_for_decorated_functions(tmp_path: Path): + endpoint_call = tmp_path / "endpoint_call.py" + endpoint_call.write_text( + "from runpod_flash import Endpoint\n\n@Endpoint(name='foo')\ndef handler():\n return 1\n" + ) + remote_plain = tmp_path / "remote_plain.py" + remote_plain.write_text( + "from runpod_flash import remote\n\n@remote\ndef handler():\n return 1\n" + ) + + assert _defines_endpoint(endpoint_call) is True + assert _defines_endpoint(remote_plain) is True + + +def test_defines_endpoint_false_for_plain_module(tmp_path: Path): + plain = tmp_path / "plain.py" + plain.write_text("def helper():\n return 1\n") + + assert _defines_endpoint(plain) is False diff --git a/tests/unit/core/test_exceptions.py b/tests/unit/core/test_exceptions.py new file mode 100644 index 00000000..d8107463 --- /dev/null +++ b/tests/unit/core/test_exceptions.py @@ -0,0 +1,16 @@ +from runpod_flash.core.exceptions import ( + LocalModulePayloadTooLargeError, + LocalModuleResolutionError, +) + + +def test_local_module_resolution_error_carries_message(): + err = LocalModuleResolutionError("cannot resolve 'utils'") + assert str(err) == "cannot resolve 'utils'" + assert isinstance(err, Exception) + + +def test_payload_too_large_error_carries_message(): + err = LocalModulePayloadTooLargeError("9 MB exceeds 8 MB cap") + assert str(err) == "9 MB exceeds 8 MB cap" + assert isinstance(err, Exception) diff --git a/tests/unit/protos/__init__.py b/tests/unit/protos/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/protos/test_remote_execution.py b/tests/unit/protos/test_remote_execution.py new file mode 100644 index 00000000..56965c2f --- /dev/null +++ b/tests/unit/protos/test_remote_execution.py @@ -0,0 +1,18 @@ +from runpod_flash.protos.remote_execution import FunctionRequest + + +def test_modules_defaults_to_empty_dict(): + req = FunctionRequest(function_name="f", function_code="def f(): pass") + assert req.modules == {} + + +def test_modules_round_trips_through_model_dump(): + req = FunctionRequest( + function_name="f", + function_code="def f(): import utils; return utils.x()", + modules={"utils.py": "def x():\n return 1\n"}, + ) + payload = req.model_dump(exclude_none=True) + assert payload["modules"] == {"utils.py": "def x():\n return 1\n"} + rebuilt = FunctionRequest(**payload) + assert rebuilt.modules == req.modules diff --git a/tests/unit/runtime/test_lb_handler_extended.py b/tests/unit/runtime/test_lb_handler_extended.py index a2cc61e8..48b3e76b 100644 --- a/tests/unit/runtime/test_lb_handler_extended.py +++ b/tests/unit/runtime/test_lb_handler_extended.py @@ -306,6 +306,38 @@ def test_execute_not_included_by_default(self): response = client.post("/execute", json={}) assert response.status_code == 404 + def test_execute_function_imports_shipped_module_inside_body(self, client): + """Regression: shipped module must still be importable when the function runs. + + The function body imports the shipped module lazily (i.e. at call + time, not at def time), which is how real endpoint code frequently + imports local siblings. If `materialized_modules` is exited before + the function is invoked, this import raises ModuleNotFoundError even + though the request is otherwise well-formed. + """ + from runpod_flash.runtime.serialization import deserialize_arg, serialize_arg + + response = client.post( + "/execute", + json={ + "function_name": "compute", + "function_code": ( + "def compute(x):\n" + " import sls360_regression_helper\n" + " return sls360_regression_helper.MULTIPLIER * x\n" + ), + "args": [serialize_arg(3)], + "kwargs": {}, + "modules": { + "sls360_regression_helper.py": "MULTIPLIER = 7\n", + }, + }, + ) + assert response.status_code == 200 + data = response.json() + assert data["success"] is True, data.get("error") + assert deserialize_arg(data["result"]) == 21 + class TestApiKeyMiddleware: """Test extract_api_key_middleware.""" diff --git a/tests/unit/runtime/test_module_loader.py b/tests/unit/runtime/test_module_loader.py new file mode 100644 index 00000000..e7201452 --- /dev/null +++ b/tests/unit/runtime/test_module_loader.py @@ -0,0 +1,81 @@ +import sys +import tempfile +from pathlib import Path + +import pytest + +from runpod_flash.runtime import module_loader +from runpod_flash.runtime.module_loader import materialized_modules + + +@pytest.fixture +def captured_tmpdirs(monkeypatch: pytest.MonkeyPatch) -> list[str]: + """Record every temp dir ``materialized_modules`` creates, to assert cleanup.""" + created: list[str] = [] + real_mkdtemp = tempfile.mkdtemp + + def _spy(*args: object, **kwargs: object) -> str: + path = real_mkdtemp(*args, **kwargs) + created.append(path) + return path + + monkeypatch.setattr(module_loader.tempfile, "mkdtemp", _spy) + return created + + +def test_empty_modules_is_noop(): + before = list(sys.path) + with materialized_modules({}) as tmpdir: + assert tmpdir is None + assert sys.path == before + assert sys.path == before + + +def test_writes_files_and_makes_them_importable(captured_tmpdirs: list[str]): + before = list(sys.path) + modules = {"pkg/__init__.py": "\n", "pkg/thing.py": "def val():\n return 42\n"} + with materialized_modules(modules) as tmpdir: + assert tmpdir in sys.path + namespace: dict = {} + exec("from pkg.thing import val\nresult = val()", namespace) + assert namespace["result"] == 42 + # sys.path restored and temp dir removed from disk after exit + assert sys.path == before + assert not Path(captured_tmpdirs[0]).exists() + + +def test_restores_sys_path_on_exception(): + before = list(sys.path) + try: + with materialized_modules({"m.py": "x = 1\n"}): + raise RuntimeError("boom") + except RuntimeError: + pass + assert sys.path == before + + +@pytest.mark.parametrize( + "rel_path", ["/etc/evil.py", "../evil.py", "pkg/../../evil.py"] +) +def test_rejects_paths_that_escape_temp_dir(rel_path: str, captured_tmpdirs: list[str]): + # ``modules`` is untrusted; a path that resolves outside the temp dir must be + # rejected before any write, sys.path left untouched, and the temp dir removed. + before = list(sys.path) + with pytest.raises(ValueError, match="escapes materialization dir"): + with materialized_modules({rel_path: "x = 1\n"}): + pass + assert sys.path == before + assert not Path(captured_tmpdirs[0]).exists() + + +def test_rejects_escape_after_valid_entry_and_cleans_up(captured_tmpdirs: list[str]): + # The guard fires mid-loop: a valid module is written first, then an escaping + # key must still raise and the partially-populated temp dir must be removed + # (regression guard for the temp-dir leak the cleanup restructure fixed). + before = list(sys.path) + modules = {"pkg/ok.py": "x = 1\n", "../evil.py": "y = 2\n"} + with pytest.raises(ValueError, match="escapes materialization dir"): + with materialized_modules(modules): + pass + assert sys.path == before + assert not Path(captured_tmpdirs[0]).exists() diff --git a/tests/unit/stubs/__init__.py b/tests/unit/stubs/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/stubs/test_live_serverless_modules.py b/tests/unit/stubs/test_live_serverless_modules.py new file mode 100644 index 00000000..78e0fc80 --- /dev/null +++ b/tests/unit/stubs/test_live_serverless_modules.py @@ -0,0 +1,29 @@ +import pytest + +from runpod_flash.core.exceptions import LocalModulePayloadTooLargeError +from runpod_flash.stubs import live_serverless +from runpod_flash.stubs.live_serverless import build_modules_map + + +def test_build_modules_map_collects_sibling(tmp_path): + (tmp_path / "utils.py").write_text("def x():\n return 1\n") + entry = tmp_path / "endpoint.py" + func_source = "def handler():\n import utils\n return utils.x()\n" + entry.write_text(func_source) + modules = build_modules_map(func_source, str(entry)) + assert modules == {"utils.py": "def x():\n return 1\n"} + + +def test_build_modules_map_enforces_size_cap(tmp_path, monkeypatch): + (tmp_path / "big.py").write_text("x = 1\n") + entry = tmp_path / "endpoint.py" + func_source = "def handler():\n import big\n return big.x\n" + entry.write_text(func_source) + monkeypatch.setattr(live_serverless, "MAX_INLINE_MODULE_BYTES", 1) + with pytest.raises(LocalModulePayloadTooLargeError): + build_modules_map(func_source, str(entry)) + + +def test_build_modules_map_no_file_returns_empty(): + # No __file__ (e.g. REPL-defined function) -> nothing to resolve, no crash. + assert build_modules_map("def handler():\n return 1\n", None) == {} diff --git a/tests/unit/stubs/test_local_modules_classify.py b/tests/unit/stubs/test_local_modules_classify.py new file mode 100644 index 00000000..9cc42c6f --- /dev/null +++ b/tests/unit/stubs/test_local_modules_classify.py @@ -0,0 +1,44 @@ +from pathlib import Path + +from runpod_flash.stubs.local_modules import ( + MAX_INLINE_MODULE_BYTES, + is_stdlib, + local_path_for, +) + + +def test_cap_is_8_mib(): + assert MAX_INLINE_MODULE_BYTES == 8 * 1024 * 1024 + + +def test_is_stdlib_true_for_os_and_dotted(): + assert is_stdlib("os") + assert is_stdlib("os.path") + + +def test_is_stdlib_false_for_unknown(): + assert not is_stdlib("utils") + + +def test_local_path_for_finds_sibling_module(tmp_path: Path): + (tmp_path / "utils.py").write_text("x = 1\n") + assert local_path_for("utils", [tmp_path]) == tmp_path / "utils.py" + + +def test_local_path_for_finds_package_init(tmp_path: Path): + pkg = tmp_path / "helpers" + pkg.mkdir() + (pkg / "__init__.py").write_text("\n") + assert local_path_for("helpers", [tmp_path]) == pkg / "__init__.py" + + +def test_local_path_for_finds_dotted_submodule(tmp_path: Path): + pkg = tmp_path / "helpers" + pkg.mkdir() + (pkg / "__init__.py").write_text("\n") + (pkg / "audio.py").write_text("y = 2\n") + assert local_path_for("helpers.audio", [tmp_path]) == pkg / "audio.py" + + +def test_local_path_for_returns_none_for_external(tmp_path: Path): + assert local_path_for("numpy", [tmp_path]) is None diff --git a/tests/unit/stubs/test_local_modules_resolve.py b/tests/unit/stubs/test_local_modules_resolve.py new file mode 100644 index 00000000..4d5c9103 --- /dev/null +++ b/tests/unit/stubs/test_local_modules_resolve.py @@ -0,0 +1,116 @@ +from pathlib import Path + +import pytest + +from runpod_flash.core.exceptions import LocalModuleResolutionError +from runpod_flash.stubs.local_modules import resolve_local_modules + + +def _entry(tmp_path: Path, body: str) -> Path: + f = tmp_path / "endpoint.py" + f.write_text(body) + return f + + +def test_includes_sibling_module(tmp_path: Path): + (tmp_path / "utils.py").write_text("def x():\n return 1\n") + entry = _entry(tmp_path, "import utils\n\ndef handler():\n return utils.x()\n") + result = resolve_local_modules(entry.read_text(), entry, tmp_path) + assert "utils.py" in result.files + + +def test_skips_stdlib_and_external(tmp_path: Path): + entry = _entry( + tmp_path, "import os\nimport numpy\n\ndef handler():\n return os.getpid()\n" + ) + result = resolve_local_modules(entry.read_text(), entry, tmp_path) + assert result.files == {} + + +def test_follows_transitive_local_imports(tmp_path: Path): + (tmp_path / "a.py").write_text("import b\n") + (tmp_path / "b.py").write_text("import c\n") + (tmp_path / "c.py").write_text("VALUE = 3\n") + entry = _entry(tmp_path, "import a\n\ndef handler():\n return a\n") + result = resolve_local_modules(entry.read_text(), entry, tmp_path) + assert set(result.files) == {"a.py", "b.py", "c.py"} + + +def test_pulls_package_init_for_submodule(tmp_path: Path): + pkg = tmp_path / "helpers" + pkg.mkdir() + (pkg / "__init__.py").write_text("\n") + (pkg / "audio.py").write_text("def load():\n return 1\n") + entry = _entry( + tmp_path, + "from helpers.audio import load\n\ndef handler():\n return load()\n", + ) + result = resolve_local_modules(entry.read_text(), entry, tmp_path) + assert set(result.files) == {"helpers/__init__.py", "helpers/audio.py"} + + +def test_in_body_import_is_discovered(tmp_path: Path): + (tmp_path / "utils.py").write_text("def x():\n return 1\n") + entry = _entry(tmp_path, "def handler():\n import utils\n return utils.x()\n") + result = resolve_local_modules(entry.read_text(), entry, tmp_path) + assert "utils.py" in result.files + + +def test_import_cycle_terminates(tmp_path: Path): + (tmp_path / "a.py").write_text("import b\n") + (tmp_path / "b.py").write_text("import a\n") + entry = _entry(tmp_path, "import a\n") + result = resolve_local_modules(entry.read_text(), entry, tmp_path) + assert set(result.files) == {"a.py", "b.py"} + + +def test_unresolved_relative_import_raises(tmp_path: Path): + entry = _entry(tmp_path, "from . import missing\n") + with pytest.raises(LocalModuleResolutionError): + resolve_local_modules(entry.read_text(), entry, tmp_path) + + +def test_dynamic_literal_is_resolved(tmp_path: Path): + (tmp_path / "plugin.py").write_text("Z = 9\n") + entry = _entry( + tmp_path, + "import importlib\n\ndef handler():\n return importlib.import_module('plugin')\n", + ) + result = resolve_local_modules(entry.read_text(), entry, tmp_path) + assert "plugin.py" in result.files + + +def test_dynamic_nonliteral_warns_not_raises(tmp_path: Path): + entry = _entry( + tmp_path, + "import importlib\n\ndef handler(name):\n return importlib.import_module(name)\n", + ) + result = resolve_local_modules(entry.read_text(), entry, tmp_path) + assert result.files == {} + assert any("dynamic import" in w for w in result.warnings) + + +def test_bare_relative_import_resolves_submodule(tmp_path): + (tmp_path / "__init__.py").write_text("\n") + (tmp_path / "helper.py").write_text("X = 1\n") + entry = _entry( + tmp_path, "from . import helper\n\ndef handler():\n return helper.X\n" + ) + result = resolve_local_modules(entry.read_text(), entry, tmp_path) + assert "helper.py" in result.files + + +def test_local_file_outside_project_root_raises(tmp_path): + proj = tmp_path / "proj" + proj.mkdir() + (tmp_path / "utils.py").write_text("x = 1\n") + entry = _entry(tmp_path, "import utils\n\ndef handler():\n return utils.x\n") + with pytest.raises(LocalModuleResolutionError): + resolve_local_modules(entry.read_text(), entry, proj) + + +def test_bare_relative_star_import_resolves_package_init(tmp_path): + (tmp_path / "__init__.py").write_text("VALUE = 1\n") + entry = _entry(tmp_path, "from . import *\n\ndef handler():\n return VALUE\n") + result = resolve_local_modules(entry.read_text(), entry, tmp_path) + assert "__init__.py" in result.files diff --git a/tests/unit/test_load_balancer_sls_stub.py b/tests/unit/test_load_balancer_sls_stub.py index f4ab4382..224539b3 100644 --- a/tests/unit/test_load_balancer_sls_stub.py +++ b/tests/unit/test_load_balancer_sls_stub.py @@ -101,6 +101,25 @@ def test_func(): assert request["dependencies"] == dependencies assert request["system_dependencies"] == system_deps + @pytest.mark.asyncio + async def test_prepare_request_includes_local_sibling_module(self): + """Local imports inside the function body are inlined into `modules`.""" + import pathlib + + stub = LoadBalancerSlsStub(test_lb_resource) + + def uses_sibling(): + from _lb_prepare_request_sibling import sibling_value + + return sibling_value() + + request = await stub._prepare_request(uses_sibling, None, None, True) + + sibling_path = pathlib.Path(__file__).parent / "_lb_prepare_request_sibling.py" + assert request["modules"] == { + "_lb_prepare_request_sibling.py": sibling_path.read_text() + } + class TestLoadBalancerSlsStubHandleResponse: """Test suite for _handle_response method.""" diff --git a/tests/unit/test_stub_live_serverless.py b/tests/unit/test_stub_live_serverless.py index 1c65f7b1..2a4e6a03 100644 --- a/tests/unit/test_stub_live_serverless.py +++ b/tests/unit/test_stub_live_serverless.py @@ -156,6 +156,31 @@ def func_with_kwargs(x, y=10): assert len(request.args) == 1 assert "y" in request.kwargs + @pytest.mark.asyncio + async def test_prepare_request_includes_local_sibling_module(self, stub): + """Local imports inside the function body are inlined into request.modules.""" + import pathlib + + def uses_sibling(): + from _live_serverless_prepare_request_sibling import sibling_value + + return sibling_value() + + with patch( + "runpod_flash.stubs.dependency_resolver.resolve_dependencies", + new_callable=AsyncMock, + return_value=[], + ): + request = await stub.prepare_request(uses_sibling, [], [], True) + + sibling_path = ( + pathlib.Path(__file__).parent + / "_live_serverless_prepare_request_sibling.py" + ) + assert request.modules == { + "_live_serverless_prepare_request_sibling.py": sibling_path.read_text() + } + def test_handle_response_success(self, stub): """handle_response returns deserialized result on success.""" result_data = {"key": "value"}