From 7291f07a3b2ceaeb6424d2daa48abb1b19e95bbf Mon Sep 17 00:00:00 2001 From: Yue Chao Qin Date: Mon, 26 Oct 2026 14:34:54 -0600 Subject: [PATCH] feat(observability): label the platform, accelerator and cluster a task ran on The status-transition histogram carries only the two statuses, so a queue that is slow because one GPU family is scarce is indistinguishable from the orchestrator being slow for everyone. Adds the labels that separate them. Each label is resolved rather than copied: the provider annotation holds whatever the client wrote, so two spellings of one cluster are collapsed through the registered target map instead of forming two series. Nothing is ever omitted -- an absent or unroutable value is labelled "none", because an attribute that disappears forms a second series that no filter on the first one covers. --- .../instrumentation/execution_annotations.py | 49 ++- .../instrumentation/metrics.py | 124 +++++++ tests/test_sql_event_listeners.py | 324 ++++++++++++++++++ 3 files changed, 494 insertions(+), 3 deletions(-) diff --git a/cloud_pipelines_backend/instrumentation/execution_annotations.py b/cloud_pipelines_backend/instrumentation/execution_annotations.py index 510f4c4..7e9cc73 100644 --- a/cloud_pipelines_backend/instrumentation/execution_annotations.py +++ b/cloud_pipelines_backend/instrumentation/execution_annotations.py @@ -7,9 +7,11 @@ """ import json +import re from typing import NamedTuple from .. import backend_types_sql as bts +from ..launchers import common_annotations from ..launchers import kubernetes_launchers # Aliased, not respelled: a launcher renaming a key cannot silently empty these readers. @@ -20,21 +22,29 @@ # Set at routing time rather than by a launcher, so no launcher module owns it. CLOUD_PROVIDER_KEY = "cloud-pipelines.net/orchestration/cloud_provider" +# Read by the deployment's pod builder, which puts such a task in a namespace of its +# own rather than the cluster's. +GOOGLE_SERVICE_ACCOUNT_KEY = "cloud-pipelines.net/launchers/google/service_account" + 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 + Target("google", "us-ce1-lt3", a cluster the config names, + "kueue-jobs-staging") with the namespace it runs in + Target("nebius", "b300-cluster", + "tangle-jobs-production") + Target("davies", "davies", "none", 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. + ``namespace`` is the cluster's own; a task carrying a service account overrides it. """ provider: str cluster: str + namespace: str = "none" cluster_from_annotation: str | None = None @@ -114,6 +124,39 @@ def target(*, execution: "bts.ExecutionNode") -> Target | None: return resolved._replace(cluster=named or "none") +def namespace(*, execution: "bts.ExecutionNode") -> str | None: + """The namespace a task's own service account puts it in, or None for its cluster's. + + reads: cloud-pipelines.net/launchers/google/service_account + + cloud-pipelines.net/pipeline_runs/created_by + + service_account="sidekick-oasis-runner@s5y-sidekick-52ee + .iam.gserviceaccount.com", created_by set + → "sidekick-oasis-runner-at-s5y-sidekick-52ee" + service_account set, created_by absent → None (a) + neither → None (a) + + (a) the pod lands in the cluster's namespace instead, which this module + cannot see; the caller falls back to Target.namespace + + Both annotations are required because the deployment's pod builder only creates a + per-account namespace when it has both. Deriving it here duplicates that rule + rather than querying Kubernetes, which no reader in this module may do. + """ + annotations = _annotations(execution=execution) + account = annotations.get(GOOGLE_SERVICE_ACCOUNT_KEY) + created_by = annotations.get( + common_annotations.PIPELINE_RUN_CREATED_BY_ANNOTATION_KEY + ) + if not account or not created_by: + return None + return re.sub( + r"[^-a-z0-9]", + "-", + account.lower().removesuffix(".iam.gserviceaccount.com").replace("@", "-at-"), + ).strip("-") + + def accelerators(*, execution: "bts.ExecutionNode") -> tuple[str, int] | None: """The accelerator model and count a task asked for. diff --git a/cloud_pipelines_backend/instrumentation/metrics.py b/cloud_pipelines_backend/instrumentation/metrics.py index 34de2ab..ee187b9 100644 --- a/cloud_pipelines_backend/instrumentation/metrics.py +++ b/cloud_pipelines_backend/instrumentation/metrics.py @@ -31,6 +31,7 @@ from sqlalchemy import orm from .. import backend_types_sql +from . import execution_annotations from . import execution_tracing _logger = logging.getLogger(__name__) @@ -86,6 +87,10 @@ class MetricUnit(str, enum.Enum): _RUNNING = backend_types_sql.ContainerExecutionStatus.RUNNING.value +# Every label an execution contributes has one of these when the task said nothing, so a +# query that filters on a label never silently drops the executions that lack it. +_UNSET = "none" + def _run_attempt(*, history: list[dict[str, str]]) -> str: """Whether the execution is waiting to START or waiting to RESUME. @@ -112,6 +117,117 @@ def _run_attempt(*, history: list[dict[str, str]]) -> str: return "repeat" if any(e["status"] == _RUNNING for e in history[:-1]) else "first" +def _cloud_provider(*, execution: backend_types_sql.ExecutionNode) -> str: + """Whose machines ran the work. + + reads: cloud-pipelines.net/orchestration/cloud_provider + + annotation target map label + "lt3" → google / lt3 → "google" + "gke-lt3" → google / lt3 → "google" (a) + "h200" → nebius / h200 → "nebius" + "davies" → davies / b300 → "davies" (b) + "retired-tpu" → not in the map → "none" (c) + no annotation → -- → "none" (c) + + (a) an alias for the row above -- one series, not two + (b) davies declares no clusters, so the task names one in a second + annotation the deployment registered + (c) labelled rather than omitted: a vanishing attribute forms a second + series that no filter on the first one covers + + The map decides, not the annotation. The config section is `gke` and the label + is `google`, because it answers whose machines, not which config stanza. + """ + resolved = execution_annotations.target(execution=execution) + return _UNSET if resolved is None else resolved.provider + + +def _namespace(*, execution: backend_types_sql.ExecutionNode) -> str: + """Which Kubernetes namespace the work ran in. + + reads: cloud-pipelines.net/launchers/google/service_account + + cloud-pipelines.net/pipeline_runs/created_by + falling back to the namespace registered with the target + + a task with its own service account → "sidekick-oasis-runner-at-s5y-…" (a) + no service account, target "lt3" → "kueue-jobs-staging" + no service account, target "eo9" → "tangle-jobs-production" + an unrouted or missing target → "none" + + (a) ~39 such namespaces exist and the set grows with onboarding: each is a + function of a submitter's service account, so none is declared anywhere + in this repository or the deployment's config + + Two origins, and only the first is visible in the annotations. A task carrying a + service account gets a namespace of its own; everything else lands in its + cluster's, which the deployment registers with the target. + """ + own = execution_annotations.namespace(execution=execution) + if own is not None: + return own + resolved = execution_annotations.target(execution=execution) + return _UNSET if resolved is None else resolved.namespace + + +def _cluster(*, execution: backend_types_sql.ExecutionNode) -> str: + """Which cluster inside that platform. + + reads: cloud-pipelines.net/orchestration/cloud_provider + + annotation target map label + "lt3" → google / us-ce1-lt3 → "us-ce1-lt3" + "eo9" → google / us-ce1-eo9 → "us-ce1-eo9" + "h200" → nebius / b300-cluster→ "b300-cluster" + "davies" → davies / named by → whatever the task named (a) + the task + "retired-tpu" → not in the map → "none" + no annotation → -- → "none" + + (a) davies declares no clusters of its own, so the value is whatever the + submitter typed, and "none" when the task named nothing + + The deployment registers these already slugified: config labels are display + strings like "B300 cluster", which is legal in a label value and unpleasant in + a query. Resolved through the map, never read off the annotation -- the + annotation names a target, which is not the same thing as a cluster. + """ + resolved = execution_annotations.target(execution=execution) + return _UNSET if resolved is None else resolved.cluster + + +def _accelerator(*, execution: backend_types_sql.ExecutionNode) -> tuple[str, str]: + """Which accelerator the task asked for, and how many. + + reads: cloud-pipelines.net/launchers/generic/resources.accelerators + + annotation type count + '{"NVIDIA-B300-SXM6-PC":"8"}' → "NVIDIA-B300-SXM6-PC" "8" + '{"NVIDIA-B300-SXM6-PC": "8"}' → "NVIDIA-B300-SXM6-PC" "8" (a) + '{"B300-SXM6-PC":"1"}' → "B300-SXM6-PC" "1" (b) + '{"NVIDIA-B300-SXM6-PC":"0"}' → "NVIDIA-B300-SXM6-PC" "0" + absent or '{}' → "none" "0" + '{"A":"1","B":"1"}' → "unparseable" "0" (c) + 'not json' → "unparseable" "0" (c) + + (a) one space from the row above -- one series, not two + (b) the capacity API's spelling of the same GPU; both kept as written + (c) labelled, not dropped, so a bad annotation shows as a spike rather + than as missing traffic + + Parsing is what collapses the spacings: labelling the raw string would split one + model across several series. + """ + try: + requested = execution_annotations.accelerators(execution=execution) + except ValueError: + return "unparseable", "0" + if requested is None: + return _UNSET, "0" + model, count = requested + return model, str(count) + + @sql_event.listens_for(orm.Session, "before_commit") def _handle_before_commit(session: orm.Session) -> None: for obj in list(session.new) + list(session.dirty): @@ -125,6 +241,9 @@ def _handle_before_commit(session: orm.Session) -> None: curr = history[-1] prev_time = datetime.datetime.fromisoformat(prev["first_observed_at"]) curr_time = datetime.datetime.fromisoformat(curr["first_observed_at"]) + accelerator_type, accelerator_count = _accelerator(execution=obj) + cluster = _cluster(execution=obj) + namespace = _namespace(execution=obj) try: execution_status_transition_duration.record( (curr_time - prev_time).total_seconds(), @@ -132,6 +251,11 @@ def _handle_before_commit(session: orm.Session) -> None: "execution.status.from": prev["status"], "execution.status.to": curr["status"], "execution.run_attempt": _run_attempt(history=history), + "execution.cloud_provider": _cloud_provider(execution=obj), + "execution.accelerator_type": accelerator_type, + "execution.accelerator_count": accelerator_count, + "execution.cluster": cluster, + "k8s.namespace.name": namespace, }, ) except Exception: diff --git a/tests/test_sql_event_listeners.py b/tests/test_sql_event_listeners.py index b984b63..67c3142 100644 --- a/tests/test_sql_event_listeners.py +++ b/tests/test_sql_event_listeners.py @@ -10,6 +10,8 @@ from cloud_pipelines_backend import ( orchestrator_sql, ) # noqa: F401 — registers set listener +from cloud_pipelines_backend.launchers import common_annotations +from cloud_pipelines_backend.instrumentation import execution_annotations from cloud_pipelines_backend.instrumentation import ( metrics, ) # noqa: F401 — registers before_commit listener @@ -67,9 +69,331 @@ def test_second_status_change_records_duration_metric( "execution.status.from": bts.ContainerExecutionStatus.QUEUED, "execution.status.to": bts.ContainerExecutionStatus.RUNNING, "execution.run_attempt": "first", + "execution.cloud_provider": "none", + "execution.accelerator_type": "none", + "execution.accelerator_count": "0", + "execution.cluster": "none", + "k8s.namespace.name": "none", } +class TestCloudProviderLabel: + """The platform label is resolved through the registered map, never labelled raw.""" + + @staticmethod + def _targets() -> dict[str, execution_annotations.Target]: + return { + "lt3": execution_annotations.Target(provider="google", cluster="lt3"), + "gke-lt3": execution_annotations.Target(provider="google", cluster="lt3"), + "h200": execution_annotations.Target(provider="nebius", cluster="h200"), + } + + @pytest.fixture(autouse=True) + def _registered(self): + execution_annotations.register_targets(targets=self._targets()) + yield + execution_annotations.register_targets(targets={}) + + @staticmethod + def _record(session: orm.Session, **annotations: str) -> dict[str, str]: + node = bts.ExecutionNode(task_spec={"annotations": annotations}) + session.add(node) + node.container_execution_status = bts.ContainerExecutionStatus.QUEUED + session.commit() + + node.container_execution_status = bts.ContainerExecutionStatus.RUNNING + with unittest.mock.patch.object( + metrics.execution_status_transition_duration, "record" + ) as mock_record: + session.commit() + return mock_record.call_args.kwargs["attributes"] + + @pytest.mark.parametrize( + "annotation_value,expected", + [ + ("lt3", "google"), + ("gke-lt3", "google"), + ("h200", "nebius"), + ("retired-cluster", "none"), + ], + ids=["canonical", "alias", "other-provider", "unrouted"], + ) + def test_the_map_decides_the_provider( + self, session: orm.Session, annotation_value: str, expected: str + ) -> None: + attributes = self._record( + session, + **{execution_annotations.CLOUD_PROVIDER_KEY: annotation_value}, + ) + assert attributes["execution.cloud_provider"] == expected + + def test_an_alias_is_not_counted_as_its_own_platform( + self, session: orm.Session + ) -> None: + """Two spellings of one cluster must not split the series.""" + by_key = self._record( + session, **{execution_annotations.CLOUD_PROVIDER_KEY: "lt3"} + ) + by_alias = self._record( + session, **{execution_annotations.CLOUD_PROVIDER_KEY: "gke-lt3"} + ) + assert ( + by_key["execution.cloud_provider"] == by_alias["execution.cloud_provider"] + ) + + def test_a_task_naming_no_platform_is_labelled_not_omitted( + self, session: orm.Session + ) -> None: + """An absent attribute forms a second series no filter on the first one covers.""" + attributes = self._record(session) + assert attributes["execution.cloud_provider"] == "none" + + +class TestNamespaceLabel: + """Two origins: a task's own service account, else the cluster's namespace.""" + + PROVIDER = execution_annotations.CLOUD_PROVIDER_KEY + GSA = execution_annotations.GOOGLE_SERVICE_ACCOUNT_KEY + CREATED_BY = common_annotations.PIPELINE_RUN_CREATED_BY_ANNOTATION_KEY + SIDEKICK = "sidekick-oasis-runner@s5y-sidekick-52ee.iam.gserviceaccount.com" + + @staticmethod + def _record(session: orm.Session, **annotations: str) -> dict[str, str]: + node = bts.ExecutionNode(task_spec={"annotations": annotations}) + session.add(node) + node.container_execution_status = bts.ContainerExecutionStatus.QUEUED + session.commit() + + node.container_execution_status = bts.ContainerExecutionStatus.RUNNING + with unittest.mock.patch.object( + metrics.execution_status_transition_duration, "record" + ) as mock_record: + session.commit() + return mock_record.call_args.kwargs["attributes"] + + @pytest.fixture(autouse=True) + def _registered(self): + execution_annotations.register_targets( + targets={ + "lt3": execution_annotations.Target( + "google", "us-ce1-lt3", "kueue-jobs-staging" + ), + "eo9": execution_annotations.Target( + "google", "us-ce1-eo9", "tangle-jobs-production" + ), + } + ) + yield + execution_annotations.register_targets(targets={}) + + def test_a_service_account_gets_a_namespace_of_its_own( + self, session: orm.Session + ) -> None: + attributes = self._record( + session, + **{ + self.PROVIDER: "lt3", + self.GSA: self.SIDEKICK, + self.CREATED_BY: "someone@shopify.com", + }, + ) + assert attributes["k8s.namespace.name"] == ( + "sidekick-oasis-runner-at-s5y-sidekick-52ee" + ) + + @pytest.mark.parametrize( + "target,expected", + [("lt3", "kueue-jobs-staging"), ("eo9", "tangle-jobs-production")], + ids=["default-cluster", "second-cluster"], + ) + def test_without_one_the_cluster_namespace_is_used( + self, session: orm.Session, target: str, expected: str + ) -> None: + attributes = self._record(session, **{self.PROVIDER: target}) + assert attributes["k8s.namespace.name"] == expected + + def test_a_service_account_without_a_creator_falls_back( + self, session: orm.Session + ) -> None: + """The pod builder only makes a per-account namespace when it has both.""" + attributes = self._record( + session, **{self.PROVIDER: "lt3", self.GSA: self.SIDEKICK} + ) + assert attributes["k8s.namespace.name"] == "kueue-jobs-staging" + + def test_an_unrouted_target_is_labelled_not_omitted( + self, session: orm.Session + ) -> None: + attributes = self._record(session, **{self.PROVIDER: "retired-tpu"}) + assert attributes["k8s.namespace.name"] == "none" + + def test_no_annotation_at_all_is_labelled_not_omitted( + self, session: orm.Session + ) -> None: + assert self._record(session)["k8s.namespace.name"] == "none" + + +class TestClusterLabel: + """The cluster comes from the registered map, never off the annotation.""" + + PROVIDER = execution_annotations.CLOUD_PROVIDER_KEY + DAVIES_CLUSTER = "example.test/byo/cluster" + + @staticmethod + def _record(session: orm.Session, **annotations: str) -> dict[str, str]: + node = bts.ExecutionNode(task_spec={"annotations": annotations}) + session.add(node) + node.container_execution_status = bts.ContainerExecutionStatus.QUEUED + session.commit() + + node.container_execution_status = bts.ContainerExecutionStatus.RUNNING + with unittest.mock.patch.object( + metrics.execution_status_transition_duration, "record" + ) as mock_record: + session.commit() + return mock_record.call_args.kwargs["attributes"] + + @pytest.fixture(autouse=True) + def _registered(self): + execution_annotations.register_targets( + targets={ + "lt3": execution_annotations.Target("google", "us-ce1-lt3"), + "gke-lt3": execution_annotations.Target("google", "us-ce1-lt3"), + "eo9": execution_annotations.Target("google", "us-ce1-eo9"), + "h200": execution_annotations.Target("nebius", "b300-cluster"), + "byo": execution_annotations.Target( + "byo", "byo", cluster_from_annotation=self.DAVIES_CLUSTER + ), + } + ) + yield + execution_annotations.register_targets(targets={}) + + @pytest.mark.parametrize( + "annotation_value,expected", + [ + ("lt3", "us-ce1-lt3"), + ("gke-lt3", "us-ce1-lt3"), + ("eo9", "us-ce1-eo9"), + ("h200", "b300-cluster"), + ("retired-tpu", "none"), + ], + ids=["default", "alias", "second-gke", "nebius", "unrouted"], + ) + def test_the_map_decides_the_cluster( + self, session: orm.Session, annotation_value: str, expected: str + ) -> None: + attributes = self._record(session, **{self.PROVIDER: annotation_value}) + assert attributes["execution.cluster"] == expected + + def test_the_annotation_is_a_target_not_a_cluster( + self, session: orm.Session + ) -> None: + """ "lt3" names a target; the cluster it resolves to is spelled differently.""" + attributes = self._record(session, **{self.PROVIDER: "lt3"}) + assert attributes["execution.cluster"] != "lt3" + assert attributes["execution.cluster"] == "us-ce1-lt3" + + def test_a_provider_with_no_clusters_takes_the_one_the_task_named( + self, session: orm.Session + ) -> None: + attributes = self._record( + session, **{self.PROVIDER: "byo", self.DAVIES_CLUSTER: "ml-control-plane"} + ) + assert attributes["execution.cluster"] == "ml-control-plane" + + def test_that_provider_naming_no_cluster_is_labelled_not_omitted( + self, session: orm.Session + ) -> None: + attributes = self._record(session, **{self.PROVIDER: "byo"}) + assert attributes["execution.cluster"] == "none" + + def test_no_annotation_is_labelled_not_omitted(self, session: orm.Session) -> None: + assert self._record(session)["execution.cluster"] == "none" + + +class TestAcceleratorLabels: + """One GPU written several ways must be one series, not several.""" + + ACCEL = execution_annotations.RESOURCES_ACCELERATORS_KEY + + @staticmethod + def _record(session: orm.Session, **annotations: str) -> dict[str, str]: + node = bts.ExecutionNode(task_spec={"annotations": annotations}) + session.add(node) + node.container_execution_status = bts.ContainerExecutionStatus.QUEUED + session.commit() + + node.container_execution_status = bts.ContainerExecutionStatus.RUNNING + with unittest.mock.patch.object( + metrics.execution_status_transition_duration, "record" + ) as mock_record: + session.commit() + return mock_record.call_args.kwargs["attributes"] + + @pytest.mark.parametrize( + "raw,expected_type,expected_count", + [ + ('{"NVIDIA-B300-SXM6-PC":"8"}', "NVIDIA-B300-SXM6-PC", "8"), + ('{"NVIDIA-B300-SXM6-PC": "8"}', "NVIDIA-B300-SXM6-PC", "8"), + ('{"B300-SXM6-PC":"1"}', "B300-SXM6-PC", "1"), + ('{"NVIDIA-B300-SXM6-PC":"0"}', "NVIDIA-B300-SXM6-PC", "0"), + ("{}", "none", "0"), + ('{"A":"1","B":"1"}', "unparseable", "0"), + ("not json", "unparseable", "0"), + ], + ids=[ + "compact", + "spaced", + "second-spelling", + "explicit-zero", + "empty-object", + "two-models", + "malformed", + ], + ) + def test_the_annotation_is_parsed_not_copied( + self, + session: orm.Session, + raw: str, + expected_type: str, + expected_count: str, + ) -> None: + attributes = self._record(session, **{self.ACCEL: raw}) + assert attributes["execution.accelerator_type"] == expected_type + assert attributes["execution.accelerator_count"] == expected_count + + def test_one_space_does_not_make_a_second_series( + self, session: orm.Session + ) -> None: + """json.dumps emits the space and a hand-typed annotation does not.""" + compact = self._record(session, **{self.ACCEL: '{"NVIDIA-H200":"4"}'}) + spaced = self._record(session, **{self.ACCEL: '{"NVIDIA-H200": "4"}'}) + assert compact["execution.accelerator_type"] == ( + spaced["execution.accelerator_type"] + ) + assert compact["execution.accelerator_count"] == ( + spaced["execution.accelerator_count"] + ) + + def test_a_task_asking_for_no_gpu_is_labelled_not_omitted( + self, session: orm.Session + ) -> None: + attributes = self._record(session) + assert attributes["execution.accelerator_type"] == "none" + assert attributes["execution.accelerator_count"] == "0" + + def test_a_malformed_annotation_does_not_lose_the_measurement( + self, session: orm.Session + ) -> None: + """The duration still records; only the accelerator labels degrade.""" + attributes = self._record(session, **{self.ACCEL: "not json"}) + assert attributes["execution.status.to"] == ( + bts.ContainerExecutionStatus.RUNNING + ) + assert attributes["execution.accelerator_type"] == "unparseable" + + class TestRunAttempt: """A preempted execution reaches RUNNING more than once; only the first is a start."""