feat: Stream DGXC logs

.vscode/settings.json

@@ -19,6 +19,7 @@
     "python.testing.pytestArgs": [
         "test"
     ],
+    "flake8.enabled": false,
     "python.testing.unittestEnabled": false,
     "python.testing.pytestEnabled": true
 }

nemo_run/core/execution/dgxcloud.py

@@ -17,13 +17,15 @@
 import json
 import logging
 import os
+import queue
 import subprocess
 import tempfile
+import threading
 import time
 from dataclasses import dataclass, field
 from enum import Enum
 from pathlib import Path
-from typing import Any, Optional, Type
+from typing import Any, Iterable, Optional
 
 import requests
 from invoke.context import Context
@@ -65,6 +67,7 @@ class DGXCloudExecutor(Executor):
     """
 
     base_url: str
+    kube_apiserver_url: str
     app_id: str
     app_secret: str
     project_name: str
@@ -359,6 +362,92 @@ def status(self, job_id: str) -> Optional[DGXCloudState]:
         r_json = response.json()
         return DGXCloudState(r_json["phase"])
 
+    def _stream_url_sync(self, url: str, headers: dict, q: queue.Queue):
+        """Stream a single URL using requests and put chunks into the queue"""
+        try:
+            with requests.get(url, stream=True, headers=headers, verify=False) as response:
+                for line in response.iter_lines(decode_unicode=True):
+                    q.put((url, f"{line}\n"))
+        except Exception as e:
+            logger.error(f"Error streaming URL {url}: {e}")
+
+        finally:
+            q.put((url, None))
+
+    def fetch_logs(
+        self,
+        job_id: str,
+        stream: bool,
+        stderr: Optional[bool] = None,
+        stdout: Optional[bool] = None,
+    ) -> Iterable[str]:
+        token = self.get_auth_token()
+        if not token:
+            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)
+        )
+        workload_name = next(
+            (
+                workload["name"]
+                for workload in response.json()["workloads"]
+                if workload["id"] == job_id
+            ),
+            None,
+        )
+        if workload_name is None:
+            logger.error(f"No workload found with id {job_id}")
+            yield ""
+
+        urls = [
+            f"{self.kube_apiserver_url}/api/v1/namespaces/runai-{self.project_name}/pods/{workload_name}-worker-{i}/log?container=pytorch"
+            for i in range(self.nodes)
+        ]
+
+        if stream:
+            urls = [url + "&follow=true" for url in urls]
+
+        while self.status(job_id) != DGXCloudState.RUNNING:
+            logger.info("Waiting for job to start...")
+            time.sleep(15)
+
+        time.sleep(10)
+
+        q = queue.Queue()
+        active_urls = set(urls)
+
+        # Start threads
+        threads = [
+            threading.Thread(
+                target=self._stream_url_sync, args=(url, self._default_headers(token=token), q)
+            )
+            for url in urls
+        ]
+        for t in threads:
+            t.start()
+
+        # Yield chunks as they arrive
+        while active_urls:
+            url, item = q.get()
+            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
+
+        # Wait for threads
+        for t in threads:
+            t.join()
+
     def cancel(self, job_id: str):
         # Retrieve the authentication token for the REST calls
         token = self.get_auth_token()
@@ -385,12 +474,6 @@ def cancel(self, job_id: str):
                 response.text,
             )
 
-    @classmethod
-    def logs(cls: Type["DGXCloudExecutor"], app_id: str, fallback_path: Optional[str]):
-        logger.warning(
-            "Logs not available for DGXCloudExecutor based jobs. Please visit the cluster UI to view the logs."
-        )
-
     def cleanup(self, handle: str): ...
 
     def assign(

nemo_run/run/logs.py

@@ -30,9 +30,7 @@
 from nemo_run.core.execution.base import LogSupportedExecutor
 from nemo_run.core.frontend.console.api import CONSOLE
 from nemo_run.run.torchx_backend.runner import Runner, get_runner
-from nemo_run.run.torchx_backend.schedulers.api import (
-    REVERSE_EXECUTOR_MAPPING,
-)
+from nemo_run.run.torchx_backend.schedulers.api import REVERSE_EXECUTOR_MAPPING
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -60,6 +58,8 @@ def print_log_lines(
                 role_name,
                 replica_id,
                 regex,
+                None,
+                None,
                 should_tail=should_tail,
                 streams=streams,
             ):

nemo_run/run/torchx_backend/schedulers/dgxcloud.py

@@ -19,8 +19,9 @@
 import shutil
 import tempfile
 from dataclasses import dataclass
+from datetime import datetime
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Iterable, Optional
 
 import fiddle as fdl
 import fiddle._src.experimental.dataclasses as fdl_dc
@@ -29,15 +30,10 @@
     DescribeAppResponse,
     ListAppResponse,
     Scheduler,
+    Stream,
+    split_lines,
 )
-from torchx.specs import (
-    AppDef,
-    AppState,
-    ReplicaStatus,
-    Role,
-    RoleStatus,
-    runopts,
-)
+from torchx.specs import AppDef, AppState, ReplicaStatus, Role, RoleStatus, runopts
 
 from nemo_run.config import get_nemorun_home
 from nemo_run.core.execution.base import Executor
@@ -189,6 +185,36 @@ def describe(self, app_id: str) -> Optional[DescribeAppResponse]:
             ui_url=f"{executor.base_url}/workloads/distributed/{job_id}",
         )
 
+    def log_iter(
+        self,
+        app_id: str,
+        role_name: str,
+        k: int = 0,
+        regex: Optional[str] = None,
+        since: Optional[datetime] = None,
+        until: Optional[datetime] = None,
+        should_tail: bool = False,
+        streams: Optional[Stream] = None,
+    ) -> Iterable[str]:
+        stored_data = _get_job_dirs()
+        job_info = stored_data.get(app_id)
+        _, _, job_id = app_id.split("___")
+        executor: Optional[DGXCloudExecutor] = job_info.get("executor", None)  # type: ignore
+        if not executor:
+            return [""]
+
+        logs = executor.fetch_logs(
+            job_id=job_id,
+            stream=should_tail,
+        )  # type: ignore
+        if isinstance(logs, str):
+            if len(logs) == 0:
+                logs = []
+            else:
+                logs = split_lines(logs)
+
+        return logs
+
     def _cancel_existing(self, app_id: str) -> None:
         """
         Cancels the job by calling the DGXExecutor's cancel method.