nshtrainer 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

nshtrainer/model/config.py

@@ -1,7 +1,6 @@
 import copy
 import os
 import re
-import signal
 import socket
 import string
 import time
@@ -21,6 +20,7 @@ from typing import (
     runtime_checkable,
 )
 
+import nshconfig as C
 import numpy as np
 import torch
 from lightning.fabric.plugins import CheckpointIO, ClusterEnvironment
@@ -39,7 +39,6 @@ from typing_extensions import Self, TypedDict, TypeVar, override
 from ..callbacks import CallbackConfig
 from ..callbacks.base import CallbackConfigBase
 from ..callbacks.wandb_watch import WandbWatchConfig
-from ..config import Field, TypedConfig
 from ..util.slurm import parse_slurm_node_list
 
 log = getLogger(__name__)
@@ -49,7 +48,7 @@ class IdSeedWarning(Warning):
     pass
 
 
-class BaseProfilerConfig(TypedConfig, ABC):
+class BaseProfilerConfig(C.Config, ABC):
     dirpath: str | Path | None = None
     """
     Directory path for the ``filename``. If ``dirpath`` is ``None`` but ``filename`` is present, the
@@ -200,11 +199,11 @@ class PyTorchProfilerConfig(BaseProfilerConfig):
 
 ProfilerConfig: TypeAlias = Annotated[
     SimpleProfilerConfig | AdvancedProfilerConfig | PyTorchProfilerConfig,
-    Field(discriminator="kind"),
+    C.Field(discriminator="kind"),
 ]
 
 
-class EnvironmentClassInformationConfig(TypedConfig):
+class EnvironmentClassInformationConfig(C.Config):
     name: str
     module: str
     full_name: str
@@ -213,7 +212,7 @@ class EnvironmentClassInformationConfig(TypedConfig):
     source_file_path: Path | None = None
 
 
-class EnvironmentSLURMInformationConfig(TypedConfig):
+class EnvironmentSLURMInformationConfig(C.Config):
     hostname: str
     hostnames: list[str]
     job_id: str
@@ -271,7 +270,7 @@ class EnvironmentSLURMInformationConfig(TypedConfig):
             return None
 
 
-class EnvironmentLSFInformationConfig(TypedConfig):
+class EnvironmentLSFInformationConfig(C.Config):
     hostname: str
     hostnames: list[str]
     job_id: str
@@ -328,7 +327,7 @@ class EnvironmentLSFInformationConfig(TypedConfig):
             return None
 
 
-class EnvironmentLinuxEnvironmentConfig(TypedConfig):
+class EnvironmentLinuxEnvironmentConfig(C.Config):
     """
     Information about the Linux environment (e.g., current user, hostname, etc.)
     """
@@ -347,9 +346,25 @@ class EnvironmentLinuxEnvironmentConfig(TypedConfig):
     load_avg: tuple[float, float, float] | None = None
 
 
-class EnvironmentConfig(TypedConfig):
+class EnvironmentSnapshotConfig(C.Config):
+    snapshot_dir: Path | None = None
+    modules: list[str] | None = None
+
+    @classmethod
+    def from_current_environment(cls):
+        draft = cls.draft()
+        if snapshot_dir := os.environ.get("NSHRUNNER_SNAPSHOT_DIR"):
+            draft.snapshot_dir = Path(snapshot_dir)
+        if modules := os.environ.get("NSHRUNNER_SNAPSHOT_MODULES"):
+            draft.modules = modules.split(",")
+        return draft.finalize()
+
+
+class EnvironmentConfig(C.Config):
     cwd: Path | None = None
 
+    snapshot: EnvironmentSnapshotConfig | None = None
+
     python_executable: Path | None = None
     python_path: list[Path] | None = None
     python_version: str | None = None
@@ -372,7 +387,7 @@ class EnvironmentConfig(TypedConfig):
     seed_workers: bool | None = None
 
 
-class BaseLoggerConfig(TypedConfig, ABC):
+class BaseLoggerConfig(C.Config, ABC):
     enabled: bool = True
     """Enable this logger."""
 
@@ -426,7 +441,7 @@ def _wandb_available():
 class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
     kind: Literal["wandb"] = "wandb"
 
-    enabled: bool = Field(default_factory=lambda: _wandb_available())
+    enabled: bool = C.Field(default_factory=lambda: _wandb_available())
     """Enable WandB logging."""
 
     priority: int = 2
@@ -543,7 +558,7 @@ def _tensorboard_available():
 class TensorboardLoggerConfig(BaseLoggerConfig):
     kind: Literal["tensorboard"] = "tensorboard"
 
-    enabled: bool = Field(default_factory=lambda: _tensorboard_available())
+    enabled: bool = C.Field(default_factory=lambda: _tensorboard_available())
     """Enable TensorBoard logging."""
 
     priority: int = 2
@@ -589,7 +604,7 @@ class TensorboardLoggerConfig(BaseLoggerConfig):
 
 LoggerConfig: TypeAlias = Annotated[
     WandbLoggerConfig | CSVLoggerConfig | TensorboardLoggerConfig,
-    Field(discriminator="kind"),
+    C.Field(discriminator="kind"),
 ]
 
 
@@ -684,7 +699,7 @@ class LoggingConfig(CallbackConfigBase):
             yield from logger.construct_callbacks(root_config)
 
 
-class GradientClippingConfig(TypedConfig):
+class GradientClippingConfig(C.Config):
     enabled: bool = True
     """Enable gradient clipping."""
     value: int | float
@@ -719,41 +734,6 @@ class OptimizationConfig(CallbackConfigBase):
         ).construct_callbacks(root_config)
 
 
-LogLevel: TypeAlias = Literal[
-    "CRITICAL", "FATAL", "ERROR", "WARN", "WARNING", "INFO", "DEBUG"
-]
-
-
-class PythonLogging(TypedConfig):
-    log_level: LogLevel | None = None
-    """Log level to use for the Python logger (or None to use the default)."""
-
-    rich: bool = False
-    """If enabled, will use the rich library to format the Python logger output."""
-    rich_tracebacks: bool = True
-    """If enabled, will use the rich library to format the Python logger tracebacks."""
-
-    lovely_tensors: bool = False
-    """If enabled, will use the lovely-tensors library to format PyTorch tensors. False by default as it causes issues when used with `torch.vmap`."""
-    lovely_numpy: bool = False
-    """If enabled, will use the lovely-numpy library to format numpy arrays. False by default as it causes some issues with other libaries."""
-
-    def pretty_(
-        self,
-        *,
-        log_level: LogLevel | None = "INFO",
-        torch: bool = True,
-        numpy: bool = True,
-        rich: bool = True,
-        rich_tracebacks: bool = True,
-    ):
-        self.log_level = log_level
-        self.lovely_tensors = torch
-        self.lovely_numpy = numpy
-        self.rich = rich
-        self.rich_tracebacks = rich_tracebacks
-
-
 TPlugin = TypeVar(
     "TPlugin",
     Precision,
@@ -813,7 +793,7 @@ StrategyLiteral: TypeAlias = Literal[
 ]
 
 
-class CheckpointLoadingConfig(TypedConfig):
+class CheckpointLoadingConfig(C.Config):
     path: Literal["best", "last", "hpc"] | str | Path | None = None
     """
     Checkpoint path to use when loading a checkpoint.
