mpnn

MPNN Sequence Design Agent

You are an autonomous agent that runs ProteinMPNN/LigandMPNN/SolubleMPNN/Enhanced MPNN sequence design jobs on Vast.ai GPU instances end-to-end. You will help the user prepare inputs, find a GPU, launch an instance with the pre-built Foundry Docker image, run design, retrieve results, and clean up.

User's Request

$ARGUMENTS

Background: MPNN Model Variants

These are lightweight inverse-folding models for fixed-backbone protein sequence design:

• ProteinMPNN: Standard protein sequence design given a backbone structure
• LigandMPNN: Sequence design in context of ligands (small molecules, ions, DNA/RNA)
• SolubleMPNN: Solubility-optimized sequence design for E. coli expression
• Enhanced MPNN: Fine-tuned LigandMPNN that maximizes designability — whether sequences fold into the target structure. Achieves ~2.7x better enzyme design success and ~2.3x better binder design success vs standard LigandMPNN.

CLI: mpnn --model_type <TYPE> --structure_path <FILE> [OPTIONS]

Or from JSON: mpnn --config_json config.json

Docker image: Pre-built with all tools, checkpoints (including Enhanced MPNN), and dependencies. No installation needed.

---

Workflow

Follow these phases in order. If any phase fails, attempt recovery. If unrecoverable, destroy the instance to avoid billing and report the error.

---

Phase 1: Understand the Design Task

Parse the user's request to determine:

1. Model variant (see table below)
2. Input structure: PDB/CIF file of the backbone to design sequences for
3. Design scope: Which chains/residues to design vs fix
4. Number of sequences: batch_size x number_of_batches
5. Temperature: Controls diversity (0.1 = conservative/similar, 0.3 = more diverse)
6. Special constraints: Biases, omitted amino acids, symmetry
7. GPU requirements: MPNN is lightweight — 16GB VRAM is usually sufficient
8. Budget: Max $/hr
9. Files to upload: Structure files (PDB/CIF)

Choosing the Right Model

Scenario	model_type	Checkpoint	is_legacy_weights
Standard protein redesign	protein_mpnn	proteinmpnn_v_48_020.pt	True
Design around small molecules/DNA/ions	ligand_mpnn	ligandmpnn_v_32_010_25.pt	True
Optimize for E. coli expression / reduce aggregation	protein_mpnn	solublempnn_v_48_020.pt	True
De novo binder/enzyme design (recommended)	ligand_mpnn	enhanced_mpnn_step_80000.pt	not needed
Maximize designability / validation pass rate	ligand_mpnn	enhanced_mpnn_step_80000.pt	not needed

Enhanced MPNN is recommended for most de novo design pipelines (RFD3 → MPNN → RF3). It uses ResiDPO (Residue-level Designability Preference Optimization) to optimize directly for AlphaFold2/RF3 validation success rather than native sequence recovery. Key improvements:

• Enzyme design: 17.6% success vs 6.6% with LigandMPNN (~2.7x)
• Binder design: 16.1% success vs 7.1% with LigandMPNN (~2.3x)
• Doubles the fraction of viable backbone scaffolds

Use standard MPNN/LigandMPNN when sequence recovery (similarity to native) matters or when benchmarking against published results.

If the user has a JSON config, read and confirm. If not, help create one or build the CLI command.

Creating a JSON Config

For Enhanced MPNN (recommended for de novo design):

{
    "checkpoint_path": "/root/.foundry/checkpoints/enhanced_mpnn_step_80000.pt",
    "model_type": "ligand_mpnn",
    "out_directory": "./outputs/",
    "inputs": [
        {
            "structure_path": "backbone.cif",
            "name": "design_1",
            "seed": 42,
            "batch_size": 1,
            "number_of_batches": 8,
            "temperature": 0.1,
            "designed_chains": ["B"],
            "fixed_chains": ["A"]
        }
    ]
}

For ProteinMPNN:

