nshtrainer 0.10.0__py3-none-any.whl → 0.10.2__py3-none-any.whl

nshtrainer/{trainer/_checkpoint_resolver.py → _checkpoint/loader.py}

@@ -10,7 +10,7 @@ from lightning.pytorch.trainer.states import TrainerFn
 from typing_extensions import assert_never
 
 from ..metrics._config import MetricConfig
-from ._checkpoint_metadata import METADATA_PATH_SUFFIX, CheckpointMetadata
+from .metadata import METADATA_PATH_SUFFIX, CheckpointMetadata
 
 if TYPE_CHECKING:
     from ..model.config import BaseConfig

nshtrainer/{trainer/_checkpoint_metadata.py → _checkpoint/metadata.py}

@@ -12,7 +12,7 @@ from ..model._environment import EnvironmentConfig
 
 if TYPE_CHECKING:
     from ..model import BaseConfig, LightningModuleBase
-    from .trainer import Trainer
+    from ..trainer.trainer import Trainer
 
 log = logging.getLogger(__name__)
 

nshtrainer/callbacks/latest_epoch_checkpoint.py

@@ -12,7 +12,7 @@ log = logging.getLogger(__name__)
 
 
 class LatestEpochCheckpointCallbackConfig(CallbackConfigBase):
-    kind: Literal["latest_epoch_checkpoint"] = "latest_epoch_checkpoint"
+    name: Literal["latest_epoch_checkpoint"] = "latest_epoch_checkpoint"
 
     dirpath: str | Path | None = None
     """Directory path to save the checkpoint file."""

nshtrainer/callbacks/model_checkpoint.py

@@ -43,7 +43,7 @@ def _convert_string(input_string: str):
 class ModelCheckpointCallbackConfig(CallbackConfigBase):
     """Arguments for the ModelCheckpoint callback."""
 
-    kind: Literal["model_checkpoint"] = "model_checkpoint"
+    name: Literal["model_checkpoint"] = "model_checkpoint"
 
     dirpath: str | Path | None = None
     """

nshtrainer/callbacks/on_exception_checkpoint.py

@@ -43,7 +43,7 @@ def _monkey_patch_disable_barrier(trainer: LightningTrainer):
 
 
 class OnExceptionCheckpointCallbackConfig(CallbackConfigBase):
-    kind: Literal["on_exception_checkpoint"] = "on_exception_checkpoint"
+    name: Literal["on_exception_checkpoint"] = "on_exception_checkpoint"
 
     dirpath: str | Path | None = None
     """Directory path to save the checkpoint file."""

nshtrainer/model/config.py

@@ -34,6 +34,7 @@ from lightning.pytorch.strategies.strategy import Strategy
 from pydantic import DirectoryPath
 from typing_extensions import Self, TypedDict, TypeVar, override
 
+from .._checkpoint.loader import CheckpointLoadingConfig
 from ..callbacks import (
     CallbackConfig,
     LatestEpochCheckpointCallbackConfig,
@@ -43,7 +44,6 @@ from ..callbacks import (
 )
 from ..callbacks.base import CallbackConfigBase
 from ..metrics import MetricConfig
-from ..trainer._checkpoint_resolver import CheckpointLoadingConfig
 from ._environment import EnvironmentConfig
 
 log = getLogger(__name__)
@@ -71,7 +71,7 @@ class BaseProfilerConfig(C.Config, ABC):
 
 
 class SimpleProfilerConfig(BaseProfilerConfig):
-    kind: Literal["simple"] = "simple"
+    name: Literal["simple"] = "simple"
 
     extended: bool = True
     """
@@ -99,7 +99,7 @@ class SimpleProfilerConfig(BaseProfilerConfig):
 
 
 class AdvancedProfilerConfig(BaseProfilerConfig):