@@ -825,7 +805,7 @@ class CheckpointLoadingConfig(TypedConfig):
     """
 
 
-class DirectoryConfig(TypedConfig):
+class DirectoryConfig(C.Config):
     project_root: Path | None = None
     """
     Root directory for this project.
@@ -901,7 +881,7 @@ class DirectoryConfig(TypedConfig):
         return log_dir
 
 
-class ReproducibilityConfig(TypedConfig):
+class ReproducibilityConfig(C.Config):
     deterministic: bool | Literal["warn"] | None = None
     """
     If ``True``, sets whether PyTorch operations must use deterministic algorithms.
@@ -1116,7 +1096,7 @@ CheckpointCallbackConfig: TypeAlias = Annotated[
     ModelCheckpointCallbackConfig
     | LatestEpochCheckpointCallbackConfig
     | OnExceptionCheckpointCallbackConfig,
-    Field(discriminator="kind"),
+    C.Field(discriminator="kind"),
 ]
 
 
@@ -1514,7 +1494,7 @@ class ActSaveConfig(CallbackConfigBase):
         return [ActSaveCallback()]
 
 
-class SanityCheckingConfig(TypedConfig):
+class SanityCheckingConfig(C.Config):
     reduce_lr_on_plateau: Literal["disable", "error", "warn"] = "error"
     """
     If enabled, will do some sanity checks if the `ReduceLROnPlateau` scheduler is used:
