fix(apps): give the account api key only to workers that call resources

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

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},

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

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")