diff --git a/server_api/main.py b/server_api/main.py index b0f5997e..a32a3564 100644 --- a/server_api/main.py +++ b/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) diff --git a/server_api/workflows/runtime_reconciliation.py b/server_api/workflows/runtime_reconciliation.py new file mode 100644 index 00000000..9b8bb2db --- /dev/null +++ b/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"], + } diff --git a/server_api/workflows/service.py b/server_api/workflows/service.py index 0bd16910..2c2a8bf9 100644 --- a/server_api/workflows/service.py +++ b/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": diff --git a/server_pytc/services/model.py b/server_pytc/services/model.py index 7edf2408..52d52730 100644 --- a/server_pytc/services/model.py +++ b/server_pytc/services/model.py @@ -324,6 +324,12 @@ def _set_runtime_error(kind: str, message: str): ) +def _terminal_runtime_error(kind: str, exit_code: int | None) -> str | None: + if exit_code in (None, 0): + return None + return f"{kind.capitalize()} subprocess exited with code {exit_code}" + + def _get_runtime_snapshot(kind: str) -> dict[str, Any]: process = _get_runtime_process(kind) is_running = bool(process and process.poll() is None) @@ -340,6 +346,12 @@ def _get_runtime_snapshot(kind: str) -> dict[str, Any]: phase = "finished" if rc == 0 else "failed" exit_code = rc ended_at = state["endedAt"] or _utc_now() + state["phase"] = phase + state["exitCode"] = exit_code + state["endedAt"] = ended_at + if not state["lastError"]: + state["lastError"] = _terminal_runtime_error(kind, exit_code) + state["lastUpdatedAt"] = _utc_now() lines = list(state["lines"]) snapshot = { @@ -1543,12 +1555,16 @@ def _log_subprocess_output(): level="WARNING", output_path=output_path, ) - _update_runtime_state( - kind, - phase="finished" if exit_code == 0 else "failed", - exitCode=exit_code, - endedAt=_utc_now(), - ) + terminal_updates = { + "phase": "finished" if exit_code == 0 else "failed", + "exitCode": exit_code, + "endedAt": _utc_now(), + } + if exit_code != 0: + terminal_updates["lastError"] = _get_runtime_snapshot(kind).get( + "lastError" + ) or _terminal_runtime_error(kind, exit_code) + _update_runtime_state(kind, **terminal_updates) _append_runtime_event( kind, f"{label} subprocess finished with exit code: {exit_code}", @@ -1801,8 +1817,11 @@ def start_training(payload: dict): "inputLabelPath": payload.get("inputLabelPath"), "configOriginPath": config_origin_path, "workflowId": payload.get("workflowId") or payload.get("workflow_id"), + "workflow_id": payload.get("workflow_id") or payload.get("workflowId"), "runId": payload.get("runId") or payload.get("run_id"), + "run_id": payload.get("run_id") or payload.get("runId"), "commandId": payload.get("command_id") or payload.get("commandId"), + "command_id": payload.get("command_id") or payload.get("commandId"), "autoParameters": auto_parameters, }, ) @@ -2167,6 +2186,11 @@ def start_inference(payload: dict): or (payload.get("arguments") or {}).get("checkpoint"), "configOriginPath": config_origin_path, "workflowId": payload.get("workflow_id") or payload.get("workflowId"), + "workflow_id": payload.get("workflow_id") or payload.get("workflowId"), + "runId": payload.get("run_id") or payload.get("runId"), + "run_id": payload.get("run_id") or payload.get("runId"), + "commandId": payload.get("command_id") or payload.get("commandId"), + "command_id": payload.get("command_id") or payload.get("commandId"), }, ) diff --git a/tests/test_pytc_runtime_routes.py b/tests/test_pytc_runtime_routes.py index 1a30ee17..269adc7f 100644 --- a/tests/test_pytc_runtime_routes.py +++ b/tests/test_pytc_runtime_routes.py @@ -23,7 +23,7 @@ _resolve_raw_image_shader, ) from server_api.main import _coerce_neuroglancer_scales -from server_api.workflows.db_models import WorkflowCommand +from server_api.workflows.db_models import WorkflowCommand, WorkflowOperation from server_api.workflows.service import encode_json from server_pytc.main import app as server_pytc_app from server_pytc.services import model as model_service @@ -344,6 +344,7 @@ def override_get_db(): server_api_app.dependency_overrides[auth_database.get_db] = override_get_db self.client = TestClient(server_api_app) + self._runtime_operation_sequence = 0 def tearDown(self): server_api_app.dependency_overrides.clear() @@ -355,6 +356,192 @@ def _workflow_id(self): self.assertEqual(response.status_code, 200) return response.json()["workflow"]["id"] + def _create_running_runtime_operation(self, runtime_kind): + """Create the durable records produced after a worker accepts a command.""" + self.assertIn(runtime_kind, {"training", "inference"}) + workflow_id = self._workflow_id() + self._runtime_operation_sequence += 1 + suffix = self._runtime_operation_sequence + run_id = f"workflow-command-{runtime_kind}-test-{suffix}" + db = self.SessionLocal() + try: + command = WorkflowCommand( + workflow_id=workflow_id, + command_type=f"start_{runtime_kind}", + status="submitted", + idempotency_key=f"test:{runtime_kind}:reconcile-command:{suffix}", + actor="agent", + input_json=encode_json({"run_id": run_id}), + attempt_count=1, + ) + db.add(command) + db.flush() + operation = WorkflowOperation( + workflow_id=workflow_id, + command_id=command.id, + operation_type=f"start_{runtime_kind}", + status="running", + idempotency_key=f"test:{runtime_kind}:reconcile-operation:{suffix}", + correlation_id=f"test:{runtime_kind}:reconcile:{suffix}", + actor="agent", + metadata_json=encode_json({"run_id": run_id}), + attempt_count=1, + lease_owner=f"server_api.{runtime_kind}_runner", + ) + db.add(operation) + db.commit() + return workflow_id, command.id, operation.id, run_id + finally: + db.close() + + @staticmethod + def _runtime_snapshot( + *, + workflow_id, + command_id, + run_id, + phase, + exit_code=None, + last_error=None, + ): + return { + "phase": phase, + "pid": 4242, + "exitCode": exit_code, + "startedAt": "2026-08-04T12:00:00+00:00", + "endedAt": "2026-08-04T12:01:00+00:00", + "lastError": last_error, + "lineCount": 12, + "metadata": { + "workflowId": workflow_id, + "commandId": command_id, + "runId": run_id, + "outputPath": "/tmp/runtime-output", + "checkpointPath": "/tmp/checkpoint.pth.tar", + "predictionPath": "/tmp/runtime-output/prediction.h5", + }, + } + + def _reconcile_runtime(self, workflow_id, operation_id, snapshot): + with patch("server_api.main._proxy_to_worker", return_value=snapshot): + return self.client.post( + f"/api/workflows/{workflow_id}/operations/{operation_id}/reconcile-runtime" + ) + + def test_reconcile_correlated_terminal_runtime_completes_running_operations(self): + for runtime_kind in ("training", "inference"): + with self.subTest(runtime_kind=runtime_kind): + workflow_id, command_id, operation_id, run_id = ( + self._create_running_runtime_operation(runtime_kind) + ) + response = self._reconcile_runtime( + workflow_id, + operation_id, + self._runtime_snapshot( + workflow_id=workflow_id, + command_id=command_id, + run_id=run_id, + phase="finished", + exit_code=0, + ), + ) + + self.assertEqual(response.status_code, 200) + payload = response.json() + self.assertTrue(payload["reconciled"]) + self.assertEqual(payload["reason"], "terminal") + self.assertEqual(payload["operation"]["status"], "succeeded") + self.assertEqual(payload["operation"]["result"]["run_id"], run_id) + + def test_reconcile_refuses_uncorrelated_snapshot_and_is_idempotent(self): + workflow_id, command_id, operation_id, run_id = ( + self._create_running_runtime_operation("inference") + ) + mismatch = self._reconcile_runtime( + workflow_id, + operation_id, + self._runtime_snapshot( + workflow_id=workflow_id, + command_id=command_id + 1, + run_id=run_id, + phase="finished", + exit_code=0, + ), + ) + self.assertEqual(mismatch.status_code, 200) + self.assertFalse(mismatch.json()["reconciled"]) + self.assertEqual(mismatch.json()["reason"], "correlation_mismatch") + self.assertEqual(mismatch.json()["operation"]["status"], "running") + + terminal_snapshot = self._runtime_snapshot( + workflow_id=workflow_id, + command_id=command_id, + run_id=run_id, + phase="finished", + exit_code=0, + ) + first = self._reconcile_runtime(workflow_id, operation_id, terminal_snapshot) + second = self._reconcile_runtime(workflow_id, operation_id, terminal_snapshot) + self.assertTrue(first.json()["reconciled"]) + self.assertFalse(second.json()["reconciled"]) + self.assertEqual(second.json()["reason"], "terminal") + + events = self.client.get(f"/api/workflows/{workflow_id}/events").json() + terminal_events = [ + event + for event in events + if event["event_type"] == "inference.completed" + and event["payload"].get("operation_id") == operation_id + ] + self.assertEqual(len(terminal_events), 1) + + def test_reconcile_failed_and_stopped_runtime_preserves_terminal_semantics(self): + workflow_id, command_id, operation_id, run_id = ( + self._create_running_runtime_operation("training") + ) + failure = self._reconcile_runtime( + workflow_id, + operation_id, + self._runtime_snapshot( + workflow_id=workflow_id, + command_id=command_id, + run_id=run_id, + phase="failed", + exit_code=1, + last_error="worker crashed", + ), + ) + self.assertEqual(failure.status_code, 200) + self.assertEqual(failure.json()["operation"]["status"], "failed") + self.assertEqual( + failure.json()["operation"]["error"]["detail"], "worker crashed" + ) + + for requested, expected_status in ((True, "cancelled"), (False, "failed")): + with self.subTest(cancellation_requested=requested): + workflow_id, command_id, operation_id, run_id = ( + self._create_running_runtime_operation("inference") + ) + if requested: + cancel = self.client.post( + f"/api/workflows/{workflow_id}/operations/{operation_id}/cancel", + json={"reason": "user requested stop"}, + ) + self.assertEqual(cancel.status_code, 200) + self.assertEqual(cancel.json()["status"], "running") + stopped = self._reconcile_runtime( + workflow_id, + operation_id, + self._runtime_snapshot( + workflow_id=workflow_id, + command_id=command_id, + run_id=run_id, + phase="stopped", + ), + ) + self.assertEqual(stopped.status_code, 200) + self.assertEqual(stopped.json()["operation"]["status"], expected_status) + def test_sync_completed_inference_runtime_materializes_prediction_run(self): workflow_id = self._workflow_id() output_dir = pathlib.Path(self.temp_dir.name) / "inference-output" @@ -542,14 +729,14 @@ def fake_worker(method, endpoint, json_body=None, **_kwargs): payload = run_response.json() self.assertEqual(payload["command"]["status"], "submitted") self.assertEqual(payload["command"]["attempt_count"], 1) - self.assertEqual(payload["operation"]["status"], "succeeded") + self.assertEqual(payload["operation"]["status"], "running") self.assertEqual(payload["operation"]["operation_type"], "start_training") self.assertEqual(payload["operation"]["command_id"], command["id"]) self.assertEqual( payload["operation"]["idempotency_key"], f"workflow-command:{command['id']}:attempt:1", ) - self.assertEqual(payload["operation"]["result"]["worker"]["pid"], 4242) + self.assertEqual(payload["operation"]["metadata"]["worker"]["pid"], 4242) self.assertEqual( duplicate_run_response.json()["operation"]["id"], payload["operation"]["id"], @@ -581,7 +768,7 @@ def fake_worker(method, endpoint, json_body=None, **_kwargs): ) self.assertEqual(operations_response.status_code, 200) self.assertEqual(len(operations_response.json()), 1) - self.assertEqual(operations_response.json()[0]["status"], "succeeded") + self.assertEqual(operations_response.json()[0]["status"], "running") events_response = self.client.get(f"/api/workflows/{workflow_id}/events") self.assertEqual(events_response.status_code, 200) @@ -658,7 +845,7 @@ def fake_worker(method, endpoint, json_body=None, **_kwargs): self.assertEqual(payload["command"]["status"], "submitted") self.assertEqual(payload["command"]["attempt_count"], 1) self.assertEqual(payload["operation"]["operation_type"], "start_inference") - self.assertEqual(payload["operation"]["status"], "succeeded") + self.assertEqual(payload["operation"]["status"], "running") self.assertEqual(payload["operation"]["command_id"], command_id) self.assertEqual( payload["operation"]["idempotency_key"], @@ -764,7 +951,7 @@ def test_durable_inference_command_retry_uses_a_new_operation_attempt(self): ) self.assertEqual(retry_response.status_code, 200) - self.assertEqual(retry_response.json()["operation"]["status"], "succeeded") + self.assertEqual(retry_response.json()["operation"]["status"], "running") self.assertEqual( retry_response.json()["operation"]["idempotency_key"], f"workflow-command:{command_id}:attempt:2", @@ -841,7 +1028,7 @@ def test_durable_training_command_failure_records_retryable_operation(self): ) self.assertEqual(retry_response.status_code, 200) - self.assertEqual(retry_response.json()["operation"]["status"], "succeeded") + self.assertEqual(retry_response.json()["operation"]["status"], "running") self.assertEqual( retry_response.json()["operation"]["idempotency_key"], f"workflow-command:{command_id}:attempt:2", diff --git a/tests/test_worker_model_service.py b/tests/test_worker_model_service.py index 930f7727..2e9c1a15 100644 --- a/tests/test_worker_model_service.py +++ b/tests/test_worker_model_service.py @@ -110,6 +110,46 @@ def test_ollama_unload_can_be_enabled_explicitly(self): self.assertEqual(unloaded, ["qwen3.6:27b"]) run_mock.assert_called_once() + def test_terminal_inference_snapshot_retains_correlators_and_artifacts(self): + with tempfile.TemporaryDirectory() as tmpdir: + output_dir = pathlib.Path(tmpdir) + prediction_path = output_dir / "result_xy.h5" + prediction_path.write_text("prediction", encoding="utf-8") + model_service._reset_runtime_state( + "inference", + phase="failed", + metadata={ + "workflowId": 17, + "workflow_id": 17, + "commandId": 23, + "command_id": 23, + "runId": "workflow-command-23", + "run_id": "workflow-command-23", + "outputPath": str(output_dir), + }, + ) + model_service._update_runtime_state( + "inference", + exitCode=2, + endedAt="2026-08-04T12:00:00+00:00", + lastError=model_service._terminal_runtime_error("inference", 2), + ) + + snapshot = model_service.get_inference_logs() + + self.assertEqual(snapshot["phase"], "failed") + self.assertEqual(snapshot["exitCode"], 2) + self.assertIn("exited with code 2", snapshot["lastError"]) + self.assertEqual(snapshot["metadata"]["workflowId"], 17) + self.assertEqual(snapshot["metadata"]["workflow_id"], 17) + self.assertEqual(snapshot["metadata"]["commandId"], 23) + self.assertEqual(snapshot["metadata"]["command_id"], 23) + self.assertEqual(snapshot["metadata"]["runId"], "workflow-command-23") + self.assertEqual(snapshot["metadata"]["run_id"], "workflow-command-23") + self.assertEqual( + snapshot["metadata"]["predictionPath"], str(prediction_path.resolve()) + ) + if __name__ == "__main__": unittest.main()