From 99eaceae49c7b7de4bbf7dc60c4c0b5ec73483c3 Mon Sep 17 00:00:00 2001 From: Justin Date: Thu, 15 Oct 2026 09:20:02 -0400 Subject: [PATCH 01/13] fix(deploy): only inject RUNPOD_API_KEY into endpoints that make remote calls flash deploy looked for flash_manifest.json in the wrong place, missed it, and fell back to giving every endpoint the deployer's API key. - carry the manifest's makes_remote_calls onto the deployed resource - detect calls through client-mode Endpoints (id= or image=) at build time - drop a stale RUNPOD_API_KEY from live templates when no longer needed --- .../cli/commands/build_utils/scanner.py | 55 ++++++++++-- src/runpod_flash/core/resources/serverless.py | 20 +++++ .../runtime/resource_provisioner.py | 5 ++ .../build_utils/test_runtime_scanner.py | 63 ++++++++++++++ tests/unit/resources/test_serverless.py | 85 +++++++++++++++++++ .../unit/runtime/test_resource_provisioner.py | 37 ++++++++ 6 files changed, 260 insertions(+), 5 deletions(-) diff --git a/src/runpod_flash/cli/commands/build_utils/scanner.py b/src/runpod_flash/cli/commands/build_utils/scanner.py index c9a383b7..2a6ac326 100644 --- a/src/runpod_flash/cli/commands/build_utils/scanner.py +++ b/src/runpod_flash/cli/commands/build_utils/scanner.py @@ -407,16 +407,48 @@ def _find_remote_decorated(module: Any) -> Dict[str, Any]: return results +def _find_client_endpoint_names(module: Any) -> Set[str]: + """find module-level names bound to client-mode Endpoints (id= or image=).""" + names: Set[str] = set() + for name in dir(module): + try: + obj = getattr(module, name) + except Exception: + continue + if isinstance(obj, Endpoint) and obj.is_client: + names.add(name) + return names + + +def _is_client_endpoint_construction(call: ast.Call) -> bool: + """match Endpoint(id=...) or Endpoint(image=...) built inline.""" + func = call.func + if isinstance(func, ast.Name): + callee = func.id + elif isinstance(func, ast.Attribute): + callee = func.attr + else: + return False + if callee != "Endpoint": + return False + return any(kw.arg in ("id", "image") for kw in call.keywords) + + def _analyze_cross_calls_ast( file_path: Path, function_names: Set[str], remote_function_names: Set[str], + client_endpoint_names: Optional[Set[str]] = None, ) -> Dict[str, List[str]]: - """find which functions call other @remote functions via AST. + """find which functions call other endpoints via AST. returns a dict of function_name -> list of called remote function names. - only matches direct calls (foo()), not attribute calls (obj.foo()). + matches direct calls to @remote functions (foo(), not obj.foo()), any + reference to a module-level client-mode Endpoint (vllm.runsync(...)), and + inline client construction (Endpoint(id="abc").runsync(...)). client-mode + Endpoints sign requests with RUNPOD_API_KEY, so the worker needs the key. """ + client_names = client_endpoint_names or set() result: Dict[str, List[str]] = {} try: tree = ast.parse(file_path.read_text(encoding="utf-8")) @@ -431,10 +463,16 @@ def _analyze_cross_calls_ast( called: List[str] = [] for child in ast.walk(node): + target: Optional[str] = None 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) + target = child.func.id + if isinstance(child, ast.Call) and _is_client_endpoint_construction(child): + target = "Endpoint" + elif isinstance(child, ast.Name) and child.id in client_names: + target = child.id + if target is not None and target not in called: + called.append(target) if called: result[node.name] = called @@ -479,6 +517,7 @@ def discover_remote_functions(self) -> List[RemoteFunctionMetadata]: functions: List[RemoteFunctionMetadata] = [] seen_functions: Set[str] = set() failed_files: List[Path] = [] + client_names_by_file: Dict[Path, Set[str]] = {} try: for py_file in py_files: @@ -498,6 +537,7 @@ def discover_remote_functions(self) -> List[RemoteFunctionMetadata]: module, module_path, py_file, seen_functions ) functions.extend(file_functions) + client_names_by_file[py_file] = _find_client_endpoint_names(module) # cross-call analysis remote_names = {f.function_name for f in functions} @@ -506,7 +546,12 @@ def discover_remote_functions(self) -> List[RemoteFunctionMetadata]: 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, + client_names_by_file.get(file_path), + ) 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/core/resources/serverless.py b/src/runpod_flash/core/resources/serverless.py index d3a56797..3fc55438 100644 --- a/src/runpod_flash/core/resources/serverless.py +++ b/src/runpod_flash/core/resources/serverless.py @@ -290,6 +290,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 @@ -889,6 +892,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 +1059,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 = not self._check_makes_remote_calls() 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 +1226,19 @@ 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 = ( + not new_config._check_makes_remote_calls() + and "RUNPOD_API_KEY" not in (new_config.env or {}) + ) 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/runtime/resource_provisioner.py b/src/runpod_flash/runtime/resource_provisioner.py index fcecbaff..56d43596 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 10b7f3da..ab82be51 100644 --- a/tests/unit/cli/commands/build_utils/test_runtime_scanner.py +++ b/tests/unit/cli/commands/build_utils/test_runtime_scanner.py @@ -315,6 +315,43 @@ 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) + """, + ) + result = _analyze_cross_calls_ast(p, {"caller"}, set(), {"vllm"}) + assert result == {"caller": ["vllm"]} + + @pytest.mark.parametrize("kwarg", ['id="abc123"', 'image="org/img:1"']) + 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_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 +425,32 @@ 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_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/resources/test_serverless.py b/tests/unit/resources/test_serverless.py index 0db5b1a2..80c3beb4 100644 --- a/tests/unit/resources/test_serverless.py +++ b/tests/unit/resources/test_serverless.py @@ -3513,3 +3513,88 @@ 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): + old_resource = ServerlessEndpoint( + name="stale-key", + imageName="test:latest", + env={"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"}, + ] + } + ) + + 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 diff --git a/tests/unit/runtime/test_resource_provisioner.py b/tests/unit/runtime/test_resource_provisioner.py index 78417a0a..a8b12e79 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" From a1aceae2b54ac3ef3d6b0b9a4174c8b8ea65561f Mon Sep 17 00:00:00 2001 From: Justin Date: Thu, 15 Oct 2026 09:33:10 -0400 Subject: [PATCH 02/13] fix(build): detect API key use through helpers, aliases, and the runpod SDK Deploy now trusts the manifest's makes_remote_calls, so a missed detection means a 401 at runtime instead of a silently over-privileged endpoint. Resolve names against the module's globals and flag a function when it: - calls a project helper (in any file) that needs the key - builds a client Endpoint under an aliased import - uses the runpod SDK or reads RUNPOD_API_KEY --- .../cli/commands/build_utils/scanner.py | 180 +++++++++++++----- .../build_utils/test_runtime_scanner.py | 84 +++++++- 2 files changed, 219 insertions(+), 45 deletions(-) diff --git a/src/runpod_flash/cli/commands/build_utils/scanner.py b/src/runpod_flash/cli/commands/build_utils/scanner.py index 2a6ac326..2886d809 100644 --- a/src/runpod_flash/cli/commands/build_utils/scanner.py +++ b/src/runpod_flash/cli/commands/build_utils/scanner.py @@ -11,6 +11,7 @@ import logging import os import sys +import textwrap import threading import types from dataclasses import dataclass, field @@ -407,48 +408,147 @@ def _find_remote_decorated(module: Any) -> Dict[str, Any]: return results -def _find_client_endpoint_names(module: Any) -> Set[str]: - """find module-level names bound to client-mode Endpoints (id= or image=).""" - names: Set[str] = set() - for name in dir(module): - try: - obj = getattr(module, name) - except Exception: - continue - if isinstance(obj, Endpoint) and obj.is_client: - names.add(name) - return names +_API_KEY_ENV = "RUNPOD_API_KEY" +_MISSING = object() -def _is_client_endpoint_construction(call: ast.Call) -> bool: - """match Endpoint(id=...) or Endpoint(image=...) built inline.""" - func = call.func - if isinstance(func, ast.Name): - callee = func.id - elif isinstance(func, ast.Attribute): - callee = func.attr - else: +def _resolve(expr: ast.expr, namespace: Dict[str, Any]) -> Any: + """resolve a Name, or an attribute of a module (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) + if isinstance(base, types.ModuleType): + return getattr(base, expr.attr, _MISSING) + 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 - if callee != "Endpoint": + return isinstance(name, str) and (name == "runpod" or name.startswith("runpod.")) + + +def _project_function(obj: Any, project_dir: Optional[Path]) -> Any: + """return obj if it is a plain function defined inside the project.""" + if project_dir is None or not isinstance(obj, types.FunctionType): + return None + if hasattr(obj, "__remote_config__"): + return None + try: + source_file = inspect.getsourcefile(obj) + except TypeError: + return None + if source_file is None: + return None + try: + Path(source_file).resolve().relative_to(project_dir.resolve()) + except ValueError: + return None + return obj + + +def _is_client_endpoint_construction(call: ast.Call, namespace: Dict[str, Any]) -> bool: + """match Endpoint(id=...) or Endpoint(image=...) built inline, incl. aliases.""" + 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) is not Endpoint: return False return any(kw.arg in ("id", "image") for kw in call.keywords) +def _remote_call_targets( + node: ast.AST, + namespace: Dict[str, Any], + remote_function_names: Set[str], + project_dir: Optional[Path], + helper_cache: Dict[int, List[str]], +) -> 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) + + 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 remote_function_names + ): + add(child.func.id) + if _is_client_endpoint_construction(child, namespace): + add("Endpoint") + helper = _project_function(_resolve(child.func, namespace), project_dir) + if helper is not None and _helper_targets( + helper, remote_function_names, project_dir, helper_cache + ): + add(helper.__name__) + elif isinstance(child, ast.Name): + obj = _resolve(child, namespace) + if _is_client_endpoint(obj) or _is_runpod_sdk(obj): + add(child.id) + elif isinstance(child, ast.Attribute) and isinstance(child.value, ast.Name): + if _is_client_endpoint(_resolve(child, namespace)): + add(f"{child.value.id}.{child.attr}") + return called + + +def _helper_targets( + func: Any, + remote_function_names: Set[str], + project_dir: Optional[Path], + helper_cache: Dict[int, List[str]], +) -> List[str]: + """analyze a project helper function, resolving names in its own module.""" + key = id(func) + if key in helper_cache: + return helper_cache[key] + helper_cache[key] = [] # break recursion cycles + try: + tree = ast.parse(textwrap.dedent(inspect.getsource(func))) + except Exception: + return [] + helper_cache[key] = _remote_call_targets( + tree, func.__globals__, remote_function_names, project_dir, helper_cache + ) + return helper_cache[key] + + def _analyze_cross_calls_ast( file_path: Path, function_names: Set[str], remote_function_names: Set[str], - client_endpoint_names: Optional[Set[str]] = None, + namespace: Optional[Dict[str, Any]] = None, + project_dir: Optional[Path] = None, ) -> Dict[str, List[str]]: - """find which functions call other endpoints via AST. - - returns a dict of function_name -> list of called remote function names. - matches direct calls to @remote functions (foo(), not obj.foo()), any - reference to a module-level client-mode Endpoint (vllm.runsync(...)), and - inline client construction (Endpoint(id="abc").runsync(...)). client-mode - Endpoints sign requests with RUNPOD_API_KEY, so the worker needs the key. + """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 calls another @remote function by name (foo(), not + obj.foo()), uses a client-mode Endpoint (id= or image=), uses the runpod + SDK, reads RUNPOD_API_KEY, or calls a project helper that does any of + these. namespace is the module's globals, used to resolve names; + project_dir limits helper analysis to project code. detection errs toward + giving the key, since a missing key fails at runtime with a 401. """ - client_names = client_endpoint_names or set() + ns = namespace or {} + helper_cache: Dict[int, List[str]] = {} result: Dict[str, List[str]] = {} try: tree = ast.parse(file_path.read_text(encoding="utf-8")) @@ -461,18 +561,9 @@ def _analyze_cross_calls_ast( if node.name not in function_names: continue - called: List[str] = [] - for child in ast.walk(node): - target: Optional[str] = None - if isinstance(child, ast.Call) and isinstance(child.func, ast.Name): - if child.func.id in remote_function_names: - target = child.func.id - if isinstance(child, ast.Call) and _is_client_endpoint_construction(child): - target = "Endpoint" - elif isinstance(child, ast.Name) and child.id in client_names: - target = child.id - if target is not None and target not in called: - called.append(target) + called = _remote_call_targets( + node, ns, remote_function_names, project_dir, helper_cache + ) if called: result[node.name] = called @@ -517,7 +608,7 @@ def discover_remote_functions(self) -> List[RemoteFunctionMetadata]: functions: List[RemoteFunctionMetadata] = [] seen_functions: Set[str] = set() failed_files: List[Path] = [] - client_names_by_file: Dict[Path, Set[str]] = {} + namespaces: Dict[Path, Dict[str, Any]] = {} try: for py_file in py_files: @@ -537,7 +628,7 @@ def discover_remote_functions(self) -> List[RemoteFunctionMetadata]: module, module_path, py_file, seen_functions ) functions.extend(file_functions) - client_names_by_file[py_file] = _find_client_endpoint_names(module) + namespaces[py_file] = vars(module) # cross-call analysis remote_names = {f.function_name for f in functions} @@ -550,7 +641,8 @@ def discover_remote_functions(self) -> List[RemoteFunctionMetadata]: file_path, func_names, remote_names, - client_names_by_file.get(file_path), + namespaces.get(file_path), + self.project_dir, ) for f in functions: if f.file_path == file_path and f.function_name in calls: 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 ab82be51..bb186c67 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 # --------------------------------------------------------------------------- @@ -324,9 +325,50 @@ async def caller(x): return await vllm.runsync(x) """, ) - result = _analyze_cross_calls_ast(p, {"caller"}, set(), {"vllm"}) + 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"']) def test_detects_inline_client_endpoint(self, tmp_path, kwarg): p = _write_worker( @@ -451,6 +493,46 @@ async def isolated(x): 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"] + assert pure_helper.calls_remote_functions is False + def test_sys_path_cleanup(self, tmp_path): """project root is removed from sys.path after scanning.""" root_str = str(tmp_path) From 5a315391bcd94a26e17b91bcbadda0f5adae6672 Mon Sep 17 00:00:00 2001 From: Justin Date: Thu, 15 Oct 2026 09:52:03 -0400 Subject: [PATCH 03/13] fix(build): close detection gaps and explain missing worker API keys Review fixes for the API key scoping change: - detect other @Endpoint functions reached through a module (tasks.square), an aliased import, or passed as a value; these were never detected and worked only because every endpoint got the key - analyze helpers only in scanned project files, so library code in an in-project .venv cannot flag every endpoint - treat a function the scanner fails to analyze as needing the key instead of crashing flash build - on a worker without RUNPOD_API_KEY, a 401 from a client Endpoint call explains why the key is missing and how to fix it - list the endpoints that received the key in the deploy summary - document what counts as a remote call --- docs/Cross_Endpoint_Routing.md | 10 +++ docs/Deployment_Architecture.md | 2 +- .../cli/commands/build_utils/scanner.py | 73 ++++++++++------ src/runpod_flash/cli/commands/deploy.py | 12 +++ src/runpod_flash/endpoint.py | 34 +++++++- .../build_utils/test_runtime_scanner.py | 86 +++++++++++++++++++ tests/unit/cli/test_deploy.py | 33 +++++++ tests/unit/test_endpoint_client.py | 52 +++++++++++ 8 files changed, 270 insertions(+), 32 deletions(-) diff --git a/docs/Cross_Endpoint_Routing.md b/docs/Cross_Endpoint_Routing.md index 18131c66..abd89c7b 100644 --- a/docs/Cross_Endpoint_Routing.md +++ b/docs/Cross_Endpoint_Routing.md @@ -183,6 +183,16 @@ 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 in your project: + +- references another `@Endpoint` function or class, including through a module (`tasks.square(x)`) or an aliased import +- uses a client `Endpoint(id=...)` or `Endpoint(image=...)`, at module level or built inline +- uses the `runpod` SDK, or reads `RUNPOD_API_KEY` + +Endpoints that do none of these do not receive the key. The deploy summary lists the endpoints that did. + +If a deployed endpoint fails with `401 Unauthorized` and the error says `RUNPOD_API_KEY is not set on this worker`, the call was not detected (for example, a client `Endpoint` passed in as an argument). Make it 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 839ae405..9924bf3f 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 their next update. 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 2886d809..03ec4f39 100644 --- a/src/runpod_flash/cli/commands/build_utils/scanner.py +++ b/src/runpod_flash/cli/commands/build_utils/scanner.py @@ -439,21 +439,29 @@ def _is_runpod_sdk(obj: Any) -> bool: return isinstance(name, str) and (name == "runpod" or name.startswith("runpod.")) -def _project_function(obj: Any, project_dir: Optional[Path]) -> Any: - """return obj if it is a plain function defined inside the project.""" - if project_dir is None or not isinstance(obj, types.FunctionType): +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 _project_function(obj: Any, project_files: Optional[Set[Path]]) -> 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 not project_files or not isinstance(obj, types.FunctionType): return None - if hasattr(obj, "__remote_config__"): + if _is_remote(obj): return None try: source_file = inspect.getsourcefile(obj) - except TypeError: - return None - if source_file is None: + except Exception: return None - try: - Path(source_file).resolve().relative_to(project_dir.resolve()) - except ValueError: + if source_file is None or Path(source_file).resolve() not in project_files: return None return obj @@ -473,7 +481,7 @@ def _remote_call_targets( node: ast.AST, namespace: Dict[str, Any], remote_function_names: Set[str], - project_dir: Optional[Path], + project_files: Optional[Set[Path]], helper_cache: Dict[int, List[str]], ) -> List[str]: """names of the things in node that need RUNPOD_API_KEY on the worker.""" @@ -494,17 +502,19 @@ def add(target: str) -> None: add(child.func.id) if _is_client_endpoint_construction(child, namespace): add("Endpoint") - helper = _project_function(_resolve(child.func, namespace), project_dir) + helper = _project_function(_resolve(child.func, namespace), project_files) if helper is not None and _helper_targets( - helper, remote_function_names, project_dir, helper_cache + helper, remote_function_names, project_files, helper_cache ): add(helper.__name__) elif isinstance(child, ast.Name): obj = _resolve(child, namespace) - if _is_client_endpoint(obj) or _is_runpod_sdk(obj): + if _is_remote(obj) or _is_client_endpoint(obj) or _is_runpod_sdk(obj): add(child.id) elif isinstance(child, ast.Attribute) and isinstance(child.value, ast.Name): - if _is_client_endpoint(_resolve(child, namespace)): + # tasks.square(x) or utils.vllm.runsync(x) through a module + obj = _resolve(child, namespace) + if _is_remote(obj) or _is_client_endpoint(obj): add(f"{child.value.id}.{child.attr}") return called @@ -512,7 +522,7 @@ def add(target: str) -> None: def _helper_targets( func: Any, remote_function_names: Set[str], - project_dir: Optional[Path], + project_files: Optional[Set[Path]], helper_cache: Dict[int, List[str]], ) -> List[str]: """analyze a project helper function, resolving names in its own module.""" @@ -525,7 +535,7 @@ def _helper_targets( except Exception: return [] helper_cache[key] = _remote_call_targets( - tree, func.__globals__, remote_function_names, project_dir, helper_cache + tree, func.__globals__, remote_function_names, project_files, helper_cache ) return helper_cache[key] @@ -535,17 +545,19 @@ def _analyze_cross_calls_ast( function_names: Set[str], remote_function_names: Set[str], namespace: Optional[Dict[str, Any]] = None, - project_dir: Optional[Path] = None, + project_files: Optional[Set[Path]] = None, ) -> Dict[str, List[str]]: """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 calls another @remote function by name (foo(), not - obj.foo()), uses a client-mode Endpoint (id= or image=), uses the runpod - SDK, reads RUNPOD_API_KEY, or calls a project helper that does any of - these. namespace is the module's globals, used to resolve names; - project_dir limits helper analysis to project code. detection errs toward - giving the key, since a missing key fails at runtime with a 401. + 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 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. detection errs toward giving the key, + since a missing key fails at runtime with a 401: a function that cannot + be analyzed is reported as needing it. """ ns = namespace or {} helper_cache: Dict[int, List[str]] = {} @@ -561,9 +573,13 @@ def _analyze_cross_calls_ast( if node.name not in function_names: continue - called = _remote_call_targets( - node, ns, remote_function_names, project_dir, helper_cache - ) + try: + called = _remote_call_targets( + node, ns, remote_function_names, project_files, helper_cache + ) + except Exception as e: + logger.debug(f"cross-call analysis failed for {node.name}: {e}") + called = [""] if called: result[node.name] = called @@ -632,6 +648,7 @@ def discover_remote_functions(self) -> List[RemoteFunctionMetadata]: # cross-call analysis remote_names = {f.function_name for f in functions} + project_files = {f.resolve() for f in py_files} files_with_functions: Dict[Path, Set[str]] = {} for f in functions: files_with_functions.setdefault(f.file_path, set()).add(f.function_name) @@ -642,7 +659,7 @@ def discover_remote_functions(self) -> List[RemoteFunctionMetadata]: func_names, remote_names, namespaces.get(file_path), - self.project_dir, + project_files, ) for f in functions: if f.file_path == file_path and f.function_name in calls: diff --git a/src/runpod_flash/cli/commands/deploy.py b/src/runpod_flash/cli/commands/deploy.py index 73939ae5..69baff1c 100644 --- a/src/runpod_flash/cli/commands/deploy.py +++ b/src/runpod_flash/cli/commands/deploy.py @@ -170,6 +170,18 @@ def _display_post_deployment_guidance( for method, path in route_list: console.print(f" {' ' * max_name} [dim]{method:6s}{path}[/dim]") + keyed = sorted( + name + for name, _url, _routes in rows + if resources.get(name, {}).get("makes_remote_calls") + ) + if keyed: + console.print() + console.print( + "[dim]RUNPOD_API_KEY added to endpoints that call other endpoints: " + f"{', '.join(keyed)}[/dim]" + ) + # one curl example curl_url = None curl_method = "POST" diff --git a/src/runpod_flash/endpoint.py b/src/runpod_flash/endpoint.py index 632369db..f0ca30d1 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/tests/unit/cli/commands/build_utils/test_runtime_scanner.py b/tests/unit/cli/commands/build_utils/test_runtime_scanner.py index bb186c67..76c38bf1 100644 --- a/tests/unit/cli/commands/build_utils/test_runtime_scanner.py +++ b/tests/unit/cli/commands/build_utils/test_runtime_scanner.py @@ -382,6 +382,15 @@ async def caller(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, @@ -533,6 +542,83 @@ async def pure_helper(x): assert via_helper.called_remote_functions == ["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_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_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 fd20a521..7db79041 100644 --- a/tests/unit/cli/test_deploy.py +++ b/tests/unit/cli/test_deploy.py @@ -643,6 +643,39 @@ 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 added to endpoints that call other endpoints: caller" + in (output) + ) + assert "caller, isolated" not in output + + 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/test_endpoint_client.py b/tests/unit/test_endpoint_client.py index 6b85a2e5..e4e0c384 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) From 363b1a59ccecec997c680ae492394bb9cf3b54d8 Mon Sep 17 00:00:00 2001 From: Justin Date: Thu, 15 Oct 2026 10:26:58 -0400 Subject: [PATCH 04/13] fix(deploy): clean stale keys on unchanged redeploys and close scanner gaps Address review findings: - remove a stale RUNPOD_API_KEY from unchanged local-only endpoints during flash deploy; reconciliation skips them, so the update path never ran and upgrading Flash then redeploying left the key in place - resolve imports made inside an endpoint function (aliased Endpoint, runpod SDK, project modules) so those callers keep the key - analyze helpers with a fresh cache per endpoint; a helper cached in the middle of a cycle could hide a remote call from a later endpoint - leave user-managed RUNPOD_API_KEY entries out of the deploy summary --- docs/Deployment_Architecture.md | 2 +- .../cli/commands/build_utils/scanner.py | 54 +++++++- src/runpod_flash/cli/commands/deploy.py | 2 + src/runpod_flash/cli/utils/deployment.py | 64 +++++++++ src/runpod_flash/core/api/runpod.py | 17 +++ .../build_utils/test_runtime_scanner.py | 61 +++++++++ tests/unit/cli/test_deploy.py | 16 +++ tests/unit/cli/utils/test_deployment.py | 129 ++++++++++++++++++ 8 files changed, 341 insertions(+), 4 deletions(-) diff --git a/docs/Deployment_Architecture.md b/docs/Deployment_Architecture.md index 9924bf3f..c264600d 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. Other endpoints do not get the key, and a key left on them by an earlier deploy is removed on their next update. See [API Key Injection](Cross_Endpoint_Routing.md#api-key-injection) for what counts as a remote call +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. 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 03ec4f39..c5df80a1 100644 --- a/src/runpod_flash/cli/commands/build_utils/scanner.py +++ b/src/runpod_flash/cli/commands/build_utils/scanner.py @@ -16,7 +16,7 @@ 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 ( @@ -436,7 +436,7 @@ def _is_runpod_sdk(obj: Any) -> bool: name = getattr(obj, "__module__", None) or "" except Exception: return False - return isinstance(name, str) and (name == "runpod" or name.startswith("runpod.")) + return isinstance(name, str) and _is_runpod_module_name(name) def _is_remote(obj: Any) -> bool: @@ -477,6 +477,46 @@ def _is_client_endpoint_construction(call: ast.Call, namespace: Dict[str, Any]) return any(kw.arg in ("id", "image") for kw in call.keywords) +def _local_imports(node: ast.AST) -> Tuple[Dict[str, Any], bool]: + """resolve imports made inside node against already-loaded modules. + + returns (name -> object overlay, whether the runpod SDK is imported). + only sys.modules is consulted, so analysis never imports new code. + """ + overlay: Dict[str, Any] = {} + uses_runpod = 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: + module = sys.modules.get(alias.name) + bound = alias.asname + else: + bound = alias.name.split(".")[0] + module = sys.modules.get(bound) + if module is not None: + overlay[bound] = module + elif isinstance(child, ast.ImportFrom) and child.module and not child.level: + if _is_runpod_module_name(child.module): + uses_runpod = True + module = sys.modules.get(child.module) + if module is None: + continue + for alias in child.names: + obj = getattr(module, alias.name, _MISSING) + if obj is _MISSING: + obj = sys.modules.get(f"{child.module}.{alias.name}", _MISSING) + if obj is not _MISSING: + overlay[alias.asname or alias.name] = obj + return overlay, uses_runpod + + +def _is_runpod_module_name(name: str) -> bool: + return name == "runpod" or name.startswith("runpod.") + + def _remote_call_targets( node: ast.AST, namespace: Dict[str, Any], @@ -491,6 +531,12 @@ def add(target: str) -> None: if target not in called: called.append(target) + overlay, uses_runpod = _local_imports(node) + 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) @@ -560,7 +606,6 @@ def _analyze_cross_calls_ast( be analyzed is reported as needing it. """ ns = namespace or {} - helper_cache: Dict[int, List[str]] = {} result: Dict[str, List[str]] = {} try: tree = ast.parse(file_path.read_text(encoding="utf-8")) @@ -573,6 +618,9 @@ def _analyze_cross_calls_ast( if node.name not in function_names: continue + # fresh per function: a helper cached mid-cycle can hold a partial + # result that is only correct within the traversal that produced it. + helper_cache: Dict[int, List[str]] = {} try: called = _remote_call_targets( node, ns, remote_function_names, project_files, helper_cache diff --git a/src/runpod_flash/cli/commands/deploy.py b/src/runpod_flash/cli/commands/deploy.py index 69baff1c..64dfb210 100644 --- a/src/runpod_flash/cli/commands/deploy.py +++ b/src/runpod_flash/cli/commands/deploy.py @@ -170,10 +170,12 @@ 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() diff --git a/src/runpod_flash/cli/utils/deployment.py b/src/runpod_flash/cli/utils/deployment.py index 1a53e80c..1fda56f8 100644 --- a/src/runpod_flash/cli/utils/deployment.py +++ b/src/runpod_flash/cli/utils/deployment.py @@ -41,6 +41,60 @@ def _manifest_without_ai_keys(manifest: Dict[str, Any]) -> Dict[str, Any]: return sanitized_manifest +async def _remove_stale_api_key(client: Any, endpoint_id: str) -> bool: + """Drop RUNPOD_API_KEY from an endpoint's live template. True if removed.""" + template_id = await client.get_endpoint_template_id(endpoint_id) + if not template_id: + return False + 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 + payload = { + key: live[key] + for key in ("name", "imageName", "containerDiskInGb", "dockerArgs", "readme") + if live.get(key) is not None + } + payload["env"] = [ + {"key": entry["key"], "value": entry.get("value", "")} for entry in kept + ] + # saveTemplate requires volumeInGb, but for sls this is always 0 + payload["volumeInGb"] = 0 + payload["id"] = 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. + Failures are logged and never fail the deploy. + """ + from runpod_flash.core.api.runpod import RunpodGraphQLClient + + try: + async with RunpodGraphQLClient() as client: + for resource_name, endpoint_id in endpoints: + try: + removed = await _remove_stale_api_key(client, endpoint_id) + except Exception as e: + log.warning( + f"{resource_name}: could not check for a stale " + f"RUNPOD_API_KEY: {e}" + ) + continue + if removed: + log.info( + f"{resource_name}: removed RUNPOD_API_KEY " + f"(endpoint does not call other endpoints)" + ) + 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 +170,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 +240,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 +259,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 b527d2ce..ee09d796 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_id(self, endpoint_id: str) -> Optional[str]: + """Return the template ID an endpoint runs, or None if not found.""" + query = """ + query getEndpointTemplate($id: String!) { + myself { + endpoint(id: $id) { + id + templateId + } + } + } + """ + + result = await self._execute_graphql(query, {"id": endpoint_id}) + endpoint = (result.get("myself") or {}).get("endpoint") or {} + return endpoint.get("templateId") + async def create_flash_auth_request(self) -> Dict[str, Any]: mutation = """ mutation createFlashAuthRequest { 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 76c38bf1..8245a7f7 100644 --- a/tests/unit/cli/commands/build_utils/test_runtime_scanner.py +++ b/tests/unit/cli/commands/build_utils/test_runtime_scanner.py @@ -586,6 +586,67 @@ async def passed(xs): 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" diff --git a/tests/unit/cli/test_deploy.py b/tests/unit/cli/test_deploy.py index 7db79041..a2df5954 100644 --- a/tests/unit/cli/test_deploy.py +++ b/tests/unit/cli/test_deploy.py @@ -665,6 +665,22 @@ def test_lists_endpoints_that_receive_api_key(self, patched_console): ) 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 diff --git a/tests/unit/cli/utils/test_deployment.py b/tests/unit/cli/utils/test_deployment.py index 334a8e24..50c28c29 100644 --- a/tests/unit/cli/utils/test_deployment.py +++ b/tests/unit/cli/utils/test_deployment.py @@ -842,3 +842,132 @@ 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_id="tpl-1"): + client = AsyncMock() + client.get_endpoint_template_id = AsyncMock(return_value=template_id) + client.get_template = AsyncMock( + return_value={ + "id": template_id, + "name": "worker-tpl", + "imageName": "runpod/flash:latest", + "containerDiskInGb": 64, + "dockerArgs": "", + "readme": None, + "env": env, + } + ) + client.update_template = AsyncMock(return_value={}) + return client + + @pytest.mark.asyncio + async def test_removes_key_and_keeps_everything_else(self): + from runpod_flash.cli.utils.deployment import _remove_stale_api_key + + client = self._client( + [ + {"key": "PORT", "value": "8080"}, + {"key": "RUNPOD_API_KEY", "value": "old-key"}, + ] + ) + + assert await _remove_stale_api_key(client, "endpoint-123") is True + + 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): + from runpod_flash.cli.utils.deployment import _remove_stale_api_key + + client = self._client([{"key": "PORT", "value": "8080"}]) + + assert await _remove_stale_api_key(client, "endpoint-123") is False + client.update_template.assert_not_awaited() + + @pytest.mark.asyncio + async def test_no_update_without_template(self): + from runpod_flash.cli.utils.deployment import _remove_stale_api_key + + client = self._client([], template_id=None) + + assert await _remove_stale_api_key(client, "endpoint-123") is False + 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")]) From 92f5ec214bac84d88a6e61025e247f2fdb8611be Mon Sep 17 00:00:00 2001 From: Justin Date: Thu, 15 Oct 2026 10:28:43 -0400 Subject: [PATCH 05/13] fix(deploy): treat a deleted endpoint as having no template during stale-key cleanup The live API answers an unknown endpoint id with an "Endpoint not found" GraphQL error, which surfaced as a confusing cleanup warning. --- src/runpod_flash/core/api/runpod.py | 8 ++++- .../core/api/test_runpod_graphql_extended.py | 34 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/src/runpod_flash/core/api/runpod.py b/src/runpod_flash/core/api/runpod.py index ee09d796..63dc220e 100644 --- a/src/runpod_flash/core/api/runpod.py +++ b/src/runpod_flash/core/api/runpod.py @@ -860,7 +860,13 @@ async def get_endpoint_template_id(self, endpoint_id: str) -> Optional[str]: } """ - result = await self._execute_graphql(query, {"id": endpoint_id}) + try: + result = await self._execute_graphql(query, {"id": endpoint_id}) + except _GraphQLErrorResponse as e: + # deleted endpoints come back as an "Endpoint not found" error + if "not found" in str(e).lower(): + return None + raise endpoint = (result.get("myself") or {}).get("endpoint") or {} return endpoint.get("templateId") diff --git a/tests/unit/core/api/test_runpod_graphql_extended.py b/tests/unit/core/api/test_runpod_graphql_extended.py index 2d6cb734..3fec6662 100644 --- a/tests/unit/core/api/test_runpod_graphql_extended.py +++ b/tests/unit/core/api/test_runpod_graphql_extended.py @@ -352,3 +352,37 @@ 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 TestGetEndpointTemplateId: + """get_endpoint_template_id returns the template or None for gone endpoints.""" + + @pytest.mark.asyncio + async def test_returns_template_id(self): + client = RunpodGraphQLClient(api_key="test") + with patch.object(client, "_execute_graphql", new_callable=AsyncMock) as mock: + mock.return_value = { + "myself": {"endpoint": {"id": "ep-1", "templateId": "tpl-1"}} + } + assert await client.get_endpoint_template_id("ep-1") == "tpl-1" + + @pytest.mark.asyncio + async def test_deleted_endpoint_returns_none(self): + from runpod_flash.core.api.runpod import _GraphQLErrorResponse + + client = RunpodGraphQLClient(api_key="test") + with patch.object(client, "_execute_graphql", new_callable=AsyncMock) as mock: + mock.side_effect = _GraphQLErrorResponse( + "GraphQL errors: Endpoint not found", [] + ) + assert await client.get_endpoint_template_id("gone") is None + + @pytest.mark.asyncio + async def test_other_errors_propagate(self): + from runpod_flash.core.api.runpod import _GraphQLErrorResponse + + client = RunpodGraphQLClient(api_key="test") + with patch.object(client, "_execute_graphql", new_callable=AsyncMock) as mock: + mock.side_effect = _GraphQLErrorResponse("GraphQL errors: forbidden", []) + with pytest.raises(_GraphQLErrorResponse): + await client.get_endpoint_template_id("ep-1") From a83b4f00a86dfe9da87e5e41b99ba8bea5b9dbc1 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 19 Oct 2026 13:19:28 -0400 Subject: [PATCH 06/13] fix(build): follow local project imports, nested attributes, and project classes Cross-call analysis now: - resolves function-local and relative imports against the scanned modules, which the scanner removes from sys.modules after importing them - walks package-qualified references such as pkg.tasks.square(x) - analyzes project classes and instances (methods, static/class methods, instance attributes) referenced by an endpoint - includes package __init__.py files in helper scope - reports for helpers whose source cannot be parsed and for project imports that cannot be resolved, so detection errs toward the key --- .../cli/commands/build_utils/scanner.py | 272 +++++++++++++----- .../build_utils/test_runtime_scanner.py | 221 ++++++++++++++ 2 files changed, 422 insertions(+), 71 deletions(-) diff --git a/src/runpod_flash/cli/commands/build_utils/scanner.py b/src/runpod_flash/cli/commands/build_utils/scanner.py index c5df80a1..71c0684e 100644 --- a/src/runpod_flash/cli/commands/build_utils/scanner.py +++ b/src/runpod_flash/cli/commands/build_utils/scanner.py @@ -409,17 +409,51 @@ def _find_remote_decorated(module: Any) -> Dict[str, Any]: _API_KEY_ENV = "RUNPOD_API_KEY" +_UNANALYZED = "" _MISSING = object() -def _resolve(expr: ast.expr, namespace: Dict[str, Any]) -> Any: - """resolve a Name, or an attribute of a module (utils.helper), to its 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) + base = _resolve(expr.value, namespace, scan) if isinstance(base, types.ModuleType): - return getattr(base, expr.attr, _MISSING) + 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 @@ -447,82 +481,149 @@ def _is_remote(obj: Any) -> bool: return False -def _project_function(obj: Any, project_files: Optional[Set[Path]]) -> Any: +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 not project_files or not isinstance(obj, types.FunctionType): - return None - if _is_remote(obj): + if not isinstance(obj, types.FunctionType) or _is_remote(obj): return None try: source_file = inspect.getsourcefile(obj) except Exception: return None - if source_file is None or Path(source_file).resolve() not in project_files: - return None - return obj + 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]) -> bool: + +def _is_client_endpoint_construction( + call: ast.Call, namespace: Dict[str, Any], scan: _CallScan +) -> bool: """match Endpoint(id=...) or Endpoint(image=...) built inline, incl. aliases.""" 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) is not Endpoint: + if not literal and _resolve(func, namespace, scan) is not Endpoint: return False return any(kw.arg in ("id", "image") for kw in call.keywords) -def _local_imports(node: ast.AST) -> Tuple[Dict[str, Any], bool]: +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). - only sys.modules is consulted, so analysis never imports new code. + 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: - module = sys.modules.get(alias.name) - bound = alias.asname + name, bound = alias.name, alias.asname else: - bound = alias.name.split(".")[0] - module = sys.modules.get(bound) + name = bound = alias.name.split(".")[0] + module = scan.lookup_module(name) if module is not None: overlay[bound] = module - elif isinstance(child, ast.ImportFrom) and child.module and not child.level: - if _is_runpod_module_name(child.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 - module = sys.modules.get(child.module) + # 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: - obj = sys.modules.get(f"{child.module}.{alias.name}", _MISSING) - if obj is not _MISSING: - overlay[alias.asname or alias.name] = obj - return overlay, uses_runpod + 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. + """ + if _is_remote(obj) or _is_client_endpoint(obj): + return True + cls, module = _project_class(obj, scan) + if cls is None: + return False + if not inspect.isclass(obj): + try: + attrs = list(vars(obj).values()) + except TypeError: + attrs = [] + if any(_is_remote(v) or _is_client_endpoint(v) for v in attrs): + return True + return bool(_class_targets(cls, module, scan)) + + def _remote_call_targets( - node: ast.AST, - namespace: Dict[str, Any], - remote_function_names: Set[str], - project_files: Optional[Set[Path]], - helper_cache: Dict[int, List[str]], + 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] = [] @@ -531,7 +632,9 @@ def add(target: str) -> None: if target not in called: called.append(target) - overlay, uses_runpod = _local_imports(node) + overlay, uses_runpod, unresolved = _local_imports(node, namespace, scan) + if unresolved: + add(_UNANALYZED) if uses_runpod: add("runpod") if overlay: @@ -543,47 +646,67 @@ def add(target: str) -> None: elif isinstance(child, ast.Call): if ( isinstance(child.func, ast.Name) - and child.func.id in remote_function_names + and child.func.id in scan.remote_function_names ): add(child.func.id) - if _is_client_endpoint_construction(child, namespace): + if _is_client_endpoint_construction(child, namespace, scan): add("Endpoint") - helper = _project_function(_resolve(child.func, namespace), project_files) - if helper is not None and _helper_targets( - helper, remote_function_names, project_files, helper_cache - ): + 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) - if _is_remote(obj) or _is_client_endpoint(obj) or _is_runpod_sdk(obj): + obj = _resolve(child, namespace, scan) + if _is_runpod_sdk(obj) or _object_needs_key(obj, scan): add(child.id) - elif isinstance(child, ast.Attribute) and isinstance(child.value, ast.Name): - # tasks.square(x) or utils.vllm.runsync(x) through a module - obj = _resolve(child, namespace) - if _is_remote(obj) or _is_client_endpoint(obj): - add(f"{child.value.id}.{child.attr}") + 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, - remote_function_names: Set[str], - project_files: Optional[Set[Path]], - helper_cache: Dict[int, List[str]], -) -> List[str]: +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 helper_cache: - return helper_cache[key] - helper_cache[key] = [] # break recursion cycles + 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: - return [] - helper_cache[key] = _remote_call_targets( - tree, func.__globals__, remote_function_names, project_files, helper_cache - ) - return helper_cache[key] + # 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( @@ -592,6 +715,7 @@ def _analyze_cross_calls_ast( 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 need RUNPOD_API_KEY on the worker, via AST. @@ -599,11 +723,12 @@ def _analyze_cross_calls_ast( 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 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. detection errs toward giving the key, - since a missing key fails at runtime with a 401: a function that cannot - be analyzed is reported as needing it. + 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]] = {} @@ -620,14 +745,12 @@ def _analyze_cross_calls_ast( # fresh per function: a helper cached mid-cycle can hold a partial # result that is only correct within the traversal that produced it. - helper_cache: Dict[int, List[str]] = {} + scan = _CallScan(remote_function_names, project_files, project_modules or {}) try: - called = _remote_call_targets( - node, ns, remote_function_names, project_files, helper_cache - ) + called = _remote_call_targets(node, ns, scan) except Exception as e: logger.debug(f"cross-call analysis failed for {node.name}: {e}") - called = [""] + called = [_UNANALYZED] if called: result[node.name] = called @@ -673,10 +796,13 @@ def discover_remote_functions(self) -> 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: @@ -687,6 +813,7 @@ 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 @@ -696,7 +823,9 @@ def discover_remote_functions(self) -> List[RemoteFunctionMetadata]: # cross-call analysis remote_names = {f.function_name for f in functions} - project_files = {f.resolve() for f in py_files} + # __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) @@ -708,6 +837,7 @@ def discover_remote_functions(self) -> List[RemoteFunctionMetadata]: 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: 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 8245a7f7..129115dd 100644 --- a/tests/unit/cli/commands/build_utils/test_runtime_scanner.py +++ b/tests/unit/cli/commands/build_utils/test_runtime_scanner.py @@ -680,6 +680,227 @@ async def lib_call(x): 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"] + + 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")) + pure = Pure() + """, + ) + _write_worker( + tmp_path, + "cls_app.py", + """\ + from runpod_flash import Endpoint + from cls_clients import Helpers, client, holder, 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="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["pure_method"].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) From 9fc641d65b916276e234700f766e6401c9bffc2c Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 19 Oct 2026 13:19:28 -0400 Subject: [PATCH 07/13] fix(deploy): keep the key on shared templates and bound stale-key cleanup - one myself.endpoints query maps every endpoint to its template, replacing the per-endpoint lookup and its error-message matching for deleted endpoints - a template shared with another endpoint keeps RUNPOD_API_KEY, in both the unchanged-redeploy cleanup and the update path - cleanup runs concurrently under a 120s timeout before provisioning - the live-template payload reuses ServerlessResource's saveTemplate builder --- src/runpod_flash/cli/utils/deployment.py | 77 ++++++++------- src/runpod_flash/core/api/runpod.py | 20 ++-- src/runpod_flash/core/resources/serverless.py | 67 +++++++++++-- tests/unit/cli/utils/test_deployment.py | 94 ++++++++++++++++--- .../core/api/test_runpod_graphql_extended.py | 38 ++++---- tests/unit/resources/test_serverless.py | 29 +++++- 6 files changed, 235 insertions(+), 90 deletions(-) diff --git a/src/runpod_flash/cli/utils/deployment.py b/src/runpod_flash/cli/utils/deployment.py index 1fda56f8..d102c604 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,27 +41,23 @@ def _manifest_without_ai_keys(manifest: Dict[str, Any]) -> Dict[str, Any]: return sanitized_manifest -async def _remove_stale_api_key(client: Any, endpoint_id: str) -> bool: - """Drop RUNPOD_API_KEY from an endpoint's live template. True if removed.""" - template_id = await client.get_endpoint_template_id(endpoint_id) - if not template_id: - return False +# 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 - payload = { - key: live[key] - for key in ("name", "imageName", "containerDiskInGb", "dockerArgs", "readme") - if live.get(key) is not None - } - payload["env"] = [ + env_payload = [ {"key": entry["key"], "value": entry.get("value", "")} for entry in kept ] - # saveTemplate requires volumeInGb, but for sls this is always 0 - payload["volumeInGb"] = 0 - payload["id"] = template_id + payload = ServerlessResource._template_update_payload( + {**live, "env": env_payload}, template_id + ) await client.update_template(payload) return True @@ -71,28 +67,45 @@ async def _remove_stale_api_keys(endpoints: list[tuple[str, str]]) -> None: Older Flash versions gave every endpoint the key. Unchanged resources skip provisioning, so the update path that drops it never runs for them. - Failures are logged and never fail the deploy. + 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 - try: + 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: - for resource_name, endpoint_id in endpoints: - try: - removed = await _remove_stale_api_key(client, endpoint_id) - except Exception as e: - log.warning( - f"{resource_name}: could not check for a stale " - f"RUNPOD_API_KEY: {e}" - ) - continue - if removed: - log.info( - f"{resource_name}: removed RUNPOD_API_KEY " - f"(endpoint does not call other endpoints)" - ) + 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 Exception as e: - log.warning(f"Could not check endpoints for a stale RUNPOD_API_KEY: {e}") + log.warning(f"Could not check endpoints for a stale RUNPOD_API_KEY: {e!r}") def _resource_config_for_compare(config: Dict[str, Any]) -> Dict[str, Any]: diff --git a/src/runpod_flash/core/api/runpod.py b/src/runpod_flash/core/api/runpod.py index 63dc220e..5c36f765 100644 --- a/src/runpod_flash/core/api/runpod.py +++ b/src/runpod_flash/core/api/runpod.py @@ -847,12 +847,12 @@ async def endpoint_exists(self, endpoint_id: str) -> bool: log.debug(f"Error checking endpoint existence: {e}") return False - async def get_endpoint_template_id(self, endpoint_id: str) -> Optional[str]: - """Return the template ID an endpoint runs, or None if not found.""" + async def get_endpoint_template_ids(self) -> Dict[str, Optional[str]]: + """Return endpoint ID -> template ID for every endpoint on the account.""" query = """ - query getEndpointTemplate($id: String!) { + query getEndpointTemplates { myself { - endpoint(id: $id) { + endpoints { id templateId } @@ -860,15 +860,9 @@ async def get_endpoint_template_id(self, endpoint_id: str) -> Optional[str]: } """ - try: - result = await self._execute_graphql(query, {"id": endpoint_id}) - except _GraphQLErrorResponse as e: - # deleted endpoints come back as an "Endpoint not found" error - if "not found" in str(e).lower(): - return None - raise - endpoint = (result.get("myself") or {}).get("endpoint") or {} - return endpoint.get("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 = """ diff --git a/src/runpod_flash/core/resources/serverless.py b/src/runpod_flash/core/resources/serverless.py index 3fc55438..d3007faf 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: @@ -856,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", @@ -873,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", []) @@ -882,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. @@ -1059,7 +1107,7 @@ 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 = not self._check_makes_remote_calls() + 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: @@ -1226,9 +1274,10 @@ 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 = ( - not new_config._check_makes_remote_calls() - and "RUNPOD_API_KEY" not in (new_config.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( diff --git a/tests/unit/cli/utils/test_deployment.py b/tests/unit/cli/utils/test_deployment.py index 50c28c29..d4865a51 100644 --- a/tests/unit/cli/utils/test_deployment.py +++ b/tests/unit/cli/utils/test_deployment.py @@ -903,12 +903,16 @@ class TestRemoveStaleApiKey: """_remove_stale_api_key drops only RUNPOD_API_KEY from the live template.""" @staticmethod - def _client(env, template_id="tpl-1"): + def _client(env, template_ids=None): client = AsyncMock() - client.get_endpoint_template_id = AsyncMock(return_value=template_id) + 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": template_id, + "id": "tpl-1", "name": "worker-tpl", "imageName": "runpod/flash:latest", "containerDiskInGb": 64, @@ -920,10 +924,17 @@ def _client(env, template_id="tpl-1"): 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): - from runpod_flash.cli.utils.deployment import _remove_stale_api_key - client = self._client( [ {"key": "PORT", "value": "8080"}, @@ -931,8 +942,9 @@ async def test_removes_key_and_keeps_everything_else(self): ] ) - assert await _remove_stale_api_key(client, "endpoint-123") is True + 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", @@ -946,20 +958,76 @@ async def test_removes_key_and_keeps_everything_else(self): @pytest.mark.asyncio async def test_no_update_when_key_absent(self): - from runpod_flash.cli.utils.deployment import _remove_stale_api_key - client = self._client([{"key": "PORT", "value": "8080"}]) - assert await _remove_stale_api_key(client, "endpoint-123") is False + 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_no_update_without_template(self): - from runpod_flash.cli.utils.deployment import _remove_stale_api_key + 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"}, + ) - client = self._client([], template_id=None) + 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): + import asyncio + + client = self._client([]) + + async def hang(): + await asyncio.sleep(10) + + client.get_endpoint_template_ids = AsyncMock(side_effect=hang) + with patch( + "runpod_flash.cli.utils.deployment.STALE_KEY_CLEANUP_TIMEOUT_SECONDS", 0.05 + ): + await self._cleanup(client) - assert await _remove_stale_api_key(client, "endpoint-123") is False client.get_template.assert_not_awaited() @pytest.mark.asyncio diff --git a/tests/unit/core/api/test_runpod_graphql_extended.py b/tests/unit/core/api/test_runpod_graphql_extended.py index 3fec6662..a9d00514 100644 --- a/tests/unit/core/api/test_runpod_graphql_extended.py +++ b/tests/unit/core/api/test_runpod_graphql_extended.py @@ -354,35 +354,29 @@ async def test_execute_rest_network_error(self): await client._execute_rest("GET", "https://example.com/api") -class TestGetEndpointTemplateId: - """get_endpoint_template_id returns the template or None for gone endpoints.""" +class TestGetEndpointTemplateIds: + """get_endpoint_template_ids maps every endpoint to its template.""" @pytest.mark.asyncio - async def test_returns_template_id(self): + 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": {"endpoint": {"id": "ep-1", "templateId": "tpl-1"}} + "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, } - assert await client.get_endpoint_template_id("ep-1") == "tpl-1" - - @pytest.mark.asyncio - async def test_deleted_endpoint_returns_none(self): - from runpod_flash.core.api.runpod import _GraphQLErrorResponse - - client = RunpodGraphQLClient(api_key="test") - with patch.object(client, "_execute_graphql", new_callable=AsyncMock) as mock: - mock.side_effect = _GraphQLErrorResponse( - "GraphQL errors: Endpoint not found", [] - ) - assert await client.get_endpoint_template_id("gone") is None @pytest.mark.asyncio - async def test_other_errors_propagate(self): - from runpod_flash.core.api.runpod import _GraphQLErrorResponse - + async def test_no_endpoints(self): client = RunpodGraphQLClient(api_key="test") with patch.object(client, "_execute_graphql", new_callable=AsyncMock) as mock: - mock.side_effect = _GraphQLErrorResponse("GraphQL errors: forbidden", []) - with pytest.raises(_GraphQLErrorResponse): - await client.get_endpoint_template_id("ep-1") + 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 80c3beb4..607aa50f 100644 --- a/tests/unit/resources/test_serverless.py +++ b/tests/unit/resources/test_serverless.py @@ -3520,7 +3520,7 @@ class TestStaleApiKeyCleanup: once the build manifest says the endpoint makes no remote calls.""" @staticmethod - async def _update(new_env, makes_remote_calls): + async def _update(new_env, makes_remote_calls, template_ids=None): old_resource = ServerlessEndpoint( name="stale-key", imageName="test:latest", @@ -3559,6 +3559,12 @@ async def _update(new_env, makes_remote_calls): ] } ) + 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" @@ -3598,3 +3604,24 @@ 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"}, {"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" From 8857d98b514933274d760eab5b5e8170aba1e0d8 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 19 Oct 2026 13:23:26 -0400 Subject: [PATCH 08/13] fix(build): follow nested project instances when checking attributes --- .../cli/commands/build_utils/scanner.py | 25 ++++++++++++------- .../build_utils/test_runtime_scanner.py | 9 ++++++- 2 files changed, 24 insertions(+), 10 deletions(-) diff --git a/src/runpod_flash/cli/commands/build_utils/scanner.py b/src/runpod_flash/cli/commands/build_utils/scanner.py index 71c0684e..5f0608a8 100644 --- a/src/runpod_flash/cli/commands/build_utils/scanner.py +++ b/src/runpod_flash/cli/commands/build_utils/scanner.py @@ -607,19 +607,26 @@ def _object_needs_key(obj: Any, scan: _CallScan) -> bool: or instance attributes would. the whole class is analyzed because calls through self cannot be resolved statically. """ + return _needs_key(obj, scan, set()) + + +def _needs_key(obj: Any, scan: _CallScan, seen: Set[int]) -> bool: if _is_remote(obj) or _is_client_endpoint(obj): return True cls, module = _project_class(obj, scan) - if cls is None: + if cls is None or id(obj) in seen: return False - if not inspect.isclass(obj): - try: - attrs = list(vars(obj).values()) - except TypeError: - attrs = [] - if any(_is_remote(v) or _is_client_endpoint(v) for v in attrs): - return True - return bool(_class_targets(cls, module, scan)) + seen.add(id(obj)) + 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 _remote_call_targets( 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 129115dd..a4f796c1 100644 --- a/tests/unit/cli/commands/build_utils/test_runtime_scanner.py +++ b/tests/unit/cli/commands/build_utils/test_runtime_scanner.py @@ -843,7 +843,9 @@ def double(self, x): client = ModelClient() holder = Holder(Endpoint(id="abc123")) + nested = Holder(Holder(Endpoint(id="abc123"))) pure = Pure() + pure.self_ref = pure """, ) _write_worker( @@ -851,7 +853,7 @@ def double(self, x): "cls_app.py", """\ from runpod_flash import Endpoint - from cls_clients import Helpers, client, holder, pure + from cls_clients import Helpers, client, holder, nested, pure @Endpoint(name="method", cpu="cpu3c-1-2") async def method(x): @@ -865,6 +867,10 @@ async def static(x): 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) @@ -875,6 +881,7 @@ async def pure_method(x): 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_unparseable_helper_assumes_key_needed(self, tmp_path): From 8574dc78cd50ad30845b105c26af1f4993d8aefb Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 19 Oct 2026 13:27:49 -0400 Subject: [PATCH 09/13] fix(build): follow bound methods, partials, containers, and cached factories Also document that Endpoint(image=...) deployments do not receive the key, since Flash does not scan their image. --- docs/Cross_Endpoint_Routing.md | 2 + .../cli/commands/build_utils/scanner.py | 41 +++++++++++++++- .../build_utils/test_runtime_scanner.py | 47 +++++++++++++++++++ 3 files changed, 88 insertions(+), 2 deletions(-) diff --git a/docs/Cross_Endpoint_Routing.md b/docs/Cross_Endpoint_Routing.md index abd89c7b..77dae363 100644 --- a/docs/Cross_Endpoint_Routing.md +++ b/docs/Cross_Endpoint_Routing.md @@ -191,6 +191,8 @@ This means your function code never needs to handle API keys directly. 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. + If a deployed endpoint fails with `401 Unauthorized` and the error says `RUNPOD_API_KEY is not set on this worker`, the call was not detected (for example, a client `Endpoint` passed in as an argument). Make it visible in one of the ways above and redeploy. ## Load-Balanced Cross-Endpoint Calls diff --git a/src/runpod_flash/cli/commands/build_utils/scanner.py b/src/runpod_flash/cli/commands/build_utils/scanner.py index 5f0608a8..9e394e6d 100644 --- a/src/runpod_flash/cli/commands/build_utils/scanner.py +++ b/src/runpod_flash/cli/commands/build_utils/scanner.py @@ -6,8 +6,10 @@ """ import ast +import functools import importlib.util import inspect +import itertools import logging import os import sys @@ -495,6 +497,9 @@ def _project_function(obj: Any, scan: _CallScan) -> Any: 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: @@ -610,13 +615,35 @@ def _object_needs_key(obj: Any, scan: _CallScan) -> bool: 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 - cls, module = _project_class(obj, scan) - if cls is None or id(obj) in seen: + 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) + cls, module = _project_class(obj, scan) + if cls is None: + return False if _class_targets(cls, module, scan): return True if inspect.isclass(obj): @@ -629,6 +656,16 @@ def _needs_key(obj: Any, scan: _CallScan, seen: Set[int]) -> bool: 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]: 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 a4f796c1..328cc0ea 100644 --- a/tests/unit/cli/commands/build_utils/test_runtime_scanner.py +++ b/tests/unit/cli/commands/build_utils/test_runtime_scanner.py @@ -884,6 +884,53 @@ async def pure_method(x): 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_unparseable_helper_assumes_key_needed(self, tmp_path): """a helper whose source cannot be read or parsed errs toward the key.""" _write_worker( From b5307193b3f8cc1b18cf2f7f148427aa5ab00dcf Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 19 Oct 2026 13:30:18 -0400 Subject: [PATCH 10/13] docs(deploy): align API key docs and deploy messages with detection rules - document classes, package paths, local imports, wrappers, and the unresolved-import rule; note shared templates and user-set keys - deploy summary says the key was set by Flash, not added - log an explicit message when stale-key cleanup times out - make the shared saveTemplate payload builder public - tests: the cleanup bound actually returns early; a user-set key survives both update branches --- docs/Cross_Endpoint_Routing.md | 10 ++++++---- docs/Deployment_Architecture.md | 2 +- src/runpod_flash/cli/commands/deploy.py | 2 +- src/runpod_flash/cli/utils/deployment.py | 9 +++++++-- src/runpod_flash/core/resources/serverless.py | 4 ++-- tests/unit/cli/test_deploy.py | 2 +- tests/unit/cli/utils/test_deployment.py | 15 ++++++++++++--- tests/unit/resources/test_serverless.py | 14 ++++++++++++++ 8 files changed, 44 insertions(+), 14 deletions(-) diff --git a/docs/Cross_Endpoint_Routing.md b/docs/Cross_Endpoint_Routing.md index 77dae363..c68cc744 100644 --- a/docs/Cross_Endpoint_Routing.md +++ b/docs/Cross_Endpoint_Routing.md @@ -183,17 +183,19 @@ 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 in your project: +`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 (`tasks.square(x)`) or an aliased import -- uses a client `Endpoint(id=...)` or `Endpoint(image=...)`, at module level or built inline +- 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. -If a deployed endpoint fails with `401 Unauthorized` and the error says `RUNPOD_API_KEY is not set on this worker`, the call was not detected (for example, a client `Endpoint` passed in as an argument). Make it visible in one of the ways above and redeploy. +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 diff --git a/docs/Deployment_Architecture.md b/docs/Deployment_Architecture.md index c264600d..81806ff9 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. 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. See [API Key Injection](Cross_Endpoint_Routing.md#api-key-injection) for what counts as a remote call +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/deploy.py b/src/runpod_flash/cli/commands/deploy.py index 64dfb210..46a7fb61 100644 --- a/src/runpod_flash/cli/commands/deploy.py +++ b/src/runpod_flash/cli/commands/deploy.py @@ -180,7 +180,7 @@ def _display_post_deployment_guidance( if keyed: console.print() console.print( - "[dim]RUNPOD_API_KEY added to endpoints that call other endpoints: " + "[dim]RUNPOD_API_KEY set by Flash on endpoints that call other endpoints: " f"{', '.join(keyed)}[/dim]" ) diff --git a/src/runpod_flash/cli/utils/deployment.py b/src/runpod_flash/cli/utils/deployment.py index d102c604..10a5d02b 100644 --- a/src/runpod_flash/cli/utils/deployment.py +++ b/src/runpod_flash/cli/utils/deployment.py @@ -55,7 +55,7 @@ async def _remove_stale_api_key(client: Any, template_id: str) -> bool: env_payload = [ {"key": entry["key"], "value": entry.get("value", "")} for entry in kept ] - payload = ServerlessResource._template_update_payload( + payload = ServerlessResource.template_update_payload( {**live, "env": env_payload}, template_id ) await client.update_template(payload) @@ -104,8 +104,13 @@ async def run() -> None: 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!r}") + 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]: diff --git a/src/runpod_flash/core/resources/serverless.py b/src/runpod_flash/core/resources/serverless.py index d3007faf..67f6eb5a 100644 --- a/src/runpod_flash/core/resources/serverless.py +++ b/src/runpod_flash/core/resources/serverless.py @@ -869,12 +869,12 @@ def _build_template_update_payload( template: Template model with desired configuration. template_id: ID of the template to update. """ - return ServerlessResource._template_update_payload( + return ServerlessResource.template_update_payload( template.model_dump(exclude_none=True, mode="json"), template_id ) @staticmethod - def _template_update_payload( + def template_update_payload( template_data: Dict[str, Any], template_id: str, ) -> Dict[str, Any]: diff --git a/tests/unit/cli/test_deploy.py b/tests/unit/cli/test_deploy.py index a2df5954..ae046b1a 100644 --- a/tests/unit/cli/test_deploy.py +++ b/tests/unit/cli/test_deploy.py @@ -660,7 +660,7 @@ def test_lists_endpoints_that_receive_api_key(self, patched_console): ) output = self._collect_output(patched_console) assert ( - "RUNPOD_API_KEY added to endpoints that call other endpoints: caller" + "RUNPOD_API_KEY set by Flash on endpoints that call other endpoints: caller" in (output) ) assert "caller, isolated" not in output diff --git a/tests/unit/cli/utils/test_deployment.py b/tests/unit/cli/utils/test_deployment.py index d4865a51..d2880bbd 100644 --- a/tests/unit/cli/utils/test_deployment.py +++ b/tests/unit/cli/utils/test_deployment.py @@ -1014,20 +1014,29 @@ async def test_one_failure_does_not_stop_the_others(self): assert client.update_template.await_count == 2 @pytest.mark.asyncio - async def test_slow_api_is_bounded(self): + 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) - with patch( - "runpod_flash.cli.utils.deployment.STALE_KEY_CLEANUP_TIMEOUT_SECONDS", 0.05 + 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 diff --git a/tests/unit/resources/test_serverless.py b/tests/unit/resources/test_serverless.py index 607aa50f..c9dd25ba 100644 --- a/tests/unit/resources/test_serverless.py +++ b/tests/unit/resources/test_serverless.py @@ -3605,6 +3605,20 @@ async def test_keeps_key_with_remote_calls(self, new_env): 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): + env = await self._update(new_env, makes_remote_calls=False) + + assert "RUNPOD_API_KEY" in env + @pytest.mark.asyncio @pytest.mark.parametrize( "new_env", From 2c54090089253bab60bc43ab5271a139cf60ed89 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 19 Oct 2026 13:34:02 -0400 Subject: [PATCH 11/13] fix(build): follow project helpers used as values in partials and containers --- .../cli/commands/build_utils/scanner.py | 4 ++ .../build_utils/test_runtime_scanner.py | 52 ++++++++++++++++++- 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/src/runpod_flash/cli/commands/build_utils/scanner.py b/src/runpod_flash/cli/commands/build_utils/scanner.py index 9e394e6d..ad39b758 100644 --- a/src/runpod_flash/cli/commands/build_utils/scanner.py +++ b/src/runpod_flash/cli/commands/build_utils/scanner.py @@ -641,6 +641,10 @@ def _needs_key(obj: Any, scan: _CallScan, seen: Set[int]) -> bool: 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 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 328cc0ea..61f1f69d 100644 --- a/tests/unit/cli/commands/build_utils/test_runtime_scanner.py +++ b/tests/unit/cli/commands/build_utils/test_runtime_scanner.py @@ -539,7 +539,7 @@ async def pure_helper(x): 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"] + 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): @@ -814,7 +814,10 @@ async def via_init(x): via_init = scanner.discover_remote_functions()[0] finally: sys.modules.pop("init_utils", None) - assert via_init.called_remote_functions == ["call_model"] + 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.""" @@ -931,6 +934,51 @@ async def plain_dict(x): 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( From 2832431eb1cb95094198567299a28270faa31087 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 19 Oct 2026 13:36:57 -0400 Subject: [PATCH 12/13] test(deploy): cover the user-set key guard on env-unchanged updates --- tests/unit/resources/test_serverless.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/unit/resources/test_serverless.py b/tests/unit/resources/test_serverless.py index c9dd25ba..def9362f 100644 --- a/tests/unit/resources/test_serverless.py +++ b/tests/unit/resources/test_serverless.py @@ -3520,11 +3520,11 @@ class TestStaleApiKeyCleanup: once the build manifest says the endpoint makes no remote calls.""" @staticmethod - async def _update(new_env, makes_remote_calls, template_ids=None): + async def _update(new_env, makes_remote_calls, template_ids=None, old_env=None): old_resource = ServerlessEndpoint( name="stale-key", imageName="test:latest", - env={"LOG_LEVEL": "INFO"}, + env=old_env or {"LOG_LEVEL": "INFO"}, flashboot=False, ) old_resource.id = "ep-stale" @@ -3615,7 +3615,8 @@ async def test_keeps_key_with_remote_calls(self, new_env): ids=["env-unchanged", "env-changed"], ) async def test_keeps_user_set_key_without_remote_calls(self, new_env): - env = await self._update(new_env, makes_remote_calls=False) + 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 From 7805547537ae27a8d9b0cd5b75e5ef6a94ea7fee Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 19 Oct 2026 17:23:33 -0400 Subject: [PATCH 13/13] fix(build): treat Endpoint(**kwargs) as a possible client endpoint --- src/runpod_flash/cli/commands/build_utils/scanner.py | 5 +++-- tests/unit/cli/commands/build_utils/test_runtime_scanner.py | 4 +++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/runpod_flash/cli/commands/build_utils/scanner.py b/src/runpod_flash/cli/commands/build_utils/scanner.py index ad39b758..95b4e243 100644 --- a/src/runpod_flash/cli/commands/build_utils/scanner.py +++ b/src/runpod_flash/cli/commands/build_utils/scanner.py @@ -527,14 +527,15 @@ def _project_class(obj: Any, scan: _CallScan) -> Tuple[Any, Any]: def _is_client_endpoint_construction( call: ast.Call, namespace: Dict[str, Any], scan: _CallScan ) -> bool: - """match Endpoint(id=...) or Endpoint(image=...) built inline, incl. aliases.""" + """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 - return any(kw.arg in ("id", "image") for kw in call.keywords) + # 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( 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 61f1f69d..2df49509 100644 --- a/tests/unit/cli/commands/build_utils/test_runtime_scanner.py +++ b/tests/unit/cli/commands/build_utils/test_runtime_scanner.py @@ -369,7 +369,9 @@ def caller(x): result = _analyze_cross_calls_ast(p, {"caller"}, set()) assert result == {"caller": ["RUNPOD_API_KEY"]} - @pytest.mark.parametrize("kwarg", ['id="abc123"', 'image="org/img:1"']) + @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,