diff --git a/docs/Cross_Endpoint_Routing.md b/docs/Cross_Endpoint_Routing.md index 18131c665..c68cc7446 100644 --- a/docs/Cross_Endpoint_Routing.md +++ b/docs/Cross_Endpoint_Routing.md @@ -183,6 +183,20 @@ Cross-endpoint calls require authentication. Flash handles this automatically: This means your function code never needs to handle API keys directly. +`flash build` sets `makes_remote_calls=True` when an endpoint function or class, directly or through a helper function, class, or object defined in your project (package `__init__.py` files included): + +- references another `@Endpoint` function or class, including through a module or package path (`tasks.square(x)`, `pkg.tasks.square(x)`), an aliased import, or an import inside the function (relative imports included) +- uses a client `Endpoint(id=...)` or `Endpoint(image=...)`, at module level, built inline, or held by a project object, container, bound method, `functools.partial`, or cached factory +- uses the `runpod` SDK, or reads `RUNPOD_API_KEY` + +A helper, class, or project import that `flash build` cannot analyze or resolve counts as needing the key. + +Endpoints that do none of these do not receive the key. The deploy summary lists the endpoints that did. + +A client `Endpoint(image=...)` deployed by Flash runs an image Flash does not scan, so it does not receive the key. If your image calls other endpoints, set `RUNPOD_API_KEY` in that endpoint's `env`; Flash never removes a key you set there. + +A call that was not detected (for example, a client `Endpoint` passed in as an argument) fails with `401 Unauthorized`. When the call goes through a client `Endpoint`, the error also says `RUNPOD_API_KEY is not set on this worker`. Make the call visible in one of the ways above and redeploy. + ## Load-Balanced Cross-Endpoint Calls LB endpoints can also participate in cross-endpoint routing: diff --git a/docs/Deployment_Architecture.md b/docs/Deployment_Architecture.md index 839ae4051..81806ff97 100644 --- a/docs/Deployment_Architecture.md +++ b/docs/Deployment_Architecture.md @@ -188,7 +188,7 @@ When deploying to an environment that already has endpoints, Flash compares the When `flash deploy` provisions endpoints: -1. Endpoints with `makes_remote_calls=True` get `RUNPOD_API_KEY` injected automatically +1. Endpoints with `makes_remote_calls=True` get `RUNPOD_API_KEY` injected automatically. Other endpoints do not get the key, and a key left on them by an earlier deploy is removed on the next `flash deploy`, even when nothing changed. A template shared with another endpoint keeps its key, and a key you set in `env` is never removed. See [API Key Injection](Cross_Endpoint_Routing.md#api-key-injection) for what counts as a remote call 2. Each endpoint gets the `flash_manifest.json` included in its artifact 3. The State Manager stores `{environment_id, resource_name} -> endpoint_id` 4. At runtime, the `ServiceRegistry` uses the manifest + State Manager to route calls diff --git a/src/runpod_flash/cli/commands/build_utils/scanner.py b/src/runpod_flash/cli/commands/build_utils/scanner.py index c9a383b7c..95b4e243b 100644 --- a/src/runpod_flash/cli/commands/build_utils/scanner.py +++ b/src/runpod_flash/cli/commands/build_utils/scanner.py @@ -6,16 +6,19 @@ """ import ast +import functools import importlib.util import inspect +import itertools import logging import os import sys +import textwrap import threading import types from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Dict, List, Optional, Set +from typing import Any, Dict, List, Optional, Set, Tuple from runpod_flash.cli.utils.ignore import get_file_tree, load_ignore_patterns from runpod_flash.core.resources.load_balancer_sls_resource import ( @@ -407,16 +410,376 @@ def _find_remote_decorated(module: Any) -> Dict[str, Any]: return results +_API_KEY_ENV = "RUNPOD_API_KEY" +_UNANALYZED = "" +_MISSING = object() + + +@dataclass +class _CallScan: + """shared state for one endpoint's cross-call analysis. + + project_modules maps each scanned module path to its module, or None when + the file failed to import. the scanner removes those modules from + sys.modules after importing them, so this is how analysis still finds + them. + """ + + remote_function_names: Set[str] + project_files: Optional[Set[Path]] + project_modules: Dict[str, Any] + cache: Dict[int, List[str]] = field(default_factory=dict) + + def lookup_module(self, name: str) -> Any: + """the module for name, or None. scanned project modules win.""" + if name in self.project_modules: + return self.project_modules[name] + return sys.modules.get(name) + + def is_project_module(self, name: str) -> bool: + """true if name is a scanned module or a package that contains one.""" + prefix = name + "." + return any( + key == name or key.startswith(prefix) for key in self.project_modules + ) + + +def _resolve(expr: ast.expr, namespace: Dict[str, Any], scan: _CallScan) -> Any: + """resolve a Name, or an attribute of a module (pkg.utils.helper), to its object.""" + if isinstance(expr, ast.Name): + return namespace.get(expr.id, _MISSING) + if isinstance(expr, ast.Attribute): + base = _resolve(expr.value, namespace, scan) + if isinstance(base, types.ModuleType): + obj = getattr(base, expr.attr, _MISSING) + if obj is _MISSING: + # a scanned submodule is not set as an attribute of its package + obj = scan.lookup_module(f"{base.__name__}.{expr.attr}") + return _MISSING if obj is None else obj + return _MISSING + + +def _is_client_endpoint(obj: Any) -> bool: + return isinstance(obj, Endpoint) and obj.is_client + + +def _is_runpod_sdk(obj: Any) -> bool: + """true for the runpod SDK module or anything defined in it.""" + try: + if isinstance(obj, types.ModuleType): + name = obj.__name__ + else: + name = getattr(obj, "__module__", None) or "" + except Exception: + return False + return isinstance(name, str) and _is_runpod_module_name(name) + + +def _is_remote(obj: Any) -> bool: + """true for an @remote / @Endpoint decorated function or class.""" + try: + return obj is not _MISSING and hasattr(obj, "__remote_config__") + except Exception: + return False + + +def _in_project(source_file: Optional[str], scan: _CallScan) -> bool: + return ( + bool(scan.project_files) + and source_file is not None + and Path(source_file).resolve() in scan.project_files # type: ignore[operator] + ) + + +def _project_function(obj: Any, scan: _CallScan) -> Any: + """return obj if it is a plain function defined in one of the scanned files. + + limiting this to scanned files keeps helper analysis out of library code, + including virtualenvs or bundled packages that live inside the project. + """ + if _is_remote(obj): + return None + obj = _unwrap(obj) + if not isinstance(obj, types.FunctionType) or _is_remote(obj): + return None + try: + source_file = inspect.getsourcefile(obj) + except Exception: + return None + return obj if _in_project(source_file, scan) else None + + +def _project_class(obj: Any, scan: _CallScan) -> Tuple[Any, Any]: + """return (class, module) if obj is a project class or an instance of one.""" + if obj is _MISSING or isinstance(obj, types.ModuleType) or _is_remote(obj): + return None, None + try: + cls = obj if inspect.isclass(obj) else type(obj) + module = scan.lookup_module(cls.__module__) + source_file = getattr(module, "__file__", None) + except Exception: + return None, None + if _is_remote(cls) or not _in_project(source_file, scan): + return None, None + return cls, module + + +def _is_client_endpoint_construction( + call: ast.Call, namespace: Dict[str, Any], scan: _CallScan +) -> bool: + """match Endpoint(id=...), Endpoint(image=...), or Endpoint(**kw) built inline.""" + func = call.func + literal = (isinstance(func, ast.Name) and func.id == "Endpoint") or ( + isinstance(func, ast.Attribute) and func.attr == "Endpoint" + ) + if not literal and _resolve(func, namespace, scan) is not Endpoint: + return False + # id and image are keyword-only; **kwargs could carry either + return any(kw.arg in ("id", "image", None) for kw in call.keywords) + + +def _local_imports( + node: ast.AST, namespace: Dict[str, Any], scan: _CallScan +) -> Tuple[Dict[str, Any], bool, bool]: + """resolve imports made inside node against already-loaded modules. + + returns (name -> object overlay, whether the runpod SDK is imported, + whether a project import could not be resolved). only loaded modules are + consulted, so analysis never imports new code. + """ + overlay: Dict[str, Any] = {} + uses_runpod = False + unresolved = False + for child in ast.walk(node): + if isinstance(child, ast.Import): + for alias in child.names: + if _is_runpod_module_name(alias.name): + uses_runpod = True + if alias.asname: + name, bound = alias.name, alias.asname + else: + name = bound = alias.name.split(".")[0] + module = scan.lookup_module(name) + if module is not None: + overlay[bound] = module + elif scan.is_project_module(alias.name): + unresolved = True + elif isinstance(child, ast.ImportFrom): + if child.level: + try: + name = importlib.util.resolve_name( + "." * child.level + (child.module or ""), + namespace.get("__package__"), + ) + except (ImportError, ValueError): + unresolved = True + continue + elif child.module: + name = child.module + else: + continue + if _is_runpod_module_name(name): + uses_runpod = True + # a relative import always names project code + in_project = bool(child.level) or scan.is_project_module(name) + module = scan.lookup_module(name) + if module is None: + unresolved = unresolved or in_project + continue + for alias in child.names: + obj = getattr(module, alias.name, _MISSING) + if obj is _MISSING: + submodule = f"{name}.{alias.name}" + obj = scan.lookup_module(submodule) + if obj is None: + unresolved = ( + unresolved + or in_project + or scan.is_project_module(submodule) + ) + continue + overlay[alias.asname or alias.name] = obj + return overlay, uses_runpod, unresolved + + +def _is_runpod_module_name(name: str) -> bool: + return name == "runpod" or name.startswith("runpod.") + + +def _object_needs_key(obj: Any, scan: _CallScan) -> bool: + """true if referencing obj means the worker needs RUNPOD_API_KEY. + + a project class, or an instance of one, needs it when any of its methods + or instance attributes would. the whole class is analyzed because calls + through self cannot be resolved statically. + """ + return _needs_key(obj, scan, set()) + + +# cap on container items followed, so a large config dict stays cheap +_MAX_CONTAINER_ITEMS = 1000 + + +def _needs_key(obj: Any, scan: _CallScan, seen: Set[int]) -> bool: + if _is_remote(obj) or _is_client_endpoint(obj): + return True + if id(obj) in seen: + return False + seen.add(id(obj)) + # vllm.runsync, functools.partial(vllm.runsync), {"vllm": vllm}, and + # wrappers that set __wrapped__ (functools.lru_cache, functools.wraps) + if isinstance(obj, types.MethodType): + return _needs_key(obj.__self__, scan, seen) + if isinstance(obj, functools.partial): + parts = [obj.func, *obj.args, *obj.keywords.values()] + return any(_needs_key(part, scan, seen) for part in parts) + if isinstance(obj, (dict, list, tuple, set, frozenset)): + items = obj.values() if isinstance(obj, dict) else obj + return any( + _needs_key(item, scan, seen) + for item in itertools.islice(items, _MAX_CONTAINER_ITEMS) + ) + unwrapped = _unwrap(obj) + if unwrapped is not obj: + return _needs_key(unwrapped, scan, seen) + # a project helper used as a value: partial(call_model), {"gen": call_model} + helper = _project_function(obj, scan) + if helper is not None: + return bool(_helper_targets(helper, scan)) + cls, module = _project_class(obj, scan) + if cls is None: + return False + if _class_targets(cls, module, scan): + return True + if inspect.isclass(obj): + return False + try: + attrs = list(vars(obj).values()) + except TypeError: + return False + # holder.inner.ep: follow instance attributes into nested project objects + return any(_needs_key(value, scan, seen) for value in attrs) + + +def _unwrap(obj: Any) -> Any: + """the function under functools.wraps / lru_cache style wrappers, else obj.""" + if inspect.isclass(obj): + return obj + try: + return inspect.unwrap(obj) + except Exception: + return obj + + +def _remote_call_targets( + node: ast.AST, namespace: Dict[str, Any], scan: _CallScan +) -> List[str]: + """names of the things in node that need RUNPOD_API_KEY on the worker.""" + called: List[str] = [] + + def add(target: str) -> None: + if target not in called: + called.append(target) + + overlay, uses_runpod, unresolved = _local_imports(node, namespace, scan) + if unresolved: + add(_UNANALYZED) + if uses_runpod: + add("runpod") + if overlay: + namespace = {**namespace, **overlay} + + for child in ast.walk(node): + if isinstance(child, ast.Constant) and child.value == _API_KEY_ENV: + add(_API_KEY_ENV) + elif isinstance(child, ast.Call): + if ( + isinstance(child.func, ast.Name) + and child.func.id in scan.remote_function_names + ): + add(child.func.id) + if _is_client_endpoint_construction(child, namespace, scan): + add("Endpoint") + helper = _project_function(_resolve(child.func, namespace, scan), scan) + if helper is not None and _helper_targets(helper, scan): + add(helper.__name__) + elif isinstance(child, ast.Name): + obj = _resolve(child, namespace, scan) + if _is_runpod_sdk(obj) or _object_needs_key(obj, scan): + add(child.id) + elif isinstance(child, ast.Attribute): + # tasks.square(x), pkg.tasks.square(x), or utils.vllm.runsync(x) + obj = _resolve(child, namespace, scan) + if _object_needs_key(obj, scan): + add(ast.unparse(child)) + return called + + +def _helper_targets(func: Any, scan: _CallScan) -> List[str]: + """analyze a project helper function, resolving names in its own module.""" + key = id(func) + if key in scan.cache: + return scan.cache[key] + scan.cache[key] = [] # break recursion cycles + try: + tree = ast.parse(textwrap.dedent(inspect.getsource(func))) + except Exception: + # e.g. a lambda inside a multi-line call: getsource returns a fragment + scan.cache[key] = [_UNANALYZED] + return scan.cache[key] + scan.cache[key] = _remote_call_targets(tree, func.__globals__, scan) + return scan.cache[key] + + +def _class_targets(cls: Any, module: Any, scan: _CallScan) -> List[str]: + """analyze every method of a project class, resolving names in its module.""" + key = id(cls) + if key in scan.cache: + return scan.cache[key] + scan.cache[key] = [] # break recursion cycles + try: + tree = ast.parse(Path(module.__file__).read_text(encoding="utf-8")) + nodes = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.ClassDef) and n.name == cls.__name__ + ] + except Exception: + nodes = [] + if not nodes: + scan.cache[key] = [_UNANALYZED] + return scan.cache[key] + targets: List[str] = [] + for class_node in nodes: + for target in _remote_call_targets(class_node, vars(module), scan): + if target not in targets: + targets.append(target) + scan.cache[key] = targets + return targets + + def _analyze_cross_calls_ast( file_path: Path, function_names: Set[str], remote_function_names: Set[str], + namespace: Optional[Dict[str, Any]] = None, + project_files: Optional[Set[Path]] = None, + project_modules: Optional[Dict[str, Any]] = None, ) -> Dict[str, List[str]]: - """find which functions call other @remote functions via AST. - - returns a dict of function_name -> list of called remote function names. - only matches direct calls (foo()), not attribute calls (obj.foo()). + """find which functions need RUNPOD_API_KEY on the worker, via AST. + + returns a dict of function_name -> list of what triggered it. a function + needs the key when it references another @remote function or class + (directly, aliased, or through a module), uses a client-mode Endpoint + (id= or image=), uses the runpod SDK, reads RUNPOD_API_KEY, or calls a + project helper or project class that does any of these. namespace is the + module's globals, used to resolve names; project_files (resolved paths) + limits helper analysis to scanned project code; project_modules maps + scanned module paths to their modules. detection errs toward giving the + key, since a missing key fails at runtime with a 401: a function, helper, + or project import that cannot be analyzed is reported as needing it. """ + ns = namespace or {} result: Dict[str, List[str]] = {} try: tree = ast.parse(file_path.read_text(encoding="utf-8")) @@ -429,12 +792,14 @@ def _analyze_cross_calls_ast( if node.name not in function_names: continue - called: List[str] = [] - for child in ast.walk(node): - if isinstance(child, ast.Call) and isinstance(child.func, ast.Name): - if child.func.id in remote_function_names: - if child.func.id not in called: - called.append(child.func.id) + # fresh per function: a helper cached mid-cycle can hold a partial + # result that is only correct within the traversal that produced it. + scan = _CallScan(remote_function_names, project_files, project_modules or {}) + try: + called = _remote_call_targets(node, ns, scan) + except Exception as e: + logger.debug(f"cross-call analysis failed for {node.name}: {e}") + called = [_UNANALYZED] if called: result[node.name] = called @@ -479,10 +844,14 @@ def discover_remote_functions(self) -> List[RemoteFunctionMetadata]: functions: List[RemoteFunctionMetadata] = [] seen_functions: Set[str] = set() failed_files: List[Path] = [] + namespaces: Dict[Path, Dict[str, Any]] = {} + # module path -> module (None if the import failed), for analysis + project_modules: Dict[str, Any] = {} try: for py_file in py_files: module_path = file_to_module_path(py_file, self.project_dir) + project_modules[module_path] = None try: module = _import_module_from_file(py_file, module_path) except Exception as e: @@ -493,20 +862,32 @@ def discover_remote_functions(self) -> List[RemoteFunctionMetadata]: if module is None: failed_files.append(py_file) continue + project_modules[module_path] = module file_functions = self._extract_from_module( module, module_path, py_file, seen_functions ) functions.extend(file_functions) + namespaces[py_file] = vars(module) # cross-call analysis remote_names = {f.function_name for f in functions} + # __init__.py files are not scanned for endpoints, but helpers + # defined in a package __init__ still need analysis + project_files = {f.resolve() for f in all_files if f.suffix == ".py"} files_with_functions: Dict[Path, Set[str]] = {} for f in functions: files_with_functions.setdefault(f.file_path, set()).add(f.function_name) for file_path, func_names in files_with_functions.items(): - calls = _analyze_cross_calls_ast(file_path, func_names, remote_names) + calls = _analyze_cross_calls_ast( + file_path, + func_names, + remote_names, + namespaces.get(file_path), + project_files, + project_modules, + ) for f in functions: if f.file_path == file_path and f.function_name in calls: f.calls_remote_functions = True diff --git a/src/runpod_flash/cli/commands/deploy.py b/src/runpod_flash/cli/commands/deploy.py index 73939ae5a..46a7fb61f 100644 --- a/src/runpod_flash/cli/commands/deploy.py +++ b/src/runpod_flash/cli/commands/deploy.py @@ -170,6 +170,20 @@ def _display_post_deployment_guidance( for method, path in route_list: console.print(f" {' ' * max_name} [dim]{method:6s}{path}[/dim]") + # a RUNPOD_API_KEY the user set in env is theirs; Flash does not add it + keyed = sorted( + name + for name, _url, _routes in rows + if resources.get(name, {}).get("makes_remote_calls") + and "RUNPOD_API_KEY" not in (resources[name].get("env") or {}) + ) + if keyed: + console.print() + console.print( + "[dim]RUNPOD_API_KEY set by Flash on endpoints that call other endpoints: " + f"{', '.join(keyed)}[/dim]" + ) + # one curl example curl_url = None curl_method = "POST" diff --git a/src/runpod_flash/cli/utils/deployment.py b/src/runpod_flash/cli/utils/deployment.py index 1a53e80cc..10a5d02b3 100644 --- a/src/runpod_flash/cli/utils/deployment.py +++ b/src/runpod_flash/cli/utils/deployment.py @@ -7,7 +7,7 @@ from typing import Dict, Any from pathlib import Path -from runpod_flash.core.resources.serverless import ServerlessResource +from runpod_flash.core.resources.serverless import ServerlessResource, template_shared from runpod_flash.core.resources.app import FlashApp from runpod_flash.core.resources.resource_manager import ResourceManager from runpod_flash.runtime.resource_provisioner import create_resource_from_manifest @@ -41,6 +41,78 @@ def _manifest_without_ai_keys(manifest: Dict[str, Any]) -> Dict[str, Any]: return sanitized_manifest +# bound on the stale-key cleanup, which runs before provisioning starts +STALE_KEY_CLEANUP_TIMEOUT_SECONDS = 120 + + +async def _remove_stale_api_key(client: Any, template_id: str) -> bool: + """Drop RUNPOD_API_KEY from a live template. True if removed.""" + live = await client.get_template(template_id) + env = live.get("env") or [] + kept = [entry for entry in env if entry.get("key") != "RUNPOD_API_KEY"] + if len(kept) == len(env): + return False + env_payload = [ + {"key": entry["key"], "value": entry.get("value", "")} for entry in kept + ] + payload = ServerlessResource.template_update_payload( + {**live, "env": env_payload}, template_id + ) + await client.update_template(payload) + return True + + +async def _remove_stale_api_keys(endpoints: list[tuple[str, str]]) -> None: + """Remove RUNPOD_API_KEY left on unchanged local-only endpoints. + + Older Flash versions gave every endpoint the key. Unchanged resources skip + provisioning, so the update path that drops it never runs for them. + A template shared with another endpoint keeps the key, since that endpoint + may call others. Failures are logged and never fail the deploy. + """ + from runpod_flash.core.api.runpod import RunpodGraphQLClient + + async def clean(client: Any, template_ids: dict, name: str, endpoint_id: str): + template_id = template_ids.get(endpoint_id) + if not template_id: + return # endpoint deleted, or no template + if template_shared(template_ids, template_id): + log.warning( + f"{name}: kept RUNPOD_API_KEY, template '{template_id}' " + f"is shared with other endpoints" + ) + return + try: + removed = await _remove_stale_api_key(client, template_id) + except Exception as e: + log.warning(f"{name}: could not check for a stale RUNPOD_API_KEY: {e}") + return + if removed: + log.info( + f"{name}: removed RUNPOD_API_KEY (endpoint does not call other endpoints)" + ) + + async def run() -> None: + async with RunpodGraphQLClient() as client: + template_ids = await client.get_endpoint_template_ids() + await asyncio.gather( + *( + clean(client, template_ids, name, endpoint_id) + for name, endpoint_id in endpoints + ) + ) + + try: + await asyncio.wait_for(run(), timeout=STALE_KEY_CLEANUP_TIMEOUT_SECONDS) + except asyncio.TimeoutError: + log.warning( + f"Stopped checking endpoints for a stale RUNPOD_API_KEY after " + f"{STALE_KEY_CLEANUP_TIMEOUT_SECONDS}s" + ) + except Exception as e: + log.warning(f"Could not check endpoints for a stale RUNPOD_API_KEY: {e}") + + def _resource_config_for_compare(config: Dict[str, Any]) -> Dict[str, Any]: compare_config = copy.deepcopy(config) for field in RUNTIME_RESOURCE_FIELDS: @@ -116,6 +188,7 @@ async def reconcile_and_provision_resources( # Create resource manager manager = ResourceManager() actions = [] + stale_key_checks: list[tuple[str, str]] = [] manifest_python_version = local_manifest.get("python_version") # Inject source fingerprint into each resource's env so that code-only @@ -185,6 +258,12 @@ async def reconcile_and_provision_resources( local_manifest["resources"][resource_name]["endpoint_id"] = ( state_config["endpoint_id"] ) + if local_config.get( + "makes_remote_calls" + ) is False and "RUNPOD_API_KEY" not in (local_config.get("env") or {}): + stale_key_checks.append( + (resource_name, state_config["endpoint_id"]) + ) if "aiKey" in state_config: local_manifest["resources"][resource_name]["aiKey"] = state_config[ "aiKey" @@ -198,6 +277,9 @@ async def reconcile_and_provision_resources( for resource_name in sorted(to_delete): log.debug(f"Resource {resource_name} marked for deletion (not implemented yet)") + if stale_key_checks: + await _remove_stale_api_keys(stale_key_checks) + # Execute all actions in parallel with timeout if actions: try: diff --git a/src/runpod_flash/core/api/runpod.py b/src/runpod_flash/core/api/runpod.py index b527d2cef..5c36f7659 100644 --- a/src/runpod_flash/core/api/runpod.py +++ b/src/runpod_flash/core/api/runpod.py @@ -847,6 +847,23 @@ async def endpoint_exists(self, endpoint_id: str) -> bool: log.debug(f"Error checking endpoint existence: {e}") return False + async def get_endpoint_template_ids(self) -> Dict[str, Optional[str]]: + """Return endpoint ID -> template ID for every endpoint on the account.""" + query = """ + query getEndpointTemplates { + myself { + endpoints { + id + templateId + } + } + } + """ + + result = await self._execute_graphql(query) + endpoints = (result.get("myself") or {}).get("endpoints") or [] + return {ep["id"]: ep.get("templateId") for ep in endpoints if ep.get("id")} + async def create_flash_auth_request(self) -> Dict[str, Any]: mutation = """ mutation createFlashAuthRequest { diff --git a/src/runpod_flash/core/resources/serverless.py b/src/runpod_flash/core/resources/serverless.py index d3a56797b..67f6eb5ae 100644 --- a/src/runpod_flash/core/resources/serverless.py +++ b/src/runpod_flash/core/resources/serverless.py @@ -54,6 +54,15 @@ ) +def template_shared(template_ids: Dict[str, Optional[str]], template_id: str) -> bool: + """True if more than one endpoint runs template_id. + + template_ids maps endpoint ID -> template ID, as returned by + RunpodGraphQLClient.get_endpoint_template_ids. + """ + return sum(1 for tid in template_ids.values() if tid == template_id) > 1 + + def _normalize_stream_log_line(line: str) -> str: normalized = line.strip() if not normalized: @@ -290,6 +299,9 @@ class ServerlessResource(DeployableResource): # === Private Attributes === _deployed_volume_ids: list[str] = PrivateAttr(default_factory=list) + # makes_remote_calls from the build manifest, set by flash deploy. + # None means unknown (e.g. flash run), so the manifest file lookup applies. + _makes_remote_calls: Optional[bool] = PrivateAttr(default=None) # === Runtime Fields === activeBuildid: Optional[str] = None @@ -853,14 +865,24 @@ def _build_template_update_payload( ) -> Dict[str, Any]: """Build saveTemplate payload from template model. - Keep this to fields supported by saveTemplate to avoid passing endpoint-only - fields to the template mutation. - Args: template: Template model with desired configuration. template_id: ID of the template to update. """ - template_data = template.model_dump(exclude_none=True, mode="json") + return ServerlessResource.template_update_payload( + template.model_dump(exclude_none=True, mode="json"), template_id + ) + + @staticmethod + def template_update_payload( + template_data: Dict[str, Any], + template_id: str, + ) -> Dict[str, Any]: + """Build saveTemplate payload from template fields, e.g. a live template. + + Keep this to fields supported by saveTemplate to avoid passing endpoint-only + fields to the template mutation. + """ allowed_fields = { "name", "imageName", @@ -870,7 +892,9 @@ def _build_template_update_payload( "readme", } payload = { - key: value for key, value in template_data.items() if key in allowed_fields + key: value + for key, value in template_data.items() + if key in allowed_fields and value is not None } # saveTemplate requires env (non-nullable) — default to empty list payload.setdefault("env", []) @@ -879,6 +903,33 @@ def _build_template_update_payload( payload["id"] = template_id return payload + async def _should_drop_api_key( + self, client: Any, template_id: str, live_env: List[Dict[str, Any]] + ) -> bool: + """True if a RUNPOD_API_KEY on the live template should be removed. + + Only for local-only endpoints, and only when no other endpoint runs + the same template: a caller sharing it would lose the key too. + """ + if self._check_makes_remote_calls(): + return False + if not any(entry.get("key") == "RUNPOD_API_KEY" for entry in live_env): + return False + try: + template_ids = await client.get_endpoint_template_ids() + except Exception as e: + log.warning( + f"{self.name}: Keeping RUNPOD_API_KEY, could not check template use: {e}" + ) + return False + if template_shared(template_ids, template_id): + log.warning( + f"{self.name}: Keeping RUNPOD_API_KEY, template '{template_id}' " + f"is shared with other endpoints" + ) + return False + return True + def _check_makes_remote_calls(self) -> bool: """Check if resource makes remote calls from build manifest. @@ -889,6 +940,9 @@ def _check_makes_remote_calls(self) -> bool: True if makes remote calls, False if local-only, True (safe default) if manifest not found. """ + if self._makes_remote_calls is not None: + return self._makes_remote_calls + try: manifest_path = Path.cwd() / "flash_manifest.json" if not manifest_path.exists(): @@ -1053,10 +1107,16 @@ async def _preserve_platform_env( new_keys = {kv.key for kv in (self.template.env or [])} old_keys = set(old_env or {}) + drop_api_key = await self._should_drop_api_key(client, template_id, live_env) for entry in live_env: key = entry.get("key", "") if not key or key in new_keys: continue + # Flash injects RUNPOD_API_KEY itself; a key left by an earlier + # deploy must not outlive the endpoint's need for it. + if key == "RUNPOD_API_KEY" and drop_api_key: + log.debug(f"{self.name}: Dropping stale RUNPOD_API_KEY") + continue # Key was in old user config — user intentionally removed it if key in old_keys: log.debug(f"{self.name}: User removed env var '{key}', not preserving") @@ -1214,11 +1274,20 @@ async def update(self, new_config: "ServerlessResource") -> "ServerlessResource" try: live = await client.get_template(resolved_template_id) live_env = live.get("env") or [] + drop_api_key = "RUNPOD_API_KEY" not in ( + new_config.env or {} + ) and await new_config._should_drop_api_key( + client, resolved_template_id, live_env + ) new_config.template.env = [ KeyValuePair( key=entry["key"], value=entry.get("value", "") ) for entry in live_env + if not ( + drop_api_key + and entry["key"] == "RUNPOD_API_KEY" + ) ] except Exception: log.warning( diff --git a/src/runpod_flash/endpoint.py b/src/runpod_flash/endpoint.py index 632369dbd..f0ca30d15 100644 --- a/src/runpod_flash/endpoint.py +++ b/src/runpod_flash/endpoint.py @@ -33,6 +33,34 @@ _POLL_MAX_CONSECUTIVE_ERRORS = 5 +_MISSING_WORKER_KEY_HINT = ( + "RUNPOD_API_KEY is not set on this worker. flash deploy only adds it to " + "endpoints that flash build detects calling other endpoints: referencing " + "another @Endpoint, a client Endpoint(id=...) or Endpoint(image=...), the " + "runpod SDK, or RUNPOD_API_KEY, directly or through a helper function in " + "the project. Make the call visible in one of those ways and redeploy." +) + + +def _raise_for_status(response: Any) -> None: + """raise_for_status, explaining a 401 caused by a missing worker API key.""" + import httpx + + from .core.credentials import get_api_key + + try: + response.raise_for_status() + except httpx.HTTPStatusError as e: + on_worker = bool(os.getenv("RUNPOD_ENDPOINT_ID")) + if response.status_code == 401 and on_worker and not get_api_key(): + raise httpx.HTTPStatusError( + f"{e}\n{_MISSING_WORKER_KEY_HINT}", + request=e.request, + response=e.response, + ) from e + raise + + class _ClientCoroutine: """wraps a coroutine from a client-mode HTTP call. @@ -946,7 +974,7 @@ async def _client_request( async with get_authenticated_httpx_client(timeout=timeout) as client: response = await client.request(method, full_url, json=data) - response.raise_for_status() + _raise_for_status(response) return response.json() async def _api_post(self, url: str, payload: Any, timeout: float = 60.0) -> dict: @@ -955,7 +983,7 @@ async def _api_post(self, url: str, payload: Any, timeout: float = 60.0) -> dict async with get_authenticated_httpx_client(timeout=timeout) as client: response = await client.post(url, json=payload) - response.raise_for_status() + _raise_for_status(response) return response.json() async def _api_get(self, url: str, timeout: float = 30.0) -> dict: @@ -964,5 +992,5 @@ async def _api_get(self, url: str, timeout: float = 30.0) -> dict: async with get_authenticated_httpx_client(timeout=timeout) as client: response = await client.get(url) - response.raise_for_status() + _raise_for_status(response) return response.json() diff --git a/src/runpod_flash/runtime/resource_provisioner.py b/src/runpod_flash/runtime/resource_provisioner.py index fcecbafff..56d43596a 100644 --- a/src/runpod_flash/runtime/resource_provisioner.py +++ b/src/runpod_flash/runtime/resource_provisioner.py @@ -201,4 +201,9 @@ def create_resource_from_manifest( # ServerlessResource (default) resource = ServerlessResource(**deployment_kwargs) + # Pin the build-time decision so template injection does not re-search + # for flash_manifest.json, which deploy keeps under .flash/. + if "makes_remote_calls" in resource_data: + resource._makes_remote_calls = bool(resource_data["makes_remote_calls"]) + return resource diff --git a/tests/unit/cli/commands/build_utils/test_runtime_scanner.py b/tests/unit/cli/commands/build_utils/test_runtime_scanner.py index 10b7f3dab..2df49509c 100644 --- a/tests/unit/cli/commands/build_utils/test_runtime_scanner.py +++ b/tests/unit/cli/commands/build_utils/test_runtime_scanner.py @@ -21,6 +21,7 @@ file_to_resource_name, file_to_url_prefix, ) +from runpod_flash.endpoint import Endpoint # --------------------------------------------------------------------------- @@ -315,6 +316,95 @@ def caller(): result = _analyze_cross_calls_ast(p, {"caller", "remote_fn"}, {"remote_fn"}) assert result["caller"] == ["remote_fn"] + def test_detects_client_endpoint_reference(self, tmp_path): + p = _write_worker( + tmp_path, + "client_ref.py", + """\ + async def caller(x): + return await vllm.runsync(x) + """, + ) + vllm = Endpoint(name="vllm", image="vllm/vllm-openai:latest") + result = _analyze_cross_calls_ast(p, {"caller"}, set(), {"vllm": vllm}) + assert result == {"caller": ["vllm"]} + + def test_detects_aliased_inline_client_endpoint(self, tmp_path): + p = _write_worker( + tmp_path, + "alias.py", + """\ + async def caller(x): + return await EP(id="abc123").runsync(x) + """, + ) + result = _analyze_cross_calls_ast(p, {"caller"}, set(), {"EP": Endpoint}) + assert result == {"caller": ["Endpoint"]} + + def test_detects_runpod_sdk_use(self, tmp_path): + import runpod + + p = _write_worker( + tmp_path, + "sdk.py", + """\ + def caller(x): + return runpod.Endpoint("abc123").run_sync(x) + """, + ) + result = _analyze_cross_calls_ast(p, {"caller"}, set(), {"runpod": runpod}) + assert result == {"caller": ["runpod"]} + + def test_detects_api_key_read(self, tmp_path): + p = _write_worker( + tmp_path, + "env_read.py", + """\ + import os + + def caller(x): + return os.environ["RUNPOD_API_KEY"] + """, + ) + result = _analyze_cross_calls_ast(p, {"caller"}, set()) + assert result == {"caller": ["RUNPOD_API_KEY"]} + + @pytest.mark.parametrize( + "kwarg", ['id="abc123"', 'image="org/img:1"', '**{"id": "abc123"}'] + ) + def test_detects_inline_client_endpoint(self, tmp_path, kwarg): + p = _write_worker( + tmp_path, + "inline_client.py", + f"""\ + async def caller(x): + return await Endpoint({kwarg}).runsync(x) + """, + ) + result = _analyze_cross_calls_ast(p, {"caller"}, set()) + assert result == {"caller": ["Endpoint"]} + + def test_analysis_error_assumes_key_needed(self, tmp_path): + p = _write_worker(tmp_path, "boom.py", "def caller():\n pass\n") + with patch( + "runpod_flash.cli.commands.build_utils.scanner._remote_call_targets", + side_effect=RuntimeError("boom"), + ): + result = _analyze_cross_calls_ast(p, {"caller"}, set()) + assert result == {"caller": [""]} + + def test_ignores_non_client_endpoint_construction(self, tmp_path): + p = _write_worker( + tmp_path, + "inline_decorator.py", + """\ + def caller(x): + return Endpoint(name="ep", gpu=None) + """, + ) + result = _analyze_cross_calls_ast(p, {"caller"}, set()) + assert result == {} + # --------------------------------------------------------------------------- # RuntimeScanner.discover_remote_functions @@ -388,6 +478,533 @@ def caller(x): assert caller.calls_remote_functions is True assert "helper" in caller.called_remote_functions + def test_client_endpoint_call_detection(self, tmp_path): + """calling a client-mode Endpoint counts as a remote call.""" + _write_worker( + tmp_path, + "wrapper.py", + """\ + from runpod_flash import Endpoint + vllm = Endpoint(name="vllm-backend", image="vllm/vllm-openai:latest") + + @Endpoint(name="wrapper", cpu="cpu3c-1-2") + async def caller(x): + return await vllm.runsync(x) + + @Endpoint(name="isolated", cpu="cpu3c-1-2") + async def isolated(x): + return x + """, + ) + scanner = RuntimeScanner(tmp_path) + functions = scanner.discover_remote_functions() + caller = next(f for f in functions if f.function_name == "caller") + isolated = next(f for f in functions if f.function_name == "isolated") + assert caller.calls_remote_functions is True + assert caller.called_remote_functions == ["vllm"] + assert isolated.calls_remote_functions is False + + def test_helper_in_other_file_detection(self, tmp_path): + """a helper in another project file that calls an endpoint counts.""" + _write_worker( + tmp_path, + "clients.py", + """\ + from runpod_flash import Endpoint + vllm = Endpoint(name="vllm-backend", image="vllm/vllm-openai:latest") + + async def generate(x): + return await vllm.runsync(x) + + def double(x): + return x * 2 + """, + ) + _write_worker( + tmp_path, + "wrapper.py", + """\ + from runpod_flash import Endpoint + import clients + from clients import double + + @Endpoint(name="via-helper", cpu="cpu3c-1-2") + async def via_helper(x): + return await clients.generate(x) + + @Endpoint(name="pure-helper", cpu="cpu3c-1-2") + async def pure_helper(x): + return double(x) + """, + ) + scanner = RuntimeScanner(tmp_path) + functions = scanner.discover_remote_functions() + via_helper = next(f for f in functions if f.function_name == "via_helper") + pure_helper = next(f for f in functions if f.function_name == "pure_helper") + assert via_helper.called_remote_functions == ["generate", "clients.generate"] + assert pure_helper.calls_remote_functions is False + + def test_remote_function_via_module_alias_or_value(self, tmp_path): + """another @Endpoint reached through a module, an alias, or as a value.""" + _write_worker( + tmp_path, + "tasks.py", + """\ + from runpod_flash import Endpoint + + @Endpoint(name="square", cpu="cpu3c-1-2") + async def square(x): + return x * x + """, + ) + _write_worker( + tmp_path, + "app.py", + """\ + from runpod_flash import Endpoint + import tasks + from tasks import square as sq + + @Endpoint(name="via-module", cpu="cpu3c-1-2") + async def via_module(x): + return await tasks.square(x) + + @Endpoint(name="aliased", cpu="cpu3c-1-2") + async def aliased(x): + return await sq(x) + + @Endpoint(name="passed", cpu="cpu3c-1-2") + async def passed(xs): + return list(map(tasks.square, xs)) + """, + ) + scanner = RuntimeScanner(tmp_path) + by_name = { + f.function_name: f + for f in scanner.discover_remote_functions() + if f.function_name != "square" + } + assert by_name["via_module"].called_remote_functions == ["tasks.square"] + assert by_name["aliased"].called_remote_functions == ["sq"] + assert by_name["passed"].called_remote_functions == ["tasks.square"] + + def test_function_local_imports(self, tmp_path): + """imports inside the endpoint function are resolved too.""" + _write_worker( + tmp_path, + "local_imports.py", + """\ + from runpod_flash import Endpoint + + @Endpoint(name="local-alias", cpu="cpu3c-1-2") + async def local_alias(x): + from runpod_flash import Endpoint as EP + return await EP(id="abc123").runsync(x) + + @Endpoint(name="local-sdk", cpu="cpu3c-1-2") + async def local_sdk(x): + import runpod + return runpod.Endpoint("abc123").run_sync(x) + + @Endpoint(name="local-stdlib", cpu="cpu3c-1-2") + async def local_stdlib(x): + import json + return json.dumps(x) + """, + ) + scanner = RuntimeScanner(tmp_path) + by_name = {f.function_name: f for f in scanner.discover_remote_functions()} + assert by_name["local_alias"].called_remote_functions == ["Endpoint"] + assert "runpod" in by_name["local_sdk"].called_remote_functions + assert by_name["local_stdlib"].calls_remote_functions is False + + def test_cyclic_helpers_shared_by_two_endpoints(self, tmp_path): + """a helper first reached mid-cycle is still detected for later callers.""" + _write_worker( + tmp_path, + "cyclic.py", + """\ + from runpod_flash import Endpoint + vllm = Endpoint(name="vllm-backend", image="vllm/vllm-openai:latest") + + async def helper_a(x, depth=0): + if depth: + await helper_b(x, depth - 1) + return await vllm.runsync(x) + + async def helper_b(x, depth=0): + return await helper_a(x, depth) + + @Endpoint(name="first", cpu="cpu3c-1-2") + async def first(x): + return await helper_a(x, 1) + + @Endpoint(name="second", cpu="cpu3c-1-2") + async def second(x): + return await helper_b(x, 1) + """, + ) + scanner = RuntimeScanner(tmp_path) + by_name = {f.function_name: f for f in scanner.discover_remote_functions()} + assert by_name["first"].calls_remote_functions is True + assert by_name["second"].calls_remote_functions is True + + def test_helpers_outside_scanned_files_are_not_analyzed(self, tmp_path): + """library code in an in-project .venv is not treated as a helper.""" + venv = tmp_path / ".venv" + venv.mkdir() + (venv / "fakelib.py").write_text( + "import os\n\ndef read_key():\n return os.environ.get('RUNPOD_API_KEY')\n", + encoding="utf-8", + ) + _write_worker( + tmp_path, + "app.py", + """\ + import sys + from pathlib import Path + sys.path.insert(0, str(Path(__file__).parent / ".venv")) + from fakelib import read_key + from runpod_flash import Endpoint + + @Endpoint(name="lib-call", cpu="cpu3c-1-2") + async def lib_call(x): + read_key() + return x + """, + ) + scanner = RuntimeScanner(tmp_path) + try: + lib_call = scanner.discover_remote_functions()[0] + finally: + sys.path[:] = [x for x in sys.path if x != str(venv)] + sys.modules.pop("fakelib", None) + assert scanner.import_errors == {} + assert lib_call.calls_remote_functions is False + + def test_function_local_project_imports(self, tmp_path): + """sibling modules imported inside the endpoint resolve after scanning.""" + _write_worker( + tmp_path, + "local_tasks.py", + """\ + from runpod_flash import Endpoint + + @Endpoint(name="square", cpu="cpu3c-1-2") + async def square(x): + return x * x + """, + ) + _write_worker( + tmp_path, + "local_app.py", + """\ + from runpod_flash import Endpoint + + @Endpoint(name="from-import", cpu="cpu3c-1-2") + async def from_import(x): + from local_tasks import square as sq + return await sq(x) + + @Endpoint(name="module-import", cpu="cpu3c-1-2") + async def module_import(x): + import local_tasks + return await local_tasks.square(x) + """, + ) + scanner = RuntimeScanner(tmp_path) + by_name = {f.function_name: f for f in scanner.discover_remote_functions()} + assert by_name["from_import"].called_remote_functions == ["sq"] + assert by_name["module_import"].called_remote_functions == [ + "local_tasks.square" + ] + + def test_relative_and_package_qualified_imports(self, tmp_path): + """subpackage endpoints reached relatively or through the package path.""" + (tmp_path / "rel_pkg").mkdir() + _write_worker( + tmp_path, + "rel_pkg/tasks.py", + """\ + from runpod_flash import Endpoint + + @Endpoint(name="square", cpu="cpu3c-1-2") + async def square(x): + return x * x + """, + ) + _write_worker( + tmp_path, + "rel_pkg/app.py", + """\ + from runpod_flash import Endpoint + + @Endpoint(name="relative", cpu="cpu3c-1-2") + async def relative(x): + from .tasks import square as sq + return await sq(x) + + @Endpoint(name="qualified", cpu="cpu3c-1-2") + async def qualified(x): + import rel_pkg.tasks + return await rel_pkg.tasks.square(x) + + @Endpoint(name="missing", cpu="cpu3c-1-2") + async def missing(x): + from .nowhere import thing + return thing(x) + """, + ) + scanner = RuntimeScanner(tmp_path) + by_name = {f.function_name: f for f in scanner.discover_remote_functions()} + assert by_name["relative"].called_remote_functions == ["sq"] + assert by_name["qualified"].called_remote_functions == ["rel_pkg.tasks.square"] + # a project import that cannot be resolved errs toward the key + assert by_name["missing"].called_remote_functions == [""] + + def test_unresolved_import_of_failed_project_module(self, tmp_path): + """a local import of a project file that failed to import errs toward the key.""" + _write_worker(tmp_path, "broken_dep.py", "raise RuntimeError('boom')\n") + _write_worker( + tmp_path, + "uses_broken.py", + """\ + from runpod_flash import Endpoint + + @Endpoint(name="uses-broken", cpu="cpu3c-1-2") + async def uses_broken(x): + from broken_dep import helper + return helper(x) + """, + ) + scanner = RuntimeScanner(tmp_path) + uses_broken = next( + f + for f in scanner.discover_remote_functions() + if f.function_name == "uses_broken" + ) + assert uses_broken.called_remote_functions == [""] + + def test_helper_in_package_init(self, tmp_path): + """helpers defined in a package __init__.py are analyzed.""" + (tmp_path / "init_utils").mkdir() + _write_worker( + tmp_path, + "init_utils/__init__.py", + """\ + from runpod_flash import Endpoint + vllm = Endpoint(id="abc123") + + def call_model(x): + return vllm.runsync(x) + """, + ) + _write_worker( + tmp_path, + "init_app.py", + """\ + from runpod_flash import Endpoint + import init_utils + + @Endpoint(name="via-init", cpu="cpu3c-1-2") + async def via_init(x): + return init_utils.call_model(x) + """, + ) + scanner = RuntimeScanner(tmp_path) + try: + via_init = scanner.discover_remote_functions()[0] + finally: + sys.modules.pop("init_utils", None) + assert via_init.called_remote_functions == [ + "call_model", + "init_utils.call_model", + ] + + def test_project_classes_and_instances(self, tmp_path): + """methods, static/class methods, and instance attributes are followed.""" + _write_worker( + tmp_path, + "cls_clients.py", + """\ + from runpod_flash import Endpoint + + class ModelClient: + def generate(self, x): + return Endpoint(id="abc123").runsync(x) + + class Holder: + def __init__(self, ep): + self.ep = ep + + class Helpers: + @staticmethod + def call(x): + return Endpoint(id="abc123").runsync(x) + + class Pure: + def double(self, x): + return x * 2 + + client = ModelClient() + holder = Holder(Endpoint(id="abc123")) + nested = Holder(Holder(Endpoint(id="abc123"))) + pure = Pure() + pure.self_ref = pure + """, + ) + _write_worker( + tmp_path, + "cls_app.py", + """\ + from runpod_flash import Endpoint + from cls_clients import Helpers, client, holder, nested, pure + + @Endpoint(name="method", cpu="cpu3c-1-2") + async def method(x): + return client.generate(x) + + @Endpoint(name="static", cpu="cpu3c-1-2") + async def static(x): + return Helpers.call(x) + + @Endpoint(name="attribute", cpu="cpu3c-1-2") + async def attribute(x): + return holder.ep.runsync(x) + + @Endpoint(name="nested-attribute", cpu="cpu3c-1-2") + async def nested_attribute(x): + return nested.ep.ep.runsync(x) + + @Endpoint(name="pure-method", cpu="cpu3c-1-2") + async def pure_method(x): + return pure.double(x) + """, + ) + scanner = RuntimeScanner(tmp_path) + by_name = {f.function_name: f for f in scanner.discover_remote_functions()} + assert by_name["method"].called_remote_functions == ["client"] + assert by_name["static"].called_remote_functions == ["Helpers"] + assert by_name["attribute"].called_remote_functions == ["holder"] + assert by_name["nested_attribute"].called_remote_functions == ["nested"] + assert by_name["pure_method"].calls_remote_functions is False + + def test_wrapped_and_contained_client_endpoints(self, tmp_path): + """bound methods, partials, containers, and cached factories are followed.""" + _write_worker( + tmp_path, + "wrapped.py", + """\ + import functools + from runpod_flash import Endpoint + + vllm = Endpoint(id="abc123") + run_vllm = vllm.runsync + CLIENTS = {"vllm": vllm} + call = functools.partial(vllm.runsync, timeout=5) + + @functools.lru_cache + def client(): + return Endpoint(id="abc123") + + @Endpoint(name="bound", cpu="cpu3c-1-2") + async def bound(x): + return await run_vllm(x) + + @Endpoint(name="container", cpu="cpu3c-1-2") + async def container(x): + return await CLIENTS["vllm"].runsync(x) + + @Endpoint(name="partial", cpu="cpu3c-1-2") + async def partial(x): + return await call(x) + + @Endpoint(name="cached", cpu="cpu3c-1-2") + async def cached(x): + return await client().runsync(x) + + @Endpoint(name="plain-dict", cpu="cpu3c-1-2") + async def plain_dict(x): + return {"a": 1}.get(x) + """, + ) + scanner = RuntimeScanner(tmp_path) + by_name = {f.function_name: f for f in scanner.discover_remote_functions()} + assert by_name["bound"].called_remote_functions == ["run_vllm"] + assert by_name["container"].called_remote_functions == ["CLIENTS"] + assert by_name["partial"].called_remote_functions == ["call"] + assert by_name["cached"].called_remote_functions == ["client"] + assert by_name["plain_dict"].calls_remote_functions is False + + def test_wrapped_and_contained_project_helpers(self, tmp_path): + """partials and containers around a project helper are followed.""" + _write_worker( + tmp_path, + "wrapped_helpers.py", + """\ + import functools + from runpod_flash import Endpoint + + vllm = Endpoint(id="abc123") + + def call_model(x, timeout=None): + return vllm.runsync(x) + + def double(x): + return x * 2 + + call = functools.partial(call_model, timeout=5) + HANDLERS = {"gen": call_model} + PURE = {"double": double} + + @Endpoint(name="helper-partial", cpu="cpu3c-1-2") + async def helper_partial(x): + return call(x) + + @Endpoint(name="helper-dict", cpu="cpu3c-1-2") + async def helper_dict(x): + return HANDLERS["gen"](x) + + @Endpoint(name="helper-map", cpu="cpu3c-1-2") + async def helper_map(xs): + return list(map(call_model, xs)) + + @Endpoint(name="pure-dict", cpu="cpu3c-1-2") + async def pure_dict(x): + return PURE["double"](x) + """, + ) + scanner = RuntimeScanner(tmp_path) + by_name = {f.function_name: f for f in scanner.discover_remote_functions()} + assert by_name["helper_partial"].called_remote_functions == ["call"] + assert by_name["helper_dict"].called_remote_functions == ["HANDLERS"] + assert by_name["helper_map"].called_remote_functions == ["call_model"] + assert by_name["pure_dict"].calls_remote_functions is False + + def test_unparseable_helper_assumes_key_needed(self, tmp_path): + """a helper whose source cannot be read or parsed errs toward the key.""" + _write_worker( + tmp_path, + "opaque_helper.py", + """\ + from runpod_flash import Endpoint + + def helper(x): + return x + + @Endpoint(name="via-helper", cpu="cpu3c-1-2") + async def via_helper(x): + return helper(x) + """, + ) + scanner = RuntimeScanner(tmp_path) + with patch( + "runpod_flash.cli.commands.build_utils.scanner.inspect.getsource", + side_effect=OSError("could not get source code"), + ): + via_helper = scanner.discover_remote_functions()[0] + assert via_helper.called_remote_functions == ["helper"] + def test_sys_path_cleanup(self, tmp_path): """project root is removed from sys.path after scanning.""" root_str = str(tmp_path) diff --git a/tests/unit/cli/test_deploy.py b/tests/unit/cli/test_deploy.py index fd20a5211..ae046b1ab 100644 --- a/tests/unit/cli/test_deploy.py +++ b/tests/unit/cli/test_deploy.py @@ -643,6 +643,55 @@ def test_qb_endpoints_shown(self, patched_console): assert "my_qb" in output assert "/runsync" in output + def test_lists_endpoints_that_receive_api_key(self, patched_console): + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={ + "caller": "https://a.api.runpod.ai", + "isolated": "https://b.api.runpod.ai", + }, + resources={ + "caller": {"makes_remote_calls": True}, + "isolated": {"makes_remote_calls": False}, + }, + routes={}, + ) + output = self._collect_output(patched_console) + assert ( + "RUNPOD_API_KEY set by Flash on endpoints that call other endpoints: caller" + in (output) + ) + assert "caller, isolated" not in output + + def test_api_key_line_skips_user_managed_key(self, patched_console): + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={"caller": "https://a.api.runpod.ai"}, + resources={ + "caller": { + "makes_remote_calls": True, + "env": {"RUNPOD_API_KEY": ""}, + } + }, + routes={}, + ) + assert "RUNPOD_API_KEY added" not in self._collect_output(patched_console) + + def test_no_api_key_line_when_no_endpoint_receives_it(self, patched_console): + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={"isolated": "https://b.api.runpod.ai"}, + resources={"isolated": {"makes_remote_calls": False}}, + routes={}, + ) + assert "RUNPOD_API_KEY added" not in self._collect_output(patched_console) + def test_routes_grouped_under_lb_endpoint(self, patched_console): from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance diff --git a/tests/unit/cli/utils/test_deployment.py b/tests/unit/cli/utils/test_deployment.py index 334a8e24e..d2880bbd9 100644 --- a/tests/unit/cli/utils/test_deployment.py +++ b/tests/unit/cli/utils/test_deployment.py @@ -842,3 +842,209 @@ async def test_source_fingerprint_skipped_for_client_mode_resources(tmp_path): client_env = local_manifest["resources"]["vllm-server"].get("env", {}) assert "_FLASH_SOURCE_FINGERPRINT" not in client_env + + +def _unchanged_manifest(**worker_fields): + worker = { + "resource_type": "LiveServerless", + "config": "same", + "endpoint_id": "endpoint-123", + **worker_fields, + } + return ( + {"resources": {"worker": dict(worker)}, "resources_endpoints": {}}, + { + "resources": {"worker": dict(worker)}, + "resources_endpoints": {"worker": "https://worker.api.runpod.ai"}, + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "worker_fields,expect_check", + [ + ({"makes_remote_calls": False}, True), + ({"makes_remote_calls": True}, False), + ({}, False), + ({"makes_remote_calls": False, "env": {"RUNPOD_API_KEY": ""}}, False), + ], + ids=["local-only", "remote-caller", "unknown", "user-managed-key"], +) +async def test_unchanged_local_only_endpoint_gets_stale_key_check( + tmp_path, monkeypatch, worker_fields, expect_check +): + """Unchanged resources skip provisioning, so stale keys are cleaned here.""" + monkeypatch.setenv("RUNPOD_API_KEY", "deployer-key") + (tmp_path / ".flash").mkdir() + local_manifest, state_manifest = _unchanged_manifest(**worker_fields) + + app = AsyncMock() + app.get_build_manifest = AsyncMock(return_value=state_manifest) + app.update_build_manifest = AsyncMock() + + with ( + patch("pathlib.Path.cwd", return_value=tmp_path), + patch("runpod_flash.cli.utils.deployment.ResourceManager"), + patch( + "runpod_flash.cli.utils.deployment._remove_stale_api_keys", + new=AsyncMock(), + ) as mock_cleanup, + ): + await reconcile_and_provision_resources(app, "build-123", "dev", local_manifest) + + if expect_check: + mock_cleanup.assert_awaited_once_with([("worker", "endpoint-123")]) + else: + mock_cleanup.assert_not_awaited() + + +class TestRemoveStaleApiKey: + """_remove_stale_api_key drops only RUNPOD_API_KEY from the live template.""" + + @staticmethod + def _client(env, template_ids=None): + client = AsyncMock() + client.get_endpoint_template_ids = AsyncMock( + return_value={"endpoint-123": "tpl-1"} + if template_ids is None + else template_ids + ) + client.get_template = AsyncMock( + return_value={ + "id": "tpl-1", + "name": "worker-tpl", + "imageName": "runpod/flash:latest", + "containerDiskInGb": 64, + "dockerArgs": "", + "readme": None, + "env": env, + } + ) + client.update_template = AsyncMock(return_value={}) + return client + + @staticmethod + async def _cleanup(client, endpoints=(("worker", "endpoint-123"),)): + from runpod_flash.cli.utils.deployment import _remove_stale_api_keys + + with patch("runpod_flash.core.api.runpod.RunpodGraphQLClient") as cls: + cls.return_value.__aenter__.return_value = client + cls.return_value.__aexit__.return_value = None + await _remove_stale_api_keys(list(endpoints)) + + @pytest.mark.asyncio + async def test_removes_key_and_keeps_everything_else(self): + client = self._client( + [ + {"key": "PORT", "value": "8080"}, + {"key": "RUNPOD_API_KEY", "value": "old-key"}, + ] + ) + + await self._cleanup(client) + + client.get_template.assert_awaited_once_with("tpl-1") + payload = client.update_template.call_args.args[0] + assert payload == { + "id": "tpl-1", + "name": "worker-tpl", + "imageName": "runpod/flash:latest", + "containerDiskInGb": 64, + "dockerArgs": "", + "env": [{"key": "PORT", "value": "8080"}], + "volumeInGb": 0, + } + + @pytest.mark.asyncio + async def test_no_update_when_key_absent(self): + client = self._client([{"key": "PORT", "value": "8080"}]) + + await self._cleanup(client) + + client.update_template.assert_not_awaited() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "template_ids", + [{}, {"endpoint-123": None}], + ids=["endpoint-deleted", "no-template"], + ) + async def test_no_update_without_template(self, template_ids): + client = self._client([], template_ids=template_ids) + + await self._cleanup(client) + + client.get_template.assert_not_awaited() + + @pytest.mark.asyncio + async def test_shared_template_keeps_key(self): + client = self._client( + [{"key": "RUNPOD_API_KEY", "value": "old-key"}], + template_ids={"endpoint-123": "tpl-1", "caller": "tpl-1"}, + ) + + await self._cleanup(client) + + client.get_template.assert_not_awaited() + client.update_template.assert_not_awaited() + + @pytest.mark.asyncio + async def test_one_lookup_for_all_endpoints(self): + client = self._client( + [{"key": "RUNPOD_API_KEY", "value": "old-key"}], + template_ids={"endpoint-1": "tpl-1", "endpoint-2": "tpl-2"}, + ) + + await self._cleanup(client, [("a", "endpoint-1"), ("b", "endpoint-2")]) + + client.get_endpoint_template_ids.assert_awaited_once() + assert client.update_template.await_count == 2 + + @pytest.mark.asyncio + async def test_one_failure_does_not_stop_the_others(self): + client = self._client( + [{"key": "RUNPOD_API_KEY", "value": "old-key"}], + template_ids={"endpoint-1": "tpl-1", "endpoint-2": "tpl-2"}, + ) + client.update_template = AsyncMock(side_effect=[RuntimeError("boom"), {}]) + + await self._cleanup(client, [("a", "endpoint-1"), ("b", "endpoint-2")]) + + assert client.update_template.await_count == 2 + + @pytest.mark.asyncio + async def test_slow_api_is_bounded(self, caplog): + import asyncio + import time + + client = self._client([]) + + async def hang(): + await asyncio.sleep(10) + return {"endpoint-123": "tpl-1"} + + client.get_endpoint_template_ids = AsyncMock(side_effect=hang) + start = time.monotonic() + with ( + patch( + "runpod_flash.cli.utils.deployment.STALE_KEY_CLEANUP_TIMEOUT_SECONDS", + 0.05, + ), + caplog.at_level("WARNING"), + ): + await self._cleanup(client) + + assert time.monotonic() - start < 2 + assert "Stopped checking endpoints" in caplog.text + client.get_template.assert_not_awaited() + + @pytest.mark.asyncio + async def test_errors_never_fail_the_deploy(self): + from runpod_flash.cli.utils.deployment import _remove_stale_api_keys + + with patch( + "runpod_flash.core.api.runpod.RunpodGraphQLClient", + side_effect=RuntimeError("no network"), + ): + await _remove_stale_api_keys([("worker", "endpoint-123")]) diff --git a/tests/unit/core/api/test_runpod_graphql_extended.py b/tests/unit/core/api/test_runpod_graphql_extended.py index 2d6cb734f..a9d005140 100644 --- a/tests/unit/core/api/test_runpod_graphql_extended.py +++ b/tests/unit/core/api/test_runpod_graphql_extended.py @@ -352,3 +352,31 @@ async def test_execute_rest_network_error(self): with pytest.raises(Exception, match="HTTP request failed"): await client._execute_rest("GET", "https://example.com/api") + + +class TestGetEndpointTemplateIds: + """get_endpoint_template_ids maps every endpoint to its template.""" + + @pytest.mark.asyncio + async def test_maps_endpoints_to_templates(self): + client = RunpodGraphQLClient(api_key="test") + with patch.object(client, "_execute_graphql", new_callable=AsyncMock) as mock: + mock.return_value = { + "myself": { + "endpoints": [ + {"id": "ep-1", "templateId": "tpl-1"}, + {"id": "ep-2", "templateId": None}, + ] + } + } + assert await client.get_endpoint_template_ids() == { + "ep-1": "tpl-1", + "ep-2": None, + } + + @pytest.mark.asyncio + async def test_no_endpoints(self): + client = RunpodGraphQLClient(api_key="test") + with patch.object(client, "_execute_graphql", new_callable=AsyncMock) as mock: + mock.return_value = {"myself": {"endpoints": None}} + assert await client.get_endpoint_template_ids() == {} diff --git a/tests/unit/resources/test_serverless.py b/tests/unit/resources/test_serverless.py index 0db5b1a2c..def9362f7 100644 --- a/tests/unit/resources/test_serverless.py +++ b/tests/unit/resources/test_serverless.py @@ -3513,3 +3513,130 @@ def test_create_new_template_env_not_in_fields_set_cpu_endpoint(self): template = resource._create_new_template() assert "env" not in template.model_fields_set + + +class TestStaleApiKeyCleanup: + """A RUNPOD_API_KEY left on a live template by an earlier deploy is dropped + once the build manifest says the endpoint makes no remote calls.""" + + @staticmethod + async def _update(new_env, makes_remote_calls, template_ids=None, old_env=None): + old_resource = ServerlessEndpoint( + name="stale-key", + imageName="test:latest", + env=old_env or {"LOG_LEVEL": "INFO"}, + flashboot=False, + ) + old_resource.id = "ep-stale" + old_resource.templateId = "tpl-stale" + + new_resource = ServerlessEndpoint( + name="stale-key", + imageName="test:latest", + env=new_env, + flashboot=False, + workersMax=5, + ) + new_resource._makes_remote_calls = makes_remote_calls + + mock_client = AsyncMock() + mock_client.save_endpoint = AsyncMock( + return_value={ + "id": "ep-stale", + "name": "stale-key", + "templateId": "tpl-stale", + "gpuIds": "AMPERE_48", + "allowedCudaVersions": "", + } + ) + mock_client.update_template = AsyncMock(return_value={}) + mock_client.get_template = AsyncMock( + return_value={ + "env": [ + {"key": "LOG_LEVEL", "value": "INFO"}, + {"key": "PORT", "value": "8080"}, + {"key": "RUNPOD_API_KEY", "value": "old-key"}, + ] + } + ) + if isinstance(template_ids, Exception): + mock_client.get_endpoint_template_ids = AsyncMock(side_effect=template_ids) + else: + mock_client.get_endpoint_template_ids = AsyncMock( + return_value=template_ids or {"ep-stale": "tpl-stale"} + ) + + with patch( + "runpod_flash.core.resources.serverless.RunpodGraphQLClient" + ) as mock_client_class: + mock_client_class.return_value.__aenter__.return_value = mock_client + mock_client_class.return_value.__aexit__.return_value = None + with patch.object( + ServerlessResource, + "_ensure_network_volume_deployed", + new=AsyncMock(), + ): + with patch.dict(os.environ, {"RUNPOD_API_KEY": "deployer-key"}): + await old_resource.update(new_resource) + + payload = mock_client.update_template.call_args.args[0] + return {e["key"]: e["value"] for e in payload["env"]} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "new_env", + [{"LOG_LEVEL": "INFO"}, {"LOG_LEVEL": "DEBUG"}], + ids=["env-unchanged", "env-changed"], + ) + async def test_drops_stale_key_without_remote_calls(self, new_env): + env = await self._update(new_env, makes_remote_calls=False) + + assert "RUNPOD_API_KEY" not in env + assert env["PORT"] == "8080" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "new_env", + [{"LOG_LEVEL": "INFO"}, {"LOG_LEVEL": "DEBUG"}], + ids=["env-unchanged", "env-changed"], + ) + async def test_keeps_key_with_remote_calls(self, new_env): + env = await self._update(new_env, makes_remote_calls=True) + + assert "RUNPOD_API_KEY" in env + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "new_env", + [ + {"LOG_LEVEL": "INFO", "RUNPOD_API_KEY": "user-key"}, + {"LOG_LEVEL": "DEBUG", "RUNPOD_API_KEY": "user-key"}, + ], + ids=["env-unchanged", "env-changed"], + ) + async def test_keeps_user_set_key_without_remote_calls(self, new_env): + old_env = {"LOG_LEVEL": "INFO", "RUNPOD_API_KEY": "user-key"} + env = await self._update(new_env, makes_remote_calls=False, old_env=old_env) + + assert "RUNPOD_API_KEY" in env + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "new_env", + [{"LOG_LEVEL": "INFO"}, {"LOG_LEVEL": "DEBUG"}], + ids=["env-unchanged", "env-changed"], + ) + @pytest.mark.parametrize( + "template_ids", + [ + {"ep-stale": "tpl-stale", "ep-caller": "tpl-stale"}, + RuntimeError("api down"), + ], + ids=["shared-template", "lookup-failed"], + ) + async def test_keeps_key_when_template_may_be_shared(self, new_env, template_ids): + env = await self._update( + new_env, makes_remote_calls=False, template_ids=template_ids + ) + + assert env["RUNPOD_API_KEY"] == "old-key" diff --git a/tests/unit/runtime/test_resource_provisioner.py b/tests/unit/runtime/test_resource_provisioner.py index 78417a0a4..a8b12e79f 100644 --- a/tests/unit/runtime/test_resource_provisioner.py +++ b/tests/unit/runtime/test_resource_provisioner.py @@ -231,6 +231,43 @@ def test_create_resource_skips_api_key_when_no_remote_calls(self): assert "RUNPOD_API_KEY" not in resource.env + @pytest.mark.parametrize("makes_remote_calls", [True, False]) + def test_template_injection_follows_manifest( + self, makes_remote_calls, tmp_path, monkeypatch + ): + """Template injection uses the manifest value, not a file lookup. + + flash deploy runs from the project root where the manifest lives at + .flash/flash_manifest.json, so a cwd lookup would miss it and fall + back to injecting the key everywhere. + """ + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("RUNPOD_API_KEY", "test-api-key-secret") + resource = create_resource_from_manifest( + "worker", + { + "resource_type": "Endpoint", + "gpuIds": "AMPERE_48", + "makes_remote_calls": makes_remote_calls, + }, + ) + + assert resource._check_makes_remote_calls() is makes_remote_calls + resource._inject_runtime_template_vars() + template_keys = {kv.key for kv in resource.template.env or []} + assert ("RUNPOD_API_KEY" in template_keys) is makes_remote_calls + + def test_missing_manifest_flag_keeps_file_lookup(self, tmp_path, monkeypatch): + """Without the field, the resource keeps the manifest file fallback.""" + monkeypatch.chdir(tmp_path) + resource = create_resource_from_manifest( + "worker", + {"resource_type": "ServerlessResource", "imageName": "img:latest"}, + ) + + assert resource._makes_remote_calls is None + assert resource._check_makes_remote_calls() is True + def test_create_resource_injects_flash_environment_id_when_makes_remote_calls(self): """Test FLASH_ENVIRONMENT_ID injected when makes_remote_calls and flash_environment_id passed.""" resource_name = "caller_worker" diff --git a/tests/unit/test_endpoint_client.py b/tests/unit/test_endpoint_client.py index 6b85a2e57..e4e0c3847 100644 --- a/tests/unit/test_endpoint_client.py +++ b/tests/unit/test_endpoint_client.py @@ -741,3 +741,55 @@ def test_in_all(self): import runpod_flash assert "EndpointJob" in runpod_flash.__all__ + + +class TestMissingWorkerKeyHint: + """a 401 on a worker without RUNPOD_API_KEY explains why the key is missing.""" + + @staticmethod + def _response(status: int): + import httpx + + request = httpx.Request("POST", "https://api.runpod.ai/v2/abc/runsync") + return httpx.Response(status, request=request) + + @pytest.mark.parametrize( + "env,has_key,expect_hint", + [ + ({"RUNPOD_ENDPOINT_ID": "ep-1"}, False, True), + ({"RUNPOD_ENDPOINT_ID": "ep-1"}, True, False), + ({}, False, False), + ], + ids=["worker-no-key", "worker-with-key", "local"], + ) + def test_401_hint(self, monkeypatch, env, has_key, expect_hint): + import httpx + + from runpod_flash.endpoint import _MISSING_WORKER_KEY_HINT, _raise_for_status + + monkeypatch.delenv("RUNPOD_ENDPOINT_ID", raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + monkeypatch.setattr( + "runpod_flash.core.credentials.get_api_key", + lambda: "key" if has_key else None, + ) + + with pytest.raises(httpx.HTTPStatusError) as exc_info: + _raise_for_status(self._response(401)) + + assert (_MISSING_WORKER_KEY_HINT in str(exc_info.value)) is expect_hint + assert exc_info.value.response.status_code == 401 + + def test_other_errors_unchanged(self, monkeypatch): + import httpx + + from runpod_flash.endpoint import _MISSING_WORKER_KEY_HINT, _raise_for_status + + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep-1") + monkeypatch.setattr("runpod_flash.core.credentials.get_api_key", lambda: None) + + with pytest.raises(httpx.HTTPStatusError) as exc_info: + _raise_for_status(self._response(500)) + + assert _MISSING_WORKER_KEY_HINT not in str(exc_info.value)