rootstock 0.5.0__tar.gz → 0.5.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (35) hide show
  1. rootstock-0.5.1/PKG-INFO +178 -0
  2. rootstock-0.5.1/README.md +158 -0
  3. {rootstock-0.5.0 → rootstock-0.5.1}/pyproject.toml +1 -1
  4. {rootstock-0.5.0 → rootstock-0.5.1}/rootstock/clusters.py +1 -1
  5. rootstock-0.5.0/PKG-INFO +0 -210
  6. rootstock-0.5.0/README.md +0 -190
  7. {rootstock-0.5.0 → rootstock-0.5.1}/.github/workflows/publish.yml +0 -0
  8. {rootstock-0.5.0 → rootstock-0.5.1}/.gitignore +0 -0
  9. {rootstock-0.5.0 → rootstock-0.5.1}/CLAUDE.md +0 -0
  10. {rootstock-0.5.0 → rootstock-0.5.1}/LICENSE.md +0 -0
  11. {rootstock-0.5.0 → rootstock-0.5.1}/benchmarks/__init__.py +0 -0
  12. {rootstock-0.5.0 → rootstock-0.5.1}/benchmarks/direct.py +0 -0
  13. {rootstock-0.5.0 → rootstock-0.5.1}/benchmarks/sanity_check.py +0 -0
  14. {rootstock-0.5.0 → rootstock-0.5.1}/benchmarks/systems.py +0 -0
  15. {rootstock-0.5.0 → rootstock-0.5.1}/claude_docs/design-doc-v2.md +0 -0
  16. {rootstock-0.5.0 → rootstock-0.5.1}/claude_docs/design-doc-v3.md +0 -0
  17. {rootstock-0.5.0 → rootstock-0.5.1}/claude_docs/design-doc-v4.md +0 -0
  18. {rootstock-0.5.0 → rootstock-0.5.1}/claude_docs/design-doc-v5.md +0 -0
  19. {rootstock-0.5.0 → rootstock-0.5.1}/claude_docs/design-doc-v6.md +0 -0
  20. {rootstock-0.5.0 → rootstock-0.5.1}/claude_docs/design-doc-v7.md +0 -0
  21. {rootstock-0.5.0 → rootstock-0.5.1}/environments/chgnet_env.py +0 -0
  22. {rootstock-0.5.0 → rootstock-0.5.1}/environments/mace_env.py +0 -0
  23. {rootstock-0.5.0 → rootstock-0.5.1}/environments/tensornet_env.py +0 -0
  24. {rootstock-0.5.0 → rootstock-0.5.1}/environments/uma_env.py +0 -0
  25. {rootstock-0.5.0 → rootstock-0.5.1}/modal_app.py +0 -0
  26. {rootstock-0.5.0 → rootstock-0.5.1}/rootstock/__init__.py +0 -0
  27. {rootstock-0.5.0 → rootstock-0.5.1}/rootstock/calculator.py +0 -0
  28. {rootstock-0.5.0 → rootstock-0.5.1}/rootstock/cli.py +0 -0
  29. {rootstock-0.5.0 → rootstock-0.5.1}/rootstock/environment.py +0 -0
  30. {rootstock-0.5.0 → rootstock-0.5.1}/rootstock/pep723.py +0 -0
  31. {rootstock-0.5.0 → rootstock-0.5.1}/rootstock/protocol.py +0 -0
  32. {rootstock-0.5.0 → rootstock-0.5.1}/rootstock/server.py +0 -0
  33. {rootstock-0.5.0 → rootstock-0.5.1}/rootstock/worker.py +0 -0
  34. {rootstock-0.5.0 → rootstock-0.5.1}/test_rootstock.sbatch +0 -0
  35. {rootstock-0.5.0 → rootstock-0.5.1}/uv.lock +0 -0
