"""Pydantic models for API request/response validation."""
from typing import Any
from pydantic import BaseModel, Field
_OUTPUTS_TEMPLATE_DESCRIPTION = (
"Partial outputs template. Pre-fill a field with an upload URL "
"to have the op runtime deliver that output there."
)
[docs]
class OutputSpecModel(BaseModel):
"""Output specification model."""
type: str
description: str | None = None
example: Any | None = None
[docs]
class ConstraintsModel(BaseModel):
"""Hardware constraints model. An unset field reserves nothing."""
ram_mb: int | None = None
vram_mb: int | None = None
cpu: int | None = None
gpu_id: int | None = None
[docs]
class RetryPolicyModel(BaseModel):
"""Retry policy model."""
max_attempts: int = 3
[docs]
class EventDeliveryModel(BaseModel):
"""Event delivery configuration model."""
topic: str = ""
attributes: dict[str, str] | None = None
[docs]
class OpDefinitionCreate(BaseModel):
"""Request model for creating/updating OpDefinition."""
id: str = Field(..., description="Unique identifier for the operation")
version: str = Field(..., description="Semantic version (e.g., '1.0.0')")
image: str = Field(..., description="Docker image identifier")
description: str | None = None
input_specs: dict[str, InputSpecModel] | None = None
output_specs: dict[str, OutputSpecModel] | None = None
constraints: ConstraintsModel | None = None
retry_policy: RetryPolicyModel | None = None
event_delivery: EventDeliveryModel | None = None
[docs]
class OpDefinitionResponse(BaseModel):
"""Response model for OpDefinition."""
id: str
version: str
image: str
description: str | None = None
input_specs: dict[str, Any]
output_specs: dict[str, Any]
constraints: dict[str, Any]
retry_policy: dict[str, Any]
event_delivery: dict[str, Any]
createdAt: str
updatedAt: str
[docs]
class OpDefinitionDeleteResponse(BaseModel):
"""Response model for an OpDefinition deletion."""
id: str
version: str | None = None
deleted: int
[docs]
class OpRunCreateFromDefinition(BaseModel):
"""Request model for creating OpRun from OpDefinition.
This is used when creating a run from a template (OpDefinition).
For creating runs from scratch, use OpRunCreate (to be implemented).
"""
inputs: dict[str, Any] = Field(default_factory=dict)
outputs: dict[str, Any] = Field(
default_factory=dict, description=_OUTPUTS_TEMPLATE_DESCRIPTION
)
priority: int = Field(default=0, ge=0, le=10)
event_delivery: EventDeliveryModel | None = Field(
default=None,
description=(
"Per-run event delivery override. Merged over the definition's: "
"a non-empty topic replaces it, and attributes are merged key by "
"key. Use it to carry a correlation id for the run, so a consumer "
"of the topic feed can tie the run back to its own record."
),
)
[docs]
class OpRunResponse(BaseModel):
"""Response model for OpRun."""
id: str
state: str
priority: int
progress: int
current_attempt: int
image: str
inputs: dict[str, Any]
outputs: dict[str, Any]
events: list[dict[str, Any]]
execution_summary: dict[str, Any]
constraints: dict[str, Any]
retry_policy: dict[str, Any]
event_delivery: dict[str, Any]
op_definition: str | None
createdAt: str
updatedAt: str
[docs]
class OpRunListItem(BaseModel):
"""OpRun as returned by the list endpoint.
``inputs``, ``outputs`` and ``events`` can be arbitrarily large (GeoJSON
inputs, one event per progress line), so the listing omits them unless
explicitly requested via ``?expand=``.
"""
id: str
state: str
priority: int
progress: int
current_attempt: int
image: str
inputs: dict[str, Any] | None = None
outputs: dict[str, Any] | None = None
events: list[dict[str, Any]] | None = None
execution_summary: dict[str, Any]
constraints: dict[str, Any]
retry_policy: dict[str, Any]
event_delivery: dict[str, Any]
op_definition: str | None
createdAt: str
updatedAt: str
[docs]
class OpRunListResponse(BaseModel):
"""Stripe-style list response with cursor-based pagination."""
object: str = "list"
data: list[OpRunListItem]
has_more: bool
url: str = "/runs"
[docs]
class RunStatsWindow(BaseModel):
"""The resolved time window a stats response covers."""
start: int = Field(
...,
description=(
"UNIX timestamp of the window start, floored to a bucket boundary. "
"Bucket i covers [start + i * bucket_seconds, +bucket_seconds)."
),
)
end: int = Field(..., description="UNIX timestamp of the window end.")
bucket_seconds: int = Field(..., description="Width of each bucket, in seconds.")
buckets: int = Field(..., description="Number of buckets in each series.")
[docs]
class RunStatsTotals(BaseModel):
"""Run counts over the whole window."""
all: int = Field(..., description="Total runs created in the window.")
by_state: dict[str, int] = Field(
default_factory=dict, description="Run count per state."
)
[docs]
class RunStatsDefinition(BaseModel):
"""Per-definition activity within the window."""
op_definition: str
count: int = Field(..., description="Total runs of this definition in the window.")
buckets: list[int] = Field(
...,
description="Run count per bucket, oldest first. Length is window.buckets.",
)
[docs]
class RunStatsResponse(BaseModel):
"""Aggregated run activity, computed server-side.
Lets a dashboard render totals and per-definition sparklines from one small
response instead of paginating every run in the window.
"""
object: str = "run_stats"
window: RunStatsWindow
totals: RunStatsTotals
definitions: list[RunStatsDefinition] = Field(
..., description="Busiest definitions in the window, most runs first."
)
[docs]
class OpRunCreate(BaseModel):
"""Request model for creating OpRun directly (ad-hoc).
This is used when creating a run from scratch without a template.
For creating runs from OpDefinition, use OpRunCreateFromDefinition.
"""
image: str = Field(..., description="Docker image identifier")
priority: int = Field(default=0, ge=0, le=10)
inputs: dict[str, Any] = Field(default_factory=dict)
outputs: dict[str, Any] = Field(
default_factory=dict, description=_OUTPUTS_TEMPLATE_DESCRIPTION
)
constraints: ConstraintsModel | None = None
retry_policy: RetryPolicyModel | None = None
event_delivery: EventDeliveryModel | None = None
[docs]
class OpRunUpdate(BaseModel):
"""Request model for updating OpRun via PATCH."""
state: str | None = None
outputs: dict[str, Any] | None = None
priority: int | None = Field(default=None, ge=0, le=10)
[docs]
class OpAttemptResponse(BaseModel):
"""Response model for OpAttempt."""
id: str
op_run_id: str
attempt: int
state: str
progress: int
events: list[dict[str, Any]]
inputs: dict[str, Any]
outputs: dict[str, Any]
execution_summary: dict[str, Any]
constraints: dict[str, Any]
started_at: str | None
finished_at: str | None
[docs]
class AddEventRequest(BaseModel):
"""Request model for adding an event to an attempt."""
type: str = Field(
..., description="Event type (STARTED, INFO, PROGRESS, COMPLETED, FAILED, etc.)"
)
msg: str | int = Field(..., description="Event message or progress value")
payload: dict[str, Any] | None = Field(
default=None, description="Optional event-specific data"
)
[docs]
class EventEntryResponse(BaseModel):
"""Response model for a single event entry."""
id: int
type: str
t: str
msg: str | int
state: str
payload: dict[str, Any] = Field(default_factory=dict)
[docs]
class EventsResponse(BaseModel):
"""Response model for polling events."""
events: list[EventEntryResponse]
[docs]
class TopicEventResponse(BaseModel):
"""A single event in the topic feed."""
cursor: str
run_id: str
attempt: int
type: str
msg: str | int
t: str
state: str
payload: dict[str, Any] = Field(default_factory=dict)
[docs]
class TopicEventsResponse(BaseModel):
"""Cursor-paginated event feed for a topic."""
events: list[TopicEventResponse]
next_cursor: str | None = None
has_more: bool = False
[docs]
class ErrorResponse(BaseModel):
"""Error response model."""
detail: str
errors: list[str] | None = None
[docs]
class GPUModel(BaseModel):
"""GPU resource model."""
id: int
vram_mb: int
tags: list[str] = Field(default_factory=list)
[docs]
class ResourcesModel(BaseModel):
"""Hardware resources model."""
cpu_cores: int
ram_mb: int
gpus: list[GPUModel] = Field(default_factory=list)
[docs]
class HeartbeatRequest(BaseModel):
"""Request model for worker heartbeat."""
worker_id: str = Field(..., description="Worker identifier")
resources: ResourcesModel
tags: list[str] = Field(default_factory=list)
running_ops: list[dict[str, Any]] | None = Field(
default=None, description="List of running OpRuns"
)
[docs]
class WorkerResponse(BaseModel):
"""Response model for Worker."""
id: str
resources: dict[str, Any]
tags: list[str]
running_ops: list[dict[str, Any]]
created_at: str
last_heartbeat_at: str
is_active: bool = Field(
...,
description="True when last heartbeat is within GYOZA_WORKER_TIMEOUT seconds.",
)
[docs]
class WorkerOpRunResponse(BaseModel):
"""Response model for WorkerOpRun."""
id: str
image: str
inputs: dict[str, Any]
outputs: dict[str, Any] = Field(default_factory=dict)
constraints: dict[str, Any]
[docs]
class ClaimOpsResponse(BaseModel):
"""Response model for claim ops endpoint."""
ops: list[WorkerOpRunResponse]
