rootstock 0.5.0__py3-none-any.whl

rootstock/worker.py

@@ -0,0 +1,273 @@
+#!/usr/bin/env python
+"""
+Rootstock worker process.
+
+This runs in an isolated subprocess and:
+1. Loads an MLIP (e.g., MACE)
+2. Connects to the server via Unix socket
+3. Receives positions, calculates forces, sends results back
+4. Persists across multiple calculations (no startup overhead per calculation)
+
+The worker is spawned via a generated wrapper script that calls run_worker().
+"""
+
+import json
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from .protocol import (
+    IPIProtocol,
+    SocketClosed,
+    connect_unix_socket,
+    create_unix_socket_path,
+)
+
+if TYPE_CHECKING:
+    from ase.calculators.calculator import Calculator
+
+
+class MLIPWorker:
+    """
+    Worker that runs an MLIP calculator and communicates via i-PI protocol.
+
+    The worker acts as an i-PI client:
+    1. Connect to server
+    2. Report READY status
+    3. Receive positions via POSDATA
+    4. Calculate energy/forces
+    5. Report HAVEDATA status
+    6. Send results via FORCEREADY
+    7. Loop back to step 2
+    """
+
+    def __init__(
+        self,
+        socket_name: str,
+        calculator: "Calculator",
+        log=None,
+    ):
+        """
+        Initialize the worker.
+
+        Args:
+            socket_name: Name of Unix socket to connect to
+            calculator: Pre-loaded ASE calculator
+            log: Optional file object for logging
+        """
+        self.socket_name = socket_name
+        self.socket_path = create_unix_socket_path(socket_name)
+        self.log = log
+
+        self._calculator = calculator
+        self._socket = None
+        self._protocol = None
+        self._atoms = None  # Cache ASE Atoms object
+
+        # Atomic species info from INIT message
+        self._atomic_numbers: list[int] | None = None
+        self._pbc: list[bool] | None = None
+
+    def _log(self, msg):
+        if self.log:
+            print(f"[Worker] {msg}", file=self.log, flush=True)
+
+    def _connect(self):
+        """Connect to the server."""
+        self._log(f"Connecting to {self.socket_path}")
+        self._socket = connect_unix_socket(self.socket_path)
+        self._protocol = IPIProtocol(self._socket, log=self.log)
+        self._log("Connected")
+
+    def _create_atoms(self, positions: np.ndarray, cell: np.ndarray):
+        """
+        Create or update ASE Atoms object.
+
+        On first call, creates a new Atoms object.
+        On subsequent calls, updates positions and cell in place.
+        """
+        from ase import Atoms
+
+        if self._atoms is None or len(self._atoms) != len(positions):
+            # Need to create new Atoms object
+            if self._atomic_numbers is None:
+                raise RuntimeError(
+                    "No atomic numbers received. Server must send INIT with species data."
+                )
+
+            self._atoms = Atoms(
+                numbers=self._atomic_numbers,
+                positions=positions,
+                cell=cell,
+                pbc=self._pbc if self._pbc is not None else [True, True, True],
+            )
+            self._atoms.calc = self._calculator
+        else:
+            # Update existing object (faster - reuses neighbor lists etc.)
+            self._atoms.positions = positions
+            self._atoms.cell = cell
+
+        return self._atoms
+
+    def _calculate(
+        self, positions: np.ndarray, cell: np.ndarray
+    ) -> tuple[float, np.ndarray, np.ndarray]:
+        """
+        Run MLIP calculation.
+
+        Returns:
+            energy: Potential energy in eV
+            forces: Nx3 forces in eV/Angstrom
+            virial: 3x3 virial tensor in eV
+        """
+        atoms = self._create_atoms(positions, cell)
+
+        energy = atoms.get_potential_energy()
+        forces = atoms.get_forces()
+
+        # Calculate virial from stress
+        # stress is in eV/ų, virial = -stress * volume
+        try:
+            stress = atoms.get_stress(voigt=False)  # 3x3 tensor
+            volume = atoms.get_volume()
+            virial = -stress * volume
+        except Exception:
+            # Some calculators don't support stress
+            virial = np.zeros((3, 3))
+
+        return energy, forces, virial
+
+    def run(self):
+        """
+        Main loop - receive positions, calculate, send results.
+
+        This implements the i-PI client state machine:
+        - NEEDINIT -> receive INIT -> READY
+        - READY -> receive POSDATA -> calculate -> HAVEDATA
+        - HAVEDATA -> receive GETFORCE -> send FORCEREADY -> NEEDINIT
+        """
+        self._connect()
+
+        state = "NEEDINIT"
+        energy = None
+        forces = None
+        virial = None
+
+        self._log("Entering main loop")
+
+        try:
+            while True:
+                # Wait for message from server
+                try:
+                    msg = self._protocol.recvmsg()
+                except SocketClosed:
+                    self._log("Server closed connection")
+                    break
+
+                if msg == "EXIT":
+                    self._log("Received EXIT, shutting down")
+                    break
+
+                elif msg == "STATUS":
+                    # Report current state
+                    if state == "NEEDINIT":
+                        self._protocol.sendmsg("NEEDINIT")
+                    elif state == "READY":
+                        self._protocol.sendmsg("READY")
+                    elif state == "HAVEDATA":
+                        self._protocol.sendmsg("HAVEDATA")
+
+                elif msg == "INIT":
+                    # Receive initialization with atomic species info
+                    bead_index, init_bytes = self._protocol.recv_init()
+
+                    # Parse JSON from init_bytes
+                    if init_bytes and init_bytes != b"\x00":
+                        try:
+                            init_data = json.loads(init_bytes.decode("utf-8"))
+                            self._atomic_numbers = init_data.get("numbers")
+                            self._pbc = init_data.get("pbc", [True, True, True])
+                            self._log(
+                                f"Received INIT (bead={bead_index}, "
+                                f"atoms={len(self._atomic_numbers) if self._atomic_numbers else 0})"
+                            )
+                        except (json.JSONDecodeError, UnicodeDecodeError) as e:
+                            self._log(f"Warning: Failed to parse INIT data: {e}")
+                    else:
+                        self._log(f"Received INIT (bead={bead_index}, no species data)")
+
+                    state = "READY"
+
+                elif msg == "POSDATA":
+                    # Receive atomic positions
+                    if state not in ("READY", "NEEDINIT"):
+                        self._log(f"Warning: POSDATA in state {state}")
+
+                    cell, positions = self._protocol.recv_posdata()
+                    self._log(f"Received POSDATA: {len(positions)} atoms")
+
+                    # Calculate energy and forces
+                    energy, forces, virial = self._calculate(positions, cell)
+                    self._log(f"Calculated: E={energy:.6f} eV")
+
+                    state = "HAVEDATA"
+
+                elif msg == "GETFORCE":
+                    # Send results
+                    if state != "HAVEDATA":
+                        raise RuntimeError(f"GETFORCE in state {state}")
+
+                    self._protocol.send_forceready(energy, forces, virial)
+                    self._log("Sent FORCEREADY")
+
+                    state = "NEEDINIT"
+
+                else:
+                    self._log(f"Unknown message: {msg}")
+
+        finally:
+            if self._socket:
+                self._socket.close()
+            self._log("Worker shutdown complete")
+
+
+def run_worker(
+    setup_fn: Callable[[str, str], "Calculator"],
+    model: str,
+    device: str,
+    socket_path: str,
+    log=None,
+):
+    """
+    Run worker with a provided setup function.
+
+    This is the entry point used by generated wrapper scripts.
+    The setup function is called once to create the calculator, which
+    is then reused for all subsequent calculations.
+
+    Args:
+        setup_fn: Function that takes (model, device) and returns an ASE calculator
+        model: Model identifier to pass to setup_fn
+        device: Device string to pass to setup_fn
+        socket_path: Full Unix socket path to connect to
+        log: Optional logging file object
+    """
+    if log:
+        print(f"[Worker] Calling setup({model!r}, {device!r})", file=log, flush=True)
+
+    # Load calculator via the setup function
+    calculator = setup_fn(model, device)
+
+    if log:
+        print(f"[Worker] Calculator loaded: {type(calculator).__name__}", file=log, flush=True)
+
+    # Extract socket name from path (e.g., /tmp/ipi_rootstock_abc -> rootstock_abc)
+    socket_name = socket_path.replace("/tmp/ipi_", "")
+
+    worker = MLIPWorker(
+        socket_name=socket_name,
+        calculator=calculator,
+        log=log,
+    )
+    worker.run()

