98 lines
2.6 KiB
Python
98 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from enum import Enum
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
|
|
class BatchStatus(str, Enum):
|
|
PENDING = "pending"
|
|
RUNNING = "running"
|
|
COMPLETED = "completed"
|
|
COMPLETED_WITH_ERRORS = "completed_with_errors"
|
|
FAILED = "failed"
|
|
|
|
|
|
class BatchItemStatus(str, Enum):
|
|
QUEUED = "queued"
|
|
DISPATCHING = "dispatching"
|
|
WAITING_CALLBACK = "waiting_callback"
|
|
SUCCEEDED = "succeeded"
|
|
FAILED = "failed"
|
|
|
|
|
|
class PluginResultStatus(str, Enum):
|
|
SUCCESS = "success"
|
|
FAILED = "failed"
|
|
|
|
|
|
class BatchSubmitItem(BaseModel):
|
|
model_config = ConfigDict(protected_namespaces=())
|
|
model_path: str
|
|
task_type: str
|
|
task_params: Dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class BatchSubmitRequest(BaseModel):
|
|
items: List[BatchSubmitItem] = Field(min_length=1)
|
|
batch_name: Optional[str] = None
|
|
metadata: Dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class BatchItem(BaseModel):
|
|
model_config = ConfigDict(protected_namespaces=())
|
|
id: str
|
|
batch_id: str
|
|
sequence: int
|
|
model_path: str
|
|
file_extension: Optional[str] = None
|
|
software_id: Optional[str] = None
|
|
task_type: str
|
|
task_params: Dict[str, Any] = Field(default_factory=dict)
|
|
|
|
status: BatchItemStatus = BatchItemStatus.QUEUED
|
|
retry_count: int = 0
|
|
max_retries: int = 0
|
|
execution_id: Optional[str] = None
|
|
|
|
created_at: datetime = Field(default_factory=datetime.now)
|
|
started_at: Optional[datetime] = None
|
|
finished_at: Optional[datetime] = None
|
|
|
|
error_message: Optional[str] = None
|
|
result: Dict[str, Any] = Field(default_factory=dict)
|
|
plugin_response: Dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class BatchJob(BaseModel):
|
|
id: str
|
|
name: Optional[str] = None
|
|
submitter_id: Optional[str] = None
|
|
status: BatchStatus = BatchStatus.PENDING
|
|
|
|
created_at: datetime = Field(default_factory=datetime.now)
|
|
started_at: Optional[datetime] = None
|
|
finished_at: Optional[datetime] = None
|
|
|
|
total_count: int = 0
|
|
queued_count: int = 0
|
|
running_count: int = 0
|
|
succeeded_count: int = 0
|
|
failed_count: int = 0
|
|
|
|
item_ids: List[str] = Field(default_factory=list)
|
|
metadata: Dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class PluginCallbackPayload(BaseModel):
|
|
execution_id: str
|
|
software_id: str
|
|
status: PluginResultStatus
|
|
error_message: Optional[str] = None
|
|
result: Dict[str, Any] = Field(default_factory=dict)
|
|
finished_at: Optional[datetime] = None
|
|
token: Optional[str] = None
|
|
signature: Optional[str] = None
|