File indexing completed on 2026-07-21 08:46:27
0001 """Publish the testbed datataking state projection to Snapper."""
0002
0003 from dataclasses import dataclass
0004 from datetime import datetime, timezone as datetime_timezone
0005 from typing import Optional
0006
0007 from django.db import transaction
0008 from django.db.models import OuterRef, Subquery
0009 from django.db.models.fields.json import KeyTextTransform
0010 from django.utils import timezone
0011
0012 from snapper_ai.services import (
0013 ComponentUpdate,
0014 publish_component,
0015 register_component,
0016 )
0017
0018 from .models import RunState
0019 from .workflow_models import WorkflowExecution
0020
0021
0022 PUBLISHER_IDENTITY = "swf-monitor:run-state"
0023 ASSESSMENT_POLICY_VERSION = "swf-datataking-state-v1"
0024 EVENT_RESOLVER = "swf-testbed-system-state-events"
0025 MAX_NAMESPACES = 128
0026 MAX_SERIALIZED_BYTES = 64 * 1024
0027
0028 DATATAKING_REGISTRATION = {
0029 "title": "Testbed datataking state by namespace",
0030 "description": (
0031 "The latest run state for each namespace sharing the testbed platform, "
0032 "forming its datataking lanes in the E0-E1 global state at an instant."
0033 ),
0034 "visibility": "public",
0035 "owning_subsystem": "SWF testbed datataking state",
0036 "assessment_policy": ASSESSMENT_POLICY_VERSION,
0037 "max_serialized_bytes": MAX_SERIALIZED_BYTES,
0038 "quantities": {
0039 "namespaces": {
0040 "path": "namespaces",
0041 "type": "object",
0042 "required": True,
0043 "kind": "bounded_map",
0044 "max_items": MAX_NAMESPACES,
0045 "description": (
0046 "Map keyed by testbed namespace; each value contains the "
0047 "current run number, phase, state, optional substate, and last "
0048 "transition time for that independently operated testbed."
0049 ),
0050 },
0051 },
0052 "event_sources": [
0053 {
0054 "name": "datataking-state-transitions",
0055 "resolver": EVENT_RESOLVER,
0056 "owner": "SWF testbed datataking state",
0057 "event_kind": "datataking-state-transition",
0058 "event_time_field": "timestamp",
0059 "visibility": "public",
0060 }
0061 ],
0062 }
0063
0064
0065 @dataclass(frozen=True)
0066 class DatatakingPublication:
0067 registration_update: ComponentUpdate
0068 update: ComponentUpdate
0069 projection: dict
0070 run_numbers: dict[str, int]
0071
0072
0073 def _timestamp(value: datetime) -> str:
0074 return (
0075 value.astimezone(datetime_timezone.utc)
0076 .isoformat()
0077 .replace("+00:00", "Z")
0078 )
0079
0080
0081 def _latest_run_states_by_namespace() -> list[RunState]:
0082 execution_key = KeyTextTransform("execution_id", "metadata")
0083 namespace = Subquery(
0084 WorkflowExecution.objects.filter(
0085 execution_id=OuterRef("execution_key")
0086 ).values("namespace")[:1]
0087 )
0088 rows = list(
0089 RunState.objects.annotate(execution_key=execution_key)
0090 .annotate(testbed_namespace=namespace)
0091 .exclude(testbed_namespace__isnull=True)
0092 .exclude(testbed_namespace="")
0093 .order_by("testbed_namespace", "-run_number")
0094 .distinct("testbed_namespace")[: MAX_NAMESPACES + 1]
0095 )
0096 if len(rows) > MAX_NAMESPACES:
0097 raise ValueError(
0098 f"datataking projection exceeds {MAX_NAMESPACES} namespaces"
0099 )
0100 return rows
0101
0102
0103 def datataking_projection() -> tuple[dict, dict[str, int]]:
0104 """Return the latest bounded run state for each testbed namespace."""
0105 run_states = _latest_run_states_by_namespace()
0106 if not run_states:
0107 raise ValueError(
0108 "cannot publish datataking state without a namespaced RunState row"
0109 )
0110
0111 namespaces = {}
0112 run_numbers = {}
0113 for run_state in run_states:
0114 namespace = run_state.testbed_namespace
0115 state = {
0116 "run_number": run_state.run_number,
0117 "phase": run_state.phase,
0118 "state": run_state.state,
0119 "last_transition_at": _timestamp(run_state.state_changed_at),
0120 }
0121 if run_state.substate is not None:
0122 state["substate"] = run_state.substate
0123 namespaces[namespace] = state
0124 run_numbers[namespace] = run_state.run_number
0125 return {"namespaces": namespaces}, run_numbers
0126
0127
0128 @transaction.atomic
0129 def publish_datataking_state(
0130 assessed_at: Optional[datetime] = None,
0131 ) -> DatatakingPublication:
0132 """Register and publish current datataking state across namespaces."""
0133 projection, run_numbers = datataking_projection()
0134 assessed_at = assessed_at or timezone.now()
0135 registration_update = register_component(
0136 scope="testbed",
0137 name="datataking",
0138 publisher_identity=PUBLISHER_IDENTITY,
0139 registration=DATATAKING_REGISTRATION,
0140 component_schema_version=2,
0141 )
0142 update = publish_component(
0143 scope="testbed",
0144 name="datataking",
0145 publisher_identity=PUBLISHER_IDENTITY,
0146 data=projection,
0147 assessed_at=assessed_at,
0148 assessment_policy_version=ASSESSMENT_POLICY_VERSION,
0149 )
0150 return DatatakingPublication(
0151 registration_update=registration_update,
0152 update=update,
0153 projection=projection,
0154 run_numbers=run_numbers,
0155 )
