refactor: refact schedule module
This commit is contained in:
tao.chen
+3
-2
@@ -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"]
|
||||
|
||||
@@ -15,7 +15,8 @@ dependencies = [
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
common = { path = "../common" }
|
||||
common = { workspace = true }
|
||||
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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()
|
||||
|
||||
|
||||
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
+108
-823
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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"]
|
||||
Reference in New Issue
Block a user