Skip to content

Commit 977f929

Browse files
mijanur132mijanur132
committed
Add OpenVLA AgentCore Orchestrator: autonomous Slurm training orchestration
Demonstrates AgentCore managing GPU training jobs on HyperPod Slurm: - MCP server exposing 6 Slurm tools (submit, status, logs, cancel, info, metrics) - Autonomous agent with anomaly detection (divergence, stall, NaN) and recovery - OpenVLA-7B LoRA fine-tuning on LIBERO as concrete workload - Container-based (Pyxis/Enroot) with pinned deps for reproducibility - All config via environment variables, no hardcoded paths/credentials Tested: 500 steps, ~10 min on 1x P5en node (8x H200), 15 GB checkpoint output.
1 parent 5bd6da1 commit 977f929

8 files changed

Lines changed: 1420 additions & 0 deletions

File tree

‎02-use-cases/openvla-agentcore-orchestrator/.env.example‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: MIT-0
3+
#
4+
# Copy this file to .env and fill in your cluster details.
5+
# Both slurm_mcp_server.py and vla_training_agent.py read from these.
6+
7+
# Cluster SSH connection
8+
CLUSTER_HOST=your-cluster-login-node.example.com
9+
CLUSTER_USER=your-username
10+
SSH_KEY_PATH=~/.ssh/id_rsa
11+
12+
# Workspace paths on the cluster (FSx)
13+
VLA_HOME=/fsx/$USER/vla
14+
15+
# Training configuration
16+
MAX_STEPS=500
17+
LEARNING_RATE=5e-4
18+
BATCH_SIZE=16
19+
LORA_RANK=32
20+
DATASET_NAME=libero_10_no_noops
21+
22+
# Weights & Biases (set to 'online' to enable logging)
23+
WANDB_MODE=disabled
24+
WANDB_ENTITY=

‎02-use-cases/openvla-agentcore-orchestrator/.gitignore‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# Python
2+
__pycache__/
3+
*.pyc
4+
*.pyo
5+
.venv/
6+
venv/
7+
8+
# Environment
9+
.env
10+
11+
# OS
12+
.DS_Store
13+
Thumbs.db
14+
15+
# Logs
16+
*.log

