From 88821a74a2194685984eaac3a1c73281dee2066f Mon Sep 17 00:00:00 2001 From: brightyorcerf Date: Mon, 7 Sep 2026 13:59:27 +0530 Subject: [PATCH] fix(trainer): don't report Complete when a container status is Unknown aggregate_status_from_containers() filtered Unknown out of the set before checking for completion, so ["Complete", "Unknown"] aggregated to Complete. Unknown means "not determined", not "no signal to weigh": get_container_status() returns it whenever the adapter raises during inspect, and container_status_to_trainjob_status() falls through to it for paused, dead and restarting. Dropping it from the check treated an undetermined container as if it agreed with the others, so on a two-node job where node-0 exited 0 and node-1 could not be inspected, wait_for_job_status(status={"Complete"}) returned on that poll with node-1's outcome never determined. Check every status instead. The empty-list guard from #562 is preserved, and Failed and Running still short-circuit first, so ["Created", "Unknown"] remains Created and the only behavior that changes is the mixed Complete case. Fixes #787 Signed-off-by: brightyorcerf --- .../backends/container/backend_test.py | 30 ++++++++++++++++++- kubeflow/trainer/backends/container/utils.py | 3 +- .../trainer/backends/container/utils_test.py | 4 +-- 3 files changed, 32 insertions(+), 5 deletions(-) diff --git a/kubeflow/trainer/backends/container/backend_test.py b/kubeflow/trainer/backends/container/backend_test.py index 37e8f92e0..4ed89bd87 100644 --- a/kubeflow/trainer/backends/container/backend_test.py +++ b/kubeflow/trainer/backends/container/backend_test.py @@ -895,13 +895,25 @@ def mock_create_with_status(*args, **kwargs): config={"wait_status": constants.TRAINJOB_COMPLETE, "container_exit_code": 1}, expected_error=RuntimeError, ), + TestCase( + name="unknown node blocks complete", + expected_status=FAILED, + config={ + "wait_status": constants.TRAINJOB_COMPLETE, + "num_nodes": 2, + "timeout": 2, + }, + expected_error=TimeoutError, + ), ], ) def test_wait_for_job_status(container_backend, test_case): """Test waiting for job status.""" print("Executing test:", test_case.name) try: - trainer = types.CustomTrainer(func=simple_train_func, num_nodes=1) + trainer = types.CustomTrainer( + func=simple_train_func, num_nodes=test_case.config.get("num_nodes", 1) + ) runtime = container_backend.get_runtime(constants.DEFAULT_TRAINING_RUNTIME) job_name = container_backend.train(runtime=runtime, trainer=trainer) @@ -936,6 +948,22 @@ def test_wait_for_job_status(container_backend, test_case): job_name, status={test_case.config["wait_status"]}, timeout=5, polling_interval=1 ) + elif test_case.name == "unknown node blocks complete": + # node-0 finishes successfully, node-1 sits in a state that maps to + # Unknown. The job must not be reported Complete on node-0 alone. + node_0, node_1 = container_backend._adapter.containers_created[:2] + container_backend._adapter.set_container_status(node_0["id"], "exited", 0) + container_backend._adapter.set_container_status(node_1["id"], "paused") + + container_backend.wait_for_job_status( + job_name, + status={test_case.config["wait_status"]}, + timeout=test_case.config["timeout"], + polling_interval=1, + ) + + assert test_case.expected_status == SUCCESS + except Exception as e: assert type(e) is test_case.expected_error if test_case.name == "job fails": diff --git a/kubeflow/trainer/backends/container/utils.py b/kubeflow/trainer/backends/container/utils.py index 80588a219..1a253a9e2 100644 --- a/kubeflow/trainer/backends/container/utils.py +++ b/kubeflow/trainer/backends/container/utils.py @@ -160,8 +160,7 @@ def aggregate_status_from_containers(container_statuses: list[str]) -> str: return constants.TRAINJOB_FAILED if constants.TRAINJOB_RUNNING in container_statuses: return constants.TRAINJOB_RUNNING - known_statuses = [s for s in container_statuses if s != UNKNOWN] - if known_statuses and all(s == constants.TRAINJOB_COMPLETE for s in known_statuses): + if container_statuses and all(s == constants.TRAINJOB_COMPLETE for s in container_statuses): return constants.TRAINJOB_COMPLETE if any(s == constants.TRAINJOB_CREATED for s in container_statuses): return constants.TRAINJOB_CREATED diff --git a/kubeflow/trainer/backends/container/utils_test.py b/kubeflow/trainer/backends/container/utils_test.py index adaa1aa15..d71fa6c05 100644 --- a/kubeflow/trainer/backends/container/utils_test.py +++ b/kubeflow/trainer/backends/container/utils_test.py @@ -133,10 +133,10 @@ def test_build_pip_install_cmd(test_case: TestCase): expected_output=constants.TRAINJOB_COMPLETE, ), TestCase( - name="complete with unknown is complete", + name="complete with unknown is not complete", expected_status=SUCCESS, config={"statuses": [constants.TRAINJOB_COMPLETE, UNKNOWN]}, - expected_output=constants.TRAINJOB_COMPLETE, + expected_output=UNKNOWN, ), TestCase( name="failed takes precedence",