be-optimumoh/temporal/temporal_workflows.py

import asyncio
from datetime import timedelta
from temporalio import activity, workflow

# with workflow.unsafe.imports_passed_through():
#     import httpx
#     from src.config import RBD_SERVICE_API
#     from src.database.core import get_main_session, get_collector_session
#     from src.optimum_time_constraint.service import (
#         create_param_and_data, 
#         get_calculation_data_by_id
#     )
#     from src.optimum_time_constraint.optimizer import run_simulation_with_spareparts
#     from src.overhaul_scope.service import get as get_scope, get_prev_oh
#     from src.calculation_time_constrains.utils import get_months_between
#     from src.calculation_target_reliability.utils import wait_for_workflow

@activity.defn
async def create_optimum_oh_calculation(args: dict) -> str:
    from src.calculation_time_constrains.schema import CalculationTimeConstrainsParametersCreate
    from src.calculation_time_constrains.service import create_param_and_data
    from src.database.core import get_main_session

    token = args["token"]
    calc_in_dict = args["calculation_in"]
    created_by = args["created_by"]
    
    calc_in = CalculationTimeConstrainsParametersCreate(**calc_in_dict)
    
    async with get_main_session() as db_session:
        calc_data = await create_param_and_data(
            db_session=db_session,
            calculation_param_in=calc_in,
            created_by=created_by
        )
        return str(calc_data.id)

@activity.defn
async def request_rbd_simulation(args: dict) -> dict:
    from src.calculation_time_constrains.service import get_calculation_data_by_id
    from src.database.core import get_main_session
    from src.overhaul_scope.service import get as get_scope, get_prev_oh
    from src.calculation_time_constrains.utils import get_months_between
    from src.calculation_target_reliability.utils import wait_for_workflow
    import httpx
    from src.config import RBD_SERVICE_API

    calc_id = args["calc_id"]
    token = args["token"]
    callback_workflow_id = args["callback_workflow_id"]
    
    async with get_main_session() as db_session:
        calc_data = await get_calculation_data_by_id(db_session=db_session, calculation_id=calc_id)
        
        scope = await get_scope(db_session=db_session, overhaul_session_id=calc_data.overhaul_session_id)
        prev_oh_scope = await get_prev_oh(db_session=db_session, overhaul_session=scope)
        
        # Calculate time window
        time_window_months = get_months_between(prev_oh_scope.end_date, scope.start_date) + 6
        sim_duration_hours = time_window_months * 720
        
        sim_input = {
            "SchematicName": "- TJB - Unit 3 -",
            "SimulationName": f"OptimumOH_Calc_{calc_id}",
            "SimDuration": sim_duration_hours,
            "DurationUnit": "UHour",
            "OffSet": 0,
            "SimSeed": 99,
            "SimNumRun": 1,
            "OverhaulInterval": sim_duration_hours + 1,
            "MaintenanceOutages": 0,
            "IsDefault": False,
            "OverhaulDuration": 1200,
            "AhmJobId": None,
            "CallbackWorkflowId": callback_workflow_id,
        }
        
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.post(
                f"{RBD_SERVICE_API}/aeros/simulation/run",
                json=sim_input,
                headers={"Authorization": f"Bearer {token}"}
            )
            resp.raise_for_status()
            sim_id = resp.json()["data"]
            
            calc_data.rbd_simulation_id = sim_id
            return sim_id

# @activity.defn
# async def wait_for_rbd_simulation(sim_id: str) -> str:

#     await wait_for_workflow(sim_id)
#     return sim_id

@activity.defn
async def run_optimum_oh_calculation(args: dict) -> dict:
    from src.database.core import get_main_session, get_collector_session
    from src.calculation_time_constrains.service import get_calculation_data_by_id, run_simulation_with_spareparts
    from src.config import RBD_SERVICE_API
    
    calc_id = args["calc_id"]
    token = args["token"]
    sim_id = args["sim_id"]
    
    async with get_main_session() as db_session:
        async with get_collector_session() as collector_db:
            calc_data = await get_calculation_data_by_id(db_session=db_session, calculation_id=calc_id)
            results = await run_simulation_with_spareparts(
                db_session=db_session,
                calculation=calc_data,
                token=token,
                collector_db_session=collector_db,
                simulation_id=sim_id,
                rbd_service_api=RBD_SERVICE_API
            )
            # return simplified result since temporal handles JSON
            return {"id": str(results["id"]), "optimum": results["optimum"]}

@workflow.defn
class OptimumOHCalculationWorkflow:
    def __init__(self):
        self.done = False

    @workflow.signal
    def notify_done(self):
        self.done = True

    @workflow.run
    async def run(self, args: dict) -> dict:
        # 1. Create calculation
        calc_id = await workflow.execute_activity(
            create_optimum_oh_calculation,
            args,
            start_to_close_timeout=timedelta(minutes=1)
        )
        
        args["calc_id"] = calc_id
        
        # 2. Request RBD simulation
        sim_id = await workflow.execute_activity(
            request_rbd_simulation,
            args,
            start_to_close_timeout=timedelta(minutes=2)
        )
        
        args["sim_id"] = sim_id
        
        # Wait for RBD simulation to finish using signal
        await workflow.wait_condition(lambda: self.done)
        
        # 4. Run Optimum OH calculation
        result = await workflow.execute_activity(
            run_optimum_oh_calculation,
            args,
            start_to_close_timeout=timedelta(minutes=30)
        )
        
        return result