Walkthrough: agents on the Spark, compute on Sunrise

Info

Sunrise is the Stockholm University Fysikum Cluster, please replace the necessary information here for your local cluster. Same is true for mentions to the Spark which is where this system was developed.

This walks through the "Spark orchestrates, cluster computes" shape end to end: the agents (opencode + Ollama) and the MCP server run on the DGX Spark as usual; heavy steps (ingest of a large catalogue, spectra, covariance, inference) run as Slurm jobs on Sunrise, reached over SSH. If instead you'll run the agents on the cluster itself (an interactive GPU-node allocation), see Same-cluster execution — the tool calls are identical, only the site config differs.

0. Prerequisites

• A Sunrise account with SSH access and an SSH key already generated (ssh-keygen, or whatever your site's onboarding gave you).
• cosmotron_mcp installed in the cosmotron conda env on the Spark (this is what drives everything — see Install).
• Sunrise reachable at all: ssh <you>@sol-login.fysik.su.se -p <port> should already work manually before any of this.

1. Write sites.yaml

Copy the Sunrise-shaped block from sites.yaml.example to ~/.cosmotron/sites.yaml (or ./sites.yaml in the repo root — gitignored) and fill in the placeholders:

sites:
  sunrise:
    host: sol-login.fysik.su.se
    port: <custom port from the welcome email>
    user: <your username>
    fallback_hosts:
      - sol-nix.fysik.su.se
    auth:
      type: static_key
      key_path: ~/.ssh/<your sunrise key>
    remote_root: /cfs/home/<your username>/cosmotron
    scratch_root: /cfs/nobackup/<your group>/<your username>
    env_setup:
      - module load conda
      - conda activate cosmotron
    env:
      FIRECROWN_DIR: <remote firecrown package dir>            # only needed for run_inference
      COSMOSIS_STD_LIB_DIR: <remote cosmosis-standard-library checkout>
    queues:
      default:
        partition: Solar
      bigmem:
        partition: CoPS

Or skip hand-editing and run the interactive wizard instead, which probes the cluster for you and fills most of this in:

cosmotron-remote add-site

2. Provision the cluster once

cosmotron-remote bootstrap sunrise

This builds a conda env on Sunrise from the committed envs/cosmotron-hpc.lock.yml, then rsyncs cosmotron_mcp/ + pyproject.toml over and pip install -es the package into that env, then runs doctor as a final check. (The env is created first, since the install step activates it.) It's idempotent — safe to re-run after a code change (or use sync-code alone for a quick code push once the env already exists).

If you'll run run_inference remotely, build the cosmosis standard library once on the cluster (cosmosis-build-standard-library) and point sites.yaml's env: FIRECROWN_DIR/COSMOSIS_STD_LIB_DIR at the resulting paths — either hand-edit the file, or set them from the CLI:

cosmotron-remote edit-site sunrise \
  --set-env FIRECROWN_DIR=/path/to/cosmosis-standard-library/likelihood/firecrown \
  --set-env COSMOSIS_STD_LIB_DIR=/path/to/cosmosis-standard-library

(doctor, below, also offers to fill these in interactively if it finds them missing.)

3. Verify

cosmotron-remote doctor sunrise

Expect all-PASS: SSH reachable, sbatch present, remote_root writable, env_setup runs, cosmotron_mcp importable, local/remote versions match, and (if configured) FIRECROWN_DIR/COSMOSIS_STD_LIB_DIR resolve. Re-run this any time something feels off — it's the single source of truth for "is this site actually usable right now".

4. Where does the data go?

Short answer: keep your raw catalogue on the Spark, in the same data/ directory you'd use for a fully-local run. bootstrap_session and ingest_to_session read it from there; when a heavy step (ingest, spectra, covariance) is dispatched with jobspec, cosmoTRON stages only the specific files that step needs to the cluster automatically — you never manually scp/rsync anything yourself for the common case.

The one exception: if your catalogue is already sitting on Sunrise's own filesystem (e.g. because it's a shared community dataset, or because it's too large to comfortably keep two copies), point bootstrap_session at it with a site:path reference instead of a bare path — cosmoTRON enumerates it remotely and runs ingest as a Slurm job so the multi-GB FITS read never touches the Spark or a login node. See Data placement & the dataset registry for the full story (including what happens to standardised products afterward, and the one-time push/pull operator commands if you do need to move something across manually).

5. Run something remotely

Any remote-capable tool (see the table in Overview) takes the same jobspec. Two concrete examples:

Covariance (assuming you already have Cls computed locally or via a prior remote step):

receipt = compute_full_covariance_from_session(
    session_dir, jobspec={"backend": "slurm", "site": "sunrise"},
)
# {"job_id": "compute_full_covariance_from_session__a1b2c3d4",
#  "status": "pending", "mode": "async", "site": "sunrise",
#  "resource_class": "covariance", "next": "check_job"}

Inference (a CosmoSIS/Firecrown chain):

receipt = run_inference(
    session_dir, jobspec={"backend": "slurm", "site": "sunrise"},
)

Both return immediately — the numbers are NOT in the receipt. Poll:

check_job(session_dir, receipt["job_id"])
# {"job_id": "...", "status": "running", "slurm_state": "RUNNING",
#  "stdout_tail": "...", "poll_error": null, "stale": false, ...}

status moves pending → running → done/failed, mirroring real Slurm states (squeue/sacct under the hood). Once done, the result — the covariance JSON, the posterior chain, whatever the tool normally writes — is already sitting in results/ under your session, pulled back automatically; re-call the same session-reading tools (validate_covariance_knox, summarise_posterior, assemble_sacc_from_session, ...) exactly as you would after a local run. Nothing about the rest of the pipeline needs to know the step ran remotely.

6. Doing this from an agent, not raw Python

From an agent you don't pass a jobspec — you declare the intent in the task and the tools attach it. The clean way is a ## remote: block in a TASK.md (or inline):

Tomographic clustering + covariance for data/my_survey_tomo/ at nside=256.

## remote:
site: sunrise
covariance

bootstrap_session resolves this to pipeline_config.remote = {site: sunrise, steps: [covariance]}; then compute_full_covariance_from_session auto-derives its own jobspec from that config (the model never picks a site, step, or queue — add an optional queue: <name> line under site: to steer every remote step in the block to a non-default sites.yaml queue), and @analysis_coder polls the resulting async job through to completion the same way it handles any background job. Everything else in that run stays local. Use ## site: sunrise for "all heavy steps there", or a sunrise:/cfs/.../cat data path to run from data that already lives on the cluster (ingest then goes remote automatically). Full grammar: Writing TASK.md.

Troubleshooting

• doctor fails at "ssh reachable" — check the key permissions (chmod 600 ~/.ssh/<key>) and that you can ssh manually with the exact host/port/user in sites.yaml.
• sbatch rejected: Invalid partition — the queues.default.partition name doesn't match what sinfo shows on Sunrise; re-run cosmotron-remote add-site and read off the real partition name.
• An MPI inference chain looks like N duplicate serial runs, no speed-up — the PMI wiring is missing; set slurm.srun_args: ["--mpi=pmix"] in sites.yaml (conda-forge Open MPI usually needs this).
• Not sure what's actually in sites.yaml for a site, or where the file is — cosmotron-remote view sunrise prints the resolved config (source file path, env, every queue) without needing to go find and read the YAML by hand. Fix what it shows with cosmotron-remote edit-site sunrise --add-queue <name> --partition <p> / --set-env KEY=VALUE / --remove-queue <name>.
• check_job shows stale: true / a poll_error — a transient SSH hiccup, not a job failure; the job keeps running on the cluster independently of whether the Spark can currently reach it. Re-poll.
• auth_required: ... in error/poll_error — only relevant for sshproxy-type auth (NERSC); see NERSC & same-cluster agents.