torch-atomic-save 0.1.0__py3-none-any.whl

torch_atomic_save/__init__.py

@@ -0,0 +1,3 @@
+from .manager import SlurmAtomicManager, install_slurm_handler
+
+__all__ = ["SlurmAtomicManager", "install_slurm_handler"]

torch_atomic_save/manager.py

@@ -0,0 +1,172 @@
+import os
+import shutil
+import tempfile
+import threading
+import concurrent.futures
+import logging
+import signal
+import torch
+from typing import Any, Optional, Dict, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+class SlurmAtomicManager:
+    """
+    Manages asynchronous, atomic checkpoint saving across different filesystems.
+    Optimized for Slurm/Lustre environments.
+    """
+
+    def __init__(self, max_workers: int = 4):
+        self._executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=max_workers,
+            thread_name_prefix="SlurmSaveWorker"
+        )
+        # path -> (future, cancellation_event)
+        self._registry: Dict[str, Tuple[concurrent.futures.Future, threading.Event]] = {}
+        self._lock = threading.Lock()
+        self._shutdown_event = threading.Event()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.shutdown()
+
+    def wait_for_all(self, timeout: Optional[float] = None):
+        """Blocks until all currently pending saves are complete."""
+        with self._lock:
+            futures = [f for f, _ in self._registry.values()]
+        concurrent.futures.wait(futures, timeout=timeout)
+
+    def save(self, model: torch.nn.Module, path: str, tmp_dir: Optional[str] = None, half_prec: bool = False):
+        """
+        Main entry point for saving models.
+        Moves tensors to CPU and clones them to prevent race conditions.
+        """
+        path = os.path.abspath(path)
+
+        # Snapshot weights: Move to CPU and Clone
+        if half_prec:
+            state_dict = {k: v.half().cpu().clone() for k, v in model.state_dict().items()}
+        else:
+            state_dict = {k: v.cpu().clone() for k, v in model.state_dict().items()}
+
+        if tmp_dir is None:
+            # Synchronous fallback
+            os.makedirs(os.path.dirname(path), exist_ok=True)
+            torch.save(state_dict, path)
+        else:
+            self._atomic_save(state_dict, path, tmp_dir)
+
+    def _atomic_save(self, obj: Any, path: str, tmp_dir: str) -> None:
+        """
+        Saves an object to a staging area, then schedules an atomic cross-FS move.
+        """
+        if self._shutdown_event.is_set():
+            logger.warning(f"Manager is shutting down. Ignoring save request for {path}")
+            return
+
+        path = os.path.abspath(path)
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+
+        if not os.path.isdir(tmp_dir):
+            raise FileNotFoundError(f"Staging directory missing: {tmp_dir}")
+
+        # 1. Setup staging dir
+        fd, tmp_src = tempfile.mkstemp(dir=tmp_dir, suffix=".pt.staging")
+        os.close(fd)
+
+        # 2. Registry Update & Cancellation
+        cancelled = threading.Event()
+        with self._lock:
+            if path in self._registry:
+                prev_f, prev_ev = self._registry[path]
+                prev_ev.set()  # Signal background thread to skip the move
+                prev_f.cancel()
+
+            future = self._executor.submit(
+                self._save_and_copy, obj, tmp_src, path, cancelled
+            )
+            self._registry[path] = (future, cancelled)
+            future.add_done_callback(lambda f: self._cleanup_registry(path, f))
+
+    def _cleanup_registry(self, path: str, future: concurrent.futures.Future) -> None:
+        """Removes the future from registry and logs errors."""
+        with self._lock:
+            current = self._registry.get(path)
+            if current and current[0] is future:
+                self._registry.pop(path, None)
+
+        if not future.cancelled():
+            exc = future.exception()
+            if exc:
+                logger.error(f"Async save to {path} failed: {exc}", exc_info=exc)
+
+    def _save_and_copy(self, obj, tmp_src, dst, cancelled):
+        """Handles both the initial write and the atomic move."""
+        try:
+            if not cancelled.is_set():
+                torch.save(obj, tmp_src)
+
+            self._atomic_copy_and_cleanup(tmp_src, dst, cancelled)
+        except Exception as e:
+            # Cleanup tmp_src if torch.save failed
+            if os.path.exists(tmp_src):
+                os.unlink(tmp_src)
+            raise
+
+    def _atomic_copy_and_cleanup(self, src: str, dst: str, cancelled: threading.Event, max_retries: int = 5) -> None:
+        """The background worker logic."""
+        dst_dir = os.path.dirname(os.path.abspath(dst))
+        try:
+            # Check if same device (ST_DEV)
+            if _same_filesystem(src, dst):
+                if not cancelled.is_set():
+                    os.replace(src, dst)
+            else:
+                # Cross-FS: Create sibling tmp in destination's parent
+                fd, sibling_tmp = tempfile.mkstemp(dir=dst_dir, suffix=".atomic_tmp")
+                os.close(fd)
+                try:
+                    shutil.copy2(src, sibling_tmp)
+                    if not cancelled.is_set():
+                        os.replace(sibling_tmp, dst)
+                    else:
+                        os.unlink(sibling_tmp)
+                except Exception:
+                    try:
+                        os.unlink(sibling_tmp)
+                    except OSError:
+                        pass
+                    raise
+        finally:
+            try:
+                os.unlink(src)
+            except OSError:
+                pass
+
+    def shutdown(self, wait: bool = True):
+        """Cleanly shut down the executor."""
+        self._shutdown_event.set()
+        self._executor.shutdown(wait=wait)
+
+
+def install_slurm_handler(manager: SlurmAtomicManager):
+    """
+    Installs a handler for SIGTERM (Slurm preemption) to ensure
+    the manager flushes pending IO before the process exits.
+    """
+
+    def handle_sigterm(signum, frame):
+        logger.info("Received SIGTERM/Preemption signal. Flushing IO...")
+        manager.shutdown(wait=True)
+        # Optional: Re-raise or exit
+        # os._exit(0)
+
+    signal.signal(signal.SIGTERM, handle_sigterm)
+
+
+def _same_filesystem(path_a: str, path_b: str) -> bool:
+    stat_b_target = path_b if os.path.exists(path_b) else os.path.dirname(os.path.abspath(path_b))
+    return os.stat(path_a).st_dev == os.stat(stat_b_target).st_dev

