diff --git a/tests/challenge/test_store_notify_dir.py b/tests/challenge/test_store_notify_dir.py new file mode 100644 index 0000000..3bc88f0 --- /dev/null +++ b/tests/challenge/test_store_notify_dir.py @@ -0,0 +1,114 @@ +"""Shared v2 notification directories: safe open/close, and the test session's leak cleanup.""" + +from __future__ import annotations + +import shutil +import socket +import tempfile +import threading +import time +from collections.abc import Iterator +from pathlib import Path + +import pytest + +from hypertrain.challenge.store import ChallengeStore +from hypertrain.data.store import LocalFSStore +from notify_cleanup import remove_new_empty_notify_dirs + +from .conftest import PARAMS + + +@pytest.fixture +def root(monkeypatch: pytest.MonkeyPatch) -> Iterator[Path]: + # A private temp root, short enough for AF_UNIX paths (tmp_path is not), so cleanup is + # observed without touching other sessions' directories in the real one. + private = Path(tempfile.mkdtemp(prefix="htnd-")) + monkeypatch.setattr(tempfile, "tempdir", str(private)) + yield private + shutil.rmtree(private) + + +def _store(state: Path, objects: Path) -> ChallengeStore: + return ChallengeStore(state, PARAMS, None, None, lambda _: None, LocalFSStore(objects)) + + +def _started(state: Path, objects: Path) -> ChallengeStore: + store = _store(state, objects) + store._start_notifications_v2() + return store + + +def _directory(store: ChallengeStore) -> Path: + assert store._ipc_path_v2 is not None + return store._ipc_path_v2.parent + + +def _snapshot(root: Path) -> set[str]: + return {path.name for path in root.glob("ht-notify-*")} + + +def test_session_cleanup_removes_directory_left_by_closed_store(root: Path, tmp_path: Path) -> None: + before = _snapshot(root) + store = _started(tmp_path / "state", tmp_path / "objects") + directory = _directory(store) + assert directory.parent == root + store.close_v2_notifications() + store._db.close() + assert directory.is_dir() # production keeps the empty directory + remove_new_empty_notify_dirs(root, before) + assert not directory.exists() + + +def test_session_cleanup_keeps_directory_with_live_socket(root: Path, tmp_path: Path) -> None: + before = _snapshot(root) + first, second = (_started(tmp_path / "state", tmp_path / f"objects-{i}") for i in range(2)) + first.close_v2_notifications() + remove_new_empty_notify_dirs(root, before) + assert second._ipc_path_v2 is not None and second._ipc_path_v2.is_socket() + second.close_v2_notifications() + for store in (first, second): + store._db.close() + + +def test_session_cleanup_keeps_directory_that_predates_snapshot(root: Path) -> None: + existing = root / "ht-notify-preexisting" + existing.mkdir(mode=0o700) + remove_new_empty_notify_dirs(root, _snapshot(root)) + assert existing.is_dir() + + +def test_shared_directory_survives_concurrent_start_close( + root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # Regression: removing the shared directory on close raced a peer's mkdir/bind + # (FileNotFoundError at bind, FileExistsError at mkdir). + class SlowBind(socket.socket): + def bind(self, address: object) -> None: + time.sleep(0.001) # widen the mkdir-to-bind window a peer's close could hit + super().bind(address) # type: ignore[arg-type] + + monkeypatch.setattr(socket, "socket", SlowBind) + threads, cycles = 8, 100 + stores = [_store(tmp_path / "state", tmp_path / f"objects-{i}") for i in range(threads)] + gate = threading.Barrier(threads) + errors: list[OSError] = [] + + def cycle(store: ChallengeStore) -> None: + gate.wait() + for _ in range(cycles): + try: + store._start_notifications_v2() + store.close_v2_notifications() + except OSError as error: # keep cycling: one lost race must not end the probe + errors.append(error) + + workers = [threading.Thread(target=cycle, args=(store,)) for store in stores] + for worker in workers: + worker.start() + for worker in workers: + worker.join() + for store in stores: + store.close_v2_notifications() + store._db.close() + assert errors == [] diff --git a/tests/conftest.py b/tests/conftest.py index f5da5a2..8968ce5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,12 +3,15 @@ import json import os import subprocess +import tempfile import uuid +from collections.abc import Iterator from pathlib import Path import numpy as np import pytest +from notify_cleanup import remove_new_empty_notify_dirs from systemd_probe import classify_systemd_probe SEQ_LEN = 1024 @@ -156,3 +159,20 @@ def systemd_manager() -> None: if verdict == "error": pytest.fail(f"systemd-run probe failed: {reason}") pytest.skip(f"system manager transient units unavailable: {reason}") + + +@pytest.fixture(scope="session", autouse=True) +def _tmp_leak_guard(tmp_path_factory: pytest.TempPathFactory) -> Iterator[None]: + """Keep the session's IPC and torchrun leftovers out of the shared temp root. + + torchrun reads PET_LOG_DIR as --log-dir, so its logs land under basetemp instead of a fresh + tempdir()/torchelastic_*; child processes inherit os.environ. Nothing outside basetemp is + deleted except empty ht-notify-* directories that appeared during the session. + """ + root = Path(tempfile.gettempdir()) + before = {path.name for path in root.glob("ht-notify-*")} + patch = pytest.MonkeyPatch() + patch.setenv("PET_LOG_DIR", str(tmp_path_factory.mktemp("torchelastic"))) + yield + patch.undo() + remove_new_empty_notify_dirs(root, before) diff --git a/tests/notify_cleanup.py b/tests/notify_cleanup.py new file mode 100644 index 0000000..62c0730 --- /dev/null +++ b/tests/notify_cleanup.py @@ -0,0 +1,16 @@ +"""Session-end cleanup of the empty ht-notify-* directories that tests leave in the shared /tmp.""" + +from __future__ import annotations + +import os +from pathlib import Path + + +def remove_new_empty_notify_dirs(root: Path, before: set[str]) -> None: + """rmdir only: a non-empty directory may be a concurrent same-uid session's live store.""" + for path in root.glob("ht-notify-*"): + if path.name not in before and path.lstat().st_uid == os.getuid(): + try: + path.rmdir() + except OSError: + continue # still in use, or a store a test never closed: leave it visible