more scheduler tests

Signed-off-by: oliver könig <[email protected]>

Author: ko3n1g

nemo_run/run/torchx_backend/schedulers/docker.py

@@ -170,6 +170,23 @@ def describe(self, app_id: str) -> Optional[DescribeAppResponse]:
                     )
                 )
                 states.append(_state)
+            else:
+                status_file = os.path.join(req.executor.job_dir, f"status_{role}.out")
+                if os.path.exists(status_file):
+                    with open(status_file, "r") as f:
+                        status = json.load(f)
+                        roles_statuses[role].replicas.append(
+                            ReplicaStatus(
+                                id=0,
+                                role=role,
+                                state=int(status["exit_code"]),
+                                hostname=container.name,
+                            )
+                        )
+                        state = (
+                            AppState.FAILED if int(status["exit_code"]) != 0 else AppState.SUCCEEDED
+                        )
+                        states.append(state)
 
         state = AppState.UNKNOWN
         if any(is_terminal(state) for state in states):
@@ -179,19 +196,6 @@ def describe(self, app_id: str) -> Optional[DescribeAppResponse]:
                 state = AppState.FAILED
         elif len(states) > 0:
             state = next(state for state in states if not is_terminal(state))
-        else:
-            status_file = os.path.join(req.executor.job_dir, f"status_{role}.out")
-            if os.path.exists(status_file):
-                with open(status_file, "r") as f:
-                    status = json.load(f)
-                    roles_statuses[role].replicas.append(
-                        ReplicaStatus(
-                            id=0, role=role, state=int(status["exit_code"]), hostname=container.name
-                        )
-                    )
-                    state = AppState.FAILED if int(status["exit_code"]) != 0 else AppState.SUCCEEDED
-            else:
-                state = AppState.UNKNOWN
 
         return DescribeAppResponse(
             app_id=app_id,

test/run/torchx_backend/schedulers/test_docker.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
+import os
 import tempfile
 from unittest import mock
 
@@ -153,6 +155,71 @@ def test_describe_running(docker_scheduler, docker_executor):
         assert len(response.roles) == 1
 
 
+def test_describe_failed(docker_scheduler, docker_executor):
+    with (
+        mock.patch.object(DockerJobRequest, "load") as mock_load,
+        mock.patch.object(DockerContainer, "get_container") as mock_get_container,
+        mock.patch.object(PersistentDockerScheduler, "_get_app_state") as mock_get_app_state,
+    ):
+        container = DockerContainer(
+            name="test_role",
+            command=["test"],
+            executor=docker_executor,
+            extra_env={},
+        )
+        req = DockerJobRequest(
+            id="test_session___test_role___test_container_id",
+            executor=docker_executor,
+            containers=[container],
+        )
+        mock_load.return_value = req
+        mock_get_container.return_value = container
+        mock_get_app_state.return_value = None
+        status_file = os.path.join(req.executor.job_dir, f"status_{req.containers[0].name}.out")
+
+        with open(status_file, "w") as f:
+            f.write(json.dumps({"exit_code": 1}))
+
+        response = docker_scheduler.describe(req.id)
+        assert response is not None
+        assert response.app_id == req.id
+        assert "FAILED" in str(response.state)
+        assert len(response.roles) == 1
+
+
+@pytest.mark.xfail
+def test_describe_failure_not_detected(docker_scheduler, docker_executor):
+    with (
+        mock.patch.object(DockerJobRequest, "load") as mock_load,
+        mock.patch.object(DockerContainer, "get_container") as mock_get_container,
+        mock.patch.object(PersistentDockerScheduler, "_get_app_state") as mock_get_app_state,
+    ):
+        container = DockerContainer(
+            name="test_role",
+            command=["test"],
+            executor=docker_executor,
+            extra_env={},
+        )
+        req = DockerJobRequest(
+            id="test_session___test_role___test_container_id",
+            executor=docker_executor,
+            containers=[container],
+        )
+        mock_load.return_value = req
+        mock_get_container.return_value = container
+        mock_get_app_state.return_value = None
+        status_file = os.path.join(req.executor.job_dir, f"status_{req.containers[0].name}.out")
+
+        with open(status_file, "w") as f:
+            f.write(json.dumps({"exit_code": 1}))
+
+        response = docker_scheduler.describe(req.id)
+        assert response is not None
+        assert response.app_id == req.id
+        assert "SUCCEEDED" in str(response.state)
+        assert len(response.roles) == 1
+
+
 def test_save_and_get_job_dirs():
     with tempfile.TemporaryDirectory() as temp_dir:
         from nemo_run.config import set_nemorun_home