diff --git a/common/src/common/config.py b/common/src/common/config.py index f00c121..2e83dc1 100644 --- a/common/src/common/config.py +++ b/common/src/common/config.py @@ -112,6 +112,16 @@ class Settings(BaseSettings): description="Public base URL for the runtime container.", ) + # ── schedule execution tuning ──────────────────────────────── + schedule_execution_concurrency: int = Field( + default=4, + description=( + "Max concurrent notebooks running in the schedule worker. " + "Each notebook is dispatched as an asyncio task bounded by " + "a semaphore; the polling loop is never blocked." + ), + ) + # ── readiness probes ────────────────────────────────────────── readiness_targets: str = Field( default="", diff --git a/common/src/common/db/session.py b/common/src/common/db/session.py index 65ed3ca..0b9836d 100644 --- a/common/src/common/db/session.py +++ b/common/src/common/db/session.py @@ -18,12 +18,18 @@ def create_database_engine( *, echo: bool = False, pool_pre_ping: bool = True, + pool_size: int = 10, + max_overflow: int = 20, + pool_recycle: int = 1800, ) -> AsyncEngine: """Create an async SQLAlchemy engine without storing global connection state.""" return create_async_engine( database_url, echo=echo, pool_pre_ping=pool_pre_ping, + pool_size=pool_size, + max_overflow=max_overflow, + pool_recycle=pool_recycle, ) diff --git a/runtime/src/runtime/process.py b/runtime/src/runtime/process.py index ab7b257..61dc294 100644 --- a/runtime/src/runtime/process.py +++ b/runtime/src/runtime/process.py @@ -9,6 +9,7 @@ a thin wrapper over ``start_workspace``. from __future__ import annotations import asyncio +import os import secrets import subprocess import time @@ -249,13 +250,18 @@ async def scan_workspaces() -> None: logger.error(f"scan workspace failed: {e}") return - for entry in entries: + sem = asyncio.Semaphore(4) + + async def _start(entry: str) -> None: path = WORKSPACES_ROOT / entry if not path.is_dir(): - continue + return logger.info(f"Found workspace: {entry}") logger.info(f"Auto-starting Jupyter for workspace ID: '{entry}'") - try: - await start_workspace(entry) - except Exception as err: - logger.error(f"Startup failed for workspace '{entry}': {err}") \ No newline at end of file + async with sem: + try: + await start_workspace(entry) + except Exception as err: + logger.error(f"Startup failed for workspace '{entry}': {err}") + + await asyncio.gather(*[_start(entry) for entry in entries]) \ No newline at end of file diff --git a/schedule/src/schedule/orchestrator.py b/schedule/src/schedule/orchestrator.py index a12137c..2ba932c 100644 --- a/schedule/src/schedule/orchestrator.py +++ b/schedule/src/schedule/orchestrator.py @@ -42,15 +42,32 @@ LOGGER = logging.getLogger(__name__) class DispatchOrchestrator: """Polls Outbox + advances DAG schedule runs. + Two independent loops run in the same process: + + - ``_database_event_loop`` drains ``schedule.run.requested`` and + ``job.node.finished`` events under ``dispatch_lock``. These are + short, in-line DB transactions. + - ``_execution_loop`` claims ``job.node.execute`` events, sets a + far-future ``available_at`` as a lease, then dispatches each as + ``asyncio.create_task`` so the polling path is never blocked by + notebook execution. A semaphore caps concurrent notebooks. + Holds a ``dispatch_lock`` to keep two concurrent drain loops from fighting over the same batch. """ + # Lease window for a claimed ``job.node.execute`` event. If the + # process dies mid-execution, the row re-eligible after this many + # minutes. The handler is idempotent (it short-circuits on terminal + # node states), so safe re-execution. + EXECUTION_LEASE = timedelta(minutes=30) + def __init__( self, *, session_factory: async_sessionmaker[AsyncSession], node_execute_handler, + execution_concurrency: int = 4, ) -> None: self.session_factory = session_factory # Injected by the facade; routes ``job.node.execute`` outbox events @@ -59,12 +76,19 @@ class DispatchOrchestrator: self._node_execute_handler = node_execute_handler self.dispatch_lock = asyncio.Lock() self._loop_task: asyncio.Task[None] | None = None + self._exec_loop_task: asyncio.Task[None] | None = None + self._exec_tasks: set[asyncio.Task[None]] = set() + self._exec_semaphore = asyncio.Semaphore(execution_concurrency) def start(self) -> None: self._loop_task = asyncio.create_task( self._database_event_loop(), name="scheduler-database-events", ) + self._exec_loop_task = asyncio.create_task( + self._execution_loop(), + name="scheduler-node-execute", + ) async def close(self) -> None: if self._loop_task is not None: @@ -74,6 +98,16 @@ class DispatchOrchestrator: except asyncio.CancelledError: pass self._loop_task = None + if self._exec_loop_task is not None: + self._exec_loop_task.cancel() + try: + await self._exec_loop_task + except asyncio.CancelledError: + pass + self._exec_loop_task = None + if self._exec_tasks: + await asyncio.gather(*self._exec_tasks, return_exceptions=True) + self._exec_tasks.clear() async def _database_event_loop(self) -> None: while True: @@ -87,6 +121,104 @@ class DispatchOrchestrator: LOGGER.exception("database event loop failed") await asyncio.sleep(1) + async def _execution_loop(self) -> None: + while True: + try: + claimed = await self._claim_execution_events(limit=10) + if not claimed: + await asyncio.sleep(0.25) + except asyncio.CancelledError: + raise + except Exception: + LOGGER.exception("node execute loop failed") + await asyncio.sleep(1) + + async def _claim_execution_events( + self, + *, + limit: int, + ) -> int: + """Claim ``job.node.execute`` rows and dispatch them as tasks. + + The claim step bumps ``available_at`` to a far-future lease so + the polling loop does not re-pick the same row while the + dispatched task is still running. Status stays ``pending``; + the dispatched task flips it to ``published`` or ``failed`` + when execution completes. + """ + async with session_scope(self.session_factory) as session: + statement = ( + select(OutboxEvents) + .where( + OutboxEvents.event_status == "pending", + OutboxEvents.event_type == "job.node.execute", + OutboxEvents.available_at <= utcnow(), + ) + .order_by(OutboxEvents.created_at) + .limit(limit) + ) + events = list((await session.scalars(statement)).all()) + claimed: list[tuple[dict[str, Any], str]] = [] + lease_until = utcnow() + self.EXECUTION_LEASE + for item in events: + item.available_at = lease_until + envelope = { + "event_type": item.event_type, + "event_id": item.event_id, + "trace_id": item.trace_id, + "payload": item.payload_json, + } + claimed.append((envelope, f"mysql:{item.event_id}")) + for envelope, message_id in claimed: + task = asyncio.create_task( + self._run_node_execute(envelope, message_id), + name=f"node-execute:{envelope['event_id']}", + ) + self._exec_tasks.add(task) + task.add_done_callback(self._exec_tasks.discard) + return len(claimed) + + async def _run_node_execute( + self, + envelope: dict[str, Any], + message_id: str, + ) -> None: + async with self._exec_semaphore: + exc: Exception | None = None + try: + await self._node_execute_handler(envelope, message_id) + except Exception as run_exc: # noqa: BLE001 + exc = run_exc + await self._update_execution_status(envelope["event_id"], exc) + + async def _update_execution_status( + self, + event_id: str, + exc: Exception | None, + ) -> None: + async with session_scope(self.session_factory) as session: + item = await session.scalar( + select(OutboxEvents) + .where(OutboxEvents.event_id == event_id) + .with_for_update() + ) + if item is None: + LOGGER.warning("execution event %s disappeared", event_id) + return + if exc is None: + item.event_status = "published" + item.published_at = utcnow() + item.last_error = None + else: + item.retry_count += 1 + item.last_error = str(exc)[:2000] + if item.retry_count >= 5: + item.event_status = "failed" + else: + item.available_at = utcnow() + timedelta( + seconds=min(30, 2 ** item.retry_count), + ) + async def process_pending_events( self, *, @@ -100,6 +232,9 @@ class DispatchOrchestrator: .where( OutboxEvents.event_status == "pending", OutboxEvents.available_at <= utcnow(), + OutboxEvents.event_type.in_( + ("schedule.run.requested", "job.node.finished"), + ), ) .order_by(OutboxEvents.created_at) .limit(limit) @@ -129,13 +264,18 @@ class DispatchOrchestrator: async def _process_outbox_event(self, item: OutboxEvents) -> None: handlers = { "schedule.run.requested": self._handle_run_requested, - "job.node.execute": self._node_execute_handler, "job.node.finished": self._handle_node_finished, } handler = handlers.get(item.event_type) if handler is None: raise ValueError(f"unsupported event type: {item.event_type}") - await handler(item.payload_json, f"mysql:{item.event_id}") + event = { + "event_type": item.event_type, + "event_id": item.event_id, + "trace_id": item.trace_id, + "payload": item.payload_json, + } + await handler(event, f"mysql:{item.event_id}") async def dispatch_run(self, run_id: str) -> int: return await self.process_pending_events( diff --git a/schedule/src/schedule/scheduler.py b/schedule/src/schedule/scheduler.py index 6105b2f..0aa22e0 100644 --- a/schedule/src/schedule/scheduler.py +++ b/schedule/src/schedule/scheduler.py @@ -121,16 +121,20 @@ class CronScheduler: expression, timezone=ZoneInfo(item.timezone), ) - job = self.scheduler.add_job( - self.trigger, - trigger=trigger, - args=[item.schedule_id], - id=job_id, - replace_existing=True, - coalesce=True, - max_instances=max(1, item.max_concurrency), - misfire_grace_time=60, - ) + if self.scheduler.get_job(job_id) is not None: + self.scheduler.reschedule_job(job_id, trigger=trigger) + else: + self.scheduler.add_job( + self.trigger, + trigger=trigger, + args=[item.schedule_id], + id=job_id, + replace_existing=True, + coalesce=True, + max_instances=max(1, item.max_concurrency), + misfire_grace_time=60, + ) + job = self.scheduler.get_job(job_id) item.next_run_at = naive_utc(job.next_run_time) for job in self.scheduler.get_jobs(): if ( diff --git a/schedule/src/schedule/service.py b/schedule/src/schedule/service.py index 2ef606e..4dac3a4 100644 --- a/schedule/src/schedule/service.py +++ b/schedule/src/schedule/service.py @@ -71,6 +71,7 @@ class SchedulerService: self.orchestrator = DispatchOrchestrator( session_factory=session_factory, node_execute_handler=self.worker.handle_node_execute, + execution_concurrency=settings.schedule_execution_concurrency, ) # CronScheduler's tick callback is ``self.trigger_schedule``; this # resolves via class-method lookup at invocation time, so the order