Update public source
This commit is contained in:
@@ -0,0 +1,258 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
JOB_STATUS_QUEUED = "QUEUED"
|
||||
JOB_STATUS_RUNNING = "RUNNING"
|
||||
JOB_STATUS_COMPLETED = "COMPLETED"
|
||||
JOB_STATUS_PARTIAL = "PARTIAL"
|
||||
JOB_STATUS_FAILED = "FAILED"
|
||||
|
||||
STEP_STATUS_PENDING = "PENDING"
|
||||
STEP_STATUS_RUNNING = "RUNNING"
|
||||
STEP_STATUS_COMPLETED = "COMPLETED"
|
||||
STEP_STATUS_PARTIAL = "PARTIAL"
|
||||
STEP_STATUS_FAILED = "FAILED"
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return (
|
||||
datetime.now(timezone.utc)
|
||||
.replace(microsecond=0)
|
||||
.isoformat()
|
||||
.replace("+00:00", "Z")
|
||||
)
|
||||
|
||||
|
||||
class StateError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def _atomic_write_json(
|
||||
path: Path,
|
||||
data: dict[str, Any],
|
||||
) -> None:
|
||||
path.parent.mkdir(
|
||||
parents=True,
|
||||
exist_ok=True,
|
||||
)
|
||||
|
||||
temp_path = path.with_name(
|
||||
f".{path.name}.tmp"
|
||||
)
|
||||
|
||||
payload = json.dumps(
|
||||
data,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
) + "\n"
|
||||
|
||||
temp_path.write_text(
|
||||
payload,
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
temp_path.replace(path)
|
||||
|
||||
|
||||
def _load_json(path: Path) -> dict[str, Any]:
|
||||
try:
|
||||
data = json.loads(
|
||||
path.read_text(
|
||||
encoding="utf-8",
|
||||
)
|
||||
)
|
||||
except FileNotFoundError as exc:
|
||||
raise StateError(
|
||||
f"Missing state file: {path}"
|
||||
) from exc
|
||||
except json.JSONDecodeError as exc:
|
||||
raise StateError(
|
||||
f"Invalid JSON in state file {path}: {exc}"
|
||||
) from exc
|
||||
|
||||
if not isinstance(data, dict):
|
||||
raise StateError(
|
||||
f"State file must contain an object: {path}"
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def state_path(job_dir: Path) -> Path:
|
||||
return job_dir / "state" / "client_state.json"
|
||||
|
||||
|
||||
def load_state(job_dir: Path) -> dict[str, Any]:
|
||||
return _load_json(
|
||||
state_path(job_dir)
|
||||
)
|
||||
|
||||
|
||||
def save_state(
|
||||
job_dir: Path,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
_atomic_write_json(
|
||||
state_path(job_dir),
|
||||
state,
|
||||
)
|
||||
|
||||
|
||||
def initialize_state(
|
||||
job_dir: Path,
|
||||
job: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
job_id = job.get("job_id")
|
||||
|
||||
if not isinstance(job_id, str) or not job_id:
|
||||
raise StateError(
|
||||
"Job is missing a valid job_id."
|
||||
)
|
||||
|
||||
steps = job.get("steps")
|
||||
|
||||
if not isinstance(steps, list) or not steps:
|
||||
raise StateError(
|
||||
f"Job {job_id} has no steps."
|
||||
)
|
||||
|
||||
now = utc_now()
|
||||
|
||||
state: dict[str, Any] = {
|
||||
"job_id": job_id,
|
||||
"status": JOB_STATUS_QUEUED,
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
"steps": {},
|
||||
}
|
||||
|
||||
for step in steps:
|
||||
step_no = step.get("step_no")
|
||||
|
||||
if not isinstance(step_no, int):
|
||||
raise StateError(
|
||||
f"Job {job_id} contains invalid step_no."
|
||||
)
|
||||
|
||||
state["steps"][str(step_no)] = {
|
||||
"step_no": step_no,
|
||||
"step_id": step.get("step_id"),
|
||||
"session_name": step.get("session_name"),
|
||||
"status": STEP_STATUS_PENDING,
|
||||
"started_at": None,
|
||||
"completed_at": None,
|
||||
"exit_code": None,
|
||||
"restore_seen": False,
|
||||
"error": None,
|
||||
}
|
||||
|
||||
save_state(
|
||||
job_dir,
|
||||
state,
|
||||
)
|
||||
|
||||
return state
|
||||
|
||||
|
||||
def update_job_state(
|
||||
job_dir: Path,
|
||||
status: str,
|
||||
*,
|
||||
error: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
state = load_state(job_dir)
|
||||
|
||||
state["status"] = status
|
||||
state["updated_at"] = utc_now()
|
||||
|
||||
if error is not None:
|
||||
state["error"] = error
|
||||
elif status != JOB_STATUS_FAILED:
|
||||
state.pop("error", None)
|
||||
|
||||
save_state(
|
||||
job_dir,
|
||||
state,
|
||||
)
|
||||
|
||||
return state
|
||||
|
||||
|
||||
def update_step_state(
|
||||
job_dir: Path,
|
||||
step_no: int,
|
||||
status: str,
|
||||
*,
|
||||
exit_code: int | None = None,
|
||||
restore_seen: bool | None = None,
|
||||
error: str | None = None,
|
||||
started: bool = False,
|
||||
completed: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
state = load_state(job_dir)
|
||||
|
||||
key = str(step_no)
|
||||
steps = state.get("steps")
|
||||
|
||||
if not isinstance(steps, dict) or key not in steps:
|
||||
raise StateError(
|
||||
f"Unknown step {step_no}."
|
||||
)
|
||||
|
||||
step_state = steps[key]
|
||||
|
||||
step_state["status"] = status
|
||||
|
||||
if started:
|
||||
step_state["started_at"] = utc_now()
|
||||
|
||||
if completed:
|
||||
step_state["completed_at"] = utc_now()
|
||||
|
||||
if exit_code is not None:
|
||||
step_state["exit_code"] = exit_code
|
||||
|
||||
if restore_seen is not None:
|
||||
step_state["restore_seen"] = restore_seen
|
||||
|
||||
if error is not None:
|
||||
step_state["error"] = error
|
||||
elif status != STEP_STATUS_FAILED:
|
||||
step_state["error"] = None
|
||||
|
||||
state["updated_at"] = utc_now()
|
||||
|
||||
save_state(
|
||||
job_dir,
|
||||
state,
|
||||
)
|
||||
|
||||
return state
|
||||
|
||||
|
||||
def get_step_state(
|
||||
job_dir: Path,
|
||||
step_no: int,
|
||||
) -> dict[str, Any]:
|
||||
state = load_state(job_dir)
|
||||
|
||||
steps = state.get("steps")
|
||||
|
||||
if not isinstance(steps, dict):
|
||||
raise StateError(
|
||||
f"Invalid steps state in {job_dir}."
|
||||
)
|
||||
|
||||
step = steps.get(str(step_no))
|
||||
|
||||
if not isinstance(step, dict):
|
||||
raise StateError(
|
||||
f"Unknown step {step_no}."
|
||||
)
|
||||
|
||||
return step
|
||||
Reference in New Issue
Block a user