From 475b8b8e89333617f6891c0e70b467f4b324be1b Mon Sep 17 00:00:00 2001 From: navidgh67 Date: Wed, 12 Aug 2026 11:09:24 -0400 Subject: [PATCH 01/17] Add feature-flagged Sim2L RCAC computations --- RCAC_SIM2L_INTEGRATION.md | 123 ++++ api/Dockerfile.worker | 25 + api/computations.env.example | 27 + api/computations/__init__.py | 5 + api/computations/artifacts.py | 118 ++++ api/computations/materials.py | 35 ++ api/computations/pseudopotentials.sha256 | 171 ++++++ api/computations/qe.py | 647 ++++++++++++++++++++ api/computations/rcac.py | 116 ++++ api/computations/results.py | 43 ++ api/computations/schemas.py | 209 +++++++ api/computations/store.py | 617 +++++++++++++++++++ api/computations/worker.py | 399 ++++++++++++ api/main.py | 7 +- api/requirements-computations.txt | 6 + api/routers/computations.py | 214 +++++++ api/scripts/add_rcac_computations.sql | 78 +++ api/tests/test_computation_schemas.py | 59 ++ api/tests/test_computation_worker.py | 166 +++++ api/tests/test_computations_postgres_e2e.py | 377 ++++++++++++ api/tests/test_computations_router.py | 28 + api/tests/test_qe_parser.py | 43 ++ geddes/k8s/09-rcac-worker.example.yaml | 55 ++ web/app/projects/[id]/page.tsx | 3 + web/components/project-computations.tsx | 272 ++++++++ web/lib/api-client.ts | 125 ++++ 26 files changed, 3967 insertions(+), 1 deletion(-) create mode 100644 RCAC_SIM2L_INTEGRATION.md create mode 100644 api/Dockerfile.worker create mode 100644 api/computations.env.example create mode 100644 api/computations/__init__.py create mode 100644 api/computations/artifacts.py create mode 100644 api/computations/materials.py create mode 100644 api/computations/pseudopotentials.sha256 create mode 100644 api/computations/qe.py create mode 100644 api/computations/rcac.py create mode 100644 api/computations/results.py create mode 100644 api/computations/schemas.py create mode 100644 api/computations/store.py create mode 100644 api/computations/worker.py create mode 100644 api/requirements-computations.txt create mode 100644 api/routers/computations.py create mode 100644 api/scripts/add_rcac_computations.sql create mode 100644 api/tests/test_computation_schemas.py create mode 100644 api/tests/test_computation_worker.py create mode 100644 api/tests/test_computations_postgres_e2e.py create mode 100644 api/tests/test_computations_router.py create mode 100644 api/tests/test_qe_parser.py create mode 100644 geddes/k8s/09-rcac-worker.example.yaml create mode 100644 web/components/project-computations.tsx diff --git a/RCAC_SIM2L_INTEGRATION.md b/RCAC_SIM2L_INTEGRATION.md new file mode 100644 index 0000000..134e5af --- /dev/null +++ b/RCAC_SIM2L_INTEGRATION.md @@ -0,0 +1,123 @@ +# CellRelaxDFT / ElasticDFT RCAC integration + +This feature adds CellRelaxDFT r99 and ElasticDFT r27 execution to the Digital +Twin without changing the existing execution queue, RCAC demo, or physical +experiment workflows. It is disabled unless `RCAC_COMPUTE_ENABLED=true`. + +## Safety boundaries + +- Purdue passwords and MFA responses are never accepted or stored. +- RCAC profiles contain only username, allocation, partition, and an opaque + reference to an operator-provisioned SSH key. +- Private keys are resolved from a read-only secret mount and must be mode 0600 + or stricter. SSH runs in batch mode with strict host-key checking. +- The downloaded r99/r27 bundles remain unchanged. Production adapters are + versioned under `api/computations/`. +- Large QE artifacts live in blob/filesystem storage; PostgreSQL stores their + URI, SHA-256 checksum, media type, and size. +- Nothing in this branch deploys itself. The example worker manifest has zero + replicas. + +## Components + +| Component | Responsibility | +| --- | --- | +| `api/computations/schemas.py` | Sim2L-compatible validation and SQUID IDs | +| `api/computations/qe.py` | QE input preparation, parsing, elastic fitting | +| `sim2l.executor.slurm` | Safe SSH/local Slurm stage/submit/status/cancel/collect | +| `api/computations/store.py` | Durable PostgreSQL run queue and project access | +| `api/computations/worker.py` | Restartable asynchronous execution lifecycle | +| `api/routers/computations.py` | Authenticated project API | +| `web/components/project-computations.tsx` | Opt-in project launch/status/results UI | + +The worker state machine is: + +```text +requested -> validating -> staging -> submitted -> running -> collecting + -> data_ingested + \-> failed / cancelled +``` + +ElasticDFT records `upstream_run_id` and inherits the converged structure and +DFT parameters from a completed CellRelaxDFT run in the same project. It never +queries nanoHUB ResultsDB. Independent strain calculations use a bounded Slurm +job array. + +## Database + +Apply `api/scripts/add_rcac_computations.sql` before enabling the feature. It +only creates: + +- `rcac_user_profiles` +- `computation_runs` +- `computation_run_events` + +It does not alter existing tables. FastAPI also runs this idempotent migration +at startup when the feature is enabled. + +## Local development + +Create an isolated environment and install the matching Sim2L checkout: + +```bash +python3 -m venv .venv-computations +.venv-computations/bin/pip install -r api/requirements-computations.txt pytest +.venv-computations/bin/pip install -e ../../Sim2L/sim2l +``` + +Copy values from `api/computations.env.example` into an ignored environment +file. For unit tests, no RCAC credentials or Materials Project key are needed. + +```bash +PYTHONPATH=api .venv-computations/bin/python -m pytest -q \ + api/tests/test_computation_schemas.py \ + api/tests/test_qe_parser.py \ + api/tests/test_computation_worker.py + +SIM2L_TEST_BOOTSTRAP_SERVICES=0 \ + .venv-computations/bin/python -m pytest -q \ + ../../Sim2L/sim2l/sim2l/tests/test_slurm_executor.py -o addopts='' +``` + +The worker image is deliberately separate from the existing API image: + +```bash +docker build --build-context sim2l=../../Sim2L/sim2l \ + -f api/Dockerfile.worker -t dt-computation-worker:dev . +``` + +The opt-in end-to-end test uses a disposable PostgreSQL database and a fake +Slurm transport. It exercises both CellRelaxDFT and ElasticDFT, including +chaining, artifacts, and caching, but never opens an RCAC connection: + +```bash +docker run --rm -d --name dt-computations-e2e \ + -e POSTGRES_PASSWORD=postgres -e POSTGRES_DB=dt_computations_test \ + -p 127.0.0.1:55439:5432 postgres:16-alpine + +cd api +DT_COMPUTATIONS_E2E=1 RCAC_COMPUTE_ENABLED=true \ + PG_HOST=127.0.0.1 PG_PORT=55439 PG_DB=dt_computations_test \ + PG_SUPER_USER=postgres PG_SUPER_PASS=postgres PYTHONPATH=. \ + ../.venv-computations/bin/pytest -q \ + tests/test_computations_postgres_e2e.py -o addopts='' + +docker stop dt-computations-e2e +``` + +## RCAC readiness gate + +Before changing the flag in any shared environment, obtain approval for the +per-user delegated authentication model and confirm: + +1. Cluster hostname and host key. +2. Allocation and partition policy. +3. Quantum ESPRESSO module/version and `pw.x` invocation. +4. Read-only shared PAW/USPP PBEsol pseudopotential location and checksums. +5. Durable artifact container and retention policy. +6. Network access from the worker to the RCAC login host and Sim2L Results + service, if used. + +The initial scientific smoke case is silicon (`mp-1265` from the source +workflow, or another approved reference record). Numerical tolerances must be +agreed per output before declaring parity with nanoHUB. diff --git a/api/Dockerfile.worker b/api/Dockerfile.worker new file mode 100644 index 0000000..77963a9 --- /dev/null +++ b/api/Dockerfile.worker @@ -0,0 +1,25 @@ +# syntax=docker/dockerfile:1.7 +# Build from the DT repository root with the matching Sim2L checkout supplied +# as a named context: +# +# docker build --build-context sim2l=../../Sim2L/sim2l \ +# -f api/Dockerfile.worker -t dt-computation-worker:dev . + +FROM python:3.12-slim +WORKDIR /app + +RUN apt-get update \ + && apt-get install -y --no-install-recommends ca-certificates openssh-client \ + && rm -rf /var/lib/apt/lists/* + +COPY api/requirements.txt api/requirements-computations.txt ./ +RUN pip install --no-cache-dir -r requirements-computations.txt + +COPY --from=sim2l / /tmp/sim2l +RUN pip install --no-cache-dir /tmp/sim2l && rm -rf /tmp/sim2l + +COPY api/ ./ +COPY domains/ /app/domains/ + +# The worker has no listening socket and must run separately from FastAPI. +CMD ["python", "-m", "computations.worker"] diff --git a/api/computations.env.example b/api/computations.env.example new file mode 100644 index 0000000..e752979 --- /dev/null +++ b/api/computations.env.example @@ -0,0 +1,27 @@ +# Off by default. Both API and worker must explicitly enable the feature. +RCAC_COMPUTE_ENABLED=false + +# One or more approved clusters. Confirm module strings with RCAC before use. +RCAC_CLUSTERS_JSON={"negishi":{"host":"negishi.rcac.purdue.edu","qe_module":"quantum-espresso/7.1","remote_root":"/scratch/{username}/dt-sim2l"}} + +# SSH mode never accepts a password. Keys are mounted read-only and profiles +# refer to a filename under this directory by an opaque credential_ref. +RCAC_TRANSPORT=ssh +RCAC_SSH_KEY_DIR=/var/run/secrets/rcac-ssh +RCAC_KNOWN_HOSTS_FILE=/etc/ssh/ssh_known_hosts + +# Source pseudo_PAW_PBEsol and pseudo_USPP_PBEsol directories. +RCAC_PSEUDO_ROOT=/depot/allocation/dt-sim2l/pseudo +RCAC_LOCAL_WORK_ROOT=/var/lib/dt-computation-worker/runs + +# Development fallback. Production should configure Azure blob storage. +RCAC_ARTIFACT_ROOT=/var/lib/dt-computation-worker/artifacts +RCAC_ARTIFACT_AZURE_CONNECTION_STRING= +RCAC_ARTIFACT_AZURE_CONTAINER= + +# Used only when structure_dict is omitted. +MATERIALS_PROJECT_API_KEY= + +# Optional Sim2L Results service. DT Postgres remains authoritative if omitted. +SIM2L_RESULTS_URL= +SIM2L_RESULTS_SESSION_ID= diff --git a/api/computations/__init__.py b/api/computations/__init__.py new file mode 100644 index 0000000..1c6812d --- /dev/null +++ b/api/computations/__init__.py @@ -0,0 +1,5 @@ +"""Feature-flagged Sim2L/RCAC computation integration.""" + +from .schemas import TOOL_DEFINITIONS, compute_squid_id, validate_tool_inputs + +__all__ = ["TOOL_DEFINITIONS", "compute_squid_id", "validate_tool_inputs"] diff --git a/api/computations/artifacts.py b/api/computations/artifacts.py new file mode 100644 index 0000000..4f89df6 --- /dev/null +++ b/api/computations/artifacts.py @@ -0,0 +1,118 @@ +"""Durable artifact stores for QE inputs, outputs, logs, and manifests.""" + +from __future__ import annotations + +import hashlib +import mimetypes +import os +import shutil +from collections.abc import Iterable +from pathlib import Path +from typing import Any +from urllib.parse import unquote, urlparse + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +class LocalArtifactStore: + """Filesystem store intended for development and fake-RCAC E2E tests.""" + + def __init__(self, root: Path): + self.root = Path(root).resolve() + self.root.mkdir(parents=True, exist_ok=True) + + def put_many( + self, run_id: str, base_dir: Path, paths: Iterable[Path] + ) -> list[dict[str, Any]]: + destination_root = self.root / run_id + destination_root.mkdir(parents=True, exist_ok=True) + base_dir = Path(base_dir).resolve() + manifest = [] + for path in paths: + path = Path(path).resolve() + relative = path.relative_to(base_dir) + destination = destination_root / relative + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(path, destination) + manifest.append(_artifact_record(path, relative, destination.as_uri())) + return manifest + + def resolve(self, uri: str) -> Path: + parsed = urlparse(uri) + if parsed.scheme != "file": + raise ValueError("artifact is not stored on the local filesystem") + path = Path(unquote(parsed.path)).resolve() + if self.root != path and self.root not in path.parents: + raise PermissionError("artifact path is outside the configured store") + if not path.is_file(): + raise FileNotFoundError(path.name) + return path + + +class AzureBlobArtifactStore: + """Azure-backed store using the platform's existing blob dependency.""" + + def __init__(self, connection_string: str, container: str, prefix: str = "qe-runs"): + from azure.storage.blob import BlobServiceClient + + self.container_name = container + self.prefix = prefix.strip("/") + client = BlobServiceClient.from_connection_string(connection_string) + self.container = client.get_container_client(container) + try: + self.container.create_container() + except Exception as exc: + # Existing containers raise ResourceExistsError. Avoid importing + # Azure exception classes during normal API startup. + if "ContainerAlreadyExists" not in str(exc) and "exists" not in str(exc): + raise + + def put_many( + self, run_id: str, base_dir: Path, paths: Iterable[Path] + ) -> list[dict[str, Any]]: + base_dir = Path(base_dir).resolve() + manifest = [] + for path in paths: + path = Path(path).resolve() + relative = path.relative_to(base_dir) + blob_name = f"{self.prefix}/{run_id}/{relative.as_posix()}" + with path.open("rb") as handle: + self.container.upload_blob(blob_name, handle, overwrite=True) + uri = f"azure://{self.container_name}/{blob_name}" + manifest.append(_artifact_record(path, relative, uri)) + return manifest + + def downloader(self, uri: str): + prefix = f"azure://{self.container_name}/" + if not uri.startswith(prefix): + raise ValueError("artifact is not in the configured Azure container") + blob_name = uri[len(prefix) :] + if not blob_name or ".." in Path(blob_name).parts: + raise ValueError("invalid Azure artifact path") + return self.container.download_blob(blob_name) + + +def _artifact_record(path: Path, relative: Path, uri: str) -> dict[str, Any]: + return { + "name": relative.as_posix(), + "uri": uri, + "sha256": _sha256(path), + "size_bytes": path.stat().st_size, + "content_type": mimetypes.guess_type(path.name)[0] + or "application/octet-stream", + } + + +def artifact_store_from_env(): + connection = os.getenv("RCAC_ARTIFACT_AZURE_CONNECTION_STRING", "").strip() + container = os.getenv("RCAC_ARTIFACT_AZURE_CONTAINER", "").strip() + if connection and container: + return AzureBlobArtifactStore(connection, container) + root = Path(os.getenv("RCAC_ARTIFACT_ROOT", "/tmp/dt-rcac-artifacts")) + return LocalArtifactStore(root) diff --git a/api/computations/materials.py b/api/computations/materials.py new file mode 100644 index 0000000..0ed29b2 --- /dev/null +++ b/api/computations/materials.py @@ -0,0 +1,35 @@ +"""Materials Project resolution kept outside request handlers.""" + +from __future__ import annotations + +import os +from typing import Any + + +def resolve_material_inputs(inputs: dict[str, Any]) -> dict[str, Any]: + """Fill structure and metallicity when the caller supplied only an MP id.""" + + if inputs.get("structure_dict"): + return inputs + api_key = os.getenv("MATERIALS_PROJECT_API_KEY", "").strip() + if not api_key: + raise RuntimeError( + "MATERIALS_PROJECT_API_KEY is required when structure_dict is omitted" + ) + try: + from mp_api.client import MPRester + except ImportError as exc: + raise RuntimeError("mp-api is required by the computation worker") from exc + + with MPRester(api_key) as rester: + documents = rester.materials.summary.search( + material_ids=[inputs["mp_id"]], + fields=["material_id", "structure", "is_metal"], + ) + if not documents: + raise ValueError(f"Materials Project record not found: {inputs['mp_id']}") + document = documents[0] + resolved = dict(inputs) + resolved["structure_dict"] = document.structure.as_dict() + resolved["is_metal"] = bool(document.is_metal) + return resolved diff --git a/api/computations/pseudopotentials.sha256 b/api/computations/pseudopotentials.sha256 new file mode 100644 index 0000000..c32f444 --- /dev/null +++ b/api/computations/pseudopotentials.sha256 @@ -0,0 +1,171 @@ +b72accf6f3c9cc0fd028cad42dc2fcbb01c71ad6f6bd522e4c234600ea8db9df pseudo_PAW_PBEsol/Ag.upf +bbf459847b29cbc32468e1a2cddeb4b1c958ca215023cc66cd40a92d2e0ed11c pseudo_PAW_PBEsol/Al.upf +6a18441a4314ab62f305745db9c09b02a63ef16023a7cfb8f7d0130037a961f3 pseudo_PAW_PBEsol/Ar.upf +fddbebe9c62aaa45c4b77abd5b90844d86b23cdd58c54c67bd2ff1b3ed62be00 pseudo_PAW_PBEsol/As.upf +07deffb7dc25da1691b821ae94a3d30fd3dfa8d5e482fc8ae422ddfdb5a373a7 pseudo_PAW_PBEsol/At.upf +c960e1484b4598715bfa33994161bdd6b2d39c83e6b6af59a2429389f1fd908a pseudo_PAW_PBEsol/Au.upf +0a9b322243b9ff02014bb87f574f5f8ab7037d9f1ed5257832f129c3c7333538 pseudo_PAW_PBEsol/B.upf +89c29de33e932c922217cdbb525a453b9ff47fc8c34161676c51e9f9b068ba8f pseudo_PAW_PBEsol/Ba.upf +db78275c1d71cc95520e02e842036ad4b95a58271ed1e997d333921984cd312f pseudo_PAW_PBEsol/Be.upf +827c550c44757b93e79cd870bf040b486a0e95e8b7864fc05b013fbe26fabec6 pseudo_PAW_PBEsol/Bi.upf +961d12e487e0460945e5d7670ada9dfe822620eb82e94fcc9666f463c0960563 pseudo_PAW_PBEsol/Br.upf +3fca0f06ef8dac54134ee149f0f17ce1990b24e9fcd0c73251b15efd19d1227a pseudo_PAW_PBEsol/C.upf +e080a39f0a6bad615296971c79f136665d63cdb4e7f854f057b854b4c2b5b12a pseudo_PAW_PBEsol/Ca.upf +8acf61688626d2db885c17e2dc07693276adabfea308b5e3e9389fac01b715d0 pseudo_PAW_PBEsol/Cd.upf +806759fbfc1af6ce93a0913a78c62eaa7530695f6ffadc41d1788eec21c89acc pseudo_PAW_PBEsol/Ce.upf +2665dbcc2c0b10b0a299617294d7248cc477ad7664ef0d3dba89daf42926fafc pseudo_PAW_PBEsol/Cl.upf +6cfc3c8e2a8140367d5442c3e03c2173f7cb6797d17f10ec5c0710c706b0258b pseudo_PAW_PBEsol/Co.upf +289d484b561678e08415b34c5c8cd627d8f634036e51e2b3f27c0eb4305227fd pseudo_PAW_PBEsol/Cr.upf +5c787d8716cd6e1a40725e014c01e9e0afec197633fb63b5ab3fddf914ddbca6 pseudo_PAW_PBEsol/Cs.upf +0e441d790ce4dd4d184c07968ae4b489814fb2d507c42bf74ae496989e333149 pseudo_PAW_PBEsol/Cu.upf +4053428d648adf04e3c6bd49033d5936c4928cf8995099313b4b838061780463 pseudo_PAW_PBEsol/Dy.upf +7f18c0d4bf647b98515858f49d30868df36c4aaa7acc673d081ef2e11364e3ae pseudo_PAW_PBEsol/Er.upf +214550cc73c804bd7056b7f39fb4a24b5809f680e3ea38a069e7e6a36569c65f pseudo_PAW_PBEsol/Eu.upf +f4759a69517310a2c93e3327cc4a511f347387e4dd3a27ebdde884c8bc3275c1 pseudo_PAW_PBEsol/F.upf +73feaa7a689b0298b4498cc13c3e9f53f4db1fde9afb569b933152ecdd5a96b7 pseudo_PAW_PBEsol/Fe.upf +9ddec49b4c381ffa1581189ff170d42febbc76c3d75efcdbc1de39dab68f87ec pseudo_PAW_PBEsol/Ga.upf +6f481a958b089fac08909ff07fc5554fe7c72a79c07ba3641035d4e6a1d48ef9 pseudo_PAW_PBEsol/Gd.upf +dedef770ac40368d94916ff9a52b3a8f184d4afae0ff470d9783dbf008f0c4ca pseudo_PAW_PBEsol/Ge.upf +8feb30fdbfad1b27e2f63f3643d876315e13827bb067214403a06ca19b06c720 pseudo_PAW_PBEsol/H.upf +cfcc03856855b7e82ab4e3d02902bdcdf6694823de273b8024b3bea6f2cc9ed2 pseudo_PAW_PBEsol/He.upf +f9506f1ea87d0b3a537977e91a19ae2268493a70c06df7d5af5c4bf1764a3481 pseudo_PAW_PBEsol/Hf.upf +19abecf9cba6ddbd95400d6d26bf059c3c9ad83c4bc8b868bc1219a84e10af66 pseudo_PAW_PBEsol/Hg.upf +62885e5b26236f2134471346d3df19f58e5682df0a24cf537aa2773397190060 pseudo_PAW_PBEsol/Ho.upf +9b235bb8a2dfcb312bfcb0e1d211beba7e44fc6ddb109a389b71e78e4576da69 pseudo_PAW_PBEsol/I.upf +c84ce1b32d8c839f81db252424bc4383ba023fcb30144c59ef801e0759f055ce pseudo_PAW_PBEsol/In.upf +18edc80f846dac7cd1d7c1541eab2069cdcdeeb8bf6808998d2bdb4ea7d64c2b pseudo_PAW_PBEsol/Ir.upf +cdf778dc665c388ef86927c05ec3f5fcd6fb55b7f2b2569a7ddbb84d68257785 pseudo_PAW_PBEsol/K.upf +b61cdcf86c1cb8bcf670746cc2b0c8402a428072f2d8780ee3d2866c98a25f46 pseudo_PAW_PBEsol/Kr.upf +0eac0fd619038c7b0fc2d1995f27702f753de246967a8b1bc9b89b601a0370fa pseudo_PAW_PBEsol/La.upf +6bb4c6c29493e25fb31ecc18d15e231a85cd8ff113965da7f87b351e3ed90e0a pseudo_PAW_PBEsol/Li.upf +525cac65df4bf0328c6476e1414ba6b80ba0958687da3cfc5b2d17e387f39326 pseudo_PAW_PBEsol/Lu.upf +0bf2bef727ce060ccbcfb36078743b5f537a9fbc4cb5df88c0ab91e343438736 pseudo_PAW_PBEsol/Mg.upf +d4cc7956efb17f03f8e89d076969952423c0fc7dd7e0f0acef0bc49c5b2aaa56 pseudo_PAW_PBEsol/Mn.upf +ac708da3dfdb552eb1ccaf0a8c1d200607d9e747516c3c6b6f9d379bfe0f1df6 pseudo_PAW_PBEsol/Mo.upf +098f4c428ee562d5eb6f5fdbb036e849c135d9c4c7ed0a8f70a8e2263df967ff pseudo_PAW_PBEsol/N.upf +f5541371f3823de7e8893534d9e05313c8a5fb17caf5c042006c9f60372f5267 pseudo_PAW_PBEsol/Na.upf +f4f3bb463efac6babaae59afa65d63ea7e11090ff754b871bb5e26c752a06803 pseudo_PAW_PBEsol/Nb.upf +7dd44b310f3eb1dcaf330574ef962921a8c14da02d49f645418db2a3924c4e20 pseudo_PAW_PBEsol/Nd.upf +0a51e12b6a3c2dd5ad7c7bdb68d27486ef93ae9c1d01fa53e8d127d244aea8bf pseudo_PAW_PBEsol/Ne.upf +28002e40e36bf4e2258f7c8c6717d391bceb23170451bd15f0a6dd905e2a2115 pseudo_PAW_PBEsol/Ni.upf +3cc26539673d44cf887ed311e96383164a1af08a44dc2d1dcad0ed77b269ca09 pseudo_PAW_PBEsol/O.upf +c45dcdc3b481bc7a393172d7dc0bfd0a30c9363ad5786dfea9d1101aac61c0eb pseudo_PAW_PBEsol/Os.upf +df6d996620772fbf473811d1616b1bc73ad5d40a856f6429d7dd9371840f7043 pseudo_PAW_PBEsol/P.upf +277d05a1ba44c0ac0a455f583f18050e0a1aeff814814592de0fedd3590a6f3f pseudo_PAW_PBEsol/Pb.upf +e34db5212290c7976b92106b6e0da102a4f9e7d4b675932b4c69c53e9f951a96 pseudo_PAW_PBEsol/Pd.upf +bfe73d26e6cce998bfe62c866b4926307fb22c88d5e8c759014d53f35b9b960d pseudo_PAW_PBEsol/Pm.upf +03ece3d0953391a41b6b2f78959bcab3b1762078d3e6db9373eee09772380ee0 pseudo_PAW_PBEsol/Po.upf +fabb6c4b175b26857b9b584e580083f0a3d13821df96801e485c2a35a6307959 pseudo_PAW_PBEsol/Pr.upf +3086672fc9640fd2bf58f53443a3e344fd0d16247e4d27f9256f2eab504ef552 pseudo_PAW_PBEsol/Pt.upf +ba043c0d342eae708ba0b7c14993d8b315f828bbccf136c28192dbc11784f688 pseudo_PAW_PBEsol/Rb.upf +9c36aaf2b179e7cf5684eb2461916b6e105553fe8fa4e76e5db137fb58950fac pseudo_PAW_PBEsol/Re.upf +5929ce9bcae3bcd9e466fee58197a7d96adcdbd79b7eb4d823cbf82a357e7817 pseudo_PAW_PBEsol/Rh.upf +bbd7611b5352504e4141eef52f8dc3c31e4b507630f0ef20aef70047fff5a702 pseudo_PAW_PBEsol/Rn.upf +5aae4d11c694a5d011742254203af06cddd881e67d9f83e54f6e413bb1f51eee pseudo_PAW_PBEsol/Ru.upf +eb356a2831fbb254e9086992f862714f844649b6849b23f0cacc2a3d9450af2b pseudo_PAW_PBEsol/S.upf +84584d07d4a6cfac3631799e2262dde986164df9ccb46ebd24847915cc0ed2dc pseudo_PAW_PBEsol/Sb.upf +65e271a90d96083809b2cd269eada462e8d2dd86c37b9823316bde1ed09d402e pseudo_PAW_PBEsol/Sc.upf +6bf49a97d32e78fa51fa39da1fc6cefb04763b15501ce72f672691927e61d41a pseudo_PAW_PBEsol/Se.upf +d2f59acaf5ed0d0c5bca36c97c99595c206b292acf70dd21c6665a1ea448c403 pseudo_PAW_PBEsol/Si.upf +d04cac215b1b345fa8fdfc17384634532ecb49b3085b74796168d88fa7f005cb pseudo_PAW_PBEsol/Sm.upf +5b18dceb51f8d5ae0a80e0e61234ad5299d2bc75f276b9fa700efab749463f56 pseudo_PAW_PBEsol/Sn.upf +bc0edb842a875086240976ef6b40a143df40d7b6ab26be969490a73c82e50643 pseudo_PAW_PBEsol/Sr.upf +ede9ce1e4f6c2aa045a5e5a41d79fd858de7307881a88a370a155fadfd60a179 pseudo_PAW_PBEsol/Ta.upf +34ea31919e7e5b57870601668f6b0ab23c0fc8f70732869933b2c5dd40f5dc7f pseudo_PAW_PBEsol/Tb.upf +0bf0cc94fe40d006f0019e65763a2295fd1383cbf6c957a1bcd33966376aae4d pseudo_PAW_PBEsol/Tc.upf +644d5d8686bb53a4d3089db27d06eb47ea850df1ddc7dd966827927b5210e195 pseudo_PAW_PBEsol/Te.upf +6c4983638781c1801f6b190030e2c38b36f4f2b6b1fd6e3b723e169044d3d500 pseudo_PAW_PBEsol/Ti.upf +60dbb0b7480bc30fdb97e1f77b9f17230cc9bc71bc9bd8825bcea752e73a402d pseudo_PAW_PBEsol/Tl.upf +f9a9811a36137aeb49d737739b04d96cd68f8760971a1c44952f8923f801d062 pseudo_PAW_PBEsol/Tm.upf +74766837b0a6dbbbc26ea9c26c7433793fb08b47c684e0cf0345af4290555ffe pseudo_PAW_PBEsol/V.upf +4404012b4ddf96cc5cb2e872f07a7328f0f1abd0d7d56f4429fe7e1de3cdcf1f pseudo_PAW_PBEsol/W.upf +8c9dd42e8e8cb9a90058c1e11e66da8b99df048905483ec192f78332126e4450 pseudo_PAW_PBEsol/Xe.upf +19817ece8cecf4d012dd1a0fee3f29538be8db75b1b88e3d47cb79f70e29a84b pseudo_PAW_PBEsol/Y.upf +416c54d3d20504cf2cf54773ff7618851840f93ebee586f31ae38b59d577c582 pseudo_PAW_PBEsol/Yb.upf +d7326e4d018653e735b266ffc023e9e3180ae9cd9efdc49571fc4bf0f8cea9d4 pseudo_PAW_PBEsol/Zn.upf +f282fba4f7fb4803678f168541ac28b9fde950734ca16c5b74193115ea901e97 pseudo_PAW_PBEsol/Zr.upf +7647cf9273cc207495764cc7f806e6e29341ff18c48b2fde7cb54f4726178c64 pseudo_USPP_PBEsol/Ag.upf +801a04431015c4326c8891f7f7b31e06ab51e5c0ad32493ee7436d28e82bd7a1 pseudo_USPP_PBEsol/Al.upf +2af5b207d0923b80542a4195082db45121d59b7f6e977616ead2d33465ae54f8 pseudo_USPP_PBEsol/Ar.upf +cdf99423e6aebe17641e1ee30c7ec68da9099903000394d47a6f55bd51dadecf pseudo_USPP_PBEsol/As.upf +af2872fcda537ac5dec334bd2759adbd1685d8b5dd9c06c58fa72c11288afbc5 pseudo_USPP_PBEsol/Au.upf +87ce207ef55d64e2f16efa9c368a1914c1dbaf1414bd204e2d72410f4bbd2641 pseudo_USPP_PBEsol/B.upf +b8114a98a94bfc0b9570acf1b4e583a75584675112ecb1972d7e521f92374ddf pseudo_USPP_PBEsol/Ba.upf +c23946724cd76b989e52863a7c4d10e40fa194a393e51d6bc847d28666aa400b pseudo_USPP_PBEsol/Be.upf +d2eea7a660cf839d1332de73d8c2845bad0d301d67d2d791f0374c37ace8ae9e pseudo_USPP_PBEsol/Bi.upf +cb6795b97049c40ca66443da5d725f9214ded6ca41a79c4eb3d2e7457713adc8 pseudo_USPP_PBEsol/Br.upf +d379b4cd2f440c357b310c3772554838c13124d7967e59e045429002efe7055b pseudo_USPP_PBEsol/C.upf +1875414aeb16344d59fef9e7959a5894899fddc21b034b239cc30a66da508e44 pseudo_USPP_PBEsol/Ca.upf +5d4602fbc2340e69c98516cb7a3f6102657778c30b687b296174f887b115d003 pseudo_USPP_PBEsol/Cd.upf +a26a6c021fa35d204c7a05491b69dc4208050a7457cbe2b3957bb90b4845f93a pseudo_USPP_PBEsol/Ce.upf +e302a65c1c989cbde50de95a81297e9f4401e0ef76d9a4178281c3abe7440441 pseudo_USPP_PBEsol/Cl.upf +05d65a6dbaba75cb4b9ad552d0113d25bb7cc91bf15ff16d3735f65d16e9aaa1 pseudo_USPP_PBEsol/Co.upf +afb3bea4a7b436f20ee6cae05e0ed6d266f6bac12c9c513de34f6f8b4ad0bda2 pseudo_USPP_PBEsol/Cr.upf +4bd6d33a2c99f681515df646712c10235d1f6ffc5809728d5af8215ae4dc4840 pseudo_USPP_PBEsol/Cs.upf +4aaa1d2269f0696ea232b3d0b8a901d66ebf58cc9ccf9cb81c383676b8dbee6e pseudo_USPP_PBEsol/Cu.upf +741b589abd977b578d36de1b984e6599dc48168ad3af2e54f5d76dded2d35d50 pseudo_USPP_PBEsol/Dy.upf +fc625713ca5c63de7084d45f664e01d4931d35af1100a94a77a1b13ef91f9a91 pseudo_USPP_PBEsol/Er.upf +26323dbc87d6fc71c48ef12dc2d7612bd1e87b20c12588cd793b1973873946bd pseudo_USPP_PBEsol/Eu.upf +3c23050eabbb40ad23379b2eb86337248eac6c6f443b74aace290bc7349c1242 pseudo_USPP_PBEsol/F.upf +c6a5cf098f72dcd22d1c4a74bbe56953619aadd64753026cb085d73c93598f01 pseudo_USPP_PBEsol/Fe.upf +129b768a4e97e04a5a33fd5ebedc858e30bd761d25562231ad0a7d3f54114042 pseudo_USPP_PBEsol/Ga.upf +5836ee0896d1fa88eb50063e3d48fb1a1d430e65e8e291c86ebc517d6765411c pseudo_USPP_PBEsol/Gd.upf +659ce92c49ab47a83ace5623561c7d4e557d34e1954565cbd53ba31406f288f8 pseudo_USPP_PBEsol/Ge.upf +5a07198cd1f0f3011bcac332163860abdd0eb2024a6b514a4f6dd12eb58c541d pseudo_USPP_PBEsol/H.upf +c93139553032cd2dacdf2b806486e92e07f643caa3bfcaee8028574f58b7a14a pseudo_USPP_PBEsol/He.upf +eb426d5a9a9a7cc8509fa5528a6bbc795b90fa83ee3ed9ce253c45d464781a9e pseudo_USPP_PBEsol/Hf.upf +08a6ba99461d7b5dd7b62800083a7dad1c51a20647b823cf85d8ac4ac0134664 pseudo_USPP_PBEsol/Hg.upf +2a56aa3df1427a6622ec295bc06fbbf3c03e48f54c9e257b9d0cb6f5b2359b64 pseudo_USPP_PBEsol/Ho.upf +4e1d523d29dc3435abc661df11e7eb71a6247c280c21f585a6cbc7de68791019 pseudo_USPP_PBEsol/I.upf +90ff08232379a881f5870dbc9e5a330f36fc45b941978907cccb2d0a963d6f7b pseudo_USPP_PBEsol/In.upf +7da93d6e375c4aa35c7fe3b47b23c1beb9164bb0e81f3f5e3fd9b78275b4c835 pseudo_USPP_PBEsol/Ir.upf +99db5c40ea33d74b0f20648b9cd61ea51cf30575a389f799266cf39d06063286 pseudo_USPP_PBEsol/K.upf +b43d9a0141156365062a5aea61fffae2f7d7bcc10e753f084d0d44ba69ad2d8f pseudo_USPP_PBEsol/Kr.upf +e3832e95b8d5e039120ef09ac7a948717c7d6ebe6b3cd17e3eed5a2df2a3c895 pseudo_USPP_PBEsol/La.upf +6cea3cb688b62fd4cbe137527d3df2efe8bebc6ef9c3bd62254b4c52a5ab749e pseudo_USPP_PBEsol/Li.upf +1f95ef6e6e1a7364f29d9515f78ff6b589999ce0b4604128c48495b55d472004 pseudo_USPP_PBEsol/Lu.upf +baa13b102c840c64e358eb64c38768b5b76765ceaef4ff3c22a54c2d5f41e791 pseudo_USPP_PBEsol/Mg.upf +74217ac6e49375a3e98397b18f98a4dc13263f0b2292818a1823495432b393bc pseudo_USPP_PBEsol/Mn.upf +eda027d30a7a8444461a2f1f5414fad4b4c4a3bb1908b0ab52b2df0ebdc750e5 pseudo_USPP_PBEsol/Mo.upf +2eea521c1eccac913da8c0c5ffc0d0ab86819c975a5dec7458e6a4fc07120354 pseudo_USPP_PBEsol/N.upf +62e0a0eb5bff3bd89e3039005a257a5e2f823b4fb5d39d4c5ff280b62114a202 pseudo_USPP_PBEsol/Na.upf +6910088685586744219051b12125da12ee78e9c70694295c3c59195ec75f4842 pseudo_USPP_PBEsol/Nb.upf +57b8177768ef07db25e4c549be370f04fc9660d7cc24636aa262bc749cc84497 pseudo_USPP_PBEsol/Nd.upf +002153d883e86f43428a0b67e09c7f6e43f0b468c620ff38c47d2c873f08d927 pseudo_USPP_PBEsol/Ne.upf +e293ca8dba4c646066dd63411b93d0fc07b1c2c3a8277a4b41f0f8966205b4cf pseudo_USPP_PBEsol/Ni.upf +fab367c8e8ebb07f7071a5256040af7f38e850cb0c326cac53b0123ebe1d1a1b pseudo_USPP_PBEsol/O.upf +c5b41389ed973dec21db2b9ff9891313c8e4876284a363279392d0fa36448418 pseudo_USPP_PBEsol/Os.upf +336bf7086784d62a8d415edafff2ba22fd62824035d3a7c593e74b0a7bf672e9 pseudo_USPP_PBEsol/P.upf +29bbae4cb95b640f7d62971ca84c96abe308886b9a8bdc910b27c2907a5bb808 pseudo_USPP_PBEsol/Pb.upf +3276e38ae2886a0bc2b8b595bae0a32bab639107047d3e4bdecebb2b4dfa9afc pseudo_USPP_PBEsol/Pd.upf +8a8d0a5060fcb83fb1dcfd0413ad8e9dc23e41314727939e525e757644ebfd02 pseudo_USPP_PBEsol/Pm.upf +0775358c854012e4945517d73304c9a775746cb7bebc483dbd59728d45f0a884 pseudo_USPP_PBEsol/Po.upf +fd6548b4810fff270fe673a679384cc62a5ca86c4e29afffe358261b8de2403f pseudo_USPP_PBEsol/Pr.upf +fa53f2ea9be66744abb5bbf503e2f863eb8bffcfb271ebbcf66b072b402292b8 pseudo_USPP_PBEsol/Pt.upf +00a0ffcccc4252b9cb2f494e061f29fc2b786bb0c2079096b30f47957c080ed0 pseudo_USPP_PBEsol/Rb.upf +bdd681aa1072daff4d2c11fcf041dbf6837d9317a5590f1106450489a9f56854 pseudo_USPP_PBEsol/Re.upf +11be90bd38aa55c771b969de4edfaf910e1f0457bdaf439c35d563e8cf424cbf pseudo_USPP_PBEsol/Rh.upf +ac145055ef8a9a358681e68a010344dceda18603698c05b023e433cb704bbacf pseudo_USPP_PBEsol/Rn.upf +4f1346e9f322c15f9ee1d7106d6efe6122ccbed6c1e574e5d1da84a83d6ea84f pseudo_USPP_PBEsol/Ru.upf +5785e57e7320fbd0292aa2f26938c4a4d595d4313fe1494b2d0e49072ac80438 pseudo_USPP_PBEsol/S.upf +50043644d41f233fd86caf966daa3d88cb1f482077dfceb1076d93f0ca9c4637 pseudo_USPP_PBEsol/Sb.upf +223e62c5acc55044f5f064778bc28ffac9ccd3be01490e8b0776c6686e08f723 pseudo_USPP_PBEsol/Sc.upf +f6f121d83feb9750517cd80b40911229a7749d3334ed89f1516da2f5417c0f4f pseudo_USPP_PBEsol/Se.upf +82e6d35f71dfa816fb8fcc7530464af3e724703d176d885ce5b35cc341ab3e4a pseudo_USPP_PBEsol/Si.upf +05f448a4e6575cf13e4f8a49f16f59dbfee12a416513f70ba26dcb0c3180f1a4 pseudo_USPP_PBEsol/Sm.upf +5fe98cd4c11ae4ae5dfd2827cb744148e3582660e9f69730b5eeb9c964c23cb1 pseudo_USPP_PBEsol/Sn.upf +c7c6fa80a86d360b51ed43e4f5154a6f1625a5d90b7215c95ef2591ebd366db8 pseudo_USPP_PBEsol/Sr.upf +34fb4d2ca944250bcdbec370146045115d5a495752b22b899777d8ba3ccd6b8d pseudo_USPP_PBEsol/Ta.upf +bb2f139cd9421de9dfc781f179180834a4758d8e0823eea9d8e24e9a2f4b2c91 pseudo_USPP_PBEsol/Tb.upf +ed06469c707087ed7e0a780671a8e01b8eba100f3c3c43030bda8d6589ef6821 pseudo_USPP_PBEsol/Tc.upf +a4fca9a03053725dd5b7e83ff29974dda64d614fb9cd8c1d25edea08fc5efd6b pseudo_USPP_PBEsol/Te.upf +9f50adf8d3f6a1b98981d5fb94205e53a523f405883dcb060c06b5eb93805f69 pseudo_USPP_PBEsol/Ti.upf +3451e912db7c2b7423f86b3a7e041cdd27de07a381d03bde0df21174b3330682 pseudo_USPP_PBEsol/Tl.upf +5ea437d39aabc01718c4a7a8638bb97b215ecad63810749623f86afbcf727694 pseudo_USPP_PBEsol/Tm.upf +b8ad92b3ca6dd6a2cf5fbc53eba1d30bd52bf311c9a9037491bb49893857689c pseudo_USPP_PBEsol/V.upf +df0c771d305ca3ace086e9eb4114e72c7e7ec8731157c866400a04feced89729 pseudo_USPP_PBEsol/W.upf +c40bc560b5a17e41c38b5c756a0054865088179a72c2d5475606672d1d131ed6 pseudo_USPP_PBEsol/Xe.upf +a8800467a40fcbbab313965aeb4869bc4c8d71446abd5c006192f984fc325e00 pseudo_USPP_PBEsol/Y.upf +4228a0dc23d687b3414efe3c844d8474d60f053152f7d8642cedd09dab353859 pseudo_USPP_PBEsol/Yb.upf +43c605383db2aa50e263a35c6a574bb2996af3ba35d8ee0abff474d18571f3aa pseudo_USPP_PBEsol/Zn.upf +6402fbcfbb6970779153fa49d988c8cfedf1b4245b295c89af1ed6b89e82164a pseudo_USPP_PBEsol/Zr.upf diff --git a/api/computations/qe.py b/api/computations/qe.py new file mode 100644 index 0000000..d36636d --- /dev/null +++ b/api/computations/qe.py @@ -0,0 +1,647 @@ +"""Quantum ESPRESSO preparation and parsing for the two adapted Sim2Ls. + +Heavy scientific dependencies are imported lazily so the existing API can +start normally while the RCAC feature is disabled. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import re +import shutil +import tarfile +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +RY_TO_EV = 13.605684958731 +RY_BOHR_TO_EV_ANGSTROM = 25.71104309541616 +BOHR_TO_ANGSTROM = 0.529177210903 + + +class QePreparationError(ValueError): + pass + + +class QeParseError(ValueError): + pass + + +@dataclass +class PreparedQeRun: + tool_name: str + local_workdir: Path + stage_files: list[Path] + command: list[str] + expected_artifacts: list[str] + job_array: str = "" + ntasks: int = 1 + time_limit: str = "24:00:00" + metadata: dict[str, Any] | None = None + + +def _pymatgen(): + try: + from pymatgen.analysis.elasticity.elastic import ElasticTensor + from pymatgen.analysis.elasticity.strain import Strain + from pymatgen.analysis.elasticity.stress import Stress + from pymatgen.core import Lattice, Structure + from pymatgen.io.pwscf import PWInput + from pymatgen.symmetry.analyzer import SpacegroupAnalyzer + except ImportError as exc: + raise QePreparationError( + "pymatgen is required by the RCAC computation worker" + ) from exc + return { + "ElasticTensor": ElasticTensor, + "Strain": Strain, + "Stress": Stress, + "Lattice": Lattice, + "Structure": Structure, + "PWInput": PWInput, + "SpacegroupAnalyzer": SpacegroupAnalyzer, + } + + +def _structure_from_dict(value: dict[str, Any]): + if not value: + raise QePreparationError( + "structure_dict is required after Materials Project resolution" + ) + try: + return _pymatgen()["Structure"].from_dict(value) + except Exception as exc: + raise QePreparationError( + "structure_dict is not a valid pymatgen structure" + ) from exc + + +def _elements(structure) -> list[str]: + return sorted( + { + element.symbol + for site in structure.sites + for element in site.species.elements + } + ) + + +def _pseudo_directory(pseudo_root: Path, pseudo: str, xc: str) -> Path: + # The source snapshots contain PBEsol pseudopotentials only. Explicitly + # reject combinations that would otherwise silently use the wrong set. + if xc != "PBEsol": + raise QePreparationError( + "The downloaded source contains only PBEsol pseudopotentials; " + f"cannot run exchange_correlation={xc!r}" + ) + directory = Path(pseudo_root) / f"pseudo_{pseudo}_PBEsol" + if not directory.is_dir(): + raise QePreparationError(f"pseudopotential directory not found: {directory}") + return directory + + +def required_pseudopotentials( + structure, pseudo_root: Path, pseudo: str, xc: str +) -> list[Path]: + directory = _pseudo_directory(pseudo_root, pseudo, xc) + paths = [] + for element in _elements(structure): + path = directory / f"{element}.upf" + if not path.is_file(): + raise QePreparationError( + f"missing {pseudo}/{xc} pseudopotential for element {element}" + ) + _verify_pseudopotential(path, directory.name) + paths.append(path) + return paths + + +def _verify_pseudopotential(path: Path, directory_name: str) -> None: + manifest_path = Path(__file__).with_name("pseudopotentials.sha256") + if not manifest_path.is_file(): + raise QePreparationError("pseudopotential checksum manifest is missing") + expected = None + relative = f"{directory_name}/{path.name}" + for line in manifest_path.read_text(encoding="utf-8").splitlines(): + digest, separator, name = line.partition(" ") + if separator and name == relative: + expected = digest + break + if expected is None: + raise QePreparationError( + f"pseudopotential is not approved by the manifest: {relative}" + ) + actual = hashlib.sha256(path.read_bytes()).hexdigest() + if actual != expected: + raise QePreparationError(f"pseudopotential checksum mismatch: {relative}") + + +def _qe_dictionaries(inputs: dict[str, Any], structure, *, calculation: str): + electrons = { + "electron_maxstep": inputs["e_steps"], + "conv_thr": inputs["scf_conv"], + } + system: dict[str, Any] = { + "ecutwfc": int(inputs["KE_cutoff"]), + "input_dft": inputs["exchange_correlation"], + "vdw_corr": inputs["vdw_corr"], + } + if inputs["is_metal"]: + system.update(occupations="smearing", smearing="gauss", degauss=0.02) + if inputs["spin_polar"]: + electrons["mixing_beta"] = 0.3 + system["nspin"] = 2 + for index in range(len(structure.symbol_set)): + system[f"starting_magnetization({index + 1})"] = 0.6 + if not inputs["is_metal"]: + system.update(occupations="smearing", smearing="gauss", degauss=0.02) + if inputs["hubbard_U_on"]: + electrons["mixing_beta"] = 0.3 + if calculation == "vc-relax" and inputs["ibrav"] != 0: + system.update( + { + "ibrav": inputs["ibrav"], + "a": structure.lattice.a, + "b": structure.lattice.b, + "c": structure.lattice.c, + # Preserve the r99 notebook's angle mapping for parity. + "cosAB": math.cos(math.radians(structure.lattice.angles[0])), + "cosAC": math.cos(math.radians(structure.lattice.angles[1])), + "cosBC": math.cos(math.radians(structure.lattice.angles[2])), + } + ) + if inputs["is_metal"]: + electrons["mixing_beta"] = 0.3 + + control = { + "pseudo_dir": "./", + "calculation": calculation, + "outdir": "./qe_tmp", + "tstress": True, + "nstep": inputs["i_steps"], + "etot_conv_thr": inputs["energy_conv"], + "forc_conv_thr": inputs["force_conv"], + "disk_io": "nowf", + } + return control, system, electrons + + +def _append_hubbard(path: Path, structure, inputs: dict[str, Any]) -> None: + if not inputs["hubbard_U_on"]: + return + lines = [f"HUBBARD {{{inputs['hubbard_projector']}}}\n"] + for element in _elements(structure): + try: + value = inputs["hubbard_U_values"][element] + orbital = inputs["hubbard_U_orbitals"][element] + except KeyError as exc: + raise QePreparationError( + f"missing Hubbard U data for element {element}" + ) from exc + lines.append(f"U {element}-{orbital} {value}\n") + with path.open("a", encoding="utf-8") as handle: + handle.writelines(lines) + + +def _write_pw_input( + path: Path, + structure, + inputs: dict[str, Any], + *, + calculation: str, + pseudo_dir: str = "./", +) -> None: + pmg = _pymatgen() + control, system, electrons = _qe_dictionaries( + inputs, structure, calculation=calculation + ) + control["pseudo_dir"] = pseudo_dir + if calculation == "vc-relax": + cell = { + "press": inputs["target_pressure"], + "press_conv_thr": inputs["pressure_conv"], + } + if inputs["ibrav"] != 0: + cell["cell_dofree"] = "ibrav" + else: + cell = None + pseudo = {element: f"{element}.upf" for element in _elements(structure)} + writer = pmg["PWInput"]( + structure=structure, + pseudo=pseudo, + control=control, + system=system, + electrons=electrons, + cell=cell, + kpoints_grid=[inputs["k_x"], inputs["k_y"], inputs["k_z"]], + ) + writer.write_file(path) + if calculation == "vc-relax" and inputs["ibrav"] != 0: + _remove_cell_parameters(path) + _append_hubbard(path, structure, inputs) + + +def _remove_cell_parameters(path: Path) -> None: + """Remove the explicit cell card when QE derives it from nonzero ibrav.""" + + lines = path.read_text(encoding="utf-8").splitlines(keepends=True) + for index, line in enumerate(lines): + if line.strip().startswith("CELL_PARAMETERS"): + del lines[index : index + 4] + path.write_text("".join(lines), encoding="utf-8") + return + raise QePreparationError("pymatgen did not emit the expected CELL_PARAMETERS card") + + +def prepare_cellrelax( + inputs: dict[str, Any], workdir: Path, pseudo_root: Path +) -> PreparedQeRun: + workdir = Path(workdir) + workdir.mkdir(parents=True, exist_ok=True) + structure = _structure_from_dict(inputs["structure_dict"]) + if inputs["target_pressure"] > 0: + # Preserve the source Sim2L's Murnaghan-style initial compression. + strain = ((1 + 4 * inputs["target_pressure"] / inputs["bulk_m"]) ** -0.25) ** ( + 1 / 3 + ) - 1 + structure.apply_strain(strain) + + input_path = workdir / "vcrelax.in" + _write_pw_input(input_path, structure, inputs, calculation="vc-relax") + pseudos = required_pseudopotentials( + structure, pseudo_root, inputs["pseudo"], inputs["exchange_correlation"] + ) + staged_pseudos = [] + for source in pseudos: + destination = workdir / source.name + shutil.copy2(source, destination) + staged_pseudos.append(destination) + + ntasks = min(64, max(1, 4 * len(structure.sites))) + return PreparedQeRun( + tool_name="cellrelaxdft", + local_workdir=workdir, + stage_files=[input_path, *staged_pseudos], + command=[ + "bash", + "-lc", + "mkdir -p qe_tmp && srun pw.x -in vcrelax.in > vcrelax.stdout", + ], + expected_artifacts=["vcrelax.in", "vcrelax.stdout"], + ntasks=ntasks, + time_limit="24:00:00", + metadata={"num_atoms": len(structure.sites), "elements": _elements(structure)}, + ) + + +def _parse_matrix( + lines: Sequence[str], start: int, columns: slice +) -> list[list[float]]: + try: + return [ + [float(value) for value in lines[index].split()[columns]] + for index in range(start, start + 3) + ] + except (IndexError, ValueError) as exc: + raise QeParseError("malformed matrix in Quantum ESPRESSO output") from exc + + +def parse_qe_output(path: Path, *, natoms: int | None = None) -> dict[str, Any]: + lines = Path(path).read_text(encoding="utf-8", errors="replace").splitlines() + if not lines: + raise QeParseError("Quantum ESPRESSO output is empty") + + ionic_energies: list[float] = [] + pressures: list[float] = [] + stresses: list[list[list[float]]] = [] + total_forces: list[float] = [] + atomic_forces: list[list[list[float]]] = [] + cell_blocks: list[tuple[str, list[list[float]]]] = [] + position_blocks: list[tuple[str, list[str], list[list[float]]]] = [] + + for index, line in enumerate(lines): + if "!" in line and "total energy" in line: + match = re.search(r"=\s*([-+0-9.Ee]+)\s+Ry", line) + if match: + ionic_energies.append(float(match.group(1))) + if "P=" in line: + match = re.search(r"P=\s*([-+0-9.Ee]+)", line) + if match: + pressures.append(float(match.group(1))) + stresses.append(_parse_matrix(lines, index + 1, slice(3, 6))) + if "Total force" in line: + match = re.search(r"Total force\s*=\s*([-+0-9.Ee]+)", line) + if match: + total_forces.append(float(match.group(1))) + if "Forces acting on atoms" in line and natoms: + forces = [] + cursor = index + 1 + while cursor < len(lines) and len(forces) < natoms: + match = re.search( + r"force\s*=\s*([-+0-9.Ee]+)\s+([-+0-9.Ee]+)\s+([-+0-9.Ee]+)", + lines[cursor], + ) + if match: + forces.append([float(match.group(i)) for i in range(1, 4)]) + cursor += 1 + if len(forces) == natoms: + atomic_forces.append(forces) + if line.strip().startswith("CELL_PARAMETERS"): + unit = line.partition("(")[2].partition(")")[0].strip().lower() or "alat" + cell_blocks.append((unit, _parse_matrix(lines, index + 1, slice(0, 3)))) + if line.strip().startswith("ATOMIC_POSITIONS"): + unit = line.partition("(")[2].partition(")")[0].strip().lower() or "alat" + species: list[str] = [] + positions: list[list[float]] = [] + cursor = index + 1 + while cursor < len(lines): + parts = lines[cursor].split() + if len(parts) < 4 or not re.fullmatch(r"[A-Za-z]{1,3}", parts[0]): + break + try: + position = [float(value) for value in parts[1:4]] + except ValueError: + break + species.append(parts[0]) + positions.append(position) + cursor += 1 + if species: + position_blocks.append((unit, species, positions)) + + if not ionic_energies or not pressures or not stresses: + tail = "\n".join(lines[-5:]) + raise QeParseError(f"Quantum ESPRESSO output is incomplete. Tail:\n{tail}") + + result: dict[str, Any] = { + "ionic_energies_ry": ionic_energies, + "pressures_kbar": pressures, + "stresses_kbar": stresses, + "total_forces_ry_bohr": total_forces, + "atomic_forces_ry_bohr": atomic_forces, + "convergence": any( + marker in line + for marker in ("End of BFGS Geometry Optimization", "JOB DONE.") + for line in lines + ), + } + + if cell_blocks and position_blocks: + unit, lattice = cell_blocks[-1] + if unit.startswith("bohr"): + lattice = [[value * BOHR_TO_ANGSTROM for value in row] for row in lattice] + elif unit.startswith("alat"): + alat = _last_alat_angstrom(lines) + lattice = [[value * alat for value in row] for row in lattice] + pos_unit, species, positions = position_blocks[-1] + coords_are_cartesian = pos_unit.startswith(("angstrom", "bohr")) + if pos_unit.startswith("bohr"): + positions = [ + [value * BOHR_TO_ANGSTROM for value in position] + for position in positions + ] + structure = _pymatgen()["Structure"]( + lattice, + species, + positions, + coords_are_cartesian=coords_are_cartesian, + ) + result["structure"] = structure + return result + + +def _last_alat_angstrom(lines: Sequence[str]) -> float: + for line in reversed(lines): + match = re.search(r"lattice parameter \(alat\)\s*=\s*([-+0-9.Ee]+)", line) + if match: + return float(match.group(1)) * BOHR_TO_ANGSTROM + raise QeParseError("CELL_PARAMETERS uses alat but no lattice parameter was found") + + +def cellrelax_outputs(inputs: dict[str, Any], output_path: Path) -> dict[str, Any]: + initial = _structure_from_dict(inputs["structure_dict"]) + parsed = parse_qe_output(output_path, natoms=len(initial.sites)) + final = parsed.get("structure") + if final is None: + raise QeParseError("final relaxed structure is missing") + pmg = _pymatgen() + n_units = final.composition.get_integer_formula_and_factor()[1] + stress = parsed["stresses_kbar"][-1] + forces = ( + parsed["atomic_forces_ry_bohr"][-1] if parsed["atomic_forces_ry_bohr"] else [] + ) + structure_dict = final.as_dict() + structure_dict.get("lattice", {}).pop("pbc", None) + return { + "num_atoms": len(initial.sites), + "chemical_system": initial.chemical_system, + "composition": initial.composition.as_dict(), + "formula_unit": initial.composition.reduced_formula, + "mp_space_group": pmg["SpacegroupAnalyzer"](initial).get_space_group_number(), + "energy": parsed["ionic_energies_ry"][-1] * RY_TO_EV / n_units, + "atomic_forces": [ + [value * RY_BOHR_TO_EV_ANGSTROM for value in row] for row in forces + ], + "qe_space_group": final.get_space_group_info()[1], + "pressure": parsed["pressures_kbar"][-1] / 10.0, + "stress": [[value / 10.0 for value in row] for row in stress], + "volume": final.volume / n_units, + "structure_dict": structure_dict, + "lattice_vectors": final.lattice.matrix.tolist(), + "lattice_parameters": list(final.lattice.abc), + "angles": list(final.lattice.angles), + "convergence": parsed["convergence"], + } + + +def inherited_elastic_parameters( + elastic_inputs: dict[str, Any], + upstream_inputs: dict[str, Any], + upstream_outputs: dict[str, Any], +) -> dict[str, Any]: + if elastic_inputs["mp_id"] != upstream_inputs["mp_id"]: + raise QePreparationError("ElasticDFT mp_id does not match its CellRelaxDFT run") + if not upstream_outputs.get("convergence") or not upstream_outputs.get( + "structure_dict" + ): + raise QePreparationError("CellRelaxDFT upstream run is not converged") + merged = dict(upstream_inputs) + merged["structure_dict"] = upstream_outputs["structure_dict"] + return merged + + +def _elastic_strains(inputs: dict[str, Any]): + Strain = _pymatgen()["Strain"] + strains = [] + for index in [(0, 0), (1, 1), (2, 2)]: + strains.extend( + Strain.from_index_amount(index, amount) + for amount in inputs["diagonal_strain_magnitudes"] + ) + for index in [(0, 1), (0, 2), (1, 2)]: + strains.extend( + Strain.from_index_amount(index, amount) + for amount in inputs["offdiagonal_strain_magnitudes"] + ) + return strains + + +def prepare_elastic( + elastic_inputs: dict[str, Any], + upstream_inputs: dict[str, Any], + upstream_outputs: dict[str, Any], + workdir: Path, + pseudo_root: Path, +) -> PreparedQeRun: + workdir = Path(workdir) + workdir.mkdir(parents=True, exist_ok=True) + inherited = inherited_elastic_parameters( + elastic_inputs, upstream_inputs, upstream_outputs + ) + structure = _structure_from_dict(inherited["structure_dict"]) + if inherited["is_metal"]: + spacing = 0.16 + else: + spacing = 0.22 + inherited.update( + { + "k_x": max(round((2 * math.pi / structure.lattice.a) / spacing), 1), + "k_y": max(round((2 * math.pi / structure.lattice.b) / spacing), 1), + "k_z": max(round((2 * math.pi / structure.lattice.c) / spacing), 1), + } + ) + + package_root = workdir / "elastic_inputs" + pseudo_dir = package_root / "pseudos" + pseudo_dir.mkdir(parents=True) + for source in required_pseudopotentials( + structure, + pseudo_root, + inherited["pseudo"], + inherited["exchange_correlation"], + ): + shutil.copy2(source, pseudo_dir / source.name) + + strains = _elastic_strains(elastic_inputs) + deformed_structures = [] + for index, strain in enumerate(strains): + deformed = strain.get_deformation_matrix().apply_to_structure(structure) + deformed_structures.append(deformed.as_dict()) + directory = package_root / f"strain_{index:03d}" + directory.mkdir() + _write_pw_input( + directory / "relax.in", + deformed, + inherited, + calculation="relax", + pseudo_dir="../pseudos", + ) + + manifest = { + "strain_matrices": [strain.tolist() for strain in strains], + "deformed_structures": deformed_structures, + "count": len(strains), + "inherited_inputs": inherited, + "external_pressure": upstream_outputs["pressure"], + } + (package_root / "manifest.json").write_text( + json.dumps(manifest, sort_keys=True), encoding="utf-8" + ) + archive = workdir / "elastic-inputs.tar.gz" + with tarfile.open(archive, "w:gz") as handle: + handle.add(package_root, arcname="elastic_inputs") + + runner = workdir / "run-elastic-task.sh" + runner.write_text( + """#!/bin/bash +set -euo pipefail +task=$(printf '%03d' "${SLURM_ARRAY_TASK_ID}") +task_root="task_${task}" +mkdir -p "${task_root}" +tar -xzf elastic-inputs.tar.gz -C "${task_root}" +cd "${task_root}/elastic_inputs/strain_${task}" +mkdir -p qe_tmp +srun pw.x -in relax.in > relax.stdout +""", + encoding="utf-8", + ) + runner.chmod(0o750) + concurrency = min(4, len(strains)) + expected = [ + f"task_{index:03d}/elastic_inputs/strain_{index:03d}/relax.stdout" + for index in range(len(strains)) + ] + return PreparedQeRun( + tool_name="elasticdft", + local_workdir=workdir, + stage_files=[archive, runner], + command=["bash", runner.name], + expected_artifacts=expected, + job_array=f"0-{len(strains) - 1}%{concurrency}", + ntasks=64, + time_limit="08:00:00", + metadata=manifest, + ) + + +def elastic_outputs(prepared: PreparedQeRun) -> dict[str, Any]: + metadata = prepared.metadata or {} + strain_matrices = metadata["strain_matrices"] + structures = metadata["deformed_structures"] + inherited = metadata["inherited_inputs"] + original = _structure_from_dict(inherited["structure_dict"]) + stresses = [] + energies = [] + pmg = _pymatgen() + for relative in prepared.expected_artifacts: + parsed = parse_qe_output(prepared.local_workdir / relative) + if not parsed["convergence"]: + raise QeParseError(f"ElasticDFT member did not converge: {relative}") + stresses.append( + pmg["Stress"]( + [ + [-value / 10.0 for value in row] + for row in parsed["stresses_kbar"][-1] + ] + ) + ) + n_units = original.composition.get_integer_formula_and_factor()[1] + energies.append(parsed["ionic_energies_ry"][-1] * RY_TO_EV / n_units) + + strain_objects = [pmg["Strain"](matrix) for matrix in strain_matrices] + tensor = pmg["ElasticTensor"].from_pseudoinverse( + strains=strain_objects, stresses=stresses + ) + voigt = tensor.voigt.tolist() + result: dict[str, Any] = { + "external_pressure": metadata["external_pressure"], + "strain_matrices": strain_matrices, + "stress_tensors": [stress.tolist() for stress in stresses], + "deformed_structures": structures, + "energies": energies, + "elastic_tensor": voigt, + "bulk_modulus_reuss": float(tensor.k_reuss), + "bulk_modulus_voigt": float(tensor.k_voigt), + "bulk_modulus_vrh": float(tensor.k_vrh), + "shear_modulus_reuss": float(tensor.g_reuss), + "shear_modulus_voigt": float(tensor.g_voigt), + "shear_modulus_vrh": float(tensor.g_vrh), + "homogeneous_poisson": float(tensor.homogeneous_poisson), + "youngs_modulus": float(tensor.y_mod * 1e-9), + "calc_params": { + "pseudopotential": inherited["pseudo"], + "exchange_correlation": inherited["exchange_correlation"], + "vdw_correction": inherited["vdw_corr"], + "kmesh": [inherited["k_x"], inherited["k_y"], inherited["k_z"]], + "kinetic_energy_cutoff": inherited["KE_cutoff"], + }, + "chemical_system": original.chemical_system, + "composition": original.composition.as_dict(), + "formula_unit": original.composition.reduced_formula, + } + for i in range(6): + for j in range(i, 6): + result[f"C{i + 1}{j + 1}"] = float(voigt[i][j]) + return result diff --git a/api/computations/rcac.py b/api/computations/rcac.py new file mode 100644 index 0000000..e96d1ec --- /dev/null +++ b/api/computations/rcac.py @@ -0,0 +1,116 @@ +"""RCAC cluster configuration and credential-safe transport construction.""" + +from __future__ import annotations + +import json +import os +import re +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import Any + + +@dataclass(frozen=True) +class ClusterConfig: + name: str + host: str + qe_module: str + remote_root: str + port: int = 22 + + +def configured_clusters() -> dict[str, ClusterConfig]: + raw = os.getenv("RCAC_CLUSTERS_JSON", "").strip() + if not raw: + return {} + try: + values = json.loads(raw) + except json.JSONDecodeError as exc: + raise RuntimeError("RCAC_CLUSTERS_JSON is invalid JSON") from exc + if not isinstance(values, dict): + raise TypeError("RCAC_CLUSTERS_JSON must be an object") + clusters = {} + for name, value in values.items(): + if not re.fullmatch(r"[A-Za-z0-9_-]+", name) or not isinstance(value, dict): + raise RuntimeError("RCAC cluster names must be safe identifiers") + try: + cluster = ClusterConfig( + name=name, + host=str(value["host"]), + qe_module=str(value["qe_module"]), + remote_root=str(value["remote_root"]), + port=int(value.get("port", 22)), + ) + except (KeyError, TypeError, ValueError) as exc: + raise RuntimeError(f"RCAC cluster {name!r} is incomplete") from exc + root = PurePosixPath(cluster.remote_root) + if not root.is_absolute() or ".." in root.parts: + raise RuntimeError(f"RCAC cluster {name!r} has an unsafe remote_root") + clusters[name] = cluster + return clusters + + +def get_cluster(name: str) -> ClusterConfig: + try: + return configured_clusters()[name] + except KeyError as exc: + raise ValueError(f"RCAC cluster is not configured: {name}") from exc + + +def resolve_identity_file(credential_ref: str) -> Path: + if not re.fullmatch(r"[A-Za-z0-9_.-]+", credential_ref): + raise ValueError("credential_ref must be an opaque safe identifier") + root_value = os.getenv("RCAC_SSH_KEY_DIR", "").strip() + if not root_value: + raise RuntimeError("RCAC_SSH_KEY_DIR is not configured") + root = Path(root_value).resolve() + path = (root / credential_ref).resolve() + if path.parent != root or not path.is_file(): + raise RuntimeError("configured RCAC credential was not provisioned") + if path.stat().st_mode & 0o077: + raise RuntimeError("RCAC private key permissions must be 0600 or stricter") + return path + + +def make_slurm_executor(profile: dict[str, Any], local_workdir: Path): + """Create the matching Sim2L executor without importing it at API startup.""" + + try: + from sim2l.executor import LocalTransport, SlurmExecutor, SshTransport + except ImportError as exc: + raise RuntimeError("the matching Sim2L package is not installed") from exc + + cluster = get_cluster(profile["cluster"]) + transport_mode = os.getenv("RCAC_TRANSPORT", "ssh").strip().lower() + if transport_mode == "local": + if os.getenv("RCAC_ALLOW_LOCAL_TRANSPORT", "false").lower() != "true": + raise RuntimeError( + "local RCAC transport is allowed only in explicit test mode" + ) + transport = LocalTransport() + elif transport_mode == "ssh": + known_hosts = os.getenv("RCAC_KNOWN_HOSTS_FILE", "").strip() + if not known_hosts: + raise RuntimeError("RCAC_KNOWN_HOSTS_FILE is required for SSH transport") + transport = SshTransport( + host=cluster.host, + username=profile["ssh_username"], + identity_file=resolve_identity_file(profile["credential_ref"]), + port=cluster.port, + known_hosts_file=Path(known_hosts), + ) + else: + raise RuntimeError(f"unsupported RCAC_TRANSPORT: {transport_mode}") + return SlurmExecutor(transport, local_workdir), cluster + + +def remote_workdir(cluster: ClusterConfig, username: str, run_id: str) -> str: + if not re.fullmatch(r"[A-Za-z0-9._-]+", username): + raise ValueError("invalid RCAC username") + if not re.fullmatch(r"[0-9a-fA-F-]{36}", run_id): + raise ValueError("invalid computation run id") + root = cluster.remote_root.replace("{username}", username) + path = PurePosixPath(root) / run_id + if not path.is_absolute() or ".." in path.parts: + raise ValueError("computed RCAC work directory is unsafe") + return str(path) diff --git a/api/computations/results.py b/api/computations/results.py new file mode 100644 index 0000000..3995d74 --- /dev/null +++ b/api/computations/results.py @@ -0,0 +1,43 @@ +"""Optional direct registration in Daniel's Sim2L Results service.""" + +from __future__ import annotations + +import os +from datetime import datetime +from typing import Any + + +def register_sim2l_result( + run: dict[str, Any], outputs: dict[str, Any] +) -> dict[str, Any]: + url = os.getenv("SIM2L_RESULTS_URL", "").strip() + if not url: + return {"status": "not_configured"} + try: + from sim2l.database import ResultsClient + except ImportError as exc: + raise RuntimeError("the matching Sim2L package is not installed") from exc + started = run.get("started_at") + duration = None + if isinstance(started, datetime): + duration = max(0.0, (datetime.now(started.tzinfo) - started).total_seconds()) + client = ResultsClient( + url, + session_id=os.getenv("SIM2L_RESULTS_SESSION_ID", "").strip() or None, + ) + return client.register_direct( + execution_id=str(run["id"]), + simulation_name=run["tool_name"], + simulation_version=run["tool_revision"], + squid_id=run["sim2l_squid_id"], + input_params=run["input_json"], + output_params=outputs, + status="completed", + duration_seconds=duration, + metadata={ + "project_id": run["project_id"], + "cluster": run["cluster"], + "scheduler_job_id": run["scheduler_job_id"], + "source": "birck-digital-twin", + }, + ) diff --git a/api/computations/schemas.py b/api/computations/schemas.py new file mode 100644 index 0000000..7edca1a --- /dev/null +++ b/api/computations/schemas.py @@ -0,0 +1,209 @@ +"""Stable input contracts derived from CellRelaxDFT r99 and ElasticDFT r27.""" + +from __future__ import annotations + +import codecs +import hashlib +import json +from typing import Any, Literal +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +class CellRelaxInputs(BaseModel): + """Inputs retained from the canonical CellRelaxDFT r99 Sim2L.""" + + model_config = ConfigDict(extra="forbid") + + mp_id: str = Field(pattern=r"^mp-[A-Za-z0-9-]+$") + structure_dict: dict[str, Any] = Field(default_factory=dict) + is_metal: bool = True + bulk_m: float = Field(default=2000.0, ge=1, le=6000) + KE_cutoff: float = Field(default=50.0, ge=10, le=300) + k_x: int = Field(default=1, ge=1, le=30) + k_y: int = Field(default=1, ge=1, le=30) + k_z: int = Field(default=1, ge=1, le=30) + i_steps: int = Field(default=50, ge=1, le=1000) + e_steps: int = Field(default=100, ge=100, le=1000) + scf_conv: float = Field(default=1e-10, gt=0) + energy_conv: float = Field(default=1e-10, ge=1e-15, le=1e-4) + force_conv: float = Field(default=1e-5, ge=1e-8, le=1e-3) + pressure_conv: float = Field(default=0.5, ge=0.01, le=10) + target_pressure: float = Field(default=0.0, ge=-100, le=10000) + vdw_corr: Literal["grimme-d2", "grimme-d3", "none"] = "none" + spin_polar: bool = False + hubbard_U_on: bool = False + hubbard_projector: Literal[ + "atomic", "ortho-atomic", "norm-atomic", "wf", "pseudo" + ] = "ortho-atomic" + hubbard_U_values: dict[str, float] = Field(default_factory=dict) + hubbard_U_orbitals: dict[str, str] = Field(default_factory=dict) + ibrav: int = 0 + exchange_correlation: Literal["LDA", "PBE", "PBEsol"] = "PBEsol" + # The source offers NC, but the r27 bundle contains only PAW and USPP. + pseudo: Literal["PAW", "USPP"] = "PAW" + + @field_validator("hubbard_U_orbitals") + @classmethod + def orbitals_are_safe(cls, value: dict[str, str]) -> dict[str, str]: + for orbital in value.values(): + if not orbital or not orbital.replace("-", "").isalnum(): + raise ValueError( + "Hubbard orbitals must contain letters, digits, or '-'" + ) + return value + + def validate_cross_fields(self) -> CellRelaxInputs: + if self.hubbard_U_on: + if not self.hubbard_U_values or not self.hubbard_U_orbitals: + raise ValueError( + "Hubbard U values and orbitals are required when hubbard_U_on is true" + ) + if set(self.hubbard_U_values) != set(self.hubbard_U_orbitals): + raise ValueError("Hubbard U value and orbital elements must match") + return self + + +class ElasticInputs(BaseModel): + """Inputs for ElasticDFT with a DT run id replacing nanoHUB squid lookup.""" + + model_config = ConfigDict(extra="forbid") + + mp_id: str = Field(pattern=r"^mp-[A-Za-z0-9-]+$") + upstream_execution_id: UUID + diagonal_strain_magnitudes: list[float] = Field( + default_factory=lambda: [-0.01, -0.005, 0.005, 0.01], + min_length=2, + max_length=12, + ) + offdiagonal_strain_magnitudes: list[float] = Field( + default_factory=lambda: [-0.01, -0.005, 0.005, 0.01], + min_length=2, + max_length=12, + ) + + @field_validator("diagonal_strain_magnitudes", "offdiagonal_strain_magnitudes") + @classmethod + def strains_are_small_and_distinct(cls, values: list[float]) -> list[float]: + values = [float(value) for value in values] + if any(value == 0 or abs(value) > 0.05 for value in values): + raise ValueError( + "strain magnitudes must be non-zero and no larger than 0.05" + ) + if len(set(values)) != len(values): + raise ValueError("strain magnitudes must be unique") + if not any(value < 0 for value in values) or not any( + value > 0 for value in values + ): + raise ValueError("strain magnitudes must include compression and tension") + return values + + +class ToolDefinition(BaseModel): + name: str + title: str + revision: str + description: str + input_model: Any + output_names: list[str] + + model_config = ConfigDict(arbitrary_types_allowed=True) + + def public_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "title": self.title, + "revision": self.revision, + "description": self.description, + "input_schema": self.input_model.model_json_schema(), + "output_names": self.output_names, + } + + +CELLRELAX_OUTPUTS = [ + "num_atoms", + "chemical_system", + "composition", + "formula_unit", + "mp_space_group", + "energy", + "atomic_forces", + "qe_space_group", + "pressure", + "stress", + "volume", + "structure_dict", + "lattice_vectors", + "lattice_parameters", + "angles", + "convergence", +] + +ELASTIC_OUTPUTS = [ + "external_pressure", + "strain_matrices", + "stress_tensors", + "deformed_structures", + "energies", + "elastic_tensor", + *[f"C{i}{j}" for i in range(1, 7) for j in range(i, 7)], + "bulk_modulus_reuss", + "bulk_modulus_voigt", + "bulk_modulus_vrh", + "shear_modulus_reuss", + "shear_modulus_voigt", + "shear_modulus_vrh", + "homogeneous_poisson", + "youngs_modulus", + "calc_params", + "chemical_system", + "composition", + "formula_unit", +] + +TOOL_DEFINITIONS = { + "cellrelaxdft": ToolDefinition( + name="cellrelaxdft", + title="Cell Relax DFT", + revision="r99-dt1", + description="Relax a crystal cell with Quantum ESPRESSO on RCAC.", + input_model=CellRelaxInputs, + output_names=CELLRELAX_OUTPUTS, + ), + "elasticdft": ToolDefinition( + name="elasticdft", + title="Elastic DFT", + revision="r27-dt1", + description="Fit an elastic tensor from strained Quantum ESPRESSO jobs.", + input_model=ElasticInputs, + output_names=ELASTIC_OUTPUTS, + ), +} + + +def validate_tool_inputs(tool_name: str, inputs: dict[str, Any]) -> dict[str, Any]: + try: + definition = TOOL_DEFINITIONS[tool_name] + except KeyError as exc: + raise ValueError(f"Unknown computation tool: {tool_name}") from exc + model = definition.input_model.model_validate(inputs) + if isinstance(model, CellRelaxInputs): + model.validate_cross_fields() + return model.model_dump(mode="json") + + +def compute_squid_id(tool_name: str, inputs: dict[str, Any]) -> str: + """Generate the same name/revision/SHA1 shape used by Daniel's Sim2L.""" + + definition = TOOL_DEFINITIONS[tool_name] + normalized = validate_tool_inputs(tool_name, inputs) + normalized = json.loads(json.dumps(normalized, sort_keys=True)) + prologue = [f"{key} {normalized[key]}" for key in sorted(normalized)] + prologue.extend( + [f"simtoolName {definition.name}", f"simtoolRevision {definition.revision}"] + ) + digest = hashlib.sha1( + memoryview(codecs.encode("\n".join(prologue), "utf-8")) + ).hexdigest() + return f"{definition.name}/{definition.revision}/{digest}" diff --git a/api/computations/store.py b/api/computations/store.py new file mode 100644 index 0000000..b3e9df6 --- /dev/null +++ b/api/computations/store.py @@ -0,0 +1,617 @@ +"""PostgreSQL persistence and project authorization for computation runs.""" + +from __future__ import annotations + +import json +import os +from datetime import datetime +from pathlib import Path +from typing import Any + +from data_loader_pg import get_pg_superuser_connection +from metadata_pg import _ensure_project_access, _ensure_project_write_access +from psycopg2.extras import Json, RealDictCursor +from security import PlatformUser + +from computations.schemas import TOOL_DEFINITIONS + +ACTIVE_STATUSES = { + "requested", + "validating", + "staging", + "submitted", + "running", + "collecting", + "cancelling", +} +TERMINAL_STATUSES = {"data_ingested", "cancelled", "failed"} +ALL_STATUSES = ACTIVE_STATUSES | TERMINAL_STATUSES + + +def feature_enabled() -> bool: + return os.getenv("RCAC_COMPUTE_ENABLED", "false").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +def ensure_tables() -> None: + path = Path(__file__).resolve().parents[1] / "scripts" / "add_rcac_computations.sql" + sql = path.read_text(encoding="utf-8") + conn = get_pg_superuser_connection() + try: + with conn.cursor() as cur: + cur.execute(sql) + conn.commit() + finally: + conn.close() + + +def _json(value: Any, default: Any) -> Any: + if value is None: + return default + if isinstance(value, str): + try: + return json.loads(value) + except json.JSONDecodeError: + return default + return value + + +def _iso(value: Any) -> str: + if isinstance(value, datetime): + return value.isoformat() + return str(value or "") + + +def _public_run(row: dict[str, Any], *, include_events: bool = False) -> dict[str, Any]: + value = dict(row) + for key in ( + "input_json", + "output_json", + "artifact_manifest_json", + "external_workflow_ref_json", + "error_json", + "prepared_manifest_json", + ): + public_key = key.removesuffix("_json") + default: Any = [] if key == "artifact_manifest_json" else {} + value[public_key] = _json(value.pop(key, None), default) + for key in ( + "created_at", + "updated_at", + "submitted_at", + "started_at", + "completed_at", + "worker_lease_expires_at", + ): + value[key] = _iso(value.get(key)) + value["id"] = str(value.get("id") or "") + value["upstream_run_id"] = str(value.get("upstream_run_id") or "") + value["artifact_manifest"] = [ + { + **artifact, + "uri": f"/api/dt/computations/runs/{value['id']}/artifacts/{index}", + } + for index, artifact in enumerate(value.get("artifact_manifest", [])) + ] + value.pop("worker_lease_owner", None) + if not include_events: + value.pop("events", None) + return value + + +def upsert_profile( + *, user: PlatformUser, cluster: str, payload: dict[str, Any] +) -> dict[str, Any]: + conn = get_pg_superuser_connection() + try: + with conn.cursor(cursor_factory=RealDictCursor) as cur: + cur.execute( + """ + INSERT INTO rcac_user_profiles ( + user_id, cluster, ssh_username, allocation, partition, + credential_ref, enabled, updated_at + ) VALUES (%s, %s, %s, %s, %s, %s, true, CURRENT_TIMESTAMP) + ON CONFLICT (user_id, cluster) DO UPDATE SET + ssh_username = EXCLUDED.ssh_username, + allocation = EXCLUDED.allocation, + partition = EXCLUDED.partition, + credential_ref = EXCLUDED.credential_ref, + enabled = true, + updated_at = CURRENT_TIMESTAMP + RETURNING user_id, cluster, ssh_username, allocation, partition, + credential_ref, enabled, created_at, updated_at + """, + ( + user.id, + cluster, + payload["ssh_username"], + payload["allocation"], + payload["partition"], + payload["credential_ref"], + ), + ) + row = dict(cur.fetchone()) + conn.commit() + finally: + conn.close() + row["credential_configured"] = bool(row.pop("credential_ref", "")) + row["created_at"] = _iso(row.get("created_at")) + row["updated_at"] = _iso(row.get("updated_at")) + return row + + +def get_profile( + user_id: str, cluster: str, *, internal: bool = False +) -> dict[str, Any] | None: + conn = get_pg_superuser_connection() + try: + with conn.cursor(cursor_factory=RealDictCursor) as cur: + cur.execute( + """ + SELECT user_id, cluster, ssh_username, allocation, partition, + credential_ref, enabled, created_at, updated_at + FROM rcac_user_profiles + WHERE user_id = %s AND cluster = %s AND enabled = true + """, + (user_id, cluster), + ) + row = cur.fetchone() + finally: + conn.close() + if not row: + return None + value = dict(row) + if not internal: + value["credential_configured"] = bool(value.pop("credential_ref", "")) + value["created_at"] = _iso(value.get("created_at")) + value["updated_at"] = _iso(value.get("updated_at")) + return value + + +def create_run( + *, + user: PlatformUser, + project_id: str, + tool_name: str, + tool_revision: str, + squid_id: str, + inputs: dict[str, Any], + cluster: str, + partition: str, + allocation: str, + upstream_run_id: str = "", +) -> dict[str, Any]: + conn = get_pg_superuser_connection() + try: + with conn.cursor(cursor_factory=RealDictCursor) as cur: + _ensure_project_write_access(cur, project_id=project_id, user=user) + _ensure_tool_catalog( + cur, user=user, project_id=project_id, tool_name=tool_name + ) + if upstream_run_id: + cur.execute( + """ + SELECT project_id, tool_name, status + FROM computation_runs WHERE id = %s + """, + (upstream_run_id,), + ) + upstream = cur.fetchone() + if not upstream or upstream["project_id"] != project_id: + raise ValueError( + "upstream computation was not found in this project" + ) + if ( + upstream["tool_name"] != "cellrelaxdft" + or upstream["status"] != "data_ingested" + ): + raise ValueError("ElasticDFT requires a completed CellRelaxDFT run") + cur.execute( + """ + INSERT INTO computation_runs ( + project_id, requested_by, tool_name, tool_revision, + sim2l_squid_id, upstream_run_id, cluster, partition, + allocation, input_json + ) VALUES (%s, %s, %s, %s, %s, NULLIF(%s, '')::uuid, %s, %s, %s, %s) + RETURNING * + """, + ( + project_id, + user.id, + tool_name, + tool_revision, + squid_id, + upstream_run_id, + cluster, + partition, + allocation, + Json(inputs), + ), + ) + row = dict(cur.fetchone()) + _insert_event(cur, str(row["id"]), "requested", "Execution requested") + conn.commit() + return _public_run(row) + finally: + conn.close() + + +def list_runs(project_id: str, user: PlatformUser) -> list[dict[str, Any]]: + conn = get_pg_superuser_connection() + try: + with conn.cursor(cursor_factory=RealDictCursor) as cur: + _ensure_project_access(cur, project_id=project_id, user=user) + cur.execute( + "SELECT * FROM computation_runs WHERE project_id = %s ORDER BY created_at DESC", + (project_id,), + ) + return [_public_run(dict(row)) for row in cur.fetchall()] + finally: + conn.close() + + +def get_run(run_id: str, user: PlatformUser | None = None, *, internal: bool = False): + conn = get_pg_superuser_connection() + try: + with conn.cursor(cursor_factory=RealDictCursor) as cur: + cur.execute("SELECT * FROM computation_runs WHERE id = %s", (run_id,)) + row = cur.fetchone() + if not row: + return None + if user is not None: + _ensure_project_access(cur, project_id=row["project_id"], user=user) + if internal: + return dict(row) + cur.execute( + """ + SELECT status, message, details_json, created_at + FROM computation_run_events WHERE run_id = %s ORDER BY created_at + """, + (run_id,), + ) + value = dict(row) + value["events"] = [ + { + "status": event["status"], + "message": event["message"], + "details": _json(event["details_json"], {}), + "created_at": _iso(event["created_at"]), + } + for event in cur.fetchall() + ] + return _public_run(value, include_events=True) + finally: + conn.close() + + +def request_cancel(run_id: str, user: PlatformUser) -> dict[str, Any] | None: + conn = get_pg_superuser_connection() + try: + with conn.cursor(cursor_factory=RealDictCursor) as cur: + cur.execute( + "SELECT * FROM computation_runs WHERE id = %s FOR UPDATE", (run_id,) + ) + row = cur.fetchone() + if not row: + return None + _ensure_project_write_access(cur, project_id=row["project_id"], user=user) + if row["status"] in TERMINAL_STATUSES: + return _public_run(dict(row)) + cur.execute( + """ + UPDATE computation_runs + SET cancel_requested = true, status = 'cancelling', updated_at = CURRENT_TIMESTAMP + WHERE id = %s RETURNING * + """, + (run_id,), + ) + updated = dict(cur.fetchone()) + _insert_event(cur, run_id, "cancelling", "Cancellation requested") + conn.commit() + return _public_run(updated) + finally: + conn.close() + + +def claim_next_run(worker_id: str, lease_seconds: int = 120) -> dict[str, Any] | None: + conn = get_pg_superuser_connection() + try: + with conn.cursor(cursor_factory=RealDictCursor) as cur: + cur.execute( + """ + WITH candidate AS ( + SELECT id FROM computation_runs + WHERE status IN ( + 'requested', 'validating', 'staging', 'submitted', 'running', + 'collecting', 'cancelling' + ) + AND (worker_lease_expires_at IS NULL OR worker_lease_expires_at < CURRENT_TIMESTAMP) + ORDER BY created_at + FOR UPDATE SKIP LOCKED + LIMIT 1 + ) + UPDATE computation_runs r + SET worker_lease_owner = %s, + worker_lease_expires_at = CURRENT_TIMESTAMP + (%s * INTERVAL '1 second'), + updated_at = CURRENT_TIMESTAMP + FROM candidate + WHERE r.id = candidate.id + RETURNING r.* + """, + (worker_id, lease_seconds), + ) + row = cur.fetchone() + conn.commit() + return dict(row) if row else None + finally: + conn.close() + + +def find_cache_hit(run: dict[str, Any]) -> dict[str, Any] | None: + conn = get_pg_superuser_connection() + try: + with conn.cursor(cursor_factory=RealDictCursor) as cur: + cur.execute( + """ + SELECT * FROM computation_runs + WHERE tool_name = %s AND sim2l_squid_id = %s + AND project_id = %s + AND status = 'data_ingested' AND id <> %s + ORDER BY completed_at DESC NULLS LAST LIMIT 1 + """, + ( + run["tool_name"], + run["sim2l_squid_id"], + run["project_id"], + run["id"], + ), + ) + row = cur.fetchone() + return dict(row) if row else None + finally: + conn.close() + + +def update_run( + run_id: str, + *, + status: str | None = None, + message: str = "", + values: dict[str, Any] | None = None, + event_details: dict[str, Any] | None = None, + release_lease: bool = True, +) -> dict[str, Any]: + if status is not None and status not in ALL_STATUSES: + raise ValueError(f"invalid computation status: {status}") + allowed = { + "scheduler_job_id", + "remote_workdir", + "input_json", + "output_json", + "artifact_manifest_json", + "external_workflow_ref_json", + "error_json", + "prepared_manifest_json", + "cancel_requested", + "attempt_count", + "submitted_at", + "started_at", + "completed_at", + } + values = values or {} + unknown = set(values) - allowed + if unknown: + raise ValueError(f"unsupported computation fields: {sorted(unknown)}") + assignments = ["updated_at = CURRENT_TIMESTAMP"] + if release_lease: + assignments.extend( + ["worker_lease_expires_at = NULL", "worker_lease_owner = ''"] + ) + params: list[Any] = [] + if status is not None: + assignments.append("status = %s") + params.append(status) + for key, value in values.items(): + assignments.append(f"{key} = %s") + if key.endswith("_json"): + value = Json(value) + params.append(value) + params.append(run_id) + conn = get_pg_superuser_connection() + try: + with conn.cursor(cursor_factory=RealDictCursor) as cur: + cur.execute( + f"UPDATE computation_runs SET {', '.join(assignments)} WHERE id = %s RETURNING *", + params, + ) + row = cur.fetchone() + if not row: + raise LookupError("computation run not found") + if status: + _insert_event(cur, run_id, status, message, event_details) + conn.commit() + return dict(row) + finally: + conn.close() + + +def _insert_event( + cur, + run_id: str, + status: str, + message: str, + details: dict[str, Any] | None = None, +) -> None: + cur.execute( + """ + INSERT INTO computation_run_events (run_id, status, message, details_json) + VALUES (%s, %s, %s, %s) + """, + (run_id, status, message, Json(details or {})), + ) + + +def _ensure_tool_catalog( + cur, *, user: PlatformUser, project_id: str, tool_name: str +) -> None: + """Register computational equipment/type when the full DT schema exists. + + The computation tables can also be tested against a minimal database, so + this integration is deliberately conditional on the platform catalog + tables being present. + """ + + cur.execute( + "SELECT to_regclass('public.equipment_metadata') AS equipment_table, " + "to_regclass('public.experiment_types') AS type_table, " + "to_regclass('public.project_experiments') AS project_table" + ) + tables = cur.fetchone() + equipment_table = tables["equipment_table"] + type_table = tables["type_table"] + project_table = tables["project_table"] + if not equipment_table or not type_table: + return + definition = TOOL_DEFINITIONS[tool_name] + equipment_id = f"sim2l_{tool_name}" + type_id = f"sim2l_{tool_name}_{definition.revision.replace('-', '_')}" + schema = definition.input_model.model_json_schema() + required = set(schema.get("required", [])) + parameters = [] + for name, field in schema.get("properties", {}).items(): + parameters.append( + { + "name": name, + "display_name": field.get("title", name.replace("_", " ").title()), + "type": field.get("type", "object"), + "required": name in required, + "default": field.get("default"), + "description": field.get("description", ""), + "source": "sim2l", + } + ) + outputs = [ + { + "name": name, + "display_name": name.replace("_", " ").title(), + "type": "object", + "source": "sim2l", + } + for name in definition.output_names + ] + config = { + "execution_backend": "sim2l+rcac", + "scheduler": "slurm", + "sim2l_name": definition.name, + "sim2l_revision": definition.revision, + "feature_flag": "RCAC_COMPUTE_ENABLED", + } + cur.execute( + """ + INSERT INTO equipment_metadata ( + domain_id, equipment_name, equipment_type, manufacturer, model, + location, description, source_type, dt_mode, outlier_method, + owner_id, owner_org, status, parameters_json, features_json, + config_json, approved_at, reviewed_by, updated_at + ) VALUES ( + %s, %s, 'Computational Software', 'Quantum ESPRESSO', %s, + 'Purdue RCAC', %s, 'rcac_slurm', 'dashboard', 'none', + %s, %s, 'approved', %s, '[]'::jsonb, %s, + CURRENT_TIMESTAMP, %s, CURRENT_TIMESTAMP + ) + ON CONFLICT (domain_id) DO UPDATE SET + equipment_name = EXCLUDED.equipment_name, + model = EXCLUDED.model, + description = EXCLUDED.description, + parameters_json = EXCLUDED.parameters_json, + config_json = EXCLUDED.config_json, + updated_at = CURRENT_TIMESTAMP + """, + ( + equipment_id, + definition.title, + definition.revision, + definition.description, + user.id, + user.organization, + Json(parameters), + Json(config), + user.id, + ), + ) + cur.execute( + """ + INSERT INTO experiment_types ( + id, name, description, scientific_objective, equipment_id, + type_parameters_json, outputs_json, additional_inputs_json, + visibility, version, author_name, author_email, owner_id, owner_org + ) VALUES ( + %s, %s, %s, %s, %s, %s, %s, '[]'::jsonb, + 'shared', 1, %s, %s, %s, %s + ) + ON CONFLICT (id) DO UPDATE SET + name = EXCLUDED.name, + description = EXCLUDED.description, + equipment_id = EXCLUDED.equipment_id, + type_parameters_json = EXCLUDED.type_parameters_json, + outputs_json = EXCLUDED.outputs_json, + updated_at = CURRENT_TIMESTAMP + """, + ( + type_id, + definition.title, + definition.description, + definition.description, + equipment_id, + Json(parameters), + Json(outputs), + user.name, + user.email, + user.id, + user.organization, + ), + ) + if not project_table: + return + upstream_type = "sim2l_cellrelaxdft_r99_dt1" if tool_name == "elasticdft" else "" + cur.execute( + """ + INSERT INTO project_experiments ( + project_id, type_id, added_by, status, sequence_index, + upstream_type_id, handoff_type, handoff_description, + output_mapping_json + ) + SELECT %s, %s, %s, 'active', + COALESCE(MAX(sequence_index), 0) + 1, %s, + 'automated_computational_job', + 'Sim2L execution on Purdue RCAC through Slurm.', %s + FROM project_experiments WHERE project_id = %s + ON CONFLICT (project_id, type_id) DO UPDATE SET + status = 'active', + upstream_type_id = EXCLUDED.upstream_type_id, + handoff_type = EXCLUDED.handoff_type, + handoff_description = EXCLUDED.handoff_description, + output_mapping_json = EXCLUDED.output_mapping_json, + updated_at = CURRENT_TIMESTAMP + """, + ( + project_id, + type_id, + user.id, + upstream_type, + Json( + { + "execution_backend": "sim2l+rcac", + "upstream_output": "structure_dict" if upstream_type else "", + "downstream_input": "upstream_execution_id" + if upstream_type + else "", + } + ), + project_id, + ), + ) diff --git a/api/computations/worker.py b/api/computations/worker.py new file mode 100644 index 0000000..858c7d7 --- /dev/null +++ b/api/computations/worker.py @@ -0,0 +1,399 @@ +"""Durable PostgreSQL-backed RCAC worker. + +Run separately from FastAPI:: + + python -m computations.worker + +The worker is safe to restart: PostgreSQL stores every scheduler identifier and +short leases prevent two replicas from handling the same run concurrently. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import socket +import time +import uuid +from collections.abc import Iterable +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from .artifacts import artifact_store_from_env +from .materials import resolve_material_inputs +from .qe import ( + PreparedQeRun, + cellrelax_outputs, + elastic_outputs, + prepare_cellrelax, + prepare_elastic, +) +from .rcac import make_slurm_executor, remote_workdir +from .results import register_sim2l_result +from .schemas import validate_tool_inputs +from .store import ( + claim_next_run, + feature_enabled, + find_cache_hit, + get_profile, + get_run, + update_run, +) + +logger = logging.getLogger("dt.computations.worker") + + +def _json(value: Any, default: Any) -> Any: + if isinstance(value, str): + try: + return json.loads(value) + except json.JSONDecodeError: + return default + return value if value is not None else default + + +def _prepared_dict(prepared: PreparedQeRun) -> dict[str, Any]: + return { + "tool_name": prepared.tool_name, + "stage_files": [path.name for path in prepared.stage_files], + "command": prepared.command, + "expected_artifacts": prepared.expected_artifacts, + "job_array": prepared.job_array, + "ntasks": prepared.ntasks, + "time_limit": prepared.time_limit, + "metadata": prepared.metadata or {}, + } + + +def _restore_prepared(run: dict[str, Any], local_workdir: Path) -> PreparedQeRun: + manifest = _json(run.get("prepared_manifest_json"), {}) + if not manifest: + raise RuntimeError("prepared run manifest is missing") + return PreparedQeRun( + tool_name=manifest["tool_name"], + local_workdir=local_workdir, + stage_files=[local_workdir / name for name in manifest.get("stage_files", [])], + command=list(manifest["command"]), + expected_artifacts=list(manifest["expected_artifacts"]), + job_array=manifest.get("job_array", ""), + ntasks=int(manifest["ntasks"]), + time_limit=manifest["time_limit"], + metadata=manifest.get("metadata", {}), + ) + + +class ComputationWorker: + def __init__(self, worker_id: str | None = None): + self.worker_id = worker_id or f"{socket.gethostname()}-{uuid.uuid4().hex[:8]}" + self.work_root = Path( + os.getenv("RCAC_LOCAL_WORK_ROOT", "/tmp/dt-rcac-worker") + ).resolve() + self.work_root.mkdir(parents=True, exist_ok=True) + pseudo_value = os.getenv("RCAC_PSEUDO_ROOT", "").strip() + self.pseudo_root = Path(pseudo_value).resolve() if pseudo_value else None + self.artifact_store = artifact_store_from_env() + + def run_once(self) -> bool: + if not feature_enabled(): + return False + run = claim_next_run(self.worker_id) + if not run: + return False + run_id = str(run["id"]) + try: + if run["status"] in {"requested", "validating", "staging"}: + self._submit(run) + else: + self._advance(run) + except Exception as exc: + logger.exception("Computation run %s failed", run_id) + update_run( + run_id, + status="failed", + message="Execution failed", + values={ + "error_json": { + "type": type(exc).__name__, + "message": str(exc), + }, + "completed_at": datetime.now(timezone.utc), + }, + ) + return True + + def _submit(self, run: dict[str, Any]) -> None: + run_id = str(run["id"]) + if run["status"] != "validating": + update_run( + run_id, + status="validating", + message="Validating Sim2L inputs", + release_lease=False, + ) + cached = find_cache_hit(run) + if cached: + update_run( + run_id, + status="data_ingested", + message="Reused a matching Sim2L result", + values={ + "output_json": _json(cached.get("output_json"), {}), + "artifact_manifest_json": _json( + cached.get("artifact_manifest_json"), [] + ), + "external_workflow_ref_json": { + "cache_hit": True, + "source_run_id": str(cached["id"]), + }, + "completed_at": datetime.now(timezone.utc), + }, + ) + return + + inputs = validate_tool_inputs(run["tool_name"], _json(run["input_json"], {})) + if run["tool_name"] == "cellrelaxdft": + inputs = resolve_material_inputs(inputs) + update_run( + run_id, + status="staging", + message="Preparing and staging Quantum ESPRESSO inputs", + values={"input_json": inputs}, + release_lease=False, + ) + if self.pseudo_root is None: + raise RuntimeError("RCAC_PSEUDO_ROOT is not configured") + + local_workdir = self.work_root / run_id + local_workdir.mkdir(parents=True, exist_ok=True) + if run["tool_name"] == "cellrelaxdft": + prepared = prepare_cellrelax(inputs, local_workdir, self.pseudo_root) + elif run["tool_name"] == "elasticdft": + upstream_id = str(run.get("upstream_run_id") or "") + upstream = get_run(upstream_id, internal=True) + if not upstream or upstream["status"] != "data_ingested": + raise RuntimeError("completed CellRelaxDFT upstream run is unavailable") + prepared = prepare_elastic( + inputs, + _json(upstream["input_json"], {}), + _json(upstream["output_json"], {}), + local_workdir, + self.pseudo_root, + ) + else: + raise ValueError(f"unsupported computation tool: {run['tool_name']}") + + profile = get_profile(run["requested_by"], run["cluster"], internal=True) + if not profile: + raise RuntimeError("an enabled RCAC profile is required") + executor, cluster = make_slurm_executor(profile, local_workdir) + remote = remote_workdir(cluster, profile["ssh_username"], run_id) + executor.stage(prepared.stage_files, remote) + + try: + from sim2l.executor import SlurmJobSpec + except ImportError as exc: + raise RuntimeError("the matching Sim2L package is not installed") from exc + output_file = "slurm-%A_%a.out" if prepared.job_array else "slurm.out" + error_file = "slurm-%A_%a.err" if prepared.job_array else "slurm.err" + spec = SlurmJobSpec( + job_name=f"{prepared.tool_name[:12]}-{run_id[:8]}", + command=prepared.command, + workdir=remote, + account=run["allocation"], + partition=run["partition"], + time_limit=prepared.time_limit, + ntasks=prepared.ntasks, + array=prepared.job_array, + modules=[cluster.qe_module], + environment={"QE_DISABLE_GGA_PBE": "0"}, + output_file=output_file, + error_file=error_file, + ) + job_id = executor.submit(spec) + update_run( + run_id, + status="submitted", + message=f"Submitted Slurm job {job_id}", + values={ + "scheduler_job_id": job_id, + "remote_workdir": remote, + "prepared_manifest_json": _prepared_dict(prepared), + "external_workflow_ref_json": { + "manager": "sim2l-slurm", + "scheduler": "slurm", + "cluster": cluster.name, + "scheduler_job_id": job_id, + "cache_hit": False, + }, + "attempt_count": int(run.get("attempt_count") or 0) + 1, + "submitted_at": datetime.now(timezone.utc), + }, + ) + + def _advance(self, run: dict[str, Any]) -> None: + try: + from sim2l.executor import SlurmState + except ImportError as exc: + raise RuntimeError("the matching Sim2L package is not installed") from exc + run_id = str(run["id"]) + if run["status"] == "collecting": + self._collect(run) + return + profile = get_profile(run["requested_by"], run["cluster"], internal=True) + if not profile: + raise RuntimeError( + "the RCAC profile used by this run is no longer available" + ) + local_workdir = self.work_root / run_id + executor, _cluster = make_slurm_executor(profile, local_workdir) + job_id = run["scheduler_job_id"] + if run["status"] == "cancelling" or run.get("cancel_requested"): + if job_id: + current = executor.status(job_id) + if current == SlurmState.COMPLETED: + update_run( + run_id, + status="collecting", + message="Job completed before cancellation; collecting results", + values={"cancel_requested": False}, + release_lease=False, + ) + refreshed = get_run(run_id, internal=True) + self._collect(refreshed) + return + if current not in {SlurmState.CANCELLED, SlurmState.FAILED}: + executor.cancel(job_id) + update_run( + run_id, + status="cancelled", + message="Slurm execution cancelled", + values={"completed_at": datetime.now(timezone.utc)}, + ) + return + + state = executor.status(job_id) + if state == SlurmState.PENDING: + update_run(run_id, values={}) + return + if state == SlurmState.RUNNING: + if run["status"] != "running": + update_run( + run_id, + status="running", + message=f"Slurm job {job_id} is running", + values={"started_at": datetime.now(timezone.utc)}, + ) + else: + update_run(run_id, values={}) + return + if state == SlurmState.CANCELLED: + update_run( + run_id, + status="cancelled", + message=f"Slurm job {job_id} was cancelled", + values={"completed_at": datetime.now(timezone.utc)}, + ) + return + if state == SlurmState.FAILED: + raise RuntimeError(f"Slurm job {job_id} failed") + if state == SlurmState.UNKNOWN: + update_run(run_id, values={}) + return + + update_run( + run_id, + status="collecting", + message="Collecting RCAC artifacts", + release_lease=False, + ) + refreshed = get_run(run_id, internal=True) + self._collect(refreshed) + + def _collect(self, run: dict[str, Any]) -> None: + run_id = str(run["id"]) + local_workdir = self.work_root / run_id + prepared = _restore_prepared(run, local_workdir) + profile = get_profile(run["requested_by"], run["cluster"], internal=True) + executor, _cluster = make_slurm_executor(profile, local_workdir) + artifacts = list(prepared.expected_artifacts) + if prepared.job_array: + count = int((prepared.metadata or {})["count"]) + artifacts.extend( + f"slurm-{run['scheduler_job_id']}_{index}.out" for index in range(count) + ) + artifacts.extend( + f"slurm-{run['scheduler_job_id']}_{index}.err" for index in range(count) + ) + else: + artifacts.extend(["slurm.out", "slurm.err"]) + executor.collect(run["remote_workdir"], artifacts) + + inputs = _json(run["input_json"], {}) + if run["tool_name"] == "cellrelaxdft": + outputs = cellrelax_outputs(inputs, local_workdir / "vcrelax.stdout") + else: + outputs = elastic_outputs(prepared) + + persistent = self._persistent_artifacts(prepared, artifacts) + manifest = self.artifact_store.put_many(run_id, local_workdir, persistent) + sim2l_registration: dict[str, Any] + try: + sim2l_registration = register_sim2l_result(run, outputs) + # Result registration is an optional secondary sink. Its client can + # surface transport, protocol, and dependency-specific exceptions; + # none should discard a successfully collected primary DT result. + except Exception as exc: # noqa: BLE001 + logger.warning("Sim2L result registration failed for %s: %s", run_id, exc) + sim2l_registration = {"status": "failed", "error": str(exc)} + external = _json(run.get("external_workflow_ref_json"), {}) + external["sim2l_result"] = sim2l_registration + update_run( + run_id, + status="data_ingested", + message="Quantum ESPRESSO outputs ingested", + values={ + "output_json": outputs, + "artifact_manifest_json": manifest, + "external_workflow_ref_json": external, + "completed_at": datetime.now(timezone.utc), + }, + ) + + @staticmethod + def _persistent_artifacts( + prepared: PreparedQeRun, collected: Iterable[str] + ) -> list[Path]: + root = prepared.local_workdir + paths = {root / relative for relative in collected} + if prepared.tool_name == "cellrelaxdft": + paths.add(root / "vcrelax.in") + else: + paths.add(root / "elastic_inputs" / "manifest.json") + paths.update((root / "elastic_inputs").glob("strain_*/relax.in")) + return sorted(path for path in paths if path.is_file()) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Birck DT Sim2L/RCAC worker") + parser.add_argument("--once", action="store_true", help="process at most one run") + parser.add_argument("--poll-seconds", type=float, default=10.0) + args = parser.parse_args() + logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO")) + if not feature_enabled(): + raise SystemExit( + "RCAC computation worker is disabled; set RCAC_COMPUTE_ENABLED=true" + ) + worker = ComputationWorker() + if args.once: + worker.run_once() + return + while True: + if not worker.run_once(): + time.sleep(max(0.25, args.poll_seconds)) + + +if __name__ == "__main__": + main() diff --git a/api/main.py b/api/main.py index 5668b9c..32409c2 100644 --- a/api/main.py +++ b/api/main.py @@ -26,7 +26,7 @@ from middleware.logging import RequestTracingMiddleware from middleware.error_handler import global_exception_handler from runtime_config import MissingRuntimeConfiguration, validate_runtime_configuration -from routers import ml, dataset, dataset_v2, equipment, upload, admin, processing +from routers import ml, dataset, dataset_v2, equipment, upload, admin, processing, computations # ── Structured logging ────────────────────────────────────────────────────── logging.basicConfig( @@ -165,6 +165,8 @@ def _run_startup_migrations() -> None: ensure_data_uploads_processing_started_at_pg, recover_stale_db3_uploads_pg, ) + from computations.store import ensure_tables as ensure_computation_tables + from computations.store import feature_enabled as computations_enabled lock_conn = get_pg_superuser_connection() try: @@ -193,6 +195,8 @@ def _run_startup_migrations() -> None: ensure_data_uploads_kind_column_pg() ensure_data_uploads_processing_started_at_pg() recover_stale_db3_uploads_pg() + if computations_enabled(): + ensure_computation_tables() finally: try: with lock_conn.cursor() as cur: @@ -257,6 +261,7 @@ async def lifespan(app: FastAPI): app.include_router(upload.router, prefix="/api") app.include_router(admin.router, prefix="/api") app.include_router(processing.router, prefix="/api") +app.include_router(computations.router, prefix="/api") # ── Health check ───────────────────────────────────────────────────────────── diff --git a/api/requirements-computations.txt b/api/requirements-computations.txt new file mode 100644 index 0000000..1f98225 --- /dev/null +++ b/api/requirements-computations.txt @@ -0,0 +1,6 @@ +-r requirements.txt + +# Scientific dependencies are isolated in the worker image so the existing +# FastAPI image remains unchanged. +pymatgen>=2025.6,<2027 +mp-api>=0.45,<0.47 diff --git a/api/routers/computations.py b/api/routers/computations.py new file mode 100644 index 0000000..31bf749 --- /dev/null +++ b/api/routers/computations.py @@ -0,0 +1,214 @@ +"""Feature-flagged API for CellRelaxDFT and ElasticDFT RCAC runs.""" + +from __future__ import annotations + +from typing import Annotated, Any + +from computations.artifacts import ( + AzureBlobArtifactStore, + LocalArtifactStore, + artifact_store_from_env, +) +from computations.rcac import configured_clusters, get_cluster +from computations.schemas import ( + TOOL_DEFINITIONS, + compute_squid_id, + validate_tool_inputs, +) +from computations.store import ( + create_run, + feature_enabled, + get_profile, + get_run, + list_runs, + request_cancel, + upsert_profile, +) +from fastapi import APIRouter, Depends, HTTPException +from fastapi.responses import FileResponse, StreamingResponse +from pydantic import BaseModel, ConfigDict, Field +from security import PlatformUser, get_platform_user + +router = APIRouter(prefix="/computations", tags=["computations"]) +CurrentUser = Annotated[PlatformUser, Depends(get_platform_user)] + + +def _require_enabled() -> None: + if not feature_enabled(): + raise HTTPException(status_code=404, detail="RCAC computations are not enabled") + + +class RunPayload(BaseModel): + model_config = ConfigDict(extra="forbid") + + cluster: str = Field(min_length=1, max_length=100, pattern=r"^[A-Za-z0-9_-]+$") + inputs: dict[str, Any] + + +class ProfilePayload(BaseModel): + model_config = ConfigDict(extra="forbid") + + ssh_username: str = Field(pattern=r"^[A-Za-z0-9._-]+$", max_length=255) + allocation: str = Field(pattern=r"^[A-Za-z0-9_.-]+$", max_length=255) + partition: str = Field( + default="standby", pattern=r"^[A-Za-z0-9_.-]+$", max_length=255 + ) + credential_ref: str = Field(pattern=r"^[A-Za-z0-9_.-]+$", max_length=255) + + +@router.get("/tools") +def tools(user: CurrentUser): + _require_enabled() + del user + return { + "tools": [definition.public_dict() for definition in TOOL_DEFINITIONS.values()], + "clusters": [ + {"name": cluster.name, "host": cluster.host, "qe_module": cluster.qe_module} + for cluster in configured_clusters().values() + ], + } + + +@router.get("/profiles/{cluster}") +def profile(cluster: str, user: CurrentUser): + _require_enabled() + get_cluster(cluster) + value = get_profile(user.id, cluster) + if not value: + raise HTTPException(status_code=404, detail="RCAC profile not configured") + return value + + +@router.put("/profiles/{cluster}") +def save_profile( + cluster: str, + payload: ProfilePayload, + user: CurrentUser, +): + _require_enabled() + try: + get_cluster(cluster) + return upsert_profile(user=user, cluster=cluster, payload=payload.model_dump()) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + +@router.post("/projects/{project_id}/runs/{tool_name}", status_code=202) +def submit_run( + project_id: str, + tool_name: str, + payload: RunPayload, + user: CurrentUser, +): + _require_enabled() + try: + definition = TOOL_DEFINITIONS[tool_name] + except KeyError as exc: + raise HTTPException( + status_code=404, detail="Computation tool not found" + ) from exc + try: + get_cluster(payload.cluster) + profile_value = get_profile(user.id, payload.cluster, internal=True) + if not profile_value: + raise ValueError("Configure an enabled RCAC profile before submitting") + inputs = validate_tool_inputs(tool_name, payload.inputs) + squid_id = compute_squid_id(tool_name, inputs) + upstream = inputs.get("upstream_execution_id", "") + return create_run( + user=user, + project_id=project_id, + tool_name=tool_name, + tool_revision=definition.revision, + squid_id=squid_id, + inputs=inputs, + cluster=payload.cluster, + partition=profile_value["partition"], + allocation=profile_value["allocation"], + upstream_run_id=str(upstream), + ) + except PermissionError as exc: + raise HTTPException(status_code=403, detail=str(exc)) from exc + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + +@router.get("/projects/{project_id}/runs") +def project_runs(project_id: str, user: CurrentUser): + _require_enabled() + try: + return list_runs(project_id, user) + except PermissionError as exc: + raise HTTPException(status_code=403, detail=str(exc)) from exc + + +@router.get("/runs/{run_id}") +def run_detail(run_id: str, user: CurrentUser): + _require_enabled() + try: + value = get_run(run_id, user) + except PermissionError as exc: + raise HTTPException(status_code=403, detail=str(exc)) from exc + if not value: + raise HTTPException(status_code=404, detail="Computation run not found") + return value + + +@router.post("/runs/{run_id}/cancel", status_code=202) +def cancel_run(run_id: str, user: CurrentUser): + _require_enabled() + try: + value = request_cancel(run_id, user) + except PermissionError as exc: + raise HTTPException(status_code=403, detail=str(exc)) from exc + if not value: + raise HTTPException(status_code=404, detail="Computation run not found") + return value + + +@router.get("/runs/{run_id}/artifacts/{artifact_index}") +def download_artifact( + run_id: str, + artifact_index: int, + user: CurrentUser, +): + _require_enabled() + try: + run = get_run(run_id, user, internal=True) + except PermissionError as exc: + raise HTTPException(status_code=403, detail=str(exc)) from exc + if not run: + raise HTTPException(status_code=404, detail="Computation run not found") + manifest = run.get("artifact_manifest_json") or [] + if isinstance(manifest, str): + import json + + manifest = json.loads(manifest) + if artifact_index < 0 or artifact_index >= len(manifest): + raise HTTPException(status_code=404, detail="Artifact not found") + artifact = manifest[artifact_index] + artifact_name = str(artifact.get("name") or "artifact") + store = artifact_store_from_env() + try: + if isinstance(store, LocalArtifactStore): + return FileResponse( + store.resolve(artifact["uri"]), + media_type=artifact.get("content_type") or "application/octet-stream", + filename=artifact_name.rsplit("/", 1)[-1], + ) + if isinstance(store, AzureBlobArtifactStore): + downloader = store.downloader(artifact["uri"]) + return StreamingResponse( + downloader.chunks(), + media_type=artifact.get("content_type") or "application/octet-stream", + headers={ + "Content-Disposition": ( + 'attachment; filename="' + + artifact_name.rsplit("/", 1)[-1].replace('"', "") + + '"' + ) + }, + ) + except (FileNotFoundError, ValueError, PermissionError) as exc: + raise HTTPException(status_code=404, detail="Artifact unavailable") from exc + raise HTTPException(status_code=500, detail="Artifact store is unsupported") diff --git a/api/scripts/add_rcac_computations.sql b/api/scripts/add_rcac_computations.sql new file mode 100644 index 0000000..1779541 --- /dev/null +++ b/api/scripts/add_rcac_computations.sql @@ -0,0 +1,78 @@ +-- Additive schema for feature-flagged Sim2L/RCAC computations. +-- Contains no changes to existing experiment or execution tables. + +CREATE TABLE IF NOT EXISTS rcac_user_profiles ( + user_id VARCHAR(255) NOT NULL REFERENCES users(id) ON DELETE CASCADE, + cluster VARCHAR(100) NOT NULL, + ssh_username VARCHAR(255) NOT NULL, + allocation VARCHAR(255) NOT NULL, + partition VARCHAR(255) NOT NULL DEFAULT 'standby', + credential_ref VARCHAR(255) NOT NULL, + enabled BOOLEAN NOT NULL DEFAULT true, + created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (user_id, cluster) +); + +CREATE TABLE IF NOT EXISTS computation_runs ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + project_id VARCHAR(50) NOT NULL REFERENCES projects(id) ON DELETE CASCADE, + requested_by VARCHAR(255) NOT NULL REFERENCES users(id), + tool_name VARCHAR(100) NOT NULL, + tool_revision VARCHAR(100) NOT NULL, + sim2l_squid_id VARCHAR(255) NOT NULL, + upstream_run_id UUID REFERENCES computation_runs(id) ON DELETE RESTRICT, + status VARCHAR(50) NOT NULL DEFAULT 'requested' + CHECK (status IN ( + 'requested', 'validating', 'staging', 'submitted', 'running', + 'collecting', 'data_ingested', 'cancelling', 'cancelled', 'failed' + )), + cluster VARCHAR(100) NOT NULL, + partition VARCHAR(255) NOT NULL, + allocation VARCHAR(255) NOT NULL, + scheduler_job_id VARCHAR(100) NOT NULL DEFAULT '', + remote_workdir TEXT NOT NULL DEFAULT '', + input_json JSONB NOT NULL DEFAULT '{}'::jsonb, + output_json JSONB NOT NULL DEFAULT '{}'::jsonb, + artifact_manifest_json JSONB NOT NULL DEFAULT '[]'::jsonb, + external_workflow_ref_json JSONB NOT NULL DEFAULT '{}'::jsonb, + error_json JSONB NOT NULL DEFAULT '{}'::jsonb, + prepared_manifest_json JSONB NOT NULL DEFAULT '{}'::jsonb, + cancel_requested BOOLEAN NOT NULL DEFAULT false, + attempt_count INTEGER NOT NULL DEFAULT 0, + worker_lease_owner VARCHAR(255) NOT NULL DEFAULT '', + worker_lease_expires_at TIMESTAMP WITH TIME ZONE, + submitted_at TIMESTAMP WITH TIME ZONE, + started_at TIMESTAMP WITH TIME ZONE, + completed_at TIMESTAMP WITH TIME ZONE, + created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE INDEX IF NOT EXISTS idx_computation_runs_project_created +ON computation_runs(project_id, created_at DESC); + +CREATE INDEX IF NOT EXISTS idx_computation_runs_worker_queue +ON computation_runs(status, worker_lease_expires_at, created_at) +WHERE status NOT IN ('data_ingested', 'cancelled', 'failed'); + +CREATE INDEX IF NOT EXISTS idx_computation_runs_squid +ON computation_runs(tool_name, sim2l_squid_id, status); + +CREATE TABLE IF NOT EXISTS computation_run_events ( + id BIGSERIAL PRIMARY KEY, + run_id UUID NOT NULL REFERENCES computation_runs(id) ON DELETE CASCADE, + status VARCHAR(50) NOT NULL, + message TEXT NOT NULL DEFAULT '', + details_json JSONB NOT NULL DEFAULT '{}'::jsonb, + created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE INDEX IF NOT EXISTS idx_computation_run_events_run +ON computation_run_events(run_id, created_at); + +-- Do not grant mutation access to PostgREST's api_client role. FastAPI owns +-- authorization for this feature and uses its existing privileged connection. +ALTER TABLE computation_runs ENABLE ROW LEVEL SECURITY; +ALTER TABLE computation_run_events ENABLE ROW LEVEL SECURITY; +ALTER TABLE rcac_user_profiles ENABLE ROW LEVEL SECURITY; diff --git a/api/tests/test_computation_schemas.py b/api/tests/test_computation_schemas.py new file mode 100644 index 0000000..ef50668 --- /dev/null +++ b/api/tests/test_computation_schemas.py @@ -0,0 +1,59 @@ +import unittest +from uuid import uuid4 + +from computations.schemas import compute_squid_id, validate_tool_inputs +from pydantic import ValidationError + + +class ComputationSchemaTests(unittest.TestCase): + def test_cellrelax_defaults_and_squid_are_stable(self): + inputs = validate_tool_inputs( + "cellrelaxdft", {"mp_id": "mp-1265", "structure_dict": {}} + ) + self.assertEqual(inputs["pseudo"], "PAW") + self.assertEqual(inputs["exchange_correlation"], "PBEsol") + first = compute_squid_id("cellrelaxdft", inputs) + second = compute_squid_id("cellrelaxdft", dict(reversed(list(inputs.items())))) + self.assertEqual(first, second) + self.assertRegex(first, r"^cellrelaxdft/r99-dt1/[0-9a-f]{40}$") + + def test_missing_nc_pseudopotentials_are_not_offered(self): + with self.assertRaises(ValidationError): + validate_tool_inputs( + "cellrelaxdft", + {"mp_id": "mp-1265", "structure_dict": {}, "pseudo": "NC"}, + ) + + def test_electronic_steps_match_r99_bounds(self): + with self.assertRaises(ValidationError): + validate_tool_inputs( + "cellrelaxdft", + {"mp_id": "mp-1265", "structure_dict": {}, "e_steps": 99}, + ) + + def test_hubbard_cross_field_validation(self): + with self.assertRaisesRegex(ValueError, "Hubbard U values"): + validate_tool_inputs( + "cellrelaxdft", + { + "mp_id": "mp-1265", + "structure_dict": {}, + "hubbard_U_on": True, + }, + ) + + def test_elastic_strains_require_compression_and_tension(self): + with self.assertRaises(ValidationError): + validate_tool_inputs( + "elasticdft", + { + "mp_id": "mp-1265", + "upstream_execution_id": str(uuid4()), + "diagonal_strain_magnitudes": [0.005, 0.01], + "offdiagonal_strain_magnitudes": [-0.01, 0.01], + }, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/api/tests/test_computation_worker.py b/api/tests/test_computation_worker.py new file mode 100644 index 0000000..97acb4e --- /dev/null +++ b/api/tests/test_computation_worker.py @@ -0,0 +1,166 @@ +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch +from uuid import uuid4 + +from computations.qe import PreparedQeRun +from computations.worker import ComputationWorker + + +class FakeArtifacts: + def put_many(self, run_id, base_dir, paths): + return [ + {"name": Path(path).name, "uri": f"fake://{run_id}/{Path(path).name}"} + for path in paths + ] + + +class FakeExecutor: + def __init__(self): + self.staged = [] + + def stage(self, files, remote): + self.staged.extend(files) + + def submit(self, spec): + return "4242" + + def status(self, job_id): + from sim2l.executor import SlurmState + + return SlurmState.COMPLETED + + def collect(self, remote, artifacts): + for relative in artifacts: + path = self.local_workdir / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("fake", encoding="utf-8") + + +class WorkerLifecycleTests(unittest.TestCase): + def test_disabled_worker_does_not_claim_a_run(self): + with ( + patch("computations.worker.feature_enabled", return_value=False), + patch("computations.worker.claim_next_run") as claim, + ): + self.assertFalse(ComputationWorker("disabled-worker").run_once()) + claim.assert_not_called() + + def test_submit_then_collect(self): + run_id = str(uuid4()) + run = { + "id": run_id, + "project_id": "P-1", + "requested_by": "user-1", + "tool_name": "cellrelaxdft", + "tool_revision": "r99-dt1", + "sim2l_squid_id": "cellrelaxdft/r99-dt1/hash", + "status": "requested", + "cluster": "fake", + "partition": "standby", + "allocation": "allocation", + "scheduler_job_id": "", + "remote_workdir": "", + "input_json": {"mp_id": "mp-1265", "structure_dict": {}}, + "attempt_count": 0, + } + updates = [] + + def update(run_id_value, status=None, values=None, **kwargs): + if status: + run["status"] = status + run.update(values or {}) + updates.append(status) + return run + + def prepared(inputs, workdir, pseudo_root): + input_path = workdir / "vcrelax.in" + input_path.write_text("input", encoding="utf-8") + return PreparedQeRun( + tool_name="cellrelaxdft", + local_workdir=workdir, + stage_files=[input_path], + command=["true"], + expected_artifacts=["vcrelax.stdout"], + ntasks=1, + metadata={"num_atoms": 1}, + ) + + fake_executor = FakeExecutor() + with tempfile.TemporaryDirectory() as directory: + fake_executor.local_workdir = Path(directory) / run_id + with ( + patch("computations.worker.claim_next_run", side_effect=[run, run]), + patch("computations.worker.feature_enabled", return_value=True), + patch("computations.worker.update_run", side_effect=update), + patch("computations.worker.find_cache_hit", return_value=None), + patch( + "computations.worker.validate_tool_inputs", + side_effect=lambda _t, value: value, + ), + patch( + "computations.worker.resolve_material_inputs", + side_effect=lambda value: { + **value, + "structure_dict": {"resolved": True}, + }, + ), + patch("computations.worker.prepare_cellrelax", side_effect=prepared), + patch( + "computations.worker.get_profile", + return_value={ + "cluster": "fake", + "ssh_username": "user", + "credential_ref": "ref", + }, + ), + patch( + "computations.worker.make_slurm_executor", + return_value=( + fake_executor, + type( + "Cluster", + (), + { + "name": "fake", + "qe_module": "qe", + "remote_root": "/tmp/{username}", + }, + )(), + ), + ), + patch( + "computations.worker.remote_workdir", + return_value=f"/tmp/user/{run_id}", + ), + patch( + "computations.worker.get_run", + side_effect=lambda *_args, **_kwargs: run, + ), + patch( + "computations.worker.cellrelax_outputs", + return_value={"convergence": True, "energy": -1.0}, + ), + patch( + "computations.worker.register_sim2l_result", + return_value={"status": "registered"}, + ), + ): + worker = ComputationWorker("test-worker") + worker.work_root = Path(directory) + worker.pseudo_root = Path(directory) + worker.artifact_store = FakeArtifacts() + self.assertTrue(worker.run_once()) + self.assertEqual(run["status"], "submitted") + self.assertEqual(run["scheduler_job_id"], "4242") + self.assertTrue(worker.run_once()) + self.assertEqual(run["status"], "data_ingested") + self.assertEqual(run["output_json"]["energy"], -1.0) + self.assertIn("validating", updates) + self.assertIn("collecting", updates) + self.assertIn("data_ingested", updates) + + +if __name__ == "__main__": + unittest.main() diff --git a/api/tests/test_computations_postgres_e2e.py b/api/tests/test_computations_postgres_e2e.py new file mode 100644 index 0000000..f10afc7 --- /dev/null +++ b/api/tests/test_computations_postgres_e2e.py @@ -0,0 +1,377 @@ +"""Opt-in API -> PostgreSQL -> worker -> API computation lifecycle test. + +Set ``DT_COMPUTATIONS_E2E=1`` and point the standard PG_* variables at a +disposable PostgreSQL database. This test never contacts RCAC. +""" + +import os +import tempfile +from pathlib import Path +from unittest.mock import patch + +import psycopg2 +import pytest +from computations.artifacts import LocalArtifactStore +from computations.qe import PreparedQeRun +from computations.store import ensure_tables +from computations.worker import ComputationWorker +from fastapi import FastAPI +from fastapi.testclient import TestClient +from routers import computations as computation_router +from security import PlatformUser, get_platform_user + +pytestmark = pytest.mark.skipif( + os.getenv("DT_COMPUTATIONS_E2E") != "1", + reason="requires an explicitly configured disposable PostgreSQL database", +) + + +class FakeExecutor: + def __init__(self, local_workdir): + self.local_workdir = Path(local_workdir) + + def stage(self, files, remote): + return [str(file) for file in files] + + def submit(self, spec): + return "9001" + + def status(self, job_id): + from sim2l.executor import SlurmState + + return SlurmState.COMPLETED + + def collect(self, remote, artifacts): + for relative in artifacts: + path = self.local_workdir / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("fake QE artifact", encoding="utf-8") + + def cancel(self, job_id): + return None + + +def _connection(): + return psycopg2.connect( + host=os.getenv("PG_HOST", "127.0.0.1"), + port=os.getenv("PG_PORT", "5432"), + dbname=os.getenv("PG_DB", "digital_twin"), + user=os.getenv("PG_SUPER_USER", "postgres"), + password=os.environ["PG_SUPER_PASS"], + ) + + +@pytest.fixture(scope="module") +def configured_database(): + os.environ["RCAC_COMPUTE_ENABLED"] = "true" + os.environ["RCAC_CLUSTERS_JSON"] = ( + '{"fake":{"host":"localhost","qe_module":"quantum-espresso/7.1",' + '"remote_root":"/tmp/{username}/dt-sim2l"}}' + ) + conn = _connection() + try: + with conn.cursor() as cur: + cur.execute('CREATE EXTENSION IF NOT EXISTS "pgcrypto"') + cur.execute( + """ + CREATE TABLE IF NOT EXISTS users ( + id VARCHAR(255) PRIMARY KEY, + name VARCHAR(255) NOT NULL, + email VARCHAR(255) NOT NULL + ); + CREATE TABLE IF NOT EXISTS projects ( + id VARCHAR(50) PRIMARY KEY, + name VARCHAR(255) NOT NULL, + access_mode VARCHAR(20) DEFAULT 'private' + ); + CREATE TABLE IF NOT EXISTS project_members ( + project_id VARCHAR(50) REFERENCES projects(id) ON DELETE CASCADE, + nanohub_user_id VARCHAR(100) NOT NULL, + role VARCHAR(20) DEFAULT 'member', + PRIMARY KEY (project_id, nanohub_user_id) + ); + CREATE TABLE IF NOT EXISTS equipment_metadata ( + domain_id VARCHAR(255) PRIMARY KEY, + equipment_name VARCHAR(255) NOT NULL, + equipment_type VARCHAR(100) DEFAULT '', + manufacturer VARCHAR(255) DEFAULT '', + model VARCHAR(255) DEFAULT '', + location VARCHAR(255) DEFAULT '', + serial_number VARCHAR(255) DEFAULT '', + description TEXT DEFAULT '', + source_type VARCHAR(100) NOT NULL DEFAULT 'postgresql', + connection_host VARCHAR(255) DEFAULT '', + connection_port VARCHAR(50) DEFAULT '5432', + connection_database VARCHAR(255) DEFAULT '', + dt_mode VARCHAR(50) NOT NULL DEFAULT 'dashboard', + outlier_method VARCHAR(100) DEFAULT 'zscore', + outlier_threshold DOUBLE PRECISION DEFAULT 3.0, + owner_id VARCHAR(255) REFERENCES users(id), + owner_org VARCHAR(255) DEFAULT '', + status VARCHAR(50) NOT NULL DEFAULT 'pending', + parameters_json JSONB NOT NULL DEFAULT '[]'::jsonb, + columns_json JSONB NOT NULL DEFAULT '[]'::jsonb, + features_json JSONB NOT NULL DEFAULT '[]'::jsonb, + primary_target_json JSONB NOT NULL DEFAULT '{}'::jsonb, + secondary_target_json JSONB NOT NULL DEFAULT '{}'::jsonb, + config_json JSONB NOT NULL DEFAULT '{}'::jsonb, + registered_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP, + approved_at TIMESTAMP WITH TIME ZONE, + reviewed_by VARCHAR(255) REFERENCES users(id), + updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP + ); + CREATE TABLE IF NOT EXISTS experiment_types ( + id VARCHAR(255) PRIMARY KEY, + name VARCHAR(255) NOT NULL, + description TEXT DEFAULT '', + scientific_objective TEXT DEFAULT '', + equipment_id VARCHAR(255) DEFAULT '', + type_parameters_json JSONB NOT NULL DEFAULT '[]'::jsonb, + outputs_json JSONB NOT NULL DEFAULT '[]'::jsonb, + additional_inputs_json JSONB NOT NULL DEFAULT '[]'::jsonb, + visibility VARCHAR(50) NOT NULL DEFAULT 'shared', + version INT NOT NULL DEFAULT 1, + pid VARCHAR(50) DEFAULT '', + parent_id VARCHAR(255) DEFAULT '', + forked_from VARCHAR(255) DEFAULT '', + author_name VARCHAR(255) DEFAULT '', + author_email VARCHAR(255) DEFAULT '', + primary_target VARCHAR(255) DEFAULT '', + primary_objective VARCHAR(50) DEFAULT '', + secondary_target VARCHAR(255) DEFAULT '', + secondary_objective VARCHAR(50) DEFAULT '', + owner_id VARCHAR(255) REFERENCES users(id), + owner_org VARCHAR(255) DEFAULT '', + created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP + ); + CREATE TABLE IF NOT EXISTS project_experiments ( + project_id VARCHAR(50) NOT NULL REFERENCES projects(id) ON DELETE CASCADE, + type_id VARCHAR(255) NOT NULL REFERENCES experiment_types(id) ON DELETE CASCADE, + added_by VARCHAR(255), + status VARCHAR(50) NOT NULL DEFAULT 'active', + sequence_index INT NOT NULL DEFAULT 1, + upstream_type_id VARCHAR(255) DEFAULT '', + handoff_type VARCHAR(100) NOT NULL DEFAULT 'manual', + handoff_description TEXT DEFAULT '', + output_mapping_json JSONB NOT NULL DEFAULT '{}'::jsonb, + created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (project_id, type_id) + ); + INSERT INTO users (id, name, email) VALUES ('e2e-user', 'E2E', 'e2e@purdue.edu') + ON CONFLICT (id) DO NOTHING; + INSERT INTO projects (id, name, access_mode) VALUES ('E2E-PROJECT', 'E2E', 'private') + ON CONFLICT (id) DO NOTHING; + INSERT INTO project_members (project_id, nanohub_user_id, role) + VALUES ('E2E-PROJECT', 'e2e-user', 'pi') + ON CONFLICT (project_id, nanohub_user_id) DO NOTHING; + """ + ) + conn.commit() + finally: + conn.close() + ensure_tables() + yield + + +@pytest.fixture +def client(configured_database): + app = FastAPI() + app.include_router(computation_router.router, prefix="/api") + user = PlatformUser( + id="e2e-user", + email="e2e@purdue.edu", + name="E2E", + role="researcher", + organization="Purdue University", + ) + app.dependency_overrides[get_platform_user] = lambda: user + return TestClient(app) + + +def test_api_worker_cache_lifecycle(client): + profile = client.put( + "/api/computations/profiles/fake", + json={ + "ssh_username": "e2e-user", + "allocation": "e2e-allocation", + "partition": "standby", + "credential_ref": "e2e-key", + }, + ) + assert profile.status_code == 200, profile.text + + payload = { + "cluster": "fake", + "inputs": {"mp_id": "mp-1265", "structure_dict": {}}, + } + submitted = client.post( + "/api/computations/projects/E2E-PROJECT/runs/cellrelaxdft", json=payload + ) + assert submitted.status_code == 202, submitted.text + run_id = submitted.json()["id"] + conn = _connection() + try: + with conn.cursor() as cur: + cur.execute( + "SELECT equipment_type FROM equipment_metadata WHERE domain_id = 'sim2l_cellrelaxdft'" + ) + assert cur.fetchone() == ("Computational Software",) + cur.execute( + "SELECT handoff_type FROM project_experiments WHERE project_id = 'E2E-PROJECT'" + ) + assert cur.fetchone() == ("automated_computational_job",) + finally: + conn.close() + + def prepare(inputs, workdir, pseudo_root): + input_path = workdir / "vcrelax.in" + input_path.write_text("QE input", encoding="utf-8") + return PreparedQeRun( + tool_name="cellrelaxdft", + local_workdir=workdir, + stage_files=[input_path], + command=["true"], + expected_artifacts=["vcrelax.stdout"], + ntasks=1, + metadata={"num_atoms": 1}, + ) + + def prepare_elastic_run( + elastic_inputs, upstream_inputs, upstream_outputs, workdir, pseudo_root + ): + package = workdir / "elastic_inputs" + strain = package / "strain_000" + strain.mkdir(parents=True) + input_path = strain / "relax.in" + input_path.write_text("elastic QE input", encoding="utf-8") + manifest_path = package / "manifest.json" + manifest_path.write_text("{}", encoding="utf-8") + return PreparedQeRun( + tool_name="elasticdft", + local_workdir=workdir, + stage_files=[manifest_path], + command=["true"], + expected_artifacts=["task_000/elastic_inputs/strain_000/relax.stdout"], + job_array="0-0%1", + ntasks=1, + metadata={"count": 1}, + ) + + with tempfile.TemporaryDirectory() as directory: + worker = ComputationWorker("postgres-e2e") + worker.work_root = Path(directory) + worker.pseudo_root = Path(directory) + artifact_root = Path(directory) / "artifacts" + os.environ["RCAC_ARTIFACT_ROOT"] = str(artifact_root) + worker.artifact_store = LocalArtifactStore(artifact_root) + + def executor_factory(profile_value, local_workdir): + cluster = type( + "Cluster", + (), + { + "name": "fake", + "qe_module": "quantum-espresso/7.1", + "remote_root": "/tmp/{username}/dt-sim2l", + }, + )() + return FakeExecutor(local_workdir), cluster + + with ( + patch( + "computations.worker.resolve_material_inputs", + side_effect=lambda value: { + **value, + "structure_dict": {"resolved": True}, + }, + ), + patch("computations.worker.prepare_cellrelax", side_effect=prepare), + patch( + "computations.worker.prepare_elastic", side_effect=prepare_elastic_run + ), + patch( + "computations.worker.make_slurm_executor", side_effect=executor_factory + ), + patch( + "computations.worker.cellrelax_outputs", + return_value={"convergence": True, "energy": -10.5}, + ), + patch( + "computations.worker.elastic_outputs", + return_value={ + "elastic_tensor": [[100.0]], + "bulk_modulus_vrh": 90.0, + }, + ), + patch( + "computations.worker.register_sim2l_result", + return_value={"status": "registered"}, + ), + ): + assert worker.run_once() + assert worker.run_once() + + detail = client.get(f"/api/computations/runs/{run_id}") + assert detail.status_code == 200, detail.text + assert detail.json()["status"] == "data_ingested" + assert detail.json()["output"]["energy"] == -10.5 + assert detail.json()["scheduler_job_id"] == "9001" + artifact_url = detail.json()["artifact_manifest"][0]["uri"].replace( + "/api/dt/", "/api/" + ) + artifact = client.get(artifact_url) + assert artifact.status_code == 200 + assert artifact.content + + cached = client.post( + "/api/computations/projects/E2E-PROJECT/runs/cellrelaxdft", + json=payload, + ) + assert cached.status_code == 202, cached.text + assert worker.run_once() + cached_detail = client.get( + f"/api/computations/runs/{cached.json()['id']}" + ).json() + assert cached_detail["status"] == "data_ingested" + assert cached_detail["external_workflow_ref"]["cache_hit"] is True + + elastic = client.post( + "/api/computations/projects/E2E-PROJECT/runs/elasticdft", + json={ + "cluster": "fake", + "inputs": { + "mp_id": "mp-1265", + "upstream_execution_id": run_id, + }, + }, + ) + assert elastic.status_code == 202, elastic.text + assert worker.run_once() + assert worker.run_once() + elastic_detail = client.get( + f"/api/computations/runs/{elastic.json()['id']}" + ) + assert elastic_detail.status_code == 200, elastic_detail.text + assert elastic_detail.json()["status"] == "data_ingested" + assert elastic_detail.json()["upstream_run_id"] == run_id + assert elastic_detail.json()["output"]["bulk_modulus_vrh"] == 90.0 + + conn = _connection() + try: + with conn.cursor() as cur: + cur.execute( + """ + SELECT upstream_type_id, handoff_type + FROM project_experiments + WHERE project_id = 'E2E-PROJECT' + AND type_id = 'sim2l_elasticdft_r27_dt1' + """ + ) + assert cur.fetchone() == ( + "sim2l_cellrelaxdft_r99_dt1", + "automated_computational_job", + ) + finally: + conn.close() diff --git a/api/tests/test_computations_router.py b/api/tests/test_computations_router.py new file mode 100644 index 0000000..a91c796 --- /dev/null +++ b/api/tests/test_computations_router.py @@ -0,0 +1,28 @@ +import unittest +from unittest.mock import patch + +from fastapi import FastAPI +from fastapi.testclient import TestClient +from routers import computations +from security import PlatformUser, get_platform_user + + +class ComputationRouterSafetyTests(unittest.TestCase): + def test_feature_disabled_returns_not_found_without_database_access(self): + app = FastAPI() + app.include_router(computations.router, prefix="/api") + app.dependency_overrides[get_platform_user] = lambda: PlatformUser( + id="test-user", + email="test@purdue.edu", + name="Test User", + role="researcher", + organization="Purdue University", + ) + with patch("routers.computations.feature_enabled", return_value=False): + response = TestClient(app).get("/api/computations/tools") + self.assertEqual(response.status_code, 404) + self.assertEqual(response.json()["detail"], "RCAC computations are not enabled") + + +if __name__ == "__main__": + unittest.main() diff --git a/api/tests/test_qe_parser.py b/api/tests/test_qe_parser.py new file mode 100644 index 0000000..de90e9b --- /dev/null +++ b/api/tests/test_qe_parser.py @@ -0,0 +1,43 @@ +import tempfile +import unittest +from pathlib import Path + +from computations.qe import QeParseError, parse_qe_output + +COMPLETE_OUTPUT = """ + ! total energy = -15.25000000 Ry + Forces acting on atoms (cartesian axes, Ry/au): + + atom 1 type 1 force = 0.00100000 -0.00200000 0.00300000 + Total force = 0.003741657 Ry/au + total stress (Ry/bohr**3) (kbar) P= 5.50 + 0.0000 0.0000 0.0000 10.0 1.0 2.0 + 0.0000 0.0000 0.0000 1.0 20.0 3.0 + 0.0000 0.0000 0.0000 2.0 3.0 30.0 + End of BFGS Geometry Optimization + JOB DONE. +""" + + +class QeParserTests(unittest.TestCase): + def test_extracts_energy_pressure_stress_and_forces(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "qe.out" + path.write_text(COMPLETE_OUTPUT, encoding="utf-8") + result = parse_qe_output(path, natoms=1) + self.assertEqual(result["ionic_energies_ry"], [-15.25]) + self.assertEqual(result["pressures_kbar"], [5.5]) + self.assertEqual(result["stresses_kbar"][0][1][1], 20.0) + self.assertEqual(result["atomic_forces_ry_bohr"][0][0][2], 0.003) + self.assertTrue(result["convergence"]) + + def test_rejects_truncated_output(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "qe.out" + path.write_text("JOB DONE.\n", encoding="utf-8") + with self.assertRaises(QeParseError): + parse_qe_output(path) + + +if __name__ == "__main__": + unittest.main() diff --git a/geddes/k8s/09-rcac-worker.example.yaml b/geddes/k8s/09-rcac-worker.example.yaml new file mode 100644 index 0000000..9f46539 --- /dev/null +++ b/geddes/k8s/09-rcac-worker.example.yaml @@ -0,0 +1,55 @@ +# Example only. It is intentionally scaled to zero and is not part of the +# deployed manifests. Complete the RCAC security review before enabling it. +apiVersion: apps/v1 +kind: Deployment +metadata: + name: dt-computation-worker + namespace: ncn-digitaltwins-zchen +spec: + replicas: 0 + selector: + matchLabels: + app: dt-computation-worker + template: + metadata: + labels: + app: dt-computation-worker + spec: + containers: + - name: worker + image: REPLACE_WITH_REVIEWED_WORKER_IMAGE + imagePullPolicy: IfNotPresent + resources: + requests: + cpu: "250m" + memory: "1Gi" + limits: + cpu: "2" + memory: "4Gi" + env: + - name: RCAC_COMPUTE_ENABLED + value: "true" + envFrom: + - secretRef: + name: dt-api-secret + - secretRef: + name: dt-rcac-computation-config + volumeMounts: + - name: rcac-ssh + mountPath: /var/run/secrets/rcac-ssh + readOnly: true + - name: known-hosts + mountPath: /etc/ssh + readOnly: true + - name: worker-data + mountPath: /var/lib/dt-computation-worker + volumes: + - name: rcac-ssh + secret: + secretName: dt-rcac-user-keys + defaultMode: 0400 + - name: known-hosts + configMap: + name: dt-rcac-known-hosts + - name: worker-data + emptyDir: {} diff --git a/web/app/projects/[id]/page.tsx b/web/app/projects/[id]/page.tsx index 09c2406..685a0c7 100644 --- a/web/app/projects/[id]/page.tsx +++ b/web/app/projects/[id]/page.tsx @@ -5,6 +5,7 @@ import Link from "next/link"; import { useRouter } from "next/navigation"; import { cn } from "@/lib/utils"; import { useAuth } from "@/lib/auth-context"; +import { ProjectComputations } from "@/components/project-computations"; import { deleteProject, createPublicationRequest, @@ -1119,6 +1120,8 @@ export default function ProjectDetailPage({ params }: { params: Promise<{ id: st )} + + {/* Unit Experiments / Process */}
diff --git a/web/components/project-computations.tsx b/web/components/project-computations.tsx new file mode 100644 index 0000000..bf36f46 --- /dev/null +++ b/web/components/project-computations.tsx @@ -0,0 +1,272 @@ +"use client"; + +import { useCallback, useEffect, useMemo, useState } from "react"; +import { Activity, Atom, Ban, ChevronDown, ChevronUp, Loader2, Play } from "lucide-react"; +import { + cancelComputationRun, + getComputationCatalog, + getProjectComputationRuns, + getRcacProfile, + saveRcacProfile, + submitComputationRun, + type ComputationCatalog, + type ComputationRun, +} from "@/lib/api-client"; + +const ACTIVE = new Set([ + "requested", + "validating", + "staging", + "submitted", + "running", + "collecting", + "cancelling", +]); + +const CELL_DEFAULTS: Record = { + is_metal: true, + bulk_m: 2000, + KE_cutoff: 50, + k_x: 8, + k_y: 8, + k_z: 8, + i_steps: 50, + e_steps: 200, + scf_conv: 1e-10, + energy_conv: 1e-10, + force_conv: 1e-5, + pressure_conv: 0.5, + target_pressure: 0, + vdw_corr: "none", + spin_polar: false, + hubbard_U_on: false, + hubbard_projector: "ortho-atomic", + hubbard_U_values: {}, + hubbard_U_orbitals: {}, + ibrav: 0, + exchange_correlation: "PBEsol", + pseudo: "PAW", +}; + +function pretty(value: unknown) { + return JSON.stringify(value, null, 2); +} + +export function ProjectComputations({ projectId }: { projectId: string }) { + const [catalog, setCatalog] = useState(null); + const [runs, setRuns] = useState([]); + const [loading, setLoading] = useState(true); + const [tool, setTool] = useState("cellrelaxdft"); + const [cluster, setCluster] = useState(""); + const [mpId, setMpId] = useState("mp-1265"); + const [upstreamId, setUpstreamId] = useState(""); + const [advancedInputs, setAdvancedInputs] = useState(pretty(CELL_DEFAULTS)); + const [submitting, setSubmitting] = useState(false); + const [error, setError] = useState(""); + const [expanded, setExpanded] = useState(null); + const [showProfile, setShowProfile] = useState(false); + const [profile, setProfile] = useState({ + ssh_username: "", + allocation: "", + partition: "standby", + credential_ref: "", + }); + + const loadRuns = useCallback(async () => { + const value = await getProjectComputationRuns(projectId); + if (value) setRuns(value); + }, [projectId]); + + useEffect(() => { + let cancelled = false; + void (async () => { + const value = await getComputationCatalog(); + if (cancelled) return; + setCatalog(value); + if (value?.clusters[0]) setCluster(value.clusters[0].name); + if (value) await loadRuns(); + setLoading(false); + })(); + return () => { cancelled = true; }; + }, [loadRuns]); + + useEffect(() => { + if (!catalog || !runs.some((run) => ACTIVE.has(run.status))) return; + const timer = window.setInterval(() => void loadRuns(), 8000); + return () => window.clearInterval(timer); + }, [catalog, loadRuns, runs]); + + useEffect(() => { + if (!cluster) return; + void getRcacProfile(cluster).then((value) => { + if (!value) return; + setProfile({ + ssh_username: value.ssh_username, + allocation: value.allocation, + partition: value.partition, + credential_ref: "", + }); + }); + }, [cluster]); + + const completedCellRuns = useMemo( + () => runs.filter((run) => run.tool_name === "cellrelaxdft" && run.status === "data_ingested"), + [runs], + ); + + if (loading || !catalog) return null; + + async function submit() { + setError(""); + setSubmitting(true); + try { + let inputs: Record; + if (tool === "cellrelaxdft") { + const advanced = JSON.parse(advancedInputs) as Record; + inputs = { ...advanced, mp_id: mpId, structure_dict: advanced.structure_dict || {} }; + } else { + if (!upstreamId) throw new Error("Select a completed CellRelaxDFT run."); + inputs = { + mp_id: mpId, + upstream_execution_id: upstreamId, + diagonal_strain_magnitudes: [-0.01, -0.005, 0.005, 0.01], + offdiagonal_strain_magnitudes: [-0.01, -0.005, 0.005, 0.01], + }; + } + const value = await submitComputationRun(projectId, tool, { cluster, inputs }); + if (!value) throw new Error("Submission failed. Check the RCAC profile and input values."); + await loadRuns(); + setExpanded(value.id); + } catch (err) { + setError(err instanceof Error ? err.message : "Submission failed"); + } finally { + setSubmitting(false); + } + } + + async function saveProfile() { + setError(""); + if (!profile.credential_ref) { + setError("Enter the credential reference provisioned by the DT operator."); + return; + } + const value = await saveRcacProfile(cluster, profile); + if (!value) setError("Could not save the RCAC profile."); + else setShowProfile(false); + } + + return ( +
+
+
+
+ +

+ Quantum ESPRESSO computations +

+
+

+ Sim2L inputs and cached results with execution on RCAC Slurm. +

+
+ +
+ + {showProfile && ( +
+ setProfile({ ...profile, ssh_username: e.target.value })} placeholder="RCAC username" className="rounded-md border bg-transparent px-3 py-2 text-sm" /> + setProfile({ ...profile, allocation: e.target.value })} placeholder="Allocation" className="rounded-md border bg-transparent px-3 py-2 text-sm" /> + setProfile({ ...profile, partition: e.target.value })} placeholder="Partition" className="rounded-md border bg-transparent px-3 py-2 text-sm" /> + setProfile({ ...profile, credential_ref: e.target.value })} placeholder="Credential reference" className="rounded-md border bg-transparent px-3 py-2 text-sm" /> + +

+ The reference identifies an operator-provisioned SSH key. Passwords and private-key contents are never entered here. +

+
+ )} + +
+ + + setMpId(e.target.value)} placeholder="Materials Project ID" className="rounded-md border bg-transparent px-3 py-2 text-sm" /> + {tool === "elasticdft" ? ( + + ) : ( + + )} +
+ + {tool === "cellrelaxdft" && ( +
+ Advanced CellRelaxDFT inputs +