done testing

Author: BalaBalaYi

dlrover/python/common/constants.py

@@ -14,6 +14,8 @@
 
 class BasicClass(object):
     LOG_LEVEL_ENV = "DLROVER_LOG_LEVEL"
+    LOG_ROOT_DIR_ENV = "DLROVER_LOG_ROOT_DIR"
+    LOG_AGENT_DIR_ENV = "DLROVER_LOG_AGENT_DIR"
 
 
 class PriorityClass(object):

dlrover/python/common/log.py

@@ -15,6 +15,7 @@
 import os
 import sys
 import typing
+from logging.handlers import RotatingFileHandler
 
 from dlrover.python.common.constants import BasicClass
 
@@ -31,7 +32,7 @@
 _ch = logging.StreamHandler(stream=sys.stderr)
 _ch.setFormatter(_DEFAULT_FORMATTER)
 
-_DEFAULT_HANDLERS = [_ch]
+_DEFAULT_HANDLERS: typing.List[logging.Handler] = [_ch]
 
 _LOGGER_CACHE: typing.Dict[str, logging.Logger] = {}
 
@@ -43,14 +44,55 @@ def get_log_level():
     return log_level
 
 
-def get_logger(name, handlers=None, update=False):
+def get_base_log_dir():
+    log_dir = os.getenv(BasicClass.LOG_ROOT_DIR_ENV, "")
+    return log_dir
+
+
+def get_agent_log_dir():
+    log_dir = os.getenv(BasicClass.LOG_AGENT_DIR_ENV, "")
+    return log_dir
+
+
+def get_base_log_file():
+    log_dir = get_base_log_dir()
+    log_file = ""
+    if log_dir:
+        log_file = os.path.join(log_dir, "dlrover.log")
+        os.makedirs(log_dir, exist_ok=True)
+    return log_file
+
+
+def get_logger(
+    name,
+    handlers: typing.Optional[typing.List[logging.Handler]] = None,
+    update=False,
+):
     __setup_extra_logger()
 
     if name in _LOGGER_CACHE and not update:
         return _LOGGER_CACHE[name]
     logger = logging.getLogger(name)
     logger.setLevel(get_log_level())
-    logger.handlers = handlers or _DEFAULT_HANDLERS
+
+    if handlers is None:
+        base_log_file = get_base_log_file()
+        if base_log_file:
+            file_handler = RotatingFileHandler(
+                base_log_file,
+                maxBytes=200 * 1024 * 1024,  # 200MB
+                backupCount=3,
+            )
+            file_handler.setFormatter(_DEFAULT_FORMATTER)
+            handlers = [file_handler] + _DEFAULT_HANDLERS
+        else:
+            handlers = _DEFAULT_HANDLERS
+    elif len(handlers) == 0:
+        handlers = _DEFAULT_HANDLERS
+    else:
+        handlers.extend(_DEFAULT_HANDLERS)
+
+    logger.handlers = list(handlers)
     logger.propagate = False
     return logger
 

dlrover/python/elastic_agent/torch/training.py

@@ -19,7 +19,6 @@
 import socket
 import subprocess
 import sys
-import tempfile
 import time
 import traceback
 import uuid
@@ -180,6 +179,79 @@ class StopWorkerTimeoutError(RuntimeError):
     pass
 
 
