feat: invoke deployed flash tasks with runtime credentials (CON-1833)

runpod/__init__.py

@@ -54,6 +54,7 @@
     "Model": "runpod.apps",
     "NetworkVolume": "runpod.apps",
     "Queue": "runpod.apps",
+    "Task": "runpod.apps",
     "Secret": "runpod.apps",
     "Volume": "runpod.apps",
     "is_local": "runpod.apps",

runpod/api/graphql.py

@@ -87,6 +87,7 @@ async def run_graphql_query_async(
     variables: Optional[Dict[str, Any]] = None,
     timeout: float = 60.0,
     anonymous: bool = False,
+    headers: Optional[Dict[str, str]] = None,
 ) -> Dict[str, Any]:
     """
     Async variant of run_graphql_query, sharing the same url, headers,
@@ -110,7 +111,7 @@ async def run_graphql_query_async(
     async with aiohttp.ClientSession(timeout=client_timeout) as session:
         async with session.post(
             _graphql_url(),
-            headers=_build_headers(effective_api_key),
+            headers={**_build_headers(effective_api_key), **(headers or {})},
             json=_build_payload(query, variables),
         ) as response:
             if response.status == HTTP_STATUS_UNAUTHORIZED:

runpod/api/mutations/apps.py

@@ -88,7 +88,7 @@
 
 MUTATION_FINALIZE_ARTIFACT_UPLOAD = """
 mutation FinalizeArtifactUpload($input: FinalizeFlashArtifactUploadInput!) {
-    finalizeFlashArtifactUpload(input: $input) { id manifest }
+    finalizeFlashArtifactUpload(input: $input) { id manifest resourceSpec }
 }
 """
 

runpod/apps/__init__.py

@@ -15,7 +15,7 @@
 from .markers import delete, get, init, patch, post, put
 from .schedule import schedule
 from .spec import ResourceKind, ResourceSpec, RouteSpec
-from .stubs import Api, Queue
+from .stubs import Api, Queue, Task
 from .model import Model
 from .secret import Secret
 from .volume import GlobalVolume, NetworkVolume, Volume
@@ -34,6 +34,7 @@
     "Model",
     "NetworkVolume",
     "Queue",
+    "Task",
     "Secret",
     "Volume",
     "RemoteExecutionError",

runpod/apps/api.py

@@ -19,6 +19,41 @@
 _TRANSPORT_RETRIES = 4
 
 
+class FlashTaskApiClient:
+    """named task lifecycle operations with runtime-scoped authentication."""
+
+    async def _execute(self, query: str, task_input: Dict[str, Any]) -> Dict[str, Any]:
+        from .utils.network import api_key, runtime_headers
+
+        headers = runtime_headers()
+        response = await run_graphql_query_async(
+            query,
+            variables={"input": task_input},
+            api_key=(headers["Authorization"][7:] if headers else api_key()),
+            headers=headers,
+        )
+        return response["data"]
+
+    async def create(self, task_input: Dict[str, Any]) -> Dict[str, Any]:
+        data = await self._execute(
+            """mutation FlashTaskCreate($input: FlashTaskCreateInput!) {
+                flashTaskCreate(input: $input) { podId flashBuildId resourceName }
+            }""",
+            task_input,
+        )
+        return data["flashTaskCreate"]
+
+    async def terminate(self, pod_id: str, token: str) -> None:
+        data = await self._execute(
+            """mutation FlashTaskTerminate($input: FlashTaskTerminateInput!) {
+                flashTaskTerminate(input: $input)
+            }""",
+            {"podId": pod_id, "executionToken": token},
+        )
+        if data.get("flashTaskTerminate") is not True:
+            raise RuntimeError("task termination was not acknowledged")
+
+
 def _rest_endpoint_input(endpoint: Dict[str, Any]) -> Optional[Dict[str, Any]]:
     """translate only saves whose capabilities rest can preserve atomically."""
     fields = {

runpod/apps/app.py

@@ -259,6 +259,12 @@ async def _resolve(self, spec: ResourceSpec) -> InvocationTarget:
         """resolve a resource spec to an invocation target."""
         spec.validate()
         if spec.kind is ResourceKind.TASK:
+            from .context import deployed_identity
+            from .targets import DeployedTaskTarget
+
+            identity = deployed_identity()
+            if identity is not None:
+                return DeployedTaskTarget(self.name, identity[1], spec.name)
             handle = self._resources.get(spec.name)
             fn = getattr(handle, "_fn", None)
             return PodTarget(

runpod/apps/context.py

@@ -7,7 +7,7 @@
 import threading
 import time
 from enum import Enum
-from typing import Any, Coroutine
+from typing import Any, Coroutine, Optional, Tuple
 
 log = logging.getLogger(__name__)
 CLEANUP_TIMEOUT = 30.0
@@ -36,6 +36,27 @@ def current_context() -> Context:
     return Context.LOCAL
 
 
+def deployed_identity() -> Optional[Tuple[str, str, str]]:
+    """app, environment and resource identity for a deployed worker."""
+    if not (
+        os.getenv("RUNPOD_RUNTIME_KIND") == "task"
+        or os.getenv("RUNPOD_BUILD_ID")
+        or os.getenv("FLASH_APP_NAME")
+        or os.getenv("FLASH_RESOURCE_NAME")
+    ):
+        return None
+    fields = tuple(
+        os.getenv(key)
+        for key in ("FLASH_APP_NAME", "FLASH_ENVIRONMENT", "FLASH_RESOURCE_NAME")
+    )
+    build = os.getenv("RUNPOD_BUILD_ID")
+    if any(not value or value != value.strip() for value in (*fields, build)):
+        raise RuntimeError(
+            "deployed flash requires app, environment, resource and build identity"
+        )
+    return fields
+
+
 def is_local() -> bool:
     """true when not running inside a runpod container.
 