@@ -1524,7 +1504,7 @@ class SanityCheckingConfig(TypedConfig):
     """
 
 
-class TrainerConfig(TypedConfig):
+class TrainerConfig(C.Config):
     checkpoint_loading: CheckpointLoadingConfig = CheckpointLoadingConfig()
     """Checkpoint loading configuration options."""
 
@@ -1739,87 +1719,7 @@ class TrainerConfig(TypedConfig):
     """If enabled, will set the torch float32 matmul precision to the specified value. Useful for faster training on Ampere+ GPUs."""
 
 
-class SeedConfig(TypedConfig):
-    seed: int
-    """Seed for the random number generator."""
-
-    seed_workers: bool = False
-    """Whether to seed the workers of the dataloader."""
-
-
-Signal: TypeAlias = Literal[
-    "SIGHUP",
-    "SIGINT",
-    "SIGQUIT",
-    "SIGILL",
-    "SIGTRAP",
-    "SIGABRT",
-    "SIGBUS",
-    "SIGFPE",
-    "SIGKILL",
-    "SIGUSR1",
-    "SIGSEGV",
-    "SIGUSR2",
-    "SIGPIPE",
-    "SIGALRM",
-    "SIGTERM",
-    "SIGCHLD",
-    "SIGCONT",
-    "SIGSTOP",
-    "SIGTSTP",
-    "SIGTTIN",
-    "SIGTTOU",
-    "SIGURG",
-    "SIGXCPU",
-    "SIGXFSZ",
-    "SIGVTALRM",
-    "SIGPROF",
-    "SIGWINCH",
-    "SIGIO",
-    "SIGPWR",
-    "SIGSYS",
-    "SIGRTMIN",
-    "SIGRTMAX",
-]
-
-
-class SubmitConfig(TypedConfig):
-    auto_requeue_signals: list[Signal] = [
-        # "SIGUSR1",
-        # On SIGURG:
-        # Important note from https://amrex-astro.github.io/workflow/olcf-workflow.html:
-        # We can also ask the job manager to send a warning signal some amount of time before the allocation expires by passing -wa 'signal' and -wt '[hour:]minute' to bsub. We can then have bash create a dump_and_stop file when it receives the signal, which will tell Castro to output a checkpoint file and exit cleanly after it finishes the current timestep. An important detail that I couldn't find documented anywhere is that the job manager sends the signal to all the processes in the job, not just the submission script, and we have to use a signal that is ignored by default so Castro doesn't immediately crash upon receiving it. SIGCHLD, SIGURG, and SIGWINCH are the only signals that fit this requirement and of these, SIGURG is the least likely to be triggered by other events.
-        "SIGURG"
-    ]
-    """Signals that will trigger an automatic requeue of the job."""
-
-    def _resolved_auto_requeue_signals(self) -> list[signal.Signals]:
-        return [getattr(signal.Signals, sig) for sig in self.auto_requeue_signals]
-
-
-class RunnerConfig(TypedConfig):
-    python_logging: PythonLogging = PythonLogging()
-    """Python logging configuration options."""
-
-    seed: SeedConfig = SeedConfig(seed=0)
-    """Seed everything configuration options."""
-
-    submit: SubmitConfig = SubmitConfig()
-    """Submit (e.g., SLURM or LSF) configuration options."""
-
-    dump_run_information: bool = True
-    """
-    If enabled, will dump different bits of run information to the output directory before starting the run.
-    This includes:
-        - Run config
-        - Full set of environment variables
-    """
-
-    additional_env_vars: dict[str, str] = {}
-    """Additional environment variables to set when running the script."""
-
-
-class MetricConfig(TypedConfig):
+class MetricConfig(C.Config):
     name: str
     """The name of the primary metric."""
 
@@ -1851,8 +1751,8 @@ class MetricConfig(TypedConfig):
 PrimaryMetricConfig: TypeAlias = MetricConfig
 
 
-class BaseConfig(TypedConfig):
-    id: str = Field(default_factory=lambda: BaseConfig.generate_id())
+class BaseConfig(C.Config):
+    id: str = C.Field(default_factory=lambda: BaseConfig.generate_id())
     """ID of the run."""
     name: str | None = None
     """Run name."""
@@ -1867,15 +1767,13 @@ class BaseConfig(TypedConfig):
 
     debug: bool = False
     """Whether to run in debug mode. This will enable debug logging and enable debug code paths."""
-    environment: Annotated[EnvironmentConfig, Field(repr=False)] = EnvironmentConfig()
+    environment: Annotated[EnvironmentConfig, C.Field(repr=False)] = EnvironmentConfig()
     """A snapshot of the current environment information (e.g. python version, slurm info, etc.). This is automatically populated by the run script."""
 
     directory: DirectoryConfig = DirectoryConfig()
     """Directory configuration options."""
     trainer: TrainerConfig = TrainerConfig()
     """PyTorch Lightning trainer configuration options. Check Lightning's `Trainer` documentation for more information."""
