From a34e5437357e943886e438f14d8a24bf27b80087 Mon Sep 17 00:00:00 2001 From: lvliang-intel Date: Tue, 8 Sep 2026 09:48:47 +0800 Subject: [PATCH 1/2] Enhance offload cleanup handling for exception cases Signed-off-by: lvliang-intel --- auto_round/utils/offload.py | 43 +++++++++++- .../unit/common/utils/test_offload_helpers.py | 68 +++++++++++++++++++ 2 files changed, 110 insertions(+), 1 deletion(-) diff --git a/auto_round/utils/offload.py b/auto_round/utils/offload.py index 4f5ee34811..0d9b2f1803 100644 --- a/auto_round/utils/offload.py +++ b/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 # ------------------------------------------------------------------ diff --git a/test/unit/common/utils/test_offload_helpers.py b/test/unit/common/utils/test_offload_helpers.py index 7d7b6c4656..163ac78319 100644 --- a/test/unit/common/utils/test_offload_helpers.py +++ b/test/unit/common/utils/test_offload_helpers.py @@ -261,3 +261,71 @@ 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() From dd4beea265647f574e9c598fc16e7973108d5dcc Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 8 Sep 2026 01:56:06 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- test/unit/common/utils/test_offload_helpers.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/unit/common/utils/test_offload_helpers.py b/test/unit/common/utils/test_offload_helpers.py index 163ac78319..63d93831a8 100644 --- a/test/unit/common/utils/test_offload_helpers.py +++ b/test/unit/common/utils/test_offload_helpers.py @@ -267,8 +267,7 @@ def test_nonexistent_path_falls_through(self, tmp_path): # OffloadManager stale-dir sweep # --------------------------------------------------------------------------- class TestOffloadStaleDirSweep: - """Startup cleanup of leftover pid-tagged offload temp dirs. - """ + """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