runpod/apps/deploy.py

@@ -11,8 +11,6 @@
 
 from pathspec import GitIgnoreSpec
 
-import runpod
-
 from .api import AppsApiClient
 from .app import App
 from .build import (
@@ -137,6 +135,17 @@ def build_manifest(
             )
 
         entry = spec.to_manifest()
+        if spec.kind is ResourceKind.TASK:
+            from .images import image_for_spec
+
+            entry["imageName"] = image_for_spec(spec, python_version=python_version)
+            entry["customImage"] = bool(spec.image)
+            for mount, volume in zip(entry.get("mounts", []), spec.mounts.values()):
+                mount["create"] = volume.create
+                if volume.kind == "network":
+                    mount["size"] = volume.size
+                    if volume.datacenter is not None:
+                        mount["datacenter"] = volume.datacenter
 
         fn = (
             getattr(handle, "_fn", None)
@@ -162,6 +171,26 @@ def build_manifest(
     return manifest
 
 
+async def _prepare_task_resources(
+    app: App, manifest: Dict[str, Any], client: AppsApiClient, events=None
+) -> None:
+    from .registry import resolve_registry_auth
+    from .volume import VolumeResolver
+
+    resolver = VolumeResolver(client, events=events)
+    specs = [handle.spec for handle in app.resources.values()]
+    for resource in manifest["resources"]:
+        if resource["kind"] != "task":
+            continue
+        spec = app.resources[resource["name"]].spec
+        bindings = await resolver.resolve_mounts(spec.mounts, specs)
+        for mount, binding in zip(resource.get("mounts", []), bindings):
+            mount["id"] = binding["id"]
+        registry_id = await resolve_registry_auth(spec.registry_auth, api=client)
+        if registry_id:
+            resource["registryAuth"] = registry_id
+
+
 def _load_ignores(path: Path) -> GitIgnoreSpec:
     if path.is_symlink() or not path.is_file():
         return GitIgnoreSpec.from_lines([])
@@ -299,21 +328,32 @@ def _deployed_endpoint_input(
     from .secret import render_env
 
     spec.validate()
+    resource_env = render_env(spec.env)
+    reserved = {
+        "RUNPOD_API_KEY",
+        "RUNPOD_AI_API_KEY",
+        "RUNPOD_ENDPOINT_ID",
+        "RUNPOD_POD_ID",
+        "RUNPOD_TASK_TOKEN",
+        "RUNPOD_TASK_PORT",
+        "FLASH_BUILD_ARTIFACT_PATH",
+        "RUNPOD_RUNTIME_KIND",
+        "RUNPOD_DEV_APP",
+        "RUNPOD_DEV_RESOURCE",
+        "RUNPOD_DEV_SESSION",
+    }
+    if reserved.intersection(resource_env):
+        raise InvalidResourceError(
+            "production endpoint env contains reserved runtime fields"
+        )
     template_env = {
+        **resource_env,
+        "FLASH_APP_NAME": app.name,
+        "FLASH_ENVIRONMENT": env_name or app.env,
         "FLASH_RESOURCE_NAME": spec.name,
-        # version-triggering: a new build recreates all workers
         "RUNPOD_BUILD_ID": build_id,
-        **render_env(spec.env),
-        "FLASH_ENVIRONMENT": env_name or app.env,
     }
 
-    # 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
-    api_key = runpod.api_key
-    if api_key and "RUNPOD_API_KEY" not in template_env:
-        template_env["RUNPOD_API_KEY"] = api_key
-
     # nested tasks use the same runtime and sdk channels as the endpoint
     for key in (
         "RUNPOD_RUNTIME_TAG",
@@ -531,6 +571,7 @@ async def deploy_app(
         python_version=python_version,
         excluded_packages=build.excluded,
     )
+    await _prepare_task_resources(app, manifest, client, events)
     _phase(events, "package")
     tar_path = package_project(project_root, manifest, env_dir=build.env_dir)
     tar_size = tar_path.stat().st_size
@@ -562,6 +603,20 @@ async def deploy_app(
         upload["uploadUrl"], str(tar_path), progress=upload_progress
     )
     build = await client.finalize_artifact_upload(app_id, upload["objectKey"], manifest)
+    definitions = build.get("resourceSpec")
+    tasks = {r["name"] for r in manifest["resources"] if r["kind"] == "task"}
+    if (
+        not isinstance(definitions, dict)
+        or not isinstance(definitions.get("tasks"), dict)
+        or set(definitions["tasks"]) != tasks
+        or any(
+            not isinstance(task, dict) or not task
+            for task in definitions["tasks"].values()
+        )
+    ):
+        raise RuntimeError(
+            "backend did not acknowledge the registered task definitions"
+        )
     log.info("uploaded build %s", build["id"])
 
     await client.deploy_build(environment["id"], build["id"])

runpod/apps/handles.py

@@ -233,6 +233,15 @@ def _guard_discovery(self) -> None:
     def _is_current_worker(self) -> bool:
         import os
 
+        from .context import deployed_identity
+
+        identity = deployed_identity()
+        if identity is not None:
+            return (
+                self.spec.kind is not ResourceKind.TASK
+                and identity[0] == self._app.name
+                and identity[2] == self.spec.name
+            )
         if dev_app := os.getenv("RUNPOD_DEV_APP"):
             return (
                 dev_app == self._app.name

runpod/apps/stubs.py

@@ -123,6 +123,33 @@ def __repr__(self) -> str:
         return f"<Queue stub {self.app_name}/{ref}>"
 
 
+class Task(_StubBase):
+    """client for an approved task deployed in a flash environment."""
+
+    def __init__(self, *, app: str, name: str, env: Optional[str] = None):
+        super().__init__(app=app, name=name, env=env)
+        if not app or not name:
+            raise ValueError("task app and name must be non-empty")
+        self.remote = Invoker(self._remote_async)
+        self.spawn = Invoker(self._spawn_async)
+
+    def _target(self):
+        from .targets import DeployedTaskTarget
+
+        return DeployedTaskTarget(self.app_name, self.env, self.resource_name)
+
+    async def _remote_async(self, *args: Any, **kwargs: Any) -> Any:
+        target = self._target()
+        return await target.invoke(target.build_payload(None, None, args, kwargs))
+
+    async def _spawn_async(self, *args: Any, **kwargs: Any) -> Any:
+        target = self._target()
+        return await target.submit(target.build_payload(None, None, args, kwargs))
+
+    def __repr__(self):
+        return f"<Task stub {self.app_name}/{self.resource_name}>"
+
+
 class _StubRouteCaller:
     __slots__ = ("_stub", "_method")