torch_atomic_save-0.1.0.dist-info/METADATA

@@ -0,0 +1,52 @@
+Metadata-Version: 2.4
+Name: torch_atomic_save
+Version: 0.1.0
+Summary: Atomic, asynchronous checkpoint saving for Pytorch for Slurm/Lustre environments
+Author-email: Patrick Carnahan <pcarnah@uwo.ca>
+License: MIT
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0.0
+Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Dynamic: license-file
+
+# torch-atomic-save
+
+[![Pytest](https://github.com/pcarnah/torch-atomic-save/actions/workflows/test.yml/badge.svg)](https://github.com/pcarnah/torch-atomic-save/actions)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+
+An asynchronous, atomic checkpointing utility for PyTorch, optimized for Slurm and Lustre/NFS environments.
+
+## Key Features
+* **Atomic Moves:** Prevents corrupted checkpoints during Slurm preemption.
+* **Non-Blocking:** Offloads I/O to a background thread pool.
+* **Race Condition Protection:** Automatically clones tensors to CPU before background saving.
+* **Cross-FS Support:** Handles moves between local SSD and network storage safely.
+
+## Installation
+```bash
+pip install torch-atomic-save
+```
+
+
+## Usage
+```python
+from torch_atomic_save import SlurmAtomicManager
+
+# Initialize the manager
+manager = SlurmAtomicManager(max_workers=4)
+
+# In your training loop:
+if epoch % save_interval == 0:
+    manager.save(model, "path/to/checkpoints/model.pt", tmp_dir="/scratch/user/tmp")
+
+# Ensure all I/O is finished before exiting
+manager.wait_for_all()
+```
+
+## Attribution
+If you use this software in your research, please cite it using the "Cite this repository" button or the provided CITATION.cff.

torch_atomic_save-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+torch_atomic_save/__init__.py,sha256=NhsLuYY4EffBDJgeShgLEsNrEnl7rXiVlDqBlHN_S8g,123
+torch_atomic_save/manager.py,sha256=XXpNCHpRP8hmxcHIYvHMQkzTnC15wPrKhPl2K9GUxbE,6383
+torch_atomic_save-0.1.0.dist-info/licenses/LICENSE,sha256=WvXpsuVMCjnZIRoaaalqisATp86XcY0_PQSxx7HRkkA,1092
+torch_atomic_save-0.1.0.dist-info/METADATA,sha256=BVcChI39xwIsEdsBsZ1rnlDEdAoPCHe1PVSbAHccZzA,1799
+torch_atomic_save-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+torch_atomic_save-0.1.0.dist-info/top_level.txt,sha256=h0Omk6Y6-mjulu1PfoOZ7TblXy4rtRT6KRgd5s8gpTA,18
+torch_atomic_save-0.1.0.dist-info/RECORD,,

torch_atomic_save-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

torch_atomic_save-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Patrick Carnahan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torch_atomic_save-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+torch_atomic_save