refactor: integrate model platform backend
This commit is contained in:
Winnie
@@ -1,2 +1 @@
|
||||
def main() -> None:
|
||||
print("Hello from schedule!")
|
||||
"""Scheduler operational application."""
|
||||
|
||||
@@ -0,0 +1,202 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path, PurePosixPath
|
||||
|
||||
|
||||
MAX_LOG_BYTES = 4 * 1024 * 1024
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExecutionResult:
|
||||
status: str
|
||||
exit_code: int | None
|
||||
logs: bytes
|
||||
result: bytes
|
||||
result_file_name: str
|
||||
result_content_type: str
|
||||
error_code: str | None = None
|
||||
error_message: str | None = None
|
||||
|
||||
|
||||
def _limited_log(value: str) -> bytes:
|
||||
encoded = value.encode("utf-8", errors="replace")
|
||||
if len(encoded) <= MAX_LOG_BYTES:
|
||||
return encoded or b"execution produced no console output\n"
|
||||
suffix = b"\n[log truncated by scheduler worker]\n"
|
||||
return encoded[: MAX_LOG_BYTES - len(suffix)] + suffix
|
||||
|
||||
|
||||
async def _execute_notebook(
|
||||
artifact: Path,
|
||||
*,
|
||||
artifact_name: str,
|
||||
arguments: list[str],
|
||||
workspace_root: Path,
|
||||
timeout_seconds: int,
|
||||
) -> ExecutionResult:
|
||||
output = artifact.with_name(f"executed-{artifact_name}")
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
sys.executable,
|
||||
"-m",
|
||||
"schedule.notebook_runner",
|
||||
"--input",
|
||||
str(artifact),
|
||||
"--output",
|
||||
str(output),
|
||||
"--workspace",
|
||||
str(workspace_root),
|
||||
"--timeout",
|
||||
str(max(1, timeout_seconds)),
|
||||
"--arguments-json",
|
||||
json.dumps(arguments, ensure_ascii=False),
|
||||
cwd=str(workspace_root),
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.STDOUT,
|
||||
)
|
||||
try:
|
||||
stdout, _ = await asyncio.wait_for(
|
||||
process.communicate(),
|
||||
timeout=max(1, timeout_seconds) + 30,
|
||||
)
|
||||
except TimeoutError:
|
||||
process.kill()
|
||||
stdout, _ = await process.communicate()
|
||||
status = "timed_out"
|
||||
exit_code = None
|
||||
error_code = "NODE_TIMEOUT"
|
||||
error_message = f"notebook exceeded timeout of {timeout_seconds} seconds"
|
||||
else:
|
||||
exit_code = process.returncode
|
||||
status = (
|
||||
"succeeded"
|
||||
if exit_code == 0
|
||||
else "timed_out"
|
||||
if exit_code == 124
|
||||
else "failed"
|
||||
)
|
||||
error_code = (
|
||||
None
|
||||
if status == "succeeded"
|
||||
else "NODE_TIMEOUT"
|
||||
if status == "timed_out"
|
||||
else "NOTEBOOK_EXECUTION_FAILED"
|
||||
)
|
||||
error_message = (
|
||||
None
|
||||
if status == "succeeded"
|
||||
else "Notebook execution timed out"
|
||||
if status == "timed_out"
|
||||
else "Notebook execution failed; see node log"
|
||||
)
|
||||
return ExecutionResult(
|
||||
status=status,
|
||||
exit_code=exit_code,
|
||||
logs=_limited_log(stdout.decode("utf-8", errors="replace")),
|
||||
result=output.read_bytes() if output.is_file() else artifact.read_bytes(),
|
||||
result_file_name=f"executed-{artifact_name}",
|
||||
result_content_type="application/x-ipynb+json",
|
||||
error_code=error_code,
|
||||
error_message=error_message,
|
||||
)
|
||||
|
||||
|
||||
async def _execute_python(
|
||||
artifact: Path,
|
||||
*,
|
||||
arguments: list[str],
|
||||
workspace_root: Path,
|
||||
timeout_seconds: int,
|
||||
) -> ExecutionResult:
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
sys.executable,
|
||||
str(artifact),
|
||||
*arguments,
|
||||
cwd=str(workspace_root),
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.STDOUT,
|
||||
)
|
||||
try:
|
||||
stdout, _ = await asyncio.wait_for(
|
||||
process.communicate(),
|
||||
timeout=max(1, timeout_seconds),
|
||||
)
|
||||
except TimeoutError:
|
||||
process.kill()
|
||||
stdout, _ = await process.communicate()
|
||||
message = f"node exceeded timeout of {timeout_seconds} seconds"
|
||||
result = json.dumps(
|
||||
{"status": "timed_out", "exit_code": None, "message": message},
|
||||
ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
return ExecutionResult(
|
||||
status="timed_out",
|
||||
exit_code=None,
|
||||
logs=_limited_log(
|
||||
stdout.decode("utf-8", errors="replace") + f"\n{message}\n"
|
||||
),
|
||||
result=result,
|
||||
result_file_name=f"{artifact.stem}-result.json",
|
||||
result_content_type="application/json",
|
||||
error_code="NODE_TIMEOUT",
|
||||
error_message=message,
|
||||
)
|
||||
exit_code = process.returncode
|
||||
status = "succeeded" if exit_code == 0 else "failed"
|
||||
message = None if status == "succeeded" else f"process exited with code {exit_code}"
|
||||
result = json.dumps(
|
||||
{"status": status, "exit_code": exit_code},
|
||||
ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
return ExecutionResult(
|
||||
status=status,
|
||||
exit_code=exit_code,
|
||||
logs=_limited_log(stdout.decode("utf-8", errors="replace")),
|
||||
result=result,
|
||||
result_file_name=f"{artifact.stem}-result.json",
|
||||
result_content_type="application/json",
|
||||
error_code=None if status == "succeeded" else "PROCESS_EXIT_NONZERO",
|
||||
error_message=message,
|
||||
)
|
||||
|
||||
|
||||
async def execute_artifact(
|
||||
source: bytes,
|
||||
*,
|
||||
run_id: str,
|
||||
node_run_id: str,
|
||||
script_type: str,
|
||||
artifact_path: str,
|
||||
arguments: list[str],
|
||||
workspace_root: Path,
|
||||
timeout_seconds: int,
|
||||
) -> ExecutionResult:
|
||||
runtime_root = workspace_root / "runtime_tmp" / "schedule-runs"
|
||||
runtime_root.mkdir(parents=True, exist_ok=True)
|
||||
suffix = ".ipynb" if script_type == "notebook" else ".py"
|
||||
raw_name = PurePosixPath(artifact_path.replace("\\", "/")).name
|
||||
artifact_name = raw_name if raw_name.endswith(suffix) else f"artifact{suffix}"
|
||||
prefix = f"{run_id[-6:]}-{node_run_id[-6:]}-"
|
||||
with tempfile.TemporaryDirectory(prefix=prefix, dir=runtime_root) as directory:
|
||||
artifact = Path(directory) / artifact_name
|
||||
artifact.write_bytes(source)
|
||||
if script_type == "notebook":
|
||||
return await _execute_notebook(
|
||||
artifact,
|
||||
artifact_name=artifact_name,
|
||||
arguments=arguments,
|
||||
workspace_root=workspace_root,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
if script_type == "python":
|
||||
return await _execute_python(
|
||||
artifact,
|
||||
arguments=arguments,
|
||||
workspace_root=workspace_root,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
raise ValueError(f"unsupported script_type: {script_type}")
|
||||
@@ -1,5 +1,51 @@
|
||||
# coding=utf-8
|
||||
"""
|
||||
@Time :2026/7/29
|
||||
@Author :tao.chen
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any, AsyncIterator
|
||||
|
||||
from common.db import create_database_engine, create_session_factory
|
||||
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
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
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()
|
||||
service = SchedulerService(
|
||||
session_factory=session_factory,
|
||||
redis=redis,
|
||||
object_store=build_object_store(),
|
||||
storage_client=SchedulerStorageClient(
|
||||
storage_http_client,
|
||||
os.environ["INTERNAL_SERVICE_TOKEN"],
|
||||
),
|
||||
workspace_root=Path(
|
||||
os.getenv("WORKSPACE_ROOT", "/workspace/workspaces")
|
||||
),
|
||||
)
|
||||
app.state.scheduler_service = service
|
||||
await service.start()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await service.close()
|
||||
await storage_http_client.aclose()
|
||||
await redis.aclose()
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
app = create_service_app(
|
||||
os.getenv("SERVICE_NAME", "scheduler-worker"),
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
|
||||
import nbformat
|
||||
from nbclient import NotebookClient
|
||||
|
||||
|
||||
def emit_outputs(notebook: object) -> None:
|
||||
for cell in notebook.cells: # type: ignore[attr-defined]
|
||||
if cell.get("cell_type") != "code":
|
||||
continue
|
||||
for output in cell.get("outputs", []):
|
||||
output_type = output.get("output_type")
|
||||
if output_type == "stream":
|
||||
text = output.get("text", "")
|
||||
print(
|
||||
"".join(text) if isinstance(text, list) else str(text),
|
||||
end="",
|
||||
flush=True,
|
||||
)
|
||||
elif output_type == "error":
|
||||
print(
|
||||
f"{output.get('ename', 'Error')}: "
|
||||
f"{output.get('evalue', '')}",
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--input", required=True)
|
||||
parser.add_argument("--output", required=True)
|
||||
parser.add_argument("--workspace", required=True)
|
||||
parser.add_argument("--timeout", required=True, type=int)
|
||||
parser.add_argument("--arguments-json", default="[]")
|
||||
args = parser.parse_args()
|
||||
|
||||
source = Path(args.input)
|
||||
output = Path(args.output)
|
||||
workspace = Path(args.workspace)
|
||||
arguments = json.loads(args.arguments_json)
|
||||
if not isinstance(arguments, list) or not all(
|
||||
isinstance(item, str) for item in arguments
|
||||
):
|
||||
raise ValueError("arguments-json must contain an array of strings")
|
||||
|
||||
notebook = nbformat.read(source, as_version=4)
|
||||
if arguments:
|
||||
notebook.cells.insert(
|
||||
0,
|
||||
nbformat.v4.new_code_cell(
|
||||
"import sys\n"
|
||||
f"sys.argv = {json.dumps([source.name, *arguments], ensure_ascii=False)}",
|
||||
metadata={"tags": ["injected-parameters"]},
|
||||
),
|
||||
)
|
||||
exit_code = 0
|
||||
try:
|
||||
client = NotebookClient(
|
||||
notebook,
|
||||
timeout=max(1, args.timeout),
|
||||
kernel_name="python3",
|
||||
allow_errors=False,
|
||||
)
|
||||
client.execute(cwd=str(workspace))
|
||||
except Exception as exc:
|
||||
traceback.print_exc()
|
||||
exit_code = 124 if "timeout" in type(exc).__name__.lower() else 1
|
||||
finally:
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
nbformat.write(notebook, output)
|
||||
emit_outputs(notebook)
|
||||
raise SystemExit(exit_code)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,903 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import traceback
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import suppress
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import boto3
|
||||
import httpx
|
||||
from redis.asyncio import Redis
|
||||
from redis.exceptions import ResponseError
|
||||
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.eventing import (
|
||||
STREAM_BY_EVENT_TYPE,
|
||||
add_outbox_event,
|
||||
event_time,
|
||||
utcnow,
|
||||
)
|
||||
from common.ids import new_ulid
|
||||
from common.db.session import session_scope
|
||||
from schedule.execution import ExecutionResult, execute_artifact
|
||||
from schedule.storage_client import SchedulerStorageClient
|
||||
|
||||
|
||||
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",
|
||||
}
|
||||
|
||||
|
||||
class SchedulerService:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
session_factory: async_sessionmaker[AsyncSession],
|
||||
redis: Redis,
|
||||
object_store: Any,
|
||||
storage_client: SchedulerStorageClient,
|
||||
workspace_root: Path,
|
||||
) -> None:
|
||||
self.session_factory = session_factory
|
||||
self.redis = redis
|
||||
self.object_store = object_store
|
||||
self.storage_client = storage_client
|
||||
self.workspace_root = workspace_root
|
||||
self.consumer_name = (
|
||||
os.getenv("SCHEDULER_CONSUMER_NAME")
|
||||
or f"{socket.gethostname()}-{os.getpid()}"
|
||||
)
|
||||
self.tasks: list[asyncio.Task[Any]] = []
|
||||
|
||||
async def start(self) -> None:
|
||||
await self._ensure_group(
|
||||
"stream:scheduler:commands",
|
||||
"schedule-orchestrator",
|
||||
)
|
||||
await self._ensure_group("stream:jobs:execute", "job-workers")
|
||||
await self._ensure_group("stream:jobs:results", "schedule-results")
|
||||
self.tasks = [
|
||||
asyncio.create_task(
|
||||
self._outbox_loop(),
|
||||
name="scheduler-outbox-publisher",
|
||||
),
|
||||
asyncio.create_task(
|
||||
self._consumer_loop(
|
||||
"stream:scheduler:commands",
|
||||
"schedule-orchestrator",
|
||||
self._handle_run_requested,
|
||||
),
|
||||
name="schedule-orchestrator",
|
||||
),
|
||||
asyncio.create_task(
|
||||
self._consumer_loop(
|
||||
"stream:jobs:execute",
|
||||
"job-workers",
|
||||
self._handle_node_execute,
|
||||
),
|
||||
name="job-worker",
|
||||
),
|
||||
asyncio.create_task(
|
||||
self._consumer_loop(
|
||||
"stream:jobs:results",
|
||||
"schedule-results",
|
||||
self._handle_node_finished,
|
||||
),
|
||||
name="schedule-results",
|
||||
),
|
||||
]
|
||||
|
||||
async def close(self) -> None:
|
||||
for task in self.tasks:
|
||||
task.cancel()
|
||||
for task in self.tasks:
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
self.tasks.clear()
|
||||
|
||||
async def _ensure_group(self, stream: str, group: str) -> None:
|
||||
try:
|
||||
await self.redis.xgroup_create(
|
||||
stream,
|
||||
group,
|
||||
id="0-0",
|
||||
mkstream=True,
|
||||
)
|
||||
except ResponseError as exc:
|
||||
if "BUSYGROUP" not in str(exc):
|
||||
raise
|
||||
|
||||
async def _outbox_loop(self) -> None:
|
||||
while True:
|
||||
try:
|
||||
published = await self._publish_outbox_batch()
|
||||
if not published:
|
||||
await asyncio.sleep(0.35)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
LOGGER.exception("outbox publisher iteration failed")
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def _publish_outbox_batch(self) -> int:
|
||||
now = utcnow()
|
||||
async with session_scope(self.session_factory) as session:
|
||||
events = list(
|
||||
(
|
||||
await session.scalars(
|
||||
select(OutboxEvents)
|
||||
.where(
|
||||
OutboxEvents.event_status == "pending",
|
||||
OutboxEvents.available_at <= now,
|
||||
)
|
||||
.order_by(OutboxEvents.created_at)
|
||||
.limit(20)
|
||||
.with_for_update(skip_locked=True)
|
||||
)
|
||||
).all()
|
||||
)
|
||||
for item in events:
|
||||
stream = STREAM_BY_EVENT_TYPE.get(item.event_type)
|
||||
if stream is None:
|
||||
item.event_status = "failed"
|
||||
item.last_error = f"unsupported event type: {item.event_type}"
|
||||
continue
|
||||
try:
|
||||
await self.redis.xadd(
|
||||
stream,
|
||||
{
|
||||
"event": json.dumps(
|
||||
item.payload_json,
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
)
|
||||
},
|
||||
)
|
||||
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]
|
||||
raise
|
||||
return len(events)
|
||||
|
||||
async def _consumer_loop(
|
||||
self,
|
||||
stream: str,
|
||||
group: str,
|
||||
handler: Callable[[dict[str, Any], str], Awaitable[None]],
|
||||
) -> None:
|
||||
while True:
|
||||
try:
|
||||
messages = await self.redis.xreadgroup(
|
||||
group,
|
||||
self.consumer_name,
|
||||
{stream: ">"},
|
||||
count=5,
|
||||
block=1000,
|
||||
)
|
||||
entries: list[tuple[str, dict[str, str]]] = []
|
||||
for _, stream_messages in messages:
|
||||
entries.extend(stream_messages)
|
||||
if not entries:
|
||||
claimed = await self.redis.xautoclaim(
|
||||
stream,
|
||||
group,
|
||||
self.consumer_name,
|
||||
min_idle_time=10_000,
|
||||
start_id="0-0",
|
||||
count=5,
|
||||
)
|
||||
if len(claimed) >= 2:
|
||||
entries.extend(claimed[1])
|
||||
for message_id, fields in entries:
|
||||
try:
|
||||
raw = fields.get("event")
|
||||
if not raw:
|
||||
raise ValueError("stream message has no event field")
|
||||
event = json.loads(raw)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except (ValueError, TypeError, json.JSONDecodeError):
|
||||
LOGGER.exception(
|
||||
"discarding malformed message %s from %s",
|
||||
message_id,
|
||||
stream,
|
||||
)
|
||||
await self.redis.xack(stream, group, message_id)
|
||||
continue
|
||||
try:
|
||||
await handler(event, message_id)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
LOGGER.exception(
|
||||
"consumer %s failed for message %s",
|
||||
group,
|
||||
message_id,
|
||||
)
|
||||
continue
|
||||
await self.redis.xack(stream, group, message_id)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
LOGGER.exception("consumer loop %s failed", group)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
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(
|
||||
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
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def build_redis_client() -> Redis:
|
||||
return Redis(
|
||||
host=os.getenv("REDIS_HOST", "redis"),
|
||||
port=int(os.getenv("REDIS_PORT", "6379")),
|
||||
password=os.getenv("REDIS_PASSWORD") or None,
|
||||
decode_responses=True,
|
||||
)
|
||||
|
||||
|
||||
def build_object_store() -> Any:
|
||||
return boto3.client(
|
||||
"s3",
|
||||
endpoint_url=os.getenv(
|
||||
"RUSTFS_INTERNAL_ENDPOINT",
|
||||
"http://rustfs:9000",
|
||||
),
|
||||
aws_access_key_id=os.environ["RUSTFS_ACCESS_KEY"],
|
||||
aws_secret_access_key=os.environ["RUSTFS_SECRET_KEY"],
|
||||
region_name="us-east-1",
|
||||
)
|
||||
|
||||
|
||||
def build_storage_http_client() -> httpx.AsyncClient:
|
||||
return httpx.AsyncClient(
|
||||
base_url=os.getenv("STORAGE_API_URL", "http://storage_api:8000"),
|
||||
timeout=httpx.Timeout(60.0),
|
||||
)
|
||||
@@ -0,0 +1,41 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
class SchedulerStorageClient:
|
||||
def __init__(self, client: httpx.AsyncClient, service_token: str) -> None:
|
||||
self.client = client
|
||||
self.headers = {"X-Service-Token": service_token}
|
||||
|
||||
async def create_object(
|
||||
self,
|
||||
*,
|
||||
workspace_id: str,
|
||||
user_id: str,
|
||||
usage_type: str,
|
||||
file_name: str,
|
||||
content_type: str,
|
||||
content: bytes,
|
||||
idempotency_key: str,
|
||||
) -> dict[str, Any]:
|
||||
response = await self.client.post(
|
||||
"/internal/v1/objects",
|
||||
headers=self.headers,
|
||||
json={
|
||||
"workspace_id": workspace_id,
|
||||
"user_id": user_id,
|
||||
"usage_type": usage_type,
|
||||
"file_name": file_name,
|
||||
"content_type": content_type,
|
||||
"content_base64": base64.b64encode(content).decode("ascii"),
|
||||
"visibility": "workspace",
|
||||
"is_immutable": True,
|
||||
"idempotency_key": idempotency_key,
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()["data"]
|
||||
Reference in New Issue
Block a user