From c50ccc5a679528f91b6b825fd387049234f286f3 Mon Sep 17 00:00:00 2001 From: Cizz22 Date: Thu, 27 Aug 2026 16:37:43 +0700 Subject: [PATCH] feat: Refactor task queue handling and enhance error messaging in simulation workflows --- run_worker.py | 3 +- src/aeros_simulation/ramp_engine.py | 126 ++++++++++++++++++++++++++-- src/aeros_simulation/router.py | 8 +- src/config.py | 2 + temporal/workflow.py | 12 ++- tests/unit/test_ramp_engine.py | 18 ++++ 6 files changed, 156 insertions(+), 13 deletions(-) diff --git a/run_worker.py b/run_worker.py index 125c32b..e393d32 100644 --- a/run_worker.py +++ b/run_worker.py @@ -5,6 +5,7 @@ from temporalio.worker import Worker from temporal.activity import calculate_plant_eaf_activity, collect_rbd_reliability_for_ramp_activity, execute_ramp_simulation_activity, execute_simulation_activity, get_simulation_engine_activity, update_contribution_bulk_mappings_activity, update_equipment_for_simulation_activity, call_callback_ahm, acquire_lock, release_lock,execute_update_node, handle_workflow_failure from temporal.workflow import SimulationWorkflow +from src.config import SIMULATION_TASK_QUEUE TEMPORAL_URL = os.environ.get("TEMPORAL_URL", "http://192.168.1.86:7233") @@ -14,7 +15,7 @@ async def main(): try: worker = Worker( client, - task_queue="simulation-task-queue", + task_queue=SIMULATION_TASK_QUEUE, workflows=[SimulationWorkflow], activities=[ update_equipment_for_simulation_activity, diff --git a/src/aeros_simulation/ramp_engine.py b/src/aeros_simulation/ramp_engine.py index 5b4cad7..2e19f84 100644 --- a/src/aeros_simulation/ramp_engine.py +++ b/src/aeros_simulation/ramp_engine.py @@ -16,6 +16,7 @@ from src.config import ( RAMP_MODEL_NAME, RAMP_POLL_INTERVAL_SECONDS, RAMP_POLL_TIMEOUT_SECONDS, + RAMP_STALE_TIMEOUT_SECONDS, RAMP_TARGET_SCHEMA, SIMULATION_ENGINE, ) @@ -255,9 +256,7 @@ def _rbd_reliability_to_ramp_override(reliability: dict) -> dict: "distribution": _clean_dict({"type": "Weibull", "shape": beta, "scale": eta}), } elif distribution == "Exponential": - exp_eta = eta - if beta and beta > 0: - exp_eta = 1 / beta + exp_eta = eta if eta is not None else beta override["model"] = { "type": "NonRepairable", "distribution": _clean_dict({"type": "Exponential", "eta": exp_eta}), @@ -374,6 +373,51 @@ async def _fetch_ramp_model_asset_names(client: httpx.AsyncClient, model_id: str return {_norm_key(asset.get("name")) for asset in assets if asset.get("name")} +def _find_ramp_task(tasks_payload: dict, task_id: str) -> dict | None: + for task in tasks_payload.get("tasks", []): + if str(task.get("id")) == str(task_id): + return task + return None + + +async def _fetch_ramp_task_status(client: httpx.AsyncClient, task_id: str) -> dict: + """Fetch task status, falling back to the lighter task-list endpoint. + + RAMP's status endpoint includes the result payload for completed tasks. When + the engine is busy that endpoint can occasionally time out even for running + tasks, while /tasks still returns the persisted progress row quickly. + """ + try: + status_response = await client.get(f"{_simulation_api_base()}/status/{task_id}") + status_response.raise_for_status() + return status_response.json() + except (httpx.ReadTimeout, httpx.ConnectTimeout, httpx.ConnectError) as status_error: + log.warning("RAMP status request failed for task %s: %s", task_id, status_error) + + tasks_response = await client.get(f"{_simulation_api_base()}/tasks") + tasks_response.raise_for_status() + task = _find_ramp_task(tasks_response.json(), task_id) + if task: + return task + raise RuntimeError(f"RAMP task {task_id} was not found in /tasks fallback.") + + +async def _fetch_completed_ramp_result(client: httpx.AsyncClient, task_id: str) -> dict: + response = await client.get(f"{_simulation_api_base()}/status/{task_id}", timeout=120.0) + response.raise_for_status() + result = response.json().get("result") + if not isinstance(result, dict): + raise RuntimeError(f"RAMP task {task_id} completed without a result payload.") + return result + + +async def _cancel_ramp_task(client: httpx.AsyncClient, task_id: str) -> None: + try: + await client.post(f"{_simulation_api_base()}/cancel/{task_id}") + except Exception as cancel_error: + log.warning("Failed to cancel stale RAMP task %s: %s", task_id, cancel_error) + + def _filter_overrides_to_model(overrides: dict, asset_names: set[str]) -> tuple[dict, list[str]]: """Keep only overrides whose asset exists in the RAMP model. @@ -612,10 +656,36 @@ async def save_ramp_simulation_result( ) simulation.status = "completed" + simulation.error = None simulation.completed_at = datetime.now() await db_session.commit() +async def record_ramp_task_start( + *, + db_session: DbSession, + simulation_id: UUID, + model_id: str, + ramp_task_id: str, + payload: dict, +) -> None: + simulation = await get_simulation_by_id(db_session=db_session, simulation_id=simulation_id) + simulation.status = "running" + simulation.input = { + **(simulation.input or {}), + "engine": "ramp", + "ramp_model_id": model_id, + "ramp_task_id": ramp_task_id, + "ramp_payload": { + key: value + for key, value in payload.items() + if key != "distribution_overrides" + }, + "ramp_distribution_override_count": len(payload.get("distribution_overrides") or {}), + } + await db_session.commit() + + async def execute_ramp_simulation( *, db_session: DbSession, @@ -653,19 +723,60 @@ async def execute_ramp_simulation( task_id = response.json().get("task_id") if not task_id: raise RuntimeError(f"RAMP execute response did not include task_id: {response.text}") + await record_ramp_task_start( + db_session=db_session, + simulation_id=simulation_id, + model_id=model_id, + ramp_task_id=task_id, + payload=payload, + ) deadline = time.monotonic() + RAMP_POLL_TIMEOUT_SECONDS last_status = None + last_progress = None + last_updated_at = None + last_change_at = time.monotonic() + last_successful_poll_at = time.monotonic() while time.monotonic() < deadline: - status_response = await client.get(f"{_simulation_api_base()}/status/{task_id}") - status_response.raise_for_status() - status_payload = status_response.json() + try: + status_payload = await _fetch_ramp_task_status(client, task_id) + last_successful_poll_at = time.monotonic() + except httpx.HTTPError as poll_error: + if time.monotonic() - last_successful_poll_at > RAMP_STALE_TIMEOUT_SECONDS: + await _cancel_ramp_task(client, task_id) + raise TimeoutError( + f"RAMP task {task_id} status polling failed for " + f"{RAMP_STALE_TIMEOUT_SECONDS}s: {poll_error}" + ) from poll_error + log.warning("Transient RAMP poll failure for task %s: %s", task_id, poll_error) + await asyncio.sleep(RAMP_POLL_INTERVAL_SECONDS) + continue + + previous_marker = (last_status, last_progress, last_updated_at) last_status = status_payload.get("status") + progress = status_payload.get("progress") + updated_at = status_payload.get("updated_at") + + if (last_status, progress, updated_at) != previous_marker: + last_change_at = time.monotonic() + last_progress = progress + last_updated_at = updated_at + + if ( + last_status == "running" + and RAMP_STALE_TIMEOUT_SECONDS > 0 + and time.monotonic() - last_change_at > RAMP_STALE_TIMEOUT_SECONDS + ): + await _cancel_ramp_task(client, task_id) + raise TimeoutError( + f"RAMP task {task_id} made no progress for " + f"{RAMP_STALE_TIMEOUT_SECONDS}s. Last progress: {progress}." + ) if last_status == "completed": result = status_payload.get("result") if not isinstance(result, dict): - raise RuntimeError(f"RAMP task {task_id} completed without a result payload.") + result = await _fetch_completed_ramp_result(client, task_id) await save_ramp_simulation_result( db_session=db_session, simulation_id=simulation_id, @@ -682,4 +793,5 @@ async def execute_ramp_simulation( await asyncio.sleep(RAMP_POLL_INTERVAL_SECONDS) + await _cancel_ramp_task(client, task_id) raise TimeoutError(f"RAMP task timed out after {RAMP_POLL_TIMEOUT_SECONDS}s. Last status: {last_status}") diff --git a/src/aeros_simulation/router.py b/src/aeros_simulation/router.py index 5ff8d0a..e1b71ee 100644 --- a/src/aeros_simulation/router.py +++ b/src/aeros_simulation/router.py @@ -15,7 +15,7 @@ from src.aeros_equipment.model import AerosEquipment from src.aeros_simulation.model import AerosSimulationCalcResult, EafContribution, AerosNode from src.aeros_simulation.utils import date_to_utc, hours_between, year_window_utc from src.auth.service import CurrentUser -from src.config import TEMPORAL_URL +from src.config import SIMULATION_TASK_QUEUE, TEMPORAL_URL from src.database.core import CollectorDbSession, DbSession from src.database.service import CommonParameters, get_params_factory from src.models import StandardResponse @@ -108,7 +108,7 @@ async def run_simulations( SimulationWorkflow.run, sim_data, id=f"simulation-{simulation_id}", - task_queue="simulation-task-queue", + task_queue=SIMULATION_TASK_QUEUE, ) return { @@ -215,7 +215,7 @@ async def run_yearly_simulation( SimulationWorkflow.run, sim_data, id=f"simulation-{simulation.id}", - task_queue="simulation-task-queue", + task_queue=SIMULATION_TASK_QUEUE, ) return { @@ -582,7 +582,7 @@ async def get_forecast_eaf( SimulationWorkflow.run, sim_data, id=f"simulation-{simulation_id}", - task_queue="simulation-task-queue", + task_queue=SIMULATION_TASK_QUEUE, ) forecast_msg = ( diff --git a/src/config.py b/src/config.py index a0642de..ac5fb3a 100644 --- a/src/config.py +++ b/src/config.py @@ -94,6 +94,7 @@ WINDOWS_AEROS_BASE_URL = config("WINDOWS_AEROS_BASE_URL", default="http://192.16 DEFAULT_PROJECT_NAME = config("DEFAULT_PROJECT_NAME", default="RBD TJB") TEMPORAL_URL = config("TEMPORAL_URL", default="http://192.168.1.86:7233") +SIMULATION_TASK_QUEUE = config("SIMULATION_TASK_QUEUE", default="simulation-task-queue") RELIABILITY_SERVICE_API = config("RELIABILITY_SERVICE_API", default="http://192.168.1.82:8000/reliability") CLAMAV_HOST = config("CLAMAV_HOST", default="192.168.1.82") @@ -109,6 +110,7 @@ RAMP_MODEL_NAME = config("RAMP_MODEL_NAME", default=None) RAMP_TARGET_SCHEMA = config("RAMP_TARGET_SCHEMA", default=None) RAMP_POLL_INTERVAL_SECONDS = config("RAMP_POLL_INTERVAL_SECONDS", cast=int, default=10) RAMP_POLL_TIMEOUT_SECONDS = config("RAMP_POLL_TIMEOUT_SECONDS", cast=int, default=7200) +RAMP_STALE_TIMEOUT_SECONDS = config("RAMP_STALE_TIMEOUT_SECONDS", cast=int, default=1800) VAULT_URL=config('VAULT_URL', default=None) ROLE_ID=config('ROLE_ID', default=None) diff --git a/temporal/workflow.py b/temporal/workflow.py index 72f7120..a961335 100644 --- a/temporal/workflow.py +++ b/temporal/workflow.py @@ -16,6 +16,16 @@ _RETRY_POLICY = RetryPolicy( ) _NO_RETRY_POLICY = RetryPolicy(maximum_attempts=1) +def _workflow_error_message(error: Exception) -> str: + messages = [] + current = error + while current: + message = str(current) + if message and message not in messages: + messages.append(message) + current = getattr(current, "__cause__", None) + return " | caused by: ".join(messages) or error.__class__.__name__ + @workflow.defn class SimulationWorkflow: def __init__(self): @@ -141,7 +151,7 @@ class SimulationWorkflow: { "simulation_id": sim_data["HubCnnId"], "stage": failed_stage, - "error": str(e), + "error": _workflow_error_message(e), "ahm_job_id": sim_data.get("AhmJobId"), }, start_to_close_timeout=timedelta(seconds=30), diff --git a/tests/unit/test_ramp_engine.py b/tests/unit/test_ramp_engine.py index 8b840a6..2ddc489 100644 --- a/tests/unit/test_ramp_engine.py +++ b/tests/unit/test_ramp_engine.py @@ -117,6 +117,24 @@ def test_build_ramp_distribution_overrides_from_rbd_reliability(): } +def test_build_ramp_distribution_overrides_uses_rbd_eta_for_exponential(): + overrides = build_ramp_distribution_overrides_from_rbd_reliability( + { + "HEATER-A": { + "distribution": "Exponential2", + "beta": 13034.0, + "eta": 8760.0, + "mttr": 6, + } + } + ) + + assert overrides["HEATER-A"]["model"]["distribution"] == { + "type": "Exponential", + "eta": 8760.0, + } + + def test_filter_overrides_keeps_only_assets_present_in_model(): overrides = { "3AH-AU501A": {"repair_duration": 4},