146 lines
3.9 KiB
Python
146 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import Annotated, Literal
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
|
|
|
from modelforge_api.domain.serving import (
|
|
EmbeddingInvokeRequest,
|
|
OCRInvokeRequest,
|
|
RerankingInvokeRequest,
|
|
SpeechTranscriptionInvokeRequest,
|
|
VisionEmbeddingInvokeRequest,
|
|
)
|
|
|
|
|
|
class LabModel(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
|
|
class LabTarget(LabModel):
|
|
deployment_candidate_id: uuid.UUID
|
|
artifact_set_id: uuid.UUID
|
|
runtime_profile_id: uuid.UUID
|
|
compute_node_id: uuid.UUID
|
|
execution_approval_id: uuid.UUID
|
|
|
|
|
|
class LabEmbeddingInput(LabModel):
|
|
capability: Literal["rag.embedding"]
|
|
payload: EmbeddingInvokeRequest
|
|
|
|
|
|
class LabRerankingInput(LabModel):
|
|
capability: Literal["rag.reranking"]
|
|
payload: RerankingInvokeRequest
|
|
|
|
|
|
class LabOCRInput(LabModel):
|
|
capability: Literal["document.ocr"]
|
|
payload: OCRInvokeRequest
|
|
|
|
|
|
class LabVisionInput(LabModel):
|
|
capability: Literal["vision.embedding"]
|
|
payload: VisionEmbeddingInvokeRequest
|
|
|
|
|
|
class LabSpeechInput(LabModel):
|
|
capability: Literal["speech.transcription"]
|
|
payload: SpeechTranscriptionInvokeRequest
|
|
|
|
|
|
LabInput = Annotated[
|
|
LabEmbeddingInput | LabRerankingInput | LabOCRInput | LabVisionInput | LabSpeechInput,
|
|
Field(discriminator="capability"),
|
|
]
|
|
|
|
|
|
class LabLimits(LabModel):
|
|
timeout_seconds: int = Field(default=60, ge=5, le=300, strict=True)
|
|
queue_timeout_seconds: int = Field(default=10, ge=1, le=60, strict=True)
|
|
retention_seconds: int = Field(default=300, ge=60, le=3600, strict=True)
|
|
max_output_bytes: int = Field(default=65536, ge=1, le=1048576, strict=True)
|
|
concurrency: Literal[1] = 1
|
|
|
|
@model_validator(mode="after")
|
|
def bounded_queue(self) -> LabLimits:
|
|
if self.queue_timeout_seconds > self.timeout_seconds:
|
|
raise ValueError("queue timeout cannot exceed total timeout")
|
|
return self
|
|
|
|
|
|
class LabPreflightRequest(LabModel):
|
|
target: LabTarget
|
|
input: LabInput
|
|
contract_version: Literal[1] = 1
|
|
limits: LabLimits = Field(default_factory=LabLimits)
|
|
|
|
|
|
class LabExecutionStartRequest(LabPreflightRequest):
|
|
"""An explicitly selected embedding experiment; body input is never echoed in status."""
|
|
|
|
deployment_id: uuid.UUID
|
|
experiment_route_id: uuid.UUID
|
|
|
|
|
|
class LabInputSummary(LabModel):
|
|
sha256: str
|
|
item_count: int
|
|
encoded_bytes: int
|
|
|
|
|
|
class LabBinding(LabTarget):
|
|
compatibility_assessment_id: uuid.UUID
|
|
runtime_probe_id: uuid.UUID
|
|
capability: str
|
|
contract_version: Literal[1] = 1
|
|
runtime_profile_fingerprint: str
|
|
assessment_fingerprint: str
|
|
approval_evidence_fingerprint: str
|
|
probe_environment_fingerprint: str
|
|
accelerator_id: uuid.UUID | None = None
|
|
|
|
|
|
class LabPreflightResponse(LabModel):
|
|
input_summary: LabInputSummary
|
|
binding: LabBinding | None = None
|
|
limits: LabLimits
|
|
blockers: list[str]
|
|
checks_passed: bool
|
|
execution_started: Literal[False] = False
|
|
permission_granted: Literal[False] = False
|
|
|
|
|
|
LabExecutionStatus = Literal[
|
|
"accepted",
|
|
"preparing",
|
|
"queued",
|
|
"running",
|
|
"cancel_requested",
|
|
"completed",
|
|
"failed",
|
|
"cancelled",
|
|
"interrupted",
|
|
]
|
|
|
|
|
|
class LabExecutionRunResponse(LabModel):
|
|
id: uuid.UUID
|
|
capability: Literal["rag.embedding@1"] = "rag.embedding@1"
|
|
status: LabExecutionStatus
|
|
target_snapshot: dict[str, object]
|
|
input_summary: LabInputSummary
|
|
output_metadata: dict[str, object] | None = None
|
|
worker_observations: dict[str, object] | None = None
|
|
failure_code: str | None = None
|
|
input_payload_expires_at: datetime
|
|
retention_seconds: int = Field(ge=60, le=3600)
|
|
created_at: datetime
|
|
started_at: datetime | None = None
|
|
cancellation_requested_at: datetime | None = None
|
|
cancellation_acknowledged_at: datetime | None = None
|
|
finished_at: datetime | None = None
|