Add log printout to announce modifier logging (#1432) (#1448)

Author: KSGulin

src/sparseml/pytorch/utils/logger.py

@@ -48,6 +48,9 @@
 from sparseml.utils import ALL_TOKEN, create_dirs
 
 
+_LOGGER = logging.getLogger(__name__)
+
+
 __all__ = [
     "BaseLogger",
     "LambdaLogger",
@@ -332,11 +335,21 @@ def __init__(
             self._logger = logger
         else:
             self._logger = logging.getLogger(__name__)
+
+            base_log_path = (
+                os.environ.get("NM_TEST_LOG_DIR")
+                if os.environ.get("NM_TEST_MODE")
+                else "sparse_logs"
+            )
             now = datetime.now()
             dt_string = now.strftime("%d-%m-%Y_%H.%M.%S")
-            os.makedirs("sparse_logs", exist_ok=True)
+            log_path = os.path.join(base_log_path, f"{dt_string}.log")
+            os.makedirs(base_log_path, exist_ok=True)
+
+            _LOGGER.info(f"Logging all SparseML modifier-level logs to {log_path}")
+
             handler = logging.FileHandler(
-                os.path.join("sparse_logs", f"{dt_string}.log"),
+                log_path,
                 delay=True,
             )
             self._logger.addHandler(handler)

tests/sparseml/conftest.py

@@ -25,7 +25,7 @@
 
 
 os.environ["NM_TEST_MODE"] = "True"
-os.environ["NM_TEST_LOG_DIR"] = "temp_test_logs"
+os.environ["NM_TEST_LOG_DIR"] = "nm_temp_test_logs"
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -35,14 +35,8 @@ def check_for_created_files():
     if wandb:
         wandb.finish()
     log_dir = os.environ.get("NM_TEST_LOG_DIR")
-    log_dir_tensorboard = os.path.join(log_dir, "tensorboard")
-    log_dir_wandb = os.path.join(log_dir, "wandb")
-    if os.path.isdir(log_dir_tensorboard):
-        shutil.rmtree(log_dir_tensorboard)
-    if os.path.isdir(log_dir_wandb):
-        shutil.rmtree(log_dir_wandb)
-    if os.path.isdir(log_dir) and len(os.listdir(log_dir)) == 0:
-        os.rmdir(log_dir)
+    if os.path.isdir(log_dir):
+        shutil.rmtree(log_dir)
     end_file_count = sum(len(files) for _, _, files in os.walk(r"."))
     assert (
         start_file_count >= end_file_count

utils/conftest.py

@@ -20,14 +20,14 @@
 # Ignore submodules
 collect_ignore_glob = ["tensorflow_v1-onnx/*"]
 
-
-FAILURE_LOG = "test_logs/failures.log"
+LOG_DIR = os.environ["NM_TEST_LOG_DIR"] or "test_logs"
+FAILURE_LOG = os.path.join(LOG_DIR, "failures.log")
 
 
 def pytest_configure(config):
     if os.path.exists(FAILURE_LOG):
         os.remove(FAILURE_LOG)
-    os.makedirs("test_logs", exist_ok=True)
+    os.makedirs(LOG_DIR, exist_ok=True)
 
 
 def write_to_failure_log(node_id, long_repr):