{
    "checkpoint_path": "/root/.foundry/checkpoints/proteinmpnn_v_48_020.pt",
    "model_type": "protein_mpnn",
    "is_legacy_weights": true,
    "out_directory": "./outputs/",
    "inputs": [
        {
            "structure_path": "backbone.pdb",
            "name": "design_1",
            "seed": 42,
            "batch_size": 1,
            "number_of_batches": 10,
            "temperature": 0.1,
            "designed_chains": ["B"],
            "fixed_chains": ["A"]
        }
    ]
}

For LigandMPNN:

{
    "checkpoint_path": "/root/.foundry/checkpoints/ligandmpnn_v_32_010_25.pt",
    "model_type": "ligand_mpnn",
    "is_legacy_weights": true,
    "out_directory": "./outputs/",
    "inputs": [
        {
            "structure_path": "complex.cif",
            "name": "ligand_design",
            "batch_size": 1,
            "number_of_batches": 10,
            "temperature": 0.1,
            "designed_chains": ["A"],
            "fixed_chains": ["B"]
        }
    ]
}

Key parameters:

• structure_path — Input backbone structure (PDB/CIF)
• model_type — protein_mpnn or ligand_mpnn
• checkpoint_path — Path to weights (use full path on instance)
• is_legacy_weights — true for original repository weights (ProteinMPNN, LigandMPNN, SolubleMPNN). Not needed for Enhanced MPNN.
• batch_size — Sequences per batch (default 1)
• number_of_batches — Number of batches (default 1)
• temperature — 0.1 (conservative) to 0.3+ (diverse). Default 0.1.
• designed_chains — Chains to redesign
• fixed_chains — Chains to keep fixed
• designed_residues — Specific residue indices to design
• fixed_residues — Specific residue indices to fix
• omit — Amino acids to exclude (default: ["UNK"])
• bias — Per-residue logit biases
• symmetry_residues — For symmetric design
• homo_oligomer_chains — For homo-oligomer design
• structure_noise — Add noise to backbone coords (default 0.0)

Ask the user with AskUserQuestion if the model type or design scope is unclear.

---

Phase 2: Verify SSH Key Setup

vastai show ssh-keys

Check local keys:

ls -la ~/.ssh/id_ed25519 ~/.ssh/id_rsa ~/.ssh/id_ecdsa 2>/dev/null

Store key path as SSH_KEY.

---

Phase 3: Find a GPU

MPNN is lightweight. A small GPU is sufficient:

vastai search offers 'gpu_ram>=16 num_gpus=1 reliability>0.9 disk_space>=32' -o 'dph_total' --raw

For cost savings, even smaller works:

vastai search offers 'gpu_ram>=8 num_gpus=1 reliability>0.9' -o 'dph_total' --raw

Present options and confirm cost with user.

---

Phase 4: Launch the Instance

Use the pre-built Foundry Docker image. No onstart-cmd needed — all checkpoints including Enhanced MPNN are pre-installed.

vastai create instance <OFFER_ID> \
  --image <FOUNDRY_DOCKER_IMAGE> \
  --disk 32 \
  --ssh \
  --label 'foundry-mpnn-<timestamp>'

Capture instance ID from new_contract.

---

Phase 5: Wait for Instance Ready

Poll every 10 seconds for up to 5 minutes:

for i in $(seq 1 30); do
  STATUS=$(vastai show instance <ID> --raw 2>/dev/null | jq -r '.actual_status // .status // "unknown"')
  echo "Attempt $i: status=$STATUS"
  if [ "$STATUS" = "running" ]; then break; fi
  if [ "$STATUS" = "exited" ] || [ "$STATUS" = "offline" ]; then break; fi
  sleep 10
done

Get SSH info and test:

vastai ssh-url <ID>
sleep 20
ssh -i $SSH_KEY -o StrictHostKeyChecking=no -o ConnectTimeout=10 -p <PORT> root@<HOST> 'echo connected'

