File size: 3,625 Bytes
e0265b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | from __future__ import annotations
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any
from uuid import uuid4
def utc_now() -> str:
return datetime.now(timezone.utc).isoformat()
class JobStatus(str, Enum):
DRAFT = "Draft"
AWAITING_CONFIRMATION = "Awaiting confirmation"
QUEUED = "Queued"
RUNNING = "Running"
PAUSED = "Paused"
FINISHED = "Finished"
FAILED = "Failed"
CANCELLED = "Cancelled"
INTERRUPTED = "Interrupted"
class StepStatus(str, Enum):
PENDING = "Pending"
RUNNING = "Running"
FINISHED = "Finished"
FAILED = "Failed"
SKIPPED = "Skipped"
@dataclass(slots=True)
class PlanStep:
tool_id: str
title: str
description: str
arguments: dict[str, Any] = field(default_factory=dict)
status: StepStatus = StepStatus.PENDING
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "PlanStep":
payload = dict(data)
payload["status"] = StepStatus(payload.get("status", StepStatus.PENDING))
return cls(**payload)
@dataclass(slots=True)
class ExecutionPlan:
request: str
summary: str
steps: list[PlanStep]
requires_confirmation: bool = False
confirmation_reason: str = ""
project_name: str = "Untitled project"
id: str = field(default_factory=lambda: uuid4().hex[:10])
created_at: str = field(default_factory=utc_now)
orion_review: dict[str, Any] = field(default_factory=dict)
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "ExecutionPlan":
payload = dict(data)
payload["steps"] = [PlanStep.from_dict(item) for item in payload["steps"]]
return cls(**payload)
@dataclass(slots=True)
class Job:
plan: ExecutionPlan
id: str = field(default_factory=lambda: uuid4().hex[:8].upper())
status: JobStatus = JobStatus.DRAFT
progress: int = 0
current_step: int = -1
created_at: str = field(default_factory=utc_now)
started_at: str | None = None
ended_at: str | None = None
output_folder: str | None = None
logs: list[str] = field(default_factory=list)
error: str | None = None
preview_path: str | None = None
preview_epoch: int = 0
preview_next_epoch: int = 0
preview_prompt: str = ""
preview_seed: int | None = None
preview_steps: int = 0
preview_kind: str = "training"
preview_current: int = 0
preview_total: int = 0
preview_image_index: int = 0
preview_image_count: int = 0
atlas_report: dict[str, Any] = field(default_factory=dict)
nova_report: dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> dict[str, Any]:
payload = asdict(self)
payload["status"] = self.status.value
for step in payload["plan"]["steps"]:
step["status"] = step["status"].value
return payload
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "Job":
payload = dict(data)
payload["status"] = JobStatus(payload["status"])
payload["plan"] = ExecutionPlan.from_dict(payload["plan"])
return cls(**payload)
@dataclass(slots=True)
class SystemSnapshot:
cpu_percent: float = 0.0
memory_percent: float = 0.0
memory_used_gb: float = 0.0
memory_total_gb: float = 0.0
disk_percent: float = 0.0
disk_used_gb: float = 0.0
disk_total_gb: float = 0.0
gpu_name: str = "No NVIDIA GPU detected"
gpu_percent: float = 0.0
vram_percent: float = 0.0
vram_used_gb: float = 0.0
vram_total_gb: float = 0.0
gpu_temperature: float | None = None
|