[Debug] cache_dir + wandb

Author: RandallBalestriero

stable_pretraining/_config.py

@@ -34,6 +34,11 @@
     "CRITICAL",
 )
 
+# Sentinel distinguishing "argument omitted" from an explicit ``None`` in
+# :func:`set`.  ``cache_dir=None`` must raise (caching is mandatory for
+# parallel-safe run isolation), while omitting it is a no-op.
+_UNSET = object()
+
 _CLEANUP_KEYS = (
     "checkpoints",
     "logs",
@@ -251,6 +256,21 @@ def cache_dir(self, value: Optional[str]) -> None:
                 )
             if not value.strip():
                 raise ValueError("cache_dir must not be empty")
+            # Forbid relative paths: the run_dir is resolved with ``.resolve()``
+            # (manager) against the *current* CWD, which Hydra's ``job.chdir``
+            # silently changes — a relative cache_dir would then land in a
+            # different (and CWD-dependent) place per job. We expanduser first
+            # so ``~/...`` (which expands to an absolute path) is accepted and
+            # stored verbatim for portability; only genuinely relative paths
+            # like ``runs`` or ``./out`` are rejected.
+            if not os.path.isabs(os.path.expanduser(value)):
+                raise ValueError(
+                    f"cache_dir must be an absolute path, got {value!r}. "
+                    "A relative path is resolved against the current working "
+                    "directory, which Hydra's job.chdir changes — making the "
+                    "run directory non-deterministic across jobs. Pass an "
+                    "absolute path (or a '~/...' path)."
+                )
         self._cache_dir = value
 
     # -- requeue_checkpoint ----------------------------------------------------
