nshtrainer 1.0.0b40__py3-none-any.whl → 1.0.0b42__py3-none-any.whl

nshtrainer/_checkpoint/metadata.py

@@ -130,7 +130,7 @@ def _write_checkpoint_metadata(trainer: Trainer, checkpoint_path: Path):
     return metadata_path
 
 
-def _remove_checkpoint_metadata(checkpoint_path: Path):
+def remove_checkpoint_metadata(checkpoint_path: Path):
     path = _metadata_path(checkpoint_path)
     try:
         path.unlink(missing_ok=True)
@@ -140,23 +140,15 @@ def _remove_checkpoint_metadata(checkpoint_path: Path):
         log.debug(f"Removed {path}")
 
 
-def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Path):
+def link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Path):
     # First, remove any existing metadata files
-    _remove_checkpoint_metadata(linked_checkpoint_path)
+    remove_checkpoint_metadata(linked_checkpoint_path)
 
     # Link the metadata files to the new checkpoint
     path = _metadata_path(checkpoint_path)
     linked_path = _metadata_path(linked_checkpoint_path)
-    try_symlink_or_copy(path, linked_path)
 
+    if not path.exists():
+        raise FileNotFoundError(f"Checkpoint path does not exist: {checkpoint_path}")
 
-def _sort_ckpts_by_metadata(
-    checkpoint_paths: list[Path],
-    key: Callable[[CheckpointMetadata, Path], Any],
-    reverse: bool = False,
-):
-    return sorted(
-        [(CheckpointMetadata.from_ckpt_path(path), path) for path in checkpoint_paths],
-        key=lambda args_tuple: key(*args_tuple),
-        reverse=reverse,
-    )
+    try_symlink_or_copy(path, linked_path)

nshtrainer/_checkpoint/saver.py

@@ -8,12 +8,12 @@ from pathlib import Path
 from lightning.pytorch import Trainer
 
 from ..util.path import try_symlink_or_copy
-from .metadata import _link_checkpoint_metadata, _remove_checkpoint_metadata
+from .metadata import link_checkpoint_metadata, remove_checkpoint_metadata
 
 log = logging.getLogger(__name__)
 
 
