Reconcile durable worker lifecycles

server_api/main.py

@@ -54,9 +54,16 @@
 )
 from server_api.workflows.operation_service import (
     create_workflow_operation,
+    get_workflow_operation_or_404,
+    heartbeat_workflow_operation,
     operation_to_dict,
     transition_workflow_operation,
 )
+from server_api.workflows.runtime_reconciliation import (
+    reconcile_runtime_operation,
+    reconciliation_response,
+    runtime_kind_for_operation,
+)
 from server_api.workflows.service import (
     append_event_for_workflow_if_present,
     command_to_dict,
@@ -903,13 +910,18 @@ def _workflow_command_run_response(
     operation: WorkflowOperation,
 ) -> dict[str, Any]:
     result = decode_json(operation.result_json)
+    command_result = decode_json(command.result_json)
+    if not result:
+        result = command_result
     return {
         "workflow_id": workflow.id,
         "command": command_to_dict(command),
         "operation": operation_to_dict(operation),
-        "worker": result.get("worker", {}),
-        "run_id": result.get("run_id"),
-        "started_event_id": result.get("started_event_id"),
+        "worker": result.get("worker", command_result.get("worker", {})),
+        "run_id": result.get("run_id", command_result.get("run_id")),
+        "started_event_id": result.get(
+            "started_event_id", command_result.get("started_event_id")
+        ),
     }
 
 
@@ -2678,21 +2690,21 @@ async def run_workflow_command(
         )
 
     if command.status == "submitted":
-        completed_operation = (
+        submitted_operation = (
             db.query(WorkflowOperation)
             .filter(
                 WorkflowOperation.workflow_id == workflow.id,
                 WorkflowOperation.command_id == command.id,
-                WorkflowOperation.status == "succeeded",
+                WorkflowOperation.status.in_({"running", "succeeded"}),
             )
             .order_by(WorkflowOperation.id.desc())
             .first()
         )
-        if completed_operation is not None:
+        if submitted_operation is not None:
             return _workflow_command_run_response(
                 workflow,
                 command,
-                completed_operation,
+                submitted_operation,
             )
         raise HTTPException(
             status_code=409, detail="Workflow command was already submitted."
@@ -2836,12 +2848,13 @@ async def run_workflow_command(
             result_payload=operation_result,
             commit=False,
         )
-        operation = transition_workflow_operation(
+        operation = heartbeat_workflow_operation(
             db,
             operation,
-            status="succeeded",
-            expected_status="running",
-            result_payload=operation_result,
+            metadata={
+                "worker": worker_data,
+                "worker_submission": {"accepted": True},
+            },
             lease_owner=runner_name,
             commit=False,
         )
@@ -2914,6 +2927,43 @@ async def run_workflow_command(
         raise HTTPException(status_code=500, detail=error_payload) from exc
 
 
+@app.post("/api/workflows/{workflow_id}/operations/{operation_id}/reconcile-runtime")
+async def reconcile_workflow_runtime_operation(
+    workflow_id: int,
+    operation_id: int,
+    current_user: models.User = Depends(get_current_user),
+    db: Session = Depends(get_db),
+):
+    """Refresh a correlated worker runtime snapshot into one operation record."""
+    workflow = get_user_workflow_or_404(
+        db, workflow_id=int(workflow_id), user_id=current_user.id
+    )
+    operation = get_workflow_operation_or_404(
+        db,
+        workflow_id=workflow.id,
+        operation_id=int(operation_id),
+    )
+    runtime_kind = runtime_kind_for_operation(operation)
+    if runtime_kind is None:
+        raise HTTPException(
+            status_code=400,
+            detail="Workflow operation does not represent a PyTC runtime.",
+        )
+    snapshot = _proxy_to_worker(
+        "get",
+        "/training_logs" if runtime_kind == "training" else "/inference_logs",
+        timeout=5,
+    )
+    return reconciliation_response(
+        reconcile_runtime_operation(
+            db,
+            workflow=workflow,
+            operation=operation,
+            snapshot=snapshot,
+        )
+    )
+
+
 @app.post("/stop_model_training")
 async def stop_model_training():
     worker_data = _proxy_to_worker("post", "/stop_model_training", timeout=30)

server_api/workflows/runtime_reconciliation.py

@@ -0,0 +1,224 @@
+"""Project correlated PyTC runtime snapshots onto durable workflow operations."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from sqlalchemy.orm import Session
+
+from .db_models import WorkflowOperation, WorkflowSession
+from .operation_service import (
+    TERMINAL_OPERATION_STATUSES,
+    operation_to_dict,
+    transition_workflow_operation,
+)
+from .service import (
+    append_event_for_workflow_if_present,
+    decode_json,
+    update_workflow_fields,
+)
+
+RUNTIME_OPERATION_KINDS = {
+    "start_training": "training",
+    "start_inference": "inference",
+}
+
+
+def runtime_kind_for_operation(operation: WorkflowOperation) -> Optional[str]:
+    return RUNTIME_OPERATION_KINDS.get(operation.operation_type)
+
+
+def _metadata_value(metadata: Dict[str, Any], *keys: str) -> Optional[str]:
+    for key in keys:
+        value = metadata.get(key)
+        if value is not None and str(value).strip():
+            return str(value)
+    return None
+
+
+def _snapshot_metadata(snapshot: Dict[str, Any]) -> Dict[str, Any]:
+    metadata = snapshot.get("metadata")
+    return metadata if isinstance(metadata, dict) else {}
+
+
+def _is_correlated(
+    operation: WorkflowOperation,
+    workflow: WorkflowSession,
+    snapshot: Dict[str, Any],
+) -> bool:
+    metadata = _snapshot_metadata(snapshot)
+    operation_metadata = decode_json(operation.metadata_json)
+    expected_run_id = _metadata_value(operation_metadata, "run_id", "runId")
+    actual_workflow_id = _metadata_value(metadata, "workflowId", "workflow_id")
+    actual_command_id = _metadata_value(metadata, "commandId", "command_id")
+    actual_run_id = _metadata_value(metadata, "runId", "run_id")
+
+    # A worker snapshot without all identifiers may belong to an older direct
+    # browser launch. Never project that global process state onto this command.
+    return (
+        actual_workflow_id == str(workflow.id)
+        and actual_command_id == str(operation.command_id)
+        and bool(expected_run_id)
+        and actual_run_id == expected_run_id
+    )
+
+
+def _runtime_details(snapshot: Dict[str, Any]) -> Dict[str, Any]:
+    metadata = _snapshot_metadata(snapshot)
+    return {
+        "runtimePhase": snapshot.get("phase"),
+        "runtimePid": snapshot.get("pid"),
+        "runtimeExitCode": snapshot.get("exitCode"),
+        "runtimeStartedAt": snapshot.get("startedAt"),
+        "runtimeEndedAt": snapshot.get("endedAt"),
+        "runtimeLastError": snapshot.get("lastError"),
+        "runtimeLineCount": snapshot.get("lineCount"),
+        "runtimeMetadata": metadata,
+    }
+
+
+def _terminal_status(
+    operation: WorkflowOperation, snapshot: Dict[str, Any]
+) -> Optional[str]:
+    phase = str(snapshot.get("phase") or "").lower()
+    exit_code = snapshot.get("exitCode")
+    if phase == "finished" and exit_code == 0:
+        return "succeeded"
+    if phase == "stopped":
+        return "cancelled" if operation.cancellation_requested_at else "failed"
+    if phase == "failed" or (exit_code is not None and exit_code != 0):
+        return "failed"
+    return None
+
+
+def reconcile_runtime_operation(
+    db: Session,
+    *,
+    workflow: WorkflowSession,
+    operation: WorkflowOperation,
+    snapshot: Optional[Dict[str, Any]],
+) -> Dict[str, Any]:
+    """Idempotently reconcile one worker snapshot into a durable operation.
+
+    The caller owns worker transport. This module intentionally accepts snapshots
+    only after verifying all durable correlators, preventing a process-global
+    worker state from completing an unrelated workflow command.
+    """
+    runtime_kind = runtime_kind_for_operation(operation)
+    if runtime_kind is None:
+        return {"operation": operation, "reconciled": False, "reason": "unsupported"}
+    if operation.status in TERMINAL_OPERATION_STATUSES:
+        return {"operation": operation, "reconciled": False, "reason": "terminal"}
+    if not isinstance(snapshot, dict):
+        return {
+            "operation": operation,
+            "reconciled": False,
+            "reason": "invalid_snapshot",
+        }
+    if not _is_correlated(operation, workflow, snapshot):
+        return {
+            "operation": operation,
+            "reconciled": False,
+            "reason": "correlation_mismatch",
+        }
+
+    terminal_status = _terminal_status(operation, snapshot)
+    if terminal_status is None:
+        return {
+            "operation": operation,
+            "reconciled": False,
+            "reason": "runtime_not_terminal",
+        }
+
+    details = _runtime_details(snapshot)
+    metadata = _snapshot_metadata(snapshot)
+    operation_metadata = decode_json(operation.metadata_json)
+    run_id = _metadata_value(operation_metadata, "run_id", "runId")
+    output_directory = _metadata_value(metadata, "outputPath", "output_path")
+    checkpoint_path = _metadata_value(
+        metadata, "checkpointPath", "latestCheckpointPath", "checkpoint"
+    )
+    prediction_path = _metadata_value(
+        metadata,
+        "predictionPath",
+        "latestPredictionPath",
+        "outputPredictionPath",
+    )
+    output_path = prediction_path if runtime_kind == "inference" else output_directory
+    event_suffix = {
+        "succeeded": "completed",
+        "failed": "failed",
+        "cancelled": "cancelled",
+    }[terminal_status]
+    event_type = f"{runtime_kind}.{event_suffix}"
+    event_payload = {
+        "source": "runtime_reconciliation",
+        "operation_id": operation.id,
+        "command_id": operation.command_id,
+        "run_id": run_id,
+        "outputPath": output_path,
+        "outputDirectory": output_directory,
+        "checkpointPath": checkpoint_path,
+        "predictionPath": prediction_path,
+        **details,
+    }
+    error_payload = None
+    if terminal_status == "failed":
+        error_payload = {
+            "error": (
+                "RuntimeStopped"
+                if snapshot.get("phase") == "stopped"
+                else "RuntimeFailed"
+            ),
+            "detail": snapshot.get("lastError")
+            or f"{runtime_kind} runtime ended with phase {snapshot.get('phase')!r}",
+            "exit_code": snapshot.get("exitCode"),
+        }
+    terminal_event_payload = (
+        event_payload
+        if terminal_status == "succeeded"
+        else {**event_payload, **(error_payload or {})}
+    )
+
+    operation = transition_workflow_operation(
+        db,
+        operation,
+        status=terminal_status,
+        expected_status=operation.status,
+        result_payload=event_payload if terminal_status == "succeeded" else None,
+        error_payload=error_payload,
+        metadata={"runtime_terminal": details},
+        lease_owner=operation.lease_owner,
+        commit=False,
+    )
+    updates: Dict[str, Any] = {}
+    if runtime_kind == "training":
+        if output_directory:
+            updates["training_output_path"] = output_directory
+        if checkpoint_path:
+            updates["checkpoint_path"] = checkpoint_path
+    elif prediction_path:
+        updates["inference_output_path"] = prediction_path
+    if updates:
+        update_workflow_fields(db, workflow, updates, commit=False)
+
+    append_event_for_workflow_if_present(
+        db,
+        workflow_id=workflow.id,
+        actor="system",
+        event_type=event_type,
+        stage=workflow.stage,
+        summary=f"Synchronized {terminal_status} {runtime_kind} runtime.",
+        payload=terminal_event_payload,
+        idempotency_key=f"workflow-operation:{operation.id}:runtime-terminal",
+    )
+    db.refresh(operation)
+    return {"operation": operation, "reconciled": True, "reason": "terminal"}
+
+
+def reconciliation_response(result: Dict[str, Any]) -> Dict[str, Any]:
+    return {
+        "operation": operation_to_dict(result["operation"]),
+        "reconciled": bool(result["reconciled"]),
+        "reason": result["reason"],
+    }

server_api/workflows/service.py

@@ -819,12 +819,16 @@ def create_or_update_model_run_from_event(
         run_type, status = "training", "completed"
     elif event_type == "training.failed":
         run_type, status = "training", "failed"
+    elif event_type == "training.cancelled":
+        run_type, status = "training", "cancelled"
     elif event_type == "inference.started":
         run_type, status = "inference", "running"
     elif event_type == "inference.completed":
         run_type, status = "inference", "completed"
     elif event_type == "inference.failed":
         run_type, status = "inference", "failed"
+    elif event_type == "inference.cancelled":
+        run_type, status = "inference", "cancelled"
     if not run_type:
         return None
 
@@ -853,7 +857,7 @@ def create_or_update_model_run_from_event(
             output_path=output_path,
             fallback_latest=not bool(run_id),
         )
-        if status in {"completed", "failed"}
+        if status in {"completed", "failed", "cancelled"}
         else None
     )
     if status == "running" and run_id:
@@ -901,7 +905,7 @@ def create_or_update_model_run_from_event(
     )
     if status == "running" and not run.started_at:
         run.started_at = now
-    if status in {"completed", "failed"}:
+    if status in {"completed", "failed", "cancelled"}:
         run.completed_at = now
 
     if output_path and run_type == "inference":