mirror of
https://github.com/rowboatlabs/rowboat.git
synced 2026-04-25 16:36:22 +02:00
120 lines
4.2 KiB
Python
120 lines
4.2 KiB
Python
import asyncio
|
||
import logging
|
||
from typing import List, Optional
|
||
|
||
# Updated imports from your new db module and scenario_types
|
||
from db import (
|
||
get_pending_run,
|
||
get_simulations_for_run,
|
||
set_run_to_completed,
|
||
get_api_key,
|
||
mark_stale_jobs_as_failed,
|
||
update_run_heartbeat
|
||
)
|
||
from scenario_types import TestRun, TestSimulation
|
||
# If you have a new simulation function, import it here.
|
||
# Otherwise, adapt the name as needed:
|
||
from simulation import simulate_simulations # or simulate_scenarios, if unchanged
|
||
|
||
logging.basicConfig(level=logging.INFO)
|
||
|
||
class JobService:
|
||
def __init__(self):
|
||
self.poll_interval = 5 # seconds
|
||
# Control concurrency of run processing
|
||
self.semaphore = asyncio.Semaphore(5)
|
||
|
||
async def poll_and_process_jobs(self, max_iterations: Optional[int] = None):
|
||
"""
|
||
Periodically checks for new runs in MongoDB and processes them.
|
||
"""
|
||
# Start the stale-run check in the background
|
||
asyncio.create_task(self.fail_stale_runs_loop())
|
||
|
||
iterations = 0
|
||
while True:
|
||
run = get_pending_run() # <--- changed to match new DB function
|
||
if run:
|
||
logging.info(f"Found new run: {run}. Processing...")
|
||
asyncio.create_task(self.process_run(run))
|
||
|
||
iterations += 1
|
||
if max_iterations is not None and iterations >= max_iterations:
|
||
break
|
||
|
||
# Sleep for the polling interval
|
||
await asyncio.sleep(self.poll_interval)
|
||
|
||
async def process_run(self, run: TestRun):
|
||
"""
|
||
Calls the simulation function and updates run status upon completion.
|
||
"""
|
||
async with self.semaphore:
|
||
# Start heartbeat in background
|
||
stop_heartbeat_event = asyncio.Event()
|
||
heartbeat_task = asyncio.create_task(self.heartbeat_loop(run.id, stop_heartbeat_event))
|
||
|
||
try:
|
||
# Fetch the simulations associated with this run
|
||
simulations = get_simulations_for_run(run)
|
||
if not simulations:
|
||
logging.info(f"No simulations found for run {run.id}")
|
||
return
|
||
|
||
# Fetch API key if needed
|
||
api_key = get_api_key(run.projectId)
|
||
|
||
# Perform your simulation logic
|
||
# adapt this call to your actual simulation function’s signature
|
||
aggregate_result = await simulate_simulations(
|
||
simulations=simulations,
|
||
run_id=run.id,
|
||
workflow_id=run.workflowId,
|
||
api_key=api_key
|
||
)
|
||
|
||
# Mark run as completed with the aggregated result
|
||
set_run_to_completed(run, aggregate_result)
|
||
logging.info(f"Run {run.id} completed.")
|
||
except Exception as exc:
|
||
logging.error(f"Run {run.id} failed: {exc}")
|
||
finally:
|
||
stop_heartbeat_event.set()
|
||
await heartbeat_task
|
||
|
||
async def fail_stale_runs_loop(self):
|
||
"""
|
||
Periodically checks for stale runs (no heartbeat) and marks them as 'failed'.
|
||
"""
|
||
while True:
|
||
count = mark_stale_jobs_as_failed()
|
||
if count > 0:
|
||
logging.warning(f"Marked {count} stale runs as failed.")
|
||
await asyncio.sleep(60) # Check every 60 seconds
|
||
|
||
async def heartbeat_loop(self, run_id: str, stop_event: asyncio.Event):
|
||
"""
|
||
Periodically updates 'lastHeartbeat' for the given run until 'stop_event' is set.
|
||
"""
|
||
try:
|
||
while not stop_event.is_set():
|
||
update_run_heartbeat(run_id)
|
||
await asyncio.sleep(10) # Heartbeat interval in seconds
|
||
except asyncio.CancelledError:
|
||
pass
|
||
|
||
def start(self):
|
||
"""
|
||
Entry point to start the service event loop.
|
||
"""
|
||
loop = asyncio.get_event_loop()
|
||
try:
|
||
loop.run_until_complete(self.poll_and_process_jobs())
|
||
except KeyboardInterrupt:
|
||
logging.info("Service stopped by user.")
|
||
finally:
|
||
loop.close()
|
||
|
||
if __name__ == "__main__":
|
||
service = JobService()
|
||
service.start()
|