fix: treat DGXCloud UNKNOWN/transient status as PENDING to avoid false failures

When a job is submitted to DGXCloud, the API may transiently return a
non-200 response or an "Unknown" phase before the workload is fully
registered. Previously this was mapped to AppState.FAILED, causing
wait_and_exit() to treat the job as terminated immediately while the
pod was still starting up on the cluster.

- DGXCloudState.UNKNOWN now maps to AppState.PENDING in DGX_STATES
- executor.status() returns None (instead of DGXCloudState.UNKNOWN)
  on non-200 HTTP responses so transient API errors don't look like
  a real "Unknown" phase reported by the scheduler
- describe() fallback for unknown keys in DGX_STATES changed to PENDING
- Tests updated and added to cover all three code paths

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: oliver könig <okoenig@nvidia.com>

Author: ko3n1g

nemo_run/core/execution/dgxcloud.py

@@ -369,7 +369,11 @@ def status(self, job_id: str) -> Optional[DGXCloudState]:
         headers = self._default_headers(token=token)
         response = requests.get(url, headers=headers)
         if response.status_code != 200:
-            return DGXCloudState("Unknown")
+            logger.warning(
+                f"Failed to get status for job {job_id}, "
+                f"status_code={response.status_code}. Treating as transient."
+            )
+            return None
 
         r_json = response.json()
         return DGXCloudState(r_json["phase"])

nemo_run/run/torchx_backend/schedulers/dgxcloud.py

@@ -59,7 +59,7 @@
     DGXCloudState.FAILED: AppState.FAILED,
     DGXCloudState.COMPLETED: AppState.SUCCEEDED,
     DGXCloudState.TERMINATING: AppState.RUNNING,
-    DGXCloudState.UNKNOWN: AppState.FAILED,
+    DGXCloudState.UNKNOWN: AppState.PENDING,
 }
 
 log = logging.getLogger(__name__)
@@ -192,7 +192,7 @@ def describe(self, app_id: str) -> Optional[DescribeAppResponse]:
             return None
 
         dgx_state = executor.status(job_id) or DGXCloudState.UNKNOWN
-        app_state = DGX_STATES.get(dgx_state, AppState.UNKNOWN)
+        app_state = DGX_STATES.get(dgx_state, AppState.PENDING)
         roles_statuses[0].replicas[0].state = app_state
 
         return DescribeAppResponse(

test/core/execution/test_dgxcloud.py

@@ -936,7 +936,7 @@ def test_status_error_response(self, mock_get):
 
             status = executor.status("job123")
 
-            assert status == DGXCloudState.UNKNOWN
+            assert status is None
 
     @patch("requests.get")
     def test_cancel(self, mock_get):

test/run/torchx_backend/schedulers/test_dgxcloud.py

@@ -19,10 +19,14 @@
 
 import pytest
 from torchx.schedulers.api import AppDryRunInfo
-from torchx.specs import AppDef, Role
+from torchx.specs import AppDef, AppState, Role
 
-from nemo_run.core.execution.dgxcloud import DGXCloudExecutor
-from nemo_run.run.torchx_backend.schedulers.dgxcloud import DGXCloudScheduler, create_scheduler
+from nemo_run.core.execution.dgxcloud import DGXCloudExecutor, DGXCloudState
+from nemo_run.run.torchx_backend.schedulers.dgxcloud import (
+    DGX_STATES,
+    DGXCloudScheduler,
+    create_scheduler,
+)
 
 
 @pytest.fixture
@@ -184,6 +188,56 @@ def test_log_iter(dgx_cloud_scheduler, dgx_cloud_executor):
         assert logs == ["log2", "log3"]
 
 
+def test_unknown_state_maps_to_pending_not_failed():
+    # DGXCloudState.UNKNOWN must map to PENDING so transient API errors during
+    # job startup do not cause wait_and_exit() to treat the job as terminal.
+    assert DGX_STATES[DGXCloudState.UNKNOWN] == AppState.PENDING
+
+
+def test_describe_returns_pending_when_status_is_none(dgx_cloud_scheduler, dgx_cloud_executor):
+    # Regression test: executor.status() returns None when the auth token is
+    # missing or the API call fails transiently right after job submission.
+    # describe() must return PENDING so the wait loop keeps polling.
+    with (
+        mock.patch(
+            "nemo_run.run.torchx_backend.schedulers.dgxcloud._get_job_dirs"
+        ) as mock_get_job_dirs,
+        mock.patch.object(DGXCloudExecutor, "status", return_value=None),
+    ):
+        mock_get_job_dirs.return_value = {
+            "test_experiment___test_role___test_job_id": {
+                "job_status": "RUNNING",
+                "executor": dgx_cloud_executor,
+            }
+        }
+
+        response = dgx_cloud_scheduler.describe("test_experiment___test_role___test_job_id")
+        assert response is not None
+        assert response.state == AppState.PENDING
+
+
+def test_describe_returns_pending_when_status_is_unknown(dgx_cloud_scheduler, dgx_cloud_executor):
+    # Regression test: the DGXCloud API transiently returns "Unknown" before a
+    # job is visible (e.g. HTTP 404 right after submission). describe() must
+    # return PENDING so the wait loop keeps polling instead of failing.
+    with (
+        mock.patch(
+            "nemo_run.run.torchx_backend.schedulers.dgxcloud._get_job_dirs"
+        ) as mock_get_job_dirs,
+        mock.patch.object(DGXCloudExecutor, "status", return_value=DGXCloudState.UNKNOWN),
+    ):
+        mock_get_job_dirs.return_value = {
+            "test_experiment___test_role___test_job_id": {
+                "job_status": "RUNNING",
+                "executor": dgx_cloud_executor,
+            }
+        }
+
+        response = dgx_cloud_scheduler.describe("test_experiment___test_role___test_job_id")
+        assert response is not None
+        assert response.state == AppState.PENDING
+
+
 def test_log_iter_str(dgx_cloud_scheduler, dgx_cloud_executor):
     with mock.patch(
         "nemo_run.run.torchx_backend.schedulers.dgxcloud._get_job_dirs"