nshtrainer 0.14.2__tar.gz → 0.15.0__tar.gz

{nshtrainer-0.14.2 → nshtrainer-0.15.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.14.2
+Version: 0.15.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -16,6 +16,7 @@ Requires-Dist: nshconfig
 Requires-Dist: nshrunner
 Requires-Dist: nshutils
 Requires-Dist: numpy
+Requires-Dist: packaging
 Requires-Dist: psutil
 Requires-Dist: pytorch-lightning
 Requires-Dist: tensorboard ; extra == "extra"

{nshtrainer-0.14.2 → nshtrainer-0.15.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.14.2"
+version = "0.15.0"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -14,6 +14,7 @@ psutil = "*"
 numpy = "*"
 torch = "*"
 typing-extensions = "*"
+packaging = "*"
 lightning = "*"
 pytorch-lightning = "*"
 torchmetrics = { version = "*", optional = true }

nshtrainer-0.15.0/src/nshtrainer/_experimental/__init__.py

@@ -0,0 +1 @@
+from lightning.fabric.utilities.throughput import measure_flops as measure_flops

{nshtrainer-0.14.2 → nshtrainer-0.15.0}/src/nshtrainer/trainer/signal_connector.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import platform
 import re
 import signal
 import subprocess
@@ -25,6 +26,7 @@ log = logging.getLogger(__name__)
 
 _SIGNUM = int | signal.Signals
 _HANDLER: TypeAlias = Callable[[_SIGNUM, FrameType], Any] | int | signal.Handlers | None
+_IS_WINDOWS = platform.system() == "Windows"
 
 
 def _resolve_requeue_signals():
@@ -57,7 +59,7 @@ class _SignalConnector(_LightningSignalConnector):
         handlers: list[_HANDLER],
         replace_existing: bool = False,
     ):
-        if self._is_on_windows():
+        if _IS_WINDOWS:
             log.info(
                 f"Signal {signum.name} has no handlers or is not supported on Windows."
             )

{nshtrainer-0.14.2 → nshtrainer-0.15.0}/src/nshtrainer/util/_environment_info.py

@@ -1,4 +1,5 @@
 import getpass
+import importlib.metadata
 import inspect
 import logging
 import os
@@ -12,6 +13,7 @@ from typing import TYPE_CHECKING, Any, cast
 import nshconfig as C
 import psutil
 import torch
+from packaging import version
 from typing_extensions import Self
 
 from .slurm import parse_slurm_node_list
@@ -398,23 +400,46 @@ class EnvironmentPackageConfig(C.Config):
 
     @classmethod
     def from_current_environment(cls):
-        # Add Python package information
         python_packages: dict[str, Self] = {}
         try:
-            import pkg_resources
-
-            for package in pkg_resources.working_set:
-                python_packages[package.key] = cls(
-                    name=package.project_name,
-                    version=package.version,
-                    path=Path(package.location) if package.location else None,
-                    summary=getattr(package, "summary", None),
-                    author=getattr(package, "author", None),
-                    license=getattr(package, "license", None),
-                    requires=[str(req) for req in package.requires()],
-                )
+            for dist in importlib.metadata.distributions():
+                try:
+                    # Get package metadata
+                    metadata = dist.metadata
+
+                    # Parse the version, stripping any local version identifier
+                    pkg_version = version.parse(dist.version)
+                    clean_version = (
+                        f"{pkg_version.major}.{pkg_version.minor}.{pkg_version.micro}"
+                    )
+
+                    # Get requirements
+                    requires = []
+                    for req in dist.requires or []:
+                        try:
+                            requires.append(str(req))
+                        except ValueError:
+                            # If there's an invalid requirement, we'll skip it
+                            log.warning(
+                                f"Skipping invalid requirement for {dist.name}: {req}"
+                            )
+
+                    python_packages[dist.name] = cls(
+                        name=dist.name,
+                        version=clean_version,
+                        path=Path(str(f)) if (f := dist.locate_file("")) else None,
+                        summary=metadata["Summary"] if "Summary" in metadata else None,
+                        author=metadata["Author"] if "Summary" in metadata else None,
+                        license=metadata["License"] if "Summary" in metadata else None,
+                        requires=requires,
+                    )
+                except Exception as e:
+                    log.warning(f"Error processing package {dist.name}: {str(e)}")
+
         except ImportError:
-            log.warning("pkg_resources not available, skipping package information")
+            log.warning(
+                "importlib.metadata not available, skipping package information"
+            )
 
         return python_packages
 

nshtrainer-0.14.2/src/nshtrainer/_experimental/__init__.py

@@ -1,2 +0,0 @@
-from .flops import MEASURE_FLOPS_AVAILABLE as MEASURE_FLOPS_AVAILABLE
-from .flops import measure_flops as measure_flops

nshtrainer-0.14.2/src/nshtrainer/_experimental/flops/__init__.py

@@ -1,48 +0,0 @@
-from collections.abc import Callable
-
-import torch
-from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_1
-
-MEASURE_FLOPS_AVAILABLE = _TORCH_GREATER_EQUAL_2_1
-
-
-def measure_flops(
-    forward_fn: Callable[[], torch.Tensor],
-    loss_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
-    display: bool = True,
-) -> int:
-    """Utility to compute the total number of FLOPs used by a module during training or during inference.
-
-    It's recommended to create a meta-device model for this:
-
-    Example::
-
-        with torch.device("meta"):
-            model = MyModel()
-            x = torch.randn(2, 32)
-
-        model_fwd = lambda: model(x)
-        fwd_flops = measure_flops(model, model_fwd)
-
-        model_loss = lambda y: y.sum()
-        fwd_and_bwd_flops = measure_flops(model, model_fwd, model_loss)
-
-    Args:
-        model: The model whose FLOPs should be measured.
-        forward_fn: A function that runs ``forward`` on the model and returns the result.
-        loss_fn: A function that computes the loss given the ``forward_fn`` output. If provided, the loss and `backward`
-            FLOPs will be included in the result.
-
-    """
-    if not MEASURE_FLOPS_AVAILABLE:
-        raise ImportError("`measure_flops` requires PyTorch >= 2.1.")
-
-    from .flop_counter import FlopCounterMode
-
-    flop_counter = FlopCounterMode(display=display)
-    with flop_counter:
-        if loss_fn is None:
-            forward_fn()
-        else:
-            loss_fn(forward_fn()).backward()
-    return flop_counter.get_total_flops()