Removing duplicate logs (#512)

Signed-off-by: Igor Gitman <igitman@nvidia.com>

Author: Kipok

nemo_skills/pipeline/utils/cluster.py

@@ -35,6 +35,10 @@
 
 LOG = logging.getLogger(get_logger_name(__file__))
 
+# Add a module-level set to track which environment variables have been logged
+_logged_required_env_vars = set()
+_logged_optional_env_vars = set()
+
 
 def get_timeout(cluster_config, partition):
     if 'timeouts' not in cluster_config:
@@ -67,20 +71,30 @@ def get_env_variables(cluster_config):
     Returns:
         dict: dictionary of environment
     """
+    global _logged_required_env_vars, _logged_optional_env_vars
+
     env_vars = {}
     # Check for user requested env variables
     required_env_vars = cluster_config.get("required_env_vars", [])
     for env_var in required_env_vars:
+        env_var_name = env_var.split('=')[0].strip() if "=" in env_var else env_var
+
         if "=" in env_var:
             if env_var.count("=") == 1:
-                env_var, value = env_var.split("=")
+                env_var_name, value = env_var.split("=")
+                env_var_name = env_var_name.strip()
+                value = value.strip()
             else:
                 raise ValueError(f"Invalid required environment variable format: {env_var}")
-            env_vars[env_var.strip()] = value.strip()
-            logging.info(f"Adding required environment variable {env_var}")
+            env_vars[env_var_name] = value
+            if env_var_name not in _logged_required_env_vars:
+                LOG.info(f"Adding required environment variable {env_var_name} from config")
+                _logged_required_env_vars.add(env_var_name)
         elif env_var in os.environ:
-            logging.info(f"Adding required environment variable {env_var} from environment")
             env_vars[env_var] = os.environ[env_var]
+            if env_var not in _logged_required_env_vars:
+                LOG.info(f"Adding required environment variable {env_var} from environment")
+                _logged_required_env_vars.add(env_var)
         else:
             raise ValueError(f"Required environment variable {env_var} not found.")
 
@@ -94,21 +108,33 @@ def get_env_variables(cluster_config):
     # Add optional env variables
     optional_env_vars = cluster_config.get("env_vars", [])
     for env_var in optional_env_vars + always_optional_env_vars:
+        env_var_name = env_var.split('=')[0].strip() if "=" in env_var else env_var
+
         if "=" in env_var:
             if env_var.count("=") == 1:
-                env_var, value = env_var.split("=")
+                env_var_name, value = env_var.split("=")
+                env_var_name = env_var_name.strip()
+                value = value.strip()
             else:
                 raise ValueError(f"Invalid optional environment variable format: {env_var}")
-            env_vars[env_var.strip()] = value.strip()
-            logging.info(f"Adding optional environment variable {env_var}")
+            env_vars[env_var_name] = value
+            if env_var_name not in _logged_optional_env_vars:
+                LOG.info(f"Adding optional environment variable {env_var_name} from config")
+                _logged_optional_env_vars.add(env_var_name)
         elif env_var in os.environ:
-            logging.info(f"Adding optional environment variable {env_var} from environment")
             env_vars[env_var] = os.environ[env_var]
+            if env_var not in _logged_optional_env_vars:
+                LOG.info(f"Adding optional environment variable {env_var} from environment")
+                _logged_optional_env_vars.add(env_var)
         elif env_var in default_factories:
             env_vars[env_var] = default_factories[env_var]()
-            logging.info(f"Adding optional environment variable {env_var} from environment")
+            if env_var not in _logged_optional_env_vars:
+                LOG.info(f"Adding optional environment variable {env_var} from environment")
+                _logged_optional_env_vars.add(env_var)
         else:
-            logging.info(f"Optional environment variable {env_var} not found in user environment; skipping.")
+            if env_var not in _logged_optional_env_vars:
+                LOG.info(f"Optional environment variable {env_var} not found in user environment; skipping.")
+                _logged_optional_env_vars.add(env_var)
 
     return env_vars
 

nemo_skills/pipeline/utils/exp.py

@@ -29,7 +29,7 @@
 from nemo_skills.pipeline.utils.mounts import get_mounts_from_config, get_unmounted_path
 from nemo_skills.pipeline.utils.packager import get_packager
 from nemo_skills.pipeline.utils.server import get_free_port, get_server_command
-from nemo_skills.utils import get_logger_name
+from nemo_skills.utils import get_logger_name, remove_handlers
 
 LOG = logging.getLogger(get_logger_name(__file__))
 
@@ -497,6 +497,8 @@ def run_exp(exp, cluster_config, sequential=None):
 
 
 def get_exp(expname, cluster_config):
+    # nemo-run redefines the handlers, so removing ours to avoid duplicate logs
+    remove_handlers()
     if cluster_config['executor'] == 'slurm':
         return run.Experiment(expname)
     # hiding all nemo-run logs otherwise as they are not useful locally

nemo_skills/pipeline/utils/packager.py

@@ -127,7 +127,7 @@ def get_packager(extra_package_dirs: tuple[str] | None = None):
     if repo_path:
         # Do we have nemo_skills package in this repo? If no, we need to pick it up from installed location
         if not (Path(repo_path) / 'nemo_skills').is_dir():
-            logging.info(
+            LOG.info(
                 "Not running from NeMo-Skills repo, trying to upload installed package. "
                 "Make sure there are no extra files in %s",
                 str(nemo_skills_dir / '*'),
@@ -147,7 +147,7 @@ def get_packager(extra_package_dirs: tuple[str] | None = None):
             check_uncommitted_changes=check_uncommited_changes,
         )
     else:
-        logging.info(
+        LOG.info(
             "Not running from a git repo, trying to upload installed package. Make sure there are no extra files in %s",
             str(nemo_skills_dir / '*'),
         )

nemo_skills/utils.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import glob
 import inspect
 import io
 import logging
@@ -86,8 +85,15 @@ def setup_logging(disable_hydra_logs: bool = True, log_level: int = logging.INFO
     if use_rich:
         handler = RichHandler(
             rich_tracebacks=True,
-            show_path=False,
-            show_time=False,
+            show_time=True,
+            show_level=True,
+            show_path=True,
+        )
+        handler.setFormatter(
+            logging.Formatter(
+                "%(message)s",
+                datefmt="[%X]",
+            )
         )
         for hdlr in logger.handlers[:]:
             logger.removeHandler(hdlr)
@@ -107,6 +113,13 @@ def setup_logging(disable_hydra_logs: bool = True, log_level: int = logging.INFO
     return logger
 
 
+def remove_handlers():
+    """Can be used to remove all nemo-skills log handlers."""
+    logger = logging.getLogger('nemo_skills')
+    for handler in logger.handlers[:]:
+        logger.removeHandler(handler)
+
+
 def get_logger_name(file):
     return 'nemo_skills' + file.split('nemo_skills')[1].replace('/', '.').replace('.py', '')
 
@@ -288,7 +301,7 @@ def get_help_message(dataclass_obj, help_message="", **kwargs):
     docstring = re.sub(r'{([^}]+(?=\s)[^}]*)}', r'{{\1}}', docstring)
     # Might need to add some other edge-case handling
     # here, so that formatting does not complain
-    docstring = dataclass_obj.__doc__ + "\n\n" +  docstring.format(**kwargs)
+    docstring = dataclass_obj.__doc__ + "\n\n" + docstring.format(**kwargs)
 
     full_help = f"{heading}\n{'-' * 75}\n{docstring}"
     if help_message: