Update dependencies and enforce dependency checking in physicsnemo

[coreyjadams]

PhysicsNeMo Pull Request

This PR has grown a little bit bigger than anticipated, so let me summarize:

1. First goal was to move zarr and xarray to optional deps instead of core deps.
2. Along the way, I realized how many optional deps in datapipes aren't protected against spurious imports. It's not a breaking issue, since we don't import those, but I added protection layers against all of them.
3. This also let me update some dependencies in pyproject.toml.

For pyproject.toml, a new optional dependency group is created for "performance" centric items. Since, in reality, the best performance is from nvidia's libraries, this section is entirely nvidia packages with some sort of cuda binding.

I updated uv.lock as well.

Description

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[greptile-apps]

Greptile Overview

Greptile Summary

This PR successfully moves zarr and xarray from core dependencies to optional dependencies and implements comprehensive import protection for all optional dependencies across datapipes.

Key Changes

• Dependency Management: Moved zarr and xarray to optional datapipes-extras group, created new perf group for NVIDIA performance libraries (transformer_engine, cuml-cu13, pylibraft-cu13, cupy-cuda13x)
• Import Protection Pattern: Consistently applied check_version_spec() with hard_fail=False and importlib.import_module() across all datapipes to defer import errors until runtime
• Protected Dependencies: Added protection for zarr, xarray, DALI, netCDF4, scipy, torch_geometric, pyvista, vtk, tfrecord, dask, wrapt, natten, and sparse_dot_mkl
• Centralized Utilities: Created physicsnemo/datapipes/healpix/utils.py to centralize xarray import protection for healpix modules
• Import Linter Update: Changed .importlinter exclusion from datapipes to metrics to enforce dependency checking in datapipes

Issues Found

The implementation is thorough and follows a consistent pattern. One minor concern identified below.

Important Files Changed

Filename	Overview
pyproject.toml	Moved zarr and xarray to optional deps, created new 'perf' group for NVIDIA performance libraries
.importlinter	Changed exclusion from datapipes to metrics for external imports, fixed whitespace formatting
physicsnemo/datapipes/readers/zarr.py	Added proper import protection for zarr using check_version_spec and importlib
physicsnemo/datapipes/climate/climate.py	Protected DALI, netCDF4, and scipy imports with version checking and helper functions
physicsnemo/datapipes/healpix/utils.py	New utility module centralizing xarray import protection for healpix datapipes
physicsnemo/domain_parallel/shard_utils/natten_patches.py	Protected wrapt and natten imports, conditional patch registration based on availability

[greptile-apps]

_{6 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Inconsistent with line 35 and the rest of the PR - should use hard_fail=False for optional dependency

Suggested change

	NATTEN_AVAILABLE = check_version_spec("natten")
	NATTEN_AVAILABLE = check_version_spec("natten", hard_fail=False)

[coreyjadams]

This PR also removes fcn_mip_plugin from models. And makes sure that all core packages are part of the import scan.

[pzharrington]

Do we want these error messages to mention the optional deps groups specified in the toml? May be something for another PR, but I feel like discoverability of those groups is not that good right now

[pzharrington]

For ex. Nick set up some import handlers/checkers that are based on deps groups and raise errors accordingly -- might clean up a lot of the dedicated check methods that have been added in this PR. See https://github.com/NVIDIA/earth2studio/blob/main/earth2studio/utils/imports.py

[coreyjadams]

It would be nice to have something like that. I was thinking somehow about how we could reduce this boiler plate... I haven't got a solution yet, I'll look at this too. Do you want this in this PR or just in some nebulous future?

[pzharrington]

In this PR there are a lot of things like the below, which in some cases led to having to create extra files or just some potentially unnecessary bloat:

def _raise_missing_xarray():
    raise ImportError(
        "xarray is required for physicsnemo.datapipes.healpix but is not available. "
        "Please install xarray with `pip install xarray`."
    )

I think they could be cleaned up nicely by instead having some sort of e.g. physicsnemo.core.raise_missing_optional_dep method or decorator that could print some ImportError message like

<package> is required for <caller/object/source of import error> but is not available. Please install <package> directly, or install the optional dependency <deps-group>

Where the caller provides the bracketed items in the src code, something like

raise_missing_optional_dep("xarray", "datapipes-extras")

Would also standardize the messaging and herd people towards paying attention to deps groups

[coreyjadams]

We have, in my mind, two problems to solve:

1. When dependencies are missing, we need to be informative about what is missing and how to get it.
2. When dependencies are missing, we need to turn ImportError into RunTimeError so that no one ever crashes out with import physicsnemo when they don't even need [healpix | transformer_engine | dali | ...].

Thinking about your suggestion, maybe we can have some sort of fake dependency factory:

class MissingImport():
  def __init__(self, library_name):
    self.lib_name = library_name

  def __getattr__(self, key):
    raise RuntimeError("Missing {self.lib_name}, please blah blah blah")

then, we could go a step further with

def maybe_import(package_name: str):
  # importlib logic to check and import
  if not available:
    return MissingImport(package_name)
  else:
    return package

It might consolidate the boiler plate and let us automate error raising when it's in python code. It won't help, though, when we do something else like use package.decorator or inherit from something. Those might not be realistic use cases.

[pzharrington]

Yeah I like that, sounds like it would get the two most important birds. Not sure, but for this:

It won't help, though, when we do something else like use package.decorator or inherit from something

We might also be able to use your idea in some decorator form that would raise the error on the __init__ of an inherited class. Not sure if it would be doable to double-decorate something successfully though

[coreyjadams]

I blew up this PR, a bit, based on some of the discussions with @pzharrington. It looks bigger than it is. Major changes:

• New OptionalImport object from core.version_spec. This can be used as a lazy import, but also will raise a runtime error at first access and not import time.
• Allow type checking ifs in import linting. So, we can use if TYPE_CHECKING for static checking without breaking the optional import paradigm. This I don't reallllly love since it's duplicating imports, but it's not necessary except in a couple places and it's not complicated.
• I made a registry of useful install hints in OptionalImport. So, we don't need to worry much about writing good error messages, since it's consolidated. It specifies the uv dependency group, where to pip install, other instruction locations, etc. We can elaborate on this - I haven't gone through closely, yet, let's make sure it's the right model first.

about 5000 lines of the +5k / -5k here are just replacing boilerplate with better syntax and un-indenting blocks of code that no longer need to be in an if/else.

[pzharrington]

dgl is deprecated so we can drop this

[dallasfoster]

While we are at it, I would argue that h5py could also be made an optional dependency.

[dallasfoster]

I would target cftime as another dependency that appears optional, it only appears to be used in examples.

[coreyjadams]

Good call. Used to be used in some functionality but it's been removed since.

[dallasfoster]

@NickGeneva do we really need to keep s3fs?

[dallasfoster]

What do we think about pandas? @pzharrington it appears that it might be used in healpix utilities which could be made optional-model dependent?

[pzharrington]

Hmm, it's also used by the insolation function and the drivaernet datapipe, so not totally clear which optional deps group it would belong in. Since it's not causing install problems (AFAIK) and not too bulky/slow to build, I wouldn't mind just leaving it in core.

[dallasfoster]

A few other suggestions for dependencies.

[coreyjadams]

Something happened in this merge to blow this up to 1400+ files. I"m going to rebase this work onto a fresh PR.