@@ -0,0 +1,178 @@
1
+ Metadata-Version: 2.4
2
+ Name: rootstock
3
+ Version: 0.5.1
4
+ Summary: MLIP calculators with isolated Python environments
5
+ License-File: LICENSE.md
6
+ Requires-Python: >=3.10
7
+ Requires-Dist: ase>=3.22
8
+ Requires-Dist: numpy>=1.24
9
+ Requires-Dist: packaging>=21.0
10
+ Requires-Dist: tomli>=2.0; python_version < '3.11'
11
+ Provides-Extra: dev
12
+ Requires-Dist: pytest>=7.0; extra == 'dev'
13
+ Requires-Dist: ruff>=0.1; extra == 'dev'
14
+ Provides-Extra: mace
15
+ Requires-Dist: mace-torch>=0.3; extra == 'mace'
16
+ Requires-Dist: torch>=2.0; extra == 'mace'
17
+ Provides-Extra: modal
18
+ Requires-Dist: modal>=0.56; extra == 'modal'
19
+ Description-Content-Type: text/markdown
20
+
21
+ # Rootstock
22
+
23
+ A proof-of-concept package for running MLIP (Machine Learning Interatomic Potential) calculators in isolated pre-built Python environments, communicating via the i-PI protocol over Unix sockets. Currently deployed on Princeton's Della cluster.
24
+
25
+ ## Quick Start
26
+
27
+ ```python
28
+ from ase.build import bulk
29
+ from rootstock import RootstockCalculator
30
+
31
+ atoms = bulk("Cu", "fcc", a=3.6) * (5, 5, 5)
32
+
33
+ with RootstockCalculator(
34
+ cluster="della",
35
+ model="mace",
36
+ checkpoint="medium",
37
+ device="cuda",
38
+ ) as calc:
39
+ atoms.calc = calc
40
+ print(atoms.get_potential_energy())
41
+ print(atoms.get_forces())
42
+ ```
43
+
44
+ **Note:** Environments must be pre-built before use. See [Administrator Setup](#administrator-setup).
45
+
46
+ ## Installation
47
+
48
+ ```bash
49
+ pip install rootstock
50
+ # or
51
+ uv pip install rootstock
52
+ ```
53
+
54
+ ## API
55
+
56
+ The `model` parameter selects the environment family. The optional `checkpoint` parameter selects specific model weights (defaults to the environment's default if omitted).
57
+
58
+ ```python
59
+ # Explicit checkpoint
60
+ RootstockCalculator(cluster="della", model="mace", checkpoint="medium")
61
+
62
+ # Default checkpoint (each environment has a sensible default)
63
+ RootstockCalculator(cluster="della", model="uma")
64
+
65
+ # Custom root path instead of a known cluster
66
+ RootstockCalculator(root="/scratch/gpfs/ROSENGROUP/common/rootstock", model="mace")
67
+ ```
68
+
69
+ ## Available Models
70
+
71
+ | Model | Checkpoint | Description |
72
+ |-------|------------|-------------|
73
+ | `mace` | `small` | MACE-MP-0 small |
74
+ | `mace` | `medium` (default) | MACE-MP-0 medium |
75
+ | `chgnet` | _(pretrained)_ | CHGNet pretrained universal model |
76
+ | `uma` | `uma-s-1p1` (default) | Meta UMA small model (FAIRChem) |
77
+ | `tensornet` | `TensorNet-MatPES-PBE-v2025.1-PES` (default) | TensorNet MatPES PBE (MatGL) |
78
+
79
+ ## Architecture
80
+
81
+ ```
82
+ Main Process Worker Process (subprocess)
83
+ +-------------------------+ +-----------------------------+
84
+ | RootstockCalculator | | Pre-built venv Python |
85
+ | (ASE-compatible) | | (mace_env/bin/python) |
86
+ | | | |
87
+ | server.py (i-PI server) |<-------->| worker.py (i-PI client) |
88
+ | - sends positions | Unix | - receives positions |
89
+ | - receives forces | socket | - calculates forces |
90
+ +-------------------------+ +-----------------------------+
91
+ ```
92
+
93
+ The worker process uses a pre-built virtual environment, providing:
94
+ - **Fast startup**: No dependency installation at runtime
95
+ - **Filesystem compatibility**: Works on NFS, Lustre, GPFS
96
+ - **Reproducibility**: Same environment every time
97
+
98
+ ## Directory Structure
99
+
100
+ ```
101
+ {root}/
102
+ ├── .python/ # uv-managed Python interpreters (portable)
103
+ ├── environments/ # Environment SOURCE files (*.py with PEP 723)
104
+ │ ├── mace_env.py
105
+ │ ├── chgnet_env.py
106
+ │ ├── uma_env.py
107
+ │ └── tensornet_env.py
108
+ ├── envs/ # Pre-built virtual environments
109
+ │ ├── mace_env/
110
+ │ │ ├── bin/python
111
+ │ │ ├── lib/python3.11/site-packages/
112
+ │ │ └── env_source.py # Copy of environment source
113
+ │ └── ...
114
+ ├── home/ # Fake HOME for build & workers
115
+ │ ├── .cache/fairchem/ # FAIRChem weights
116
+ │ └── .matgl/ # MatGL weights
117
+ └── cache/ # XDG_CACHE_HOME for well-behaved libs
118
+ ├── mace/
119
+ └── huggingface/
120
+ ```
121
+
122
+ ## CLI Commands
123
+
124
+ ```bash
125
+ # Build a pre-built environment
126
+ rootstock build <env_name> --root <path> [--models m1,m2] [--force]
127
+
128
+ # Show status
129
+ rootstock status --root <path>
130
+
131
+ # Register an environment source file
132
+ rootstock register <env_file> --root <path>
133
+
134
+ # List environments
135
+ rootstock list --root <path>
136
+ ```
137
+
138
+ ## Administrator Setup
139
+
140
+ Environments must be pre-built before users can run calculations.
141
+
142
+ ### 1. Create Directory Structure
143
+
144
+ ```bash
145
+ mkdir -p /scratch/gpfs/ROSENGROUP/common/rootstock/{environments,envs,cache,home}
146
+ ```
147
+
148
+ ### 2. Register Environment Source Files
149
+
150
+ ```bash
151
+ rootstock register environments/mace_env.py --root /scratch/gpfs/ROSENGROUP/common/rootstock
152
+ rootstock register environments/chgnet_env.py --root /scratch/gpfs/ROSENGROUP/common/rootstock
153
+ rootstock register environments/uma_env.py --root /scratch/gpfs/ROSENGROUP/common/rootstock
154
+ rootstock register environments/tensornet_env.py --root /scratch/gpfs/ROSENGROUP/common/rootstock
155
+ ```
156
+
157
+ ### 3. Build Environments
158
+
159
+ ```bash
160
+ ROOT=/scratch/gpfs/ROSENGROUP/common/rootstock
161
+
162
+ rootstock build mace_env --root $ROOT --models small,medium
163
+ rootstock build chgnet_env --root $ROOT
164
+ rootstock build uma_env --root $ROOT --models uma-s-1p1
165
+ rootstock build tensornet_env --root $ROOT
166
+
167
+ # Verify
168
+ rootstock status --root $ROOT
169
+ ```
170
+
171
+ ## Local Development
172
+
173
+ ```bash
174
+ uv venv && source .venv/bin/activate
175
+ uv pip install -e ".[dev]"
176
+ ruff check rootstock/
177
+ ruff format rootstock/
178
+ ```
@@ -0,0 +1,158 @@
1
+ # Rootstock
2
+
3
+ A proof-of-concept package for running MLIP (Machine Learning Interatomic Potential) calculators in isolated pre-built Python environments, communicating via the i-PI protocol over Unix sockets. Currently deployed on Princeton's Della cluster.
4
+
5
+ ## Quick Start
6
+
7
+ ```python
8
+ from ase.build import bulk
9
+ from rootstock import RootstockCalculator
10
+
11
+ atoms = bulk("Cu", "fcc", a=3.6) * (5, 5, 5)
12
+
13
+ with RootstockCalculator(
14
+ cluster="della",
15
+ model="mace",
16
+ checkpoint="medium",
17
+ device="cuda",
18
+ ) as calc:
19
+ atoms.calc = calc
20
+ print(atoms.get_potential_energy())
21
+ print(atoms.get_forces())
22
+ ```
23
+
24
+ **Note:** Environments must be pre-built before use. See [Administrator Setup](#administrator-setup).
25
+
26
+ ## Installation
27
+
28
+ ```bash
29
+ pip install rootstock
30
+ # or
31
+ uv pip install rootstock
32
+ ```
33
+
34
+ ## API
35
+
36
+ The `model` parameter selects the environment family. The optional `checkpoint` parameter selects specific model weights (defaults to the environment's default if omitted).
37
+
38
+ ```python
39
+ # Explicit checkpoint
40
+ RootstockCalculator(cluster="della", model="mace", checkpoint="medium")
41
+
42
+ # Default checkpoint (each environment has a sensible default)
43
+ RootstockCalculator(cluster="della", model="uma")
44
+
45
+ # Custom root path instead of a known cluster
46
+ RootstockCalculator(root="/scratch/gpfs/ROSENGROUP/common/rootstock", model="mace")
47
+ ```
48
+
49
+ ## Available Models
50
+
51
+ | Model | Checkpoint | Description |
52
+ |-------|------------|-------------|
53
+ | `mace` | `small` | MACE-MP-0 small |
54
+ | `mace` | `medium` (default) | MACE-MP-0 medium |
55
+ | `chgnet` | _(pretrained)_ | CHGNet pretrained universal model |
56
+ | `uma` | `uma-s-1p1` (default) | Meta UMA small model (FAIRChem) |
57
+ | `tensornet` | `TensorNet-MatPES-PBE-v2025.1-PES` (default) | TensorNet MatPES PBE (MatGL) |
58
+
59
+ ## Architecture
60
+
61
+ ```
62
+ Main Process Worker Process (subprocess)
63
+ +-------------------------+ +-----------------------------+
64
+ | RootstockCalculator | | Pre-built venv Python |
65
+ | (ASE-compatible) | | (mace_env/bin/python) |
66
+ | | | |
67
+ | server.py (i-PI server) |<-------->| worker.py (i-PI client) |
68
+ | - sends positions | Unix | - receives positions |
69
+ | - receives forces | socket | - calculates forces |
70
+ +-------------------------+ +-----------------------------+
71
+ ```
72
+
73
+ The worker process uses a pre-built virtual environment, providing:
74
+ - **Fast startup**: No dependency installation at runtime
75
+ - **Filesystem compatibility**: Works on NFS, Lustre, GPFS
76
+ - **Reproducibility**: Same environment every time
77
+
78
+ ## Directory Structure
79
+
80
+ ```
81
+ {root}/
82
+ ├── .python/ # uv-managed Python interpreters (portable)
83
+ ├── environments/ # Environment SOURCE files (*.py with PEP 723)
84
+ │ ├── mace_env.py
85
+ │ ├── chgnet_env.py
86
+ │ ├── uma_env.py
87
+ │ └── tensornet_env.py
88
+ ├── envs/ # Pre-built virtual environments
89
+ │ ├── mace_env/
90
+ │ │ ├── bin/python
91
+ │ │ ├── lib/python3.11/site-packages/
92
+ │ │ └── env_source.py # Copy of environment source
93
+ │ └── ...
94
+ ├── home/ # Fake HOME for build & workers
95
+ │ ├── .cache/fairchem/ # FAIRChem weights
96
+ │ └── .matgl/ # MatGL weights
97
+ └── cache/ # XDG_CACHE_HOME for well-behaved libs
98
+ ├── mace/
99
+ └── huggingface/
100
+ ```
101
+
102
+ ## CLI Commands
103
+
104
+ ```bash
105
+ # Build a pre-built environment
106
+ rootstock build <env_name> --root <path> [--models m1,m2] [--force]
107
+
108
+ # Show status
109
+ rootstock status --root <path>
110
+
111
+ # Register an environment source file
112
+ rootstock register <env_file> --root <path>
113
+
114
+ # List environments
115
+ rootstock list --root <path>
116
+ ```
117
+
118
+ ## Administrator Setup
119
+
120
+ Environments must be pre-built before users can run calculations.
121
+
122
+ ### 1. Create Directory Structure
123
+
124
+ ```bash
125
+ mkdir -p /scratch/gpfs/ROSENGROUP/common/rootstock/{environments,envs,cache,home}
126
+ ```
127
+
128
+ ### 2. Register Environment Source Files
129
+
130
+ ```bash
131
+ rootstock register environments/mace_env.py --root /scratch/gpfs/ROSENGROUP/common/rootstock
132
+ rootstock register environments/chgnet_env.py --root /scratch/gpfs/ROSENGROUP/common/rootstock
133
+ rootstock register environments/uma_env.py --root /scratch/gpfs/ROSENGROUP/common/rootstock
134
+ rootstock register environments/tensornet_env.py --root /scratch/gpfs/ROSENGROUP/common/rootstock
135
+ ```
136
+
137
+ ### 3. Build Environments
138
+
139
+ ```bash
140
+ ROOT=/scratch/gpfs/ROSENGROUP/common/rootstock
141
+
142
+ rootstock build mace_env --root $ROOT --models small,medium
143
+ rootstock build chgnet_env --root $ROOT
144
+ rootstock build uma_env --root $ROOT --models uma-s-1p1
145
+ rootstock build tensornet_env --root $ROOT
146
+
147
+ # Verify
148
+ rootstock status --root $ROOT
149
+ ```
150
+
151
+ ## Local Development
152
+
153
+ ```bash
154
+ uv venv && source .venv/bin/activate
155
+ uv pip install -e ".[dev]"
156
+ ruff check rootstock/
157
+ ruff format rootstock/
158
+ ```
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "rootstock"
3
- version = "0.5.0"
3
+ version = "0.5.1"
4
4
  description = "MLIP calculators with isolated Python environments"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.10"
@@ -12,7 +12,7 @@ from pathlib import Path
12
12
  # Registry of known clusters and their rootstock root directories
13
13
  CLUSTER_REGISTRY: dict[str, str] = {
14
14
  "modal": "/vol/rootstock",
15
- "della": "/scratch/gpfs/SHARED/rootstock",
15
+ "della": "/scratch/gpfs/ROSENGROUP/common/rootstock",
16
16
  }
17
17
 
18
18
  # Known environment families (used for validation)
rootstock-0.5.0/PKG-INFO DELETED
@@ -1,210 +0,0 @@
1
- Metadata-Version: 2.4
2
- Name: rootstock
3
- Version: 0.5.0
4
- Summary: MLIP calculators with isolated Python environments
5
- License-File: LICENSE.md
6
- Requires-Python: >=3.10
7
- Requires-Dist: ase>=3.22
8
- Requires-Dist: numpy>=1.24
9
- Requires-Dist: packaging>=21.0
10
- Requires-Dist: tomli>=2.0; python_version < '3.11'
11
- Provides-Extra: dev
12
- Requires-Dist: pytest>=7.0; extra == 'dev'
13
- Requires-Dist: ruff>=0.1; extra == 'dev'
14
- Provides-Extra: mace
15
- Requires-Dist: mace-torch>=0.3; extra == 'mace'
16
- Requires-Dist: torch>=2.0; extra == 'mace'
17
- Provides-Extra: modal
18
- Requires-Dist: modal>=0.56; extra == 'modal'
19
- Description-Content-Type: text/markdown
20
-
21
- # Rootstock
22
-
23
- Run MLIP (Machine Learning Interatomic Potential) calculators in isolated pre-built Python environments, communicating via the i-PI protocol over Unix sockets.
24
-
25
- ## Quick Start
26
-
27
- ```python
28
- from ase.build import bulk
29
- from rootstock import RootstockCalculator
30
-
31
- atoms = bulk("Cu", "fcc", a=3.6) * (5, 5, 5)
32
-
33
- # Using a known cluster
34
- with RootstockCalculator(
35
- cluster="modal", # or "della"
36
- model="mace-medium", # or "chgnet", "mace-small", etc.
37
- device="cuda",
38
- ) as calc:
39
- atoms.calc = calc
40
- print(atoms.get_potential_energy())
41
- print(atoms.get_forces())
42
-
43
- # Or with an explicit root path
44
- with RootstockCalculator(
45
- root="/scratch/gpfs/SHARED/rootstock",
46
- model="mace-medium",
47
- device="cuda",
48
- ) as calc:
49
- atoms.calc = calc
50
- print(atoms.get_potential_energy())
51
- ```
52
-
53
- **Note:** Environments must be pre-built before use. See [Administrator Setup](#administrator-setup).
54
-
55
- ## Installation
56
-
57
- ```bash
58
- pip install rootstock
59
- # or
60
- uv pip install rootstock
61
- ```
62
-
63
- ## Model String Format
64
-
65
- The `model` parameter encodes both the environment and model-specific argument:
66
-
67
- | `model=` | Environment | Model Arg |
68
- |---------------------|----------------|---------------------|
69
- | `"mace-medium"` | mace_env | `"medium"` |
70
- | `"mace-small"` | mace_env | `"small"` |
71
- | `"mace-large"` | mace_env | `"large"` |
72
- | `"chgnet"` | chgnet_env | `""` (default) |
73
- | `"mace-/path/to/weights.pt"` | mace_env | `"/path/to/weights.pt"` |
74
-
75
- ## Known Clusters
76
-
77
- | Cluster | Root Path |
78
- |---------|-----------|
79
- | `modal` | `/vol/rootstock` |
80
- | `della` | `/scratch/gpfs/SHARED/rootstock` |
81
-
82
- For other clusters, use `root="/path/to/rootstock"` directly.
83
-
84
- ## Administrator Setup
85
-
86
- Environments must be pre-built before users can run calculations.
87
-
88
- ### 1. Create Directory Structure
89
-
90
- ```bash
91
- mkdir -p /scratch/gpfs/SHARED/rootstock/{environments,envs,cache}
92
- ```
93
-
94
- ### 2. Create Environment Source Files
95
-
96
- ```bash
97
- # mace_env.py
98
- cat > /scratch/gpfs/SHARED/rootstock/environments/mace_env.py << 'EOF'
99
- # /// script
100
- # requires-python = ">=3.10"
101
- # dependencies = ["mace-torch>=0.3.0", "ase>=3.22", "torch>=2.0"]
102
- # ///
103
- """MACE environment for Rootstock."""
104
-
105
- def setup(model: str, device: str = "cuda"):
106
- from mace.calculators import mace_mp
107
- return mace_mp(model=model, device=device, default_dtype="float32")
108
- EOF
109
- ```
110
-
111
- ### 3. Build Environments
112
-
113
- ```bash
114
- # Build MACE environment with model pre-download
115
- rootstock build mace_env --root /scratch/gpfs/SHARED/rootstock --models small,medium,large
116
-
117
- # Build CHGNet environment
118
- rootstock build chgnet_env --root /scratch/gpfs/SHARED/rootstock
119
-
120
- # Verify
121
- rootstock status --root /scratch/gpfs/SHARED/rootstock
122
- ```
123
-
124
- ## Architecture
125
-
126
- ```
127
- Main Process Worker Process (subprocess)
128
- +-------------------------+ +-----------------------------+
129
- | RootstockCalculator | | Pre-built venv Python |
130
- | (ASE-compatible) | | (mace_env/bin/python) |
131
- | | | |
132
- | server.py (i-PI server) |<-------->| worker.py (i-PI client) |
133
- | - sends positions | Unix | - receives positions |
134
- | - receives forces | socket | - calculates forces |
135
- +-------------------------+ +-----------------------------+
136
- ```
137
-
138
- The worker process uses a pre-built virtual environment, providing:
139
- - **Fast startup**: No dependency installation at runtime
140
- - **Filesystem compatibility**: Works on NFS, Lustre, GPFS, Modal volumes
141
- - **Reproducibility**: Same environment every time
142
-
143
- ## Directory Structure
144
-
145
- ```
146
- {root}/
147
- ├── environments/ # Environment SOURCE files (*.py with PEP 723)
148
- │ ├── mace_env.py
149
- │ └── chgnet_env.py
150
- ├── envs/ # Pre-built virtual environments
151
- │ ├── mace_env/
152
- │ │ ├── bin/python
153
- │ │ ├── lib/python3.11/site-packages/
154
- │ │ └── env_source.py # Copy of environment source
155
- │ └── chgnet_env/
156
- └── cache/ # XDG_CACHE_HOME for model weights
157
- ├── mace/ # MACE models
158
- └── huggingface/ # HuggingFace models
159
- ```
160
-
161
- ## CLI Commands
162
-
163
- ```bash
164
- # Build a pre-built environment
165
- rootstock build <env_name> --root <path> [--models m1,m2] [--force]
166
-
167
- # Show status
168
- rootstock status --root <path>
169
-
170
- # Register an environment source file
171
- rootstock register <env_file> --root <path>
172
-
173
- # List environments
174
- rootstock list --root <path>
175
- ```
176
-
177
- ## Running on Modal
178
-
179
- ```bash
180
- # Initialize volume and build environments (takes ~10-15 min)
181
- modal run modal_app.py::init_rootstock_volume
182
-
183
- # Test pre-built environments
184
- modal run modal_app.py::test_prebuilt
185
-
186
- # Show status
187
- modal run modal_app.py::inspect_status
188
-
189
- # Run benchmarks
190
- modal run modal_app.py::benchmark_v4
191
- ```
192
-
193
- ## Performance
194
-
195
- IPC overhead is <5% for systems with 1000+ atoms compared to direct in-process execution.
196
-
197
- | System Size | Atoms | Typical Overhead |
198
- |-------------|-------|------------------|
199
- | Small | 64 | ~10-15% |
200
- | Medium | 256 | ~5-8% |
201
- | Large | 1000 | <5% |
202
-
203
- ## Local Development
204
-
205
- ```bash
206
- uv venv && source .venv/bin/activate
207
- uv pip install -e ".[dev]"
208
- ruff check rootstock/
209
- ruff format rootstock/
210
- ```
rootstock-0.5.0/README.md DELETED
@@ -1,190 +0,0 @@
1
- # Rootstock
2
-
3
- Run MLIP (Machine Learning Interatomic Potential) calculators in isolated pre-built Python environments, communicating via the i-PI protocol over Unix sockets.
4
-
5
- ## Quick Start
6
-
7
- ```python
8
- from ase.build import bulk
9
- from rootstock import RootstockCalculator
10
-
11
- atoms = bulk("Cu", "fcc", a=3.6) * (5, 5, 5)
12
-
13
- # Using a known cluster
14
- with RootstockCalculator(
15
- cluster="modal", # or "della"
16
- model="mace-medium", # or "chgnet", "mace-small", etc.
17
- device="cuda",
18
- ) as calc:
19
- atoms.calc = calc
20
- print(atoms.get_potential_energy())
21
- print(atoms.get_forces())
22
-
23
- # Or with an explicit root path
24
- with RootstockCalculator(
25
- root="/scratch/gpfs/SHARED/rootstock",
26
- model="mace-medium",
27
- device="cuda",
28
- ) as calc:
29
- atoms.calc = calc
30
- print(atoms.get_potential_energy())
31
- ```
32
-
33
- **Note:** Environments must be pre-built before use. See [Administrator Setup](#administrator-setup).
34
-
35
- ## Installation
36
-
37
- ```bash
38
- pip install rootstock
39
- # or
40
- uv pip install rootstock
41
- ```
42
-
43
- ## Model String Format
44
-
45
- The `model` parameter encodes both the environment and model-specific argument:
46
-
47
- | `model=` | Environment | Model Arg |
48
- |---------------------|----------------|---------------------|
49
- | `"mace-medium"` | mace_env | `"medium"` |
50
- | `"mace-small"` | mace_env | `"small"` |
51
- | `"mace-large"` | mace_env | `"large"` |
52
- | `"chgnet"` | chgnet_env | `""` (default) |
53
- | `"mace-/path/to/weights.pt"` | mace_env | `"/path/to/weights.pt"` |
54
-
55
- ## Known Clusters
56
-
57
- | Cluster | Root Path |
58
- |---------|-----------|
59
- | `modal` | `/vol/rootstock` |
60
- | `della` | `/scratch/gpfs/SHARED/rootstock` |
61
-
62
- For other clusters, use `root="/path/to/rootstock"` directly.
63
-
64
- ## Administrator Setup
65
-
66
- Environments must be pre-built before users can run calculations.
67
-
68
- ### 1. Create Directory Structure
69
-
70
- ```bash
71
- mkdir -p /scratch/gpfs/SHARED/rootstock/{environments,envs,cache}
72
- ```
73
-
74
- ### 2. Create Environment Source Files
75
-
76
- ```bash
77
- # mace_env.py
78
- cat > /scratch/gpfs/SHARED/rootstock/environments/mace_env.py << 'EOF'
79
- # /// script
80
- # requires-python = ">=3.10"
81
- # dependencies = ["mace-torch>=0.3.0", "ase>=3.22", "torch>=2.0"]
82
- # ///
83
- """MACE environment for Rootstock."""
84
-
85
- def setup(model: str, device: str = "cuda"):
86
- from mace.calculators import mace_mp
87
- return mace_mp(model=model, device=device, default_dtype="float32")
88
- EOF
89
- ```
90
-
91
- ### 3. Build Environments
92
-
93
- ```bash
94
- # Build MACE environment with model pre-download
95
- rootstock build mace_env --root /scratch/gpfs/SHARED/rootstock --models small,medium,large
96
-
97
- # Build CHGNet environment
98
- rootstock build chgnet_env --root /scratch/gpfs/SHARED/rootstock
99
-
100
- # Verify
101
- rootstock status --root /scratch/gpfs/SHARED/rootstock
102
- ```
103
-
104
- ## Architecture
105
-
106
- ```
107
- Main Process Worker Process (subprocess)
108
- +-------------------------+ +-----------------------------+
109
- | RootstockCalculator | | Pre-built venv Python |
110
- | (ASE-compatible) | | (mace_env/bin/python) |
111
- | | | |
112
- | server.py (i-PI server) |<-------->| worker.py (i-PI client) |
113
- | - sends positions | Unix | - receives positions |
114
- | - receives forces | socket | - calculates forces |
115
- +-------------------------+ +-----------------------------+
116
- ```
117
-
118
- The worker process uses a pre-built virtual environment, providing:
119
- - **Fast startup**: No dependency installation at runtime
120
- - **Filesystem compatibility**: Works on NFS, Lustre, GPFS, Modal volumes
121
- - **Reproducibility**: Same environment every time
122
-
123
- ## Directory Structure
124
-
125
- ```
126
- {root}/
127
- ├── environments/ # Environment SOURCE files (*.py with PEP 723)
128
- │ ├── mace_env.py
129
- │ └── chgnet_env.py
130
- ├── envs/ # Pre-built virtual environments
131
- │ ├── mace_env/
132
- │ │ ├── bin/python
133
- │ │ ├── lib/python3.11/site-packages/
134
- │ │ └── env_source.py # Copy of environment source
135
- │ └── chgnet_env/
136
- └── cache/ # XDG_CACHE_HOME for model weights
137
- ├── mace/ # MACE models
138
- └── huggingface/ # HuggingFace models
139
- ```
140
-
141
- ## CLI Commands
142
-
143
- ```bash
144
- # Build a pre-built environment
145
- rootstock build <env_name> --root <path> [--models m1,m2] [--force]
146
-
147
- # Show status
148
- rootstock status --root <path>
149
-
150
- # Register an environment source file
151
- rootstock register <env_file> --root <path>
152
-
153
- # List environments
154
- rootstock list --root <path>
155
- ```
156
-
157
- ## Running on Modal
158
-
159
- ```bash
160
- # Initialize volume and build environments (takes ~10-15 min)
161
- modal run modal_app.py::init_rootstock_volume
162
-
163
- # Test pre-built environments
164
- modal run modal_app.py::test_prebuilt
165
-
166
- # Show status
167
- modal run modal_app.py::inspect_status
168
-
169
- # Run benchmarks
170
- modal run modal_app.py::benchmark_v4
171
- ```
172
-
173
- ## Performance
174
-
175
- IPC overhead is <5% for systems with 1000+ atoms compared to direct in-process execution.
176
-
177
- | System Size | Atoms | Typical Overhead |
178
- |-------------|-------|------------------|
179
- | Small | 64 | ~10-15% |
180
- | Medium | 256 | ~5-8% |
181
- | Large | 1000 | <5% |
182
-
183
- ## Local Development
184
-
185
- ```bash
186
- uv venv && source .venv/bin/activate
187
- uv pip install -e ".[dev]"
188
- ruff check rootstock/
189
- ruff format rootstock/
190
- ```
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes