|
14 | 14 | # limitations under the License. |
15 | 15 |
|
16 | 16 | import base64 |
| 17 | +import glob |
17 | 18 | import json |
18 | 19 | import logging |
19 | 20 | import os |
20 | | -import queue |
21 | 21 | import subprocess |
22 | 22 | import tempfile |
23 | | -import threading |
24 | 23 | import time |
25 | 24 | from dataclasses import dataclass, field |
26 | 25 | from enum import Enum |
@@ -323,7 +322,8 @@ def launch(self, name: str, cmd: list[str]) -> tuple[str, str]: |
323 | 322 | launch_script = f""" |
324 | 323 | ln -s {self.pvc_job_dir}/ /nemo_run |
325 | 324 | cd /nemo_run/code |
326 | | -{" ".join(cmd)} |
| 325 | +mkdir -p {self.pvc_job_dir}/logs |
| 326 | +{" ".join(cmd)} 2>&1 | tee -a {self.pvc_job_dir}/log_$HOSTNAME.out {self.pvc_job_dir}/log-allranks_0.out |
327 | 327 | """ |
328 | 328 | with open(os.path.join(self.job_dir, "launch_script.sh"), "w+") as f: |
329 | 329 | f.write(launch_script) |
@@ -371,91 +371,69 @@ def status(self, job_id: str) -> Optional[DGXCloudState]: |
371 | 371 | r_json = response.json() |
372 | 372 | return DGXCloudState(r_json["phase"]) |
373 | 373 |
|
374 | | - def _stream_url_sync(self, url: str, headers: dict, q: queue.Queue): |
375 | | - """Stream a single URL using requests and put chunks into the queue""" |
376 | | - try: |
377 | | - with requests.get(url, stream=True, headers=headers, verify=False) as response: |
378 | | - for line in response.iter_lines(decode_unicode=True): |
379 | | - q.put((url, f"{line}\n")) |
380 | | - except Exception as e: |
381 | | - logger.error(f"Error streaming URL {url}: {e}") |
382 | | - |
383 | | - finally: |
384 | | - q.put((url, None)) |
385 | | - |
386 | 374 | def fetch_logs( |
387 | 375 | self, |
388 | 376 | job_id: str, |
389 | 377 | stream: bool, |
390 | 378 | stderr: Optional[bool] = None, |
391 | 379 | stdout: Optional[bool] = None, |
392 | 380 | ) -> Iterable[str]: |
393 | | - token = self.get_auth_token() |
394 | | - if not token: |
395 | | - logger.error("Failed to retrieve auth token for fetch logs request.") |
396 | | - yield "" |
397 | | - |
398 | | - response = requests.get( |
399 | | - f"{self.base_url}/workloads", headers=self._default_headers(token=token) |
400 | | - ) |
401 | | - workload_name = next( |
402 | | - ( |
403 | | - workload["name"] |
404 | | - for workload in response.json()["workloads"] |
405 | | - if workload["id"] == job_id |
406 | | - ), |
407 | | - None, |
408 | | - ) |
409 | | - if workload_name is None: |
410 | | - logger.error(f"No workload found with id {job_id}") |
411 | | - yield "" |
412 | | - |
413 | | - urls = [ |
414 | | - f"{self.kube_apiserver_url}/api/v1/namespaces/runai-{self.project_name}/pods/{workload_name}-worker-{i}/log?container=pytorch" |
415 | | - for i in range(self.nodes) |
416 | | - ] |
417 | | - |
418 | | - if stream: |
419 | | - urls = [url + "&follow=true" for url in urls] |
420 | | - |
421 | 381 | while self.status(job_id) != DGXCloudState.RUNNING: |
422 | 382 | logger.info("Waiting for job to start...") |
423 | 383 | time.sleep(15) |
424 | 384 |
|
425 | | - time.sleep(10) |
| 385 | + cmd = ["tail"] |
| 386 | + |
| 387 | + if stream: |
| 388 | + cmd.append("-f") |
426 | 389 |
|
427 | | - q = queue.Queue() |
428 | | - active_urls = set(urls) |
| 390 | + # setting linked PVC job directory |
| 391 | + nemo_run_home = get_nemorun_home() |
| 392 | + job_subdir = self.job_dir[len(nemo_run_home) + 1 :] # +1 to remove the initial backslash |
| 393 | + self.pvc_job_dir = os.path.join(self.pvc_nemo_run_dir, job_subdir) |
429 | 394 |
|
430 | | - # Start threads |
431 | | - threads = [ |
432 | | - threading.Thread( |
433 | | - target=self._stream_url_sync, args=(url, self._default_headers(token=token), q) |
| 395 | + files = [] |
| 396 | + while len(files) < self.nodes: |
| 397 | + files = list(glob.glob(f"{self.pvc_job_dir}/log_*.out")) |
| 398 | + files = [f for f in files if "log-allranks_0" not in f] |
| 399 | + logger.info( |
| 400 | + f"Waiting for {self.nodes + 1 - len(files)} log files to be created in {self.pvc_job_dir}..." |
434 | 401 | ) |
435 | | - for url in urls |
436 | | - ] |
437 | | - for t in threads: |
438 | | - t.start() |
439 | | - |
440 | | - # Yield chunks as they arrive |
441 | | - while active_urls: |
442 | | - url, item = q.get() |
443 | | - if item is None or self.status(job_id) in [ |
444 | | - DGXCloudState.DELETING, |
445 | | - DGXCloudState.STOPPED, |
446 | | - DGXCloudState.STOPPING, |
447 | | - DGXCloudState.DEGRADED, |
448 | | - DGXCloudState.FAILED, |
449 | | - DGXCloudState.COMPLETED, |
450 | | - DGXCloudState.TERMINATING, |
451 | | - ]: |
452 | | - active_urls.discard(url) |
453 | | - else: |
454 | | - yield item |
455 | | - |
456 | | - # Wait for threads |
457 | | - for t in threads: |
458 | | - t.join() |
| 402 | + time.sleep(3) |
| 403 | + |
| 404 | + cmd.extend(files) |
| 405 | + |
| 406 | + logger.info(f"Attempting to stream logs with command: {cmd}") |
| 407 | + |
| 408 | + proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, text=True, bufsize=1) |
| 409 | + |
| 410 | + if stream: |
| 411 | + while True: |
| 412 | + try: |
| 413 | + for line in iter(proc.stdout.readline, ""): |
| 414 | + if ( |
| 415 | + line |
| 416 | + and not line.rstrip("\n").endswith(".out <==") |
| 417 | + and line.rstrip("\n") != "" |
| 418 | + ): |
| 419 | + yield f"{line}" |
| 420 | + if proc.poll() is not None: |
| 421 | + break |
| 422 | + except Exception as e: |
| 423 | + logger.error(f"Error streaming logs: {e}") |
| 424 | + time.sleep(3) |
| 425 | + continue |
| 426 | + |
| 427 | + else: |
| 428 | + try: |
| 429 | + for line in iter(proc.stdout.readline, ""): |
| 430 | + if line: |
| 431 | + yield line.rstrip("\n") |
| 432 | + if proc.poll() is not None: |
| 433 | + break |
| 434 | + finally: |
| 435 | + proc.terminate() |
| 436 | + proc.wait(timeout=2) |
459 | 437 |
|
460 | 438 | def cancel(self, job_id: str): |
461 | 439 | # Retrieve the authentication token for the REST calls |
|
0 commit comments