nshtrainer 0.26.0__tar.gz → 0.26.1__tar.gz

{nshtrainer-0.26.0 → nshtrainer-0.26.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.26.0
+Version: 0.26.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.26.0 → nshtrainer-0.26.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.26.0"
+version = "0.26.1"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshtrainer-0.26.0 → nshtrainer-0.26.1}/src/nshtrainer/_checkpoint/metadata.py

@@ -1,7 +1,6 @@
 import copy
 import datetime
 import logging
-import shutil
 from collections.abc import Callable
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, cast
@@ -11,7 +10,7 @@ import numpy as np
 import torch
 
 from ..util._environment_info import EnvironmentConfig
-from ..util.path import compute_file_checksum, get_relative_path
+from ..util.path import compute_file_checksum, try_symlink_or_copy
 
 if TYPE_CHECKING:
     from ..model import BaseConfig, LightningModuleBase
@@ -142,21 +141,7 @@ def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Pat
     # Link the metadata files to the new checkpoint
     path = _metadata_path(checkpoint_path)
     linked_path = _metadata_path(linked_checkpoint_path)
-    try:
-        try:
-            # linked_path.symlink_to(path)
-            # We should store the path as a relative path
-            # to the metadata file to avoid issues with
-            # moving the checkpoint directory
-            linked_path.symlink_to(get_relative_path(linked_path, path))
-        except OSError:
-            # on Windows, special permissions are required to create symbolic links as a regular user
-            # fall back to copying the file
-            shutil.copy(path, linked_path)
-    except Exception:
-        log.exception(f"Failed to link {path} to {linked_path}")
-    else:
-        log.debug(f"Linked {path} to {linked_path}")
+    try_symlink_or_copy(path, linked_path)
 
 
 def _sort_ckpts_by_metadata(

{nshtrainer-0.26.0 → nshtrainer-0.26.1}/src/nshtrainer/_checkpoint/saver.py

@@ -5,7 +5,7 @@ from pathlib import Path
 
 from lightning.pytorch import Trainer
 
-from ..util.path import get_relative_path
+from ..util.path import try_symlink_or_copy
 from .metadata import _link_checkpoint_metadata, _remove_checkpoint_metadata
 
 log = logging.getLogger(__name__)
@@ -34,13 +34,7 @@ def _link_checkpoint(
         if metadata:
             _remove_checkpoint_metadata(linkpath)
 
-    try:
-        linkpath.symlink_to(get_relative_path(linkpath, filepath))
-    except OSError:
-        # on Windows, special permissions are required to create symbolic links as a regular user
-        # fall back to copying the file
-        shutil.copy(filepath, linkpath)
-
+    try_symlink_or_copy(filepath, linkpath)
     if metadata:
         _link_checkpoint_metadata(filepath, linkpath)
 

{nshtrainer-0.26.0 → nshtrainer-0.26.1}/src/nshtrainer/util/path.py

@@ -1,8 +1,13 @@
 import hashlib
+import logging
 import os
+import platform
+import shutil
 from pathlib import Path
 from typing import TypeAlias
 
+log = logging.getLogger(__name__)
+
 _Path: TypeAlias = str | Path | os.PathLike
 
 
@@ -68,3 +73,32 @@ def compute_file_checksum(file_path: Path) -> str:
         for byte_block in iter(lambda: f.read(4096), b""):
             sha256_hash.update(byte_block)
     return sha256_hash.hexdigest()
+
+
+def try_symlink_or_copy(
+    file_path: Path,
+    link_path: Path,
+    target_is_directory: bool = False,
+    relative: bool = True,
+):
+    """
+    Symlinks on Unix, copies on Windows.
+    """
+
+    symlink_target = get_relative_path(link_path, file_path) if relative else file_path
+    try:
+        if platform.system() == "Windows":
+            if target_is_directory:
+                shutil.copytree(file_path, link_path)
+            else:
+                shutil.copy(file_path, link_path)
+        else:
+            link_path.symlink_to(
+                symlink_target, target_is_directory=target_is_directory
+            )
+    except Exception:
+        log.exception(f"Failed to create symlink or copy {file_path} to {link_path}")
+        return False
+    else:
+        log.debug(f"Created symlink or copied {file_path} to {link_path}")
+        return True