diff --git a/schedule/Dockerfile b/schedule/Dockerfile index 42f343b..c88ed6c 100644 --- a/schedule/Dockerfile +++ b/schedule/Dockerfile @@ -5,7 +5,8 @@ WORKDIR /app COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv COPY common ./common COPY schedule ./schedule -RUN uv pip install --system ./common ./schedule +COPY contracts ./contracts +RUN UV_DEFAULT_INDEX="https://pypi.tuna.tsinghua.edu.cn/simple/" uv sync --frozen --no-dev --no-editable --package schedule EXPOSE 8000 -CMD ["uvicorn", "schedule.main:app", "--host", "0.0.0.0", "--port", "8000"] +CMD ["uv", "run", "uvicorn", "schedule.main:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/schedule/pyproject.toml b/schedule/pyproject.toml index 89ffac9..b116545 100644 --- a/schedule/pyproject.toml +++ b/schedule/pyproject.toml @@ -15,7 +15,8 @@ dependencies = [ ] [tool.uv.sources] -common = { path = "../common" } +common = { workspace = true } + [build-system] requires = ["hatchling"] diff --git a/schedule/src/schedule/context.py b/schedule/src/schedule/context.py new file mode 100644 index 0000000..5f67123 --- /dev/null +++ b/schedule/src/schedule/context.py @@ -0,0 +1,42 @@ +"""Shared constants and timezone helpers for the scheduler components. + +Designed to be import-side-effect-free: no logging, no I/O, no model imports. +Used by ``scheduler``, ``orchestrator`` and ``worker`` modules. +""" + +from __future__ import annotations + +from datetime import UTC, datetime + + +TERMINAL_NODE_STATES = frozenset({ + "succeeded", + "failed", + "skipped", + "cancelled", + "timed_out", +}) +FAILED_NODE_STATES = frozenset({"failed", "cancelled", "timed_out"}) +TERMINAL_RUN_STATES = frozenset({ + "succeeded", + "failed", + "cancelled", + "timed_out", +}) + + +def naive_utc(value: datetime | None) -> datetime | None: + """Normalize a datetime to naive UTC; pass through ``None``.""" + if value is None: + return None + if value.tzinfo is None: + value = value.replace(tzinfo=UTC) + return value.astimezone(UTC).replace(tzinfo=None) + + +__all__ = [ + "TERMINAL_NODE_STATES", + "FAILED_NODE_STATES", + "TERMINAL_RUN_STATES", + "naive_utc", +] diff --git a/schedule/src/schedule/main.py b/schedule/src/schedule/main.py index ece1c3b..7caa80a 100644 --- a/schedule/src/schedule/main.py +++ b/schedule/src/schedule/main.py @@ -10,7 +10,6 @@ from common.service_app import create_service_app from schedule.service import ( SchedulerService, build_object_store, - build_redis_client, build_storage_http_client, ) from schedule.storage_client import SchedulerStorageClient @@ -20,16 +19,16 @@ from schedule.storage_client import SchedulerStorageClient async def lifespan(app: Any) -> AsyncIterator[None]: engine = create_database_engine(os.environ["DATABASE_URL"]) session_factory = create_session_factory(engine) - redis = build_redis_client() - storage_http_client = build_storage_http_client() + backend_http_client = build_storage_http_client() service = SchedulerService( session_factory=session_factory, - redis=redis, + backend_http_client=backend_http_client, object_store=build_object_store(), - storage_client=SchedulerStorageClient(storage_http_client), + storage_client=SchedulerStorageClient(backend_http_client), workspace_root=Path( os.getenv("WORKSPACE_ROOT", "/workspace/workspaces") ), + database_url=os.environ["DATABASE_URL"], ) app.state.scheduler_service = service await service.start() @@ -37,8 +36,7 @@ async def lifespan(app: Any) -> AsyncIterator[None]: yield finally: await service.close() - await storage_http_client.aclose() - await redis.aclose() + await backend_http_client.aclose() await engine.dispose() diff --git a/schedule/src/schedule/orchestrator.py b/schedule/src/schedule/orchestrator.py new file mode 100644 index 0000000..a12137c --- /dev/null +++ b/schedule/src/schedule/orchestrator.py @@ -0,0 +1,443 @@ +"""Outbox-driven DAG orchestrator. + +Polls the MySQL ``OutboxEvents`` table for ``schedule.run.requested`` and +``job.node.finished`` events and advances schedule runs accordingly. New +nodes are dispatched by writing ``job.node.execute`` rows to the Outbox +and letting the worker component consume them. + +This module owns no HTTP / boto3 / notebook execution dependencies — +those belong to the worker (see ``schedule.worker``) and the cron +post-back (see ``schedule.service.trigger_schedule``). +""" + +from __future__ import annotations + +import asyncio +import logging +from datetime import timedelta +from typing import Any + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from common.db import session_scope +from common.db.models import ( + ConsumerInbox, + OutboxEvents, + ScheduleNodeRuns, + ScheduleRuns, +) +from common.eventing import add_outbox_event, utcnow +from common.ids import new_ulid + +from schedule.context import ( + FAILED_NODE_STATES, + TERMINAL_NODE_STATES, + TERMINAL_RUN_STATES, +) + +LOGGER = logging.getLogger(__name__) + + +class DispatchOrchestrator: + """Polls Outbox + advances DAG schedule runs. + + Holds a ``dispatch_lock`` to keep two concurrent drain loops from + fighting over the same batch. + """ + + def __init__( + self, + *, + session_factory: async_sessionmaker[AsyncSession], + node_execute_handler, + ) -> None: + self.session_factory = session_factory + # Injected by the facade; routes ``job.node.execute`` outbox events + # to ``NodeExecutor.handle_node_execute``. Kept as a callable so this + # module can stay independent of worker module imports. + self._node_execute_handler = node_execute_handler + self.dispatch_lock = asyncio.Lock() + self._loop_task: asyncio.Task[None] | None = None + + def start(self) -> None: + self._loop_task = asyncio.create_task( + self._database_event_loop(), + name="scheduler-database-events", + ) + + async def close(self) -> None: + if self._loop_task is not None: + self._loop_task.cancel() + try: + await self._loop_task + except asyncio.CancelledError: + pass + self._loop_task = None + + async def _database_event_loop(self) -> None: + while True: + try: + processed = await self.process_pending_events(limit=20) + if not processed: + await asyncio.sleep(0.25) + except asyncio.CancelledError: + raise + except Exception: + LOGGER.exception("database event loop failed") + await asyncio.sleep(1) + + async def process_pending_events( + self, + *, + limit: int = 20, + aggregate_id: str | None = None, + ) -> int: + async with self.dispatch_lock: + async with session_scope(self.session_factory) as session: + statement = ( + select(OutboxEvents) + .where( + OutboxEvents.event_status == "pending", + OutboxEvents.available_at <= utcnow(), + ) + .order_by(OutboxEvents.created_at) + .limit(limit) + ) + if aggregate_id: + statement = statement.where( + OutboxEvents.aggregate_id == aggregate_id + ) + events = list((await session.scalars(statement)).all()) + for item in events: + try: + await self._process_outbox_event(item) + item.event_status = "published" + item.published_at = utcnow() + item.last_error = None + except Exception as exc: + 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) + ) + return len(events) + + 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}") + + async def dispatch_run(self, run_id: str) -> int: + return await self.process_pending_events( + limit=50, + aggregate_id=run_id, + ) + + async def _start_inbox( + self, + session: AsyncSession, + *, + consumer_name: str, + event_id: str, + message_id: str, + ) -> tuple[ConsumerInbox, bool]: + item = await session.get( + ConsumerInbox, + (consumer_name, event_id), + with_for_update=True, + ) + if item is not None and item.process_status == "succeeded": + return item, False + if item is None: + item = ConsumerInbox( + consumer_name=consumer_name, + event_id=event_id, + process_status="processing", + message_id=message_id, + ) + session.add(item) + else: + item.process_status = "processing" + item.message_id = message_id + item.error_message = None + return item, True + + @staticmethod + def _finish_inbox(item: ConsumerInbox) -> None: + item.process_status = "succeeded" + item.processed_at = utcnow() + item.error_message = None + + async def _handle_run_requested( + self, + event: dict[str, Any], + message_id: str, + ) -> None: + if event.get("event_type") != "schedule.run.requested": + raise ValueError("unexpected event type") + payload = event["payload"] + async with session_scope(self.session_factory) as session: + inbox, should_process = await self._start_inbox( + session, + consumer_name="schedule-orchestrator", + event_id=event["event_id"], + message_id=message_id, + ) + if not should_process: + return + run = await session.scalar( + select(ScheduleRuns) + .where(ScheduleRuns.run_id == payload["run_id"]) + .with_for_update() + ) + if run is None: + raise ValueError("schedule run does not exist") + if run.run_status not in TERMINAL_RUN_STATES: + if run.run_status == "queued": + run.run_status = "running" + run.started_at = utcnow() + run.state_version += 1 + await self._advance_run( + session, + run, + trace_id=event["trace_id"], + ) + self._finish_inbox(inbox) + + async def _dispatch_node( + self, + session: AsyncSession, + *, + run: ScheduleRuns, + node: dict[str, Any], + attempt_no: int, + trace_id: str, + delay_seconds: int = 0, + ) -> ScheduleNodeRuns: + node_run = ScheduleNodeRuns( + node_run_id=new_ulid(), + run_id=run.run_id, + node_id=node["node_id"], + versions_id=node["versions_id"], + attempt_no=attempt_no, + node_status="queued", + state_version=0, + message=( + f"等待重试（{delay_seconds} 秒）" + if delay_seconds + else "等待 Worker 执行" + ), + ) + session.add(node_run) + await add_outbox_event( + session, + event_type="job.node.execute", + producer="schedule-orchestrator", + trace_id=trace_id, + aggregate_type="schedule_node_run", + aggregate_id=node_run.node_run_id, + idempotency_key=f"{node_run.node_run_id}:{attempt_no}", + available_at=( + utcnow() + timedelta(seconds=delay_seconds) + if delay_seconds + else None + ), + payload={ + "workspace_id": run.workspace_id, + "run_id": run.run_id, + "node_run_id": node_run.node_run_id, + "node_id": node["node_id"], + "versions_id": node["versions_id"], + "attempt_no": attempt_no, + "script_type": node["script_type"], + "artifact_object_id": node["artifact_object_id"], + "artifact_path": node["artifact_path"], + "timeout_seconds": node["timeout_seconds"], + "arguments": node.get("arguments", []), + }, + ) + return node_run + + async def _advance_run( + self, + session: AsyncSession, + run: ScheduleRuns, + *, + trace_id: str, + ) -> None: + snapshot = run.schedule_snapshot + nodes = snapshot.get("nodes", []) + node_by_id = {node["node_id"]: node for node in nodes} + parents: dict[str, set[str]] = {node_id: set() for node_id in node_by_id} + for edge in snapshot.get("edges", []): + parents.setdefault(edge["target_node_id"], set()).add( + edge["source_node_id"] + ) + rows = list( + ( + await session.scalars( + select(ScheduleNodeRuns) + .where(ScheduleNodeRuns.run_id == run.run_id) + .order_by(ScheduleNodeRuns.attempt_no) + ) + ).all() + ) + latest: dict[str, ScheduleNodeRuns] = {} + for row in rows: + current = latest.get(row.node_id) + if current is None or row.attempt_no >= current.attempt_no: + latest[row.node_id] = row + + max_concurrency = max(1, int(snapshot.get("max_concurrency", 1))) + while True: + changed = False + active_count = sum( + item.node_status in {"queued", "running"} + for item in latest.values() + ) + for node in nodes: + current = latest.get(node["node_id"]) + if ( + current is not None + and current.node_status in FAILED_NODE_STATES + and current.attempt_no <= int(node.get("retry_count", 0)) + and active_count < max_concurrency + ): + retried = await self._dispatch_node( + session, + run=run, + node=node, + attempt_no=current.attempt_no + 1, + trace_id=trace_id, + delay_seconds=int(node.get("retry_interval_sec", 0)), + ) + latest[node["node_id"]] = retried + active_count += 1 + changed = True + + exhausted_failure = any( + item.node_status in FAILED_NODE_STATES + and item.attempt_no + > int(node_by_id[item.node_id].get("retry_count", 0)) + for item in latest.values() + ) + stop_all = ( + snapshot.get("failure_policy", "stop") == "stop" + and exhausted_failure + ) + for node in nodes: + node_id = node["node_id"] + if node_id in latest: + continue + parent_runs = [latest.get(parent) for parent in parents[node_id]] + parent_failed = any( + item is not None + and item.node_status in TERMINAL_NODE_STATES + and item.node_status != "succeeded" + for item in parent_runs + ) + if stop_all or parent_failed: + skipped = ScheduleNodeRuns( + node_run_id=new_ulid(), + run_id=run.run_id, + node_id=node_id, + versions_id=node["versions_id"], + attempt_no=1, + node_status="skipped", + state_version=1, + finished_at=utcnow(), + duration_ms=0, + message=( + "调度失败策略为 stop，未再启动" + if stop_all + else "上游节点未成功，已跳过" + ), + ) + session.add(skipped) + latest[node_id] = skipped + changed = True + elif ( + all( + item is not None and item.node_status == "succeeded" + for item in parent_runs + ) + and active_count < max_concurrency + ): + dispatched = await self._dispatch_node( + session, + run=run, + node=node, + attempt_no=1, + trace_id=trace_id, + ) + latest[node_id] = dispatched + active_count += 1 + changed = True + if not changed: + break + + if nodes and len(latest) == len(nodes) and all( + item.node_status in TERMINAL_NODE_STATES + for item in latest.values() + ): + now = utcnow() + succeeded = all( + item.node_status == "succeeded" for item in latest.values() + ) + run.run_status = "succeeded" if succeeded else "failed" + run.error_code = None if succeeded else "SCHEDULE_NODE_FAILED" + run.error_message = ( + None if succeeded else "one or more schedule nodes did not succeed" + ) + run.finished_at = now + if run.started_at: + run.duration_ms = max( + 0, + int((now - run.started_at).total_seconds() * 1000), + ) + run.state_version += 1 + + async def _handle_node_finished( + self, + event: dict[str, Any], + message_id: str, + ) -> None: + if event.get("event_type") != "job.node.finished": + raise ValueError("unexpected event type") + payload = event["payload"] + async with session_scope(self.session_factory) as session: + inbox, should_process = await self._start_inbox( + session, + consumer_name="schedule-results", + event_id=event["event_id"], + message_id=message_id, + ) + if not should_process: + return + run = await session.scalar( + select(ScheduleRuns) + .where(ScheduleRuns.run_id == payload["run_id"]) + .with_for_update() + ) + if run is None: + raise ValueError("schedule run does not exist") + if run.run_status not in TERMINAL_RUN_STATES: + await self._advance_run( + session, + run, + trace_id=event["trace_id"], + ) + self._finish_inbox(inbox) + + +__all__ = ["DispatchOrchestrator"] diff --git a/schedule/src/schedule/scheduler.py b/schedule/src/schedule/scheduler.py new file mode 100644 index 0000000..6105b2f --- /dev/null +++ b/schedule/src/schedule/scheduler.py @@ -0,0 +1,143 @@ +"""APScheduler-backed cron trigger layer. + +Owns the ``AsyncIOScheduler`` instance plus the periodic sync loop that +reconciles in-memory APScheduler jobs against the ``Schedules`` table in +MySQL. The cron tick callback delegates to ``SchedulerService.trigger_schedule`` +(via the ``on_trigger`` callable injected at construction), which posts +back to Backend; Backend then writes a ``schedule.run.requested`` Outbox +row that the orchestrator consumes. + +The two layers (this cron scheduler, the orchestrator) are decoupled +through MySQL — they share no in-memory state and survive independent +restarts. +""" + +from __future__ import annotations + +import asyncio +import logging +from datetime import UTC +from typing import Awaitable, Callable +from zoneinfo import ZoneInfo + +from apscheduler.schedulers.asyncio import AsyncIOScheduler +from apscheduler.triggers.cron import CronTrigger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from common.db import session_scope +from common.db.models import Schedules +from common.scheduler import build_sqlalchemy_jobstore + +from schedule.context import naive_utc + +LOGGER = logging.getLogger(__name__) + + +class CronScheduler: + """Manages cron triggers in APScheduler, backed by a MySQL jobstore.""" + + def __init__( + self, + *, + session_factory: async_sessionmaker[AsyncSession], + database_url: str, + on_trigger: Callable[[str], Awaitable[None]], + ) -> None: + self.session_factory = session_factory + self.scheduler = AsyncIOScheduler( + jobstores={"default": build_sqlalchemy_jobstore(database_url)}, + timezone=UTC, + ) + self._on_trigger = on_trigger + self._sync_task: asyncio.Task[None] | None = None + + def start(self) -> None: + """Start APScheduler and spawn the periodic sync loop.""" + self.scheduler.start() + self._sync_task = asyncio.create_task( + self._sync_loop(), + name="scheduler-cron-sync", + ) + + async def close(self) -> None: + """Cancel the sync loop and shut APScheduler down.""" + if self._sync_task is not None: + self._sync_task.cancel() + try: + await self._sync_task + except asyncio.CancelledError: + pass + self._sync_task = None + if self.scheduler.running: + self.scheduler.shutdown(wait=False) + + async def trigger(self, schedule_id: str) -> None: + """APScheduler cron tick callback. + + Wired via ``add_job(self.trigger, args=[schedule_id], ...)``. The + standard on_trigger is ``SchedulerService.trigger_schedule`` which + posts back to Backend; Backend then writes the Outbox row that the + orchestrator picks up. + """ + await self._on_trigger(schedule_id) + + async def _sync_loop(self) -> None: + while True: + try: + await self._sync_once() + except asyncio.CancelledError: + raise + except Exception: + LOGGER.exception("cron job synchronization failed") + await asyncio.sleep(5) + + async def _sync_once(self) -> None: + """Reconcile APScheduler jobs against ``Schedules.cron_expression``. + + - adds jobs for enabled cron schedules present in MySQL + - removes jobs whose schedule has been disabled / soft-deleted + - updates ``Schedules.next_run_at`` from the next APScheduler tick + """ + async with session_scope(self.session_factory) as session: + schedules = list( + ( + await session.scalars( + select(Schedules).where( + Schedules.deleted_at.is_(None), + Schedules.enabled == 1, + Schedules.trigger_type == "cron", + Schedules.cron_expression.is_not(None), + ) + ) + ).all() + ) + active_job_ids: set[str] = set() + for item in schedules: + job_id = f"schedule:{item.schedule_id}" + active_job_ids.add(job_id) + expression = (item.cron_expression or "").strip() + trigger = CronTrigger.from_crontab( + 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, + ) + item.next_run_at = naive_utc(job.next_run_time) + for job in self.scheduler.get_jobs(): + if ( + job.id.startswith("schedule:") + and job.id not in active_job_ids + ): + self.scheduler.remove_job(job.id) + + +__all__ = ["CronScheduler"] diff --git a/schedule/src/schedule/service.py b/schedule/src/schedule/service.py index 33e5202..4f62776 100644 --- a/schedule/src/schedule/service.py +++ b/schedule/src/schedule/service.py @@ -1,247 +1,117 @@ +"""Scheduler service facade. + +Composes three single-purpose components into one bootable service: + + - :class:`schedule.scheduler.CronScheduler` — APScheduler + cron sync loop + - :class:`schedule.orchestrator.DispatchOrchestrator` — Outbox polling + DAG + - :class:`schedule.worker.NodeExecutor` — node-level execution + +This module also exposes the two factory functions (``build_object_store`` / +``build_storage_http_client``) consumed by ``schedule.main`` to construct +the backing resources that flow into the facade. + +The ``SchedulerService`` itself stays small: it wires the three components +together and implements ``trigger_schedule``, the cron post-back to +Backend that ``CronScheduler`` calls at every cron tick. +""" + from __future__ import annotations -import asyncio -import hashlib -import json import logging import os -import traceback -from contextlib import suppress -from datetime import UTC, datetime, timedelta -from pathlib import Path +from datetime import UTC, datetime from typing import Any -from zoneinfo import ZoneInfo -import boto3 import httpx -from apscheduler.schedulers.asyncio import AsyncIOScheduler -from apscheduler.triggers.cron import CronTrigger -from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker -from common.db.models import ( - ConsumerInbox, - OutboxEvents, - ScheduleNodeRuns, - ScheduleRuns, - Schedules, - StorageObjects, - Versions, - Workspaces, -) -from common.scheduler import build_sqlalchemy_jobstore +from common.ids import new_ulid +from schedule.orchestrator import DispatchOrchestrator +from schedule.scheduler import CronScheduler +from schedule.worker import NodeExecutor LOGGER = logging.getLogger(__name__) -TERMINAL_NODE_STATES = { - "succeeded", - "failed", - "skipped", - "cancelled", - "timed_out", -} -FAILED_NODE_STATES = {"failed", "cancelled", "timed_out"} -TERMINAL_RUN_STATES = { - "succeeded", - "failed", - "cancelled", - "timed_out", -} - -_ACTIVE_SERVICE: "SchedulerService | None" = None - - -def _naive_utc(value: datetime | None) -> datetime | None: - if value is None: - return None - if value.tzinfo is None: - value = value.replace(tzinfo=UTC) - return value.astimezone(UTC).replace(tzinfo=None) - - -async def run_scheduled_job(schedule_id: str) -> None: - service = _ACTIVE_SERVICE - if service is None: - LOGGER.warning("scheduler job skipped because service is not ready") - return - await service.trigger_schedule(schedule_id) class SchedulerService: + """Composes cron / orchestrator / worker into one bootable service. + + Lifecycle:: + + service = SchedulerService(...) + await service.start() # spawns all loops + ... + await service.close() # cancels all loops, shuts APScheduler down + """ + def __init__( self, *, session_factory: async_sessionmaker[AsyncSession], - object_store: Any, - storage_client: SchedulerStorageClient, backend_http_client: httpx.AsyncClient, - workspace_root: Path, + object_store: Any, + storage_client: Any, + workspace_root: Any, database_url: str, ) -> None: self.session_factory = session_factory + self.backend_http_client = backend_http_client self.object_store = object_store self.storage_client = storage_client - self.backend_http_client = backend_http_client self.workspace_root = workspace_root - self.tasks: list[asyncio.Task[Any]] = [] - self.dispatch_lock = asyncio.Lock() - self.scheduler = AsyncIOScheduler( - jobstores={ - "default": build_sqlalchemy_jobstore(database_url) - }, - timezone=UTC, + self.database_url = database_url + + # Wire worker BEFORE orchestrator: orchestrator's dispatch table + # references ``self.worker.handle_node_execute`` directly, so the + # worker instance must exist on ``self`` at injection time. + self.worker = NodeExecutor( + session_factory=session_factory, + object_store=object_store, + storage_client=storage_client, + workspace_root=workspace_root, + ) + self.orchestrator = DispatchOrchestrator( + session_factory=session_factory, + node_execute_handler=self.worker.handle_node_execute, + ) + # CronScheduler's tick callback is ``self.trigger_schedule``; this + # resolves via class-method lookup at invocation time, so the order + # relative to ``self.cron`` itself doesn't matter. + self.cron = CronScheduler( + session_factory=session_factory, + database_url=database_url, + on_trigger=self.trigger_schedule, ) async def start(self) -> None: - global _ACTIVE_SERVICE - _ACTIVE_SERVICE = self - self.scheduler.start() - await self._sync_cron_jobs() - self.tasks = [ - asyncio.create_task( - self._database_event_loop(), - name="scheduler-database-events", - ), - asyncio.create_task( - self._schedule_sync_loop(), - name="scheduler-cron-sync", - ), - ] + """Start cron scheduler and outbox polling loop. + + Neither component's ``start()`` is itself async — they both just + ``create_task`` for their inner loops and return. We keep this + method ``async def`` so the FastAPI lifespan can ``await`` it as a + coroutine, but we don't ``await`` either inner call. + """ + self.cron.start() + self.orchestrator.start() async def close(self) -> None: - global _ACTIVE_SERVICE - for task in self.tasks: - task.cancel() - for task in self.tasks: - with suppress(asyncio.CancelledError): - await task - self.tasks.clear() - if self.scheduler.running: - self.scheduler.shutdown(wait=False) - _ACTIVE_SERVICE = None - - async def _database_event_loop(self) -> None: - while True: - try: - processed = await self.process_pending_events(limit=20) - if not processed: - await asyncio.sleep(0.25) - except asyncio.CancelledError: - raise - except Exception: - LOGGER.exception("database event loop failed") - await asyncio.sleep(1) - - async def process_pending_events( - self, - *, - limit: int = 20, - aggregate_id: str | None = None, - ) -> int: - async with self.dispatch_lock: - async with session_scope(self.session_factory) as session: - statement = ( - select(OutboxEvents) - .where( - OutboxEvents.event_status == "pending", - OutboxEvents.available_at <= utcnow(), - ) - .order_by(OutboxEvents.created_at) - .limit(limit) - ) - if aggregate_id: - statement = statement.where( - OutboxEvents.aggregate_id == aggregate_id - ) - events = list((await session.scalars(statement)).all()) - for item in events: - try: - await self._process_outbox_event(item) - item.event_status = "published" - item.published_at = utcnow() - item.last_error = None - except Exception as exc: - 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) - ) - LOGGER.exception( - "failed to process database event %s", - item.event_id, - ) - return len(events) - - async def _process_outbox_event(self, item: OutboxEvents) -> None: - handlers = { - "schedule.run.requested": self._handle_run_requested, - "job.node.execute": self._handle_node_execute, - "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}") - - async def dispatch_run(self, run_id: str) -> int: - return await self.process_pending_events( - limit=50, - aggregate_id=run_id, - ) - - async def _schedule_sync_loop(self) -> None: - while True: - try: - await self._sync_cron_jobs() - except asyncio.CancelledError: - raise - except Exception: - LOGGER.exception("cron job synchronization failed") - await asyncio.sleep(5) - - async def _sync_cron_jobs(self) -> None: - async with session_scope(self.session_factory) as session: - schedules = list( - ( - await session.scalars( - select(Schedules).where( - Schedules.deleted_at.is_(None), - Schedules.enabled == 1, - Schedules.trigger_type == "cron", - Schedules.cron_expression.is_not(None), - ) - ) - ).all() - ) - active_job_ids: set[str] = set() - for item in schedules: - job_id = f"schedule:{item.schedule_id}" - active_job_ids.add(job_id) - expression = (item.cron_expression or "").strip() - trigger = CronTrigger.from_crontab( - expression, - timezone=ZoneInfo(item.timezone), - ) - job = self.scheduler.add_job( - run_scheduled_job, - 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, - ) - item.next_run_at = _naive_utc(job.next_run_time) - for job in self.scheduler.get_jobs(): - if job.id.startswith("schedule:") and job.id not in active_job_ids: - self.scheduler.remove_job(job.id) + """Cancel orchestrator loop, then shut APScheduler down.""" + await self.orchestrator.close() + await self.cron.close() async def trigger_schedule(self, schedule_id: str) -> None: + """Cron tick callback: post back to Backend to register a new run. + + Backend writes the ``schedule.run.requested`` Outbox row in the same + transaction as the ``ScheduleRuns`` insert; the orchestrator's polling + loop will pick it up and start advancing the DAG. + """ async with self.session_factory() as session: + from sqlalchemy import select + + from common.db.models import Schedules + item = await session.get(Schedules, schedule_id) if ( item is None @@ -272,629 +142,36 @@ class SchedulerService: f"{response.text[:500]}" ) - async def _start_inbox( - self, - session: AsyncSession, - *, - consumer_name: str, - event_id: str, - message_id: str, - ) -> tuple[ConsumerInbox, bool]: - item = await session.get( - ConsumerInbox, - (consumer_name, event_id), - with_for_update=True, - ) - if item is not None and item.process_status == "succeeded": - return item, False - if item is None: - item = ConsumerInbox( - consumer_name=consumer_name, - event_id=event_id, - process_status="processing", - message_id=message_id, - ) - session.add(item) - else: - item.process_status = "processing" - item.message_id = message_id - item.error_message = None - return item, True - - @staticmethod - def _finish_inbox(item: ConsumerInbox) -> None: - item.process_status = "succeeded" - item.processed_at = utcnow() - item.error_message = None - - async def _handle_run_requested( - self, - event: dict[str, Any], - message_id: str, - ) -> None: - if event.get("event_type") != "schedule.run.requested": - raise ValueError("unexpected event type") - payload = event["payload"] - async with session_scope(self.session_factory) as session: - inbox, should_process = await self._start_inbox( - session, - consumer_name="schedule-orchestrator", - event_id=event["event_id"], - message_id=message_id, - ) - if not should_process: - return - run = await session.scalar( - select(ScheduleRuns) - .where(ScheduleRuns.run_id == payload["run_id"]) - .with_for_update() - ) - if run is None: - raise ValueError("schedule run does not exist") - if run.run_status not in TERMINAL_RUN_STATES: - if run.run_status == "queued": - run.run_status = "running" - run.started_at = utcnow() - run.state_version += 1 - await self._advance_run( - session, - run, - trace_id=event["trace_id"], - ) - self._finish_inbox(inbox) - - async def _dispatch_node( - self, - session: AsyncSession, - *, - run: ScheduleRuns, - node: dict[str, Any], - attempt_no: int, - trace_id: str, - delay_seconds: int = 0, - ) -> ScheduleNodeRuns: - node_run = ScheduleNodeRuns( - node_run_id=new_ulid(), - run_id=run.run_id, - node_id=node["node_id"], - versions_id=node["versions_id"], - attempt_no=attempt_no, - node_status="queued", - state_version=0, - message=( - f"等待重试（{delay_seconds} 秒）" - if delay_seconds - else "等待 Worker 执行" - ), - ) - session.add(node_run) - await add_outbox_event( - session, - event_type="job.node.execute", - producer="schedule-orchestrator", - trace_id=trace_id, - aggregate_type="schedule_node_run", - aggregate_id=node_run.node_run_id, - idempotency_key=f"{node_run.node_run_id}:{attempt_no}", - available_at=( - utcnow() + timedelta(seconds=delay_seconds) - if delay_seconds - else None - ), - payload={ - "workspace_id": run.workspace_id, - "run_id": run.run_id, - "node_run_id": node_run.node_run_id, - "node_id": node["node_id"], - "versions_id": node["versions_id"], - "attempt_no": attempt_no, - "script_type": node["script_type"], - "artifact_object_id": node["artifact_object_id"], - "artifact_path": node["artifact_path"], - "timeout_seconds": node["timeout_seconds"], - "arguments": node.get("arguments", []), - }, - ) - return node_run - - async def _advance_run( - self, - session: AsyncSession, - run: ScheduleRuns, - *, - trace_id: str, - ) -> None: - snapshot = run.schedule_snapshot - nodes = snapshot.get("nodes", []) - node_by_id = {node["node_id"]: node for node in nodes} - parents: dict[str, set[str]] = {node_id: set() for node_id in node_by_id} - for edge in snapshot.get("edges", []): - parents.setdefault(edge["target_node_id"], set()).add( - edge["source_node_id"] - ) - rows = list( - ( - await session.scalars( - select(ScheduleNodeRuns) - .where(ScheduleNodeRuns.run_id == run.run_id) - .order_by(ScheduleNodeRuns.attempt_no) - ) - ).all() - ) - latest: dict[str, ScheduleNodeRuns] = {} - for row in rows: - current = latest.get(row.node_id) - if current is None or row.attempt_no >= current.attempt_no: - latest[row.node_id] = row - - max_concurrency = max(1, int(snapshot.get("max_concurrency", 1))) - while True: - changed = False - active_count = sum( - item.node_status in {"queued", "running"} - for item in latest.values() - ) - for node in nodes: - current = latest.get(node["node_id"]) - if ( - current is not None - and current.node_status in FAILED_NODE_STATES - and current.attempt_no <= int(node.get("retry_count", 0)) - and active_count < max_concurrency - ): - retried = await self._dispatch_node( - session, - run=run, - node=node, - attempt_no=current.attempt_no + 1, - trace_id=trace_id, - delay_seconds=int(node.get("retry_interval_sec", 0)), - ) - latest[node["node_id"]] = retried - active_count += 1 - changed = True - - exhausted_failure = any( - item.node_status in FAILED_NODE_STATES - and item.attempt_no - > int(node_by_id[item.node_id].get("retry_count", 0)) - for item in latest.values() - ) - stop_all = ( - snapshot.get("failure_policy", "stop") == "stop" - and exhausted_failure - ) - for node in nodes: - node_id = node["node_id"] - if node_id in latest: - continue - parent_runs = [latest.get(parent) for parent in parents[node_id]] - parent_failed = any( - item is not None - and item.node_status in TERMINAL_NODE_STATES - and item.node_status != "succeeded" - for item in parent_runs - ) - if stop_all or parent_failed: - skipped = ScheduleNodeRuns( - node_run_id=new_ulid(), - run_id=run.run_id, - node_id=node_id, - versions_id=node["versions_id"], - attempt_no=1, - node_status="skipped", - state_version=1, - finished_at=utcnow(), - duration_ms=0, - message=( - "调度失败策略为 stop，未再启动" - if stop_all - else "上游节点未成功，已跳过" - ), - ) - session.add(skipped) - latest[node_id] = skipped - changed = True - elif ( - all( - item is not None and item.node_status == "succeeded" - for item in parent_runs - ) - and active_count < max_concurrency - ): - dispatched = await self._dispatch_node( - session, - run=run, - node=node, - attempt_no=1, - trace_id=trace_id, - ) - latest[node_id] = dispatched - active_count += 1 - changed = True - if not changed: - break - - if nodes and len(latest) == len(nodes) and all( - item.node_status in TERMINAL_NODE_STATES - for item in latest.values() - ): - now = utcnow() - succeeded = all( - item.node_status == "succeeded" for item in latest.values() - ) - run.run_status = "succeeded" if succeeded else "failed" - run.error_code = None if succeeded else "SCHEDULE_NODE_FAILED" - run.error_message = ( - None if succeeded else "one or more schedule nodes did not succeed" - ) - run.finished_at = now - if run.started_at: - run.duration_ms = max( - 0, - int((now - run.started_at).total_seconds() * 1000), - ) - run.state_version += 1 - - async def _execution_context( - self, - payload: dict[str, Any], - ) -> dict[str, Any]: - async with self.session_factory() as session: - row = ( - await session.execute( - select( - ScheduleNodeRuns, - ScheduleRuns, - Versions, - StorageObjects, - Workspaces, - Schedules, - ) - .join( - ScheduleRuns, - ScheduleRuns.run_id == ScheduleNodeRuns.run_id, - ) - .join( - Versions, - Versions.versions_id == ScheduleNodeRuns.versions_id, - ) - .join( - StorageObjects, - StorageObjects.storage_object_id - == Versions.artifact_object_id, - ) - .join( - Workspaces, - Workspaces.workspace_id == ScheduleRuns.workspace_id, - ) - .join( - Schedules, - Schedules.schedule_id == ScheduleRuns.schedule_id, - ) - .where( - ScheduleNodeRuns.node_run_id == payload["node_run_id"], - ) - ) - ).one_or_none() - if row is None: - raise ValueError("node execution metadata not found") - node_run, run, version, storage, workspace, schedule = row - if storage.object_status != "available": - raise ValueError("stable version artifact is not available") - if storage.storage_backend != "rustfs": - raise ValueError("stable version artifact is not stored in RustFS") - if not storage.bucket_name or not storage.object_key: - raise ValueError("stable version artifact location is incomplete") - user_id = run.triggered_by or schedule.created_by - return { - "node_status": node_run.node_status, - "workspace_id": run.workspace_id, - "workspace_code": workspace.workspace_code, - "user_id": user_id, - "bucket_name": storage.bucket_name, - "object_key": storage.object_key, - "content_hash": version.content_hash, - } - - async def _fallback_execution_context( - self, - payload: dict[str, Any], - ) -> dict[str, Any]: - async with self.session_factory() as session: - row = ( - await session.execute( - select(ScheduleRuns, Workspaces, Schedules) - .join( - Workspaces, - Workspaces.workspace_id == ScheduleRuns.workspace_id, - ) - .join( - Schedules, - Schedules.schedule_id == ScheduleRuns.schedule_id, - ) - .where(ScheduleRuns.run_id == payload["run_id"]) - ) - ).one_or_none() - if row is None: - raise ValueError("schedule run execution context not found") - run, workspace, schedule = row - return { - "workspace_id": run.workspace_id, - "workspace_code": workspace.workspace_code, - "user_id": run.triggered_by or schedule.created_by, - } - - async def _download_artifact( + async def process_pending_events( self, *, - bucket_name: str, - object_key: str, - content_hash: str, - ) -> bytes: - def read() -> bytes: - response = self.object_store.get_object( - Bucket=bucket_name, - Key=object_key, - ) - body = response["Body"] - try: - return body.read() - finally: - body.close() - - content = await asyncio.to_thread(read) - if hashlib.sha256(content).hexdigest() != content_hash: - raise ValueError("stable version artifact hash mismatch") - return content - - async def _set_node_running( - self, - event: dict[str, Any], - message_id: str, - ) -> bool: - payload = event["payload"] - async with session_scope(self.session_factory) as session: - inbox, should_process = await self._start_inbox( - session, - consumer_name="job-workers", - event_id=event["event_id"], - message_id=message_id, - ) - if not should_process: - return False - node_run = await session.scalar( - select(ScheduleNodeRuns) - .where( - ScheduleNodeRuns.node_run_id == payload["node_run_id"], - ) - .with_for_update() - ) - if node_run is None: - raise ValueError("schedule node run does not exist") - if node_run.node_status in TERMINAL_NODE_STATES: - self._finish_inbox(inbox) - return False - if node_run.node_status == "queued": - node_run.node_status = "running" - node_run.started_at = utcnow() - node_run.message = "Worker 正在执行稳定版本" - node_run.state_version += 1 - return True - - async def _upload_execution_artifacts( - self, - *, - payload: dict[str, Any], - context: dict[str, Any], - result: ExecutionResult, - ) -> tuple[str | None, str | None, str | None]: - log_id: str | None = None - result_id: str | None = None - upload_error: str | None = None - try: - log_object = await self.storage_client.create_object( - workspace_id=context["workspace_id"], - user_id=context["user_id"], - usage_type="run_log", - file_name=f"{payload['node_run_id']}.log", - content_type="text/plain; charset=utf-8", - content=result.logs, - idempotency_key=f"{payload['node_run_id']}:log", - ) - log_id = log_object["storage_object_id"] - result_object = await self.storage_client.create_object( - workspace_id=context["workspace_id"], - user_id=context["user_id"], - usage_type="run_result", - file_name=result.result_file_name, - content_type=result.result_content_type, - content=result.result, - idempotency_key=f"{payload['node_run_id']}:result", - ) - result_id = result_object["storage_object_id"] - except Exception as exc: - upload_error = f"result upload failed: {exc}"[:2000] - LOGGER.exception("failed to upload node execution artifacts") - return log_id, result_id, upload_error - - async def _handle_node_execute( - self, - event: dict[str, Any], - message_id: str, - ) -> None: - if event.get("event_type") != "job.node.execute": - raise ValueError("unexpected event type") - payload = event["payload"] - if not await self._set_node_running(event, message_id): - return - started_at = utcnow() - context: dict[str, Any] | None = None - try: - context = await self._execution_context(payload) - content = await self._download_artifact( - bucket_name=context["bucket_name"], - object_key=context["object_key"], - content_hash=context["content_hash"], - ) - workspace_root = ( - self.workspace_root / context["workspace_code"] - ).resolve() - workspace_root.mkdir(parents=True, exist_ok=True) - result = await execute_artifact( - content, - run_id=payload["run_id"], - node_run_id=payload["node_run_id"], - script_type=payload["script_type"], - artifact_path=payload["artifact_path"], - arguments=[str(item) for item in payload.get("arguments", [])], - workspace_root=workspace_root, - timeout_seconds=int(payload["timeout_seconds"]), - ) - except Exception as exc: - trace = traceback.format_exc() - result = ExecutionResult( - status="failed", - exit_code=1, - logs=trace.encode("utf-8", errors="replace"), - result=json.dumps( - {"status": "failed", "error": str(exc)}, - ensure_ascii=False, - ).encode("utf-8"), - result_file_name=f"{payload['node_run_id']}-result.json", - result_content_type="application/json", - error_code="WORKER_EXECUTION_FAILED", - error_message=str(exc)[:2000], - ) - if context is None: - context = await self._fallback_execution_context(payload) - - log_id, result_id, upload_error = await self._upload_execution_artifacts( - payload=payload, - context=context, - result=result, - ) - finished_at = utcnow() - duration_ms = max( - 0, - int((finished_at - started_at).total_seconds() * 1000), - ) - error_message = result.error_message - if upload_error: - error_message = ( - f"{error_message}; {upload_error}" - if error_message - else upload_error - )[:2000] - final_status = "failed" if upload_error else result.status - final_error_code = ( - "ARTIFACT_UPLOAD_FAILED" if upload_error else result.error_code + limit: int = 20, + aggregate_id: str | None = None, + ) -> int: + """Pass-through to ``orchestrator.process_pending_events``.""" + return await self.orchestrator.process_pending_events( + limit=limit, + aggregate_id=aggregate_id, ) - async with session_scope(self.session_factory) as session: - node_run = await session.scalar( - select(ScheduleNodeRuns) - .where( - ScheduleNodeRuns.node_run_id == payload["node_run_id"], - ) - .with_for_update() - ) - if node_run is None: - raise ValueError("schedule node run disappeared") - inbox = await session.get( - ConsumerInbox, - ("job-workers", event["event_id"]), - with_for_update=True, - ) - if inbox is None: - raise ValueError("job worker inbox record disappeared") - if node_run.node_status not in TERMINAL_NODE_STATES: - node_run.node_status = final_status - node_run.finished_at = finished_at - node_run.duration_ms = duration_ms - node_run.exit_code = result.exit_code - node_run.message = ( - "节点执行成功" - if final_status == "succeeded" - else (error_message or "节点执行失败") - )[:2000] - node_run.metrics_json = { - "log_size_bytes": len(result.logs), - "result_size_bytes": len(result.result), - } - node_run.logs_object_id = log_id - node_run.result_object_id = result_id - node_run.state_version += 1 - await add_outbox_event( - session, - event_type="job.node.finished", - producer="job-worker", - trace_id=event["trace_id"], - aggregate_type="schedule_node_run", - aggregate_id=node_run.node_run_id, - idempotency_key=( - f"{node_run.node_run_id}:{node_run.attempt_no}:finished" - ), - payload={ - "workspace_id": context["workspace_id"], - "run_id": node_run.run_id, - "node_run_id": node_run.node_run_id, - "node_id": node_run.node_id, - "versions_id": node_run.versions_id, - "attempt_no": node_run.attempt_no, - "node_status": node_run.node_status, - "exit_code": node_run.exit_code, - "started_at": event_time( - node_run.started_at or started_at - ), - "finished_at": event_time(finished_at), - "duration_ms": duration_ms, - "logs_object_id": log_id, - "result_object_id": result_id, - "error_code": final_error_code, - "error_message": error_message, - }, - ) - self._finish_inbox(inbox) - - async def _handle_node_finished( - self, - event: dict[str, Any], - message_id: str, - ) -> None: - if event.get("event_type") != "job.node.finished": - raise ValueError("unexpected event type") - payload = event["payload"] - async with session_scope(self.session_factory) as session: - inbox, should_process = await self._start_inbox( - session, - consumer_name="schedule-results", - event_id=event["event_id"], - message_id=message_id, - ) - if not should_process: - return - run = await session.scalar( - select(ScheduleRuns) - .where(ScheduleRuns.run_id == payload["run_id"]) - .with_for_update() - ) - if run is None: - raise ValueError("schedule run does not exist") - if run.run_status not in TERMINAL_RUN_STATES: - await self._advance_run( - session, - run, - trace_id=event["trace_id"], - ) - self._finish_inbox(inbox) + async def dispatch_run(self, run_id: str) -> int: + """Pass-through to ``orchestrator.dispatch_run``.""" + return await self.orchestrator.dispatch_run(run_id) def build_object_store() -> Any: + """Construct a boto3 S3 client pointed at RustFS. + + Reads ``RUSTFS_ENDPOINT`` (full URL), ``RUSTFS_ACCESS_KEY``, and + ``RUSTFS_SECRET_KEY``. Falls back to ``http://rustfs:9000`` for the + endpoint — that's the default docker-compose service name. + """ + import boto3 + return boto3.client( "s3", endpoint_url=os.getenv( - "RUSTFS_INTERNAL_ENDPOINT", + "RUSTFS_ENDPOINT", "http://rustfs:9000", ), aws_access_key_id=os.environ["RUSTFS_ACCESS_KEY"], @@ -904,7 +181,15 @@ def build_object_store() -> Any: def build_storage_http_client() -> httpx.AsyncClient: + """Construct the httpx client that talks to Backend's HTTP API.""" return httpx.AsyncClient( base_url=os.getenv("BACKEND_API_URL", "http://backend:8000"), timeout=httpx.Timeout(60.0), ) + + +__all__ = [ + "SchedulerService", + "build_object_store", + "build_storage_http_client", +] diff --git a/schedule/src/schedule/worker.py b/schedule/src/schedule/worker.py new file mode 100644 index 0000000..60231f8 --- /dev/null +++ b/schedule/src/schedule/worker.py @@ -0,0 +1,417 @@ +"""Node-level worker: executes one ``job.node.execute`` event. + +The orchestrator (see ``schedule.orchestrator``) writes a ``job.node.execute`` +Outbox row with all the metadata needed to run the node (script type, +artifact location, timeout, arguments ...). The polling loop picks those up +and calls :meth:`NodeExecutor.handle_node_execute`. This module owns the +artifact download + subprocess invocation + result-upload side of things. + +`DispatchOrchestrator` writes the node's lifecycle row + outbox event; the +worker only mutates ``ScheduleNodeRuns`` columns related to execution +(started_at / finished_at / exit_code / result_object_id ...). +""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +import logging +import traceback +from pathlib import Path +from typing import Any + +from sqlalchemy import select + +from common.db import session_scope +from common.db.models import ( + ConsumerInbox, + ScheduleNodeRuns, + ScheduleRuns, + Schedules, + StorageObjects, + Versions, + Workspaces, +) +from common.eventing import add_outbox_event, event_time, utcnow + +from schedule.context import TERMINAL_NODE_STATES +from schedule.execution import ExecutionResult, execute_artifact + +LOGGER = logging.getLogger(__name__) + + +class NodeExecutor: + """Owns the actual execution of one schedule node (notebook / python).""" + + def __init__( + self, + *, + session_factory, + object_store: Any, + storage_client: Any, + workspace_root: Path, + ) -> None: + self.session_factory = session_factory + self.object_store = object_store + self.storage_client = storage_client + self.workspace_root = workspace_root + + async def handle_node_execute( + self, + event: dict[str, Any], + message_id: str, + ) -> None: + if event.get("event_type") != "job.node.execute": + raise ValueError("unexpected event type") + payload = event["payload"] + if not await self._set_node_running(event, message_id): + return + started_at = utcnow() + context: dict[str, Any] | None = None + try: + context = await self._execution_context(payload) + content = await self._download_artifact( + bucket_name=context["bucket_name"], + object_key=context["object_key"], + content_hash=context["content_hash"], + ) + workspace_root = ( + self.workspace_root / context["workspace_code"] + ).resolve() + workspace_root.mkdir(parents=True, exist_ok=True) + result = await execute_artifact( + content, + run_id=payload["run_id"], + node_run_id=payload["node_run_id"], + script_type=payload["script_type"], + artifact_path=payload["artifact_path"], + arguments=[str(item) for item in payload.get("arguments", [])], + workspace_root=workspace_root, + timeout_seconds=int(payload["timeout_seconds"]), + ) + except Exception as exc: + trace = traceback.format_exc() + result = ExecutionResult( + status="failed", + exit_code=1, + logs=trace.encode("utf-8", errors="replace"), + result=json.dumps( + {"status": "failed", "error": str(exc)}, + ensure_ascii=False, + ).encode("utf-8"), + result_file_name=f"{payload['node_run_id']}-result.json", + result_content_type="application/json", + error_code="WORKER_EXECUTION_FAILED", + error_message=str(exc)[:2000], + ) + if context is None: + context = await self._fallback_execution_context(payload) + + log_id, result_id, upload_error = await self._upload_execution_artifacts( + payload=payload, + context=context, + result=result, + ) + finished_at = utcnow() + duration_ms = max( + 0, + int((finished_at - started_at).total_seconds() * 1000), + ) + error_message = result.error_message + if upload_error: + error_message = ( + f"{error_message}; {upload_error}" + if error_message + else upload_error + )[:2000] + final_status = "failed" if upload_error else result.status + final_error_code = ( + "ARTIFACT_UPLOAD_FAILED" if upload_error else result.error_code + ) + + async with session_scope(self.session_factory) as session: + node_run = await session.scalar( + select(ScheduleNodeRuns) + .where( + ScheduleNodeRuns.node_run_id == payload["node_run_id"], + ) + .with_for_update() + ) + if node_run is None: + raise ValueError("schedule node run disappeared") + inbox = await session.get( + ConsumerInbox, + ("job-workers", event["event_id"]), + with_for_update=True, + ) + if inbox is None: + raise ValueError("job worker inbox record disappeared") + if node_run.node_status not in TERMINAL_NODE_STATES: + node_run.node_status = final_status + node_run.finished_at = finished_at + node_run.duration_ms = duration_ms + node_run.exit_code = result.exit_code + node_run.message = ( + "节点执行成功" + if final_status == "succeeded" + else (error_message or "节点执行失败") + )[:2000] + node_run.metrics_json = { + "log_size_bytes": len(result.logs), + "result_size_bytes": len(result.result), + } + node_run.logs_object_id = log_id + node_run.result_object_id = result_id + node_run.state_version += 1 + await add_outbox_event( + session, + event_type="job.node.finished", + producer="job-worker", + trace_id=event["trace_id"], + aggregate_type="schedule_node_run", + aggregate_id=node_run.node_run_id, + idempotency_key=( + f"{node_run.node_run_id}:{node_run.attempt_no}:finished" + ), + payload={ + "workspace_id": context["workspace_id"], + "run_id": node_run.run_id, + "node_run_id": node_run.node_run_id, + "node_id": node_run.node_id, + "versions_id": node_run.versions_id, + "attempt_no": node_run.attempt_no, + "node_status": node_run.node_status, + "exit_code": node_run.exit_code, + "started_at": event_time( + node_run.started_at or started_at + ), + "finished_at": event_time(finished_at), + "duration_ms": duration_ms, + "logs_object_id": log_id, + "result_object_id": result_id, + "error_code": final_error_code, + "error_message": error_message, + }, + ) + self._finish_inbox(inbox) + + async def _set_node_running( + self, + event: dict[str, Any], + message_id: str, + ) -> bool: + payload = event["payload"] + async with session_scope(self.session_factory) as session: + inbox, should_process = await self._start_inbox( + session, + consumer_name="job-workers", + event_id=event["event_id"], + message_id=message_id, + ) + if not should_process: + return False + node_run = await session.scalar( + select(ScheduleNodeRuns) + .where( + ScheduleNodeRuns.node_run_id == payload["node_run_id"], + ) + .with_for_update() + ) + if node_run is None: + raise ValueError("schedule node run does not exist") + if node_run.node_status in TERMINAL_NODE_STATES: + self._finish_inbox(inbox) + return False + if node_run.node_status == "queued": + node_run.node_status = "running" + node_run.started_at = utcnow() + node_run.message = "Worker 正在执行稳定版本" + node_run.state_version += 1 + return True + + async def _execution_context( + self, + payload: dict[str, Any], + ) -> dict[str, Any]: + async with self.session_factory() as session: + row = ( + await session.execute( + select( + ScheduleNodeRuns, + ScheduleRuns, + Versions, + StorageObjects, + Workspaces, + Schedules, + ) + .join( + ScheduleRuns, + ScheduleRuns.run_id == ScheduleNodeRuns.run_id, + ) + .join( + Versions, + Versions.versions_id == ScheduleNodeRuns.versions_id, + ) + .join( + StorageObjects, + StorageObjects.storage_object_id + == Versions.artifact_object_id, + ) + .join( + Workspaces, + Workspaces.workspace_id == ScheduleRuns.workspace_id, + ) + .join( + Schedules, + Schedules.schedule_id == ScheduleRuns.schedule_id, + ) + .where( + ScheduleNodeRuns.node_run_id == payload["node_run_id"], + ) + ) + ).one_or_none() + if row is None: + raise ValueError("node execution metadata not found") + node_run, run, version, storage, workspace, schedule = row + if storage.object_status != "available": + raise ValueError("stable version artifact is not available") + if storage.storage_backend != "rustfs": + raise ValueError("stable version artifact is not stored in RustFS") + if not storage.bucket_name or not storage.object_key: + raise ValueError("stable version artifact location is incomplete") + user_id = run.triggered_by or schedule.created_by + return { + "node_status": node_run.node_status, + "workspace_id": run.workspace_id, + "workspace_code": workspace.workspace_code, + "user_id": user_id, + "bucket_name": storage.bucket_name, + "object_key": storage.object_key, + "content_hash": version.content_hash, + } + + async def _fallback_execution_context( + self, + payload: dict[str, Any], + ) -> dict[str, Any]: + async with self.session_factory() as session: + row = ( + await session.execute( + select(ScheduleRuns, Workspaces, Schedules) + .join( + Workspaces, + Workspaces.workspace_id == ScheduleRuns.workspace_id, + ) + .join( + Schedules, + Schedules.schedule_id == ScheduleRuns.schedule_id, + ) + .where(ScheduleRuns.run_id == payload["run_id"]) + ) + ).one_or_none() + if row is None: + raise ValueError("schedule run execution context not found") + run, workspace, schedule = row + return { + "workspace_id": run.workspace_id, + "workspace_code": workspace.workspace_code, + "user_id": run.triggered_by or schedule.created_by, + } + + async def _download_artifact( + self, + *, + bucket_name: str, + object_key: str, + content_hash: str, + ) -> bytes: + def read() -> bytes: + response = self.object_store.get_object( + Bucket=bucket_name, + Key=object_key, + ) + body = response["Body"] + try: + return body.read() + finally: + body.close() + + content = await asyncio.to_thread(read) + if hashlib.sha256(content).hexdigest() != content_hash: + raise ValueError("stable version artifact hash mismatch") + return content + + async def _upload_execution_artifacts( + self, + *, + payload: dict[str, Any], + context: dict[str, Any], + result: ExecutionResult, + ) -> tuple[str | None, str | None, str | None]: + log_id: str | None = None + result_id: str | None = None + upload_error: str | None = None + try: + log_object = await self.storage_client.create_object( + workspace_id=context["workspace_id"], + user_id=context["user_id"], + usage_type="run_log", + file_name=f"{payload['node_run_id']}.log", + content_type="text/plain; charset=utf-8", + content=result.logs, + idempotency_key=f"{payload['node_run_id']}:log", + ) + log_id = log_object["storage_object_id"] + result_object = await self.storage_client.create_object( + workspace_id=context["workspace_id"], + user_id=context["user_id"], + usage_type="run_result", + file_name=result.result_file_name, + content_type=result.result_content_type, + content=result.result, + idempotency_key=f"{payload['node_run_id']}:result", + ) + result_id = result_object["storage_object_id"] + except Exception as exc: + upload_error = f"result upload failed: {exc}"[:2000] + LOGGER.exception("failed to upload node execution artifacts") + return log_id, result_id, upload_error + + @staticmethod + def _start_inbox( + session, + *, + consumer_name: str, + event_id: str, + message_id: str, + ) -> tuple[ConsumerInbox, bool]: + item = session.get( + ConsumerInbox, + (consumer_name, event_id), + with_for_update=True, + ) + if item is not None and item.process_status == "succeeded": + return item, False + if item is None: + item = ConsumerInbox( + consumer_name=consumer_name, + event_id=event_id, + process_status="processing", + message_id=message_id, + ) + session.add(item) + else: + item.process_status = "processing" + item.message_id = message_id + item.error_message = None + return item, True + + @staticmethod + def _finish_inbox(item: ConsumerInbox) -> None: + item.process_status = "succeeded" + item.processed_at = utcnow() + item.error_message = None + + +__all__ = ["NodeExecutor"]