nshtrainer 1.0.0b42__tar.gz → 1.0.0b44__tar.gz

{nshtrainer-1.0.0b42 → nshtrainer-1.0.0b44}/PKG-INFO

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

{nshtrainer-1.0.0b42 → nshtrainer-1.0.0b44}/pyproject.toml

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

{nshtrainer-1.0.0b42 → nshtrainer-1.0.0b44}/src/nshtrainer/_checkpoint/metadata.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import copy
 import datetime
 import logging
-from collections.abc import Callable
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar
 
@@ -115,7 +114,7 @@ def _metadata_path(checkpoint_path: Path):
     return checkpoint_path.with_suffix(CheckpointMetadata.PATH_SUFFIX)
 
 
-def _write_checkpoint_metadata(trainer: Trainer, checkpoint_path: Path):
+def write_checkpoint_metadata(trainer: Trainer, checkpoint_path: Path):
     metadata_path = _metadata_path(checkpoint_path)
     metadata = _generate_checkpoint_metadata(trainer, checkpoint_path, metadata_path)
 

{nshtrainer-1.0.0b42 → nshtrainer-1.0.0b44}/src/nshtrainer/_checkpoint/saver.py

@@ -25,7 +25,10 @@ def link_checkpoint(
 
     if remove_existing:
         try:
-            if linkpath.exists():
+            if linkpath.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 linkpath.is_dir():
                     shutil.rmtree(linkpath)
                 else:

{nshtrainer-1.0.0b42 → nshtrainer-1.0.0b44}/src/nshtrainer/callbacks/__init__.py

@@ -40,6 +40,10 @@ from .log_epoch import LogEpochCallback as LogEpochCallback
 from .log_epoch import LogEpochCallbackConfig as LogEpochCallbackConfig
 from .lr_monitor import LearningRateMonitor as LearningRateMonitor
 from .lr_monitor import LearningRateMonitorConfig as LearningRateMonitorConfig
+from .metric_validation import MetricValidationCallback as MetricValidationCallback
+from .metric_validation import (
+    MetricValidationCallbackConfig as MetricValidationCallbackConfig,
+)
 from .norm_logging import NormLoggingCallback as NormLoggingCallback
 from .norm_logging import NormLoggingCallbackConfig as NormLoggingCallbackConfig
 from .print_table import PrintTableMetricsCallback as PrintTableMetricsCallback

nshtrainer-1.0.0b44/src/nshtrainer/callbacks/metric_validation.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import logging
+from typing import Literal
+
+from lightning.pytorch.utilities.exceptions import MisconfigurationException
+from typing_extensions import final, override, assert_never
+
+from .._callback import NTCallbackBase
+from ..metrics import MetricConfig
+from .base import CallbackConfigBase, callback_registry
+
+log = logging.getLogger(__name__)
+
+
+@final
+@callback_registry.register
+class MetricValidationCallbackConfig(CallbackConfigBase):
+    name: Literal["metric_validation"] = "metric_validation"
+
+    error_behavior: Literal["raise", "warn"] = "raise"
+    """
+    Behavior when an error occurs during validation:
+    - "raise": Raise an error and stop the training.
+    - "warn": Log a warning and continue the training.
+    """
+
+    validate_default_metric: bool = True
+    """Whether to validate the default metric from the root config."""
+
+    metrics: list[MetricConfig] = []
+    """List of metrics to validate."""
+
+    @override
+    def create_callbacks(self, trainer_config):
+        metrics = self.metrics.copy()
+        if (
+            self.validate_default_metric
+            and (default_metric := trainer_config.primary_metric) is not None
+        ):
+            metrics.append(default_metric)
+
+        yield MetricValidationCallback(self, metrics)
+
+
+class MetricValidationCallback(NTCallbackBase):
+    def __init__(
+        self, config: MetricValidationCallbackConfig, metrics: list[MetricConfig]
+    ):
+        super().__init__()
+
+        self.config = config
+        self.metrics = metrics
+
+    @override
+    def on_sanity_check_end(self, trainer, pl_module):
+        super().on_sanity_check_end(trainer, pl_module)
+
+        log.debug("Validating metrics...")
+        logged_metrics = set(trainer.logged_metrics.keys())
+        for metric in self.metrics:
+            if metric.validation_monitor in logged_metrics:
+                continue
+
+            match self.config.error_behavior:
+                case "raise":
+                    raise MisconfigurationException(
+                        f"Metric '{metric.validation_monitor}' not found in logged metrics."
+                    )
+                case "warn":
+                    log.warning(
+                        f"Metric '{metric.validation_monitor}' not found in logged metrics."
+                    )
+                case _:
+                    assert_never(self.config.error_behavior)

{nshtrainer-1.0.0b42 → nshtrainer-1.0.0b44}/src/nshtrainer/configs/__init__.py

@@ -39,6 +39,9 @@ from nshtrainer.callbacks import (
 )
 from nshtrainer.callbacks import LearningRateMonitorConfig as LearningRateMonitorConfig
 from nshtrainer.callbacks import LogEpochCallbackConfig as LogEpochCallbackConfig
+from nshtrainer.callbacks import (
+    MetricValidationCallbackConfig as MetricValidationCallbackConfig,
+)
 from nshtrainer.callbacks import NormLoggingCallbackConfig as NormLoggingCallbackConfig
 from nshtrainer.callbacks import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
@@ -287,6 +290,7 @@ __all__ = [
     "MPIEnvironmentPlugin",
     "MPSAcceleratorConfig",
     "MetricConfig",
+    "MetricValidationCallbackConfig",
     "MishNonlinearityConfig",
     "MixedPrecisionPluginConfig",
     "NonlinearityConfig",

{nshtrainer-1.0.0b42 → nshtrainer-1.0.0b44}/src/nshtrainer/configs/callbacks/__init__.py

@@ -28,6 +28,9 @@ from nshtrainer.callbacks import (
 )
 from nshtrainer.callbacks import LearningRateMonitorConfig as LearningRateMonitorConfig
 from nshtrainer.callbacks import LogEpochCallbackConfig as LogEpochCallbackConfig
+from nshtrainer.callbacks import (
+    MetricValidationCallbackConfig as MetricValidationCallbackConfig,
+)
 from nshtrainer.callbacks import NormLoggingCallbackConfig as NormLoggingCallbackConfig
 from nshtrainer.callbacks import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
@@ -65,6 +68,7 @@ from . import finite_checks as finite_checks
 from . import gradient_skipping as gradient_skipping
 from . import log_epoch as log_epoch
 from . import lr_monitor as lr_monitor
+from . import metric_validation as metric_validation
 from . import norm_logging as norm_logging
 from . import print_table as print_table
 from . import rlp_sanity_checks as rlp_sanity_checks
@@ -91,6 +95,7 @@ __all__ = [
     "LearningRateMonitorConfig",
     "LogEpochCallbackConfig",
     "MetricConfig",
+    "MetricValidationCallbackConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
     "PrintTableMetricsCallbackConfig",
@@ -110,6 +115,7 @@ __all__ = [
     "gradient_skipping",
     "log_epoch",
     "lr_monitor",
+    "metric_validation",
     "norm_logging",
     "print_table",
     "rlp_sanity_checks",

nshtrainer-1.0.0b44/src/nshtrainer/configs/callbacks/metric_validation/__init__.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.callbacks.metric_validation import (
+    CallbackConfigBase as CallbackConfigBase,
+)
+from nshtrainer.callbacks.metric_validation import MetricConfig as MetricConfig
+from nshtrainer.callbacks.metric_validation import (
+    MetricValidationCallbackConfig as MetricValidationCallbackConfig,
+)
+from nshtrainer.callbacks.metric_validation import (
+    callback_registry as callback_registry,
+)
+
+__all__ = [
+    "CallbackConfigBase",
+    "MetricConfig",
+    "MetricValidationCallbackConfig",
+    "callback_registry",
+]

{nshtrainer-1.0.0b42 → nshtrainer-1.0.0b44}/src/nshtrainer/configs/trainer/__init__.py

@@ -38,6 +38,9 @@ from nshtrainer.trainer._config import LogEpochCallbackConfig as LogEpochCallbac
 from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
 from nshtrainer.trainer._config import LoggerConfigBase as LoggerConfigBase
 from nshtrainer.trainer._config import MetricConfig as MetricConfig
+from nshtrainer.trainer._config import (
+    MetricValidationCallbackConfig as MetricValidationCallbackConfig,
+)
 from nshtrainer.trainer._config import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
 )
@@ -164,6 +167,7 @@ __all__ = [
     "MPIEnvironmentPlugin",
     "MPSAcceleratorConfig",
     "MetricConfig",
+    "MetricValidationCallbackConfig",
     "MixedPrecisionPluginConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",

{nshtrainer-1.0.0b42 → nshtrainer-1.0.0b44}/src/nshtrainer/configs/trainer/_config/__init__.py

@@ -34,6 +34,9 @@ from nshtrainer.trainer._config import LogEpochCallbackConfig as LogEpochCallbac
 from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
 from nshtrainer.trainer._config import LoggerConfigBase as LoggerConfigBase
 from nshtrainer.trainer._config import MetricConfig as MetricConfig
+from nshtrainer.trainer._config import (
+    MetricValidationCallbackConfig as MetricValidationCallbackConfig,
+)
 from nshtrainer.trainer._config import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
 )
@@ -77,6 +80,7 @@ __all__ = [
     "LoggerConfig",
     "LoggerConfigBase",
     "MetricConfig",
+    "MetricValidationCallbackConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
     "PluginConfig",

{nshtrainer-1.0.0b42 → nshtrainer-1.0.0b44}/src/nshtrainer/trainer/_config.py

@@ -40,6 +40,7 @@ from ..callbacks.base import CallbackConfigBase
 from ..callbacks.debug_flag import DebugFlagCallbackConfig
 from ..callbacks.log_epoch import LogEpochCallbackConfig
 from ..callbacks.lr_monitor import LearningRateMonitorConfig
+from ..callbacks.metric_validation import MetricValidationCallbackConfig
 from ..callbacks.rlp_sanity_checks import RLPSanityChecksCallbackConfig
 from ..callbacks.shared_parameters import SharedParametersCallbackConfig
 from ..loggers import (
@@ -697,6 +698,10 @@ class TrainerConfig(C.Config):
     - The trainer is running in fast_dev_run mode.
     - The trainer is running a sanity check (which happens before starting the training routine).
     """
+    auto_validate_metrics: MetricValidationCallbackConfig | None = (
+        MetricValidationCallbackConfig()
+    )
+    """If enabled, will automatically validate the metrics before starting the training routine."""
 
     lightning_kwargs: LightningTrainerKwargs = LightningTrainerKwargs()
     """
@@ -768,6 +773,7 @@ class TrainerConfig(C.Config):
         yield self.shared_parameters
         yield self.reduce_lr_on_plateau_sanity_checking
         yield self.auto_set_debug_flag
+        yield self.auto_validate_metrics
         yield from self.callbacks
 
     def _nshtrainer_all_logger_configs(self) -> Iterable[LoggerConfigBase | None]:

{nshtrainer-1.0.0b42 → nshtrainer-1.0.0b44}/src/nshtrainer/trainer/trainer.py

@@ -18,7 +18,7 @@ from lightning.pytorch.trainer.states import TrainerFn
 from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT, _PREDICT_OUTPUT
 from typing_extensions import Never, Unpack, assert_never, deprecated, override
 
-from .._checkpoint.metadata import _write_checkpoint_metadata
+from .._checkpoint.metadata import write_checkpoint_metadata
 from ..callbacks.base import resolve_all_callbacks
 from ..util._environment_info import EnvironmentConfig
 from ..util.bf16 import is_bf16_supported_no_emulation
@@ -478,7 +478,7 @@ class Trainer(LightningTrainer):
         metadata_path = None
         if self.hparams.save_checkpoint_metadata and self.is_global_zero:
             # Generate the metadata and write to disk
-            metadata_path = _write_checkpoint_metadata(self, filepath)
+            metadata_path = write_checkpoint_metadata(self, filepath)
 
         # Call the `on_checkpoint_saved` method on all callbacks
         from .. import _callback

{nshtrainer-1.0.0b42 → nshtrainer-1.0.0b44}/src/nshtrainer/util/path.py

@@ -85,11 +85,16 @@ def try_symlink_or_copy(
     target_is_directory: bool = False,
     relative: bool = True,
     remove_existing: bool = True,
+    throw_on_invalid_target: bool = False,
 ):
     """
     Symlinks on Unix, copies on Windows.
     """
 
+    # Check if the target file exists
+    if throw_on_invalid_target and not file_path.exists():
+        raise FileNotFoundError(f"File not found: {file_path}")
+
     # If the link already exists, remove it
     if remove_existing:
         try:

nshtrainer-1.0.0b42/src/nshtrainer/nn/tests/test_mlp.py

@@ -1,120 +0,0 @@
-from __future__ import annotations
-
-from typing import cast
-
-import pytest
-import torch
-
-from nshtrainer.nn.mlp import MLP, custom_seed_context
-
-
-def test_mlp_seed_reproducibility():
-    """Test that the seed parameter in MLP ensures reproducible weights."""
-
-    # Test dimensions
-    dims = [10, 20, 5]
-
-    # Create two MLPs with the same seed
-    seed1 = 42
-    mlp1 = MLP(dims, activation=torch.nn.ReLU(), seed=seed1)
-    mlp2 = MLP(dims, activation=torch.nn.ReLU(), seed=seed1)
-
-    # Create an MLP with a different seed
-    seed2 = 123
-    mlp3 = MLP(dims, activation=torch.nn.ReLU(), seed=seed2)
-
-    # Check first layer weights
-    layer1_weights1 = cast(torch.Tensor, mlp1[0].weight)
-    layer1_weights2 = cast(torch.Tensor, mlp2[0].weight)
-    layer1_weights3 = cast(torch.Tensor, mlp3[0].weight)
-
-    # Same seed should produce identical weights
-    assert torch.allclose(layer1_weights1, layer1_weights2)
-
-    # Different seeds should produce different weights
-    assert not torch.allclose(layer1_weights1, layer1_weights3)
-
-    # Check second layer weights
-    layer2_weights1 = cast(torch.Tensor, mlp1[2].weight)
-    layer2_weights2 = cast(torch.Tensor, mlp2[2].weight)
-    layer2_weights3 = cast(torch.Tensor, mlp3[2].weight)
-
-    # Same seed should produce identical weights for all layers
-    assert torch.allclose(layer2_weights1, layer2_weights2)
-
-    # Different seeds should produce different weights for all layers
-    assert not torch.allclose(layer2_weights1, layer2_weights3)
-
-    # Test that not providing a seed gives different results each time
-    mlp4 = MLP(dims, activation=torch.nn.ReLU(), seed=None)
-    mlp5 = MLP(dims, activation=torch.nn.ReLU(), seed=None)
-
-    # Without seeds, weights should be different
-    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)