+class LogConfig:
+    _log_dir: Optional[str] = None
+    _redirects: Union[Std, Dict[int, Std]] = Std.NONE
+    _tee: Union[Std, Dict[int, Std]] = Std.NONE
+
+    @classmethod
+    def _parse_std_value(cls, val) -> Std:
+        if val is None:
+            return Std.NONE
+
+        if isinstance(val, str):
+            try:
+                return Std.from_str(val)
+            except ValueError:
+                return Std.NONE
+        elif isinstance(val, Std):
+            return val
+
+        return Std.NONE
+
+    def setup(self, log_dir, redirects=None, tee=None):
+        if not log_dir:
+            return
+
+        self._log_dir = log_dir
+
+        redirects = self._parse_std_value(redirects)
+        tee = self._parse_std_value(tee)
+
+        if redirects == Std.NONE and tee == Std.NONE:
+            # override default when log dir is specified
+            self._redirects = Std.ALL
+            self._tee = Std.ALL
+        else:
+            self._redirects = redirects
+            self._tee = tee
+
+        logger.debug(
+            f"Log config: {self._log_dir}-{self._redirects}-{self._tee}"
+        )
+
+    @property
+    def log_dir(self) -> Optional[str]:
+        return self._log_dir
+
+    @property
+    def redirects(self) -> Union[Std, Dict[int, Std]]:
+        return self._redirects
+
+    @property
+    def tee(self) -> Union[Std, Dict[int, Std]]:
+        return self._tee
+
+    @property
+    def logs_specs(self):
+        if version_less_than_230():
+            return {
+                "log_dir": self.log_dir,
+                "redirects": self.redirects,
+                "tee": self.tee,
+            }
+        else:
+            from torch.distributed.elastic.multiprocessing import (
+                DefaultLogsSpecs,
+            )
+
+            log_specs = DefaultLogsSpecs(
+                log_dir=self.log_dir, redirects=self.redirects, tee=self.tee
+            )
+            log_specs._run_log_dir = self.log_dir
+            return log_specs
+
+
 @dataclass
 class ElasticLaunchConfig(LaunchConfig):
     """
@@ -216,15 +288,33 @@ class ElasticLaunchConfig(LaunchConfig):
     exclude_straggler: bool = False
     save_at_breakpoint: bool = False
     accelerator: str = ""
-    log_dir: Optional[str] = None  # Keep Compatibility with PyTorch>=2.3.0
-    redirects: Union[Std, Dict[int, Std]] = Std.NONE
-    tee: Union[Std, Dict[int, Std]] = Std.NONE
+    log_config: LogConfig = LogConfig()
     training_log_file: str = ""
     failure_node_errors: str = ""
     numa_affinity: bool = False
     membind_policy: str = "none"
     ucp_device_type: str = "cpu"
 
+    def get_log_dir(self):
+        return self.log_config.log_dir
+
+    def get_log_tee(self):
+        return self.log_config.tee
+
+    def get_log_redirects(self):
+        return self.log_config.redirects
+
+    def get_log_specs(self):
+        return self.log_config.logs_specs
+
+    def setup_log(self, log_dir, redirects=None, tee=None):
+        if log_dir:
+            logger.info(f"Initiate specified log directory: {log_dir}.")
+
+            self.log_config.setup(log_dir, redirects=redirects, tee=tee)
+        else:
+            logger.info("No specified log directory is configured.")
+
     def set_node_unit(self, node_unit):
         """Set the number unit of nodes."""
         self.node_unit = node_unit
@@ -554,7 +644,6 @@ def __init__(
         spec: WorkerSpec,
         start_method="spawn",
         exit_barrier_timeout: float = 300,
-        log_dir: Optional[str] = None,
         training_log_file: str = "",
         failure_node_errors: str = "",
         with_diagnostician: bool = True,
@@ -564,18 +653,26 @@ def __init__(
                 spec=spec,
                 exit_barrier_timeout=exit_barrier_timeout,
             )
+            # compatible
+            # https://github.com/pytorch/pytorch/blob/39901f229520a5256505ec24782f716ee7ddc843/torch/distributed/elastic/agent/server/local_elastic_agent.py#L148C9-L148C22
+            self._log_dir = config.get_log_dir()
         else:
+            logger.info(
+                "Setup logging configuration for torch version>=230 with "
+                f"log_dir: {config.get_log_dir()}, "
+                f"redirections: {config.get_log_redirects()}, "
+                f"tee: {config.get_log_tee()}, log_specs: {config.get_log_specs().__dict__}"
+            )
             super().__init__(
                 spec=spec,
-                logs_specs=config.logs_specs,
+                logs_specs=config.get_log_specs(),
                 exit_barrier_timeout=exit_barrier_timeout,
             )
         self._node_rank = node_rank
         self._config = config
         self._entrypoint = entrypoint
         self._start_method = start_method
         self._pcontext: Optional[PContext] = None
-        self._log_dir = log_dir or tempfile.mkdtemp(prefix="torchelastic_")
         self._worker_watchdog: Optional[timer.FileTimerServer] = None
         self._restart_count = 0
         self._remaining_failovers = self._remaining_restarts
@@ -1682,7 +1779,7 @@ def launch_agent(
         f"  rdzv_configs     : {config.rdzv_configs}\n"
         f"  max_restarts     : {config.max_restarts}\n"
         f"  monitor_interval : {config.monitor_interval}\n"
-        f"  log_dir          : {config.log_dir}\n"
+        f"  log_dir          : {config.get_log_dir()}\n"
         f"  metrics_cfg      : {config.metrics_cfg}\n"
         f"  training_log     : {config.training_log_file}\n"
         f"  failure_errors   : {config.failure_node_errors}\n"
@@ -1712,7 +1809,6 @@ def launch_agent(
         entrypoint=entrypoint,
         spec=spec,
         start_method=config.start_method,
-        log_dir=config.log_dir,
         training_log_file=config.training_log_file,
         failure_node_errors=config.failure_node_errors,
         exit_barrier_timeout=900,
@@ -1826,10 +1922,16 @@ def _create_worker_spec(
         master_addr=master_addr,
     )
 
+    # for torch < 230, the tee and redirects config for log is located in spec
     if version_less_than_230():
-        spec.redirects = config.redirects
-        spec.tee = config.tee
-
+        spec.redirects = config.get_log_redirects()
+        spec.tee = config.get_log_tee()
+        logger.info(
+            "Setup logging configuration for torch version<230 with "
+            f"log_dir: {config.get_log_dir()}, "
+            f"redirections: {config.get_log_redirects()}, "
+            f"tee: {config.get_log_tee()}"
+        )
     return spec
 
 
@@ -1858,7 +1960,6 @@ def __init__(
         spec: WorkerSpec,
         start_method="spawn",
         exit_barrier_timeout: float = 300,
-        log_dir: Optional[str] = None,
         check_round=1,
     ):
         super().__init__(
@@ -1868,10 +1969,8 @@ def __init__(
             spec,
             start_method,
             exit_barrier_timeout,
-            log_dir,
             with_diagnostician=False,
         )
-        self._log_dir = log_dir or tempfile.mkdtemp(prefix="node_check_")
         self._check_round = check_round
         self._config: ElasticLaunchConfig = config
 
@@ -2063,7 +2162,7 @@ def _create_check_agent(
         f"  rdzv_configs     : {config.rdzv_configs}\n"
         f"  max_restarts     : {config.max_restarts}\n"
         f"  monitor_interval : {config.monitor_interval}\n"
-        f"  log_dir          : {config.log_dir}\n"
+        f"  log_dir          : {config.get_log_dir()}\n"
         f"  metrics_cfg      : {config.metrics_cfg}\n"
     )