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() == []