diff --git a/.agents/rules/mcp-environments.md b/.agents/rules/mcp-environments.md index cfef1693..52ad64e8 100644 --- a/.agents/rules/mcp-environments.md +++ b/.agents/rules/mcp-environments.md @@ -11,7 +11,7 @@ AtomisticSkills runs all skill scripts and MCP servers through the unified launc | Environment | Primary Packages | GPU | MCP Servers | | :--- | :--- | :--- | :--- | | `cpu` | ASE, pymatgen, RDKit, pycalphad, MDAnalysis (no torch) | No | `base`, `atomate2`, `drugdisc`, `smol` | -| `mlip` | PyTorch 2.14.1, MACE-torch 0.3.16, MatGL 4.1.2, nvalchemi-toolkit 0.2.0 | Yes | `mace`, `matgl` | +| `mlip` | PyTorch 2.14.1, MACE-torch 0.3.16 + cuEquivariance, MatGL 4.1.2, nvalchemi-toolkit 0.2.0 | Yes | `mace`, `matgl` | | `fairchem` | PyTorch 2.13.0, fairchem-core 2.24.0, nvalchemi-toolkit 0.2.0 | Yes | `fairchem` | ### Research Stacks diff --git a/src/mcp_server/mace_server.py b/src/mcp_server/mace_server.py index ff464d01..0b6e5c4d 100644 --- a/src/mcp_server/mace_server.py +++ b/src/mcp_server/mace_server.py @@ -41,6 +41,7 @@ def load_model( model_name: str = "MACE-OMAT-0-small", device: str = "auto", task_name: Optional[str] = None, + accelerator: str = "auto", ) -> str: """ Load a MACE model. @@ -59,9 +60,15 @@ def load_model( task_name: Optional task name that sets the model's head. Supported options for multi-head models (MACE-MH): 'omat_pbe' (default), 'matpes_r2scan', 'omol', 'spice_wB97M', 'oc20_usemppbe'. + accelerator: 'auto' (default) uses CuEquivariance on a CUDA device when + it is installed, and runs the plain e3nn path otherwise — no + GPU, no cuequivariance, or a card its kernels do not cover + all still work, just slower. 'cueq' asks for it and degrades + the same way; 'none' pins the plain path for reproducible + timings. Returns: - Confirmation message. + Confirmation message naming the accelerator actually in use. CRITICAL: This tool must be called before using any other tool to load the model into memory. """ @@ -74,7 +81,12 @@ def load_model( try: from src.utils.mlips.mace.mace_wrapper import MACEWrapper - wrapper = MACEWrapper(model_name=model_name, device=device, head=task_name) + wrapper = MACEWrapper( + model_name=model_name, + device=device, + head=task_name, + accelerator=accelerator, + ) # Load the model - this might take a while and produce output wrapper.load(model_path=model_name if os.path.exists(model_name) else None) @@ -88,7 +100,13 @@ def load_model( f"Model wrapper does not support setting head to {task_name}" ) - return f"Successfully loaded MACE model: {model_name}" + # The calculator, and with it the accelerator, is built lazily on the + # first structure; report what it will attempt rather than forcing a + # build here, which would download and convert the checkpoint. + return ( + f"Successfully loaded MACE model: {model_name}. " + f"Acceleration: {wrapper.describe_accelerator()}." + ) except Exception as e: logger.error(f"Failed to load model {model_name}: {e}") return f"Error loading model: {str(e)}" diff --git a/src/utils/mlips/mace/mace_wrapper.py b/src/utils/mlips/mace/mace_wrapper.py index 39809f80..f5ebf456 100644 --- a/src/utils/mlips/mace/mace_wrapper.py +++ b/src/utils/mlips/mace/mace_wrapper.py @@ -12,6 +12,32 @@ logger = logging.getLogger(__name__) +# How a calculator may be accelerated. CuEquivariance ("cueq") swaps MACE's +# e3nn tensor-product kernels for fused CUDA ones, which needs both its Python +# front end and its compiled ops package, and only helps on a GPU. +# +# "auto" is deliberately forgiving. A host with no GPU, no cuequivariance +# installed (the CPU-only environments, or a run that did not sync the cu130 +# extra), or a card that cuequivariance has no kernels for must still be able +# to run MACE, so an accelerator that cannot be used is a log line and a slower +# path, never an exception out of load(). Asking for "cueq" is a request, not a +# guarantee, for the same reason: it is honoured where it works. +ACCELERATORS = ("auto", "cueq", "none") +_CUEQ_MODULES = ("cuequivariance_torch", "cuequivariance_ops_torch") + + +def cueq_installed() -> bool: + """Whether both halves of CuEquivariance can be imported.""" + from importlib.util import find_spec + + return all(find_spec(name) is not None for name in _CUEQ_MODULES) + + +def _is_cuda(device: str) -> bool: + """Whether an ASE/``torch`` device string names a CUDA device.""" + return str(device).split(":")[0] == "cuda" + + # Available MACE checkpoints and models AVAILABLE_MACE_MODELS = { # MACE-MH (Multi-Head) models - MACE-MH-1 is the latest recommended default @@ -76,6 +102,7 @@ def __init__( model_version: str = "latest", device: str = "auto", head: Optional[str] = None, + accelerator: str = "auto", ): """ Initialize MACE wrapper. @@ -85,10 +112,26 @@ def __init__( model_version: Version of the model device: Device to run the model on ("auto", "cpu", "cuda") head: Optional head name for multi-head models (e.g. "omat_pbe") + accelerator: One of ``ACCELERATORS``. "auto" uses CuEquivariance + where it can help and quietly runs the plain e3nn path + elsewhere; "cueq" asks for it but still degrades; "none" pins + the plain path for reproducible timings. """ + if accelerator not in ACCELERATORS: + raise ValueError( + f"Unknown accelerator {accelerator!r}; expected one of " + f"{', '.join(ACCELERATORS)}." + ) super().__init__(model_name, model_version) self.device = device self.head = head + self.accelerator = accelerator + # What actually ended up on the calculator; set by _build_calculator. + self.accelerator_effective = "none" + self._accelerator_warned = False + # Set once the accelerated attempt has failed on this device, so a + # per-structure rebuild does not retry a conversion that cannot work. + self._accelerator_failed = False self.model_class = None self.calculator_class = None self.model = None @@ -131,6 +174,107 @@ def load(self, model_path: Optional[str] = None) -> None: self.load_checkpoint(model_path) logger.info(f"Loaded model from checkpoint: {model_path}") + def _accepts_accelerator(self) -> bool: + """ + Whether this model's calculator path can take an accelerator at all. + + MACE-ANI-CC is built by mace_anicc, which fixes its own argument list, + so it cannot be handed enable_cueq. Reporting the cuEquivariance + speedup for it would promise a fast path that never runs — the same + mistake as announcing a Nitro fallback for a checkpoint that was never + eligible for Nitro. + """ + return "ANI-CC" not in self.model_name.upper() + + def _accelerator_kwargs(self) -> Dict[str, Any]: + """ + Calculator kwargs for the accelerated path, or empty when it is out. + + Decides once per call and explains the decision the first time it is + a "no", so a user who expected the speedup can tell whether the host + lacks a GPU, the package, or a compatible card. + """ + if self.accelerator == "none" or self._accelerator_failed: + return {} + if not self._accepts_accelerator(): + return {} + if not _is_cuda(self.device): + if not self._accelerator_warned: + logger.info( + "CuEquivariance needs a CUDA device; running %s on %s " + "without it.", + self.model_name, + self.device, + ) + self._accelerator_warned = True + return {} + if not cueq_installed(): + if not self._accelerator_warned: + logger.info( + "CuEquivariance is not installed; running %s without it. " + "It ships with the mlip cu126/cu130 extras.", + self.model_name, + ) + self._accelerator_warned = True + return {} + return {"enable_cueq": True} + + def describe_accelerator(self) -> str: + """ + One line on what the calculator will attempt, and why. + + Cheap and side-effect free: it answers from the device, the installed + packages and the requested mode, without building a calculator or + touching the model. Whether cueq then holds up on the card is only + known when the calculator is built, and _build_calculator logs it. + """ + if self.accelerator == "none": + return "none (disabled by request)" + if not self._accepts_accelerator(): + return f"none ({self.model_name} has no accelerated calculator path)" + if not _is_cuda(self.device): + return f"none (device is {self.device}, and cueq needs CUDA)" + if not cueq_installed(): + return "none (cuequivariance is not installed)" + return "cueq (falls back to the plain path if this card has no kernels)" + + def _build_calculator(self, factory, /, **kwargs): + """ + Build a calculator, preferring the accelerated path. + + MACE raises rather than degrading when cueqvariance is missing or has + no kernels for the card in use, and doing that here would take MACE + away from a user whose GPU merely is not one cuequivariance supports. + So the accelerated attempt is the only thing in the try: if it fails, + rebuild on the plain path and keep the reason. The plain path is not + guarded, so a real error (a bad checkpoint, an unreadable model) still + surfaces instead of being absorbed. + """ + extra = self._accelerator_kwargs() + if extra: + try: + calculator = factory(**kwargs, **extra) + except Exception as exc: + logger.warning( + "CuEquivariance could not run %s on %s (%s: %s); " + "continuing on the plain e3nn path, which is slower.", + self.model_name, + self.device, + type(exc).__name__, + exc, + ) + self.accelerator_effective = "none" + self._accelerator_failed = True + else: + self.accelerator_effective = "cueq" + logger.info( + "CuEquivariance acceleration enabled for %s.", self.model_name + ) + return calculator + else: + self.accelerator_effective = "none" + return factory(**kwargs) + def create_calculator(self) -> Calculator: """ Create an ASE calculator from the model. @@ -145,16 +289,22 @@ def create_calculator(self) -> Calculator: if self.model_path and Path(self.model_path).exists(): from mace.calculators import MACECalculator - return MACECalculator( - model_paths=str(self.model_path), device=self.device, head=self.head + return self._build_calculator( + MACECalculator, + model_paths=str(self.model_path), + device=self.device, + head=self.head, ) # 2. Check if model_name is a path (fallback) if Path(self.model_name).exists(): from mace.calculators import MACECalculator - return MACECalculator( - model_paths=str(self.model_name), device=self.device, head=self.head + return self._build_calculator( + MACECalculator, + model_paths=str(self.model_name), + device=self.device, + head=self.head, ) # 3. Standard Pretrained mapping @@ -199,7 +349,9 @@ def create_calculator(self) -> Calculator: ): head = "omat_pbe" - return MACECalculator(model_paths=model_path, device=self.device, head=head) + return self._build_calculator( + MACECalculator, model_paths=model_path, device=self.device, head=head + ) if "OFF" in model_name_upper: # Use MACE-OFF calculator @@ -211,25 +363,29 @@ def create_calculator(self) -> Calculator: "MACE-OFF23-LARGE": "large", } off_model_size = off_model_mapping.get(model_name_upper, "medium") - return mace_off(model=off_model_size, device=self.device) + return self._build_calculator( + mace_off, model=off_model_size, device=self.device + ) elif "ANI-CC" in model_name_upper: # Use MACE-ANI-CC from mace.calculators import mace_anicc + # mace_anicc takes a fixed argument list, so it cannot be handed + # enable_cueq the way the other helpers can. return mace_anicc(device=self.device) elif "OMOL" in model_name_upper: # Use MACE-OMOL from mace.calculators import mace_omol - return mace_omol(model=model_id, device=self.device) + return self._build_calculator(mace_omol, model=model_id, device=self.device) else: # Use MACE-MP calculator for materials models (including MPA, MATPES, OMAT) from mace.calculators import mace_mp - return mace_mp(model=model_id, device=self.device) + return self._build_calculator(mace_mp, model=model_id, device=self.device) def predict_atomic_features(self, structure_data: Any) -> Dict[str, Any]: """ diff --git a/tests/mace/test_mace_accelerator.py b/tests/mace/test_mace_accelerator.py new file mode 100644 index 00000000..89fa3d0c --- /dev/null +++ b/tests/mace/test_mace_accelerator.py @@ -0,0 +1,167 @@ +""" +Tests for MACE accelerator selection and its fallbacks. + +CuEquivariance is an optimisation, never a requirement: a CPU-only host, an +environment that did not sync the cu130 extra, or a card its kernels do not +cover must all still be able to run MACE. These tests pin that behaviour with +a fake calculator factory, so they need no GPU, no cuequivariance install and +no checkpoint download. + +Run in: venv/run mlip python -m pytest tests/mace/test_mace_accelerator.py +""" + +import logging + +import pytest + +from src.utils.mlips.mace import mace_wrapper as mw + + +def fake_calculator(**kwargs): + """Stand-in for a MACE calculator that records how it was built.""" + return {"kwargs": kwargs} + + +@pytest.fixture +def wrapper_cuda(): + """A wrapper pointed at CUDA, without touching any hardware.""" + w = mw.MACEWrapper(model_name="MACE-MH-1", device="cuda") + w.is_loaded = True + return w + + +@pytest.mark.mace +class TestAcceleratorSelection: + def test_auto_on_cpu_never_requests_cueq(self): + w = mw.MACEWrapper(model_name="MACE-MH-1", device="cpu") + built = w._build_calculator(fake_calculator, device="cpu") + assert "enable_cueq" not in built["kwargs"] + assert w.accelerator_effective == "none" + + def test_auto_without_cueq_installed_degrades(self, wrapper_cuda, monkeypatch): + monkeypatch.setattr(mw, "cueq_installed", lambda: False) + built = wrapper_cuda._build_calculator(fake_calculator, device="cuda") + assert "enable_cueq" not in built["kwargs"] + assert wrapper_cuda.accelerator_effective == "none" + + def test_auto_with_cueq_requests_it(self, wrapper_cuda, monkeypatch): + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + built = wrapper_cuda._build_calculator(fake_calculator, device="cuda") + assert built["kwargs"]["enable_cueq"] is True + assert wrapper_cuda.accelerator_effective == "cueq" + + def test_none_never_attempts_even_when_installed(self, monkeypatch): + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + w = mw.MACEWrapper(model_name="MACE-MH-1", device="cuda", accelerator="none") + built = w._build_calculator(fake_calculator, device="cuda") + assert "enable_cueq" not in built["kwargs"] + assert w.accelerator_effective == "none" + + def test_explicit_cueq_request_still_degrades(self, wrapper_cuda, monkeypatch): + """Asking for cueq is a request, not a promise.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: False) + w = mw.MACEWrapper(model_name="MACE-MH-1", device="cuda", accelerator="cueq") + built = w._build_calculator(fake_calculator, device="cuda") + assert "enable_cueq" not in built["kwargs"] + assert w.accelerator_effective == "none" + + def test_unknown_accelerator_rejected(self): + with pytest.raises(ValueError, match="Unknown accelerator"): + mw.MACEWrapper(model_name="MACE-MH-1", accelerator="fastest") + + +@pytest.mark.mace +class TestAcceleratorFallback: + def test_incompatible_card_degrades_instead_of_raising( + self, wrapper_cuda, monkeypatch, caplog + ): + """A GPU cuequivariance has no kernels for must not take MACE down.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + + def factory(**kwargs): + if kwargs.get("enable_cueq"): + raise RuntimeError("CUDA error: no kernel image is available") + return fake_calculator(**kwargs) + + with caplog.at_level(logging.WARNING): + built = wrapper_cuda._build_calculator(factory, device="cuda") + + assert "enable_cueq" not in built["kwargs"] + assert wrapper_cuda.accelerator_effective == "none" + assert any("no kernel image" in r.getMessage() for r in caplog.records) + + def test_missing_module_degrades(self, wrapper_cuda, monkeypatch): + """mace raises ImportError when cuequivariance is absent; absorb it.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + + def factory(**kwargs): + if kwargs.get("enable_cueq"): + raise ImportError("cuequivariance is not installed") + return fake_calculator(**kwargs) + + built = wrapper_cuda._build_calculator(factory, device="cuda") + assert "enable_cueq" not in built["kwargs"] + + def test_real_error_on_the_plain_path_surfaces(self, wrapper_cuda, monkeypatch): + """The plain path is unguarded, so genuine bugs still raise.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + + def factory(**kwargs): + raise TypeError("model_paths must be a path") + + with pytest.raises(TypeError, match="model_paths"): + wrapper_cuda._build_calculator(factory, device="cuda") + + def test_fallback_is_remembered_for_later_calls(self, wrapper_cuda, monkeypatch): + """One failed attempt must not cost a rebuild on every structure.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + attempts = [] + + def factory(**kwargs): + attempts.append(kwargs.get("enable_cueq", False)) + if kwargs.get("enable_cueq"): + raise RuntimeError("no kernel image") + return fake_calculator(**kwargs) + + wrapper_cuda._build_calculator(factory, device="cuda") + wrapper_cuda._build_calculator(factory, device="cuda") + assert attempts == [True, False, False] + + +@pytest.mark.mace +class TestModelsWithoutAcceleratedPath: + def test_ani_cc_never_requests_cueq(self, monkeypatch): + """mace_anicc fixes its own arguments, so cueq cannot be passed.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + w = mw.MACEWrapper(model_name="MACE-ANI-CC", device="cuda") + assert w._accelerator_kwargs() == {} + assert w.accelerator_effective == "none" + + def test_ani_cc_does_not_promise_the_speedup(self, monkeypatch): + """Do not announce an accelerator the model can never use.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + w = mw.MACEWrapper(model_name="MACE-ANI-CC", device="cuda") + described = w.describe_accelerator() + assert described.startswith("none") + assert "no accelerated calculator path" in described + + def test_other_models_do_accept_it(self, monkeypatch): + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + for name in ("MACE-MH-1", "MACE-OFF23-medium", "MACE-OMOL-extra-large"): + w = mw.MACEWrapper(model_name=name, device="cuda") + assert w._accelerator_kwargs() == {"enable_cueq": True}, name + + +@pytest.mark.mace +class TestDescribeAccelerator: + def test_reports_cpu_as_none(self): + w = mw.MACEWrapper(model_name="MACE-MH-1", device="cpu") + assert "CPU" in w.describe_accelerator() or "cpu" in w.describe_accelerator() + + def test_reports_missing_package(self, wrapper_cuda, monkeypatch): + monkeypatch.setattr(mw, "cueq_installed", lambda: False) + assert "not installed" in wrapper_cuda.describe_accelerator() + + def test_reports_cueq_available(self, wrapper_cuda, monkeypatch): + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + assert wrapper_cuda.describe_accelerator().startswith("cueq") diff --git a/tests/test_uv_projects.py b/tests/test_uv_projects.py index 13285a86..87ea94bf 100644 --- a/tests/test_uv_projects.py +++ b/tests/test_uv_projects.py @@ -125,7 +125,30 @@ def test_extras_are_the_same_everywhere(project): # from each (build support in mlip, the LAMMPS wheel in fairchem); everything # else matches cpu. if project in ("mlip", "fairchem"): - assert extras.pop("cu126") == extras.pop("cu130") == ["torch", "torchvision"] + cu126, cu130 = extras.pop("cu126"), extras.pop("cu130") + # torch and torchvision are the same two names in both flavours, told + # apart by their index (see test_gpu_projects_lock_both_torch_cuda_builds). + assert cu126[:2] == cu130[:2] == ["torch", "torchvision"] + if project == "mlip": + # cuEquivariance is the one dependency whose wheel is compiled per + # CUDA major under a different *distribution name* + # (…-ops-torch-cu12 / …-ops-torch-cu13) with no unsuffixed + # meta-package, so unlike torch it cannot be switched by index. + # Spell both halves out: anything else appearing here, or the two + # flavours drifting apart, still fails. + def split_cueq(extra): + tail = extra[2:] + ops = [e for e in tail if e.startswith("cuequivariance-ops-torch-cu")] + assert len(ops) == 1, extra + return set(tail) - set(ops), ops[0] + + shared126, ops126 = split_cueq(cu126) + shared130, ops130 = split_cueq(cu130) + assert shared126 == shared130 == {"cuequivariance-torch>=0.2.0"} + assert ops126 == "cuequivariance-ops-torch-cu12>=0.2.0" + assert ops130 == "cuequivariance-ops-torch-cu13>=0.2.0" + else: + assert cu126 == cu130 == ["torch", "torchvision"] assert extras.pop("lammps"), project assert extras == load("cpu")["project"]["optional-dependencies"] assert set(extras) == {"openmm", "pymol", "docking", "void", "transport"} diff --git a/venv/README.md b/venv/README.md index c8aa633c..b942cc23 100644 --- a/venv/README.md +++ b/venv/README.md @@ -51,7 +51,7 @@ Floors (not strict pins) keep the scientific packages of the shared projects cur | Package | mlip | fairchem | cpu | | :--- | :--- | :--- | :--- | | torch | 2.14.1 (CUDA 12.6 or 13) | 2.13.0 (CUDA 12.6 or 13) | — | -| mace-torch / e3nn | 0.3.16 / 0.4.4 | — | — | +| mace-torch / e3nn | 0.3.16 / 0.4.4 + cuEquivariance (cu126/cu130 extras) | — | — | | matgl | 4.1.2 (PyG only; DGL removed) | — | — | | fairchem-core / e3nn | — | 2.24.0 / 0.6.0 | — | | nvalchemi-toolkit | 0.2.0 | 0.2.0 | — | diff --git a/venv/mlip/pyproject.toml b/venv/mlip/pyproject.toml index d3d2328a..19013a13 100644 --- a/venv/mlip/pyproject.toml +++ b/venv/mlip/pyproject.toml @@ -137,8 +137,18 @@ lammps = ["cython", "mkl-devel ; platform_machine == 'x86_64'"] # torchvision (via timm, physicsnemo) is compiled against one torch build, so it # comes from the same index; a PyPI torchvision with a cu126 torch fails with # "operator torchvision::nms does not exist". -cu126 = ["torch", "torchvision"] -cu130 = ["torch", "torchvision"] +# cuEquivariance ships the same pair of packages for both CUDA majors, so it +# rides along with the torch build it must match: the ops wheels are compiled +# against one CUDA major, and cuequivariance_torch imports them under the bare +# module name cuequivariance_ops_torch either way. Upstream mace-torch declares +# only the cu12 pair (cueq-cuda-11/12); cu13 exists and is what a driver >= 580 +# host needs. Roughly 30 MB per architecture. +# +# Installing it is not the same as using it: MaceWrapper enables CuEquivariance +# only on a CUDA device where the kernels import, and falls back to the plain +# e3nn path otherwise (see _build_calculator). +cu126 = ["torch", "torchvision", "cuequivariance-torch>=0.2.0", "cuequivariance-ops-torch-cu12>=0.2.0"] +cu130 = ["torch", "torchvision", "cuequivariance-torch>=0.2.0", "cuequivariance-ops-torch-cu13>=0.2.0"] [dependency-groups] dev = ["pytest", "pytest-timeout"] diff --git a/venv/mlip/uv.lock b/venv/mlip/uv.lock index 6253698e..522378ca 100644 --- a/venv/mlip/uv.lock +++ b/venv/mlip/uv.lock @@ -345,10 +345,14 @@ dependencies = [ [package.optional-dependencies] cu126 = [ + { name = "cuequivariance-ops-torch-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "cuequivariance-torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "torch", version = "2.14.1+cu126", source = { registry = "https://download.pytorch.org/whl/cu126" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "torchvision", version = "0.29.1+cu126", source = { registry = "https://download.pytorch.org/whl/cu126" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] cu130 = [ + { name = "cuequivariance-ops-torch-cu13", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "cuequivariance-torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "torch", version = "2.14.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "torchvision", version = "0.29.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] @@ -401,6 +405,10 @@ requires-dist = [ { name = "beautifulsoup4" }, { name = "biopython" }, { name = "cachetools", marker = "extra == 'openmm'" }, + { name = "cuequivariance-ops-torch-cu12", marker = "extra == 'cu126'", specifier = ">=0.2.0" }, + { name = "cuequivariance-ops-torch-cu13", marker = "extra == 'cu130'", specifier = ">=0.2.0" }, + { name = "cuequivariance-torch", marker = "extra == 'cu126'", specifier = ">=0.2.0" }, + { name = "cuequivariance-torch", marker = "extra == 'cu130'", specifier = ">=0.2.0" }, { name = "cvxpy" }, { name = "cython", marker = "extra == 'lammps'" }, { name = "dara-xrd" }, @@ -1123,6 +1131,92 @@ nvtx = [ { name = "nvidia-nvtx", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-25-atomisticskills-venv-mlip-cu126' and extra == 'extra-25-atomisticskills-venv-mlip-cu130') or (platform_machine == 'aarch64' and sys_platform == 'linux' and extra != 'extra-25-atomisticskills-venv-mlip-cu126') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-25-atomisticskills-venv-mlip-cu126') or (platform_machine == 'aarch64' and extra == 'extra-25-atomisticskills-venv-mlip-cu126' and extra == 'extra-25-atomisticskills-venv-mlip-cu130') or (platform_machine == 'x86_64' and extra == 'extra-25-atomisticskills-venv-mlip-cu126' and extra == 'extra-25-atomisticskills-venv-mlip-cu130') or (sys_platform != 'linux' and extra == 'extra-25-atomisticskills-venv-mlip-cu126' and extra == 'extra-25-atomisticskills-venv-mlip-cu130')" }, ] +[[package]] +name = "cuequivariance" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "networkx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opt-einsum", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "scipy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "sympy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ee/91/c78169c11dfe89c8a200f9eab178b7f637dde39673be6815af3007fa8922/cuequivariance-0.13.0.tar.gz", hash = "sha256:4227705313bdd89650077c5b723416381e76add1bec4aa103a26347946f53efd", size = 219499, upload-time = "2026-10-15T19:53:14.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/df/e1/40915811975b95dc31dde5c227e3a1e4a6af264cf01b32cd2ea05d2f0e58/cuequivariance-0.13.0-py3-none-any.whl", hash = "sha256:72cc774e2f74bb2666c431fb6b5e6e612317474d38b9114240fc78585f8c1f81", size = 250801, upload-time = "2026-10-15T19:53:13.402Z" }, +] + +[[package]] +name = "cuequivariance-ops-cu12" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-ml-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "platformdirs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/a1/4ab8138424855d55c4e123132d60322987eb0f0177ff4b41dd72fe7ccb2a/cuequivariance_ops_cu12-0.13.0-py3-none-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7be711daa7313f37ed1bff288f4edd718d006ed84429ab62cde90843d639f55", size = 41111805, upload-time = "2026-10-15T19:22:07.153Z" }, + { url = "https://files.pythonhosted.org/packages/5b/30/db90ba304fc07b300c87ff1593acffcd39a17162f5d73359ba8d7171f425/cuequivariance_ops_cu12-0.13.0-py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:54fcc31a63b0768946d6ccb83580796dec50de6c795a78022436ca9f47e1bdd2", size = 41168895, upload-time = "2026-10-15T19:22:11.354Z" }, +] + +[[package]] +name = "cuequivariance-ops-cu13" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-ml-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "platformdirs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/29/ff/7d956b9bb5cee909f7ac2b653fcf69572bb0658a808b556d2970b16415bc/cuequivariance_ops_cu13-0.13.0-py3-none-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:62d267f500919231f964384762e4eddef2f5bd260608e3932cc3f95b84249fc0", size = 44889148, upload-time = "2026-10-15T19:22:15.7Z" }, + { url = "https://files.pythonhosted.org/packages/7b/07/f5319a6039b9ef5d05ee11125ad2378e4da099117fd35b24e24a787dc960/cuequivariance_ops_cu13-0.13.0-py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:12b0d8b736bb195447946771a0092d69f9c597789a917d8609dcaa16220ffa34", size = 44958162, upload-time = "2026-10-15T19:22:20.411Z" }, +] + +[[package]] +name = "cuequivariance-ops-torch-cu12" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuequivariance-ops-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "scipy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", version = "2.14.1+cu126", source = { registry = "https://download.pytorch.org/whl/cu126" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/0e/dd/c3cbedcaef043456154df7226117c72d1db3c2fa343d9702ab6510250728/cuequivariance_ops_torch_cu12-0.13.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5c4c42f16380639e4282725e204c72643350af3246b5335bb8d42fb5e7426f58", size = 1193257, upload-time = "2026-10-15T19:22:32.672Z" }, + { url = "https://files.pythonhosted.org/packages/f6/c2/bdc3ca4ac9560dca4b70aeaa5f39ef19a03f80a3ce9abcbd39ef2a9aa3ac/cuequivariance_ops_torch_cu12-0.13.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:909d49ac15a4c132b110a20ae755a9f956d2466a19c2f36e0274e03a9558784a", size = 1202213, upload-time = "2026-10-15T19:22:33.875Z" }, +] + +[[package]] +name = "cuequivariance-ops-torch-cu13" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuequivariance-ops-cu13", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "scipy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", version = "2.14.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/bd/7740b6bd8ab680d3ba9edb7d6bc90bcee60aa2f4e2bff4ed4ef1dca53f33/cuequivariance_ops_torch_cu13-0.13.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:28121445776a9a268f55a5a38c9f48f76512cdd051300a7574d9c2eaf82d7507", size = 28512268, upload-time = "2026-10-15T19:23:25.913Z" }, + { url = "https://files.pythonhosted.org/packages/c4/8d/176e83eb4fb411e6b5ff5bda06e5a7799bc6369aea19112277319ced1f8a/cuequivariance_ops_torch_cu13-0.13.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:181e22ab2644b2ac0020aa0654ff5a1468b99d7d6409a8a82269dd6e8941ca09", size = 28530554, upload-time = "2026-10-15T19:23:33.065Z" }, +] + +[[package]] +name = "cuequivariance-torch" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuequivariance", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/96/fd/221132e3f204a5bfe4efbc531242bbbb84a568e8aa834648a371a47f969b/cuequivariance_torch-0.13.0.tar.gz", hash = "sha256:f5ffb513454d64cd7cf29915153ae029a27538ead7672c552dab0f005984141e", size = 170540, upload-time = "2026-10-15T19:53:20.036Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e3/1b/975e2b9703ead34b24c3b3226074255e0ef02b171ee8ae98921094a57783/cuequivariance_torch-0.13.0-py3-none-any.whl", hash = "sha256:f256fe9ab5c3a8f37e08345ba3c734d37a6ccd1a59b176479d61a4126c56acd0", size = 184761, upload-time = "2026-10-15T19:53:19.001Z" }, +] + [[package]] name = "custodian" version = "2025.12.14" @@ -3441,6 +3535,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/31/83/f3647ce26916c94a6ca4ff1810623e2c405cff2dea6e78d29516b2514df9/nvidia_cusparselt_cu13-0.8.1-py3-none-win_amd64.whl", hash = "sha256:dccbd362f91a7b9024d1f55ee9f548ac065027ff15d8c8b0db889ab3a8f31215", size = 156885108, upload-time = "2025-09-05T18:51:35.958Z" }, ] +[[package]] +name = "nvidia-ml-py" +version = "13.615.71" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/30/b25216758be3d3e2834825d8193609e2d71c770a8bd7984438c058c90268/nvidia_ml_py-13.615.71.tar.gz", hash = "sha256:bebe4e48f51b1dc75028c0815cb7bfa14a31a5bb80be70c9d980c6036953fc3d", size = 57485, upload-time = "2026-09-25T15:15:28.226Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/53/a1/1681dfa1c904d4e3e72e51b55a0ff012d50b766843ef832d584abe2113c6/nvidia_ml_py-13.615.71-py3-none-any.whl", hash = "sha256:959bf4adf6fe1308e4bd739e722236b0d1ec8392e2cefad33ff70c311380b9b6", size = 58132, upload-time = "2026-09-25T15:15:26.54Z" }, +] + [[package]] name = "nvidia-nccl-cu12" version = "2.29.3"