From b23383364178182a626ca7b1b860f74bd3b424f1 Mon Sep 17 00:00:00 2001 From: Morgan Wowk Date: Tue, 3 Nov 2026 11:02:00 -0700 Subject: [PATCH] feat(orchestrator): Refresh a container execution as soon as its pod or job changes The running queue visits container executions round-robin by last_processed_at, so a pod that finished waits for every other active container execution to be visited first. ContainerTaskLauncher gets pop_changed_container_execution_ids(), which returns the container executions whose containers changed since the last call; it returns nothing unless a launcher overrides it. The Kubernetes launchers now record the container execution id annotation on the pods they launch (and so on job pod templates), and their pod and job informers record each changed object's container execution, keeping at most 10,000 ids between calls. Before selecting a running container execution, the orchestrator moves the changed ones that are still PENDING or RUNNING to the front of the queue with one UPDATE. Pods and jobs launched before this change carry no id and keep the round-robin. --- .../launchers/interfaces.py | 4 + .../launchers/kubernetes_launchers.py | 67 +++++ cloud_pipelines_backend/orchestrator_sql.py | 25 ++ ..._launchers_changed_container_executions.py | 246 ++++++++++++++++++ ...chestrator_changed_container_executions.py | 147 +++++++++++ 5 files changed, 489 insertions(+) create mode 100644 tests/launchers/kubernetes/test_kubernetes_launchers_changed_container_executions.py create mode 100644 tests/test_orchestrator_changed_container_executions.py diff --git a/cloud_pipelines_backend/launchers/interfaces.py b/cloud_pipelines_backend/launchers/interfaces.py index 433e09dc..c90dc09b 100644 --- a/cloud_pipelines_backend/launchers/interfaces.py +++ b/cloud_pipelines_backend/launchers/interfaces.py @@ -80,6 +80,10 @@ def get_refreshed_launched_container_from_dict( ) -> _TLaunchedContainer: raise NotImplementedError() + def pop_changed_container_execution_ids(self) -> list[str]: + """Returns and forgets the ids of the container executions whose containers changed since the last call; empty when the launcher cannot tell.""" + return [] + class ContainerStatus(str, enum.Enum): PENDING = "PENDING" diff --git a/cloud_pipelines_backend/launchers/kubernetes_launchers.py b/cloud_pipelines_backend/launchers/kubernetes_launchers.py index 25db318e..72721567 100644 --- a/cloud_pipelines_backend/launchers/kubernetes_launchers.py +++ b/cloud_pipelines_backend/launchers/kubernetes_launchers.py @@ -79,6 +79,8 @@ # Multi-node constants _MULTI_NODE_MAX_NUMBER_OF_NODES = 16 +_MAX_CHANGED_CONTAINER_EXECUTION_IDS = 10_000 + # dynamicData argument keys _MULTI_NODE_NUMBER_OF_NODES_DYNAMIC_DATA_KEY = "system/multi_node/number_of_nodes" _MULTI_NODE_NODE_INDEX_DYNAMIC_DATA_KEY = "system/multi_node/node_index" @@ -203,6 +205,8 @@ def __init__( ) self._pod_postprocessor = pod_postprocessor self._create_volume_and_volume_mount = _create_volume_and_volume_mount + self._changed_container_execution_ids: set[str] = set() + self._changed_container_execution_ids_lock = threading.Lock() try: k8s_client_lib.VersionApi(self._api_client).get_code( @@ -231,6 +235,7 @@ def __init__( # But Tangle has only recently started labeling it's pods/jobs, so if we add filtering too soon, this will become a breaking change. # label_selector=f"{_TANGLE_KUBERNETES_LABEL_KEY}=true", label_selector=None, + on_change=self._mark_container_execution_changed, ) self._pod_informer.start() except Exception: @@ -488,6 +493,14 @@ def get_output_path(output_name: str) -> str: service_account_name=pod_service_account, ) + container_execution_id = (annotations or {}).get( + common_annotations.CONTAINER_EXECUTION_ID_ANNOTATION_KEY + ) + if container_execution_id: + pod_annotations = (pod_annotations or {}) | { + common_annotations.CONTAINER_EXECUTION_ID_ANNOTATION_KEY: container_execution_id + } + pod = k8s_client_lib.V1Pod( api_version="v1", kind="Pod", @@ -513,6 +526,27 @@ def _choose_namespace( del annotations return self._namespace + def pop_changed_container_execution_ids(self) -> list[str]: + """Returns and forgets the ids of the container executions whose pods or jobs changed since the last call. + + At most `_MAX_CHANGED_CONTAINER_EXECUTION_IDS` ids are kept between calls, so a launcher nobody asks stays bounded. + """ + with self._changed_container_execution_ids_lock: + container_execution_ids = list(self._changed_container_execution_ids) + self._changed_container_execution_ids = set() + return container_execution_ids + + def _mark_container_execution_changed(self, obj: Any) -> None: + """Records the container execution that launched `obj`, a pod or a job, as changed.""" + container_execution_id = _get_container_execution_id(obj) + if container_execution_id: + with self._changed_container_execution_ids_lock: + if ( + len(self._changed_container_execution_ids) + < _MAX_CHANGED_CONTAINER_EXECUTION_IDS + ): + self._changed_container_execution_ids.add(container_execution_id) + class _KubernetesPodLauncher( _KubernetesContainerLauncherBase, @@ -1106,6 +1140,7 @@ def __init__( # But Tangle has only recently started labeling it's pods/jobs, so if we add filtering too soon, this will become a breaking change. # label_selector=f"{_TANGLE_KUBERNETES_LABEL_KEY}=true", label_selector=None, + on_change=self._mark_container_execution_changed, ) self._job_informer.start() except Exception: @@ -2047,6 +2082,13 @@ def get_refreshed_launched_container_from_dict( ) return launched_container.get_refreshed() + def pop_changed_container_execution_ids(self) -> list[str]: + """Returns and forgets the ids of the container executions whose pods or jobs changed since the last call.""" + return ( + self._pod_launcher.pop_changed_container_execution_ids() + + self._job_launcher.pop_changed_container_execution_ids() + ) + def deserialize_launched_container_from_dict( self, launched_container_dict: dict ) -> "LaunchedKubernetesContainer | LaunchedKubernetesJob": @@ -2119,6 +2161,22 @@ def __init__( ) +def _get_container_execution_id(obj: Any) -> str | None: + """Reads the container execution id annotation off a pod, or off a job's pod template.""" + pod_template = getattr(getattr(obj, "spec", None), "template", None) + for metadata in ( + getattr(obj, "metadata", None), + getattr(pod_template, "metadata", None), + ): + annotations = getattr(metadata, "annotations", None) or {} + container_execution_id = annotations.get( + common_annotations.CONTAINER_EXECUTION_ID_ANNOTATION_KEY + ) + if container_execution_id: + return container_execution_id + return None + + def _get_job_pod_completion_index(pod: k8s_client_lib.V1Pod) -> str | None: if not (pod.metadata and pod.metadata.annotations): return None @@ -2219,12 +2277,14 @@ def __init__( label_selector=None, field_selector=None, key_func=None, + on_change: typing.Callable[[Any], None] | None = None, ): self._list_func = list_func self._resync_period = resync_period self._label_selector = label_selector self._field_selector = field_selector self._key_func = key_func + self._on_change = on_change self._informers: dict[str, k8s_informer_lib.SharedInformer] = {} self._lock = threading.Lock() self._is_active = False @@ -2250,6 +2310,13 @@ def _get_or_create_informer( ) # Prefill the cache informer._initial_list() + if self._on_change: + for event_type in ( + k8s_informer_lib.ADDED, + k8s_informer_lib.MODIFIED, + k8s_informer_lib.DELETED, + ): + informer.add_event_handler(event_type, self._on_change) self._informers[namespace] = informer if self._is_active: informer.start() diff --git a/cloud_pipelines_backend/orchestrator_sql.py b/cloud_pipelines_backend/orchestrator_sql.py index 178fb3bb..0c02756a 100644 --- a/cloud_pipelines_backend/orchestrator_sql.py +++ b/cloud_pipelines_backend/orchestrator_sql.py @@ -36,6 +36,10 @@ DYNAMIC_DATA_SECRET_KEY = "secret" DYNAMIC_DATA_SECRET_NAME_KEY = "name" +_FRONT_OF_RUNNING_QUEUE_LAST_PROCESSED_AT = datetime.datetime( + 1970, 1, 1, tzinfo=datetime.timezone.utc +) + class OrchestratorError(RuntimeError): pass @@ -199,6 +203,7 @@ def internal_process_queued_executions_queue(self, session: orm.Session): return False def internal_process_running_executions_queue(self, session: orm.Session): + self._move_changed_container_executions_to_front(session=session) _logger.debug("Selecting running container execution to process.)") # Select only the ID to avoid loading large JSON columns into the sort buffer. @@ -347,6 +352,26 @@ def internal_process_running_executions_queue(self, session: orm.Session): self._running_executions_queue_idle = True return False + def _move_changed_container_executions_to_front(self, session: orm.Session) -> None: + """Moves every pending or running container execution whose container the launcher saw change to the front of the running queue.""" + container_execution_ids = self._launcher.pop_changed_container_execution_ids() + if not container_execution_ids: + return + session.execute( + sql.update(bts.ContainerExecution) + .where( + bts.ContainerExecution.id.in_(container_execution_ids), + bts.ContainerExecution.status.in_( + ( + bts.ContainerExecutionStatus.PENDING, + bts.ContainerExecutionStatus.RUNNING, + ) + ), + ) + .values(last_processed_at=_FRONT_OF_RUNNING_QUEUE_LAST_PROCESSED_AT) + ) + session.commit() + def internal_process_one_queued_execution( self, session: orm.Session, execution: bts.ExecutionNode ): diff --git a/tests/launchers/kubernetes/test_kubernetes_launchers_changed_container_executions.py b/tests/launchers/kubernetes/test_kubernetes_launchers_changed_container_executions.py new file mode 100644 index 00000000..771d1339 --- /dev/null +++ b/tests/launchers/kubernetes/test_kubernetes_launchers_changed_container_executions.py @@ -0,0 +1,246 @@ +"""Tests for the Kubernetes launchers reporting which container executions' pods or jobs changed. + +The Kubernetes API calls are faked so that the tests run offline. +""" + +from __future__ import annotations + +import pathlib + +import pytest +from kubernetes import client as k8s_client_lib + +from cloud_pipelines_backend import component_structures as structures +from cloud_pipelines_backend.launchers import common_annotations +from cloud_pipelines_backend.launchers import kubernetes_launchers + +_CONTAINER_EXECUTION_ID_KEY = common_annotations.CONTAINER_EXECUTION_ID_ANNOTATION_KEY + + +class _FakeVersionApi: + def __init__(self, api_client=None): + pass + + def get_code(self, **kwargs): + return None + + +class _FakeCluster: + """Records the created pods and jobs.""" + + def __init__(self): + self.created_pods: list[k8s_client_lib.V1Pod] = [] + self.created_jobs: list[k8s_client_lib.V1Job] = [] + + def make_core_api(self): + cluster = self + + class _FakeCoreV1Api: + def __init__(self, api_client=None): + pass + + def create_namespaced_pod(self, namespace, body, **kwargs): + cluster.created_pods.append(body) + body.metadata.name = (body.metadata.generate_name or "") + "abcde" + body.metadata.namespace = namespace + return body + + def create_namespaced_service(self, namespace, body, **kwargs): + return body + + return _FakeCoreV1Api + + def make_batch_api(self): + cluster = self + + class _FakeBatchV1Api: + def __init__(self, api_client=None): + pass + + def create_namespaced_job(self, namespace, body, **kwargs): + cluster.created_jobs.append(body) + body.metadata.uid = "job-uid" + return body + + return _FakeBatchV1Api + + +class _FakeSharedInformer: + """Records the event handlers it is given and whether they were added after the cache was prefilled.""" + + def __init__(self, **kwargs): + self.handlers: list[tuple[str, object, bool]] = [] + self.listed = False + + def _initial_list(self): + self.listed = True + + def add_event_handler(self, event_type, handler): + self.handlers.append((event_type, handler, self.listed)) + + def start(self): + pass + + +@pytest.fixture +def cluster(monkeypatch) -> _FakeCluster: + cluster = _FakeCluster() + monkeypatch.setattr(k8s_client_lib, "VersionApi", _FakeVersionApi) + monkeypatch.setattr(k8s_client_lib, "CoreV1Api", cluster.make_core_api()) + monkeypatch.setattr(k8s_client_lib, "BatchV1Api", cluster.make_batch_api()) + + def _no_informer(**kwargs): + raise RuntimeError("Informers are disabled in tests.") + + monkeypatch.setattr(kubernetes_launchers, "MultiNamespaceInformer", _no_informer) + return cluster + + +def _make_launcher() -> ( + kubernetes_launchers.Local_Kubernetes_UsingHostPathStorage_KubernetesPodOrJobLauncher +): + return kubernetes_launchers.Local_Kubernetes_UsingHostPathStorage_KubernetesPodOrJobLauncher( + api_client=k8s_client_lib.ApiClient(), + ) + + +def _launch(launcher, tmp_path: pathlib.Path, annotations: dict[str, str]) -> None: + launcher.launch_container_task( + component_spec=structures.ComponentSpec( + name="test", + implementation=structures.ContainerImplementation( + container=structures.ContainerSpec(image="alpine") + ), + ), + input_arguments={}, + output_uris={}, + log_uri=str(tmp_path / "log.txt"), + annotations={_CONTAINER_EXECUTION_ID_KEY: "123", **annotations}, + ) + + +def test_launched_pod_carries_the_container_execution_id(cluster, tmp_path): + _launch(_make_launcher(), tmp_path, {}) + + [pod] = cluster.created_pods + assert pod.metadata.annotations[_CONTAINER_EXECUTION_ID_KEY] == "123" + + +def test_launched_job_pod_template_carries_the_container_execution_id( + cluster, tmp_path +): + _launch( + _make_launcher(), + tmp_path, + {kubernetes_launchers.RETRIES_MAX_RETRIES_ANNOTATION_KEY: "1"}, + ) + + [job] = cluster.created_jobs + assert job.spec.template.metadata.annotations[_CONTAINER_EXECUTION_ID_KEY] == "123" + + +def test_launching_does_not_change_the_launcher_pod_annotations(cluster, tmp_path): + launcher = _make_launcher() + + _launch(launcher, tmp_path, {}) + + assert _CONTAINER_EXECUTION_ID_KEY not in launcher._pod_launcher._pod_annotations + + +def test_reports_each_changed_pod_and_job_once(cluster): + launcher = _make_launcher() + launcher._pod_launcher._mark_container_execution_changed( + k8s_client_lib.V1Pod( + metadata=k8s_client_lib.V1ObjectMeta( + annotations={_CONTAINER_EXECUTION_ID_KEY: "pod-execution"} + ) + ) + ) + launcher._job_launcher._mark_container_execution_changed( + k8s_client_lib.V1Job( + metadata=k8s_client_lib.V1ObjectMeta(), + spec=k8s_client_lib.V1JobSpec( + template=k8s_client_lib.V1PodTemplateSpec( + metadata=k8s_client_lib.V1ObjectMeta( + annotations={_CONTAINER_EXECUTION_ID_KEY: "job-execution"} + ) + ) + ), + ) + ) + + assert sorted(launcher.pop_changed_container_execution_ids()) == [ + "job-execution", + "pod-execution", + ] + assert launcher.pop_changed_container_execution_ids() == [] + + +def test_ignores_objects_without_a_container_execution_id(cluster): + launcher = _make_launcher() + + launcher._pod_launcher._mark_container_execution_changed( + k8s_client_lib.V1Pod(metadata=k8s_client_lib.V1ObjectMeta()) + ) + + assert launcher.pop_changed_container_execution_ids() == [] + + +def test_records_changes_before_the_first_call(cluster): + launcher = _make_launcher() + launcher._pod_launcher._mark_container_execution_changed( + k8s_client_lib.V1Pod( + metadata=k8s_client_lib.V1ObjectMeta( + annotations={_CONTAINER_EXECUTION_ID_KEY: "pod-execution"} + ) + ) + ) + + assert launcher.pop_changed_container_execution_ids() == ["pod-execution"] + + +def test_keeps_at_most_the_maximum_number_of_ids_between_calls(cluster, monkeypatch): + monkeypatch.setattr(kubernetes_launchers, "_MAX_CHANGED_CONTAINER_EXECUTION_IDS", 2) + launcher = _make_launcher() + for container_execution_id in ("first", "second", "third"): + launcher._pod_launcher._mark_container_execution_changed( + k8s_client_lib.V1Pod( + metadata=k8s_client_lib.V1ObjectMeta( + annotations={_CONTAINER_EXECUTION_ID_KEY: container_execution_id} + ) + ) + ) + + assert sorted(launcher.pop_changed_container_execution_ids()) == [ + "first", + "second", + ] + + +def test_informer_reports_changes_after_prefilling_its_cache(monkeypatch): + monkeypatch.setattr( + kubernetes_launchers.k8s_informer_lib, "SharedInformer", _FakeSharedInformer + ) + on_change = object() + + informer = kubernetes_launchers.MultiNamespaceInformer( + list_func=None, initial_namespaces=["default"], on_change=on_change + ) + + assert informer._informers["default"].handlers == [ + (kubernetes_launchers.k8s_informer_lib.ADDED, on_change, True), + (kubernetes_launchers.k8s_informer_lib.MODIFIED, on_change, True), + (kubernetes_launchers.k8s_informer_lib.DELETED, on_change, True), + ] + + +def test_informer_without_on_change_adds_no_handlers(monkeypatch): + monkeypatch.setattr( + kubernetes_launchers.k8s_informer_lib, "SharedInformer", _FakeSharedInformer + ) + + informer = kubernetes_launchers.MultiNamespaceInformer( + list_func=None, initial_namespaces=["default"] + ) + + assert informer._informers["default"].handlers == [] diff --git a/tests/test_orchestrator_changed_container_executions.py b/tests/test_orchestrator_changed_container_executions.py new file mode 100644 index 00000000..1bf1e223 --- /dev/null +++ b/tests/test_orchestrator_changed_container_executions.py @@ -0,0 +1,147 @@ +"""Tests for refreshing the container executions the launcher saw change ahead of the round-robin.""" + +from typing import Any, Callable +from unittest import mock + +from sqlalchemy import orm +from sqlalchemy import sql + +from cloud_pipelines_backend import api_server_sql +from cloud_pipelines_backend import backend_types_sql as bts +from cloud_pipelines_backend import component_structures as structures +from cloud_pipelines_backend import database_ops +from cloud_pipelines_backend import orchestrator_sql +from cloud_pipelines_backend.launchers import interfaces as launcher_interfaces + + +def _make_launched_container(launcher_data: dict[str, Any]) -> mock.MagicMock: + return mock.MagicMock( + status=launcher_interfaces.ContainerStatus.PENDING, + to_dict=lambda: dict(launcher_data), + ) + + +def _container_task(image: str) -> structures.TaskSpec: + return structures.TaskSpec( + component_ref=structures.ComponentReference( + spec=structures.ComponentSpec( + implementation=structures.ContainerImplementation( + container=structures.ContainerSpec(image=image) + ) + ) + ) + ) + + +def _make_orchestrator( + session_factory: Callable[[], orm.Session], launcher: mock.MagicMock +) -> orchestrator_sql.OrchestratorService_Sql: + return orchestrator_sql.OrchestratorService_Sql( + session_factory=session_factory, + launcher=launcher, + storage_provider=mock.MagicMock(), + data_root_uri="file:///tmp/artifacts", + logs_root_uri="file:///tmp/logs", + ) + + +def _create_launched_container_executions( + task_count: int, +) -> tuple[Callable[[], orm.Session], mock.MagicMock]: + """A pipeline run with ``task_count`` launched, PENDING container executions, and the launcher that launched them.""" + db_engine = database_ops.create_db_engine_and_migrate_db(database_uri="sqlite://") + session_factory = lambda: orm.Session(bind=db_engine) + api_server_sql.PipelineRunsApiService_Sql().create( + session=session_factory(), + root_task=structures.TaskSpec( + component_ref=structures.ComponentReference( + spec=structures.ComponentSpec( + implementation=structures.GraphImplementation( + graph=structures.GraphSpec( + tasks={ + f"task{i}": _container_task(image=f"python:3.{i}") + for i in range(task_count) + } + ) + ), + ) + ) + ), + created_by="user1", + ) + launch_count = iter(range(task_count)) + launcher = mock.MagicMock( + launch_container_task=mock.MagicMock( + side_effect=lambda *a, **kw: _make_launched_container( + {"pod": f"pod-{next(launch_count)}"} + ) + ), + deserialize_launched_container_from_dict=mock.MagicMock( + side_effect=_make_launched_container + ), + get_refreshed_launched_container_from_dict=mock.MagicMock( + side_effect=_make_launched_container + ), + pop_changed_container_execution_ids=mock.MagicMock(return_value=[]), + ) + orchestrator = _make_orchestrator(session_factory, launcher) + session = session_factory() + while orchestrator.internal_process_queued_executions_queue(session=session): + pass + return session_factory, launcher + + +def _container_execution_ids_by_pod( + session_factory: Callable[[], orm.Session], +) -> dict[str, str]: + return { + container_execution.launcher_data["pod"]: container_execution.id + for container_execution in session_factory().scalars( + sql.select(bts.ContainerExecution) + ) + } + + +def test_refreshes_a_changed_container_execution_first(): + session_factory, launcher = _create_launched_container_executions(task_count=3) + changed_id = _container_execution_ids_by_pod(session_factory)["pod-2"] + launcher.pop_changed_container_execution_ids.return_value = [changed_id] + + _make_orchestrator( + session_factory, launcher + ).internal_process_running_executions_queue(session=session_factory()) + + launcher.get_refreshed_launched_container_from_dict.assert_called_once_with( + {"pod": "pod-2"} + ) + + +def test_leaves_an_ended_container_execution_out_of_the_running_queue(): + session_factory, launcher = _create_launched_container_executions(task_count=1) + container_execution_id = _container_execution_ids_by_pod(session_factory)["pod-0"] + with session_factory() as session: + session.get(bts.ContainerExecution, container_execution_id).status = ( + bts.ContainerExecutionStatus.SUCCEEDED + ) + session.commit() + launcher.pop_changed_container_execution_ids.return_value = [container_execution_id] + + _make_orchestrator( + session_factory, launcher + ).internal_process_running_executions_queue(session=session_factory()) + + last_processed_at = ( + session_factory() + .get(bts.ContainerExecution, container_execution_id) + .last_processed_at + ) + assert last_processed_at.year != 1970 + launcher.get_refreshed_launched_container_from_dict.assert_not_called() + + +def test_launchers_report_no_changes_by_default(): + class _Launcher(launcher_interfaces.ContainerTaskLauncher): + def launch_container_task(self, **kwargs): + raise NotImplementedError() + + assert _Launcher().pop_changed_container_execution_ids() == []