Enhance offload cleanup handling for exception cases

auto_round/utils/offload.py

@@ -819,7 +819,11 @@ def _ensure_dir(self) -> str:
                 self._tempdir = os.path.join(base_dir, f"{self._prefix}_resume")
                 os.makedirs(self._tempdir, exist_ok=True)
             else:
-                self._tempdir = tempfile.mkdtemp(prefix=f"{self._prefix}_", dir=base_dir)
+                # A fresh temp dir is unique to this process. Before creating
+                # it, sweep leftover pid-tagged dirs from prior processes that
+                # were killed/crashed before their cleanup() could run.
+                self._sweep_stale_dirs(base_dir)
+                self._tempdir = tempfile.mkdtemp(prefix=f"{self._prefix}_{os.getpid()}_", dir=base_dir)
             logger.info(f"OffloadManager ({self._prefix}): tempdir = {self._tempdir}")
         return self._tempdir
 
@@ -901,6 +905,43 @@ def _cleanup_tempdir(self) -> None:
                     pass  # not empty or already removed
         self._tempdir = None
 
+    @staticmethod
+    def _pid_alive(pid: int) -> bool:
+        """Return *True* if a process with *pid* is still running."""
+        if pid <= 0:
+            return False
+        try:
+            os.kill(pid, 0)
+        except ProcessLookupError:
+            return False  # no such process
+        except OSError:
+            return True  # e.g. PermissionError -> exists, owned by another user
+        return True
+
+    def _sweep_stale_dirs(self, base_dir: str) -> None:
+        """Remove leftover pid-tagged temp dirs whose owning process is dead."""
+        if not os.path.isdir(base_dir):
+            return
+        tag = f"{self._prefix}_"
+        for entry in sorted(os.listdir(base_dir)):
+            full = os.path.join(base_dir, entry)
+            if not os.path.isdir(full) or not entry.startswith(tag):
+                continue
+            if entry == f"{self._prefix}_resume":
+                continue  # resume dir is intentionally persistent
+            pid_str = entry[len(tag) :].split("_", 1)[0]
+            if not pid_str.isdigit():
+                continue  # legacy dir without a pid tag -- cannot prove it is stale
+            pid = int(pid_str)
+            if self._pid_alive(pid):
+                continue
+            try:
+                shutil.rmtree(full)
+            except OSError as e:
+                logger.warning(f"OffloadManager ({self._prefix}): could not remove stale dir {full}: {e}")
+                continue
+            logger.info(f"OffloadManager ({self._prefix}): removed stale temp dir {full} (dead pid {pid})")
+
     # ------------------------------------------------------------------
     # Internal: clearing
     # ------------------------------------------------------------------

test/unit/common/utils/test_offload_helpers.py

@@ -261,3 +261,70 @@ def test_nonexistent_path_falls_through(self, tmp_path):
         ):
             result = _resolve_model_dir(str(tmp_path / "missing"))
             assert result == str(tmp_path / "missing")
+
+
+# ---------------------------------------------------------------------------
+# OffloadManager stale-dir sweep
+# ---------------------------------------------------------------------------
+class TestOffloadStaleDirSweep:
+    """Startup cleanup of leftover pid-tagged offload temp dirs."""
+
+    def _offload_base(self, tmp_path, monkeypatch) -> str:
+        workspace = str(tmp_path).lower()  # envs.AR_WORK_SPACE is lowercased
+        monkeypatch.setenv("AR_WORK_SPACE", workspace)
+        monkeypatch.delenv("AR_RESUME_DIR", raising=False)
+        base = os.path.join(workspace, "offload")
+        os.makedirs(base, exist_ok=True)
+        return base
+
+    def test_ensure_dir_is_pid_tagged_and_sweeps_stale(self, tmp_path, monkeypatch):
+        from auto_round.utils.offload import OffloadManager
+
+        base = self._offload_base(tmp_path, monkeypatch)
+        stale = os.path.join(base, "compressor_424242_deadbeef")
+        os.makedirs(stale)
+        mgr = OffloadManager(enabled=True, mode="offload", offload_dir_prefix="compressor")
+        # Only our own pid counts as alive.
+        monkeypatch.setattr(OffloadManager, "_pid_alive", staticmethod(lambda pid: pid == os.getpid()))
+        tempdir = mgr._ensure_dir()
+        try:
+            assert not os.path.exists(stale), "stale dir from a dead pid should have been swept"
+            assert os.path.basename(tempdir).startswith(f"compressor_{os.getpid()}_")
+        finally:
+            mgr._cleanup_tempdir()
+        assert not os.path.exists(tempdir)
+
+    def test_sweep_keeps_live_resume_and_legacy_dirs(self, tmp_path, monkeypatch):
+        from auto_round.utils.offload import OffloadManager
+
+        base = self._offload_base(tmp_path, monkeypatch)
+        dead = os.path.join(base, "compressor_424242_deadbeef")
+        live = os.path.join(base, f"compressor_{os.getpid()}_cafebabe")
+        resume = os.path.join(base, "compressor_resume")
+        legacy = os.path.join(base, "compressor_0gbcamv0")  # pre-pid-tag naming
+        for d in (dead, live, resume, legacy):
+            os.makedirs(d)
+        mgr = OffloadManager(enabled=True, mode="offload", offload_dir_prefix="compressor")
+        monkeypatch.setattr(OffloadManager, "_pid_alive", staticmethod(lambda pid: pid == os.getpid()))
+        mgr._sweep_stale_dirs(base)
+        assert not os.path.exists(dead)
+        assert os.path.isdir(live)
+        assert os.path.isdir(resume)
+        assert os.path.isdir(legacy)
+
+    def test_ensure_dir_resume_mode_never_sweeps(self, tmp_path, monkeypatch):
+        from auto_round.utils.offload import OffloadManager
+
+        base = self._offload_base(tmp_path, monkeypatch)
+        monkeypatch.setenv("AR_RESUME_DIR", str(tmp_path / "resume"))
+        stale = os.path.join(base, "compressor_424242_deadbeef")
+        os.makedirs(stale)
+        mgr = OffloadManager(enabled=True, mode="offload", offload_dir_prefix="compressor")
+        # Even if every pid were dead, resume mode must not delete anything.
+        monkeypatch.setattr(OffloadManager, "_pid_alive", staticmethod(lambda pid: False))
+        tempdir = mgr._ensure_dir()
+        try:
+            assert tempdir == os.path.join(base, "compressor_resume")
+            assert os.path.exists(stale)
+        finally:
+            mgr._cleanup_tempdir()