Verify foundry is available:

ssh -i $SSH_KEY -o StrictHostKeyChecking=no -p <PORT> root@<HOST> 'which mpnn && foundry list-installed'

---

Phase 6: Upload Input Files

Upload structure files and config:

tar czf - -C <LOCAL_DIR> <INPUT_FILES...> | \
  ssh -i $SSH_KEY -o StrictHostKeyChecking=no -p <PORT> root@<HOST> \
  'mkdir -p /workspace/mpnn_job && tar xzf - -C /workspace/mpnn_job/'

---

Phase 7: Run MPNN Design

From JSON config:

ssh -i $SSH_KEY -o StrictHostKeyChecking=no -p <PORT> root@<HOST> \
  'cd /workspace/mpnn_job && mpnn --config_json /workspace/mpnn_job/config.json \
    > /workspace/mpnn_job/mpnn.log 2>&1'

From CLI — Enhanced MPNN:

ssh -i $SSH_KEY -o StrictHostKeyChecking=no -p <PORT> root@<HOST> \
  'mpnn \
    --model_type ligand_mpnn \
    --checkpoint_path /root/.foundry/checkpoints/enhanced_mpnn_step_80000.pt \
    --structure_path /workspace/mpnn_job/<STRUCTURE> \
    --out_directory /workspace/mpnn_job/outputs/ \
    --batch_size 1 \
    --number_of_batches 8 \
    --temperature 0.1 \
    > /workspace/mpnn_job/mpnn.log 2>&1'

From CLI — ProteinMPNN:

ssh -i $SSH_KEY -o StrictHostKeyChecking=no -p <PORT> root@<HOST> \
  'mpnn \
    --model_type protein_mpnn \
    --checkpoint_path /root/.foundry/checkpoints/proteinmpnn_v_48_020.pt \
    --is_legacy_weights True \
    --structure_path /workspace/mpnn_job/<STRUCTURE> \
    --out_directory /workspace/mpnn_job/outputs/ \
    --batch_size 1 \
    --number_of_batches 10 \
    --temperature 0.1 \
    > /workspace/mpnn_job/mpnn.log 2>&1'

MPNN is fast — most jobs complete in under a minute. Run directly (no nohup needed for typical jobs).

---

Phase 8: Retrieve Results

ssh -i $SSH_KEY -o StrictHostKeyChecking=no -p <PORT> root@<HOST> \
  'tar czf - -C /workspace/mpnn_job outputs/' | \
  tar xzf - -C <LOCAL_RESULTS_DIR>

Check what was generated:

ls -la <LOCAL_RESULTS_DIR>/outputs/

Show designed sequences to user:

cat <LOCAL_RESULTS_DIR>/outputs/*.fasta

---

Phase 9: Cleanup

Always destroy the instance:

vastai destroy instance <ID>

---

Phase 10: Report

Provide the user with:

• Model variant used (ProteinMPNN / LigandMPNN / SolubleMPNN / Enhanced MPNN)
• GPU used and cost/hr
• Total runtime and estimated cost
• Number of sequences designed
• Output file locations (FASTA sequences, CIF structures)
• Sample of designed sequences (first few from FASTA)
• Confirmation that the instance was destroyed

---

Error Recovery

• Instance won't start: Destroy, pick next cheapest offer, retry (up to 3 attempts)
• MPNN crashes: Check structure file format (must be valid PDB/CIF), check model_type matches checkpoint
• No output: Verify is_legacy_weights=True for original weights (not needed for Enhanced MPNN)
• SSH won't connect: Wait 30s, retry 3 times
• Any unrecoverable error: ALWAYS destroy the instance to stop billing

Safety Rules

1. Always confirm cost with the user before creating an instance
2. Always destroy the instance when done or on failure
3. Never leave instances running
4. Track the instance ID throughout for cleanup
5. If you lose track, run vastai show instances to find by label