-    kind: Literal["advanced"] = "advanced"
+    name: Literal["advanced"] = "advanced"
 
     line_count_restriction: float = 1.0
     """
@@ -128,7 +128,7 @@ class AdvancedProfilerConfig(BaseProfilerConfig):
 
 
 class PyTorchProfilerConfig(BaseProfilerConfig):
-    kind: Literal["pytorch"] = "pytorch"
+    name: Literal["pytorch"] = "pytorch"
 
     group_by_input_shapes: bool = False
     """Include operator input shapes and group calls by shape."""
@@ -204,7 +204,7 @@ class PyTorchProfilerConfig(BaseProfilerConfig):
 
 ProfilerConfig: TypeAlias = Annotated[
     SimpleProfilerConfig | AdvancedProfilerConfig | PyTorchProfilerConfig,
-    C.Field(discriminator="kind"),
+    C.Field(discriminator="name"),
 ]
 
 
@@ -260,7 +260,7 @@ def _wandb_available():
 
 
 class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
-    kind: Literal["wandb"] = "wandb"
+    name: Literal["wandb"] = "wandb"
 
     enabled: bool = C.Field(default_factory=lambda: _wandb_available())
     """Enable WandB logging."""
@@ -319,7 +319,7 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
 
 
 class CSVLoggerConfig(BaseLoggerConfig):
-    kind: Literal["csv"] = "csv"
+    name: Literal["csv"] = "csv"
 
     enabled: bool = True
     """Enable CSV logging."""
@@ -373,7 +373,7 @@ def _tensorboard_available():
 
 
 class TensorboardLoggerConfig(BaseLoggerConfig):
-    kind: Literal["tensorboard"] = "tensorboard"
+    name: Literal["tensorboard"] = "tensorboard"
 
     enabled: bool = C.Field(default_factory=lambda: _tensorboard_available())
     """Enable TensorBoard logging."""
@@ -419,7 +419,7 @@ class TensorboardLoggerConfig(BaseLoggerConfig):
 
 LoggerConfig: TypeAlias = Annotated[
     WandbLoggerConfig | CSVLoggerConfig | TensorboardLoggerConfig,
-    C.Field(discriminator="kind"),
+    C.Field(discriminator="name"),
 ]
 
 
@@ -717,9 +717,9 @@ class DirectoryConfig(C.Config):
         if (log_dir := logger.log_dir) is not None:
             return log_dir
 
-        # Save to nshtrainer/{id}/log/{logger kind}
+        # Save to nshtrainer/{id}/log/{logger name}
         log_dir = self.resolve_subdirectory(run_id, "log")
-        log_dir = log_dir / logger.kind
+        log_dir = log_dir / logger.name
         log_dir.mkdir(exist_ok=True)
 
         return log_dir
@@ -738,7 +738,7 @@ CheckpointCallbackConfig: TypeAlias = Annotated[
     ModelCheckpointCallbackConfig
     | LatestEpochCheckpointCallbackConfig
     | OnExceptionCheckpointCallbackConfig,
-    C.Field(discriminator="kind"),
+    C.Field(discriminator="name"),
 ]
 
 

nshtrainer/trainer/checkpoint_connector.py

@@ -8,7 +8,7 @@ from lightning.pytorch.trainer.connectors.checkpoint_connector import (
 from lightning.pytorch.trainer.states import TrainerFn
 from typing_extensions import override
 
-from ._checkpoint_resolver import CheckpointLoadingConfig, _resolve_checkpoint
+from .._checkpoint.loader import CheckpointLoadingConfig, _resolve_checkpoint
 
 if TYPE_CHECKING:
     from ..model.config import BaseConfig

nshtrainer/trainer/trainer.py

@@ -17,6 +17,7 @@ from lightning.pytorch.trainer.states import TrainerFn
 from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT, _PREDICT_OUTPUT
 from typing_extensions import Unpack, assert_never, override
 
+from .._checkpoint.metadata import _write_checkpoint_metadata
 from ..callbacks.base import resolve_all_callbacks
 from ..model.config import (
     AcceleratorConfigProtocol,
@@ -25,7 +26,6 @@ from ..model.config import (
     LightningTrainerKwargs,
     StrategyConfigProtocol,
 )
-from ._checkpoint_metadata import _write_checkpoint_metadata
 from ._runtime_callback import RuntimeTrackerCallback, Stage
 from .signal_connector import _SignalConnector
 

{nshtrainer-0.10.0.dist-info → nshtrainer-0.10.2.dist-info}/METADATA

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

{nshtrainer-0.10.0.dist-info → nshtrainer-0.10.2.dist-info}/RECORD

@@ -1,4 +1,6 @@
 nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
+nshtrainer/_checkpoint/loader.py,sha256=48flPr1XgQHOgIPaCrRqOEvRuG0SZuV3cQ1vgHLqFqI,11025
+nshtrainer/_checkpoint/metadata.py,sha256=C7je_soYyEbZjiq7p2_pSVFkgcXnz2J2H5sMy8oskx0,3051
 nshtrainer/_experimental/__init__.py,sha256=2tQIcrWT8U8no_AeBTYnozaTmxN40kuAJdGQ4b-PoWM,120
 nshtrainer/_experimental/flops/__init__.py,sha256=edo9Ez3LlrnxkNRX9W6YBhPkRPKYGLpkpnl5gx7sEX8,1550
 nshtrainer/_experimental/flops/flop_counter.py,sha256=-sL0Fy6poXa__hyzUMdZScjPULp4coQELQpPU6p6dXU,25736
@@ -12,11 +14,11 @@ nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,1
 nshtrainer/callbacks/finite_checks.py,sha256=AO5fa51uANAjAkeJfTquOjK6W_4RSU5Kky3f5jmAPlQ,2084
 nshtrainer/callbacks/gradient_skipping.py,sha256=fSJpjgHbztFKz7w3qFuCHZpmbEt9BCLAy-sU0B4xJQI,3474
 nshtrainer/callbacks/interval.py,sha256=smz5Zl8cN6X6yHKVsMRS2e3SEkzRCP3LvwE1ONvLfaw,8080
-nshtrainer/callbacks/latest_epoch_checkpoint.py,sha256=p0zeDK3PLWWl485e9o08ywEEARCfuZ5it47tNCtR4ec,2838
+nshtrainer/callbacks/latest_epoch_checkpoint.py,sha256=UnwgGIc2reD7cTnUeIlDHo1LeAkgLEZFNvy2NGvUfRQ,2838
 nshtrainer/callbacks/log_epoch.py,sha256=fTa_K_Y8A7g09630cG4YkDE6AzSMPkjb9bpPm4gtqos,1120
-nshtrainer/callbacks/model_checkpoint.py,sha256=4zYycpXHGRyL4svWLP6GmG3WJs5m3B5PRCOzXC3m_qg,5955
+nshtrainer/callbacks/model_checkpoint.py,sha256=N0raLsHlCVSbO3QU5eNFUXUDqxxW3C73oQwceMnFE_k,5955
 nshtrainer/callbacks/norm_logging.py,sha256=EWyrfkp8iHjQi9iAAXHxb0xStw2RwkdpKG2_gLarQRA,6281
-nshtrainer/callbacks/on_exception_checkpoint.py,sha256=zna_QF_x4HwD7Es5XxrHLDED43NU1GpcDNoL139HEOs,3355
+nshtrainer/callbacks/on_exception_checkpoint.py,sha256=x42BYZ2ejf2rhqPLCmT5nyWKhA9qBEosiV8ZNhhZ6lI,3355
 nshtrainer/callbacks/print_table.py,sha256=_FdAHhqylWGk4Z0c2FrLFeiMA4jhfA_beZRK_BHpzmE,2837
 nshtrainer/callbacks/throughput_monitor.py,sha256=4EF3b79HdHiRgBGIFDyD4O1oywb5h1tV8nml7NuuDjU,1845
 nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50,4607
@@ -50,7 +52,7 @@ nshtrainer/metrics/_config.py,sha256=hWWS4IXENRyH3RmJ7z1Wx1n3Lt1sNMlGOrcU6PW15o0
 nshtrainer/model/__init__.py,sha256=TbexTxiE20WHYg5q3L88Hysk4LlHeKk_isv33aSBREA,1918
 nshtrainer/model/_environment.py,sha256=s3JFnigbssFRJTwH33K7DcAYVhLOFCC1OZgFNXJgjuw,22317
 nshtrainer/model/base.py,sha256=Bmw-t70TydDbE9P0ee-lTibGoUhrCx5Qke-upa7FGVM,17512
-nshtrainer/model/config.py,sha256=f8gbTaIi02U8EyooC1vv2ElZfXPgMIAVtU0n-LnkNE4,53187
+nshtrainer/model/config.py,sha256=OsVba02cmEYVf6V-A6ljV7VMAW5XZO6GWNRk8ktUw2o,53177
 nshtrainer/model/modules/callback.py,sha256=JF59U9-CjJsAIspEhTJbVaGN0wGctZG7UquE3IS7R8A,6408
 nshtrainer/model/modules/debug.py,sha256=DTVty8cKnzj1GCULRyGx_sWTTsq9NLi30dzqjRTnuCU,1127
 nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
@@ -68,17 +70,15 @@ nshtrainer/runner.py,sha256=6qfE5FBONzD79kVHuWYKEvK0J_Qi5dMBbHQhRMmnIhE,3649
 nshtrainer/scripts/check_env.py,sha256=IMl6dSqsLYppI0XuCsVq8lK4bYqXwY9KHJkzsShz4Kg,806
 nshtrainer/scripts/find_packages.py,sha256=FbdlfmAefttFSMfaT0A46a-oHLP_ioaQKihwBfBeWeA,1467
 nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3E,40
-nshtrainer/trainer/_checkpoint_metadata.py,sha256=dj3g0rUZLWfohIRFAhhLqB4qh1fJsquQ5-EZ0Zbl5ZE,3042
-nshtrainer/trainer/_checkpoint_resolver.py,sha256=kfIccBLWAMwn-Bw1pbj3XTXXaCdO_taUEUp3RdwFuLY,11037
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
-nshtrainer/trainer/checkpoint_connector.py,sha256=9DrZliK95BIZfwVFxL06Uf7DbfbQ5UAWd0xckH-LU6U,2125
+nshtrainer/trainer/checkpoint_connector.py,sha256=xoqI2dcPnlNFPPLVIU6dBOvRPC9PtfX5qu__xV1lx0Y,2124
 nshtrainer/trainer/signal_connector.py,sha256=llwc8pdKAWxREFpjdi14Bpy8rGVMEJsmJx_s2p4gI8E,10689
-nshtrainer/trainer/trainer.py,sha256=qHwerfdQUCU21IkWf50d_qZAIHb2d8qOLfqTszBpzks,16784
+nshtrainer/trainer/trainer.py,sha256=n3T9Iz3eaDostxEdjapWImAsVMxyU9WBdhlPl0THX-g,16785
 nshtrainer/util/environment.py,sha256=_SEtiQ_s5bL5pllUlf96AOUv15kNvCPvocVC13S7mIk,4166
 nshtrainer/util/seed.py,sha256=HEXgVs-wldByahOysKwq7506OHxdYTEgmP-tDQVAEkQ,287
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.10.0.dist-info/METADATA,sha256=GslAMAaEXDbMxDd4ijoqjQKYBjb0iAnEGkZ3pAF_sOQ,695
-nshtrainer-0.10.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.10.0.dist-info/RECORD,,
+nshtrainer-0.10.2.dist-info/METADATA,sha256=bLP7xa9qV9BqgqGZr2RayfaL2eFB9eBSTHVlMpslm5s,695
+nshtrainer-0.10.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.10.2.dist-info/RECORD,,