Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions cloud_pipelines_backend/launchers/interfaces.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,10 @@ def get_refreshed_launched_container_from_dict(
) -> _TLaunchedContainer:
raise NotImplementedError()

def pop_changed_container_execution_ids(self) -> list[str]:

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Interface change

Only the launcher knows which containers changed (its pod and job informers already see every change), while only the orchestrator decides which container executions to refresh next, so a single method on the launcher interface is the smallest way to pass that signal from one to the other. It defaults to [], so launchers that can't tell (Davies, local and Docker launchers, routers that don't delegate) keep the current round-robin with no changes, and it doesn't commit the interface to anything Kubernetes-specific.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alternatives that leave ContainerTaskLauncher unchanged:

  • Orchestrator kwarg: OrchestratorService_Sql takes an optional function that returns the changed ids, and the deployment passes in the Kubernetes launcher's method.
  • Push callback: Kubernetes launchers take an on_container_execution_changed callback that feeds a set the orchestrator owns.
  • Optional capability: a separate Protocol, or a getattr check, that only Kubernetes launchers implement.
  • Launcher writes to the DB: the informer moves the row to the front itself, but that breaks layering because launchers have no DB session.
  • Downstream patch: keep it as a deployment-side monkeypatch, which is fragile because it depends on informer internals.

Why an interface addition?

It's the only option where the launcher reports changes and the orchestrator consumes them through the same typed contract as the rest of the launcher API. The default [] means no other launcher or deployment needs any change, while the alternatives either spread the wiring across each deployment or hide the same method behind a looser contract.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Ark-kun Do you want any alternative?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I gave this a critical review and propose a different approach.

Decision doc:

https://docs.google.com/document/d/1hH0iTcVH2g8biNEoEQX--OPL5A6y0iCv8xmwekoUnro/edit?usp=sharing

"""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"
Expand Down
67 changes: 67 additions & 0 deletions cloud_pipelines_backend/launchers/kubernetes_launchers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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 {}) | {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(AI Assisted)

N1 · ▫️ nit — the new annotation write duplicates what oasis's pod postprocessor already does

▫️ Severity: nit — cosmetic, no behaviour change; ignore freely.

Issue. kubernetes_launchers.py:498 copies the container-execution id from the launch
annotations onto the pod's annotations. In oasis that annotation is already written, with the same
key and the same value, by the shared pod postprocessor at oasis-backend launchers/k8s.py:112 —
reached by both the GKE launchers (launchers/gke.py:41) and Nebius (launchers/nebius.py:210).

The duplication is harmless — same key, same value, merged with | — and it is the right call
upstream, where the postprocessor does not exist and the informer path would otherwise never fire.
Worth stating in the PR body so nobody later "cleans up" the oasis one and assumes the OSS write
has been there all along.

Counter-argument, stated honestly. Two writers of the same constant is better than a fast path
that silently depends on a downstream repository. Nit.

common_annotations.CONTAINER_EXECUTION_ID_ANNOTATION_KEY: container_execution_id
}

pod = k8s_client_lib.V1Pod(
api_version="v1",
kind="Pod",
Expand All @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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":
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand Down
25 changes: 25 additions & 0 deletions cloud_pipelines_backend/orchestrator_sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
):
Expand Down
Loading
Loading