nshtrainer 0.14.4__tar.gz → 0.15.1__tar.gz

{nshtrainer-0.14.4 → nshtrainer-0.15.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.14.4
+Version: 0.15.1
 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.4 → nshtrainer-0.15.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.14.4"
+version = "0.15.1"
 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.14.4 → nshtrainer-0.15.1}/src/nshtrainer/loggers/wandb.py

@@ -1,9 +1,10 @@
+import importlib.metadata
 import logging
 from typing import TYPE_CHECKING, Literal
 
 import nshconfig as C
-import pkg_resources
 from lightning.pytorch import Callback, LightningModule, Trainer
+from packaging import version
 from typing_extensions import override
 
 from ..callbacks import WandbWatchConfig
@@ -120,9 +121,8 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
                 import wandb  # type: ignore
 
                 # The minimum version that supports the new backend is 0.17.5
-                if pkg_resources.parse_version(
-                    wandb.__version__
-                ) < pkg_resources.parse_version("0.17.5"):
+                wandb_version = version.parse(importlib.metadata.version("wandb"))
+                if wandb_version < version.parse("0.17.5"):
                     log.warning(
                         "The version of WandB installed does not support the `wandb-core` backend "
                         f"(expected version >= 0.17.5, found version {wandb.__version__}). "

{nshtrainer-0.14.4 → nshtrainer-0.15.1}/src/nshtrainer/scripts/find_packages.py

@@ -40,7 +40,6 @@ def main():
 
     if args.exclude_std:
         std_libs = set(sys.stdlib_module_names)
-        std_libs.update({"pkg_resources"})
         all_imports = all_imports - std_libs
 
     for package in sorted(all_imports):

{nshtrainer-0.14.4 → nshtrainer-0.15.1}/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