From 70e678673a8e946465af3705bf2ad95e97d6b84f Mon Sep 17 00:00:00 2001 From: Yue Chao Qin Date: Mon, 26 Oct 2026 12:13:25 -0600 Subject: [PATCH] refactor(observability): one home for the task annotations telemetry reads The cloud_provider key was declared twice in instrumentation/, and the three resource keys were reached by importing a launcher purely to borrow its constants. Collects every read into instrumentation/execution_annotations.py: the resource keys are aliased from kubernetes_launchers rather than respelled, and cloud_provider -- which no launcher owns, since routing sets it -- is declared once. Also adds Target and register_targets(), the lookup that turns whatever a task's provider annotation says into a provider and a cluster. A provider that declares no clusters of its own registers the annotation key its tasks use, so no deployment-specific key appears in this repository. Pure readers: task_spec is a loaded column, never a relationship, so these are safe inside before_commit where a lazy load would put a SELECT in the commit path. No behaviour change. --- .../instrumentation/contextual_logging.py | 7 +- .../instrumentation/execution_annotations.py | 141 ++++++++++++++++++ .../instrumentation/execution_tracing.py | 15 +- .../test_execution_annotations.py | 105 +++++++++++++ 4 files changed, 254 insertions(+), 14 deletions(-) create mode 100644 cloud_pipelines_backend/instrumentation/execution_annotations.py create mode 100644 tests/instrumentation/test_execution_annotations.py diff --git a/cloud_pipelines_backend/instrumentation/contextual_logging.py b/cloud_pipelines_backend/instrumentation/contextual_logging.py index f95f8087..44a925d9 100644 --- a/cloud_pipelines_backend/instrumentation/contextual_logging.py +++ b/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) diff --git a/cloud_pipelines_backend/instrumentation/execution_annotations.py b/cloud_pipelines_backend/instrumentation/execution_annotations.py new file mode 100644 index 00000000..510f4c4b --- /dev/null +++ b/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=) 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=)} + + 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) diff --git a/cloud_pipelines_backend/instrumentation/execution_tracing.py b/cloud_pipelines_backend/instrumentation/execution_tracing.py index 095a1a8c..63f6f59d 100644 --- a/cloud_pipelines_backend/instrumentation/execution_tracing.py +++ b/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 diff --git a/tests/instrumentation/test_execution_annotations.py b/tests/instrumentation/test_execution_annotations.py new file mode 100644 index 00000000..0378a3d1 --- /dev/null +++ b/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)