@@ -334,7 +354,7 @@ def set(
     log_rank: Optional[Union[int, Literal["all"]]] = None,
     default_callbacks: Optional[Dict[str, bool]] = None,
     default_loggers: Optional[Dict[str, bool]] = None,
-    cache_dir: Optional[str] = None,
+    cache_dir: Union[str, object] = _UNSET,
     requeue_checkpoint: Optional[bool] = None,
     requeue_checkpoint_every_n_steps: Optional[int] = None,
     exclude_bias_norm: Optional[bool] = None,
@@ -382,8 +402,9 @@ def set(
             ensuring no path collisions across parallel sweep jobs.
             Defaults to ``~/.cache/stable-pretraining``.  Can be
             overridden via the ``SPT_CACHE_DIR`` environment variable.
-            Set to ``None`` to disable and preserve the standard
-            Lightning / Hydra directory behavior.
+            ``None`` is **not** allowed — a cache directory is mandatory
+            so every run gets a unique, parallel-safe output directory.
+            Passing ``cache_dir=None`` raises ``ValueError``.
 
             .. note::
                 SLURM ``.out`` / ``.err`` files are created by the
@@ -446,7 +467,16 @@ def set(
     if default_loggers is not None:
         cfg.default_loggers = default_loggers
 
-    if cache_dir is not None:
+    if cache_dir is not _UNSET:
+        if cache_dir is None:
+            raise ValueError(
+                "cache_dir cannot be None — a cache directory is mandatory so "
+                "that every run gets a unique, parallel-safe output directory "
+                "(without it, concurrent jobs collide on files such as "
+                "wandb_resume.json in the working directory). Pass a path, set "
+                "the SPT_CACHE_DIR environment variable, or leave it unset to "
+                "use the default (~/.cache/stable-pretraining)."
+            )
         cfg.cache_dir = cache_dir
 
     if requeue_checkpoint is not None:

stable_pretraining/callbacks/checkpoint_sklearn.py

@@ -292,15 +292,6 @@ def on_save_checkpoint(self, trainer, pl_module, checkpoint):
             sidecar.parent.mkdir(parents=True, exist_ok=True)
             sidecar.write_text(json.dumps(resume_info))
             logging.info(f"  Wrote {sidecar.resolve()}")
-            # Also write to CWD if it differs (backward compat),
-            # but skip when cache_dir is active to avoid polluting CWD.
-            from stable_pretraining._config import get_config
-
-            if get_config().cache_dir is None:
-                cwd_sidecar = Path(_WANDB_RESUME_FILENAME)
-                if cwd_sidecar.resolve() != sidecar.resolve():
-                    cwd_sidecar.write_text(json.dumps(resume_info))
-                    logging.info(f"  Wrote {cwd_sidecar.resolve()} (compat)")
 
     def on_load_checkpoint(self, trainer, pl_module, checkpoint):
         if "wandb" not in checkpoint:

stable_pretraining/callbacks/checkpoint_swanlab.py

@@ -80,14 +80,6 @@ def on_save_checkpoint(
             sidecar.write_text(json.dumps(resume_info))
             logging.info(f"  Wrote {sidecar.resolve()}")
 
-            from stable_pretraining._config import get_config
-
-            if get_config().cache_dir is None:
-                cwd_sidecar = Path(_SWANLAB_RESUME_FILENAME)
-                if cwd_sidecar.resolve() != sidecar.resolve():
-                    cwd_sidecar.write_text(json.dumps(resume_info))
-                    logging.info(f"  Wrote {cwd_sidecar.resolve()} (compat)")
-
     def on_load_checkpoint(
         self,
         trainer: Trainer,

stable_pretraining/callbacks/checkpoint_trackio.py

@@ -79,14 +79,6 @@ def on_save_checkpoint(
             sidecar.write_text(json.dumps(resume_info))
             logging.info(f"  Wrote {sidecar.resolve()}")
 
-            from stable_pretraining._config import get_config
-
-            if get_config().cache_dir is None:
-                cwd_sidecar = Path(_TRACKIO_RESUME_FILENAME)
-                if cwd_sidecar.resolve() != sidecar.resolve():
-                    cwd_sidecar.write_text(json.dumps(resume_info))
-                    logging.info(f"  Wrote {cwd_sidecar.resolve()} (compat)")
-
     def on_load_checkpoint(
         self,
         trainer: Trainer,

stable_pretraining/manager.py

@@ -1636,6 +1636,22 @@ def init_and_sync_wandb(self):
         if wandb_logger is None:
             return
         log_header("Wandb")
+
+        # Anchor wandb's local files (the ``wandb/`` run tree) to the unique
+        # run_dir *before* the experiment is created — accessing
+        # ``.experiment`` below triggers ``wandb.init(**_wandb_init)``.
+        # Lightning's WandbLogger defaults ``dir``/``save_dir`` to CWD, which
+        # is NOT covered by ``default_root_dir`` (unlike CSV/TensorBoard
+        # loggers). Without this, parallel jobs sharing a CWD all dump their
+        # ``wandb/`` trees into the same place instead of each job's run_dir.
+        run_dir = getattr(self, "_run_dir", None)
+        already_running = WANDB_AVAILABLE and wandb.run is not None
+        if run_dir is not None and not already_running:
+            run_dir_str = str(run_dir)
+            wandb_logger._wandb_init["dir"] = run_dir_str
+            wandb_logger._save_dir = run_dir_str
+            logging.info(f"  Anchored wandb dir to {run_dir_str}")
+
         exp = wandb_logger.experiment
 
         if exp.offline:

stable_pretraining/tests/unit/test_cache_dir.py

@@ -68,12 +68,24 @@ def test_set_via_property(self, tmp_path):
         cfg.cache_dir = str(tmp_path)
         assert cfg.cache_dir == str(tmp_path)
 
-    def test_set_to_none(self, tmp_path):
+    def test_set_to_none_via_property(self, tmp_path):
+        # The property setter still permits None — the JAX backend uses it to
+        # disable checkpointing (see jax/manager.py). Only the public spt.set()
+        # API prohibits it.
         spt_set(cache_dir=str(tmp_path))
         cfg = get_config()
         cfg.cache_dir = None
         assert cfg.cache_dir is None
 
+    def test_spt_set_none_raises(self, tmp_path):
+        # Public API must reject an explicit None: caching is mandatory so each
+        # run gets a unique, parallel-safe output dir (no CWD collisions).
+        spt_set(cache_dir=str(tmp_path))
+        with pytest.raises(ValueError, match="cache_dir cannot be None"):
+            spt_set(cache_dir=None)
+        # The failed call must not have mutated the existing value.
+        assert get_config().cache_dir == str(tmp_path)
+
     def test_rejects_empty_string(self):
         with pytest.raises(ValueError, match="must not be empty"):
             spt_set(cache_dir="")
@@ -86,6 +98,18 @@ def test_rejects_non_string(self):
         with pytest.raises(TypeError, match="must be a str"):
             get_config().cache_dir = 123
 
+    def test_rejects_relative_path(self):
+        with pytest.raises(ValueError, match="must be an absolute path"):
+            spt_set(cache_dir="runs")
+        with pytest.raises(ValueError, match="must be an absolute path"):
+            get_config().cache_dir = "./out"
+
+    def test_accepts_tilde_path(self):
+        # '~/...' expands to an absolute path, so it is accepted and stored
+        # verbatim (expansion happens later in _resolve_run_dir).
+        spt_set(cache_dir="~/spt_cache_test")
+        assert get_config().cache_dir == "~/spt_cache_test"
+
     def test_reset_restores_default_cache_dir(self, tmp_path, monkeypatch):
         monkeypatch.delenv("SPT_CACHE_DIR", raising=False)
         spt_set(cache_dir=str(tmp_path))