-def _link_checkpoint(
+def link_checkpoint(
     filepath: str | Path | os.PathLike,
     linkpath: str | Path | os.PathLike,
     *,
@@ -36,14 +36,14 @@ def _link_checkpoint(
             log.debug(f"Removed {linkpath=}")
 
         if metadata:
-            _remove_checkpoint_metadata(linkpath)
+            remove_checkpoint_metadata(linkpath)
 
     try_symlink_or_copy(filepath, linkpath)
     if metadata:
-        _link_checkpoint_metadata(filepath, linkpath)
+        link_checkpoint_metadata(filepath, linkpath)
 
 
-def _remove_checkpoint(
+def remove_checkpoint(
     trainer: Trainer,
     filepath: str | Path | os.PathLike,
     *,
@@ -54,4 +54,4 @@ def _remove_checkpoint(
     trainer.strategy.remove_checkpoint(filepath)
 
     if metadata:
-        _remove_checkpoint_metadata(filepath)
+        remove_checkpoint_metadata(filepath)

nshtrainer/callbacks/checkpoint/_base.py

@@ -12,7 +12,7 @@ from lightning.pytorch.callbacks import Checkpoint
 from typing_extensions import TypeVar, override
 
 from ..._checkpoint.metadata import CheckpointMetadata
-from ..._checkpoint.saver import _link_checkpoint, _remove_checkpoint
+from ..._checkpoint.saver import link_checkpoint, remove_checkpoint
 from ..base import CallbackConfigBase
 
 if TYPE_CHECKING:
@@ -122,7 +122,7 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
                 )
                 continue
 
-            _remove_checkpoint(trainer, old_ckpt_path, metadata=True)
+            remove_checkpoint(trainer, old_ckpt_path, metadata=True)
             log.debug(f"Removed old checkpoint: {old_ckpt_path}")
 
     def current_metrics(self, trainer: Trainer) -> dict[str, Any]:
@@ -167,7 +167,7 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
             # Create the latest symlink
             if (symlink_filename := self.symlink_path()) is not None:
                 symlink_path = self.dirpath / symlink_filename
-                _link_checkpoint(filepath, symlink_path, metadata=True)
+                link_checkpoint(filepath, symlink_path, metadata=True)
                 log.debug(f"Created latest symlink: {symlink_path}")
 
         # Barrier to ensure all processes have saved the checkpoint,

nshtrainer/nn/__init__.py

@@ -4,6 +4,7 @@ from .mlp import MLP as MLP
 from .mlp import MLPConfig as MLPConfig
 from .mlp import MLPConfigDict as MLPConfigDict
 from .mlp import ResidualSequential as ResidualSequential
+from .mlp import custom_seed_context as custom_seed_context
 from .module_dict import TypedModuleDict as TypedModuleDict
 from .module_list import TypedModuleList as TypedModuleList
 from .nonlinearity import ELUNonlinearityConfig as ELUNonlinearityConfig

nshtrainer/nn/mlp.py

@@ -99,6 +99,18 @@ class MLPConfig(C.Config):
         )
 
 
+@contextlib.contextmanager
+def custom_seed_context(seed: int | None):
+    with contextlib.ExitStack() as stack:
+        if seed is not None:
+            stack.enter_context(
+                torch.random.fork_rng(devices=range(torch.cuda.device_count()))
+            )
+            torch.manual_seed(seed)
+
+        yield
+
+
 def MLP(
     dims: Sequence[int],
     activation: NonlinearityConfigBase
@@ -140,13 +152,7 @@ def MLP(
         nn.Sequential: The constructed MLP.
     """
 
-    with contextlib.ExitStack() as stack:
-        if seed is not None:
-            stack.enter_context(
-                torch.random.fork_rng(devices=range(torch.cuda.device_count()))
-            )
-            torch.manual_seed(seed)
-
+    with custom_seed_context(seed):
         if activation is None:
             activation = nonlinearity
 

nshtrainer/nn/tests/test_mlp.py

@@ -5,7 +5,7 @@ from typing import cast
 import pytest
 import torch
 
-from nshtrainer.nn.mlp import MLP
+from nshtrainer.nn.mlp import MLP, custom_seed_context
 
 
 def test_mlp_seed_reproducibility():
@@ -53,3 +53,68 @@ def test_mlp_seed_reproducibility():
     assert not torch.allclose(
         cast(torch.Tensor, mlp4[0].weight), cast(torch.Tensor, mlp5[0].weight)
     )
+
+
+def test_custom_seed_context():
+    """Test that custom_seed_context properly controls random number generation."""
+
+    # Test that the same seed produces the same random numbers
+    with custom_seed_context(42):
+        tensor1 = torch.randn(10)
+
+    with custom_seed_context(42):
+        tensor2 = torch.randn(10)
+
+    # Same seed should produce identical random tensors
+    assert torch.allclose(tensor1, tensor2)
+
+    # Test that different seeds produce different random numbers
+    with custom_seed_context(123):
+        tensor3 = torch.randn(10)
+
+    # Different seeds should produce different random tensors
+    assert not torch.allclose(tensor1, tensor3)
+
+
+def test_custom_seed_context_preserves_state():
+    """Test that custom_seed_context preserves the original random state."""
+
+    # Set a known seed for the test
+    original_seed = 789
+    torch.manual_seed(original_seed)
+
+    # Generate a tensor with the original seed
+    original_tensor = torch.randn(10)
+
+    # Use a different seed in the context
+    with custom_seed_context(42):
+        # This should use the temporary seed
+        context_tensor = torch.randn(10)
+
+    # After exiting the context, we should be back to the original seed state
+    # Reset the generator to get the same sequence again
+    torch.manual_seed(original_seed)
+    expected_tensor = torch.randn(10)
+
+    # The tensor generated after the context should match what we would get
+    # if we had just set the original seed again
+    assert torch.allclose(original_tensor, expected_tensor)
+
+    # And it should be different from the tensor generated inside the context
+    assert not torch.allclose(original_tensor, context_tensor)
+
+
+def test_custom_seed_context_with_none():
+    """Test that custom_seed_context with None seed doesn't affect randomization."""
+
+    # Set a known seed
+    torch.manual_seed(555)
+    expected_tensor = torch.randn(10)
+
+    # Reset and use None seed in context
+    torch.manual_seed(555)
+    with custom_seed_context(None):
+        actual_tensor = torch.randn(10)
+
+    # With None seed, the context should not affect the random state
+    assert torch.allclose(expected_tensor, actual_tensor)

nshtrainer/util/path.py

@@ -81,6 +81,7 @@ def compute_file_checksum(file_path: Path) -> str:
 def try_symlink_or_copy(
     file_path: Path,
     link_path: Path,
+    *,
     target_is_directory: bool = False,
     relative: bool = True,
     remove_existing: bool = True,
@@ -92,7 +93,10 @@ def try_symlink_or_copy(
     # If the link already exists, remove it
     if remove_existing:
         try:
-            if link_path.exists():
+            if link_path.exists(follow_symlinks=False):
+                # follow_symlinks=False is EXTREMELY important here
+                # Otherwise, we've already deleted the file that the symlink
+                # used to point to, so this always returns False
                 if link_path.is_dir():
                     shutil.rmtree(link_path)
                 else:

{nshtrainer-1.0.0b40.dist-info → nshtrainer-1.0.0b42.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.0.0b40
+Version: 1.0.0b42
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.0.0b40.dist-info → nshtrainer-1.0.0b42.dist-info}/RECORD

@@ -1,8 +1,8 @@
 nshtrainer/.nshconfig.generated.json,sha256=yZd6cn1RhvNNJUgiUTRYut8ofZYvbulnpPG-rZIRhi4,106
 nshtrainer/__init__.py,sha256=g_moPnfQxSxFZX5NB9ILQQOJrt4RTRuiFt9N0STIpxM,874
 nshtrainer/_callback.py,sha256=tXQCDzS6CvMTuTY5lQSH5qZs1pXUi-gt9bQdpXMVdEs,12715
-nshtrainer/_checkpoint/metadata.py,sha256=PHy-54Cg-o3OtCffAqrVv6ZVMU7zhRo_-sZiSEEno1Y,5019
-nshtrainer/_checkpoint/saver.py,sha256=LOP8jjKF0Dw9x9H-BKrLMWlEp1XTan2DUK0zQUCWw5U,1360
+nshtrainer/_checkpoint/metadata.py,sha256=Kp1BK1iHCsOzjjy_B0Rbs0yIyd_yV9gZl4uym7rNQ_E,4796
+nshtrainer/_checkpoint/saver.py,sha256=qG8dEetyqwLJS0MkdWDz4DCAInaSyafW8-Oe8x0UVnE,1353
 nshtrainer/_directory.py,sha256=TJR9ccyuzRlAVfVjGyeQ3E2AFAcz-XbBCxWfiXo2SlY,3191
 nshtrainer/_experimental/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 nshtrainer/_hf_hub.py,sha256=4OsCbIITnZk_YLyoMrVyZ0SIN04FBxlC0ig2Et8UAdo,14287
@@ -10,7 +10,7 @@ nshtrainer/callbacks/__init__.py,sha256=4giOYT8A709UOLRtQEt16QbOAFUHCjJ_aLB7ITTw
 nshtrainer/callbacks/actsave.py,sha256=NSXIIu62MNYe5gz479SMW33bdoKYoYtWtd_iTWFpKpc,3881
 nshtrainer/callbacks/base.py,sha256=Alaou1IHAIlMEM7g58d_02ozY2xWlshBN7fsw5Ee21s,3683
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=l8tkHc83_mLiU0-wT09SWdRzwpm2ulbkLzcuCmuTwzE,620
-nshtrainer/callbacks/checkpoint/_base.py,sha256=ZVEUVl5kjCSSe69Q0rMUbKBNNUog0pxBwWkeyuxG2w0,6304
+nshtrainer/callbacks/checkpoint/_base.py,sha256=wCJBRI0pQYZc3GBu0b-aUBlBDhd39AdL82VvFgKmv3k,6300
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=2CQuhPJ3Fi7lDw7z-J8kXXXuDU8-4HcU48oZxR49apk,2667
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=vn-as3ex7kaTRcKsIurVtM6kUSHYNwHJeYG82j2dMcc,3554
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=nljzETqkHwA-4g8mxaeFK5HxA8My0dlIPzIUscSMWyk,3525
@@ -117,12 +117,12 @@ nshtrainer/model/base.py,sha256=JL3AmH17GQjQIoMrZl3O0vUI7dj5ZsO5iEJgoLPyzHw,1035
 nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0dXLvqHw,3154
 nshtrainer/model/mixins/debug.py,sha256=1LX9KzeFX9JDPs_a6YCdYDZXLhEk_5rBO2aCqlfBy7w,2087
 nshtrainer/model/mixins/logger.py,sha256=27H99FuLaxc6_dDLG2pid4E_5E0-eLGnc2Ifpt0HYIM,6066
-nshtrainer/nn/__init__.py,sha256=7KCs-GDOynCXAIdwkgAQacc0p3FHLEION50UtrvgAOc,1463
-nshtrainer/nn/mlp.py,sha256=OatI_pzAqa2awKfvMY4CnznO-LltLT8NHQPJKSNn8IM,6979
+nshtrainer/nn/__init__.py,sha256=0FgeoaLYtRiSLT8fdPigLD8t-d8DKR8IQDw16JA9lT4,1523
+nshtrainer/nn/mlp.py,sha256=_a8rJJniSCvM08gyQGO-5MUoO18U9_FSGGn3tZL2_U4,7101
 nshtrainer/nn/module_dict.py,sha256=9plb8aQUx5TUEPhX5jI9u8LrpTeKe7jZAHi8iIqcN8w,2365
 nshtrainer/nn/module_list.py,sha256=UB43pcwD_3nUke_DyLQt-iXKhWdKM6Zjm84lRC1hPYA,1755
 nshtrainer/nn/nonlinearity.py,sha256=xmaL4QCRvCxqmaGIOwetJeKK-6IK4m2OV7D3SjxSwJQ,6322
-nshtrainer/nn/tests/test_mlp.py,sha256=xBPiHlBvOCn67EbpzzKL-2FU7ikGxHT3i6CMSp1wk7M,1840
+nshtrainer/nn/tests/test_mlp.py,sha256=COw88Oc4FMvj_mGdQf6F2UtgZ39FE2lbQjTmMAwtCWE,4031
 nshtrainer/optimizer.py,sha256=u968GRNPUNn3f_9BEY2RBNuJq5O3wJWams3NG0dkrOA,1738
 nshtrainer/profiler/__init__.py,sha256=RjaNBoVcTFu8lF0dNlFp-2LaPYdonoIbDy2_KhgF0Ek,594
 nshtrainer/profiler/_base.py,sha256=kFcSVn9gJuMwgDxbfyHh46CmEAIPZjxw3yjPbKgzvwA,950
@@ -148,11 +148,11 @@ nshtrainer/util/config/__init__.py,sha256=Z39JJufSb61Lhn2GfVcv3eFW_eorOrN9-9llDW
 nshtrainer/util/config/dtype.py,sha256=Fn_MhhQoHPyFAnFPSwvcvLiGR3yWFIszMba02CJiC4g,2213
 nshtrainer/util/config/duration.py,sha256=mM-UfU_HvhXwW33TYEDg0x58n80tnle2e6VaWtxZTjk,764
 nshtrainer/util/environment.py,sha256=s-B5nY0cKYXdFMdNYumvC_xxacMATiI4DvV2gUDu20k,4195
-nshtrainer/util/path.py,sha256=L-Nh9tlXSUfoP19TFbQq8I0AfS5ugCfGYTYFeddDHcs,3516
+nshtrainer/util/path.py,sha256=IL8MS29MzP2aBdQhqMl2Eh4RvyuZbzaIbAVZzsdLDSk,3754
 nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.0.0b40.dist-info/METADATA,sha256=m_MxoWKJnesp95-EClArQl90vTm____cgqMdNJ6n-Ng,988
-nshtrainer-1.0.0b40.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-nshtrainer-1.0.0b40.dist-info/RECORD,,
+nshtrainer-1.0.0b42.dist-info/METADATA,sha256=m8vNeWow5AfufFEd6uCjm6w-0NOCTTzCFjXDed0eq6U,988
+nshtrainer-1.0.0b42.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+nshtrainer-1.0.0b42.dist-info/RECORD,,