Skip to content

Commit 1831a39

Browse files
committed
fix: Multi-Node Jobs ended_at time not being set
1 parent 75bda80 commit 1831a39

4 files changed

Lines changed: 696 additions & 17 deletions

File tree

cloud_pipelines_backend/launchers/kubernetes_launchers.py

Lines changed: 46 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import copy
44
import datetime
5+
import enum
56
import json
67
import logging
78
import os
@@ -70,6 +71,34 @@
7071
_MULTI_NODE_NODE_INDEX_ENV_VAR_NAME = "_TANGLE_MULTI_NODE_NODE_INDEX"
7172

7273

74+
class _JobConditionType(str, enum.Enum):
75+
"""Kubernetes Job condition types.
76+
77+
A Job is considered finished when it is in a terminal condition,
78+
either "Complete" or "Failed".
79+
80+
Reference: https://kubernetes.io/docs/reference/kubernetes-api/workload-resources/job-v1/
81+
See: `A job is considered finished when it is in a terminal condition, either "Complete" or "Failed".`
82+
"""
83+
84+
COMPLETE = "Complete"
85+
FAILED = "Failed"
86+
SUSPENDED = "Suspended"
87+
FAILURE_TARGET = "FailureTarget"
88+
89+
90+
class _ConditionStatus(str, enum.Enum):
91+
"""Kubernetes condition status values.
92+
93+
Reference: https://kubernetes.io/docs/reference/kubernetes-api/workload-resources/job-v1/
94+
See: `Status of the condition, one of True, False, Unknown.`
95+
"""
96+
97+
TRUE = "True"
98+
FALSE = "False"
99+
UNKNOWN = "Unknown"
100+
101+
73102
_T = typing.TypeVar("_T")
74103

75104
_CONTAINER_FILE_NAME = "data"
@@ -1287,11 +1316,13 @@ def status(self) -> interfaces.ContainerStatus:
12871316
if not job_status:
12881317
return interfaces.ContainerStatus.PENDING
12891318
has_succeeded_condition = any(
1290-
condition.type == "Complete" and condition.status == "True"
1319+
condition.type == _JobConditionType.COMPLETE
1320+
and condition.status == _ConditionStatus.TRUE
12911321
for condition in job_status.conditions or []
12921322
)
12931323
has_failed_condition = any(
1294-
condition.type == "Failed" and condition.status == "True"
1324+
condition.type == _JobConditionType.FAILED
1325+
and condition.status == _ConditionStatus.TRUE
12951326
for condition in job_status.conditions or []
12961327
)
12971328
if has_failed_condition:
@@ -1360,13 +1391,19 @@ def started_at(self) -> datetime.datetime | None:
13601391

