From a3799505d8aa777695d7439fc58ece56cf4f30d8 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 16 Oct 2026 20:24:42 -0400 Subject: [PATCH 1/4] fix(apps): give the account api key only to workers that call resources Deploy and dev copied the client's account API key into every endpoint's template env, so any endpoint's workers (and anyone who can read its template) held full account credentials, including endpoints that never call another resource. runpod-flash injected the key only for endpoints whose functions make remote calls. Analyze each resource's module source: the resource body (not its decorators) and the module-level helpers it reaches. The key is added when that code references another resource, calls .remote or .spawn, uses anything imported from runpod, or uses a project-local import the analysis cannot follow. Unreadable source and nested definitions also get the key, so a worker that needs it never goes without. --- runpod/apps/deploy.py | 13 ++- runpod/apps/dev.py | 12 +- runpod/apps/remote_calls.py | 160 ++++++++++++++++++++++++++ tests/test_apps/test_remote_calls.py | 163 +++++++++++++++++++++++++++ 4 files changed, 343 insertions(+), 5 deletions(-) create mode 100644 runpod/apps/remote_calls.py create mode 100644 tests/test_apps/test_remote_calls.py diff --git a/runpod/apps/deploy.py b/runpod/apps/deploy.py index a0ef47ca7..b968f972c 100644 --- a/runpod/apps/deploy.py +++ b/runpod/apps/deploy.py @@ -308,10 +308,17 @@ def _deployed_endpoint_input( } # cross-resource calls from inside a worker go through the sentinel, - # which authenticates with the account api key; without it any - # .remote() from worker code would fail + # which authenticates with the account api key. only resources that + # may make such a call get it (the analysis fails open) + from .remote_calls import makes_remote_calls + api_key = runpod.api_key - if api_key and "RUNPOD_API_KEY" not in template_env: + handle = app.resources.get(spec.name) + if ( + api_key + and "RUNPOD_API_KEY" not in template_env + and (handle is None or makes_remote_calls(app, handle)) + ): template_env["RUNPOD_API_KEY"] = api_key # nested tasks use the same runtime and sdk channels as the endpoint diff --git a/runpod/apps/dev.py b/runpod/apps/dev.py index c4e95ad1c..43a49e9a5 100644 --- a/runpod/apps/dev.py +++ b/runpod/apps/dev.py @@ -107,6 +107,13 @@ def _cpu_locations(instance_ids: List[str]) -> str: GENERATION_ENV = "RUNPOD_DEV_GENERATION" +def _needs_api_key(app: App, spec: ResourceSpec) -> bool: + from .remote_calls import makes_remote_calls + + handle = app.resources.get(spec.name) + return handle is None or makes_remote_calls(app, handle) + + def _endpoint_input(app: App, spec: ResourceSpec, generation: int = 1) -> Dict: """build a live endpoint payload with an atomically bound container config.""" @@ -133,10 +140,11 @@ def _endpoint_input(app: App, spec: ResourceSpec, generation: int = 1) -> Dict: {"key": GENERATION_ENV, "value": str(generation)}, # nested .remote() calls from inside a dev worker need # credentials and the dev-session marker to resolve - # sibling dev endpoints by name + # sibling dev endpoints by name. only resources that may + # make one get the key (the analysis fails open) *( [{"key": "RUNPOD_API_KEY", "value": key}] - if (key := _client_api_key()) + if _needs_api_key(app, spec) and (key := _client_api_key()) else [] ), {"key": "RUNPOD_DEV_APP", "value": app.name}, diff --git a/runpod/apps/remote_calls.py b/runpod/apps/remote_calls.py new file mode 100644 index 000000000..c81434146 --- /dev/null +++ b/runpod/apps/remote_calls.py @@ -0,0 +1,160 @@ +"""decide which resources call other resources from inside a worker. + +workers authenticate nested .remote() calls with the account api key, so +deploy and dev give the key to every resource that may make one. the +analysis reads the resource's module source and fails open: anything it +cannot follow counts as a remote call, so no worker that needs the key +goes without it. +""" + +import ast +import inspect +import logging +from pathlib import Path +from typing import Dict, Optional, Set + +log = logging.getLogger(__name__) + +# handle methods that start a call on another resource +_REMOTE_ATTRS = frozenset({"remote", "spawn"}) + + +def resource_target(handle): + """the user function or class behind a resource handle.""" + return ( + getattr(handle, "_fn", None) + or getattr(handle, "_cls", None) + or getattr(handle, "__wrapped__", None) + ) + + +class _BodyNames(ast.NodeVisitor): + """names and attributes used in a definition's body, not its decorators.""" + + def __init__(self): + self.names: Set[str] = set() + self.attrs: Set[str] = set() + + def visit_FunctionDef(self, node): + self._visit_def(node) + + visit_AsyncFunctionDef = visit_FunctionDef + + def visit_ClassDef(self, node): + for base in node.bases: + self.visit(base) + for stmt in node.body: + self.visit(stmt) + + def visit_Lambda(self, node): + self.visit(node.body) + + def _visit_def(self, node): + for default in node.args.defaults + node.args.kw_defaults: + if default is not None: + self.visit(default) + for stmt in node.body: + self.visit(stmt) + + def visit_Name(self, node): + self.names.add(node.id) + + def visit_Attribute(self, node): + self.attrs.add(node.attr) + self.generic_visit(node) + + +def _body_names(node) -> _BodyNames: + visitor = _BodyNames() + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + visitor.visit(node) + else: + visitor.generic_visit(node) + return visitor + + +def _module_tree(target) -> Optional[ast.Module]: + try: + path = inspect.getsourcefile(target) + return ast.parse(Path(path).read_text(encoding="utf-8")) if path else None + except (OSError, TypeError, SyntaxError, ValueError): + return None + + +def makes_remote_calls(app, handle) -> bool: + """true unless the resource provably cannot reach another resource.""" + target = resource_target(handle) + name = getattr(target, "__name__", None) + if target is None or getattr(target, "__qualname__", name) != name: + # nested or anonymous definitions are not analyzed + return True + tree = _module_tree(target) + if tree is None: + log.debug("source for %s unavailable; giving it the api key", name) + return True + + defs: Dict[str, ast.AST] = {} + assigns: Dict[str, ast.AST] = {} + runpod_names: Set[str] = set() + local_imports: Set[str] = set() + module_file = inspect.getsourcefile(target) or "" + project_dir = Path(module_file).resolve().parent if module_file else None + + def _is_local(module_name: str) -> bool: + if not module_name or project_dir is None: + return True + top = module_name.split(".")[0] + return (project_dir / f"{top}.py").exists() or (project_dir / top).is_dir() + + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + defs[node.name] = node + elif isinstance(node, (ast.Assign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + for t in targets: + for n in ast.walk(t): + if isinstance(n, ast.Name): + assigns[n.id] = node + elif isinstance(node, ast.Import): + for alias in node.names: + bound = alias.asname or alias.name.split(".")[0] + if alias.name.split(".")[0] == "runpod": + runpod_names.add(bound) + elif _is_local(alias.name): + local_imports.add(bound) + elif isinstance(node, ast.ImportFrom): + module_name = node.module or "" + for alias in node.names: + bound = alias.asname or alias.name + if module_name.split(".")[0] == "runpod": + runpod_names.add(bound) + elif node.level or _is_local(module_name): + local_imports.add(bound) + + if name not in defs: + return True + + other_resources = { + resource_target(h).__name__ + for h in app.resources.values() + if h is not handle and resource_target(h) is not None + } + + seen: Set[str] = set() + pending = [name] + while pending: + current = pending.pop() + if current in seen: + continue + seen.add(current) + used = _body_names(defs.get(current) or assigns[current]) + if used.attrs & _REMOTE_ATTRS: + return True + for ref in used.names: + if ref in other_resources or ref in runpod_names or ref in local_imports: + return True + if ref in seen: + continue + if ref in defs or ref in assigns: + pending.append(ref) + return False diff --git a/tests/test_apps/test_remote_calls.py b/tests/test_apps/test_remote_calls.py new file mode 100644 index 000000000..1d2b1745b --- /dev/null +++ b/tests/test_apps/test_remote_calls.py @@ -0,0 +1,163 @@ +"""tests for the worker api key scoping analysis.""" + +import importlib.util +import sys +import textwrap + +import pytest + +from runpod.apps.app import _clear_registry +from runpod.apps.remote_calls import makes_remote_calls + + +@pytest.fixture(autouse=True) +def clean_registry(): + _clear_registry() + yield + _clear_registry() + + +@pytest.fixture +def load(tmp_path, monkeypatch): + monkeypatch.syspath_prepend(str(tmp_path)) + + def _load(source, name="remote_calls_app"): + path = tmp_path / f"{name}.py" + path.write_text(textwrap.dedent(source)) + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + return module + + return _load + + +SOURCE = ''' +import json + +import runpod + +app = runpod.App("calls") + + +def fmt(value): + return json.dumps({"value": value}) + + +def fan_out(items): + return [leaf.remote(x=i) for i in items] + + +@app.queue(name="leaf", cpu="cpu3c-1-2") +def leaf(x: int) -> str: + return fmt(x) + + +@app.queue(name="caller", cpu="cpu3c-1-2") +def caller(n: int) -> list: + return [leaf(x=i) for i in range(n)] + + +@app.queue(name="indirect", cpu="cpu3c-1-2") +def indirect(n: int) -> list: + return fan_out(range(n)) + + +@app.queue(name="client", cpu="cpu3c-1-2") +def client(n: int): + return runpod.Queue(app="other", name="x").remote(n=n) + + +@app.api(name="service", cpu="cpu3c-1-2") +class Service: + @runpod.init + def setup(self): + self.ready = True + + @runpod.post("/echo") + def echo(self, body: dict) -> dict: + return body + + +@runpod.local_entrypoint +def main(): + print(caller.remote(n=2)) +''' + + +class TestMakesRemoteCalls: + def test_classification(self, load): + mod = load(SOURCE) + app = mod.app + result = { + name: makes_remote_calls(app, handle) + for name, handle in app.resources.items() + } + assert result == { + "leaf": False, # helper only, no other resource + "caller": True, # references another resource + "indirect": True, # through a module-level helper + "client": True, # runpod client inside the body + "service": False, # @runpod.* decorators are not calls + } + + def test_local_project_import_counts(self, tmp_path, load): + (tmp_path / "helpers_mod.py").write_text("def run():\n return 1\n") + mod = load( + """ + import runpod + from helpers_mod import run + + app = runpod.App("local") + + @app.queue(name="uses", cpu="cpu3c-1-2") + def uses() -> int: + return run() + """, + name="remote_calls_local", + ) + assert makes_remote_calls(mod.app, mod.app.resources["uses"]) is True + + def test_unanalyzable_targets_fail_open(self): + import runpod + + app = runpod.App("nested") + + @app.queue(name="inner", cpu="cpu3c-1-2") + def inner() -> int: + return 1 + + # defined inside a function: qualname differs, not analyzed + assert makes_remote_calls(app, app.resources["inner"]) is True + + +class TestPayloadsScopeTheKey: + def test_deploy_gives_key_only_to_callers(self, load, monkeypatch): + import runpod + from runpod.apps.deploy import _deployed_endpoint_input + + monkeypatch.setattr(runpod, "api_key", "acct-key") + mod = load(SOURCE, name="remote_calls_deploy") + + def env_keys(name): + payload = _deployed_endpoint_input( + mod.app, mod.app.resources[name].spec, "env-1", "b1", "3.12" + ) + return {e["key"] for e in payload["template"]["env"]} + + assert "RUNPOD_API_KEY" not in env_keys("leaf") + assert "RUNPOD_API_KEY" in env_keys("caller") + + def test_dev_gives_key_only_to_callers(self, load, monkeypatch): + from runpod.apps import dev + + monkeypatch.setattr(dev, "_client_api_key", lambda: "acct-key") + mod = load(SOURCE, name="remote_calls_dev") + + def env_keys(name): + payload = dev._endpoint_input(mod.app, mod.app.resources[name].spec) + return {e["key"] for e in payload["template"]["env"]} + + assert "RUNPOD_API_KEY" not in env_keys("leaf") + assert "RUNPOD_API_KEY" in env_keys("indirect") From 3e3a16c6405552b33b8ed6f614d50ed7bae6df70 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 19 Oct 2026 23:09:08 -0400 Subject: [PATCH 2/4] fix(apps): give the api key for imports, init hooks and other apps' handles the analysis missed four ways a worker reaches another resource: an import inside the resource body, a package import resolved from inside the package, a queue init hook, and a handle from another app in the module. each now gives the key. --- runpod/apps/remote_calls.py | 121 +++++++++++++++++++++++---- tests/test_apps/test_remote_calls.py | 95 +++++++++++++++++++++ 2 files changed, 199 insertions(+), 17 deletions(-) diff --git a/runpod/apps/remote_calls.py b/runpod/apps/remote_calls.py index c81434146..69c2265a6 100644 --- a/runpod/apps/remote_calls.py +++ b/runpod/apps/remote_calls.py @@ -8,10 +8,14 @@ """ import ast +import importlib.util import inspect import logging +import site +import sys +import sysconfig from pathlib import Path -from typing import Dict, Optional, Set +from typing import Dict, Optional, Set, Tuple log = logging.getLogger(__name__) @@ -34,6 +38,8 @@ class _BodyNames(ast.NodeVisitor): def __init__(self): self.names: Set[str] = set() self.attrs: Set[str] = set() + # (module, level) of imports inside the body + self.imports: Set[Tuple[str, int]] = set() def visit_FunctionDef(self, node): self._visit_def(node) @@ -63,6 +69,13 @@ def visit_Attribute(self, node): self.attrs.add(node.attr) self.generic_visit(node) + def visit_Import(self, node): + for alias in node.names: + self.imports.add((alias.name, 0)) + + def visit_ImportFrom(self, node): + self.imports.add((node.module or "", node.level)) + def _body_names(node) -> _BodyNames: visitor = _BodyNames() @@ -73,6 +86,75 @@ def _body_names(node) -> _BodyNames: return visitor +def _installed_dirs() -> Tuple[Path, ...]: + """directories holding the stdlib and installed distributions.""" + paths = {sysconfig.get_paths().get(k) for k in ("stdlib", "platstdlib", "purelib", "platlib")} + try: + paths.update(site.getsitepackages()) + except AttributeError: # some virtualenv site modules lack it + pass + paths.add(site.getusersitepackages()) + return tuple(Path(p).resolve() for p in paths if p) + + +def _project_dirs(module_file: str) -> Tuple[Path, ...]: + """where the module's project-local imports live: its directory, the + directory above its package, and the working directory.""" + here = Path(module_file).resolve().parent + top = here + while (top / "__init__.py").exists() and top.parent != top: + top = top.parent + return tuple(dict.fromkeys((here, top, Path.cwd().resolve()))) + + +def _is_local(module_name: str, project_dirs: Tuple[Path, ...]) -> bool: + """true unless the module is the stdlib or an installed distribution. + + project modules, and modules that resolve outside the installed + locations, count as local: the analysis does not follow them. a + module found nowhere (torch imported inside a gpu function, absent + from the client) is a third-party package, not project code. + """ + top = module_name.split(".")[0] + if not top: + return True + if any((d / f"{top}.py").is_file() or (d / top).is_dir() for d in project_dirs): + return True + if top in sys.builtin_module_names or top in getattr(sys, "stdlib_module_names", ()): + return False + try: + spec = importlib.util.find_spec(top) + except (ImportError, ValueError): + return True + if spec is None: + return False + locations = [spec.origin] if spec.origin and spec.has_location else [] + locations += list(spec.submodule_search_locations or []) + installed = _installed_dirs() + return any( + not any(Path(loc).resolve().is_relative_to(d) for d in installed) + for loc in locations + ) + + +def _resource_names(target, app, handle) -> Set[str]: + """names bound to resource handles other than this one: the app's + resources and any handle in the module's globals (other apps).""" + from .handles import ApiHandle, FunctionHandle + + names = { + getattr(resource_target(h), "__name__", None) + for h in app.resources.values() + if h is not handle + } + module = inspect.getmodule(target) + for key, value in vars(module).items() if module else (): + if isinstance(value, (FunctionHandle, ApiHandle)) and value is not handle: + names.add(key) + names.discard(None) + return names + + def _module_tree(target) -> Optional[ast.Module]: try: path = inspect.getsourcefile(target) @@ -97,14 +179,13 @@ def makes_remote_calls(app, handle) -> bool: assigns: Dict[str, ast.AST] = {} runpod_names: Set[str] = set() local_imports: Set[str] = set() - module_file = inspect.getsourcefile(target) or "" - project_dir = Path(module_file).resolve().parent if module_file else None + project_dirs = _project_dirs(inspect.getsourcefile(target)) - def _is_local(module_name: str) -> bool: - if not module_name or project_dir is None: - return True - top = module_name.split(".")[0] - return (project_dir / f"{top}.py").exists() or (project_dir / top).is_dir() + def _is_local_import(module_name: str, level: int) -> bool: + return bool(level) or _is_local(module_name, project_dirs) + + def _uncertain_import(module_name: str, level: int) -> bool: + return module_name.split(".")[0] == "runpod" or _is_local_import(module_name, level) for node in tree.body: if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): @@ -120,7 +201,7 @@ def _is_local(module_name: str) -> bool: bound = alias.asname or alias.name.split(".")[0] if alias.name.split(".")[0] == "runpod": runpod_names.add(bound) - elif _is_local(alias.name): + elif _is_local_import(alias.name, 0): local_imports.add(bound) elif isinstance(node, ast.ImportFrom): module_name = node.module or "" @@ -128,20 +209,24 @@ def _is_local(module_name: str) -> bool: bound = alias.asname or alias.name if module_name.split(".")[0] == "runpod": runpod_names.add(bound) - elif node.level or _is_local(module_name): + elif _is_local_import(module_name, node.level): local_imports.add(bound) - if name not in defs: + # worker startup hooks run on the worker too + roots = [name] + init_fn = getattr(handle, "_init_fn", None) + if init_fn is not None: + init_name = getattr(init_fn, "__name__", None) + if getattr(init_fn, "__qualname__", init_name) != init_name: + return True + roots.append(init_name) + if any(root not in defs for root in roots): return True - other_resources = { - resource_target(h).__name__ - for h in app.resources.values() - if h is not handle and resource_target(h) is not None - } + other_resources = _resource_names(target, app, handle) seen: Set[str] = set() - pending = [name] + pending = list(roots) while pending: current = pending.pop() if current in seen: @@ -150,6 +235,8 @@ def _is_local(module_name: str) -> bool: used = _body_names(defs.get(current) or assigns[current]) if used.attrs & _REMOTE_ATTRS: return True + if any(_uncertain_import(m, level) for m, level in used.imports): + return True for ref in used.names: if ref in other_resources or ref in runpod_names or ref in local_imports: return True diff --git a/tests/test_apps/test_remote_calls.py b/tests/test_apps/test_remote_calls.py index 1d2b1745b..1266e2e4e 100644 --- a/tests/test_apps/test_remote_calls.py +++ b/tests/test_apps/test_remote_calls.py @@ -119,6 +119,101 @@ def uses() -> int: ) assert makes_remote_calls(mod.app, mod.app.resources["uses"]) is True + def test_import_inside_the_body_counts(self, tmp_path, load): + (tmp_path / "helpers_body.py").write_text("def run():\n return 1\n") + mod = load( + """ + import runpod + + app = runpod.App("bodyimport") + + @app.queue(name="uses", cpu="cpu3c-1-2") + def uses() -> int: + from helpers_body import run + return run() + + @app.queue(name="gpu", gpu="NVIDIA GeForce RTX 4090") + def gpu() -> str: + import json + import not_installed_anywhere + return json.dumps(not_installed_anywhere.x) + """, + name="remote_calls_body", + ) + assert makes_remote_calls(mod.app, mod.app.resources["uses"]) is True + # stdlib and packages absent from the client (torch on a gpu + # worker) are not project code + assert makes_remote_calls(mod.app, mod.app.resources["gpu"]) is False + + def test_package_import_from_inside_the_package_counts(self, tmp_path, load): + pkg = tmp_path / "mypkg_rc" + pkg.mkdir() + (pkg / "__init__.py").write_text("") + (pkg / "helpers.py").write_text("def run():\n return 1\n") + source = textwrap.dedent( + """ + import runpod + from mypkg_rc.helpers import run + + app = runpod.App("pkgimport") + + @app.queue(name="uses", cpu="cpu3c-1-2") + def uses() -> int: + return run() + """ + ) + (pkg / "worker.py").write_text(source) + import importlib + + mod = importlib.import_module("mypkg_rc.worker") + assert makes_remote_calls(mod.app, mod.app.resources["uses"]) is True + + def test_init_hook_calls_count(self, load): + mod = load( + """ + import runpod + + app = runpod.App("inithook") + + @app.queue(name="sibling", cpu="cpu3c-1-2") + def sibling() -> int: + return 1 + + @app.queue(name="leaf", cpu="cpu3c-1-2") + def leaf() -> int: + return 2 + + @leaf.init + def warm(): + sibling.remote() + """, + name="remote_calls_init", + ) + assert makes_remote_calls(mod.app, mod.app.resources["leaf"]) is True + assert makes_remote_calls(mod.app, mod.app.resources["sibling"]) is False + + def test_handles_of_other_apps_count(self, load): + mod = load( + """ + import runpod + + other = runpod.App("other-app") + app = runpod.App("main-app") + + @other.api(name="svc", cpu="cpu3c-1-2") + class Svc: + @runpod.post("/x") + def x(self, body: dict) -> dict: + return body + + @app.queue(name="calls_api", cpu="cpu3c-1-2") + def calls_api() -> dict: + return Svc.post("/x", json={}) + """, + name="remote_calls_other", + ) + assert makes_remote_calls(mod.app, mod.app.resources["calls_api"]) is True + def test_unanalyzable_targets_fail_open(self): import runpod From df0399f191eee09da1377c3e4cf92bae490670b4 Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 20 Oct 2026 12:28:31 -0400 Subject: [PATCH 3/4] refactor(apps): split the remote-call analysis into named steps --- runpod/apps/deploy.py | 7 +- runpod/apps/dev.py | 12 +- runpod/apps/remote_calls.py | 226 ++++++++++++++++++++---------------- 3 files changed, 134 insertions(+), 111 deletions(-) diff --git a/runpod/apps/deploy.py b/runpod/apps/deploy.py index b968f972c..e7e32028a 100644 --- a/runpod/apps/deploy.py +++ b/runpod/apps/deploy.py @@ -309,15 +309,14 @@ def _deployed_endpoint_input( # cross-resource calls from inside a worker go through the sentinel, # which authenticates with the account api key. only resources that - # may make such a call get it (the analysis fails open) - from .remote_calls import makes_remote_calls + # may make such a call get it + from .remote_calls import may_call_resources api_key = runpod.api_key - handle = app.resources.get(spec.name) if ( api_key and "RUNPOD_API_KEY" not in template_env - and (handle is None or makes_remote_calls(app, handle)) + and may_call_resources(app, spec) ): template_env["RUNPOD_API_KEY"] = api_key diff --git a/runpod/apps/dev.py b/runpod/apps/dev.py index 43a49e9a5..7d1f04cf2 100644 --- a/runpod/apps/dev.py +++ b/runpod/apps/dev.py @@ -14,6 +14,7 @@ from .app import App from .errors import AppError from .handles import ApiHandle, FunctionHandle +from .remote_calls import may_call_resources from .shim import runtime_launcher from .spec import ResourceKind, ResourceSpec from .targets import LiveTarget @@ -107,13 +108,6 @@ def _cpu_locations(instance_ids: List[str]) -> str: GENERATION_ENV = "RUNPOD_DEV_GENERATION" -def _needs_api_key(app: App, spec: ResourceSpec) -> bool: - from .remote_calls import makes_remote_calls - - handle = app.resources.get(spec.name) - return handle is None or makes_remote_calls(app, handle) - - def _endpoint_input(app: App, spec: ResourceSpec, generation: int = 1) -> Dict: """build a live endpoint payload with an atomically bound container config.""" @@ -141,10 +135,10 @@ def _endpoint_input(app: App, spec: ResourceSpec, generation: int = 1) -> Dict: # nested .remote() calls from inside a dev worker need # credentials and the dev-session marker to resolve # sibling dev endpoints by name. only resources that may - # make one get the key (the analysis fails open) + # make one get the key *( [{"key": "RUNPOD_API_KEY", "value": key}] - if _needs_api_key(app, spec) and (key := _client_api_key()) + if may_call_resources(app, spec) and (key := _client_api_key()) else [] ), {"key": "RUNPOD_DEV_APP", "value": app.name}, diff --git a/runpod/apps/remote_calls.py b/runpod/apps/remote_calls.py index 69c2265a6..ba1a336c5 100644 --- a/runpod/apps/remote_calls.py +++ b/runpod/apps/remote_calls.py @@ -1,25 +1,27 @@ """decide which resources call other resources from inside a worker. workers authenticate nested .remote() calls with the account api key, so -deploy and dev give the key to every resource that may make one. the -analysis reads the resource's module source and fails open: anything it -cannot follow counts as a remote call, so no worker that needs the key -goes without it. +deploy and dev give the key only to resources that may make one. the +analysis reads the resource's module source and errs toward the key: a +resource gets it when its source is unavailable, when it is defined +inside another function, or when it uses runpod or project code that the +analysis does not follow. """ import ast +import functools import importlib.util import inspect import logging import site import sys import sysconfig +from dataclasses import dataclass, field from pathlib import Path from typing import Dict, Optional, Set, Tuple log = logging.getLogger(__name__) -# handle methods that start a call on another resource _REMOTE_ATTRS = frozenset({"remote", "spawn"}) @@ -32,17 +34,38 @@ def resource_target(handle): ) +def may_call_resources(app, spec) -> bool: + """true unless the resource named by spec provably makes no remote calls.""" + handle = app.resources.get(spec.name) + return handle is None or makes_remote_calls(app, handle) + + +def _top_module(module_name: str) -> str: + return module_name.split(".")[0] + + +def _is_module_level(fn) -> bool: + """true for a plain module-level definition the analysis can find.""" + name = getattr(fn, "__name__", None) + return name is not None and getattr(fn, "__qualname__", name) == name + + class _BodyNames(ast.NodeVisitor): - """names and attributes used in a definition's body, not its decorators.""" + """names, attributes, and imports used in a definition's body, not + its decorators.""" def __init__(self): self.names: Set[str] = set() self.attrs: Set[str] = set() - # (module, level) of imports inside the body + # (module, relative import level) self.imports: Set[Tuple[str, int]] = set() def visit_FunctionDef(self, node): - self._visit_def(node) + for default in node.args.defaults + node.args.kw_defaults: + if default is not None: + self.visit(default) + for stmt in node.body: + self.visit(stmt) visit_AsyncFunctionDef = visit_FunctionDef @@ -55,13 +78,6 @@ def visit_ClassDef(self, node): def visit_Lambda(self, node): self.visit(node.body) - def _visit_def(self, node): - for default in node.args.defaults + node.args.kw_defaults: - if default is not None: - self.visit(default) - for stmt in node.body: - self.visit(stmt) - def visit_Name(self, node): self.names.add(node.id) @@ -70,14 +86,13 @@ def visit_Attribute(self, node): self.generic_visit(node) def visit_Import(self, node): - for alias in node.names: - self.imports.add((alias.name, 0)) + self.imports.update((alias.name, 0) for alias in node.names) def visit_ImportFrom(self, node): self.imports.add((node.module or "", node.level)) -def _body_names(node) -> _BodyNames: +def _body_names(node: ast.AST) -> _BodyNames: visitor = _BodyNames() if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): visitor.visit(node) @@ -86,9 +101,13 @@ def _body_names(node) -> _BodyNames: return visitor +@functools.lru_cache(maxsize=1) def _installed_dirs() -> Tuple[Path, ...]: """directories holding the stdlib and installed distributions.""" - paths = {sysconfig.get_paths().get(k) for k in ("stdlib", "platstdlib", "purelib", "platlib")} + paths = { + sysconfig.get_paths().get(key) + for key in ("stdlib", "platstdlib", "purelib", "platlib") + } try: paths.update(site.getsitepackages()) except AttributeError: # some virtualenv site modules lack it @@ -107,6 +126,11 @@ def _project_dirs(module_file: str) -> Tuple[Path, ...]: return tuple(dict.fromkeys((here, top, Path.cwd().resolve()))) +def _is_installed(path: str) -> bool: + resolved = Path(path).resolve() + return any(resolved.is_relative_to(d) for d in _installed_dirs()) + + def _is_local(module_name: str, project_dirs: Tuple[Path, ...]) -> bool: """true unless the module is the stdlib or an installed distribution. @@ -115,12 +139,14 @@ def _is_local(module_name: str, project_dirs: Tuple[Path, ...]) -> bool: module found nowhere (torch imported inside a gpu function, absent from the client) is a third-party package, not project code. """ - top = module_name.split(".")[0] + top = _top_module(module_name) if not top: return True if any((d / f"{top}.py").is_file() or (d / top).is_dir() for d in project_dirs): return True - if top in sys.builtin_module_names or top in getattr(sys, "stdlib_module_names", ()): + if top in sys.builtin_module_names or top in getattr( + sys, "stdlib_module_names", () + ): return False try: spec = importlib.util.find_spec(top) @@ -130,11 +156,7 @@ def _is_local(module_name: str, project_dirs: Tuple[Path, ...]) -> bool: return False locations = [spec.origin] if spec.origin and spec.has_location else [] locations += list(spec.submodule_search_locations or []) - installed = _installed_dirs() - return any( - not any(Path(loc).resolve().is_relative_to(d) for d in installed) - for loc in locations - ) + return not all(_is_installed(loc) for loc in locations) def _resource_names(target, app, handle) -> Set[str]: @@ -148,83 +170,65 @@ def _resource_names(target, app, handle) -> Set[str]: if h is not handle } module = inspect.getmodule(target) - for key, value in vars(module).items() if module else (): - if isinstance(value, (FunctionHandle, ApiHandle)) and value is not handle: - names.add(key) + module_globals = vars(module).items() if module else () + names.update( + key + for key, value in module_globals + if isinstance(value, (FunctionHandle, ApiHandle)) and value is not handle + ) names.discard(None) return names -def _module_tree(target) -> Optional[ast.Module]: - try: - path = inspect.getsourcefile(target) - return ast.parse(Path(path).read_text(encoding="utf-8")) if path else None - except (OSError, TypeError, SyntaxError, ValueError): - return None - - -def makes_remote_calls(app, handle) -> bool: - """true unless the resource provably cannot reach another resource.""" - target = resource_target(handle) - name = getattr(target, "__name__", None) - if target is None or getattr(target, "__qualname__", name) != name: - # nested or anonymous definitions are not analyzed - return True - tree = _module_tree(target) - if tree is None: - log.debug("source for %s unavailable; giving it the api key", name) - return True - - defs: Dict[str, ast.AST] = {} - assigns: Dict[str, ast.AST] = {} - runpod_names: Set[str] = set() - local_imports: Set[str] = set() - project_dirs = _project_dirs(inspect.getsourcefile(target)) +@dataclass +class _ModuleIndex: + """a module's top-level definitions, assignments, and the names it + imports from runpod or project code.""" - def _is_local_import(module_name: str, level: int) -> bool: - return bool(level) or _is_local(module_name, project_dirs) + project_dirs: Tuple[Path, ...] + defs: Dict[str, ast.AST] = field(default_factory=dict) + assigns: Dict[str, ast.AST] = field(default_factory=dict) + call_capable: Set[str] = field(default_factory=set) - def _uncertain_import(module_name: str, level: int) -> bool: - return module_name.split(".")[0] == "runpod" or _is_local_import(module_name, level) + def import_may_call(self, module_name: str, level: int) -> bool: + """true for runpod and project imports, which can reach a resource.""" + return ( + bool(level) + or _top_module(module_name) == "runpod" + or _is_local(module_name, self.project_dirs) + ) - for node in tree.body: + def add(self, node: ast.stmt) -> None: if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): - defs[node.name] = node + self.defs[node.name] = node elif isinstance(node, (ast.Assign, ast.AnnAssign)): targets = node.targets if isinstance(node, ast.Assign) else [node.target] - for t in targets: - for n in ast.walk(t): + for target in targets: + for n in ast.walk(target): if isinstance(n, ast.Name): - assigns[n.id] = node + self.assigns[n.id] = node elif isinstance(node, ast.Import): - for alias in node.names: - bound = alias.asname or alias.name.split(".")[0] - if alias.name.split(".")[0] == "runpod": - runpod_names.add(bound) - elif _is_local_import(alias.name, 0): - local_imports.add(bound) - elif isinstance(node, ast.ImportFrom): - module_name = node.module or "" - for alias in node.names: - bound = alias.asname or alias.name - if module_name.split(".")[0] == "runpod": - runpod_names.add(bound) - elif _is_local_import(module_name, node.level): - local_imports.add(bound) - - # worker startup hooks run on the worker too - roots = [name] - init_fn = getattr(handle, "_init_fn", None) - if init_fn is not None: - init_name = getattr(init_fn, "__name__", None) - if getattr(init_fn, "__qualname__", init_name) != init_name: - return True - roots.append(init_name) - if any(root not in defs for root in roots): - return True + self.call_capable.update( + alias.asname or _top_module(alias.name) + for alias in node.names + if self.import_may_call(alias.name, 0) + ) + elif isinstance(node, ast.ImportFrom) and self.import_may_call( + node.module or "", node.level + ): + self.call_capable.update(alias.asname or alias.name for alias in node.names) + + +def _source_tree(path: Optional[str]) -> Optional[ast.Module]: + try: + return ast.parse(Path(path).read_text(encoding="utf-8")) if path else None + except (OSError, SyntaxError, ValueError): + return None - other_resources = _resource_names(target, app, handle) +def _reaches_resource(index: _ModuleIndex, roots, other_resources: Set[str]) -> bool: + """walk the roots and every module-level helper they use, looking for + a remote call, a resource, or a runpod or project import.""" seen: Set[str] = set() pending = list(roots) while pending: @@ -232,16 +236,42 @@ def _uncertain_import(module_name: str, level: int) -> bool: if current in seen: continue seen.add(current) - used = _body_names(defs.get(current) or assigns[current]) + used = _body_names(index.defs.get(current) or index.assigns[current]) if used.attrs & _REMOTE_ATTRS: return True - if any(_uncertain_import(m, level) for m, level in used.imports): + if any(index.import_may_call(m, level) for m, level in used.imports): + return True + if used.names & (other_resources | index.call_capable): return True - for ref in used.names: - if ref in other_resources or ref in runpod_names or ref in local_imports: - return True - if ref in seen: - continue - if ref in defs or ref in assigns: - pending.append(ref) + pending.extend( + ref + for ref in used.names - seen + if ref in index.defs or ref in index.assigns + ) return False + + +def makes_remote_calls(app, handle) -> bool: + """true unless the resource provably cannot reach another resource.""" + target = resource_target(handle) + init_fn = getattr(handle, "_init_fn", None) + # worker startup hooks run on the worker too + hooks = [init_fn] if init_fn is not None else [] + if target is None or not all(_is_module_level(fn) for fn in [target, *hooks]): + return True + try: + source_file = inspect.getsourcefile(target) + except TypeError: + source_file = None + tree = _source_tree(source_file) + if tree is None: + log.debug("source for %s unavailable; giving it the api key", target.__name__) + return True + + index = _ModuleIndex(_project_dirs(source_file)) + for node in tree.body: + index.add(node) + roots = [fn.__name__ for fn in [target, *hooks]] + if any(root not in index.defs for root in roots): + return True + return _reaches_resource(index, roots, _resource_names(target, app, handle)) From 4b37af54e692106c1236b956dd984ec14d4825ec Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 20 Oct 2026 12:50:11 -0400 Subject: [PATCH 4/4] refactor(apps): read makes_remote_calls as three named steps --- runpod/apps/remote_calls.py | 171 +++++++++++++++++++++++------------- 1 file changed, 108 insertions(+), 63 deletions(-) diff --git a/runpod/apps/remote_calls.py b/runpod/apps/remote_calls.py index ba1a336c5..a4929f28e 100644 --- a/runpod/apps/remote_calls.py +++ b/runpod/apps/remote_calls.py @@ -18,11 +18,12 @@ import sysconfig from dataclasses import dataclass, field from pathlib import Path -from typing import Dict, Optional, Set, Tuple +from typing import Callable, Dict, List, Optional, Set, Tuple log = logging.getLogger(__name__) _REMOTE_ATTRS = frozenset({"remote", "spawn"}) +_DEF_NODES = (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef) def resource_target(handle): @@ -94,7 +95,7 @@ def visit_ImportFrom(self, node): def _body_names(node: ast.AST) -> _BodyNames: visitor = _BodyNames() - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + if isinstance(node, _DEF_NODES): visitor.visit(node) else: visitor.generic_visit(node) @@ -131,23 +132,21 @@ def _is_installed(path: str) -> bool: return any(resolved.is_relative_to(d) for d in _installed_dirs()) -def _is_local(module_name: str, project_dirs: Tuple[Path, ...]) -> bool: - """true unless the module is the stdlib or an installed distribution. +def _in_project(top: str, project_dirs: Tuple[Path, ...]) -> bool: + return any((d / f"{top}.py").is_file() or (d / top).is_dir() for d in project_dirs) - project modules, and modules that resolve outside the installed - locations, count as local: the analysis does not follow them. a - module found nowhere (torch imported inside a gpu function, absent - from the client) is a third-party package, not project code. - """ - top = _top_module(module_name) - if not top: - return True - if any((d / f"{top}.py").is_file() or (d / top).is_dir() for d in project_dirs): - return True - if top in sys.builtin_module_names or top in getattr( + +def _is_stdlib(top: str) -> bool: + return top in sys.builtin_module_names or top in getattr( sys, "stdlib_module_names", () - ): - return False + ) + + +def _resolves_outside_installed(top: str) -> bool: + """true when the module resolves outside the stdlib and installed + distributions, or fails to resolve. a module found nowhere (torch + imported inside a gpu function, absent from the client) is a + third-party package, so it does not count.""" try: spec = importlib.util.find_spec(top) except (ImportError, ValueError): @@ -159,6 +158,15 @@ def _is_local(module_name: str, project_dirs: Tuple[Path, ...]) -> bool: return not all(_is_installed(loc) for loc in locations) +def _is_local(module_name: str, project_dirs: Tuple[Path, ...]) -> bool: + """true unless the module is the stdlib or an installed distribution. + local modules are not followed, so a resource using one gets the key.""" + top = _top_module(module_name) + if not top or _in_project(top, project_dirs): + return True + return not _is_stdlib(top) and _resolves_outside_installed(top) + + def _resource_names(target, app, handle) -> Set[str]: """names bound to resource handles other than this one: the app's resources and any handle in the module's globals (other apps).""" @@ -199,24 +207,36 @@ def import_may_call(self, module_name: str, level: int) -> bool: ) def add(self, node: ast.stmt) -> None: - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + if isinstance(node, _DEF_NODES): self.defs[node.name] = node elif isinstance(node, (ast.Assign, ast.AnnAssign)): - targets = node.targets if isinstance(node, ast.Assign) else [node.target] - for target in targets: - for n in ast.walk(target): - if isinstance(n, ast.Name): - self.assigns[n.id] = node - elif isinstance(node, ast.Import): - self.call_capable.update( + self.assigns.update(dict.fromkeys(_assigned_names(node), node)) + elif isinstance(node, (ast.Import, ast.ImportFrom)): + self.call_capable.update(self._call_capable_imports(node)) + + def _call_capable_imports(self, node) -> List[str]: + """names an import binds when it comes from runpod or project code.""" + if isinstance(node, ast.Import): + return [ alias.asname or _top_module(alias.name) for alias in node.names if self.import_may_call(alias.name, 0) - ) - elif isinstance(node, ast.ImportFrom) and self.import_may_call( - node.module or "", node.level - ): - self.call_capable.update(alias.asname or alias.name for alias in node.names) + ] + if not self.import_may_call(node.module or "", node.level): + return [] + return [alias.asname or alias.name for alias in node.names] + + def node(self, name: str) -> ast.AST: + return self.defs.get(name) or self.assigns[name] + + def helpers(self, names: Set[str]) -> List[str]: + """the names that refer to module-level definitions or values.""" + return [n for n in names if n in self.defs or n in self.assigns] + + +def _assigned_names(node) -> List[str]: + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + return [n.id for t in targets for n in ast.walk(t) if isinstance(n, ast.Name)] def _source_tree(path: Optional[str]) -> Optional[ast.Module]: @@ -226,52 +246,77 @@ def _source_tree(path: Optional[str]) -> Optional[ast.Module]: return None +def _uses_resource( + used: _BodyNames, index: _ModuleIndex, other_resources: Set[str] +) -> bool: + """true if a body calls .remote()/.spawn(), names another resource, or + uses runpod or project code.""" + return bool( + used.attrs & _REMOTE_ATTRS + or used.names & (other_resources | index.call_capable) + or any(index.import_may_call(m, level) for m, level in used.imports) + ) + + def _reaches_resource(index: _ModuleIndex, roots, other_resources: Set[str]) -> bool: - """walk the roots and every module-level helper they use, looking for - a remote call, a resource, or a runpod or project import.""" + """true if the roots, or any module-level helper they use, use a resource.""" seen: Set[str] = set() pending = list(roots) while pending: - current = pending.pop() - if current in seen: + name = pending.pop() + if name in seen: continue - seen.add(current) - used = _body_names(index.defs.get(current) or index.assigns[current]) - if used.attrs & _REMOTE_ATTRS: - return True - if any(index.import_may_call(m, level) for m, level in used.imports): + seen.add(name) + used = _body_names(index.node(name)) + if _uses_resource(used, index, other_resources): return True - if used.names & (other_resources | index.call_capable): - return True - pending.extend( - ref - for ref in used.names - seen - if ref in index.defs or ref in index.assigns - ) + pending.extend(index.helpers(used.names - seen)) return False -def makes_remote_calls(app, handle) -> bool: - """true unless the resource provably cannot reach another resource.""" +def _worker_entry_points(handle) -> Optional[List[Callable]]: + """the functions a worker runs for this resource: its target and any + startup hook. None when one is nested or anonymous, since the + analysis cannot find it in the module source.""" target = resource_target(handle) init_fn = getattr(handle, "_init_fn", None) - # worker startup hooks run on the worker too - hooks = [init_fn] if init_fn is not None else [] - if target is None or not all(_is_module_level(fn) for fn in [target, *hooks]): - return True + entry_points = [fn for fn in (target, init_fn) if fn is not None] + if target is None or not all(_is_module_level(fn) for fn in entry_points): + return None + return entry_points + + +def _index_module(fn: Callable) -> Optional[_ModuleIndex]: + """index the module that defines fn, or None without its source.""" try: - source_file = inspect.getsourcefile(target) - except TypeError: - source_file = None - tree = _source_tree(source_file) + path = inspect.getsourcefile(fn) + except TypeError: # builtins and c extensions + path = None + tree = _source_tree(path) if tree is None: - log.debug("source for %s unavailable; giving it the api key", target.__name__) - return True - - index = _ModuleIndex(_project_dirs(source_file)) + log.debug("source for %s unavailable; giving it the api key", fn.__name__) + return None + index = _ModuleIndex(_project_dirs(path)) for node in tree.body: index.add(node) - roots = [fn.__name__ for fn in [target, *hooks]] - if any(root not in index.defs for root in roots): + return index + + +def makes_remote_calls(app, handle) -> bool: + """true unless the resource provably cannot reach another resource. + + whenever the analysis cannot follow the resource, it answers true, so + a worker that needs the api key never goes without it. + """ + entry_points = _worker_entry_points(handle) + if entry_points is None: return True - return _reaches_resource(index, roots, _resource_names(target, app, handle)) + + target = entry_points[0] + index = _index_module(target) + names = [fn.__name__ for fn in entry_points] + if index is None or not all(name in index.defs for name in names): + return True + + other_resources = _resource_names(target, app, handle) + return _reaches_resource(index, names, other_resources)