‎02-use-cases/openvla-agentcore-orchestrator/README.md‎

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
# OpenVLA AgentCore Orchestrator — Autonomous Training on HyperPod Slurm
2+
3+
## Overview
4+
5+
The ML Training Agent demonstrates how Amazon Bedrock AgentCore can autonomously manage GPU training jobs on a SageMaker HyperPod Slurm cluster. It uses a Slurm MCP server to expose cluster operations as tools, and an autonomous agent that submits OpenVLA fine-tuning jobs, monitors training metrics, detects anomalies (loss divergence, stalls, NaN), and takes corrective action — all without human intervention.
6+
7+
The concrete workload is [OpenVLA-7B](https://github.com/openvla/openvla) LoRA fine-tuning on the LIBERO robotics benchmark, but the MCP server and agent pattern are workload-agnostic.
8+
9+
### Use case details
10+
11+
| Information | Details |
12+
|---------------------|----------------------------------------------------------------------|
13+
| Use case type | autonomous orchestration |
14+
| Agent type | Single-agent with MCP tools |
15+
| Use case components | MCP tools (Slurm operations over SSH), anomaly detection, recovery |
16+
| Use case vertical | MLOps / ML Training |
17+
| Example complexity | Intermediate |
18+
| SDK used | MCP (Model Context Protocol) |
19+
20+
### Architecture
21+
22+
```
23+
┌─────────────────────────────────────────────────────────────┐
24+
│ Local Machine / AgentCore Runtime │
25+
│ │
26+
│ ┌──────────────────┐ ┌─────────────────────────────┐ │
27+
│ │ VLA Training │ │ Slurm MCP Server │ │
28+
│ │ Agent │──────▶│ (6 tools: submit, status, │ │
29+
│ │ │ MCP │ logs, cancel, info, metrics)│ │
30+
│ └──────────────────┘ └──────────────┬──────────────┘ │
31+
└─────────────────────────────────────────────┼───────────────┘
32+
│ SSH
33+
▼
34+
┌──────────────────────────────┐
35+
│ HyperPod Slurm Cluster │
36+
│ (P5/P5en nodes, FSx, EFA) │
37+
│ │
38+
│ sbatch → torchrun → GPUs │
39+
└──────────────────────────────┘
40+
```
41+
42+
### Key Features
43+
44+
- **MCP-based Slurm integration**: 6 tools (submit, status, logs, cancel, cluster info, metrics) exposed via the Model Context Protocol
45+
- **Anomaly detection**: Monitors loss curves for divergence, stalls, and NaN values
46+
- **Automatic recovery**: Cancels failing jobs, adjusts learning rate, resubmits
47+
- **Input validation**: Job IDs validated against injection; SSH with configurable keys
48+
- **Container-based training**: Pinned Dockerfile with EFA/NCCL for reproducibility
49+
- **HyperPod auto-resume**: Handles node failures via `srun --auto-resume`
50+
- **Environment-driven config**: All paths/credentials via `.env` file — no hardcoded values
51+
52+
## Prerequisites
53+
54+
| Requirement | Description |
55+
|-------------|-------------|
56+
| Python 3.10+ | For the MCP server (`mcp` library requires 3.10+) |
57+
| SageMaker HyperPod cluster | Slurm scheduler with P5/P5en GPU nodes |
58+
| SSH access | Key-based SSH to the cluster login node |
59+
| FSx for Lustre | Shared filesystem mounted at `/fsx` |
60+
| Docker + Enroot/Pyxis | For building and running the training container |
61+
62+
## File Structure
63+
64+
```
65+
openvla-agentcore-orchestrator/
66+
├── README.md # This file
67+
├── .env.example # Configuration template
68+
├── .gitignore
69+
├── requirements.txt # MCP server dependency
70+
├── slurm_mcp_server.py # MCP server (6 Slurm tools over SSH)
71+
├── vla_training_agent.py # Autonomous training agent
72+
├── openvla.Dockerfile # Training container (pinned deps)
73+
└── slurm/
74+
└── finetune_openvla.sbatch # Slurm batch script
75+
```
76+
77+
## Setup
78+
79+
### 1. Clone this repository
80+
81+
```bash
82+
git clone https://github.com/awslabs/agentcore-samples.git
83+
cd agentcore-samples/02-use-cases/openvla-agentcore-orchestrator
84+
```
85+
86+
### 2. Install dependencies
87+
88+
```bash
89+
python3 -m venv .venv
90+
source .venv/bin/activate
91+
pip install -r requirements.txt
92+
```
93+
94+
### 3. Configure environment
95+
96+
```bash
97+
cp .env.example .env
98+
# Edit .env with your cluster details:
99+
# CLUSTER_HOST, CLUSTER_USER, SSH_KEY_PATH, VLA_HOME
100+
```
101+
102+
### 4. Prepare the training environment on the cluster
103+
104+
Build the container and set up the workspace:
105+
106+
```bash
107+
# Build container image
108+
docker build -t openvla-finetune -f openvla.Dockerfile .
109+
110+
# Import as Enroot squashfs (on cluster or transfer after build)
111+
enroot import -o /fsx/$USER/vla/openvla-finetune.sqsh dockerd://openvla-finetune:latest
112+
113+
# Create workspace
114+
ssh $CLUSTER_USER@$CLUSTER_HOST "mkdir -p /fsx/$USER/vla/{models,data,checkpoints,logs,scripts}"
115+
```
116+
117+
### 5. Download model and data
118+
119+
On the cluster:
120+
121+
```bash
122+
# Download OpenVLA-7B weights
123+
python3 -c "
124+
from huggingface_hub import snapshot_download
125+
snapshot_download('openvla/openvla-7b-prismatic', local_dir='/fsx/$USER/vla/models/openvla-7b')
126+
"
127+
128+
# Download LIBERO dataset (RLDS format from Hugging Face Hub — not in stock TFDS catalog)
129+
git lfs install
130+
git clone https://huggingface.co/datasets/openvla/modified_libero_rlds /fsx/$USER/vla/data/libero_rlds
131+
```
132+
133+
### 6. Copy the sbatch script to the cluster
134+
135+
```bash
136+
scp slurm/finetune_openvla.sbatch $CLUSTER_USER@$CLUSTER_HOST:/fsx/$USER/vla/scripts/
137+
```
138+
139+
## Usage
140+
141+
### Option A: Run the MCP Server (for any MCP client)
142+
143+
Start the MCP server as a long-running process. Any MCP-compatible client (Claude Desktop, Cursor, a custom agent, or AgentCore Runtime) can connect:
144+
145+
```bash
146+
python3 slurm_mcp_server.py
147+
```
148+
149+
The server exposes 6 tools over stdio:
150+
- `slurm_submit` — Submit a training job
151+
- `slurm_status` — Check job state
152+
- `slurm_logs` — Read stdout/stderr
153+
- `slurm_cancel` — Cancel a job
154+
- `slurm_info` — Cluster/partition info
155+
- `slurm_metrics` — Parse training metrics from logs
156+
157+
### Option B: Run the Autonomous Agent (full loop)
158+
159+
The agent submits a job and monitors it through completion, handling anomalies:
160+
161+
```bash
162+
python3 vla_training_agent.py
163+
```
164+
165+
This will:
166+
1. Submit `finetune_openvla.sbatch` via sbatch
167+
2. Poll every 30 seconds for status and metrics
168+
3. Detect divergence/stall/NaN and recover (cancel + resubmit with adjusted LR)
169+
4. Report final results
170+
171+
Expected runtime: ~12–15 minutes for 500 steps on 1 node (8x H200).
172+
173+
### Option C: Monitor an Existing Job
174+
175+
If you already have a running/completed job:
176+
177+
```bash
178+
python3 vla_training_agent.py --monitor-job <JOB_ID> --check-interval 10 --max-checks 5
179+
```
180+
181+
## Configuration
182+
183+
All settings are read from environment variables (or `.env` file):
184+
185+
| Variable | Default | Description |
186+
|----------|---------|-------------|
187+
| `CLUSTER_HOST` | *(required)* | Cluster login node hostname |
188+
| `CLUSTER_USER` | *(required)* | SSH username |
189+
| `SSH_KEY_PATH` | `~/.ssh/id_rsa` | Path to SSH private key |
190+
| `VLA_HOME` | `/fsx/$USER/vla` | Base workspace on FSx |
191+
| `MAX_STEPS` | `500` | Training steps |
192+
| `LEARNING_RATE` | `5e-4` | Initial learning rate |
193+
| `BATCH_SIZE` | `16` | Per-device batch size |
194+
| `LORA_RANK` | `32` | LoRA adapter rank |
195+
196+
## Production Deployment with AgentCore
197+
198+
In production, this same MCP toolset can be hosted on Amazon Bedrock AgentCore Runtime:
199+
200+
- **AgentCore Gateway** exposes the Slurm MCP tools with authentication, rate limiting, and audit logging
201+
- **AgentCore Memory** stores experiment history across sessions
202+
- **AgentCore Policy** enforces cost guardrails (max GPU hours, budget caps)
203+
- **AgentCore Observability** logs all agent decisions for review
204+
205+
The local agent (`vla_training_agent.py`) demonstrates the orchestration logic; in production, an LLM model on Bedrock invokes the same tools through the Gateway.
206+
207+
## Security Notes
208+
209+
- SSH uses `StrictHostKeyChecking=no` for demo convenience. In production, pin host keys or use a jump host with managed credentials.
210+
- Job IDs are validated (digits only) to prevent shell injection through the MCP tools.
211+
- The `.env` file contains credentials — never commit it to version control.
212+
213+
## References
214+
215+
- [Amazon Bedrock AgentCore](https://docs.aws.amazon.com/bedrock/latest/userguide/agentcore.html)
216+
- [Model Context Protocol (MCP)](https://modelcontextprotocol.io/)
217+
- [OpenVLA](https://github.com/openvla/openvla) — Vision-Language-Action model
218+
- [LIBERO](https://libero-project.github.io/) — Robotic manipulation benchmark
219+
- [SageMaker HyperPod](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod.html)
220+
- [SRE Agent (sibling sample)](../SRE-agent/) — Similar architecture for infrastructure operations

‎02-use-cases/openvla-agentcore-orchestrator/openvla.Dockerfile‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: MIT-0
3+
#
4+
# OpenVLA fine-tuning container for HyperPod Slurm (P5/P5en)
5+
# Build:
6+
# docker build -t openvla-finetune -f openvla.Dockerfile .
7+
# Import for Pyxis/Enroot:
8+
# enroot import -o /fsx/$USER/openvla-finetune.sqsh dockerd://openvla-finetune:latest
9+
10+
# Base image: AWS HPC container with CUDA, EFA, NCCL, and aws-ofi-nccl pre-installed.
11+
# Same base used by sibling test cases (nanoVLM). Provides:
12+
# CUDA 12.8.1, NCCL 2.27.7, EFA 1.43.2 (libfabric), aws-ofi-nccl 1.16.3
13+
# Note: No published tag meets the CI floor (NCCL>=2.28, CUDA>=13.0) yet.
14+
# PyTorch wheels bundle their own CUDA runtime (cu124), so training itself is unaffected.
15+
FROM public.ecr.aws/hpc-cloud/nccl-tests:cuda12.8.1-efa1.43.2-ofiv1.16.3-ncclv2.27.7-1-testsv2.16.9
16+
17+
RUN apt-get update && apt-get install -y --no-install-recommends \
18+
git \
19+
git-lfs \
20+
nvtop \
21+
&& rm -rf /var/lib/apt/lists/*
22+
23+
# ---------------------------------------------------------------------------
24+
# Python training dependencies — pinned to the exact versions validated on
25+
# HyperPod P5en (8x H200), job 5631, ~10 min for 500 LoRA steps.
26+
# ---------------------------------------------------------------------------
27+
RUN pip install --no-cache-dir \
28+
torch==2.6.0 \
29+
torchvision==0.21.0 \
30+
transformers==4.44.2 \
31+
peft==0.13.2 \
32+
accelerate==1.2.1 \
33+
"datasets>=2.14.0,<3.0.0" \
34+
tensorflow-datasets==4.9.3 \
35+
"tensorflow>=2.15.0,<2.18.0" \
36+
"tensorflow-graphics==2021.12.3" \
37+
"huggingface-hub>=0.20.0,<1.0.0" \
38+
"wandb>=0.16.0,<1.0.0" \
39+
"pillow>=10.0.0,<11.0.0" \
40+
"scipy>=1.11.0,<2.0.0" \
41+
"einops>=0.7.0,<1.0.0" \
42+
timm==0.9.10 \
43+
draccus==0.8.0 \
44+
sentencepiece==0.1.99 \
45+
tokenizers==0.19.1
46+
47+
# dlimp (data loading for RLDS) — no-deps to avoid pulling conflicting versions
48+
RUN pip install --no-cache-dir --no-deps \
49+
git+https://github.com/kvablack/dlimp.git@5edaa4691567873d495633f2708982b42edf1972
50+
51+
# dlimp's fork used by OpenVLA for RLDS dataset loading
52+
RUN pip install --no-cache-dir --no-deps \
53+
git+https://github.com/moojink/dlimp_openvla.git@040105d256bd28866cc6620621a3d5f7b6b91b46
54+
55+
# ---------------------------------------------------------------------------
56+
# Clone OpenVLA and install in editable mode WITHOUT dependencies.
57+
# This ensures our explicit pins above are not overwritten by OpenVLA's
58+
# pyproject.toml (which hard-pins torch==2.2.0, transformers==4.40.1, etc.).
59+
# ---------------------------------------------------------------------------
60+
ARG OPENVLA_COMMIT=c8f03f48af692657d3060c19588038c7220e9af9
61+
RUN git clone https://github.com/openvla/openvla.git /openvla \
62+
&& cd /openvla && git checkout ${OPENVLA_COMMIT}
63+
64+
RUN cd /openvla && pip install --no-cache-dir --no-deps -e .
65+
66+
# Symlink python for convenience
67+
RUN ln -sf /usr/bin/python3 /usr/bin/python
68+
69+
WORKDIR /openvla

‎02-use-cases/openvla-agentcore-orchestrator/requirements.txt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: MIT-0
3+
#
4+
# MCP server dependency (for slurm_mcp_server.py)
5+
# The training agent (vla_training_agent.py) uses only stdlib.
6+
mcp>=1.0.0,<2.0.0

0 commit comments

Comments
 (0)