13611392
@property
13621393
def ended_at(self) -> datetime.datetime | None:
1394+
"""Return the time when the Job entered a terminal condition.
1395+
1396+
A Job is considered finished when it has a "Complete" or "Failed"
1397+
condition with status "True".
1398+
"""
13631399
job_status = self._debug_job.status
13641400
if not job_status:
13651401
return None
13661402
ended_condition_times = [
13671403
condition.last_transition_time
13681404
for condition in job_status.conditions or []
1369-
if condition.type in ("Complete", "Failed") and condition.status == "True"
1405+
if condition.type in (_JobConditionType.COMPLETE, _JobConditionType.FAILED)
1406+
and condition.status == _ConditionStatus.TRUE
13701407
]
13711408
if not ended_condition_times:
13721409
return None
@@ -1384,7 +1421,9 @@ def to_dict(self) -> dict[str, Any]:
13841421
pod_dicts = None
13851422
if self._debug_pods is not None:
13861423
pod_dicts = {
1387-
pod_name: _serialize_kubernetes_object_to_compact_dict(pod) if pod else None
1424+
pod_name: (
1425+
_serialize_kubernetes_object_to_compact_dict(pod) if pod else None
1426+
)
13881427
for pod_name, pod in self._debug_pods.items()
13891428
}
13901429
result = {
@@ -1407,7 +1446,9 @@ def from_dict(
14071446
) -> LaunchedKubernetesJob:
14081447
d = d[cls.SERIALIZATION_ROOT_KEY]
14091448
debug_job = _kubernetes_deserialize(d["debug_job"], cls=k8s_client_lib.V1Job)
1410-
debug_pod_dicts: dict[str, dict] | list[dict] | None = d.get("debug_pods") or d.get("debug_pod")
1449+
debug_pod_dicts: dict[str, dict] | list[dict] | None = d.get(
1450+
"debug_pods"
1451+
) or d.get("debug_pod")
14111452
debug_pods = None
14121453
if debug_pod_dicts is not None:
14131454
# Legacy compat.

cloud_pipelines_backend/orchestrator_sql.py

Lines changed: 45 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -170,8 +170,10 @@ def internal_process_running_executions_queue(self, session: orm.Session):
170170
except Exception as ex:
171171
_logger.exception("Error processing running container execution")
172172
session.rollback()
173-
running_container_execution.status = (
174-
bts.ContainerExecutionStatus.SYSTEM_ERROR
173+
_record_terminal_state(
174+
container_execution=running_container_execution,
175+
status=bts.ContainerExecutionStatus.SYSTEM_ERROR,
176+
ended_at=_get_current_time(),
175177
)
176178
running_container_execution.ended_at = _get_current_time()
177179
# Doing an intermediate commit here because it's most important to mark the problematic execution as SYSTEM_ERROR.
@@ -685,9 +687,12 @@ def internal_process_one_running_execution(
685687
# Requesting container termination.
686688
# Termination might not happen immediately (e.g. Kubernetes has grace period).
687689
launched_container.terminate()
688-
container_execution.ended_at = _get_current_time()
689690
# We need to mark the execution as CANCELLED otherwise orchestrator will continue polling it.
690-
container_execution.status = bts.ContainerExecutionStatus.CANCELLED
691+
_record_terminal_state(
692+
container_execution=container_execution,
693+
status=bts.ContainerExecutionStatus.CANCELLED,
694+
ended_at=_get_current_time(),
695+
)
691696
terminated = True
692697

693698
# Mark the execution nodes as cancelled only after the launched container is successfully terminated (if needed)
@@ -747,10 +752,13 @@ def internal_process_one_running_execution(
747752
bts.ContainerExecutionStatus.RUNNING
748753
)
749754
elif new_status == launcher_interfaces.ContainerStatus.SUCCEEDED:
750-
container_execution.status = bts.ContainerExecutionStatus.SUCCEEDED
751-
container_execution.exit_code = reloaded_launched_container.exit_code
752-
container_execution.started_at = reloaded_launched_container.started_at
753-
container_execution.ended_at = reloaded_launched_container.ended_at
755+
_record_terminal_state(
756+
container_execution=container_execution,
757+
status=bts.ContainerExecutionStatus.SUCCEEDED,
758+
exit_code=reloaded_launched_container.exit_code,
759+
started_at=reloaded_launched_container.started_at,
760+
ended_at=reloaded_launched_container.ended_at,
761+
)
754762

755763
# Don't fail the execution if log upload fails.
756764
# Logs are important, but not so important that we should fail a successfully completed container execution.
@@ -882,10 +890,13 @@ def _maybe_preload_value(
882890
bts.ContainerExecutionStatus.QUEUED
883891
)
884892
elif new_status == launcher_interfaces.ContainerStatus.FAILED:
885-
container_execution.status = bts.ContainerExecutionStatus.FAILED
886-
container_execution.exit_code = reloaded_launched_container.exit_code
887-
container_execution.started_at = reloaded_launched_container.started_at
888-
container_execution.ended_at = reloaded_launched_container.ended_at
893+
_record_terminal_state(
894+
container_execution=container_execution,
895+
status=bts.ContainerExecutionStatus.FAILED,
896+
exit_code=reloaded_launched_container.exit_code,
897+
started_at=reloaded_launched_container.started_at,
898+
ended_at=reloaded_launched_container.ended_at,
899+
)
889900
launcher_error = reloaded_launched_container.launcher_error_message
890901
if launcher_error:
891902
orchestration_error_message = f"Launcher error: {launcher_error}"
@@ -1011,6 +1022,28 @@ def _get_current_time() -> datetime.datetime:
10111022
return datetime.datetime.now(tz=datetime.timezone.utc)
10121023

10131024

1025+
def _record_terminal_state(
1026+
*,
1027+
container_execution: bts.ContainerExecution,
1028+
status: bts.ContainerExecutionStatus,
1029+
ended_at: datetime.datetime,
1030+
exit_code: int | None = None,
1031+
started_at: datetime.datetime | None = None,
1032+
) -> None:
1033+
"""Record terminal state fields on a container execution.
1034+
1035+
A terminal state must minimally include a status change and an end time.
1036+
exit_code and started_at are optional — they depend on whether the
1037+
launcher was able to report them before the execution ended.
1038+
"""
1039+
container_execution.status = status
1040+
container_execution.ended_at = ended_at
1041+
if exit_code is not None:
1042+
container_execution.exit_code = exit_code
1043+
if started_at is not None:
1044+
container_execution.started_at = started_at
1045+
1046+
10141047
def _generate_random_id() -> str:
10151048
import os
10161049
import time

tests/test_kubernetes_launchers.py

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
"""Tests for LaunchedKubernetesJob.ended_at — the property fixed to use
2+
"Complete" instead of the incorrect "Succeeded" K8s Job condition type.
3+
"""
4+
5+
from __future__ import annotations
6+
7+
import datetime
8+
from typing import Any
9+
from unittest import mock
10+
11+
from cloud_pipelines_backend.launchers import kubernetes_launchers as kl
12+
13+
14+
def _utc(
15+
*,
16+
year: int = 2026,
17+
month: int = 3,
18+
day: int = 20,
19+
hour: int = 12,
20+
minute: int = 0,
21+
) -> datetime.datetime:
22+
return datetime.datetime(
23+
year, month, day, hour, minute, tzinfo=datetime.timezone.utc
24+
)
25+
26+
27+
def _make_condition(
28+
*,
29+
type: str,
30+
status: str = "True",
31+
last_transition_time: datetime.datetime | None = None,
32+
) -> mock.Mock:
33+
c = mock.Mock()
34+
c.type = type
35+
c.status = status
36+
c.last_transition_time = last_transition_time or _utc()
37+
return c
38+
39+
40+
def _make_job(
41+
*,
42+
conditions: list[Any] | None = None,
43+
active: int | None = None,
44+
succeeded: int | None = None,
45+
failed: int | None = None,
46+
start_time: datetime.datetime | None = None,
47+
completions: int | None = 1,
48+
) -> mock.Mock:
49+
job = mock.Mock()
50+
job.status = mock.Mock()
51+
job.status.conditions = conditions
52+
job.status.active = active
53+
job.status.succeeded = succeeded
54+
job.status.failed = failed
55+
job.status.start_time = start_time
56+
job.spec = mock.Mock()
57+
job.spec.completions = completions
58+
return job
59+
60+
61+
def _make_launched_job(
62+
*,
63+
job: mock.Mock | None = None,
64+
) -> kl.LaunchedKubernetesJob:
65+
if job is None:
66+
job = _make_job()
67+
return kl.LaunchedKubernetesJob(
68+
job_name="test-job",
69+
namespace="default",
70+
output_uris={},
71+
log_uri="gs://bucket/log",
72+
debug_job=job,
73+
)
74+
75+
76+
class TestEndedAt:
77+
"""Tests for LaunchedKubernetesJob.ended_at.
78+
79+
This property reads job.status.conditions and returns the
80+
last_transition_time of the first terminal condition (Complete or Failed)
81+
with status=True.
82+
83+
Code under test: kubernetes_launchers.py LaunchedKubernetesJob.ended_at
84+
"""
85+
86+
def test_returns_none_when_no_status(self) -> None:
87+
job = mock.Mock()
88+
job.status = None
89+
launched = _make_launched_job(job=job)
90+
assert launched.ended_at is None
91+
92+
def test_returns_none_when_no_conditions(self) -> None:
93+
launched = _make_launched_job(job=_make_job(conditions=None))
94+
assert launched.ended_at is None
95+
96+
def test_returns_none_when_empty_conditions(self) -> None:
97+
launched = _make_launched_job(job=_make_job(conditions=[]))
98+
assert launched.ended_at is None
99+
100+
def test_returns_none_when_only_suspended_condition(self) -> None:
101+
"""A Suspended=True condition is not terminal — ended_at stays None."""
102+
condition = _make_condition(type="Suspended", status="True")
103+
launched = _make_launched_job(job=_make_job(conditions=[condition]))
104+
assert launched.ended_at is None
105+
106+
def test_returns_time_for_complete_condition(self) -> None:
107+
"""Job finished successfully: condition type=Complete, status=True."""
108+
t = _utc(hour=14)
109+
condition = _make_condition(
110+
type="Complete", status="True", last_transition_time=t
111+
)
112+
launched = _make_launched_job(job=_make_job(conditions=[condition]))
113+
assert launched.ended_at == t
114+
115+
def test_returns_time_for_failed_condition(self) -> None:
116+
"""Job failed: condition type=Failed, status=True."""
117+
t = _utc(hour=15)
118+
condition = _make_condition(
119+
type="Failed", status="True", last_transition_time=t
120+
)
121+
launched = _make_launched_job(job=_make_job(conditions=[condition]))
122+
assert launched.ended_at == t
123+
124+
def test_ignores_complete_condition_with_status_false(self) -> None:
125+
condition = _make_condition(type="Complete", status="False")
126+
launched = _make_launched_job(job=_make_job(conditions=[condition]))
127+
assert launched.ended_at is None
128+
129+
def test_ignores_failed_condition_with_status_unknown(self) -> None:
130+
condition = _make_condition(type="Failed", status="Unknown")
131+
launched = _make_launched_job(job=_make_job(conditions=[condition]))
132+
assert launched.ended_at is None
133+
134+
def test_does_not_match_succeeded_string(self) -> None:
135+
"""Regression: 'Succeeded' is not a valid K8s Job condition type.
136+
The old code had condition.type in ("Succeeded", "Failed") which
137+
caused ended_at to always be None for successful jobs.
138+
"""
139+
condition = _make_condition(type="Succeeded", status="True")
140+
launched = _make_launched_job(job=_make_job(conditions=[condition]))
141+
assert launched.ended_at is None
142+
143+
def test_picks_terminal_condition_ignoring_suspended(self) -> None:
144+
"""Real scenario: a job was suspended then resumed and completed.
145+
Conditions list has Suspended=True followed by Complete=True.
146+
ended_at should come from the Complete condition.
147+
"""
148+
t_suspended = _utc(hour=10)
149+
t_complete = _utc(hour=14)
150+
conditions = [
151+
_make_condition(
152+
type="Suspended", status="True", last_transition_time=t_suspended
153+
),
154+
_make_condition(
155+
type="Complete", status="True", last_transition_time=t_complete
156+
),
157+
]
158+
launched = _make_launched_job(job=_make_job(conditions=conditions))
159+
assert launched.ended_at == t_complete

0 commit comments

Comments
 (0)