Querying SciFMs Hosted by Nomad¶

This notebook walks through a complete Nomad inference workflow against a local MCP endpoint launched from this repository’s container/deploy/demo configuration. You will start the server with five MIST models and the Janus HEAT PLI model, discover the SciFMs it exposes, inspect a model card, and run a PDE surrogate that returns a full solution field.

By the end, you will have confirmed connectivity to the local endpoint, fetched model metadata, and compared an input state with the model’s predicted final state.

The first server startup can take a while because uv may need to resolve the environment for container/deploy/demo. You also need access to the public model sources referenced by container/deploy/demo/nomad.yml.

Launch the Nomad Server¶

This notebook starts nomad serve in the background using the demo server project under container/deploy/demo. That project includes a public demo config containing five MIST models and the Janus HEAT PLI model. The hidden cells below define the notebook-specific startup helpers, and the visible cell after them launches the server.

You need uv installed because the launcher uses uv run --directory container/deploy/demo ....

Show code cell source

Hide code cell source

import atexit
import os
import socket
import tempfile
import time
from dataclasses import dataclass

from fastmcp import Client


def get_free_port(host: str = "127.0.0.1") -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind((host, 0))
        return int(sock.getsockname()[1])


GUIDE_DIR = REPO_ROOT / "docs" / "guides"


@dataclass
class NomadServerHandle:
    server_dir: Path
    config_path: Path
    host: str
    port: int
    url: str
    log_path: Path
    process: subprocess.Popen[str]


def resolve_server_dir(server_dir: str | Path) -> Path:
    candidate = Path(server_dir).expanduser()
    if not candidate.is_absolute():
        candidate = REPO_ROOT / candidate
    candidate = candidate.resolve()
    if not candidate.is_dir():
        raise FileNotFoundError(f"Server directory does not exist: {candidate}")
    if not (candidate / "pyproject.toml").is_file():
        raise FileNotFoundError(f"Expected a pyproject.toml file in: {candidate}")
    if not (candidate / "nomad.yml").is_file():
        raise FileNotFoundError(f"Expected a nomad.yml file in: {candidate}")
    return candidate


def read_server_log(log_path: Path) -> str:
    if not log_path.is_file():
        return ""
    return log_path.read_text(encoding="utf-8", errors="replace")


def tail_server_log(log_path: Path, max_lines: int = 40) -> str:
    text = read_server_log(log_path).strip()
    if not text:
        return "No server log was captured."
    lines = text.splitlines()
    return "\n".join(lines[-max_lines:])


def require_nomad_server(server_dir: str | Path | None = None) -> NomadServerHandle:
    target_dir = resolve_server_dir(server_dir) if server_dir is not None else None
    server = globals().get("NOMAD_SERVER")

    if server is None:
        if target_dir is None:
            raise RuntimeError(
                "Nomad server is not running. Run start_nomad_server(...) first."
            )
        start_nomad_server(target_dir)
        server = globals().get("NOMAD_SERVER")

    if server is not None and server.process.poll() is not None:
        if target_dir is None:
            raise RuntimeError(
                "Nomad server is not running. Restart it with start_nomad_server(...).\n\n"
                f"{tail_server_log(server.log_path)}"
            )
        start_nomad_server(target_dir)
        server = globals().get("NOMAD_SERVER")

    if (
        server is not None
        and target_dir is not None
        and server.server_dir != target_dir
    ):
        start_nomad_server(target_dir)
        server = globals().get("NOMAD_SERVER")

    if server is None:
        raise RuntimeError("Nomad server failed to start.")
    return server


def stop_nomad_server() -> None:
    server = globals().get("NOMAD_SERVER")
    if server is None:
        return

    process = server.process
    if process.poll() is None:
        process.terminate()
        try:
            process.wait(timeout=10)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait(timeout=10)

    log_handle = globals().pop("_nomad_server_log_handle", None)
    if log_handle is not None and not log_handle.closed:
        log_handle.close()

    globals()["NOMAD_SERVER"] = None
    globals().pop("NOMAD_URL", None)


def start_nomad_server(
    server_dir: str | Path,
    host: str = "127.0.0.1",
    port: int | None = None,
    timeout_seconds: float = 300.0,
):
    if shutil.which("uv") is None:
        raise RuntimeError(
            "uv is required to launch the local server. Install it, then rerun this cell."
        )

    resolved_server_dir = resolve_server_dir(server_dir)
    config_path = resolved_server_dir / "nomad.yml"

    existing = globals().get("NOMAD_SERVER")
    if (
        existing is not None
        and existing.process.poll() is None
        and existing.server_dir == resolved_server_dir
        and existing.host == host
        and (port is None or existing.port == port)
    ):
        return {
            "server_dir": str(existing.server_dir.relative_to(REPO_ROOT)),
            "config": str(existing.config_path.relative_to(REPO_ROOT)),
            "url": existing.url,
            "log_path": str(existing.log_path),
        }

    stop_nomad_server()

    selected_port = port or get_free_port(host)
    log_path = Path(tempfile.gettempdir()) / f"nomad_inference_{selected_port}.log"
    log_handle = open(log_path, "w", encoding="utf-8")

    env = os.environ.copy()
    env["MPLBACKEND"] = "Agg"
    env.pop("VIRTUAL_ENV", None)
    env.pop("UV_FROZEN", None)

    command = [
        "uv",
        "run",
        "--directory",
        str(resolved_server_dir),
        "nomad",
        "serve",
        "--transport",
        "streamable-http",
        "--host",
        host,
        "--port",
        str(selected_port),
        "--log-level",
        "INFO",
        str(config_path),
    ]
    process = subprocess.Popen(
        command,
        cwd=resolved_server_dir,
        env=env,
        stdout=log_handle,
        stderr=subprocess.STDOUT,
        text=True,
    )

    server = NomadServerHandle(
        server_dir=resolved_server_dir,
        config_path=config_path,
        host=host,
        port=selected_port,
        url=f"http://{host}:{selected_port}/mcp",
        log_path=log_path,
        process=process,
    )
    globals()["NOMAD_SERVER"] = server
    globals()["NOMAD_URL"] = server.url
    globals()["_nomad_server_log_handle"] = log_handle

    deadline = time.perf_counter() + timeout_seconds
    while time.perf_counter() < deadline:
        if process.poll() is not None:
            error_log = tail_server_log(log_path)
            stop_nomad_server()
            raise RuntimeError(
                f"Local Nomad server exited before it became ready.\n\n{error_log}"
            )

        log_text = read_server_log(log_path)
        if (
            "Application startup complete." in log_text
            and "Uvicorn running on" in log_text
        ):
            return {
                "server_dir": str(server.server_dir.relative_to(REPO_ROOT)),
                "config": str(server.config_path.relative_to(REPO_ROOT)),
                "url": server.url,
                "log_path": str(server.log_path),
            }

        time.sleep(1)

    error_log = tail_server_log(log_path)
    stop_nomad_server()
    raise TimeoutError(
        f"Timed out waiting for the local Nomad server at {server.url}.\n\n{error_log}"
    )


atexit.register(stop_nomad_server)

SERVER_DIR = "container/deploy/demo"


def get_client():
    """Function for connecting to a local Nomad server
    Update to point to a hosted one, by just returning a
    StreamableHttpTransport(url=nomad_mcp_server_url, headers=...)
    """
    from fastmcp.client.transports import StreamableHttpTransport

    server = require_nomad_server(SERVER_DIR)
    return StreamableHttpTransport(url=server.url)


start_nomad_server(SERVER_DIR)

List Available Tools¶

This is the fastest way to confirm that the local endpoint is reachable and that the Nomad tools are visible through MCP. If you want more detail on the transport layer itself, see the Code-Mode Gateway reference and the MCP documentation.

async def list_tools(name_filter: str | None = None) -> list[str]:
    # Ask the local server for the exposed tools, then optionally narrow the list.
    async with Client(get_client()) as client:
        tools = await client.list_tools()

    names = sorted(tool.name for tool in tools)
    if name_filter is None:
        return names
    return [name for name in names if name_filter in name]


tool_names = await list_tools()
model_card_tool_name = "get_model_card"
mist_qm9_tool_name = "mist_models---mist_26p9M_kkgx0omx_qm9"
heat_pli_tool_name = "diffunet2_heat_pli_f2l"

missing_tool_names = [
    name
    for name in (model_card_tool_name, mist_qm9_tool_name, heat_pli_tool_name)
    if name not in tool_names
]
if missing_tool_names:
    raise LookupError(
        f"Expected tools were not exposed: {missing_tool_names}.\n"
        f"Available tools: {tool_names}"
    )

{
    "model_card_tool_name": model_card_tool_name,
    "mist_qm9_tool_name": mist_qm9_tool_name,
    "heat_pli_tool_name": heat_pli_tool_name,
    "tool_names": tool_names,
}

Fetch a Model Card¶

Nomad provides a get_model_card(tool_name: str) tool for retrieving the model card, when one is available, for each SciFM it serves. The model card is the README.md in the SciFM’s model directory. It is the quickest way to inspect a deployed model’s metadata, expected inputs, and intended behavior before running inference.

If you want the builder-side view of how SciFMs are served by Nomad, the Model Builders guide is the next place to read. You can use the template model card as a starting point for documenting your own model.

model_card = await get_model_card(tool_name=model_card_tool_name)

display_model_card(
    model_card=model_card,
    max_height="28rem",
)

Run Inference on a PDE SciFM¶

Now that we are connected to a Nomad endpoint, let’s query a PDE surrogate trained on the HEAT PLI dataset. These models take an initial state as an AutoRegressiveInput and return a predicted trajectory, or a final timestep, as a WellFormat.

We will start by loading an initial state with WellFormat.from_file, inspecting its metadata, and visualizing the input field before sending it to the local model.

import matplotlib.pyplot as plt
import numpy as np

from nomad.well_format import AutoRegressiveInput, WellFormat

INPUT_FILE = GUIDE_DIR / "angled_half_res.h5"
input_state = WellFormat.from_file(INPUT_FILE)

{
    "input_file": str(INPUT_FILE.relative_to(REPO_ROOT)),
    "dataset_name": input_state.dataset_name,
    "grid_type": input_state.grid_type,
    "n_spatial_dims": input_state.n_spatial_dims,
    "field_count": len(input_state.t0_fields),
}

plot_density(
    input_state=input_state,
    field="av_density",
    input_time_index=0,  # Change this to look at different timesteps of the simulation
);

Run Local Inference¶

We can now query the locally hosted SciFM to predict the final state of the HEAT PLI simulation. The request packages the initial condition into AutoRegressiveInput, sends it to the local tool, and parses the returned WellFormat response.

The final plot uses a shared color scale so you can compare the initial field with the predicted final state at a glance.

The same pattern works in scripts, larger pipelines, and builder-side evaluation loops; see Getting Started and the CLI reference for adjacent workflows.

prediction_state, elapsed_seconds = await run_inference(
    initial_state=input_state,
    duration=100,
)

print(f"Inference completed in {elapsed_seconds:.2f} seconds")

plot_density(
    input_state=input_state,
    predicted_state=prediction_state,
    field="av_density",
    input_time_index=0,
    prediction_time_index=-1,
);

Troubleshooting¶

Missing uv: install uv, then rerun the install and server-start cells.
Kernel has no pip: the setup cell uses uv pip --python <current-kernel> automatically, so rerun that cell after switching kernels.
Startup failures: start_nomad_server(...) returns the server log path; inspect it if the launch fails or times out.
Changing the local deployment: edit container/deploy/demo/nomad.yml, then rerun the server-start and tool-discovery cells.
Stopping the background server: call stop_nomad_server() if you want to shut it down before the kernel exits.
Need more detail? See Getting Started, the Well Format reference, the Code-Mode Gateway reference, the CLI reference, and the Model Builders guide.