nshtrainer 1.0.0b56__tar.gz → 1.1.1b1__tar.gz

{nshtrainer-1.0.0b56 → nshtrainer-1.1.1b1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.0.0b56
+Version: 1.1.1b1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -14,13 +14,12 @@ Provides-Extra: extra
 Requires-Dist: GitPython ; extra == "extra"
 Requires-Dist: huggingface-hub ; extra == "extra"
 Requires-Dist: lightning
-Requires-Dist: nshconfig
+Requires-Dist: nshconfig (>0.39)
 Requires-Dist: nshrunner
 Requires-Dist: nshutils ; extra == "extra"
 Requires-Dist: numpy
 Requires-Dist: packaging
 Requires-Dist: psutil
-Requires-Dist: pytorch-lightning
 Requires-Dist: tensorboard ; extra == "extra"
 Requires-Dist: torch
 Requires-Dist: torchmetrics

{nshtrainer-1.0.0b56 → nshtrainer-1.1.1b1}/pyproject.toml

@@ -1,28 +1,32 @@
-[tool.poetry]
+[project]
 name = "nshtrainer"
-version = "1.0.0-beta56"
+version = "1.1.1-beta.1"
 description = ""
-authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
+authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
+requires-python = ">=3.10,<4.0"
 readme = "README.md"
 
-[tool.poetry.dependencies]
-python = "^3.10"
-nshrunner = { version = "*" }
-nshconfig = "*"
-nshutils = { version = "*", optional = true }
-psutil = "*"
-numpy = "*"
-torch = "*"
-typing-extensions = "*"
-packaging = "*"
-lightning = "*"
-pytorch-lightning = "*"
-torchmetrics = "*"
-wrapt = { version = "*", optional = true }
-GitPython = { version = "*", optional = true }
-wandb = { version = "*", optional = true }
-tensorboard = { version = "*", optional = true }
-huggingface-hub = { version = "*", optional = true }
+dependencies = [
+    "nshconfig>0.39",
+    "nshrunner",
+    "psutil",
+    "numpy",
+    "torch",
+    "typing-extensions",
+    "packaging",
+    "lightning",
+    "torchmetrics",
+]
+
+[project.optional-dependencies]
+extra = [
+    "wrapt",
+    "GitPython",
+    "wandb",
+    "tensorboard",
+    "huggingface-hub",
+    "nshutils",
+]
 
 [tool.poetry.group.dev.dependencies]
 basedpyright = "*"
@@ -51,16 +55,6 @@ ignore = ["F722", "F821", "E731", "E741"]
 [tool.ruff.lint.isort]
 required-imports = ["from __future__ import annotations"]
 
-[tool.poetry.extras]
-extra = [
-    "wrapt",
-    "GitPython",
-    "wandb",
-    "tensorboard",
-    "huggingface-hub",
-    "nshutils",
-]
-
 [tool.pytest]
 testpaths = ["tests"]
 python_files = "test_*.py"

{nshtrainer-1.0.0b56 → nshtrainer-1.1.1b1}/src/nshtrainer/__init__.py

@@ -14,9 +14,6 @@ from .metrics import MetricConfig as MetricConfig
 from .model import LightningModuleBase as LightningModuleBase
 from .trainer import Trainer as Trainer
 from .trainer import TrainerConfig as TrainerConfig
-from .trainer import accelerator_registry as accelerator_registry
-from .trainer import callback_registry as callback_registry
-from .trainer import plugin_registry as plugin_registry
 
 try:
     from . import configs as configs

{nshtrainer-1.0.0b56 → nshtrainer-1.1.1b1}/src/nshtrainer/configs/__init__.py

@@ -4,9 +4,6 @@ __codegen__ = True
 
 from nshtrainer import MetricConfig as MetricConfig
 from nshtrainer import TrainerConfig as TrainerConfig
-from nshtrainer import accelerator_registry as accelerator_registry
-from nshtrainer import callback_registry as callback_registry
-from nshtrainer import plugin_registry as plugin_registry
 from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
 from nshtrainer._directory import DirectoryConfig as DirectoryConfig
 from nshtrainer._hf_hub import CallbackConfigBase as CallbackConfigBase
@@ -14,6 +11,7 @@ from nshtrainer._hf_hub import (
     HuggingFaceHubAutoCreateConfig as HuggingFaceHubAutoCreateConfig,
 )
 from nshtrainer._hf_hub import HuggingFaceHubConfig as HuggingFaceHubConfig
+from nshtrainer._hf_hub import callback_registry as callback_registry
 from nshtrainer.callbacks import ActSaveConfig as ActSaveConfig
 from nshtrainer.callbacks import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
@@ -106,6 +104,8 @@ from nshtrainer.profiler import BaseProfilerConfig as BaseProfilerConfig
 from nshtrainer.profiler import ProfilerConfig as ProfilerConfig
 from nshtrainer.profiler import PyTorchProfilerConfig as PyTorchProfilerConfig
 from nshtrainer.profiler import SimpleProfilerConfig as SimpleProfilerConfig
+from nshtrainer.trainer import accelerator_registry as accelerator_registry
+from nshtrainer.trainer import plugin_registry as plugin_registry
 from nshtrainer.trainer._config import AcceleratorConfig as AcceleratorConfig
 from nshtrainer.trainer._config import (
     CheckpointCallbackConfig as CheckpointCallbackConfig,

{nshtrainer-1.0.0b56 → nshtrainer-1.1.1b1}/src/nshtrainer/data/datamodule.py

@@ -30,9 +30,9 @@ class LightningDataModuleBase(
 
     @property
     @override
-    def hparams_initial(self):  # pyright: ignore[reportIncompatibleMethodOverride]
+    def hparams_initial(self) -> THparams:  # pyright: ignore[reportIncompatibleMethodOverride]
         hparams = cast(THparams, super().hparams_initial)
-        return cast(Never, {"datamodule": hparams.model_dump(mode="json")})
+        return hparams
 
     @property
     @deprecated("Use `hparams` instead")

{nshtrainer-1.0.0b56 → nshtrainer-1.1.1b1}/src/nshtrainer/model/base.py

@@ -134,16 +134,9 @@ class LightningModuleBase(
 
     @property
     @override
-    def hparams_initial(self):  # pyright: ignore[reportIncompatibleMethodOverride]
+    def hparams_initial(self) -> THparams:  # pyright: ignore[reportIncompatibleMethodOverride]
         hparams = cast(THparams, super().hparams_initial)
-        hparams_dict = {"model": hparams.model_dump(mode="json")}
-        if (trainer := self._trainer) is not None:
-            from ..trainer import Trainer
-
-            if isinstance(trainer, Trainer):
-                hparams_dict["trainer"] = trainer.hparams.model_dump(mode="json")
-
-        return cast(Never, hparams_dict)
+        return hparams
 
     @property
     @deprecated("Use `hparams` instead")

nshtrainer-1.1.1b1/src/nshtrainer/trainer/_log_hparams.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any, cast
+
+import nshconfig as C
+from lightning.pytorch import LightningDataModule, Trainer
+
+log = logging.getLogger(__name__)
+
+
+def _dict(obj: Any):
+    if isinstance(obj, C.Config):
+        return obj.model_dump(mode="json")
+
+    try:
+        return dict(obj)
+    except Exception:
+        return json.loads(
+            json.dumps(obj, default=lambda o: str(o), indent=4, sort_keys=True)
+        )
+
+
+def _dict_and_clean(obj: Any):
+    d = _dict(obj)
+
+    # Remove LightningCLI's internal hparam
+    d = {k: v for k, v in d.items() if k != "_class_path"}
+    return d
+
+
+def _log_hyperparams(trainer: Trainer) -> None:
+    if not trainer.loggers:
+        return
+
+    hparams_to_log: dict[str, Any] = {}
+
+    from .trainer import Trainer
+
+    if isinstance(trainer, Trainer):
+        hparams_to_log["trainer"] = _dict_and_clean(trainer.hparams)
+
+    if (
+        pl_module := trainer.lightning_module
+    ) is not None and pl_module._log_hyperparams:
+        hparams_to_log["model"] = _dict_and_clean(pl_module.hparams_initial)
+
+    if (
+        datamodule := cast(
+            LightningDataModule | None, getattr(trainer, "datamodule", None)
+        )
+    ) is not None and (datamodule._log_hyperparams):
+        hparams_to_log["datamodule"] = _dict_and_clean(datamodule.hparams_initial)
+
+    for logger in trainer.loggers:
+        logger.log_hyperparams(hparams_to_log)
+        logger.log_graph(pl_module)
+        logger.save()
+
+
+def patch_log_hparams_function():
+    try:
+        import lightning.pytorch.loggers.utilities
+        import lightning.pytorch.trainer.trainer
+
+        lightning.pytorch.loggers.utilities._log_hyperparams = _log_hyperparams
+        lightning.pytorch.trainer.trainer._log_hyperparams = _log_hyperparams
+        log.info(
+            "Patched lightning.pytorch's _log_hyperparams to use nshtrainer's version"
+        )
+    except ImportError:
+        pass
+
+    try:
+        import pytorch_lightning.loggers.utilities
+        import pytorch_lightning.trainer.trainer
+
+        pytorch_lightning.loggers.utilities._log_hyperparams = _log_hyperparams
+        pytorch_lightning.trainer.trainer._log_hyperparams = _log_hyperparams
+        log.info(
+            "Patched pytorch_lightning's _log_hyperparams to use nshtrainer's version"
+        )
+    except ImportError:
+        pass

{nshtrainer-1.0.0b56 → nshtrainer-1.1.1b1}/src/nshtrainer/trainer/trainer.py

@@ -23,6 +23,7 @@ from ..callbacks.base import resolve_all_callbacks
 from ..util._environment_info import EnvironmentConfig
 from ..util.bf16 import is_bf16_supported_no_emulation
 from ._config import LightningTrainerKwargs, TrainerConfig
+from ._log_hparams import patch_log_hparams_function
 from ._runtime_callback import RuntimeTrackerCallback, Stage
 from .accelerator import AcceleratorConfigBase
 from .signal_connector import _SignalConnector
@@ -31,6 +32,9 @@ from .strategy import StrategyConfigBase
 log = logging.getLogger(__name__)
 
 
+patch_log_hparams_function()
+
+
 class Trainer(LightningTrainer):
     CHECKPOINT_HYPER_PARAMS_KEY = "trainer_hyper_parameters"