diff --git a/src/runpod_flash/cli/commands/build.py b/src/runpod_flash/cli/commands/build.py index 23e99506..5a29c057 100644 --- a/src/runpod_flash/cli/commands/build.py +++ b/src/runpod_flash/cli/commands/build.py @@ -564,9 +564,9 @@ def build_command( None, "--python-version", help=( - "Target Python version for worker images (3.10, 3.11, or 3.12). " - "Overrides per-resource python_version declarations. " - "Defaults to the version declared on resource configs, or 3.12 if none set." + "Target Python version for worker images (3.10, 3.11, 3.12, or 3.13). " + "Defaults to the version declared on resource configs, or 3.12 if none " + "set. Fails if it conflicts with a declared python_version." ), ), ): diff --git a/src/runpod_flash/cli/commands/deploy.py b/src/runpod_flash/cli/commands/deploy.py index 46a7fb61..870491a1 100644 --- a/src/runpod_flash/cli/commands/deploy.py +++ b/src/runpod_flash/cli/commands/deploy.py @@ -2,13 +2,17 @@ from __future__ import annotations +import ast import asyncio +import json import logging import shutil +from pathlib import Path from typing import Any import typer from rich.console import Console +from rich.markup import escape from runpod_flash.cli.utils.formatting import print_error from runpod_flash.core.exceptions import RunpodAPIKeyError @@ -47,8 +51,9 @@ def deploy_command( None, "--python-version", help=( - "Target Python version for worker images (3.10, 3.11, or 3.12). " - "Overrides per-resource python_version declarations." + "Target Python version for worker images (3.10, 3.11, 3.12, or 3.13). " + "Defaults to the version declared on resource configs, or 3.12 if none " + "set. Fails if it conflicts with a declared python_version." ), ), ): @@ -116,15 +121,102 @@ def _handle_deploy_error(exc: Exception) -> None: raise typer.Exit(1) -def _print_curl_example(url: str, method: str = "POST") -> None: +def _required_params( + file_path: str, func_name: str, class_name: str | None = None +) -> list[str] | None: + """names of a function's parameters that have no default, read from source. + + looks only at top-level definitions (or methods of a top-level class), + which is what the generated handler imports. returns None when the + source can't be read or the definition isn't found. + """ + try: + tree = ast.parse(Path(file_path).read_text()) + except (OSError, SyntaxError, ValueError): + return None + + body: list[ast.stmt] = tree.body + if class_name is not None: + classes = [ + n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == class_name + ] + if not classes: + return None + body = classes[-1].body + + # the last definition wins at import time, as in python itself + matches = [ + node + for node in body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == func_name + ] + if not matches: + return None + + func = matches[-1] + args = func.args + positional = args.posonlyargs + args.args + required = positional[: len(positional) - len(args.defaults)] + required += [ + a for a, default in zip(args.kwonlyargs, args.kw_defaults) if default is None + ] + # a method's receiver (whatever it's named) is bound by the handler, not + # sent as input. static methods have no receiver. + is_static = any( + (isinstance(d, ast.Name) and d.id == "staticmethod") + or (isinstance(d, ast.Attribute) and d.attr == "staticmethod") + for d in func.decorator_list + ) + if class_name is not None and not is_static and required: + if required[0] is positional[0]: + required = required[1:] + return [a.arg for a in required] + + +def _qb_example_input(resource: dict[str, Any]) -> dict[str, Any]: + """example runsync input for a queue-based endpoint. + + runpod drops a job whose input is an empty dict, so the job never + reaches the worker and times out. the deployed handler strips the + "__empty" marker, so use it when the call needs no arguments. + """ + functions = resource.get("functions") or [] + if not functions: + return {"__empty": True} + func = functions[0] + file_path = resource.get("file_path", "") + + if func.get("is_class"): + methods = func.get("class_methods") or [] + if not methods: + return {"method": ""} + method = methods[0] + params = _required_params(file_path, method, class_name=func.get("name")) + body: dict[str, Any] = {name: f"<{name}>" for name in params or []} + # class handlers dispatch on "method" only when there's more than one + if len(methods) > 1: + body = {"method": method, **body} + return body or {"__empty": True} + + params = _required_params(file_path, func.get("name", "")) + if not params: + return {"__empty": True} + return {name: f"<{name}>" for name in params} + + +def _print_curl_example( + url: str, method: str = "POST", body: dict[str, Any] | None = None +) -> None: """Print a curl example for the given URL.""" lines = [f"curl -X {method} {url}"] if method == "POST": lines.append(' -H "Content-Type: application/json"') lines.append(' -H "Authorization: Bearer $RUNPOD_API_KEY"') if method == "POST": - lines.append(""" -d '{"input": {}}'""") - console.print("[dim]" + " \\\n".join(lines) + "[/dim]") + payload = json.dumps(body if body is not None else {"input": {}}) + lines.append(f" -d '{payload}'") + console.print("[dim]" + escape(" \\\n".join(lines)) + "[/dim]") def _display_post_deployment_guidance( @@ -187,8 +279,11 @@ def _display_post_deployment_guidance( # one curl example curl_url = None curl_method = "POST" + curl_body = None if qb_entries: - curl_url = f"{qb_entries[0][1]}/runsync" + qb_name, qb_url = qb_entries[0] + curl_url = f"{qb_url}/runsync" + curl_body = {"input": _qb_example_input(resources.get(qb_name, {}))} elif lb_entries: _name, base_url, lb_routes = lb_entries[0] for k in sorted(lb_routes.keys()): @@ -205,7 +300,7 @@ def _display_post_deployment_guidance( if curl_url: console.print() - _print_curl_example(curl_url, method=curl_method) + _print_curl_example(curl_url, method=curl_method, body=curl_body) def _launch_preview(project_dir): diff --git a/src/runpod_flash/cli/docs/README.md b/src/runpod_flash/cli/docs/README.md index 0feca0e7..39f410cb 100644 --- a/src/runpod_flash/cli/docs/README.md +++ b/src/runpod_flash/cli/docs/README.md @@ -281,9 +281,11 @@ Optional: ```bash FLASH_APP=my-project # defaults to current directory name FLASH_ENV=staging # defaults to "production" -FLASH_SENTINEL_TIMEOUT=120 # sentinel request timeout in seconds (default: 90) +FLASH_SENTINEL_TIMEOUT=120 # timeout for each sentinel HTTP request in seconds (default: 90) ``` +`FLASH_SENTINEL_TIMEOUT` doesn't cap how long a call waits. A job that's still queued or running is polled until it finishes. A job no worker picks up within 15 minutes is cancelled. To bound a running job, set `execution_timeout_ms` on the `Endpoint`. The default of 0 sets no limit. + ## Testing your dev server ```bash diff --git a/src/runpod_flash/core/resources/serverless.py b/src/runpod_flash/core/resources/serverless.py index 67f6eb5a..5f4af5c5 100644 --- a/src/runpod_flash/core/resources/serverless.py +++ b/src/runpod_flash/core/resources/serverless.py @@ -276,7 +276,7 @@ class ServerlessResource(DeployableResource): datacenter: Optional[List[DataCenter] | DataCenter] = Field(default=None) python_version: Optional[str] = Field( default=None, - description="Python version for runtime image selection. Defaults to the local interpreter version at build time.", + description="Python version for runtime image selection. Defaults to 3.12 when no resource declares one.", ) # === Input Fields === diff --git a/src/runpod_flash/flash_sentinel.py b/src/runpod_flash/flash_sentinel.py index e5a83219..8f54e946 100644 --- a/src/runpod_flash/flash_sentinel.py +++ b/src/runpod_flash/flash_sentinel.py @@ -10,27 +10,66 @@ inspect.signature and sends the result as the runsync input body. """ +import asyncio import base64 import inspect import logging +import math import os +import time from typing import Any, Callable, Dict, Optional import cloudpickle +import httpx import runpod from .core.urls import ENDPOINT_DOMAIN from .core.utils import http as _http +from .core.utils.backoff import get_backoff_delay from .protos.remote_execution import FunctionRequest log = logging.getLogger(__name__) FLASH_SENTINEL_ID = "flash" -# default timeout for sentinel requests. configurable via -# FLASH_SENTINEL_TIMEOUT env var (seconds). +# default timeout for each sentinel HTTP request. configurable via +# FLASH_SENTINEL_TIMEOUT env var (seconds). it doesn't cap the total wait: +# a pending job is polled until it finishes. DEFAULT_SENTINEL_TIMEOUT = 90 +# runsync returns one of these when the job outlives the server-side wait +# (cold starts, model downloads). the job keeps running, so poll /status. +_PENDING_STATUSES = frozenset({"IN_QUEUE", "IN_PROGRESS"}) + +# longest gap between /status polls, in seconds. +_STATUS_POLL_MAX_SECONDS = 5.0 + +# backoff exponent at which the poll delay reaches _STATUS_POLL_MAX_SECONDS +# (with a 1s base). capping the exponent there keeps 2**attempt from +# overflowing float after ~1000 polls on long jobs. +_STATUS_POLL_BASE_SECONDS = 1.0 +_STATUS_POLL_MAX_ATTEMPT = math.ceil( + math.log2(_STATUS_POLL_MAX_SECONDS / _STATUS_POLL_BASE_SECONDS) +) + +# give up (and cancel) when no worker picks the job up for this long. +# once the job is IN_PROGRESS, only the endpoint's execution_timeout_ms +# bounds it. flash defaults that to 0 (no limit), so a hung handler is +# polled until the job's 24h TTL expires. +_QUEUE_WAIT_MAX_SECONDS = 900.0 + +# polling raises after this many consecutive network errors or 429/5xx +# responses. +_STATUS_POLL_MAX_TRANSPORT_ERRORS = 5 + +# runsync holds the request for 90s by default, the same as the default +# client timeout, so the client could give up before runsync hands back a +# pending status. ask runsync to return this many seconds before the client +# timeout. ai-api accepts a wait of 1s to 300s. +_RUNSYNC_WAIT_MARGIN_SECONDS = 10.0 +_RUNSYNC_WAIT_MIN_SECONDS = 1.0 +_RUNSYNC_WAIT_MAX_SECONDS = 300.0 + def _get_timeout(override: Optional[float] = None) -> float: """resolve the sentinel request timeout. @@ -49,6 +88,13 @@ def _get_timeout(override: Optional[float] = None) -> float: return DEFAULT_SENTINEL_TIMEOUT +def _runsync_wait_ms(timeout: float) -> int: + """server-side runsync wait in ms, kept below the client timeout.""" + wait = timeout - _RUNSYNC_WAIT_MARGIN_SECONDS + wait = max(_RUNSYNC_WAIT_MIN_SECONDS, min(wait, _RUNSYNC_WAIT_MAX_SECONDS)) + return int(wait * 1000) + + def _flash_headers(app: str, env: str, endpoint: str) -> Dict[str, str]: """build the flash resolution headers.""" return { @@ -80,18 +126,36 @@ async def _sentinel_qb_post( payload: Dict[str, Any], timeout: Optional[float] = None, ) -> Dict[str, Any]: - """post a payload to the sentinel runsync URL and return the raw response dict.""" - url = f"{runpod.endpoint_url_base}/{FLASH_SENTINEL_ID}/runsync" + """submit a payload through the sentinel and return the final response dict. + + uses /runsync with a server-side wait below the client timeout. a + timeout too short to leave room for runsync's minimum wait submits + with /run instead. either way, a pending job is polled to completion. + """ headers = _flash_headers(app, env, endpoint_name) effective_timeout = _get_timeout(timeout) + base = f"{runpod.endpoint_url_base}/{FLASH_SENTINEL_ID}" + if effective_timeout > _RUNSYNC_WAIT_MARGIN_SECONDS + _RUNSYNC_WAIT_MIN_SECONDS: + url = f"{base}/runsync" + params: Dict[str, Any] = {"wait": _runsync_wait_ms(effective_timeout)} + else: + url = f"{base}/run" + params = {} log.debug("sentinel QB -> %s/%s/%s", app, env, endpoint_name) + submitted_at = time.monotonic() + try: async with _http.get_authenticated_httpx_client( timeout=effective_timeout ) as client: - response = await client.post(url, json=payload, headers=headers) + response = await client.post( + url, + json=payload, + headers=headers, + params=params, + ) except Exception as exc: if "timeout" in type(exc).__name__.lower() or "timeout" in str(exc).lower(): raise RuntimeError( @@ -108,7 +172,111 @@ async def _sentinel_qb_post( f"environment '{env}'. deploy it first with 'flash deploy'." ) response.raise_for_status() - return response.json() + data: Dict[str, Any] = response.json() + + if data.get("status") in _PENDING_STATUSES: + if not data.get("id"): + raise RuntimeError( + f"endpoint '{endpoint_name}' returned status " + f"{data['status']} without a job id: {data}" + ) + data = await _poll_until_done( + app, env, endpoint_name, data["id"], effective_timeout, submitted_at + ) + return data + + +async def _poll_until_done( + app: str, + env: str, + endpoint_name: str, + job_id: str, + request_timeout: float, + submitted_at: float, +) -> Dict[str, Any]: + """poll the sentinel /status route until the job leaves IN_QUEUE/IN_PROGRESS. + + a running job is bounded only by the endpoint's execution_timeout_ms, + which ends it server-side with TIMED_OUT. the default of 0 sets no limit. + a job no worker picks up is cancelled after _QUEUE_WAIT_MAX_SECONDS, + counted from submission. + """ + base = f"{runpod.endpoint_url_base}/{FLASH_SENTINEL_ID}" + headers = _flash_headers(app, env, endpoint_name) + + log.debug("sentinel QB job %s still pending, polling /status", job_id) + + queued_since = submitted_at + attempt = 0 + transport_errors = 0 + async with _http.get_authenticated_httpx_client(timeout=request_timeout) as client: + while True: + await asyncio.sleep( + get_backoff_delay( + attempt, + base=_STATUS_POLL_BASE_SECONDS, + max_seconds=_STATUS_POLL_MAX_SECONDS, + ) + ) + attempt = min(attempt + 1, _STATUS_POLL_MAX_ATTEMPT) + try: + response = await client.get(f"{base}/status/{job_id}", headers=headers) + if response.status_code == 429 or response.status_code >= 500: + response.raise_for_status() + except (httpx.TransportError, httpx.HTTPStatusError) as exc: + transport_errors += 1 + if transport_errors >= _STATUS_POLL_MAX_TRANSPORT_ERRORS: + raise RuntimeError( + f"lost contact with endpoint '{endpoint_name}' while " + f"waiting for job {job_id}: {exc}" + ) from exc + log.debug("polling job %s failed (%s), retrying", job_id, exc) + continue + transport_errors = 0 + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + raise RuntimeError( + f"checking job {job_id} on endpoint '{endpoint_name}' failed: {exc}" + ) from exc + try: + data: Dict[str, Any] = response.json() + except ValueError as exc: + raise RuntimeError( + f"checking job {job_id} on endpoint '{endpoint_name}' " + f"returned a non-JSON response: {response.text[:200]!r}" + ) from exc + status = data.get("status") + if status not in _PENDING_STATUSES: + return data + if status != "IN_QUEUE": + queued_since = time.monotonic() + elif time.monotonic() - queued_since > _QUEUE_WAIT_MAX_SECONDS: + cancelled = await _cancel_job(client, base, headers, job_id) + outcome = ( + "it was cancelled" + if cancelled + else "cancelling it failed, so it may still run later" + ) + raise RuntimeError( + f"job {job_id} on endpoint '{endpoint_name}' was still " + f"queued after {_QUEUE_WAIT_MAX_SECONDS:.0f}s and {outcome}. " + f"no worker picked it up; check that GPUs of the requested " + f"type are available." + ) + + +async def _cancel_job( + client: httpx.AsyncClient, base: str, headers: Dict[str, str], job_id: str +) -> bool: + """best-effort cancel so an abandoned job doesn't run later.""" + try: + response = await client.post(f"{base}/cancel/{job_id}", headers=headers) + response.raise_for_status() + except Exception as exc: + log.warning("could not cancel job %s: %s", job_id, exc) + return False + return True def _handle_sentinel_response(data: Dict[str, Any]) -> Any: @@ -122,6 +290,11 @@ def _handle_sentinel_response(data: Dict[str, Any]) -> Any: err = data.get("error") or data.get("output", {}).get("error", "unknown") raise RuntimeError(f"remote execution failed: {err}") + if data.get("status") in ("CANCELLED", "TIMED_OUT"): + raise RuntimeError( + f"remote execution {data['status'].lower()} (job {data.get('id', 'unknown')})" + ) + if "output" not in data and "status" not in data: raise RuntimeError( f"unexpected response from sentinel (no 'output' or 'status' key): {data}" diff --git a/tests/unit/cli/test_deploy.py b/tests/unit/cli/test_deploy.py index ae046b1a..ecbd37e1 100644 --- a/tests/unit/cli/test_deploy.py +++ b/tests/unit/cli/test_deploy.py @@ -801,6 +801,238 @@ def test_qb_curl_uses_runsync(self, patched_console): assert "curl" in output assert "curl -X POST https://def.api.runpod.ai/runsync" in output + def test_qb_curl_sends_required_params(self, patched_console, tmp_path): + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + src = tmp_path / "gpu_demo.py" + src.write_text( + "async def gpu_matrix_multiply(size, *, scale, seed=0, verbose=False):\n" + " return size\n" + ) + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={"my_qb": "https://api.runpod.ai/v2/abc123"}, + resources={ + "my_qb": { + "file_path": str(src), + "functions": [{"name": "gpu_matrix_multiply", "is_class": False}], + } + }, + routes={}, + ) + output = self._collect_output(patched_console) + assert '{"input": {"size": "", "scale": ""}}' in output + assert '{"input": {}}' not in output + + def test_qb_curl_marks_empty_input_for_no_arg_function( + self, patched_console, tmp_path + ): + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + src = tmp_path / "jobs.py" + src.write_text("def ping(verbose=False):\n return 'pong'\n") + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={"my_qb": "https://api.runpod.ai/v2/abc123"}, + resources={ + "my_qb": { + "file_path": str(src), + "functions": [{"name": "ping", "is_class": False}], + } + }, + routes={}, + ) + output = self._collect_output(patched_console) + assert '{"input": {"__empty": true}}' in output + + def test_qb_curl_never_sends_empty_input(self, patched_console): + """runpod drops jobs with an empty input dict, so the example must not.""" + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={"my_qb": "https://api.runpod.ai/v2/abc123"}, + resources={ + "my_qb": { + "file_path": "missing.py", + "functions": [{"name": "f", "is_class": False}], + } + }, + routes={}, + ) + output = self._collect_output(patched_console) + assert '{"input": {"__empty": true}}' in output + assert '{"input": {}}' not in output + + def _class_resource(self, src, methods): + return { + "my_qb": { + "file_path": str(src), + "functions": [ + {"name": "Model", "is_class": True, "class_methods": methods} + ], + } + } + + def test_qb_curl_for_single_method_class_omits_method( + self, patched_console, tmp_path + ): + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + src = tmp_path / "model.py" + src.write_text( + "class Model:\n def predict(self, image, k=1):\n ...\n" + ) + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={"my_qb": "https://api.runpod.ai/v2/abc123"}, + resources=self._class_resource(src, ["predict"]), + routes={}, + ) + output = self._collect_output(patched_console) + assert '{"input": {"image": ""}}' in output + + def test_qb_curl_for_multi_method_class_names_real_method( + self, patched_console, tmp_path + ): + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + src = tmp_path / "model.py" + src.write_text( + "class Model:\n" + " def predict(self, image):\n ...\n" + " def embed(self, text):\n ...\n" + ) + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={"my_qb": "https://api.runpod.ai/v2/abc123"}, + resources=self._class_resource(src, ["predict", "embed"]), + routes={}, + ) + output = self._collect_output(patched_console) + assert '{"input": {"method": "predict", "image": ""}}' in output + + @pytest.mark.parametrize( + ("source", "expected"), + [ + ( + "class Model:\n def predict(this, image):\n ...\n", + '{"input": {"image": ""}}', + ), + ( + "class Model:\n @staticmethod\n" + " def predict(image):\n ...\n", + '{"input": {"image": ""}}', + ), + ( + "import builtins\n\nclass Model:\n @builtins.staticmethod\n" + " def predict(image):\n ...\n", + '{"input": {"image": ""}}', + ), + ], + ) + def test_qb_curl_drops_only_the_method_receiver( + self, patched_console, tmp_path, source, expected + ): + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + src = tmp_path / "model.py" + src.write_text(source) + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={"my_qb": "https://api.runpod.ai/v2/abc123"}, + resources=self._class_resource(src, ["predict"]), + routes={}, + ) + assert expected in self._collect_output(patched_console) + + def test_qb_curl_uses_last_definition_of_redefined_class( + self, patched_console, tmp_path + ): + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + src = tmp_path / "model.py" + src.write_text( + "class Model:\n def predict(self, text):\n ...\n\n" + "class Model:\n def predict(self, image):\n ...\n" + ) + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={"my_qb": "https://api.runpod.ai/v2/abc123"}, + resources=self._class_resource(src, ["predict"]), + routes={}, + ) + output = self._collect_output(patched_console) + assert '{"input": {"image": ""}}' in output + + def test_qb_curl_reads_top_level_function_not_nested_namesake( + self, patched_console, tmp_path + ): + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + src = tmp_path / "jobs.py" + src.write_text( + "class Helper:\n def predict(self, image):\n ...\n\n" + "def predict(prompt):\n ...\n" + ) + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={"my_qb": "https://api.runpod.ai/v2/abc123"}, + resources={ + "my_qb": { + "file_path": str(src), + "functions": [{"name": "predict", "is_class": False}], + } + }, + routes={}, + ) + output = self._collect_output(patched_console) + assert '{"input": {"prompt": ""}}' in output + + def test_qb_curl_keeps_self_param_on_plain_function( + self, patched_console, tmp_path + ): + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + src = tmp_path / "jobs.py" + src.write_text("def run(self):\n ...\n") + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={"my_qb": "https://api.runpod.ai/v2/abc123"}, + resources={ + "my_qb": { + "file_path": str(src), + "functions": [{"name": "run", "is_class": False}], + } + }, + routes={}, + ) + output = self._collect_output(patched_console) + assert '{"input": {"self": ""}}' in output + + def test_qb_curl_uses_last_definition_of_redefined_function( + self, patched_console, tmp_path + ): + from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance + + src = tmp_path / "jobs.py" + src.write_text( + "def predict(image):\n ...\n\ndef predict(image, scale):\n ...\n" + ) + _display_post_deployment_guidance( + env_name="production", + resources_endpoints={"my_qb": "https://api.runpod.ai/v2/abc123"}, + resources={ + "my_qb": { + "file_path": str(src), + "functions": [{"name": "predict", "is_class": False}], + } + }, + routes={}, + ) + output = self._collect_output(patched_console) + assert '{"input": {"image": "", "scale": ""}}' in output + def test_lb_endpoint_shows_routes(self, patched_console): from runpod_flash.cli.commands.deploy import _display_post_deployment_guidance diff --git a/tests/unit/test_flash_sentinel.py b/tests/unit/test_flash_sentinel.py index cfaf4051..a9520d16 100644 --- a/tests/unit/test_flash_sentinel.py +++ b/tests/unit/test_flash_sentinel.py @@ -154,6 +154,490 @@ async def my_func(): await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) +class TestHandleSentinelTerminalStatuses: + @pytest.mark.parametrize("status", ["CANCELLED", "TIMED_OUT"]) + def test_cancelled_or_timed_out_raises(self, status): + from runpod_flash.flash_sentinel import _handle_sentinel_response + + data = {"id": "job-1", "status": status} + with pytest.raises(RuntimeError, match=status.lower()): + _handle_sentinel_response(data) + + +def _json_response(payload): + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + response.json.return_value = payload + return response + + +def _http_error_response(status_code): + import httpx + + response = MagicMock() + response.status_code = status_code + response.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError( + str(status_code), request=MagicMock(), response=MagicMock() + ) + ) + return response + + +class TestSentinelQBPolling: + """runsync returns IN_QUEUE/IN_PROGRESS when the job outlives the server wait.""" + + @pytest.fixture(autouse=True) + def no_sleep(self): + with patch("runpod_flash.flash_sentinel.asyncio.sleep", new=AsyncMock()): + yield + + @pytest.mark.asyncio + async def test_polls_status_until_completed(self, mock_httpx): + from runpod_flash.flash_sentinel import FLASH_SENTINEL_ID, sentinel_qb_execute + + mock_client = _make_mock_client( + _json_response({"id": "job-1", "status": "IN_QUEUE"}) + ) + mock_client.get = AsyncMock( + side_effect=[ + _json_response({"id": "job-1", "status": "IN_QUEUE"}), + _json_response({"id": "job-1", "status": "IN_PROGRESS"}), + _json_response( + {"id": "job-1", "status": "COMPLETED", "output": {"image": "b64"}} + ), + ] + ) + mock_httpx.return_value = mock_client + + async def my_func(prompt): + pass + + result = await sentinel_qb_execute( + "myapp", "prod", "gpu-worker", my_func, "a cat" + ) + + assert result == {"image": "b64"} + assert mock_client.get.call_count == 3 + status_call = mock_client.get.call_args + assert status_call.args[0].endswith(f"/{FLASH_SENTINEL_ID}/status/job-1") + assert status_call.kwargs["headers"]["X-Flash-Endpoint"] == "gpu-worker" + + @pytest.mark.asyncio + async def test_polled_failure_raises(self, mock_httpx): + from runpod_flash.flash_sentinel import sentinel_qb_execute + + mock_client = _make_mock_client( + _json_response({"id": "job-1", "status": "IN_PROGRESS"}) + ) + mock_client.get = AsyncMock( + return_value=_json_response( + {"id": "job-1", "status": "FAILED", "error": "CUDA out of memory"} + ) + ) + mock_httpx.return_value = mock_client + + async def my_func(): + pass + + with pytest.raises(RuntimeError, match="CUDA out of memory"): + await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) + + @pytest.mark.asyncio + async def test_polled_timeout_raises(self, mock_httpx): + from runpod_flash.flash_sentinel import sentinel_qb_execute + + mock_client = _make_mock_client( + _json_response({"id": "job-1", "status": "IN_QUEUE"}) + ) + mock_client.get = AsyncMock( + return_value=_json_response({"id": "job-1", "status": "TIMED_OUT"}) + ) + mock_httpx.return_value = mock_client + + async def my_func(): + pass + + with pytest.raises(RuntimeError, match="timed_out"): + await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) + + @pytest.mark.asyncio + async def test_class_execute_polls_too(self, mock_httpx): + from runpod_flash.flash_sentinel import sentinel_qb_class_execute + from runpod_flash.protos.remote_execution import FunctionRequest + + mock_client = _make_mock_client( + _json_response({"id": "job-2", "status": "IN_PROGRESS"}) + ) + mock_client.get = AsyncMock( + return_value=_json_response( + {"id": "job-2", "status": "COMPLETED", "output": {"ok": True}} + ) + ) + mock_httpx.return_value = mock_client + + request = FunctionRequest( + execution_type="class", class_name="MyModel", method_name="predict" + ) + result = await sentinel_qb_class_execute("myapp", "prod", "gpu-worker", request) + + assert result == {"ok": True} + mock_client.get.assert_awaited_once() + + @pytest.mark.asyncio + async def test_pending_without_id_raises(self, mock_httpx): + from runpod_flash.flash_sentinel import _sentinel_qb_post + + mock_client = _make_mock_client(_json_response({"status": "IN_QUEUE"})) + mock_client.get = AsyncMock() + mock_httpx.return_value = mock_client + + with pytest.raises(RuntimeError, match="without a job id"): + await _sentinel_qb_post("myapp", "prod", "gpu-worker", {"input": {}}) + mock_client.get.assert_not_awaited() + + @pytest.mark.asyncio + async def test_transient_poll_error_is_retried(self, mock_httpx): + import httpx + + from runpod_flash.flash_sentinel import sentinel_qb_execute + + mock_client = _make_mock_client( + _json_response({"id": "job-1", "status": "IN_PROGRESS"}) + ) + mock_client.get = AsyncMock( + side_effect=[ + httpx.ReadTimeout("slow"), + _json_response( + {"id": "job-1", "status": "COMPLETED", "output": {"ok": True}} + ), + ] + ) + mock_httpx.return_value = mock_client + + async def my_func(): + pass + + result = await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) + + assert result == {"ok": True} + assert mock_client.get.call_count == 2 + + @pytest.mark.asyncio + async def test_repeated_poll_errors_raise(self, mock_httpx): + import httpx + + from runpod_flash.flash_sentinel import ( + _STATUS_POLL_MAX_TRANSPORT_ERRORS, + sentinel_qb_execute, + ) + + mock_client = _make_mock_client( + _json_response({"id": "job-1", "status": "IN_PROGRESS"}) + ) + mock_client.get = AsyncMock(side_effect=httpx.ConnectError("down")) + mock_httpx.return_value = mock_client + + async def my_func(): + pass + + with pytest.raises(RuntimeError, match="lost contact.*job-1"): + await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) + assert mock_client.get.call_count == _STATUS_POLL_MAX_TRANSPORT_ERRORS + + @pytest.mark.asyncio + async def test_poll_client_error_propagates(self, mock_httpx): + import httpx + + from runpod_flash.flash_sentinel import sentinel_qb_execute + + error_response = _http_error_response(404) + mock_client = _make_mock_client( + _json_response({"id": "job-1", "status": "IN_PROGRESS"}) + ) + mock_client.get = AsyncMock(return_value=error_response) + mock_httpx.return_value = mock_client + + async def my_func(): + pass + + with pytest.raises(RuntimeError, match="job job-1.*gpu-worker") as exc_info: + await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) + assert isinstance(exc_info.value.__cause__, httpx.HTTPStatusError) + mock_client.get.assert_awaited_once() + + @pytest.mark.asyncio + async def test_poll_non_json_response_raises(self, mock_httpx): + import json + + from runpod_flash.flash_sentinel import sentinel_qb_execute + + html_response = _json_response(None) + html_response.json.side_effect = json.JSONDecodeError("Expecting value", "", 0) + html_response.text = "bad gateway" + mock_client = _make_mock_client( + _json_response({"id": "job-1", "status": "IN_PROGRESS"}) + ) + mock_client.get = AsyncMock(return_value=html_response) + mock_httpx.return_value = mock_client + + async def my_func(): + pass + + with pytest.raises(RuntimeError, match="job job-1.*gpu-worker.*non-JSON"): + await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) + + def test_poll_attempt_cap_reaches_max_delay(self): + from runpod_flash.flash_sentinel import ( + _STATUS_POLL_BASE_SECONDS, + _STATUS_POLL_MAX_ATTEMPT, + _STATUS_POLL_MAX_SECONDS, + ) + + capped = _STATUS_POLL_BASE_SECONDS * 2**_STATUS_POLL_MAX_ATTEMPT + below = _STATUS_POLL_BASE_SECONDS * 2 ** (_STATUS_POLL_MAX_ATTEMPT - 1) + assert capped >= _STATUS_POLL_MAX_SECONDS > below + + @pytest.mark.asyncio + @pytest.mark.parametrize("status_code", [429, 503]) + async def test_poll_recovers_from_transient_status(self, mock_httpx, status_code): + from runpod_flash.flash_sentinel import sentinel_qb_execute + + mock_client = _make_mock_client( + _json_response({"id": "job-1", "status": "IN_PROGRESS"}) + ) + mock_client.get = AsyncMock( + side_effect=[ + _http_error_response(status_code), + _json_response( + {"id": "job-1", "status": "COMPLETED", "output": {"ok": True}} + ), + ] + ) + mock_httpx.return_value = mock_client + + async def my_func(): + pass + + result = await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) + + assert result == {"ok": True} + assert mock_client.get.call_count == 2 + + @pytest.mark.asyncio + async def test_job_stuck_in_queue_is_cancelled(self, mock_httpx): + from runpod_flash.flash_sentinel import ( + _QUEUE_WAIT_MAX_SECONDS, + FLASH_SENTINEL_ID, + sentinel_qb_execute, + ) + + queued = _json_response({"id": "job-1", "status": "IN_QUEUE"}) + mock_client = _make_mock_client(queued) + mock_client.get = AsyncMock(return_value=queued) + mock_httpx.return_value = mock_client + + clock = iter([0.0, 10.0, _QUEUE_WAIT_MAX_SECONDS + 1]) + + async def my_func(): + pass + + with patch( + "runpod_flash.flash_sentinel.time.monotonic", + side_effect=lambda: next(clock), + ): + with pytest.raises(RuntimeError, match="still queued.*cancelled"): + await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) + + cancel_url = mock_client.post.call_args.args[0] + assert cancel_url.endswith(f"/{FLASH_SENTINEL_ID}/cancel/job-1") + + @pytest.mark.asyncio + async def test_failed_cancel_is_reported(self, mock_httpx): + from runpod_flash.flash_sentinel import ( + _QUEUE_WAIT_MAX_SECONDS, + sentinel_qb_execute, + ) + + queued = _json_response({"id": "job-1", "status": "IN_QUEUE"}) + mock_client = _make_mock_client(queued) + mock_client.post = AsyncMock(side_effect=[queued, _http_error_response(500)]) + mock_client.get = AsyncMock(return_value=queued) + mock_httpx.return_value = mock_client + + clock = iter([0.0, _QUEUE_WAIT_MAX_SECONDS + 1]) + + async def my_func(): + pass + + with patch( + "runpod_flash.flash_sentinel.time.monotonic", + side_effect=lambda: next(clock), + ): + with pytest.raises(RuntimeError, match="cancelling it failed"): + await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) + + @pytest.mark.asyncio + async def test_queue_deadline_counts_time_spent_in_runsync(self, mock_httpx): + from runpod_flash.flash_sentinel import ( + _QUEUE_WAIT_MAX_SECONDS, + sentinel_qb_execute, + ) + + queued = _json_response({"id": "job-1", "status": "IN_QUEUE"}) + mock_client = _make_mock_client(queued) + mock_client.get = AsyncMock(return_value=queued) + mock_httpx.return_value = mock_client + + # submitted at 0, runsync returns pending, first poll is already late + clock = iter([0.0, _QUEUE_WAIT_MAX_SECONDS + 1]) + + async def my_func(): + pass + + with patch( + "runpod_flash.flash_sentinel.time.monotonic", + side_effect=lambda: next(clock), + ): + with pytest.raises(RuntimeError, match="still queued"): + await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) + mock_client.get.assert_awaited_once() + + @pytest.mark.asyncio + async def test_long_running_job_polls_past_backoff_overflow(self, mock_httpx): + """an uncapped exponent overflows float after ~1024 polls.""" + from runpod_flash.flash_sentinel import sentinel_qb_execute + + running = _json_response({"id": "job-1", "status": "IN_PROGRESS"}) + done = _json_response( + {"id": "job-1", "status": "COMPLETED", "output": {"ok": True}} + ) + mock_client = _make_mock_client(running) + mock_client.get = AsyncMock(side_effect=[running] * 1100 + [done]) + mock_httpx.return_value = mock_client + + async def my_func(): + pass + + result = await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) + + assert result == {"ok": True} + assert mock_client.get.call_count == 1101 + + @pytest.mark.asyncio + async def test_running_job_is_not_cancelled_by_queue_deadline(self, mock_httpx): + from runpod_flash.flash_sentinel import ( + _QUEUE_WAIT_MAX_SECONDS, + sentinel_qb_execute, + ) + + mock_client = _make_mock_client( + _json_response({"id": "job-1", "status": "IN_QUEUE"}) + ) + mock_client.get = AsyncMock( + side_effect=[ + _json_response({"id": "job-1", "status": "IN_PROGRESS"}), + _json_response({"id": "job-1", "status": "IN_PROGRESS"}), + _json_response( + {"id": "job-1", "status": "COMPLETED", "output": {"ok": 1}} + ), + ] + ) + mock_httpx.return_value = mock_client + + far = _QUEUE_WAIT_MAX_SECONDS * 10 + clock = iter([0.0, far, far * 2, far * 3]) + + async def my_func(): + pass + + with patch( + "runpod_flash.flash_sentinel.time.monotonic", + side_effect=lambda: next(clock), + ): + result = await sentinel_qb_execute("myapp", "prod", "gpu-worker", my_func) + + assert result == {"ok": 1} + assert mock_client.post.await_count == 1 # the runsync call, no cancel + + +class TestRunsyncWait: + """runsync must return a pending status before the client times out.""" + + @pytest.mark.parametrize( + ("timeout", "expected_ms"), + [(90, 80_000), (30, 20_000), (5, 1_000), (900, 300_000)], + ) + def test_wait_stays_below_client_timeout(self, timeout, expected_ms): + from runpod_flash.flash_sentinel import _runsync_wait_ms + + assert _runsync_wait_ms(timeout) == expected_ms + + @pytest.mark.asyncio + async def test_runsync_sends_wait_param(self, mock_httpx, monkeypatch): + from runpod_flash.flash_sentinel import _sentinel_qb_post + + monkeypatch.delenv("FLASH_SENTINEL_TIMEOUT", raising=False) + mock_client = _make_mock_client( + _json_response({"status": "COMPLETED", "output": {}}) + ) + mock_httpx.return_value = mock_client + + await _sentinel_qb_post("myapp", "prod", "gpu-worker", {"input": {"a": 1}}) + + assert mock_client.post.call_args.kwargs["params"] == {"wait": 80_000} + assert mock_httpx.call_args.kwargs["timeout"] == 90 + + +class TestShortTimeoutUsesRun: + @pytest.fixture(autouse=True) + def no_sleep(self): + with patch("runpod_flash.flash_sentinel.asyncio.sleep", new=AsyncMock()): + yield + + @pytest.mark.asyncio + @pytest.mark.parametrize("timeout", [1, 11]) + async def test_short_timeout_submits_with_run(self, mock_httpx, timeout): + from runpod_flash.flash_sentinel import FLASH_SENTINEL_ID, _sentinel_qb_post + + mock_client = _make_mock_client( + _json_response({"id": "job-1", "status": "IN_QUEUE"}) + ) + mock_client.get = AsyncMock( + return_value=_json_response( + {"id": "job-1", "status": "COMPLETED", "output": {"ok": True}} + ) + ) + mock_httpx.return_value = mock_client + + data = await _sentinel_qb_post( + "myapp", "prod", "gpu-worker", {"input": {"a": 1}}, timeout=timeout + ) + + assert data["output"] == {"ok": True} + call = mock_client.post.call_args + assert call.args[0].endswith(f"/{FLASH_SENTINEL_ID}/run") + assert call.kwargs["params"] == {} + + @pytest.mark.asyncio + async def test_timeout_above_margin_uses_runsync(self, mock_httpx): + from runpod_flash.flash_sentinel import FLASH_SENTINEL_ID, _sentinel_qb_post + + mock_client = _make_mock_client( + _json_response({"status": "COMPLETED", "output": {}}) + ) + mock_httpx.return_value = mock_client + + await _sentinel_qb_post( + "myapp", "prod", "gpu-worker", {"input": {"a": 1}}, timeout=12 + ) + + call = mock_client.post.call_args + assert call.args[0].endswith(f"/{FLASH_SENTINEL_ID}/runsync") + assert call.kwargs["params"] == {"wait": 2_000} + + class TestSentinelQBClassExecute: @pytest.mark.asyncio async def test_dispatches_on_method_name(self, mock_httpx):