"""Run command for the gyoza CLI."""
from __future__ import annotations
import ast
import importlib.util
import json
import os
import sys
import time
from collections.abc import Iterator
from pathlib import Path
from types import ModuleType
from typing import TYPE_CHECKING, Any
from urllib.parse import urlsplit, urlunsplit
import typer
from gyoza.cli.errors import fail
if TYPE_CHECKING:
import httpx
_DEFAULT_INPUT_PATH = "/data/input.json"
_DEFAULT_OUTPUT_PATH = "/data/output.json"
_SEPARATOR = "━" * 50
# Seconds to wait before each retry of an upload. The URL names a fixed key,
# so sending the bytes again only overwrites the same object.
_UPLOAD_RETRY_DELAYS = (2, 5, 15, 30)
_RETRYABLE_STATUS = frozenset({408, 429, 500, 502, 503, 504})
# One PUT stops at 5 GiB on S3 and R2; a larger file needs part URLs.
_SINGLE_PUT_MAX = 5 * 1024**3
_READ_CHUNK = 8 * 1024 * 1024
def _load_module(path: Path) -> ModuleType:
"""Load and return a Python module from a file path.
The module goes into ``sys.modules`` before it executes. Ops declare their
contracts under ``from __future__ import annotations``, which leaves every
annotation as a string, and pydantic resolves those strings by looking the
defining module up in ``sys.modules``. Absent from there, a field like
``dict[str, Any] | None`` cannot be resolved — ``Any`` is not a builtin —
and the op dies at validation time with "Input is not fully defined",
naming a type the file imports at the top.
Registering before ``exec_module`` is also what the import system does, so
a module that refers to itself behaves the same way here. The entry is
removed again if execution fails, leaving no half-built module behind.
"""
name = f"_gyoza_cli_{path.stem}_{abs(hash(path))}"
spec = importlib.util.spec_from_file_location(name, path)
if spec is None or spec.loader is None:
fail(f"unable to import module at {path}")
module = importlib.util.module_from_spec(spec)
sys.modules[name] = module
try:
spec.loader.exec_module(module)
except BaseException:
sys.modules.pop(name, None)
raise
return module
def _find_gyoza_op(module: ModuleType, func_name: str | None) -> Any:
"""
Return the gyoza op function from *module*.
Parameters
----------
module : ModuleType
Loaded Python module.
func_name : str or None
Explicit function name to look up. When ``None`` the first callable
with a ``.run`` attribute defined in the module is used.
Returns
-------
Any
The decorated function.
Raises
------
typer.Exit
On any lookup failure.
"""
if func_name:
func = getattr(module, func_name, None)
if func is None:
fail(f"function '{func_name}' not found in {module.__file__}")
return func
ops = [
obj
for obj in module.__dict__.values()
if callable(obj)
and hasattr(obj, "run")
and getattr(obj, "__module__", None) == module.__name__
]
if not ops:
fail(f"no @gyoza_op function found in {module.__file__}")
return ops[0]
def _parse_inline_input(raw: str) -> dict[str, Any]:
"""
Parse an inline string into a dict.
Parameters
----------
raw : str
JSON or Python-literal dict string.
Returns
-------
dict[str, Any]
Parsed input data.
Raises
------
typer.Exit
If the string cannot be parsed or is not a dict.
"""
try:
parsed = json.loads(raw)
except json.JSONDecodeError:
try:
parsed = ast.literal_eval(raw)
except (ValueError, SyntaxError) as exc:
fail(f"invalid --inline-input: {exc}")
if not isinstance(parsed, dict):
fail("--inline-input must be a JSON object (dictionary)")
return parsed
def _read_json_dict(path: str) -> dict[str, Any]:
"""
Read a JSON object from a file.
Parameters
----------
path : str
Path to the JSON file.
Returns
-------
dict[str, Any]
Parsed data.
Raises
------
OSError
If the file cannot be read.
json.JSONDecodeError
If the file is not valid JSON.
ValueError
If the parsed content is not a JSON object.
"""
with Path(path).open() as fh:
data = json.load(fh)
if not isinstance(data, dict):
msg = "content is not a JSON object (dictionary)"
raise ValueError(msg)
return data
def _read_input_file(path: str) -> dict[str, Any]:
"""
Read a JSON input file into a dict.
Parameters
----------
path : str
Path to the JSON file.
Returns
-------
dict[str, Any]
Parsed input data.
Raises
------
typer.Exit
If the file cannot be read or parsed.
"""
try:
return _read_json_dict(path)
except (OSError, json.JSONDecodeError) as exc:
fail(f"failed to read input file '{path}': {exc}")
except ValueError:
fail(f"input file '{path}' must contain a JSON object (dictionary)")
def _read_output_template(path: str) -> dict[str, Any]:
"""
Read the pre-filled outputs template from the output path, if present.
The worker pre-writes the output file with the caller's partial outputs
(fields whose value is an upload URL) so this runtime can deliver the
produced files there.
Parameters
----------
path : str
Output file path.
Returns
-------
dict[str, Any]
The template, or an empty dict when the file is absent or invalid.
"""
try:
return _read_json_dict(path)
except (OSError, json.JSONDecodeError, ValueError):
return {}
def _is_upload_url(value: Any) -> bool:
"""Return True when *value* is an http(s) URL string."""
return isinstance(value, str) and value.startswith(("http://", "https://"))
def _canonical_url(url: str) -> str:
"""Strip the query string (signature/expiry) from a presigned URL."""
parts = urlsplit(url)
return urlunsplit((parts.scheme, parts.netloc, parts.path, "", ""))
def _is_parts_target(value: Any) -> bool:
"""Return True when *value* is ``{"url", "parts", "part_size"}``."""
return (
isinstance(value, dict)
and _is_upload_url(value.get("url"))
and isinstance(value.get("parts"), list)
and isinstance(value.get("part_size"), int)
)
def _has_upload_url(template: Any) -> bool:
"""Return True when *template* holds an upload URL at any depth."""
if _is_upload_url(template) or _is_parts_target(template):
return True
if isinstance(template, dict):
return any(_has_upload_url(v) for v in template.values())
return False
def _deliver_outputs(
outputs: dict[str, Any], template: dict[str, Any]
) -> dict[str, Any]:
"""
Upload output files to the destinations pre-filled in the template.
For each template field holding an upload URL, the op must have
produced a local file path in that output field. The file is uploaded
via HTTP PUT and the field's final value becomes the canonical URL
(without the presigned query string). Fields without a template URL
pass through unchanged.
A template field may also be a dict of upload URLs, for an op whose
files are keyed by the caller (one per requested item). The op then
returns a dict under the same field; every key it fills must have an
upload URL in the template, and each is delivered like a top-level
field.
A field the op left as ``None`` is an **optional output that was not
produced**, not an error: the caller cannot know at submission time
which optional layers a given run will emit, so it pre-fills the whole
template. Such a field is skipped and stays ``None`` in the result.
Producing something that is not a local file is still a failure.
Parameters
----------
outputs : dict[str, Any]
Outputs produced by the op function.
template : dict[str, Any]
Partial outputs pre-filled by the caller.
Returns
-------
dict[str, Any]
Final outputs with delivered fields replaced by canonical URLs.
Raises
------
typer.Exit
If a field with an upload URL has no local file, a dict field does
not match its template, or an upload fails.
"""
if not _has_upload_url(template):
return outputs
import httpx # noqa: PLC0415 — deferred so CLI startup doesn't pay for it
# No timeout: uploads can be arbitrarily large.
with httpx.Client(timeout=None) as client:
return _deliver_level(client, outputs, template, "")
def _deliver_level(
client: httpx.Client, outputs: dict[str, Any], template: dict[str, Any], prefix: str
) -> dict[str, Any]:
"""Deliver one level of *template*; nested dicts recurse with a dotted *prefix*."""
import httpx # noqa: PLC0415
delivered = dict(outputs)
for key, target in template.items():
name = f"{prefix}{key}"
produced = outputs.get(key)
if (
isinstance(target, dict)
and not _is_parts_target(target)
and _has_upload_url(target)
):
if produced is None:
typer.echo(f" ⏭️ Output '{name}' not produced (optional), skipping.")
continue
if not isinstance(produced, dict):
fail(
f"output '{name}' has a dict of upload URLs but the op "
f"produced {produced!r}"
)
unsigned = [
k
for k in produced
if produced[k] is not None and not _is_upload_url(target.get(k))
]
if unsigned:
fail(
f"output '{name}' produced keys without an upload URL: "
+ ", ".join(map(str, unsigned))
)
delivered[key] = _deliver_level(client, produced, target, f"{name}.")
continue
if not (_is_upload_url(target) or _is_parts_target(target)):
continue
if produced is None:
# Optional output the op chose not to produce. The field stays
# None, which is what the op declared; failing here would kill a
# run that succeeded, after uploading the earlier keys.
typer.echo(f" ⏭️ Output '{name}' not produced (optional), skipping.")
continue
if not isinstance(produced, str) or not Path(produced).is_file():
fail(
f"output '{name}' has an upload URL but the op did not "
f"produce a local file (got: {produced!r})"
)
typer.echo(f" 📦 Uploading output '{name}' ({produced})...")
if _is_parts_target(target):
_upload_in_parts(client, target, Path(produced), name)
target = target["url"]
else:
_upload(client, target, Path(produced), name)
delivered[key] = _canonical_url(target)
return delivered
def _upload_in_parts(
client: httpx.Client, target: dict[str, Any], path: Path, name: str
) -> None:
"""PUT *path* to ``url`` when one PUT takes it, else in ``part_size`` slices."""
size = path.stat().st_size
if size <= _SINGLE_PUT_MAX:
_upload(client, target["url"], path, name)
return
part_size, parts = target["part_size"], target["parts"]
count = -(-size // part_size)
if count > len(parts):
fail(
f"output '{name}' needs {count} parts of {part_size} bytes "
f"but the run has {len(parts)} part URLs"
)
for i in range(count):
offset = i * part_size
typer.echo(f" 📦 Part {i + 1}/{count} of '{name}'...")
_upload(client, parts[i], path, name, offset, min(part_size, size - offset))
def _upload(
client: httpx.Client,
target: str,
path: Path,
name: str,
offset: int = 0,
length: int | None = None,
) -> None:
"""PUT *path*, or *length* bytes of it from *offset*, retrying a storage outage."""
import httpx # noqa: PLC0415
attempts = len(_UPLOAD_RETRY_DELAYS) + 1
for attempt in range(attempts):
try:
with path.open("rb") as fh:
if length is None:
client.put(target, content=fh).raise_for_status()
else:
fh.seek(offset)
client.put(
target,
content=_read_slice(fh, length),
headers={"Content-Length": str(length)},
).raise_for_status()
return
except httpx.HTTPError as exc:
if attempt == attempts - 1 or not _retryable(exc):
fail(f"failed to upload output '{name}': {exc}")
delay = _UPLOAD_RETRY_DELAYS[attempt]
typer.echo(
f" ↻ Upload of '{name}' failed ({exc}), retrying in {delay}s..."
)
time.sleep(delay)
def _read_slice(fh: Any, length: int) -> Iterator[bytes]:
"""Yield the next *length* bytes of *fh* without holding a whole part in memory."""
while length > 0:
chunk = fh.read(min(_READ_CHUNK, length))
if not chunk:
return
length -= len(chunk)
yield chunk
def _retryable(exc: httpx.HTTPError) -> bool:
import httpx # noqa: PLC0415
if isinstance(exc, httpx.HTTPStatusError):
return exc.response.status_code in _RETRYABLE_STATUS
return isinstance(exc, httpx.TransportError)
def _write_output_file(path: str, data: dict[str, Any]) -> None:
"""
Write the output dict to a JSON file.
Parameters
----------
path : str
Destination file path.
data : dict[str, Any]
Output data to serialise.
Raises
------
typer.Exit
If the file cannot be written.
"""
out = Path(path)
payload = json.dumps(data, indent=2)
try:
out.parent.mkdir(parents=True, exist_ok=True)
try:
_replace_output(out, payload)
except OSError:
# The rename needs permission on the directory, opening the file
# needs it on the file. Neither covers every case, so try both.
out.write_text(payload)
except OSError as exc:
fail(f"failed to write output file '{path}': {exc}")
def _replace_output(out: Path, payload: str) -> None:
"""Write *payload* beside *out* and rename it over.
The rename never truncates what is already there, so a write that dies
halfway leaves the previous file intact instead of a half-written one.
"""
tmp = out.with_name(f".{out.name}.{os.getpid()}.tmp")
try:
tmp.write_text(payload)
# The rename swaps the inode, so the mode comes from this process's
# umask, not from the file being replaced. The worker reads the result
# back under a different uid, so set it explicitly.
os.chmod(tmp, 0o644)
os.replace(tmp, out)
finally:
# The worker only clears the run directory once the run is finalised,
# so debris here outlives the run that left it.
tmp.unlink(missing_ok=True)
[docs]
def run(
file: str | None = typer.Option(
None,
"--file",
"-f",
help="Python file that defines one @gyoza_op function.",
),
function: str | None = typer.Option(
None,
"--function",
"-F",
help=(
"Name of the @gyoza_op function to execute. "
"Defaults to the first one found in the file."
),
),
input: str | None = typer.Option(
None,
"--input",
"-i",
help=(
"JSON input file path. "
"Falls back to $GYOZA_INPUT_PATH. "
"Ignored when --inline-input is provided."
),
),
output: str | None = typer.Option(
None,
"--output",
"-o",
help=("JSON output file path. Falls back to $GYOZA_OUTPUT_PATH."),
),
inline_input: str | None = typer.Option(
None,
"--inline-input",
help=(
'Inline JSON dict string for inputs, e.g. \'{"a": 1, "b": 2}\'. '
"Mutually exclusive with --input."
),
),
) -> None:
"""
Execute a @gyoza_op function from a Python file.
Three input modes are supported:
1. ``--inline-input`` — parse input directly from the CLI string.
2. ``--input`` — read input from a JSON file at the given path.
3. default — read input from the path in $GYOZA_INPUT_PATH.
In all modes the output dict is written as JSON to ``--output`` or
$GYOZA_OUTPUT_PATH.
"""
if not file:
fail("--file is required")
if inline_input and input:
fail("--input and --inline-input are mutually exclusive")
code_path = Path(file).resolve()
if not code_path.exists():
fail(f"file not found: {file}")
resolved_output = output or os.getenv("GYOZA_OUTPUT_PATH", _DEFAULT_OUTPUT_PATH)
typer.echo(f"\n{_SEPARATOR}")
typer.echo(" ▶️ GYOZA RUN")
typer.echo(f"{_SEPARATOR}\n")
typer.echo(f" 📄 File: {code_path}")
typer.echo(f" 🎯 Function: {function or '(auto-detect)'}")
try:
module = _load_module(code_path)
func = _find_gyoza_op(module, function)
except typer.Exit:
raise
except Exception as exc: # noqa: BLE001
fail(f"failed to load module: {exc}")
if inline_input is not None:
mode = "inline"
input_dict = _parse_inline_input(inline_input)
typer.echo(f" 📝 Mode: {mode}")
typer.echo(f" 📥 Input: {inline_input}")
else:
resolved_input = input or os.getenv("GYOZA_INPUT_PATH", _DEFAULT_INPUT_PATH)
mode = "path" if input else "env"
typer.echo(f" 📝 Mode: {mode}")
typer.echo(f" 📥 Input: {resolved_input}")
input_dict = _read_input_file(resolved_input)
typer.echo(f" 📤 Output: {resolved_output}")
typer.echo(f"\n{_SEPARATOR}")
typer.echo(" 🚀 Executing...\n")
output_template = _read_output_template(resolved_output)
try:
output_dict = func.run(input_dict)
except Exception as exc: # noqa: BLE001
fail(f"op failed: {exc}")
output_dict = _deliver_outputs(output_dict, output_template)
_write_output_file(resolved_output, output_dict)
typer.echo(f"\n{_SEPARATOR}")
typer.echo(" ✅ Done.")
typer.echo(f"{_SEPARATOR}\n")
