feat: run dataflows through durable workers
This commit is contained in:
@@ -0,0 +1,631 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timedelta, timezone
|
||||
import logging
|
||||
import socket
|
||||
|
||||
from sqlalchemy import func, or_, select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from govoplan_core.auth import ApiPrincipal
|
||||
from govoplan_core.core.automation import (
|
||||
AutomationPrincipalRequest,
|
||||
automation_principal_provider,
|
||||
)
|
||||
from govoplan_core.core.notifications import (
|
||||
NotificationDispatchRequest,
|
||||
notification_dispatch_provider,
|
||||
)
|
||||
from govoplan_core.db.base import utcnow
|
||||
from govoplan_dataflow.backend.db.models import (
|
||||
DataflowPipeline,
|
||||
DataflowPipelineRevision,
|
||||
DataflowRun,
|
||||
DataflowTrigger,
|
||||
)
|
||||
from govoplan_dataflow.backend.governance import require_definition_action
|
||||
from govoplan_dataflow.backend.service import (
|
||||
_execute_pipeline_run,
|
||||
_require_deployed_revision,
|
||||
pipeline_run_request,
|
||||
)
|
||||
|
||||
|
||||
MAX_ACTIVE_RUNS_PER_TENANT = 4
|
||||
DEFAULT_LEASE_SECONDS = 120
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DataflowWorkerError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def dispatch_pending_runs(
|
||||
session: Session,
|
||||
*,
|
||||
registry: object | None,
|
||||
now: datetime | None = None,
|
||||
limit: int = 10,
|
||||
worker_id: str | None = None,
|
||||
) -> dict[str, object]:
|
||||
current = _as_utc(now or utcnow())
|
||||
resolved_worker_id = (
|
||||
str(worker_id).strip()
|
||||
if worker_id and str(worker_id).strip()
|
||||
else socket.gethostname()
|
||||
)
|
||||
recovered = _recover_expired_leases(session, now=current)
|
||||
session.commit()
|
||||
summary: dict[str, object] = {
|
||||
"claimed": 0,
|
||||
"succeeded": 0,
|
||||
"retrying": 0,
|
||||
"failed": 0,
|
||||
"cancelled": 0,
|
||||
"recovered": recovered,
|
||||
"runs": [],
|
||||
}
|
||||
for _ in range(max(1, min(int(limit), 100))):
|
||||
run_id = _claim_next_run(
|
||||
session,
|
||||
now=current,
|
||||
worker_id=resolved_worker_id,
|
||||
)
|
||||
if run_id is None:
|
||||
session.rollback()
|
||||
break
|
||||
session.commit()
|
||||
summary["claimed"] = int(summary["claimed"]) + 1
|
||||
outcome = _execute_claimed_run(
|
||||
session,
|
||||
run_id=run_id,
|
||||
registry=registry,
|
||||
now=current,
|
||||
)
|
||||
session.commit()
|
||||
summary[outcome] = int(summary[outcome]) + 1
|
||||
runs = list(summary["runs"])
|
||||
runs.append({"ref": f"dataflow-run:{run_id}", "status": outcome})
|
||||
summary["runs"] = runs
|
||||
return summary
|
||||
|
||||
|
||||
def purge_expired_runs(
|
||||
session: Session,
|
||||
*,
|
||||
now: datetime | None = None,
|
||||
limit: int = 500,
|
||||
) -> dict[str, object]:
|
||||
current = _as_utc(now or utcnow())
|
||||
runs = list(
|
||||
session.scalars(
|
||||
select(DataflowRun)
|
||||
.where(
|
||||
DataflowRun.status.in_(
|
||||
("succeeded", "failed", "cancelled")
|
||||
),
|
||||
DataflowRun.retention_until.is_not(None),
|
||||
DataflowRun.retention_until <= current,
|
||||
DataflowRun.purged_at.is_(None),
|
||||
)
|
||||
.order_by(DataflowRun.retention_until, DataflowRun.id)
|
||||
.limit(max(1, min(int(limit), 5_000)))
|
||||
.with_for_update(skip_locked=True)
|
||||
)
|
||||
)
|
||||
for run in runs:
|
||||
run.request_ = {}
|
||||
run.source_fingerprints = []
|
||||
run.diagnostics = []
|
||||
authorization = dict(run.authorization_)
|
||||
authorization.pop("submitted_principal", None)
|
||||
authorization["personal_data_purged"] = True
|
||||
run.authorization_ = authorization
|
||||
run.purged_at = current
|
||||
session.flush()
|
||||
return {
|
||||
"purged": len(runs),
|
||||
"run_refs": [f"dataflow-run:{run.id}" for run in runs],
|
||||
}
|
||||
|
||||
|
||||
def run_metrics(
|
||||
session: Session,
|
||||
*,
|
||||
tenant_id: str,
|
||||
now: datetime | None = None,
|
||||
) -> dict[str, object]:
|
||||
current = _as_utc(now or utcnow())
|
||||
rows = session.execute(
|
||||
select(DataflowRun.status, func.count())
|
||||
.where(DataflowRun.tenant_id == tenant_id)
|
||||
.group_by(DataflowRun.status)
|
||||
).all()
|
||||
statuses = {str(status): int(count) for status, count in rows}
|
||||
recent = list(
|
||||
session.execute(
|
||||
select(
|
||||
DataflowRun.status,
|
||||
DataflowRun.started_at,
|
||||
DataflowRun.finished_at,
|
||||
).where(
|
||||
DataflowRun.tenant_id == tenant_id,
|
||||
DataflowRun.finished_at.is_not(None),
|
||||
DataflowRun.finished_at >= current - timedelta(hours=24),
|
||||
)
|
||||
)
|
||||
)
|
||||
durations = [
|
||||
(
|
||||
_as_utc(finished_at) - _as_utc(started_at)
|
||||
).total_seconds()
|
||||
for _status, started_at, finished_at in recent
|
||||
if finished_at is not None and started_at is not None
|
||||
]
|
||||
oldest_queued_at = session.scalar(
|
||||
select(func.min(DataflowRun.queued_at)).where(
|
||||
DataflowRun.tenant_id == tenant_id,
|
||||
DataflowRun.status.in_(("queued", "retrying")),
|
||||
)
|
||||
)
|
||||
return {
|
||||
"statuses": statuses,
|
||||
"active": sum(
|
||||
statuses.get(status, 0)
|
||||
for status in ("queued", "retrying", "running")
|
||||
),
|
||||
"completed_last_24_hours": len(recent),
|
||||
"failed_last_24_hours": sum(
|
||||
1 for status, _started_at, _finished_at in recent
|
||||
if status == "failed"
|
||||
),
|
||||
"average_duration_seconds": (
|
||||
round(sum(durations) / len(durations), 3)
|
||||
if durations
|
||||
else None
|
||||
),
|
||||
"oldest_queued_at": oldest_queued_at,
|
||||
}
|
||||
|
||||
|
||||
def _recover_expired_leases(
|
||||
session: Session,
|
||||
*,
|
||||
now: datetime,
|
||||
) -> int:
|
||||
runs = list(
|
||||
session.scalars(
|
||||
select(DataflowRun)
|
||||
.where(
|
||||
DataflowRun.status == "running",
|
||||
DataflowRun.lease_expires_at.is_not(None),
|
||||
DataflowRun.lease_expires_at < now,
|
||||
)
|
||||
.with_for_update(skip_locked=True)
|
||||
)
|
||||
)
|
||||
for run in runs:
|
||||
run.worker_id = None
|
||||
run.claimed_at = None
|
||||
run.lease_expires_at = None
|
||||
run.heartbeat_at = None
|
||||
if run.cancellation_requested_at is not None:
|
||||
_cancel_run(run, now=now)
|
||||
elif run.attempts < run.max_attempts:
|
||||
run.status = "retrying"
|
||||
run.available_at = now
|
||||
run.progress_phase = "retrying_after_worker_loss"
|
||||
run.error = "The previous worker lease expired."
|
||||
else:
|
||||
_fail_run(
|
||||
run,
|
||||
now=now,
|
||||
message="The worker lease expired after the final attempt.",
|
||||
)
|
||||
return len(runs)
|
||||
|
||||
|
||||
def _claim_next_run(
|
||||
session: Session,
|
||||
*,
|
||||
now: datetime,
|
||||
worker_id: str,
|
||||
) -> str | None:
|
||||
candidates = list(
|
||||
session.scalars(
|
||||
select(DataflowRun)
|
||||
.where(
|
||||
DataflowRun.status.in_(("queued", "retrying")),
|
||||
or_(
|
||||
DataflowRun.available_at.is_(None),
|
||||
DataflowRun.available_at <= now,
|
||||
),
|
||||
)
|
||||
.order_by(DataflowRun.available_at, DataflowRun.created_at)
|
||||
.limit(40)
|
||||
.with_for_update(skip_locked=True)
|
||||
)
|
||||
)
|
||||
for run in candidates:
|
||||
active = int(
|
||||
session.scalar(
|
||||
select(func.count())
|
||||
.select_from(DataflowRun)
|
||||
.where(
|
||||
DataflowRun.tenant_id == run.tenant_id,
|
||||
DataflowRun.status == "running",
|
||||
)
|
||||
)
|
||||
or 0
|
||||
)
|
||||
if active >= MAX_ACTIVE_RUNS_PER_TENANT:
|
||||
continue
|
||||
run.status = "running"
|
||||
run.attempts += 1
|
||||
run.worker_id = worker_id[:255]
|
||||
run.claimed_at = now
|
||||
run.heartbeat_at = now
|
||||
run.lease_expires_at = now + timedelta(
|
||||
seconds=_lease_seconds(run)
|
||||
)
|
||||
run.started_at = run.started_at or now
|
||||
run.finished_at = None
|
||||
run.progress_percent = max(run.progress_percent, 5)
|
||||
run.progress_phase = "authorizing"
|
||||
session.flush()
|
||||
return run.id
|
||||
return None
|
||||
|
||||
|
||||
def _execute_claimed_run(
|
||||
session: Session,
|
||||
*,
|
||||
run_id: str,
|
||||
registry: object | None,
|
||||
now: datetime,
|
||||
) -> str:
|
||||
run = session.get(DataflowRun, run_id)
|
||||
if run is None:
|
||||
return "failed"
|
||||
if run.cancellation_requested_at is not None:
|
||||
_cancel_run(run, now=now)
|
||||
return "cancelled"
|
||||
try:
|
||||
principal, provenance = _resolve_principal(
|
||||
session,
|
||||
run=run,
|
||||
registry=registry,
|
||||
)
|
||||
run.authorization_ = {
|
||||
**dict(run.authorization_),
|
||||
"last_resolution": provenance,
|
||||
"resolved_at": now.isoformat(),
|
||||
}
|
||||
pipeline = session.get(DataflowPipeline, run.pipeline_id)
|
||||
revision = session.get(
|
||||
DataflowPipelineRevision,
|
||||
run.pipeline_revision_id,
|
||||
)
|
||||
if pipeline is None or revision is None:
|
||||
raise DataflowWorkerError(
|
||||
"The pipeline or pinned revision no longer exists."
|
||||
)
|
||||
action = (
|
||||
"run"
|
||||
if run.invocation_kind in {"manual", "api", "backfill"}
|
||||
else "automate"
|
||||
)
|
||||
require_definition_action(
|
||||
pipeline,
|
||||
principal=principal,
|
||||
registry=registry,
|
||||
action=action,
|
||||
)
|
||||
_require_deployed_revision(
|
||||
session,
|
||||
tenant_id=run.tenant_id,
|
||||
pipeline=pipeline,
|
||||
revision=revision,
|
||||
environment=run.environment,
|
||||
)
|
||||
_notify_run(
|
||||
session,
|
||||
registry=registry,
|
||||
run=run,
|
||||
event_kind="dataflow.run.started",
|
||||
subject=f"Dataflow run started: {pipeline.name}",
|
||||
)
|
||||
retryable = _execute_pipeline_run(
|
||||
session,
|
||||
run=run,
|
||||
pipeline=pipeline,
|
||||
revision=revision,
|
||||
request=pipeline_run_request(run),
|
||||
principal=principal,
|
||||
registry=registry,
|
||||
)
|
||||
if (
|
||||
run.status == "failed"
|
||||
and retryable
|
||||
and run.attempts < run.max_attempts
|
||||
):
|
||||
_retry_run(run, now=now)
|
||||
_finish_claim(run)
|
||||
_update_trigger_status(session, run)
|
||||
_notify_terminal_run(
|
||||
session,
|
||||
registry=registry,
|
||||
run=run,
|
||||
pipeline_name=pipeline.name,
|
||||
)
|
||||
return run.status
|
||||
except (DataflowWorkerError, PermissionError, ValueError) as exc:
|
||||
_fail_run(run, now=now, message=str(exc))
|
||||
_finish_claim(run)
|
||||
_update_trigger_status(session, run)
|
||||
_notify_terminal_run(
|
||||
session,
|
||||
registry=registry,
|
||||
run=run,
|
||||
pipeline_name="Dataflow",
|
||||
)
|
||||
return "failed"
|
||||
except Exception as exc:
|
||||
logger.exception(
|
||||
"Unexpected Dataflow worker failure for run %s",
|
||||
run_id,
|
||||
)
|
||||
session.rollback()
|
||||
run = session.get(DataflowRun, run_id)
|
||||
if run is None:
|
||||
return "failed"
|
||||
_fail_run(
|
||||
run,
|
||||
now=now,
|
||||
message=(
|
||||
"Unexpected worker failure "
|
||||
f"({type(exc).__name__})."
|
||||
),
|
||||
)
|
||||
_finish_claim(run)
|
||||
return "failed"
|
||||
|
||||
|
||||
def _resolve_principal(
|
||||
session: Session,
|
||||
*,
|
||||
run: DataflowRun,
|
||||
registry: object | None,
|
||||
) -> tuple[ApiPrincipal, Mapping[str, object]]:
|
||||
provider = automation_principal_provider(registry)
|
||||
if provider is None:
|
||||
raise DataflowWorkerError(
|
||||
"Queued Dataflow execution requires the automation-principal "
|
||||
"capability."
|
||||
)
|
||||
authorization = dict(run.authorization_)
|
||||
subject_kind = str(
|
||||
authorization.get("subject_kind") or "delegated_user"
|
||||
)
|
||||
common = {
|
||||
"tenant_id": run.tenant_id,
|
||||
"authorization_ref": str(
|
||||
authorization.get("authorization_ref")
|
||||
or f"dataflow-run:{run.id}"
|
||||
),
|
||||
"grant_scopes": tuple(
|
||||
str(scope)
|
||||
for scope in authorization.get("grant_scopes") or ()
|
||||
),
|
||||
"context": {
|
||||
"run_ref": f"dataflow-run:{run.id}",
|
||||
"pipeline_ref": f"pipeline:{run.pipeline_id}",
|
||||
"environment": run.environment,
|
||||
"attempt": run.attempts,
|
||||
},
|
||||
}
|
||||
if subject_kind == "service_account":
|
||||
service_account_id = str(
|
||||
authorization.get("service_account_id") or ""
|
||||
).strip()
|
||||
request = AutomationPrincipalRequest.service_account(
|
||||
service_account_id=service_account_id,
|
||||
**common,
|
||||
)
|
||||
else:
|
||||
request = AutomationPrincipalRequest.delegated_user(
|
||||
account_id=str(authorization.get("account_id") or ""),
|
||||
membership_id=str(
|
||||
authorization.get("membership_id") or ""
|
||||
),
|
||||
**common,
|
||||
)
|
||||
resolution = provider.resolve_automation_principal(
|
||||
session,
|
||||
request=request,
|
||||
)
|
||||
if not resolution.allowed or not isinstance(
|
||||
resolution.principal,
|
||||
ApiPrincipal,
|
||||
):
|
||||
raise DataflowWorkerError(
|
||||
resolution.reason or "Dataflow run authorization was denied."
|
||||
)
|
||||
return resolution.principal, resolution.provenance
|
||||
|
||||
|
||||
def _retry_run(run: DataflowRun, *, now: datetime) -> None:
|
||||
delay_seconds = min(15 * (2 ** max(0, run.attempts - 1)), 900)
|
||||
run.status = "retrying"
|
||||
run.available_at = now + timedelta(seconds=delay_seconds)
|
||||
run.finished_at = None
|
||||
run.progress_percent = 0
|
||||
run.progress_phase = "retrying"
|
||||
|
||||
|
||||
def _finish_claim(run: DataflowRun) -> None:
|
||||
run.worker_id = None
|
||||
run.claimed_at = None
|
||||
run.lease_expires_at = None
|
||||
run.heartbeat_at = None
|
||||
|
||||
|
||||
def _cancel_run(run: DataflowRun, *, now: datetime) -> None:
|
||||
run.status = "cancelled"
|
||||
run.finished_at = now
|
||||
run.error = "Cancelled by request."
|
||||
run.progress_phase = "cancelled"
|
||||
_finish_claim(run)
|
||||
|
||||
|
||||
def _fail_run(
|
||||
run: DataflowRun,
|
||||
*,
|
||||
now: datetime,
|
||||
message: str,
|
||||
) -> None:
|
||||
run.status = "failed"
|
||||
run.finished_at = now
|
||||
run.error = message
|
||||
run.progress_phase = "failed"
|
||||
|
||||
|
||||
def _update_trigger_status(session: Session, run: DataflowRun) -> None:
|
||||
if not run.trigger_id:
|
||||
return
|
||||
trigger = session.get(DataflowTrigger, run.trigger_id)
|
||||
if trigger is None:
|
||||
return
|
||||
trigger.last_status = run.status
|
||||
trigger.last_error = run.error
|
||||
|
||||
|
||||
def _notify_terminal_run(
|
||||
session: Session,
|
||||
*,
|
||||
registry: object | None,
|
||||
run: DataflowRun,
|
||||
pipeline_name: str,
|
||||
) -> None:
|
||||
if run.status == "retrying":
|
||||
event_kind = "dataflow.run.retrying"
|
||||
subject = f"Dataflow run will retry: {pipeline_name}"
|
||||
elif run.status == "succeeded":
|
||||
event_kind = "dataflow.run.completed"
|
||||
subject = f"Dataflow run completed: {pipeline_name}"
|
||||
elif run.status == "cancelled":
|
||||
event_kind = "dataflow.run.cancelled"
|
||||
subject = f"Dataflow run cancelled: {pipeline_name}"
|
||||
else:
|
||||
event_kind = "dataflow.run.failed"
|
||||
subject = f"Dataflow run failed: {pipeline_name}"
|
||||
_notify_run(
|
||||
session,
|
||||
registry=registry,
|
||||
run=run,
|
||||
event_kind=event_kind,
|
||||
subject=subject,
|
||||
)
|
||||
|
||||
|
||||
def _notify_run(
|
||||
session: Session,
|
||||
*,
|
||||
registry: object | None,
|
||||
run: DataflowRun,
|
||||
event_kind: str,
|
||||
subject: str,
|
||||
) -> None:
|
||||
provider = notification_dispatch_provider(registry)
|
||||
authorization = dict(run.authorization_)
|
||||
recipient_id = str(authorization.get("account_id") or "").strip()
|
||||
if provider is None or not recipient_id:
|
||||
return
|
||||
try:
|
||||
provider.enqueue_notification(
|
||||
session,
|
||||
NotificationDispatchRequest(
|
||||
tenant_id=run.tenant_id,
|
||||
source_module="dataflow",
|
||||
source_resource_type="dataflow_run",
|
||||
source_resource_id=run.id,
|
||||
event_kind=event_kind,
|
||||
recipient_type="account",
|
||||
recipient_id=recipient_id,
|
||||
subject=subject,
|
||||
body_text=run.error,
|
||||
action_url="/dataflow",
|
||||
payload={
|
||||
"run_ref": f"dataflow-run:{run.id}",
|
||||
"pipeline_ref": f"pipeline:{run.pipeline_id}",
|
||||
"status": run.status,
|
||||
"environment": run.environment,
|
||||
"attempt": run.attempts,
|
||||
},
|
||||
),
|
||||
)
|
||||
except Exception:
|
||||
# Notification delivery is optional and cannot change run outcome.
|
||||
logger.warning(
|
||||
"Could not enqueue Dataflow notification for run %s",
|
||||
run.id,
|
||||
exc_info=True,
|
||||
)
|
||||
return
|
||||
|
||||
|
||||
def _lease_seconds(run: DataflowRun) -> int:
|
||||
budget = dict(run.resource_budget)
|
||||
wall_seconds = float(budget.get("max_wall_seconds") or 30.0)
|
||||
return max(DEFAULT_LEASE_SECONDS, int(wall_seconds) + 60)
|
||||
|
||||
|
||||
def _as_utc(value: datetime | None) -> datetime:
|
||||
if value is None:
|
||||
return datetime.now(timezone.utc)
|
||||
if value.tzinfo is None:
|
||||
return value.replace(tzinfo=timezone.utc)
|
||||
return value.astimezone(timezone.utc)
|
||||
|
||||
|
||||
class SqlDataflowRunWorker:
|
||||
def __init__(self, *, registry: object | None = None) -> None:
|
||||
self._registry = registry
|
||||
|
||||
def dispatch_pending(
|
||||
self,
|
||||
session: object,
|
||||
*,
|
||||
now: datetime | None = None,
|
||||
limit: int = 10,
|
||||
worker_id: str | None = None,
|
||||
) -> Mapping[str, object]:
|
||||
if not isinstance(session, Session):
|
||||
raise TypeError("Dataflow run dispatch requires a Session.")
|
||||
return dispatch_pending_runs(
|
||||
session,
|
||||
registry=self._registry,
|
||||
now=now,
|
||||
limit=limit,
|
||||
worker_id=worker_id,
|
||||
)
|
||||
|
||||
def purge_expired(
|
||||
self,
|
||||
session: object,
|
||||
*,
|
||||
now: datetime | None = None,
|
||||
limit: int = 500,
|
||||
) -> Mapping[str, object]:
|
||||
if not isinstance(session, Session):
|
||||
raise TypeError("Dataflow retention requires a Session.")
|
||||
return purge_expired_runs(session, now=now, limit=limit)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SqlDataflowRunWorker",
|
||||
"dispatch_pending_runs",
|
||||
"purge_expired_runs",
|
||||
"run_metrics",
|
||||
]
|
||||
Reference in New Issue
Block a user