refactor(observability): one home for the task annotations telemetry reads

cloud_pipelines_backend/instrumentation/contextual_logging.py

@@ -26,10 +26,11 @@
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, Optional
 
+from . import execution_annotations
+
 if TYPE_CHECKING:
     from .. import backend_types_sql as bts
 
-_CLOUD_PROVIDER_ANNOTATION_KEY = "cloud-pipelines.net/orchestration/cloud_provider"
 
 # Single context variable to store all metadata as a dictionary
 _context_metadata: contextvars.ContextVar[dict[str, Any]] = contextvars.ContextVar(
@@ -141,9 +142,7 @@ def execution_logging_context(execution: "bts.ExecutionNode"):
     on the task spec.
     """
     ctx: dict[str, str] = {"execution_id": execution.id}
-    cloud_provider = ((execution.task_spec or {}).get("annotations") or {}).get(
-        _CLOUD_PROVIDER_ANNOTATION_KEY
-    )
+    cloud_provider = execution_annotations.cloud_provider(execution=execution)
     if cloud_provider is not None:
         ctx["cloud_provider"] = cloud_provider
     return logging_context(**ctx)

cloud_pipelines_backend/instrumentation/execution_annotations.py

@@ -0,0 +1,141 @@
+"""Routing facts read off an ExecutionNode's annotations, without touching the database.
+
+One home for each annotation key, so the span emitter, the log enricher and the metric
+emitter cannot drift apart. Every reader is pure -- ``task_spec`` is a loaded column,
+never a relationship -- so these are safe to call inside ``before_commit``, where a lazy
+load would put a SELECT in the commit path.
+"""
+
+import json
+from typing import NamedTuple
+
+from .. import backend_types_sql as bts
+from ..launchers import kubernetes_launchers
+
+# Aliased, not respelled: a launcher renaming a key cannot silently empty these readers.
+RESOURCES_CPU_KEY = kubernetes_launchers.RESOURCES_CPU_ANNOTATION_KEY
+RESOURCES_MEMORY_KEY = kubernetes_launchers.RESOURCES_MEMORY_ANNOTATION_KEY
+RESOURCES_ACCELERATORS_KEY = kubernetes_launchers.RESOURCES_ACCELERATORS_ANNOTATION_KEY
+
+# Set at routing time rather than by a launcher, so no launcher module owns it.
+CLOUD_PROVIDER_KEY = "cloud-pipelines.net/orchestration/cloud_provider"
+
+
+class Target(NamedTuple):
+    """Where a task runs, once its annotation has been resolved.
+
+        Target("google", "lt3")                   a cluster the config names
+        Target("nebius", "h200")                  ditto
+        Target("davies", "davies",                a provider with no clusters of
+               cluster_from_annotation=<key>)     its own; each task names one
+
+    ``cluster_from_annotation`` lets the deployment register which annotation key its
+    tasks use for that, so a deployment-specific key never appears in this repository.
+    """
+
+    provider: str
+    cluster: str
+    cluster_from_annotation: str | None = None
+
+
+# Filled once at startup from the launcher config, which names either spelling of every
+# target. Reading that config validates a file, so it must not happen per status change.
+_TARGETS: dict[str, Target] = {}
+
+
+def register_targets(*, targets: dict[str, Target]) -> None:
+    """Install the alias-and-key → Target map. Called once, at startup.
+
+        {"lt3":     Target("google", "lt3"),
+         "gke-lt3": Target("google", "lt3"),   ← an alias for the same cluster
+         "h200":    Target("nebius", "h200"),
+         "davies":  Target("davies", "davies", cluster_from_annotation=<key>)}
+
+    Replaces the map rather than adding to it, so a reload cannot leave behind a
+    target the config no longer names.
+    """
+    _TARGETS.clear()
+    _TARGETS.update(targets)
+
+
+def _annotations(*, execution: "bts.ExecutionNode") -> dict[str, str]:
+    """Every annotation on the task, or an empty dict.
+
+        reads: task_spec["annotations"]
+
+    A loaded column, never a relationship, which is what makes every reader here
+    safe inside ``before_commit``.
+    """
+    return ((execution.task_spec or {}).get("annotations")) or {}
+
+
+def cloud_provider(*, execution: "bts.ExecutionNode") -> str | None:
+    """The provider annotation exactly as written -- "gke-lt3", never "google".
+
+        reads: cloud-pipelines.net/orchestration/cloud_provider
+
+        cloud_provider="lt3"       → "lt3"
+        cloud_provider="gke-lt3"   → "gke-lt3"    an alias is NOT resolved here
+        no annotation              → None
+
+    Kept for readers that want what the task asked for rather than where it landed.
+    Use ``target`` to resolve an alias to a provider and a cluster.
+    """
+    return _annotations(execution=execution).get(CLOUD_PROVIDER_KEY)
+
+
+def target(*, execution: "bts.ExecutionNode") -> Target | None:
+    """The provider and cluster a task selected, however its annotation spelled it.
+
+        reads: cloud-pipelines.net/orchestration/cloud_provider
+               + the cluster key registered by a provider that declares none (a)
+
+        cloud_provider="lt3"                      → Target("google", "lt3")
+        cloud_provider="gke-lt3"                  → Target("google", "lt3")
+        cloud_provider="h200"                     → Target("nebius", "h200")
+        cloud_provider="davies" + cluster="b300"  → Target("davies", "b300")
+        cloud_provider="davies", no cluster named → Target("davies", "none")
+        cloud_provider="retired-tpu"              → None    not in the map
+        no annotation at all                      → None
+
+        (a) that key is deployment-specific, so it is registered rather than
+            spelled here; davies is the only provider using one today
+
+    None means "this deployment does not route that", not "unlabelled": callers turn
+    it into "none" so a filter cannot drop the executions that lack it.
+    """
+    selected = cloud_provider(execution=execution)
+    if selected is None:
+        return None
+    resolved = _TARGETS.get(selected)
+    if resolved is None or resolved.cluster_from_annotation is None:
+        return resolved
+    named = _annotations(execution=execution).get(resolved.cluster_from_annotation)
+    return resolved._replace(cluster=named or "none")
+
+
+def accelerators(*, execution: "bts.ExecutionNode") -> tuple[str, int] | None:
+    """The accelerator model and count a task asked for.
+
+        reads: cloud-pipelines.net/launchers/generic/resources.accelerators
+
+        '{"NVIDIA-B300-SXM6-PC": "8"}'  → ("NVIDIA-B300-SXM6-PC", 8)
+        '{"NVIDIA-B300-SXM6-PC":"8"}'   → ("NVIDIA-B300-SXM6-PC", 8)   (a)
+        absent  or  '{}'                → None
+        '{"A": "1", "B": "1"}'          → ValueError, two models
+        'not json'                      → ValueError
+
+        (a) one space from the row above -- one request, two spellings
+
+    Parsing is what makes the two spacings one series instead of two.
+    """
+    raw = _annotations(execution=execution).get(RESOURCES_ACCELERATORS_KEY)
+    if not raw:
+        return None
+    requested = json.loads(raw)  # ValueError on malformed input
+    if not isinstance(requested, dict) or not requested:
+        return None
+    if len(requested) > 1:
+        raise ValueError(f"Multiple accelerator types were specified: {requested=}")
+    model, quantity = next(iter(requested.items()))
+    return model, int(quantity)

cloud_pipelines_backend/instrumentation/execution_tracing.py

@@ -14,12 +14,11 @@
 from opentelemetry.trace import StatusCode
 
 from .. import backend_types_sql as bts
-from ..launchers import kubernetes_launchers
+from . import execution_annotations
 
 _logger = logging.getLogger(__name__)
 _tracer = trace.get_tracer("tangle.orchestrator")
 
-_CLOUD_PROVIDER_ANNOTATION_KEY = "cloud-pipelines.net/orchestration/cloud_provider"
 
 _HISTORY_KEY = bts.EXECUTION_NODE_EXTRA_DATA_STATUS_HISTORY_KEY
 _TERMINAL_STATUSES = frozenset(s.value for s in bts.CONTAINER_STATUSES_ENDED)
@@ -116,11 +115,7 @@ def _cloud_provider_attrs(*, execution: bts.ExecutionNode) -> dict[str, object]:
     set this annotation at routing time; tangle launchers with a fixed cloud affinity
     can set it too.
     """
-    provider = (
-        (execution.task_spec or {})
-        .get("annotations", {})
-        .get(_CLOUD_PROVIDER_ANNOTATION_KEY)
-    )
+    provider = execution_annotations.cloud_provider(execution=execution)
     if provider is None:
         return {}
     return {"execution.cloud_provider": provider}
@@ -151,12 +146,12 @@ def _resource_attrs(*, execution: bts.ExecutionNode, status: str) -> dict[str, o
         return {}
     annotations: dict = (execution.task_spec or {}).get("annotations", {})
     attrs: dict[str, object] = {}
-    if cpu := annotations.get(kubernetes_launchers.RESOURCES_CPU_ANNOTATION_KEY):
+    if cpu := annotations.get(execution_annotations.RESOURCES_CPU_KEY):
         attrs["execution.resources.cpu"] = cpu
-    if memory := annotations.get(kubernetes_launchers.RESOURCES_MEMORY_ANNOTATION_KEY):
+    if memory := annotations.get(execution_annotations.RESOURCES_MEMORY_KEY):
         attrs["execution.resources.memory"] = memory
     if accelerators := annotations.get(
-        kubernetes_launchers.RESOURCES_ACCELERATORS_ANNOTATION_KEY
+        execution_annotations.RESOURCES_ACCELERATORS_KEY
     ):
         attrs["execution.resources.accelerators"] = accelerators
     return attrs

tests/instrumentation/test_execution_annotations.py

@@ -0,0 +1,105 @@
+"""Readers for the annotations that say where a task runs."""
+
+import json
+
+import pytest
+
+from cloud_pipelines_backend import backend_types_sql as bts
+from cloud_pipelines_backend.instrumentation import execution_annotations as ea
+
+
+def _execution(**annotations: str) -> bts.ExecutionNode:
+    return bts.ExecutionNode(task_spec={"annotations": annotations})
+
+
+class TestCloudProvider:
+    def test_absent_annotation_reads_none(self) -> None:
+        assert ea.cloud_provider(execution=_execution()) is None
+
+    def test_value_is_returned_as_written(self) -> None:
+        execution = _execution(**{ea.CLOUD_PROVIDER_KEY: "gke-lt3"})
+        assert ea.cloud_provider(execution=execution) == "gke-lt3"
+
+    def test_absent_task_spec_is_not_an_error(self) -> None:
+        assert ea.cloud_provider(execution=bts.ExecutionNode(task_spec=None)) is None
+
+
+class TestTarget:
+    @pytest.fixture(autouse=True)
+    def _targets(self) -> None:
+        ea.register_targets(
+            targets={
+                "gke-lt3": ea.Target(provider="google", cluster="gke-lt3"),
+                "lt3": ea.Target(provider="google", cluster="gke-lt3"),
+                "byo": ea.Target(
+                    provider="byo",
+                    cluster="",
+                    cluster_from_annotation="example.test/byo/cluster",
+                ),
+            }
+        )
+
+    def test_no_annotation_resolves_to_none(self) -> None:
+        assert ea.target(execution=_execution()) is None
+
+    def test_unknown_value_resolves_to_none(self) -> None:
+        execution = _execution(**{ea.CLOUD_PROVIDER_KEY: "not-in-the-config"})
+        assert ea.target(execution=execution) is None
+
+    def test_an_alias_resolves_to_the_same_target_as_its_key(self) -> None:
+        by_key = _execution(**{ea.CLOUD_PROVIDER_KEY: "gke-lt3"})
+        by_alias = _execution(**{ea.CLOUD_PROVIDER_KEY: "lt3"})
+        assert ea.target(execution=by_key) == ea.target(execution=by_alias)
+
+    def test_provider_naming_its_own_cluster_reads_the_registered_key(self) -> None:
+        execution = _execution(
+            **{
+                ea.CLOUD_PROVIDER_KEY: "byo",
+                "example.test/byo/cluster": "rack-7",
+            }
+        )
+        assert ea.target(execution=execution) == ea.Target(
+            provider="byo",
+            cluster="rack-7",
+            cluster_from_annotation="example.test/byo/cluster",
+        )
+
+    def test_provider_naming_no_cluster_reads_none_not_empty(self) -> None:
+        execution = _execution(**{ea.CLOUD_PROVIDER_KEY: "byo"})
+        assert ea.target(execution=execution).cluster == "none"
+
+
+class TestAccelerators:
+    def test_absent_annotation_reads_none(self) -> None:
+        assert ea.accelerators(execution=_execution()) is None
+
+    def test_empty_string_reads_none(self) -> None:
+        execution = _execution(**{ea.RESOURCES_ACCELERATORS_KEY: ""})
+        assert ea.accelerators(execution=execution) is None
+
+    def test_empty_object_reads_none(self) -> None:
+        execution = _execution(**{ea.RESOURCES_ACCELERATORS_KEY: "{}"})
+        assert ea.accelerators(execution=execution) is None
+
+    @pytest.mark.parametrize("raw", ['{"nvidia-h100":"8"}', '{"nvidia-h100": "8"}'])
+    def test_whitespace_does_not_make_two_requests(self, raw: str) -> None:
+        execution = _execution(**{ea.RESOURCES_ACCELERATORS_KEY: raw})
+        assert ea.accelerators(execution=execution) == ("nvidia-h100", 8)
+
+    def test_integer_quantity_is_accepted(self) -> None:
+        execution = _execution(
+            **{ea.RESOURCES_ACCELERATORS_KEY: json.dumps({"nvidia-h100": 4})}
+        )
+        assert ea.accelerators(execution=execution) == ("nvidia-h100", 4)
+
+    def test_two_models_is_an_error(self) -> None:
+        execution = _execution(
+            **{ea.RESOURCES_ACCELERATORS_KEY: '{"nvidia-h100":"8","tpu":"1"}'}
+        )
+        with pytest.raises(ValueError, match="Multiple accelerator types"):
+            ea.accelerators(execution=execution)
+
+    def test_malformed_json_is_an_error(self) -> None:
+        execution = _execution(**{ea.RESOURCES_ACCELERATORS_KEY: "{not json"})
+        with pytest.raises(ValueError):
+            ea.accelerators(execution=execution)