diff --git a/runpod/apps/deploy.py b/runpod/apps/deploy.py index a0ef47ca..b968f972 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 c4e95ad1..43a49e9a 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 00000000..c8143414 --- /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 00000000..1d2b1745 --- /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")