-    runner: RunnerConfig = RunnerConfig()
-    """`ll.Runner` configuration options."""
 
     primary_metric: PrimaryMetricConfig | None = None
     """Primary metric configuration options. This is used in the following ways:

nshtrainer/nn/nonlinearity.py

@@ -1,14 +1,13 @@
 from abc import ABC, abstractmethod
 from typing import Annotated, Literal
 
+import nshconfig as C
 import torch
 import torch.nn as nn
 from typing_extensions import override
 
-from ..config import Field, TypedConfig
 
-
-class BaseNonlinearityConfig(TypedConfig, ABC):
+class BaseNonlinearityConfig(C.Config, ABC):
     @abstractmethod
     def create_module(self) -> nn.Module:
         pass
@@ -153,5 +152,5 @@ NonlinearityConfig = Annotated[
     | SiLUNonlinearityConfig
     | MishNonlinearityConfig
     | SwiGLUNonlinearityConfig,
-    Field(discriminator="name"),
+    C.Field(discriminator="name"),
 ]

nshtrainer/optimizer.py

@@ -2,14 +2,13 @@ from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from typing import Annotated, Any, Literal, TypeAlias
 
+import nshconfig as C
 import torch.nn as nn
 from torch.optim import Optimizer
 from typing_extensions import override
 
-from .config import Field, TypedConfig
 
-
-class OptimizerConfigBase(TypedConfig, ABC):
+class OptimizerConfigBase(C.Config, ABC):
     @abstractmethod
     def create_optimizer(
         self,
@@ -56,7 +55,4 @@ class AdamWConfig(OptimizerConfigBase):
         )
 
 
-OptimizerConfig: TypeAlias = Annotated[
-    AdamWConfig,
-    Field(discriminator="name"),
-]
+OptimizerConfig: TypeAlias = Annotated[AdamWConfig, C.Field(discriminator="name")]

nshtrainer/runner.py

@@ -1,21 +1,31 @@
-from dataclasses import dataclass
 from typing import Generic
 
+from nshrunner import RunInfo
 from nshrunner import Runner as _Runner
-from typing_extensions import Concatenate, TypeVar, TypeVarTuple, Unpack, override
+from typing_extensions import TypeVar, TypeVarTuple, Unpack, override
 
 from .model.config import BaseConfig
 
 TConfig = TypeVar("TConfig", bound=BaseConfig, infer_variance=True)
-TArguments = TypeVarTuple("TArguments")
+TArguments = TypeVarTuple("TArguments", default=Unpack[tuple[()]])
 TReturn = TypeVar("TReturn", infer_variance=True)
 
 
-@dataclass(frozen=True)
 class Runner(
-    _Runner[Unpack[tuple[TConfig, Unpack[TArguments]]], TReturn],
-    Generic[TConfig, Unpack[TArguments], TReturn],
+    _Runner[TReturn, TConfig, Unpack[TArguments]],
+    Generic[TReturn, TConfig, Unpack[TArguments]],
 ):
     @override
-    def default_validate_fn():
-        pass
+    @classmethod
+    def default_validate_fn(cls, config: TConfig, *args: Unpack[TArguments]) -> None:
+        super().default_validate_fn(config, *args)
+
+    @override
+    @classmethod
+    def default_info_fn(cls, config: TConfig, *args: Unpack[TArguments]) -> RunInfo:
+        run_info = super().default_info_fn(config, *args)
+        return {
+            **run_info,
+            "id": config.id,
+            "base_dir": config.directory.project_root,
+        }

nshtrainer/trainer/signal_connector.py

@@ -25,14 +25,21 @@ _SIGNUM = int | signal.Signals
 _HANDLER: TypeAlias = Callable[[_SIGNUM, FrameType], Any] | int | signal.Handlers | None
 
 
-class _SignalConnector(_LightningSignalConnector):
-    def _auto_requeue_signals(self) -> list[signal.Signals]:
-        from ..model.base import BaseConfig
+def _resolve_requeue_signals():
+    signals: list[signal.Signals] = []
+
+    if timeout_signal_name := os.environ.get("NSHRUNNER_TIMEOUT_SIGNAL"):
+        signals.append(signal.Signals[timeout_signal_name])
 
-        if not isinstance(config := self.trainer.lightning_module.hparams, BaseConfig):
-            return []
+    if preempt_signal_name := os.environ.get("NSHRUNNER_PREEMPT_SIGNAL"):
+        signals.append(signal.Signals[preempt_signal_name])
 
-        signals = config.runner.submit._resolved_auto_requeue_signals()
+    return signals
+
+
+class _SignalConnector(_LightningSignalConnector):
+    def _auto_requeue_signals(self) -> list[signal.Signals]:
+        signals = _resolve_requeue_signals()
         signals_set = set(signals)
         valid_signals: set[signal.Signals] = signal.valid_signals()
         assert signals_set.issubset(
@@ -42,25 +49,29 @@ class _SignalConnector(_LightningSignalConnector):
 
     def _compose_and_register(
         self,
-        signum: _SIGNUM,
+        signum: signal.Signals,
         handlers: list[_HANDLER],
         replace_existing: bool = False,
     ):
         if self._is_on_windows():
-            log.info(f"Signal {signum} has no handlers or is not supported on Windows.")
+            log.info(
+                f"Signal {signum.name} has no handlers or is not supported on Windows."
+            )
             return
 
         if self._has_already_handler(signum):
             if not replace_existing:
                 log.info(
-                    f"Signal {signum} already has a handler. Adding ours to the existing one."
+                    f"Signal {signum.name} already has a handler. Adding ours to the existing one."
                 )
                 handlers.append(signal.getsignal(signum))
             else:
-                log.info(f"Replacing existing handler for signal {signum} with ours.")
+                log.info(
+                    f"Replacing existing handler for signal {signum.name} with ours."
+                )
 
         self._register_signal(signum, _HandlersCompose(handlers))
-        log.info(f"Registered {len(handlers)} handlers for signal {signum}.")
+        log.info(f"Registered {len(handlers)} handlers for signal {signum.name}.")
 
     @override
     def register_signal_handlers(self) -> None:

nshtrainer/trainer/trainer.py

@@ -31,7 +31,7 @@ log = logging.getLogger(__name__)
 
 def _is_bf16_supported_no_emulation():
     r"""Return a bool indicating if the current CUDA/ROCm device supports dtype bfloat16."""
-    version = cast(Any, torch.version)
+    version = getattr(torch, "version")
 
     # Check for ROCm, if true return true, no ROCM_VERSION check required,
     # since it is supported on AMD GPU archs.

nshtrainer/typecheck.py

@@ -82,6 +82,7 @@ def typecheck_this_module(additional_modules: Sequence[str] = ()):
     frame = get_frame(1)
     assert frame is not None, "frame is None"
     calling_module_name = get_frame_package_name(frame)
+    assert calling_module_name is not None, "calling_module_name is None"
 
     # Typecheck the calling module + any additional modules.
     typecheck_modules((calling_module_name, *additional_modules))

{nshtrainer-0.1.0.dist-info → nshtrainer-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.1.0
+Version: 0.2.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -9,10 +9,21 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: beartype (>=0.18.5,<0.19.0)
+Requires-Dist: jaxtyping (>=0.2.33,<0.3.0)
+Requires-Dist: lightning
+Requires-Dist: lovely-numpy (>=0.2.13,<0.3.0)
+Requires-Dist: lovely-tensors (>=0.1.16,<0.2.0)
 Requires-Dist: nshconfig (>=0.2.0,<0.3.0)
-Requires-Dist: nshrunner (>=0.1.0,<0.2.0)
+Requires-Dist: nshrunner (>=0.5.4,<0.6.0)
+Requires-Dist: numpy
+Requires-Dist: pysnooper
+Requires-Dist: pytorch-lightning
+Requires-Dist: rich
 Requires-Dist: torch
+Requires-Dist: torchmetrics
 Requires-Dist: typing-extensions
+Requires-Dist: wrapt
 Description-Content-Type: text/markdown
 
 

{nshtrainer-0.1.0.dist-info → nshtrainer-0.2.0.dist-info}/RECORD

@@ -1,22 +1,16 @@
-nshtrainer/__init__.py,sha256=o39TbnjwUYzE4POcncUiDx02Ey-Hzx8UGuwJDjMcKZU,2971
+nshtrainer/__init__.py,sha256=_r7kBmgGSLVfActlqQeupNolrmBu45xUuSS8odt3HL8,2208
 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
 nshtrainer/_experimental/flops/module_tracker.py,sha256=bUL-IRTd0aF_DwmXkZjHZAA31p4ZEhyqhc26XWKQUUY,4922
 nshtrainer/_snoop.py,sha256=Rofv1Rd92E0LY40G3A-o9Hu0ZI73RR59wJD5l4Q3PDM,7022
-nshtrainer/_submit/print_environment_info.py,sha256=enbJGl_iHIlhKN8avzKnoZSb0zUQ_fUdnsQ8a_9tbYk,963
-nshtrainer/_submit/session/_output.py,sha256=CNGH5W6_XxAC5-TRvMAMxOHd3fjGpJhK-7RGTDyvMu4,245
-nshtrainer/_submit/session/_script.py,sha256=0AeBgBduDsoIEBrY9kebARiBUEGc50JAD9oE_IDiLnA,3775
-nshtrainer/_submit/session/lsf.py,sha256=p19EP6OhROZxqfRhzeTD7GDmfYaREIKMXMOI8G933FE,14307
-nshtrainer/_submit/session/slurm.py,sha256=JpAjQvck4LjGN8o8fOvIeMuFqrg1cioANoVsX5hU-3g,17594
-nshtrainer/_submit/session/unified.py,sha256=gfh-AtnMyFHzcQOUlhlAR__vaWDk1r9XCivz_t_lHKk,11695
 nshtrainer/actsave/__init__.py,sha256=G1T-fELuGWkVqdhdyoePtj2dTOUtcIOW4VgsXv9JNTA,338
 nshtrainer/actsave/_callback.py,sha256=QoTa60F70f1RxB41VKixN9l5_htfFQxXDPHHSNFreuk,2770
 nshtrainer/actsave/_loader.py,sha256=fAhD32DrJa4onkYfcwc21YIeGEYzOSXCK_HVo9SZLgQ,4604
 nshtrainer/actsave/_saver.py,sha256=0EHmQDhqVxQWRWWSyt03eP1K9ETiACMQYmsZkDMt6HY,9451
-nshtrainer/callbacks/__init__.py,sha256=ohE_MO_kX1o4SZwcipIXUA9m7XYcijEKJtGcoU8dTkY,1667
+nshtrainer/callbacks/__init__.py,sha256=I6W33ityL9Ko8jjqHh3WH_8miV59SAe9LxInhoqX5XE,1665
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
-nshtrainer/callbacks/base.py,sha256=WESZz1VSTl1xSGVXBmxFqWwbLxXcJp97jpg9zrE0EsY,3560
+nshtrainer/callbacks/base.py,sha256=LrcRUV02bZEKXRIRvhHT9qsvw_kwoWiAdQkVMyKc5NU,3542
 nshtrainer/callbacks/early_stopping.py,sha256=jriSU761wf_qTJ9Bos0D3h5aDvZHYpRqK62Ne8aWp5I,3768
 nshtrainer/callbacks/ema.py,sha256=zKCtvzZFo0ORlwNZHjaMk-sJoxrlTtFWOzR-yGy95W0,12134
 nshtrainer/callbacks/finite_checks.py,sha256=kX3TIJsxyqx0GuLJfYsqVgKU27zwjG9Z8324lyCFtwM,2087
@@ -30,17 +24,17 @@ nshtrainer/callbacks/print_table.py,sha256=FcA-CBWwMf9c1NNRinvYpZC400RNQxuP28bJf
 nshtrainer/callbacks/throughput_monitor.py,sha256=YQLdpX3LGybIiD814yT9yCCVSEXRWf8WwsvVaN5aDBE,1848
 nshtrainer/callbacks/timer.py,sha256=sDXPPcdDKu5xnuK_bjr8plIq9MBuluNJ42Mt9LvPZzc,4610
 nshtrainer/callbacks/wandb_watch.py,sha256=pUpMsNxd03ex1rzOmFw2HzGOXjnQGaH84m8cc2dXo4g,2937
-nshtrainer/config.py,sha256=0Fj5w-ry0BRl2_zJI6jwCnmMWE3p_eD8_Wn-NyFkTqU,10442
+nshtrainer/config.py,sha256=IXOAl_JWFNX9kPTo_iw4Nc3qXqkKrbA6-ZrvTAjqu6A,104
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
 nshtrainer/data/balanced_batch_sampler.py,sha256=bcJBcQjh1hB1yKF_xSlT9AtEWv0BJjYc1CuH2BF-ea8,4392
 nshtrainer/data/transform.py,sha256=JeGxvytQly8hougrsdMmKG8gJ6qvFPDglJCO4Tp6STk,1795
-nshtrainer/lr_scheduler/__init__.py,sha256=GNGmkcJD3jgCMk7pfaanAYrKz9957qkx6_Q0rssiHK0,738
-nshtrainer/lr_scheduler/_base.py,sha256=1tWMABevKZAuGhJN8Me2E9eqEyqoLtsG0bADPjED7a4,3752
-nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=VhsxZJ_Mw9zjkAGunFQ1KRub5_QM5NRqaEFWtmedFp8,5212
-nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=Ct-uLo8Q4t7lJ_HwoLRhNmudnCw4cSnblpBEg22aVTI,2691
-nshtrainer/model/__init__.py,sha256=PdvZkpAVkqvCLipGJvEHFU3WxnSMxYpvtuOkvLIenxg,2078
-nshtrainer/model/base.py,sha256=bhngGHxr0suQB9Ezi_3d5JgDWYqS_yPgGJZrGmc1TnI,23571
-nshtrainer/model/config.py,sha256=RMDdrbtvwm5vTFPxQ2x1hqiBIEEE-OAknhF6KTWfkkk,70293
+nshtrainer/lr_scheduler/__init__.py,sha256=uEvgaFAs-4s_bAEMaildy0GT6OvgpgOEKTuzqutESHE,736
+nshtrainer/lr_scheduler/_base.py,sha256=7xOIuxQ86YHbFWG5a3gX46emQj1WN_LaY4-i0Q1TDBg,3659
+nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=mn6cyizyI_stkXtg6zxIEGF9btIxMRWigUHUTlUYCSw,5221
+nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=2ZdlV0RUMwg2DClzqYHr8_EKT1jZBUlSD39e-XlCsC4,2764
+nshtrainer/model/__init__.py,sha256=y32Hla-5whpzLL2BtCJpBakSp8o-1nQbpO0j_-xq_Po,1864
+nshtrainer/model/base.py,sha256=EMkOtp4YWGPHM0HPSTLbx75T9vlYmXO4XyD725xU70w,21453
+nshtrainer/model/config.py,sha256=6lATW6-Z1SIDgQ1IWrGBVQKTr8DhL5b_rFbJHQz0d5o,66796
 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
@@ -52,21 +46,20 @@ nshtrainer/nn/__init__.py,sha256=57LPaP3G-BBGD2eGxbBUABNgYl3s_oASwrtOSS4bzTs,133
 nshtrainer/nn/mlp.py,sha256=i-dHk0tomO_XlU6cKN4CC4HxTaYb-ukBCAgY1ySXl4I,3963
 nshtrainer/nn/module_dict.py,sha256=NOY0B6WDTnktyWH4GthsprMQo0bpehC-hCq9SfD8paE,2329
 nshtrainer/nn/module_list.py,sha256=fb2u5Rqdjff8Pekyr9hkCPkBorQ-fldzzFAjsgWAm30,1719
-nshtrainer/nn/nonlinearity.py,sha256=IhIR8NCTY3Np9dMDnUouERR8ZhWpK3S0hTbT0i8HezU,3645
-nshtrainer/optimizer.py,sha256=JiLNRtcfYxyhAab1Z1QcEzmrX9S_JyrBS67TXy12kXI,1557
-nshtrainer/runner.py,sha256=9HsYB58aasY9RVvya_gPECDs_MBhM1fl4cbM3iJYTDc,600
+nshtrainer/nn/nonlinearity.py,sha256=owtU4kh4G98psD0axOJWVfBhm-OtJVgFM-TXSHmbNPU,3625
+nshtrainer/optimizer.py,sha256=kuJEA1pvB3y1FcsfhAoOJujVqEZqFHlmYO8GW6JeA1g,1527
+nshtrainer/runner.py,sha256=vyHr0EZ0PBOWZh09BtOOxio-FRQZFbVoL4cdBlI97vY,991
 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/signal_connector.py,sha256=aGg6kRiHiqtAdGlEvEvGLmOy7AvRHTSkXdTmZpRXbjU,8435
-nshtrainer/trainer/trainer.py,sha256=oi8KdHF1AdZ54KFbCFAEI7W-C7qRtRe-KtOjNwBuS3M,14033
-nshtrainer/typecheck.py,sha256=CFkmPIxCU24CHk_7_pykb-Y1PRNhpLgsVZw1zuuOS_U,4614
+nshtrainer/trainer/signal_connector.py,sha256=QAoPM_C5JJOVQebcrJOimUUD3GHyoeZUqCEAvzZlT4U,8710
+nshtrainer/trainer/trainer.py,sha256=eYEYfY9v70MuorHcSf8nqM7f2CkmUHhpPcjCk4FJD7k,14034
+nshtrainer/typecheck.py,sha256=RGYHxDBcs97E6ayl6Olc43JBZXQolCtMxcLBniVCVBg,4688
 nshtrainer/util/environment.py,sha256=_SEtiQ_s5bL5pllUlf96AOUv15kNvCPvocVC13S7mIk,4166
 nshtrainer/util/seed.py,sha256=HEXgVs-wldByahOysKwq7506OHxdYTEgmP-tDQVAEkQ,287
-nshtrainer/util/singleton.py,sha256=nLhpuMZxl0zdNsnvS97o4ASUnKzCWYEKLzR_j9oP_xs,2208
 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.1.0.dist-info/METADATA,sha256=3zdNPxyB-I6Gudq2gTaU0crdgmDCcGCp6Zudef0DtuM,529
-nshtrainer-0.1.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.1.0.dist-info/RECORD,,
+nshtrainer-0.2.0.dist-info/METADATA,sha256=cwb3IbKGyJ9HbNSvsORYhCiI61nrDMb1dVm5nE1q_XA,882
+nshtrainer-0.2.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.2.0.dist-info/RECORD,,

nshtrainer/_submit/print_environment_info.py

@@ -1,31 +0,0 @@
-import logging
-import os
-import sys
-
-
-def print_environment_info(log: logging.Logger | None = None):
-    if log is None:
-        logging.basicConfig(level=logging.INFO)
-        log = logging.getLogger(__name__)
-
-    log_message_lines: list[str] = []
-    log_message_lines.append("Python executable: " + sys.executable)
-    log_message_lines.append("Python version: " + sys.version)
-    log_message_lines.append("Python prefix: " + sys.prefix)
-    log_message_lines.append("Python path:")
-    for path in sys.path:
-        log_message_lines.append(f"  {path}")
-
-    log_message_lines.append("Environment variables:")
-    for key, value in os.environ.items():
-        log_message_lines.append(f"  {key}={value}")
-
-    log_message_lines.append("Command line arguments:")
-    for i, arg in enumerate(sys.argv):
-        log_message_lines.append(f"  {i}: {arg}")
-
-    log.critical("\n".join(log_message_lines))
-
-
-if __name__ == "__main__":
-    print_environment_info()