refactor
This commit is contained in:
@@ -39,10 +39,6 @@ async def lifespan(app: Any) -> AsyncIterator[None]:
|
||||
"RUSTFS_INTERNAL_ENDPOINT",
|
||||
"http://rustfs:9000",
|
||||
),
|
||||
public_endpoint=os.getenv(
|
||||
"RUSTFS_PUBLIC_ENDPOINT",
|
||||
"http://localhost:9000",
|
||||
),
|
||||
access_key=os.environ["RUSTFS_ACCESS_KEY"],
|
||||
secret_key=os.environ["RUSTFS_SECRET_KEY"],
|
||||
)
|
||||
|
||||
@@ -78,7 +78,6 @@ async def create_resource_upload(
|
||||
"expected_size_bytes": payload.expected_size_bytes,
|
||||
"expected_hash": payload.expected_hash,
|
||||
"idempotency_key": idempotency_key,
|
||||
"url_scope": "public",
|
||||
}
|
||||
)
|
||||
return {"request_id": context.request_id, "data": data, "meta": {}}
|
||||
|
||||
@@ -1,14 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from common.clients.base import BaseInternalClient, InternalClientError
|
||||
|
||||
|
||||
class RuntimeClientError(InternalClientError):
|
||||
"""Backward-compatible alias for the runtime error type."""
|
||||
@dataclass(frozen=True)
|
||||
class RuntimeClientError(Exception):
|
||||
status_code: int
|
||||
detail: Any
|
||||
|
||||
|
||||
_RUNTIME_TRANSPORT_ERROR = RuntimeClientError(
|
||||
@@ -22,11 +23,9 @@ _RUNTIME_TRANSPORT_ERROR = RuntimeClientError(
|
||||
)
|
||||
|
||||
|
||||
class RuntimeClient(BaseInternalClient):
|
||||
error_class = RuntimeClientError
|
||||
|
||||
class RuntimeClient:
|
||||
def __init__(self, client: httpx.AsyncClient) -> None:
|
||||
super().__init__(client)
|
||||
self.client = client
|
||||
|
||||
async def _request(
|
||||
self,
|
||||
@@ -34,12 +33,17 @@ class RuntimeClient(BaseInternalClient):
|
||||
path: str,
|
||||
payload: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
return await super()._request(
|
||||
method,
|
||||
path,
|
||||
payload=payload,
|
||||
on_transport_error=_RUNTIME_TRANSPORT_ERROR,
|
||||
)
|
||||
try:
|
||||
response = await self.client.request(method, path, json=payload)
|
||||
except httpx.RequestError as exc:
|
||||
raise _RUNTIME_TRANSPORT_ERROR from exc
|
||||
if response.is_error:
|
||||
try:
|
||||
detail = response.json().get("detail", response.text)
|
||||
except ValueError:
|
||||
detail = response.text
|
||||
raise RuntimeClientError(response.status_code, detail)
|
||||
return response.json()
|
||||
|
||||
async def get_workspace(
|
||||
self,
|
||||
@@ -72,3 +76,6 @@ class RuntimeClient(BaseInternalClient):
|
||||
"/api/v1/jupyter",
|
||||
{"action": "start", "workspace_id": workspace_id},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["RuntimeClient", "RuntimeClientError"]
|
||||
|
||||
@@ -90,9 +90,6 @@ async def lifespan(app: Any) -> AsyncIterator[None]:
|
||||
internal_endpoint=os.getenv(
|
||||
"RUSTFS_INTERNAL_ENDPOINT",
|
||||
"http://rustfs:9000"),
|
||||
public_endpoint=os.getenv(
|
||||
"RUSTFS_PUBLIC_ENDPOINT",
|
||||
"http://localhost:9000"),
|
||||
access_key=os.environ["RUSTFS_ACCESS_KEY"],
|
||||
secret_key=os.environ["RUSTFS_SECRET_KEY"])
|
||||
app.state.default_bucket = os.getenv(
|
||||
@@ -218,18 +215,43 @@ async def create_upload_record(
|
||||
object_key=upload.object_key,
|
||||
content_type=upload.content_type or "application/octet-stream",
|
||||
expected_hash=upload.expected_hash,
|
||||
expires_seconds=900,
|
||||
public=payload.url_scope == "public")
|
||||
expires_seconds=900)
|
||||
presigned_url = request.app.state.object_store.rewrite_to_public_path(
|
||||
url,
|
||||
public_base_url=_public_base_url(request),
|
||||
)
|
||||
return {
|
||||
"upload_id": upload.upload_id,
|
||||
"status": upload.upload_status,
|
||||
"method": "PUT",
|
||||
"presigned_url": url,
|
||||
"presigned_url": presigned_url,
|
||||
"required_headers": headers,
|
||||
"expires_at": upload.expires_at.isoformat(),
|
||||
}
|
||||
|
||||
|
||||
def _public_base_url(request: Request) -> str:
|
||||
"""Return the public base URL the client should use.
|
||||
|
||||
Falls back to the inbound request's ``Host`` header and the scheme
|
||||
Nginx forwards via ``X-Forwarded-Proto`` so the resulting
|
||||
presigned URL always points at the public edge rather than the
|
||||
in-cluster RustFS endpoint.
|
||||
"""
|
||||
forwarded_proto = request.headers.get("x-forwarded-proto", "").strip()
|
||||
scheme = forwarded_proto or request.url.scheme or "http"
|
||||
host = (
|
||||
request.headers.get("x-forwarded-host", "").strip()
|
||||
or request.headers.get("host", "").strip()
|
||||
)
|
||||
if not host:
|
||||
raise HTTPException(
|
||||
status.HTTP_400_BAD_REQUEST,
|
||||
"cannot determine public host for presigned URL",
|
||||
)
|
||||
return f"{scheme}://{host}"
|
||||
|
||||
|
||||
async def complete_upload_record(
|
||||
upload_id: str,
|
||||
payload: CompleteUploadRequest,
|
||||
@@ -408,8 +430,7 @@ async def create_server_object(
|
||||
content_type=payload.content_type,
|
||||
expected_size_bytes=len(content),
|
||||
expected_hash=content_hash,
|
||||
idempotency_key=payload.idempotency_key,
|
||||
url_scope="internal"),
|
||||
idempotency_key=payload.idempotency_key),
|
||||
session,
|
||||
request)
|
||||
if upload_result.get("status") == "completed":
|
||||
@@ -544,10 +565,14 @@ async def create_download_url(
|
||||
object_key=item.object_key,
|
||||
file_name=item.file_name,
|
||||
expires_seconds=payload.expires_seconds)
|
||||
presigned_url = request.app.state.object_store.rewrite_to_public_path(
|
||||
url,
|
||||
public_base_url=_public_base_url(request),
|
||||
)
|
||||
return {
|
||||
"data": {
|
||||
"storage_object_id": item.storage_object_id,
|
||||
"presigned_url": url,
|
||||
"presigned_url": presigned_url,
|
||||
"method": "GET",
|
||||
"expires_in_seconds": payload.expires_seconds,
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ version = "0.2.0"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"SQLAlchemy==2.0.51",
|
||||
"apscheduler>=3.11.3",
|
||||
"asyncmy==0.2.11",
|
||||
"boto3>=1.34,<2",
|
||||
"fastapi==0.116.1",
|
||||
@@ -15,3 +16,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/common"]
|
||||
|
||||
[[tool.uv.index]]
|
||||
url = "https://pypi.tuna.tsinghua.edu.cn/simple/"
|
||||
default = true
|
||||
|
||||
@@ -9,13 +9,6 @@ from common.db.models import OutboxEvents
|
||||
from common.ids import new_ulid
|
||||
|
||||
|
||||
STREAM_BY_EVENT_TYPE = {
|
||||
"schedule.run.requested": "stream:scheduler:commands",
|
||||
"job.node.execute": "stream:jobs:execute",
|
||||
"job.node.finished": "stream:jobs:results",
|
||||
}
|
||||
|
||||
|
||||
def utcnow() -> datetime:
|
||||
return datetime.now(UTC).replace(tzinfo=None)
|
||||
|
||||
@@ -39,28 +32,14 @@ async def add_outbox_event(
|
||||
payload: dict[str, Any],
|
||||
available_at: datetime | None = None,
|
||||
) -> OutboxEvents:
|
||||
if event_type not in STREAM_BY_EVENT_TYPE:
|
||||
raise ValueError(f"unsupported event type: {event_type}")
|
||||
event_id = new_ulid()
|
||||
envelope = {
|
||||
"event_id": event_id,
|
||||
"event_type": event_type,
|
||||
"schema_version": 1,
|
||||
"occurred_at": event_time(),
|
||||
"producer": producer,
|
||||
"trace_id": trace_id,
|
||||
"aggregate_type": aggregate_type,
|
||||
"aggregate_id": aggregate_id,
|
||||
"idempotency_key": idempotency_key,
|
||||
"payload": payload,
|
||||
}
|
||||
item = OutboxEvents(
|
||||
event_id=event_id,
|
||||
aggregate_type=aggregate_type,
|
||||
aggregate_id=aggregate_id,
|
||||
event_type=event_type,
|
||||
schema_version=1,
|
||||
payload_json=envelope,
|
||||
payload_json=payload,
|
||||
event_status="pending",
|
||||
available_at=available_at or utcnow(),
|
||||
retry_count=0,
|
||||
@@ -69,3 +48,7 @@ async def add_outbox_event(
|
||||
)
|
||||
session.add(item)
|
||||
return item
|
||||
|
||||
|
||||
__all__ = ["add_outbox_event", "event_time", "utcnow"]
|
||||
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Shared APScheduler jobstore helpers.
|
||||
|
||||
The same ``apscheduler_jobs`` table is consumed by the FastAPI backend
|
||||
(to enqueue cron jobs through ``add_job``) and by the schedule executor
|
||||
(to run them). Both sides agree on the table name and the URL form
|
||||
that :class:`SQLAlchemyJobStore` expects, so this module is the single
|
||||
source of truth for that contract.
|
||||
|
||||
``apscheduler`` is declared as an optional peer dependency: only the
|
||||
schedule service imports it at runtime. This module therefore exposes
|
||||
the table name and URL helpers but defers actual jobstore construction
|
||||
to the caller.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover - typing only
|
||||
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
|
||||
|
||||
|
||||
JOBSTORE_TABLE = "apscheduler_jobs"
|
||||
|
||||
|
||||
def to_sync_database_url(database_url: str) -> str:
|
||||
"""Rewrite ``mysql+asyncmy://`` to ``mysql+pymysql://``.
|
||||
|
||||
APScheduler's ``SQLAlchemyJobStore`` uses a synchronous engine; the
|
||||
project's default async URL needs to be downgraded before it can
|
||||
drive the jobstore.
|
||||
"""
|
||||
return database_url.replace("mysql+asyncmy://", "mysql+pymysql://", 1)
|
||||
|
||||
|
||||
def build_sqlalchemy_jobstore(
|
||||
database_url: str,
|
||||
*,
|
||||
tablename: str = JOBSTORE_TABLE,
|
||||
) -> "SQLAlchemyJobStore":
|
||||
"""Instantiate a :class:`SQLAlchemyJobStore` for the canonical table.
|
||||
|
||||
The caller is responsible for ensuring APScheduler and its sync
|
||||
driver (``pymysql``) are installed; ``common`` does not depend on
|
||||
either.
|
||||
"""
|
||||
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
|
||||
|
||||
return SQLAlchemyJobStore(
|
||||
url=to_sync_database_url(database_url),
|
||||
tablename=tablename,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"JOBSTORE_TABLE",
|
||||
"build_sqlalchemy_jobstore",
|
||||
"to_sync_database_url",
|
||||
]
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Shared Pydantic base models for service-level request validation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
|
||||
class StrictModel(BaseModel):
|
||||
"""Base model that rejects unknown fields."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
__all__ = ["StrictModel"]
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from typing import Any, BinaryIO
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
import boto3
|
||||
from botocore.client import Config
|
||||
@@ -13,7 +14,6 @@ class RustFSObjectStore:
|
||||
self,
|
||||
*,
|
||||
internal_endpoint: str,
|
||||
public_endpoint: str,
|
||||
access_key: str,
|
||||
secret_key: str,
|
||||
) -> None:
|
||||
@@ -31,11 +31,7 @@ class RustFSObjectStore:
|
||||
endpoint_url=internal_endpoint.rstrip("/"),
|
||||
**common,
|
||||
)
|
||||
self.public = boto3.client(
|
||||
"s3",
|
||||
endpoint_url=public_endpoint.rstrip("/"),
|
||||
**common,
|
||||
)
|
||||
self._internal_endpoint = internal_endpoint.rstrip("/")
|
||||
|
||||
def ensure_bucket(self, bucket_name: str) -> None:
|
||||
try:
|
||||
@@ -54,7 +50,6 @@ class RustFSObjectStore:
|
||||
content_type: str,
|
||||
expected_hash: str | None,
|
||||
expires_seconds: int,
|
||||
public: bool,
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
params: dict[str, Any] = {
|
||||
"Bucket": bucket_name,
|
||||
@@ -65,8 +60,7 @@ class RustFSObjectStore:
|
||||
if expected_hash:
|
||||
params["Metadata"] = {"sha256": expected_hash}
|
||||
headers["x-amz-meta-sha256"] = expected_hash
|
||||
client = self.public if public else self.internal
|
||||
url = client.generate_presigned_url(
|
||||
url = self.internal.generate_presigned_url(
|
||||
"put_object",
|
||||
Params=params,
|
||||
ExpiresIn=expires_seconds,
|
||||
@@ -81,7 +75,7 @@ class RustFSObjectStore:
|
||||
file_name: str,
|
||||
expires_seconds: int,
|
||||
) -> str:
|
||||
return self.public.generate_presigned_url(
|
||||
return self.internal.generate_presigned_url(
|
||||
"get_object",
|
||||
Params={
|
||||
"Bucket": bucket_name,
|
||||
@@ -93,6 +87,36 @@ class RustFSObjectStore:
|
||||
ExpiresIn=expires_seconds,
|
||||
)
|
||||
|
||||
def rewrite_to_public_path(
|
||||
self,
|
||||
url: str,
|
||||
*,
|
||||
public_base_url: str,
|
||||
) -> str:
|
||||
"""Rewrite the host of a presigned URL to the public edge.
|
||||
|
||||
Replaces the scheme + host (and strips any trailing slash) with
|
||||
``public_base_url``; the bucket prefix is moved under the
|
||||
``/storage/`` path so Nginx can forward the call to RustFS
|
||||
without exposing its port.
|
||||
"""
|
||||
parsed = urlsplit(url)
|
||||
public = urlsplit(public_base_url.rstrip("/"))
|
||||
netloc = public.netloc
|
||||
prefix = public.path.rstrip("/")
|
||||
# Presigned URLs are generated against the internal S3 endpoint
|
||||
# and always start with ``/{bucket}/...``; relocate the bucket
|
||||
# segment under ``/storage/...`` so the public edge can route
|
||||
# the request to the right bucket.
|
||||
rewritten_path = f"{prefix}/storage{parsed.path}"
|
||||
return urlunsplit((
|
||||
public.scheme,
|
||||
netloc,
|
||||
rewritten_path,
|
||||
parsed.query,
|
||||
parsed.fragment,
|
||||
))
|
||||
|
||||
def put_bytes(
|
||||
self,
|
||||
*,
|
||||
|
||||
@@ -23,7 +23,6 @@ services:
|
||||
RUNTIME_API_URL: http://runtime:8000
|
||||
RUNTIME_BASE_URL: http://runtime:8000
|
||||
RUSTFS_INTERNAL_ENDPOINT: http://rustfs:9000
|
||||
RUSTFS_PUBLIC_ENDPOINT: http://localhost:${RUSTFS_API_PORT:-9010}
|
||||
RUSTFS_ACCESS_KEY: ${RUSTFS_ACCESS_KEY:-modelplatform}
|
||||
RUSTFS_SECRET_KEY: ${RUSTFS_SECRET_KEY:-modelplatformsecret}
|
||||
RUSTFS_DEFAULT_BUCKET: model-platform
|
||||
|
||||
+182
-175
@@ -5,18 +5,17 @@ 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 datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
import boto3
|
||||
import httpx
|
||||
from redis.asyncio import Redis
|
||||
from redis.exceptions import ResponseError
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
from apscheduler.triggers.cron import CronTrigger
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
|
||||
@@ -30,16 +29,7 @@ from common.db.models import (
|
||||
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
|
||||
from common.scheduler import build_sqlalchemy_jobstore
|
||||
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
@@ -58,203 +48,229 @@ TERMINAL_RUN_STATES = {
|
||||
"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:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
session_factory: async_sessionmaker[AsyncSession],
|
||||
redis: Redis,
|
||||
object_store: Any,
|
||||
storage_client: SchedulerStorageClient,
|
||||
backend_http_client: httpx.AsyncClient,
|
||||
workspace_root: Path,
|
||||
database_url: str,
|
||||
) -> None:
|
||||
self.session_factory = session_factory
|
||||
self.redis = redis
|
||||
self.object_store = object_store
|
||||
self.storage_client = storage_client
|
||||
self.backend_http_client = backend_http_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]] = []
|
||||
self.dispatch_lock = asyncio.Lock()
|
||||
self.scheduler = AsyncIOScheduler(
|
||||
jobstores={
|
||||
"default": build_sqlalchemy_jobstore(database_url)
|
||||
},
|
||||
timezone=UTC,
|
||||
)
|
||||
|
||||
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")
|
||||
global _ACTIVE_SERVICE
|
||||
_ACTIVE_SERVICE = self
|
||||
self.scheduler.start()
|
||||
await self._sync_cron_jobs()
|
||||
self.tasks = [
|
||||
asyncio.create_task(
|
||||
self._outbox_loop(),
|
||||
name="scheduler-outbox-publisher",
|
||||
self._database_event_loop(),
|
||||
name="scheduler-database-events",
|
||||
),
|
||||
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",
|
||||
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 _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:
|
||||
async def _database_event_loop(self) -> None:
|
||||
while True:
|
||||
try:
|
||||
published = await self._publish_outbox_batch()
|
||||
if not published:
|
||||
await asyncio.sleep(0.35)
|
||||
processed = await self.process_pending_events(limit=20)
|
||||
if not processed:
|
||||
await asyncio.sleep(0.25)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
LOGGER.exception("outbox publisher iteration failed")
|
||||
LOGGER.exception("database event loop failed")
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def _publish_outbox_batch(self) -> int:
|
||||
now = utcnow()
|
||||
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:
|
||||
events = list(
|
||||
schedules = list(
|
||||
(
|
||||
await session.scalars(
|
||||
select(OutboxEvents)
|
||||
.where(
|
||||
OutboxEvents.event_status == "pending",
|
||||
OutboxEvents.available_at <= now,
|
||||
select(Schedules).where(
|
||||
Schedules.deleted_at.is_(None),
|
||||
Schedules.enabled == 1,
|
||||
Schedules.trigger_type == "cron",
|
||||
Schedules.cron_expression.is_not(None),
|
||||
)
|
||||
.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,
|
||||
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),
|
||||
)
|
||||
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)
|
||||
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,
|
||||
@@ -874,15 +890,6 @@ class SchedulerService:
|
||||
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",
|
||||
@@ -898,6 +905,6 @@ def build_object_store() -> Any:
|
||||
|
||||
def build_storage_http_client() -> httpx.AsyncClient:
|
||||
return httpx.AsyncClient(
|
||||
base_url=os.getenv("STORAGE_API_URL", "http://storage_api:8000"),
|
||||
base_url=os.getenv("BACKEND_API_URL", "http://backend:8000"),
|
||||
timeout=httpx.Timeout(60.0),
|
||||
)
|
||||
|
||||
@@ -60,6 +60,18 @@ wheels = [
|
||||
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/81/29/5ecc3a15d5a33e31b26c11426c45c501e439cb865d0bff96315d86443b78/appnope-0.1.4-py2.py3-none-any.whl", hash = "sha256:502575ee11cd7a28c0205f379b525beefebab9d161b7c964670864014ed7213c", size = 4321, upload-time = "2024-02-06T09:43:09.663Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "apscheduler"
|
||||
version = "3.11.3"
|
||||
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
|
||||
dependencies = [
|
||||
{ name = "tzlocal" },
|
||||
]
|
||||
sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8c/6b/eeff360196bb20b312c9e762a820fd1b2c6d809466c755ef57863478e454/apscheduler-3.11.3.tar.gz", hash = "sha256:cd2fcc9330039a81a5893472ad49facf23a6d5604cbe1d918c835c6de7834d5a", size = 110312, upload-time = "2026-06-28T19:39:22.493Z" }
|
||||
wheels = [
|
||||
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/42/c9/8638db32514dbb9157b3d82680c6faea89283523edf9ed2415ea3884f2ae/apscheduler-3.11.3-py3-none-any.whl", hash = "sha256:bbeb2ec02d23d3c06a6c07ed7f0f3939ada6680eb121fae809a69bb42c537a30", size = 66024, upload-time = "2026-06-28T19:39:20.982Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "argon2-cffi"
|
||||
version = "25.1.0"
|
||||
@@ -449,6 +461,7 @@ name = "common"
|
||||
version = "0.2.0"
|
||||
source = { editable = "common" }
|
||||
dependencies = [
|
||||
{ name = "apscheduler" },
|
||||
{ name = "asyncmy" },
|
||||
{ name = "boto3" },
|
||||
{ name = "fastapi" },
|
||||
@@ -457,6 +470,7 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "apscheduler", specifier = ">=3.11.3" },
|
||||
{ name = "asyncmy", specifier = "==0.2.11" },
|
||||
{ name = "boto3", specifier = ">=1.34,<2" },
|
||||
{ name = "fastapi", specifier = "==0.116.1" },
|
||||
@@ -2052,6 +2066,18 @@ wheels = [
|
||||
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/e5/6d/b53b99a9f2766d095985947a5782f1702cabb129a34f7a802d7197af832f/tzdata-2026.3-py2.py3-none-any.whl", hash = "sha256:dc096730c87af6cab1b171c9d532be840741ff5d459015e7f6947bd7d7e54931", size = 348168, upload-time = "2026-07-10T08:50:36.46Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tzlocal"
|
||||
version = "5.4.4"
|
||||
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
|
||||
dependencies = [
|
||||
{ name = "tzdata", marker = "sys_platform == 'win32'" },
|
||||
]
|
||||
sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/81/5b/879b2f932adfa7a053c360d50bc896c977fa6426109185f7c12ebdd0cb9d/tzlocal-5.4.4.tar.gz", hash = "sha256:8dbb8660838688a7b6ba4fed31d18dedf842afb4d47ca050d6d891c2c15f3be4", size = 31170, upload-time = "2026-06-29T08:03:40.026Z" }
|
||||
wheels = [
|
||||
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/9e/a4/017a7a6cbe387d961a688ec31364ae60a5c4e22c96ae9921b79a947c855d/tzlocal-5.4.4-py3-none-any.whl", hash = "sha256:aae09f0126a8a86fa736be266eb4a471380d26a0de3bc14844e7821fee3e2a15", size = 18115, upload-time = "2026-06-29T08:03:38.666Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "uri-template"
|
||||
version = "1.3.0"
|
||||
|
||||
Reference in New Issue
Block a user