more tests

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

Author: ko3n1g

nemo_run/core/execution/dgxcloud.py

@@ -368,8 +368,8 @@ def fetch_logs(
     ) -> Iterable[str]:
         token = self.get_auth_token()
         if not token:
-            logger.error("Failed to retrieve auth token for cancellation request.")
-            return
+            logger.error("Failed to retrieve auth token for fetch logs request.")
+            yield ""
 
         response = requests.get(
             f"{self.base_url}/workloads", headers=self._default_headers(token=token)
@@ -383,6 +383,7 @@ def fetch_logs(
             None,
         )
         if workload_name is None:
+            logger.error(f"No workload found with id {job_id}")
             yield ""
 
         urls = [
@@ -415,7 +416,15 @@ def fetch_logs(
         # Yield chunks as they arrive
         while active_urls:
             url, item = q.get()
-            if item is None:
+            if item is None or self.status(job_id) in [
+                DGXCloudState.DELETING,
+                DGXCloudState.STOPPED,
+                DGXCloudState.STOPPING,
+                DGXCloudState.DEGRADED,
+                DGXCloudState.FAILED,
+                DGXCloudState.COMPLETED,
+                DGXCloudState.TERMINATING,
+            ]:
                 active_urls.discard(url)
             else:
                 yield item

test/core/execution/test_dgxcloud.py

@@ -13,7 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import os
+import queue
 import subprocess
 import tempfile
 from unittest.mock import MagicMock, mock_open, patch
@@ -101,8 +103,67 @@ def test_get_auth_token_failure(self, mock_post):
 
         assert token is None
 
+    def test_fetch_no_token(self, caplog):
+        with (
+            patch.object(DGXCloudExecutor, "get_auth_token", return_value=None),
+            caplog.at_level(logging.ERROR),
+        ):
+            executor = DGXCloudExecutor(
+                base_url="https://dgxapi.example.com",
+                kube_apiserver_url="https://127.0.0.1:443",
+                app_id="test_app_id",
+                app_secret="test_app_secret",
+                project_name="test_project",
+                container_image="nvcr.io/nvidia/test:latest",
+                pvc_nemo_run_dir="/workspace/nemo_run",
+            )
+
+            logs_iter = executor.fetch_logs("123", stream=True)
+            assert next(logs_iter) == ""
+            assert (
+                caplog.records[-1].message
+                == "Failed to retrieve auth token for fetch logs request."
+            )
+            assert caplog.records[-1].levelname == "ERROR"
+            caplog.clear()
+
     @patch("nemo_run.core.execution.dgxcloud.requests.get")
-    def test_fetch_logs(self, mock_requests_get):
+    def test_fetch_no_workload_with_name(self, mock_requests_get, caplog):
+        mock_workloads_response = MagicMock(spec=requests.Response)
+        mock_workloads_response.json.return_value = {
+            "workloads": [{"name": "hello-world", "id": "123"}]
+        }
+
+        mock_requests_get.side_effect = [mock_workloads_response]
+
+        with (
+            patch.object(DGXCloudExecutor, "get_auth_token", return_value="test_token"),
+            caplog.at_level(logging.ERROR),
+        ):
+            executor = DGXCloudExecutor(
+                base_url="https://dgxapi.example.com",
+                kube_apiserver_url="https://127.0.0.1:443",
+                app_id="test_app_id",
+                app_secret="test_app_secret",
+                project_name="test_project",
+                container_image="nvcr.io/nvidia/test:latest",
+                pvc_nemo_run_dir="/workspace/nemo_run",
+            )
+
+            logs_iter = executor.fetch_logs("this-workload-does-not-exist", stream=True)
+            assert next(logs_iter) == ""
+            assert (
+                caplog.records[-1].message
+                == "No workload found with id this-workload-does-not-exist"
+            )
+            assert caplog.records[-1].levelname == "ERROR"
+            caplog.clear()
+
+    @patch("nemo_run.core.execution.dgxcloud.requests.get")
+    @patch("nemo_run.core.execution.dgxcloud.time.sleep")
+    @patch("nemo_run.core.execution.dgxcloud.threading.Thread")
+    @patch("nemo_run.core.execution.dgxcloud.queue.Queue")
+    def test_fetch_logs(self, mock_queue, mock_threading_Thread, mock_sleep, mock_requests_get):
         # --- 1. Setup Primitives for the *live* test ---
         mock_log_response = MagicMock(spec=requests.Response)
 
@@ -117,12 +178,29 @@ def test_fetch_logs(self, mock_requests_get):
             "workloads": [{"name": "hello-world", "id": "123"}]
         }
 
+        mock_queue_instance = MagicMock()
+        mock_queue_instance.get.side_effect = [
+            (
+                "https://127.0.0.1:443/api/v1/namespaces/runai-test_project/pods/hello-world-worker-0/log?container=pytorch&follow=true",
+                "this is a static log\n",
+            ),
+            (
+                "https://127.0.0.1:443/api/v1/namespaces/runai-test_project/pods/hello-world-worker-0/log?container=pytorch&follow=true",
+                None,
+            ),
+            (
+                "https://127.0.0.1:443/api/v1/namespaces/runai-test_project/pods/hello-world-worker-1/log?container=pytorch&follow=true",
+                None,
+            ),
+        ]
+
         mock_requests_get.side_effect = [mock_workloads_response, mock_log_response]
 
         # --- 4. Setup Executor (inside the patch) ---
         with (
             patch.object(DGXCloudExecutor, "get_auth_token", return_value="test_token"),
             patch.object(DGXCloudExecutor, "status", return_value=DGXCloudState.RUNNING),
+            patch("nemo_run.core.execution.dgxcloud.queue.Queue", return_value=mock_queue_instance),
         ):
             executor = DGXCloudExecutor(
                 base_url="https://dgxapi.example.com",
@@ -132,24 +210,31 @@ def test_fetch_logs(self, mock_requests_get):
                 project_name="test_project",
                 container_image="nvcr.io/nvidia/test:latest",
                 pvc_nemo_run_dir="/workspace/nemo_run",
+                nodes=2,
             )
 
             logs_iter = executor.fetch_logs("123", stream=True)
 
             assert next(logs_iter) == "this is a static log\n"
-            assert next(logs_iter) == "this is the last static log\n"
 
-            mock_requests_get.assert_any_call(
-                "https://127.0.0.1:443/api/v1/namespaces/runai-test_project/pods/hello-world-worker-0/log?container=pytorch&follow=true",
-                headers={
-                    "Accept": "application/json",
-                    "Content-Type": "application/json",
-                    "Authorization": "Bearer test_token",
-                },
-                verify=False,
-                stream=True,
-            )
+            mock_sleep.assert_called_once_with(10)
 
+            mock_threading_Thread.assert_any_call(
+                target=executor._stream_url_sync,
+                args=(
+                    "https://127.0.0.1:443/api/v1/namespaces/runai-test_project/pods/hello-world-worker-0/log?container=pytorch&follow=true",
+                    executor._default_headers(token="test_token"),
+                    mock_queue_instance,
+                ),
+            )
+            mock_threading_Thread.assert_any_call(
+                target=executor._stream_url_sync,
+                args=(
+                    "https://127.0.0.1:443/api/v1/namespaces/runai-test_project/pods/hello-world-worker-1/log?container=pytorch&follow=true",
+                    executor._default_headers(token="test_token"),
+                    mock_queue_instance,
+                ),
+            )
             with pytest.raises(StopIteration):
                 next(logs_iter)
 
@@ -1007,3 +1092,4 @@ def test_default_headers_with_token(self):
         assert headers["Content-Type"] == "application/json"
         assert "Authorization" in headers
         assert headers["Authorization"] == "Bearer test_token"
+        assert headers["Authorization"] == "Bearer test_token"