feat(package): wire gridfm-graphkit vLLM serving (PowerFlow reconstruction)

packages/gridfm-graphkit/models/gridfm-pf-reconstruction/model.yaml

@@ -0,0 +1,15 @@
+model:
+  # Trained GridFM PowerFlow checkpoint (case14_ieee, GNS_heterogeneous),
+  # exported via gridfm_graphkit.vllm.export.export_model and published to the Hub.
+  id: gridfm/gridfm-pf-reconstruction
+
+  vllm:
+    enabled: true
+    plugins:
+      # Registers the custom pooling-model class (architecture "GridFMGNS").
+      # Installed by gridfm-graphkit[vllm] via the vllm.general_plugins entry point.
+      general: gridfm_gns
+      io_processors:
+        # Builds a normalized HeteroData graph from a power-grid case, serves the
+        # PowerFlow reconstruction, and returns per-node embeddings + Vm/Va/Pg.
+        - gridfm_pf_reconstruction

packages/gridfm-graphkit/models/gridfm-pf-reconstruction/usage.md

@@ -0,0 +1,124 @@
+# Using the GridFM PowerFlow Reconstruction Model
+
+This guide explains how to serve the GridFM heterogeneous GNN for PowerFlow
+reconstruction through vLLM's `/pooling` endpoint.
+
+## Overview
+
+GridFM (`gridfm/gridfm-pf-reconstruction`) is a heterogeneous graph neural
+network (`GNS_heterogeneous`) over power-grid graphs. Buses and generators are
+nodes; branches and bus↔generator connections are typed edges. Served for the
+**PowerFlow reconstruction** task, it returns, per request:
+
+- per-node latent **embeddings** (bus and generator), and
+- **predictions**: bus voltage magnitude/angle (Vm, Va) and generator active
+  power (Pg), denormalized back to physical units.
+
+Because the model consumes a graph rather than text, it is served as a vLLM
+**pooling** model paired with an I/O processor plugin that translates a
+power-grid case into the graph tensors the model expects and translates the
+pooled output back into a structured response.
+
+## Prerequisites
+
+### Installation
+
+Install gridfm-graphkit with its vLLM extra (pulls in `vllm==0.26.x`):
+
+```bash
+pip install "gridfm-graphkit[vllm]"
+```
+
+This installs two entry points that vLLM discovers at startup:
+
+- `vllm.general_plugins` → `gridfm_gns` — registers the `GridFMGNS` pooling
+  model architecture.
+- `vllm.io_processor_plugins` → `gridfm_pf_reconstruction` — the PowerFlow
+  reconstruction I/O processor.
+
+## Starting the vLLM Server
+
+```bash
+vllm serve gridfm/gridfm-pf-reconstruction \
+  --runner pooling \
+  --trust-remote-code \
+  --skip-tokenizer-init \
+  --enforce-eager \
+  --io-processor-plugin gridfm_pf_reconstruction \
+  --enable-mm-embeds
+```
+
+### Command Arguments Explained
+
+- `--runner pooling`: serve as a pooling/encoder model (no token generation).
+- `--trust-remote-code`: allow loading the model's config/architecture.
+- `--skip-tokenizer-init`: no tokenizer is needed for a graph model.
+- `--enforce-eager`: eager execution (the model is attention-free).
+- `--io-processor-plugin gridfm_pf_reconstruction`: enable the PowerFlow
+  reconstruction I/O processor.
+- `--enable-mm-embeds`: enable the multimodal path used to ship graph tensors.
+
+The server starts on `http://localhost:8000` by default.
+
+## Making Inference Requests
+
+The I/O processor accepts a power-grid case as record lists for buses,
+generators and branches (the same columns GridFM's dataset builder consumes).
+Send it to the `/pooling` endpoint.
+
+### Python Client Example
+
+```python
+import requests
+
+VLLM_SERVER_ENDPOINT = "http://localhost:8000/pooling"
+
+request_payload = {
+    "data": {
+        "case": {
+            "bus": [ /* per-bus records */ ],
+            "gen": [ /* per-generator records */ ],
+            "branch": [ /* per-branch records */ ],
+        },
+        "return_embeddings": True,
+        "return_predictions": True,
+    },
+    "priority": 0,
+    "model": "gridfm/gridfm-pf-reconstruction",
+}
+
+response = requests.post(VLLM_SERVER_ENDPOINT, json=request_payload)
+response.raise_for_status()
+result = response.json()["data"]
+
+print("num buses:", result["num_buses"])
+print("num gens: ", result["num_gens"])
+# result["bus_predictions"]  -> per-bus [Vm, Va, ...] in physical units
+# result["gen_predictions"]  -> per-generator [Pg, ...] in physical units
+# result["bus_embeddings"]   -> per-bus latent vectors
+# result["gen_embeddings"]   -> per-generator latent vectors
+```
+
+### Request Payload Structure
+
+- **`data`**: the input configuration
+    - **`case`**: the power-grid case as `bus` / `gen` / `branch` record lists.
+    - **`return_embeddings`**: include per-node latent embeddings in the response.
+    - **`return_predictions`**: include per-node Vm/Va/Pg predictions.
+- **`priority`**: request priority (0 = normal).
+- **`model`**: the model identifier.
+
+## Output
+
+The response `data` object contains:
+
+- **`num_buses`**, **`num_gens`**: node counts for the served case.
+- **`bus_predictions`**, **`gen_predictions`**: denormalized per-node
+  predictions (bus Vm/Va, generator Pg).
+- **`bus_embeddings`**, **`gen_embeddings`**: per-node latent embeddings.
+
+## Additional Resources
+
+- [gridfm-graphkit](https://github.com/gridfm/gridfm-graphkit)
+- [vLLM pooling models](https://docs.vllm.ai/en/latest/models/pooling_models.html)
+- [vLLM I/O processor plugins](https://docs.vllm.ai/en/latest/design/io_processor_plugins.html)

pyproject.toml

@@ -20,6 +20,11 @@ nexus = "algorithm_nexus.cli:main"
 candidate = [
   "biomed-multi-alignment[vllm]",
   "bmfm-targets[vllm]>=0.20.0",
+  # gridfm-graphkit 0.9.0.3 ships the vLLM serving subpackage + `vllm` extra +
+  # entry points (gridfm_gns, gridfm_pf_reconstruction). Its `vllm` extra targets
+  # vllm>=0.29,<0.30, matching this candidate variant (vllm==0.29.0). It is
+  # intentionally NOT added to `product` (vllm>=0.28,<0.29) until validated there.
+  "gridfm-graphkit[vllm]==0.9.0.3",
   # CP 02/07/2026
   # Scanners are flagging opencv-python-headless==5.0.0.93 as
   # potentially malicious. While we think it's a false positive,
@@ -40,7 +45,7 @@ cli = [
 ecosystem = [
   "biomed-multi-alignment",
   "bmfm-targets",
-  "gridfm-graphkit==0.8.1",
+  "gridfm-graphkit==0.9.0.3",
   "terrakit>=0.2.0",
   "terratorch>=1.2.7",
   "tokamind",
@@ -87,6 +92,13 @@ conflicts = [
   ],
 ]
 
+# gridfm-graphkit (>=0.9) depends on torch-scatter, whose sdist imports torch at
+# build time without declaring it as a build dependency. Inject torch into its
+# build environment so `uv lock` can generate its metadata. (Must sit after all
+# flat [tool.uv] keys above, or it orphans them.)
+[tool.uv.extra-build-dependencies]
+torch-scatter = ["torch"]
+
 [tool.uv.sources]
 bmfm-targets = { git = "https://github.com/BiomedSciAI/biomed-multi-omic.git" }
 tokamind = { git = "https://github.com/UKAEA-IBM-STFC-Fusion-FMs/tokamind", tag = "OS_v2.0" }