重构模型平台前后端并移除Redis依赖
This commit is contained in:
Winnie
@@ -0,0 +1 @@
|
||||
3.12
|
||||
@@ -0,0 +1,11 @@
|
||||
FROM python:3.12-slim-bookworm
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1 PYTHONPATH=/app
|
||||
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
|
||||
|
||||
EXPOSE 8000
|
||||
CMD ["uvicorn", "schedule.main:app", "--host", "0.0.0.0", "--port", "8000"]
|
||||
@@ -0,0 +1,9 @@
|
||||
# Schedule Executor
|
||||
|
||||
独立调度执行服务，内置 APScheduler。
|
||||
|
||||
- Cron 任务持久化到 MySQL 的 `apscheduler_jobs` 表；
|
||||
- FastAPI Backend 创建运行记录和 Outbox 事件后，通过 HTTP 尝试立即推送；
|
||||
- HTTP 推送失败时，Executor 继续轮询 MySQL `outbox_events`，保证任务不会丢失；
|
||||
- Executor 负责 DAG 节点派发、稳定版本执行、重试、状态推进和结果回写；
|
||||
- 不依赖 Redis，MySQL 是调度状态与幂等状态的唯一权威。
|
||||
@@ -0,0 +1,25 @@
|
||||
[project]
|
||||
name = "schedule"
|
||||
version = "0.2.0"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"common",
|
||||
"fastapi==0.116.1",
|
||||
"uvicorn[standard]==0.35.0",
|
||||
"httpx==0.28.1",
|
||||
"apscheduler==3.11.3",
|
||||
"pymysql==1.2.0",
|
||||
"nbclient==0.10.2",
|
||||
"nbformat==5.10.4",
|
||||
"ipykernel==6.29.5",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
common = { path = "../common" }
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/schedule"]
|
||||
@@ -0,0 +1 @@
|
||||
"""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}")
|
||||
@@ -0,0 +1,77 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import secrets
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any, AsyncIterator
|
||||
|
||||
from fastapi import Depends, Header, HTTPException, Request, status
|
||||
|
||||
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_storage_http_client,
|
||||
)
|
||||
from schedule.storage_client import SchedulerStorageClient
|
||||
|
||||
|
||||
def verify_internal_service(
|
||||
x_service_token: str = Header(alias="X-Service-Token"),
|
||||
) -> None:
|
||||
expected = os.environ.get("INTERNAL_SERVICE_TOKEN", "")
|
||||
if not expected or not secrets.compare_digest(expected, x_service_token):
|
||||
raise HTTPException(
|
||||
status.HTTP_401_UNAUTHORIZED,
|
||||
"invalid internal service identity",
|
||||
)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: Any) -> AsyncIterator[None]:
|
||||
database_url = os.environ["DATABASE_URL"]
|
||||
engine = create_database_engine(database_url)
|
||||
session_factory = create_session_factory(engine)
|
||||
backend_http_client = build_storage_http_client()
|
||||
service = SchedulerService(
|
||||
session_factory=session_factory,
|
||||
object_store=build_object_store(),
|
||||
storage_client=SchedulerStorageClient(
|
||||
backend_http_client,
|
||||
os.environ["INTERNAL_SERVICE_TOKEN"],
|
||||
),
|
||||
backend_http_client=backend_http_client,
|
||||
workspace_root=Path(
|
||||
os.getenv("WORKSPACE_ROOT", "/workspace/workspaces")
|
||||
),
|
||||
database_url=database_url,
|
||||
)
|
||||
app.state.scheduler_service = service
|
||||
await service.start()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await service.close()
|
||||
await backend_http_client.aclose()
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
app = create_service_app(
|
||||
os.getenv("SERVICE_NAME", "schedule-executor"),
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
|
||||
@app.post(
|
||||
"/internal/v1/runs/{run_id}/dispatch",
|
||||
dependencies=[Depends(verify_internal_service)],
|
||||
)
|
||||
async def dispatch_run(run_id: str, request: Request) -> dict[str, Any]:
|
||||
processed = await request.app.state.scheduler_service.dispatch_run(run_id)
|
||||
return {
|
||||
"status": "accepted",
|
||||
"run_id": run_id,
|
||||
"processed_events": processed,
|
||||
}
|
||||
@@ -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,922 @@
|
||||
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 typing import Any
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
import boto3
|
||||
import httpx
|
||||
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
|
||||
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.eventing import 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",
|
||||
}
|
||||
|
||||
_ACTIVE_SERVICE: "SchedulerService | None" = None
|
||||
|
||||
|
||||
def _sync_database_url(value: str) -> str:
|
||||
return value.replace("mysql+asyncmy://", "mysql+pymysql://", 1)
|
||||
|
||||
|
||||
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:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
session_factory: async_sessionmaker[AsyncSession],
|
||||
object_store: Any,
|
||||
storage_client: SchedulerStorageClient,
|
||||
backend_http_client: httpx.AsyncClient,
|
||||
workspace_root: Path,
|
||||
database_url: str,
|
||||
) -> None:
|
||||
self.session_factory = session_factory
|
||||
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": SQLAlchemyJobStore(
|
||||
url=_sync_database_url(database_url),
|
||||
tablename="apscheduler_jobs",
|
||||
)
|
||||
},
|
||||
timezone=UTC,
|
||||
)
|
||||
|
||||
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",
|
||||
),
|
||||
]
|
||||
|
||||
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)
|
||||
|
||||
async def trigger_schedule(self, schedule_id: str) -> None:
|
||||
async with self.session_factory() as session:
|
||||
item = await session.get(Schedules, schedule_id)
|
||||
if (
|
||||
item is None
|
||||
or item.deleted_at is not None
|
||||
or not item.enabled
|
||||
or item.trigger_type != "cron"
|
||||
):
|
||||
return
|
||||
user_id = item.created_by
|
||||
workspace_id = item.workspace_id
|
||||
now = datetime.now(UTC)
|
||||
idempotency_key = (
|
||||
f"cron:{schedule_id}:{now.strftime('%Y%m%d%H%M')}"
|
||||
)
|
||||
response = await self.backend_http_client.post(
|
||||
f"/api/v1/schedules/{schedule_id}/run",
|
||||
headers={
|
||||
"X-User-ID": user_id,
|
||||
"X-Workspace-ID": workspace_id,
|
||||
"X-Request-ID": new_ulid(),
|
||||
"Idempotency-Key": idempotency_key,
|
||||
},
|
||||
json={"reason": "cron"},
|
||||
)
|
||||
if response.is_error:
|
||||
raise RuntimeError(
|
||||
f"backend rejected cron run: {response.status_code} "
|
||||
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(
|
||||
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_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("BACKEND_API_URL", "http://backend: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