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."""