rootstock-0.5.0.dist-info/METADATA

@@ -0,0 +1,210 @@
+Metadata-Version: 2.4
+Name: rootstock
+Version: 0.5.0
+Summary: MLIP calculators with isolated Python environments
+License-File: LICENSE.md
+Requires-Python: >=3.10
+Requires-Dist: ase>=3.22
+Requires-Dist: numpy>=1.24
+Requires-Dist: packaging>=21.0
+Requires-Dist: tomli>=2.0; python_version < '3.11'
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Requires-Dist: ruff>=0.1; extra == 'dev'
+Provides-Extra: mace
+Requires-Dist: mace-torch>=0.3; extra == 'mace'
+Requires-Dist: torch>=2.0; extra == 'mace'
+Provides-Extra: modal
+Requires-Dist: modal>=0.56; extra == 'modal'
+Description-Content-Type: text/markdown
+
+# Rootstock
+
+Run MLIP (Machine Learning Interatomic Potential) calculators in isolated pre-built Python environments, communicating via the i-PI protocol over Unix sockets.
+
+## Quick Start
+
+```python
+from ase.build import bulk
+from rootstock import RootstockCalculator
+
+atoms = bulk("Cu", "fcc", a=3.6) * (5, 5, 5)
+
+# Using a known cluster
+with RootstockCalculator(
+    cluster="modal",       # or "della"
+    model="mace-medium",   # or "chgnet", "mace-small", etc.
+    device="cuda",
+) as calc:
+    atoms.calc = calc
+    print(atoms.get_potential_energy())
+    print(atoms.get_forces())
+
+# Or with an explicit root path
+with RootstockCalculator(
+    root="/scratch/gpfs/SHARED/rootstock",
+    model="mace-medium",
+    device="cuda",
+) as calc:
+    atoms.calc = calc
+    print(atoms.get_potential_energy())
+```
+
+**Note:** Environments must be pre-built before use. See [Administrator Setup](#administrator-setup).
+
+## Installation
+
+```bash
+pip install rootstock
+# or
+uv pip install rootstock
+```
+
+## Model String Format
+
+The `model` parameter encodes both the environment and model-specific argument:
+
+| `model=`            | Environment    | Model Arg           |
+|---------------------|----------------|---------------------|
+| `"mace-medium"`     | mace_env       | `"medium"`          |
+| `"mace-small"`      | mace_env       | `"small"`           |
+| `"mace-large"`      | mace_env       | `"large"`           |
+| `"chgnet"`          | chgnet_env     | `""` (default)      |
+| `"mace-/path/to/weights.pt"` | mace_env | `"/path/to/weights.pt"` |
+
+## Known Clusters
+
+| Cluster | Root Path |
+|---------|-----------|
+| `modal` | `/vol/rootstock` |
+| `della` | `/scratch/gpfs/SHARED/rootstock` |
+
+For other clusters, use `root="/path/to/rootstock"` directly.
+
+## Administrator Setup
+
+Environments must be pre-built before users can run calculations.
+
+### 1. Create Directory Structure
+
+```bash
+mkdir -p /scratch/gpfs/SHARED/rootstock/{environments,envs,cache}
+```
+
+### 2. Create Environment Source Files
+
+```bash
+# mace_env.py
+cat > /scratch/gpfs/SHARED/rootstock/environments/mace_env.py << 'EOF'
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["mace-torch>=0.3.0", "ase>=3.22", "torch>=2.0"]
+# ///
+"""MACE environment for Rootstock."""
+
+def setup(model: str, device: str = "cuda"):
+    from mace.calculators import mace_mp
+    return mace_mp(model=model, device=device, default_dtype="float32")
+EOF
+```
+
+### 3. Build Environments
+
+```bash
+# Build MACE environment with model pre-download
+rootstock build mace_env --root /scratch/gpfs/SHARED/rootstock --models small,medium,large
+
+# Build CHGNet environment
+rootstock build chgnet_env --root /scratch/gpfs/SHARED/rootstock
+
+# Verify
+rootstock status --root /scratch/gpfs/SHARED/rootstock
+```
+
+## Architecture
+
+```
+Main Process                          Worker Process (subprocess)
++-------------------------+          +-----------------------------+
+| RootstockCalculator     |          | Pre-built venv Python       |
+| (ASE-compatible)        |          | (mace_env/bin/python)       |
+|                         |          |                             |
+| server.py (i-PI server) |<-------->| worker.py (i-PI client)     |
+| - sends positions       |   Unix   | - receives positions        |
+| - receives forces       |  socket  | - calculates forces         |
++-------------------------+          +-----------------------------+
+```
+
+The worker process uses a pre-built virtual environment, providing:
+- **Fast startup**: No dependency installation at runtime
+- **Filesystem compatibility**: Works on NFS, Lustre, GPFS, Modal volumes
+- **Reproducibility**: Same environment every time
+
+## Directory Structure
+
+```
+{root}/
+├── environments/           # Environment SOURCE files (*.py with PEP 723)
+│   ├── mace_env.py
+│   └── chgnet_env.py
+├── envs/                   # Pre-built virtual environments
+│   ├── mace_env/
+│   │   ├── bin/python
+│   │   ├── lib/python3.11/site-packages/
+│   │   └── env_source.py   # Copy of environment source
+│   └── chgnet_env/
+└── cache/                  # XDG_CACHE_HOME for model weights
+    ├── mace/               # MACE models
+    └── huggingface/        # HuggingFace models
+```
+
+## CLI Commands
+
+```bash
+# Build a pre-built environment
+rootstock build <env_name> --root <path> [--models m1,m2] [--force]
+
+# Show status
+rootstock status --root <path>
+
+# Register an environment source file
+rootstock register <env_file> --root <path>
+
+# List environments
+rootstock list --root <path>
+```
+
+## Running on Modal
+
+```bash
+# Initialize volume and build environments (takes ~10-15 min)
+modal run modal_app.py::init_rootstock_volume
+
+# Test pre-built environments
+modal run modal_app.py::test_prebuilt
+
+# Show status
+modal run modal_app.py::inspect_status
+
+# Run benchmarks
+modal run modal_app.py::benchmark_v4
+```
+
+## Performance
+
+IPC overhead is <5% for systems with 1000+ atoms compared to direct in-process execution.
+
+| System Size | Atoms | Typical Overhead |
+|-------------|-------|------------------|
+| Small       | 64    | ~10-15%          |
+| Medium      | 256   | ~5-8%            |
+| Large       | 1000  | <5%              |
+
+## Local Development
+
+```bash
+uv venv && source .venv/bin/activate
+uv pip install -e ".[dev]"
+ruff check rootstock/
+ruff format rootstock/
+```

rootstock-0.5.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+rootstock/__init__.py,sha256=Qu570ne-AeWn6IT4Us00iU43A0yoRlyYDygYvKHduVQ,956
+rootstock/calculator.py,sha256=0KpW3coJ0ealEIZkxq9i4kZV5AxqCX1xZqDI1HkUPA4,6257
+rootstock/cli.py,sha256=6YtUXs38Zc5BC0AjEaAygtmct6JAsfaGTX7Zn2K_imY,13815
+rootstock/clusters.py,sha256=JQdPgBl_eyyP3nWxg_JOvD0Kh7YTVVXsNq038ZusVJg,1207
+rootstock/environment.py,sha256=e7AAKd7cECwnX-KHCdfelHYVqm3xa1mIvc2MGaTJwBs,6583
+rootstock/pep723.py,sha256=hgTcP7Tp8_z_DCrmX4VPAitijvujmbGmz2ng1VgNYks,4616
+rootstock/protocol.py,sha256=CgfYNc0aKOkAQ-D8tvUrVO5q1mg0Xwcl2F02ZvYOJCI,10169
+rootstock/server.py,sha256=H7kn1FHGGyrrPZtgRKdfagbAF-SxI89H3u85GSToYmY,9210
+rootstock/worker.py,sha256=ty13OPDcKd_Zy95MAnTwa063Glz_RfUTBDja9Lxy4SM,8961
+rootstock-0.5.0.dist-info/METADATA,sha256=AD3xqByFrIWK-j8LdWfAxJXFKvbsdgro7xX1x9PPp-Y,5993
+rootstock-0.5.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+rootstock-0.5.0.dist-info/entry_points.txt,sha256=rPiVll-qj1wq7wZQTwSl2aF22GLnnpDz37hkPLvqlh0,49
+rootstock-0.5.0.dist-info/licenses/LICENSE.md,sha256=ORJAYeKSWpOYZ89KWT8ETWFb2u6MvKK3AhrMReDMWrA,1072
+rootstock-0.5.0.dist-info/RECORD,,

rootstock-0.5.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

rootstock-0.5.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+rootstock = rootstock.cli:main

rootstock-0.5.0.dist-info/licenses/LICENSE.md

@@ -0,0 +1,7 @@
+Copyright 2026 The University of Chicago
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.