rootstock 0.5.0__tar.gz

rootstock-0.5.0/.github/workflows/publish.yml

@@ -0,0 +1,30 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  publish:
+    name: Build and publish to PyPI
+    runs-on: ubuntu-latest
+    
+    permissions:
+      id-token: write  # Required for trusted publishing
+    
+    steps:
+      - uses: actions/checkout@v4
+      
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+      
+      - name: Build package
+        run: uv build
+      
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

rootstock-0.5.0/.gitignore

@@ -0,0 +1,46 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+.venv/
+venv/
+ENV/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# Modal
+.modal/
+
+# Benchmark outputs
+results/
+*.json
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Sockets (shouldn't be committed but just in case)
+/tmp/ipi_*

rootstock-0.5.0/CLAUDE.md

@@ -0,0 +1,160 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project Overview
+
+Rootstock is a proof-of-concept for running MLIP (Machine Learning Interatomic Potential) calculators in isolated pre-built Python environments, communicating via the i-PI protocol over Unix sockets.
+
+**Current version: v0.5** - Pre-built environments with UMA and TensorNet support.
+
+## Commands
+
+### 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
+```
+
+### Local Development
+```bash
+uv venv && source .venv/bin/activate
+uv pip install -e ".[dev]"
+```
+
+### CLI Commands
+```bash
+# Build a pre-built environment
+rootstock build <env_name> --root <path> [--models m1,m2] [--force]
+
+# Show status
+rootstock status --root <path>
+
+# List environments
+rootstock list --root <path>
+```
+
+### Linting
+```bash
+ruff check rootstock/
+ruff format 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         │
+└─────────────────────────┘          └─────────────────────────────┘
+```
+
+**Key design**: v0.4 uses pre-built virtual environments instead of dynamic `uv run`. This provides:
+- Fast startup (no pip install at runtime)
+- Works on any filesystem (no lock files or hardlinks needed)
+- Reproducible environments
+
+### Core Files
+
+- `rootstock/cli.py` - CLI commands (`build`, `status`, `list`, `register`)
+- `rootstock/calculator.py` - ASE Calculator interface (main entry point)
+- `rootstock/server.py` - Spawns worker subprocess, manages socket lifecycle
+- `rootstock/worker.py` - i-PI client state machine
+- `rootstock/environment.py` - Pre-built environment management, wrapper generation
+- `rootstock/clusters.py` - Cluster registry and known environments
+
+### Directory Structure
+
+```
+{root}/
+├── .python/                # uv-managed Python interpreters (portable)
+│   └── cpython-3.10.19-linux-x86_64-gnu/
+├── environments/           # Environment SOURCE files (*.py with PEP 723)
+│   ├── mace_env.py
+│   ├── chgnet_env.py
+│   ├── uma_env.py
+│   └── tensornet_env.py
+├── envs/                   # Pre-built virtual environments
+│   ├── mace_env/
+│   │   ├── bin/python      # Symlinks to .python/
+│   │   ├── lib/python3.11/site-packages/
+│   │   └── env_source.py   # Copy of source for imports
+│   └── chgnet_env/
+└── cache/                  # XDG_CACHE_HOME for model weights
+    ├── mace/
+    └── huggingface/
+```
+
+The entire `{root}/` directory is self-contained and portable. Python interpreters
+are stored in `.python/` so venv symlinks resolve correctly on any machine where
+the root directory is mounted (Modal Volume, HPC shared filesystem, etc.).
+
+### Known Clusters
+
+| Cluster | Root Path |
+|---------|-----------|
+| `modal` | `/vol/rootstock` |
+| `della` | `/scratch/gpfs/SHARED/rootstock` |
+
+## API
+
+```python
+# v0.5 API: explicit model and checkpoint parameters
+with RootstockCalculator(
+    cluster="della",
+    model="mace",              # Environment family -> mace_env
+    checkpoint="medium",       # Specific checkpoint (optional, uses default if omitted)
+    device="cuda",
+) as calc:
+    atoms.calc = calc
+    energy = atoms.get_potential_energy()
+
+# Checkpoint defaults to environment's default if omitted
+with RootstockCalculator(cluster="della", model="uma") as calc:
+    ...  # Uses uma-s-1p1 by default
+```
+
+### Available Models
+
+| Model | Environment | Default Checkpoint | Other Checkpoints |
+|-------|-------------|-------------------|-------------------|
+| `mace` | mace_env | `medium` | `small`, `large` |
+| `chgnet` | chgnet_env | (pretrained) | — |
+| `uma` | uma_env | `uma-s-1p1` | — |
+| `tensornet` | tensornet_env | `TensorNet-MatPES-PBE-v2025.1-PES` | Other MatGL models |
+
+## Build Process
+
+```bash
+# 1. Create environment source file
+cat > environments/mace_env.py << 'EOF'
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["mace-torch>=0.3.0", "ase>=3.22", "torch>=2.0"]
+# ///
+
+def setup(model: str, device: str = "cuda"):
+    from mace.calculators import mace_mp
+    return mace_mp(model=model, device=device, default_dtype="float32")
+EOF
+
+# 2. Build pre-built environment
+rootstock build mace_env --root /path/to/rootstock --models small,medium
+
+# 3. Verify
+rootstock status --root /path/to/rootstock
+```

rootstock-0.5.0/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.

rootstock-0.5.0/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,190 @@
+# 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/benchmarks/__init__.py

@@ -0,0 +1,3 @@
+"""
+Benchmarks for measuring IPC overhead in Rootstock.
+"""