From 977f929c4cde34292fb4dd32e6e6e04f52178b74 Mon Sep 17 00:00:00 2001 From: Mijanur Palash Date: Thu, 28 May 2026 22:20:42 -0400 Subject: [PATCH] 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. --- .../.env.example | 24 + .../openvla-agentcore-orchestrator/.gitignore | 16 + .../openvla-agentcore-orchestrator/README.md | 220 ++++++++ .../openvla.Dockerfile | 69 +++ .../requirements.txt | 6 + .../slurm/finetune_openvla.sbatch | 138 +++++ .../slurm_mcp_server.py | 438 +++++++++++++++ .../vla_training_agent.py | 509 ++++++++++++++++++ 8 files changed, 1420 insertions(+) create mode 100644 02-use-cases/openvla-agentcore-orchestrator/.env.example create mode 100644 02-use-cases/openvla-agentcore-orchestrator/.gitignore create mode 100644 02-use-cases/openvla-agentcore-orchestrator/README.md create mode 100644 02-use-cases/openvla-agentcore-orchestrator/openvla.Dockerfile create mode 100644 02-use-cases/openvla-agentcore-orchestrator/requirements.txt create mode 100644 02-use-cases/openvla-agentcore-orchestrator/slurm/finetune_openvla.sbatch create mode 100644 02-use-cases/openvla-agentcore-orchestrator/slurm_mcp_server.py create mode 100644 02-use-cases/openvla-agentcore-orchestrator/vla_training_agent.py diff --git a/02-use-cases/openvla-agentcore-orchestrator/.env.example b/02-use-cases/openvla-agentcore-orchestrator/.env.example new file mode 100644 index 000000000..286398bba --- /dev/null +++ b/02-use-cases/openvla-agentcore-orchestrator/.env.example @@ -0,0 +1,24 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# +# Copy this file to .env and fill in your cluster details. +# Both slurm_mcp_server.py and vla_training_agent.py read from these. + +# Cluster SSH connection +CLUSTER_HOST=your-cluster-login-node.example.com +CLUSTER_USER=your-username +SSH_KEY_PATH=~/.ssh/id_rsa + +# Workspace paths on the cluster (FSx) +VLA_HOME=/fsx/$USER/vla + +# Training configuration +MAX_STEPS=500 +LEARNING_RATE=5e-4 +BATCH_SIZE=16 +LORA_RANK=32 +DATASET_NAME=libero_10_no_noops + +# Weights & Biases (set to 'online' to enable logging) +WANDB_MODE=disabled +WANDB_ENTITY= diff --git a/02-use-cases/openvla-agentcore-orchestrator/.gitignore b/02-use-cases/openvla-agentcore-orchestrator/.gitignore new file mode 100644 index 000000000..40c09c067 --- /dev/null +++ b/02-use-cases/openvla-agentcore-orchestrator/.gitignore @@ -0,0 +1,16 @@ +# Python +__pycache__/ +*.pyc +*.pyo +.venv/ +venv/ + +# Environment +.env + +# OS +.DS_Store +Thumbs.db + +# Logs +*.log diff --git a/02-use-cases/openvla-agentcore-orchestrator/README.md b/02-use-cases/openvla-agentcore-orchestrator/README.md new file mode 100644 index 000000000..3a5c7f52e --- /dev/null +++ b/02-use-cases/openvla-agentcore-orchestrator/README.md @@ -0,0 +1,220 @@ +# OpenVLA AgentCore Orchestrator — Autonomous Training on HyperPod Slurm + +## Overview + +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. + +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. + +### Use case details + +| Information | Details | +|---------------------|----------------------------------------------------------------------| +| Use case type | autonomous orchestration | +| Agent type | Single-agent with MCP tools | +| Use case components | MCP tools (Slurm operations over SSH), anomaly detection, recovery | +| Use case vertical | MLOps / ML Training | +| Example complexity | Intermediate | +| SDK used | MCP (Model Context Protocol) | + +### Architecture + +``` +┌─────────────────────────────────────────────────────────────┐ +│ Local Machine / AgentCore Runtime │ +│ │ +│ ┌──────────────────┐ ┌─────────────────────────────┐ │ +│ │ VLA Training │ │ Slurm MCP Server │ │ +│ │ Agent │──────▶│ (6 tools: submit, status, │ │ +│ │ │ MCP │ logs, cancel, info, metrics)│ │ +│ └──────────────────┘ └──────────────┬──────────────┘ │ +└─────────────────────────────────────────────┼───────────────┘ + │ SSH + ▼ + ┌──────────────────────────────┐ + │ HyperPod Slurm Cluster │ + │ (P5/P5en nodes, FSx, EFA) │ + │ │ + │ sbatch → torchrun → GPUs │ + └──────────────────────────────┘ +``` + +### Key Features + +- **MCP-based Slurm integration**: 6 tools (submit, status, logs, cancel, cluster info, metrics) exposed via the Model Context Protocol +- **Anomaly detection**: Monitors loss curves for divergence, stalls, and NaN values +- **Automatic recovery**: Cancels failing jobs, adjusts learning rate, resubmits +- **Input validation**: Job IDs validated against injection; SSH with configurable keys +- **Container-based training**: Pinned Dockerfile with EFA/NCCL for reproducibility +- **HyperPod auto-resume**: Handles node failures via `srun --auto-resume` +- **Environment-driven config**: All paths/credentials via `.env` file — no hardcoded values + +## Prerequisites + +| Requirement | Description | +|-------------|-------------| +| Python 3.10+ | For the MCP server (`mcp` library requires 3.10+) | +| SageMaker HyperPod cluster | Slurm scheduler with P5/P5en GPU nodes | +| SSH access | Key-based SSH to the cluster login node | +| FSx for Lustre | Shared filesystem mounted at `/fsx` | +| Docker + Enroot/Pyxis | For building and running the training container | + +## File Structure + +``` +openvla-agentcore-orchestrator/ +├── README.md # This file +├── .env.example # Configuration template +├── .gitignore +├── requirements.txt # MCP server dependency +├── slurm_mcp_server.py # MCP server (6 Slurm tools over SSH) +├── vla_training_agent.py # Autonomous training agent +├── openvla.Dockerfile # Training container (pinned deps) +└── slurm/ + └── finetune_openvla.sbatch # Slurm batch script +``` + +## Setup + +### 1. Clone this repository + +```bash +git clone https://github.com/awslabs/agentcore-samples.git +cd agentcore-samples/02-use-cases/openvla-agentcore-orchestrator +``` + +### 2. Install dependencies + +```bash +python3 -m venv .venv +source .venv/bin/activate +pip install -r requirements.txt +``` + +### 3. Configure environment + +```bash +cp .env.example .env +# Edit .env with your cluster details: +# CLUSTER_HOST, CLUSTER_USER, SSH_KEY_PATH, VLA_HOME +``` + +### 4. Prepare the training environment on the cluster + +Build the container and set up the workspace: + +```bash +# Build container image +docker build -t openvla-finetune -f openvla.Dockerfile . + +# Import as Enroot squashfs (on cluster or transfer after build) +enroot import -o /fsx/$USER/vla/openvla-finetune.sqsh dockerd://openvla-finetune:latest + +# Create workspace +ssh $CLUSTER_USER@$CLUSTER_HOST "mkdir -p /fsx/$USER/vla/{models,data,checkpoints,logs,scripts}" +``` + +### 5. Download model and data + +On the cluster: + +```bash +# Download OpenVLA-7B weights +python3 -c " +from huggingface_hub import snapshot_download +snapshot_download('openvla/openvla-7b-prismatic', local_dir='/fsx/$USER/vla/models/openvla-7b') +" + +# Download LIBERO dataset (RLDS format from Hugging Face Hub — not in stock TFDS catalog) +git lfs install +git clone https://huggingface.co/datasets/openvla/modified_libero_rlds /fsx/$USER/vla/data/libero_rlds +``` + +### 6. Copy the sbatch script to the cluster + +```bash +scp slurm/finetune_openvla.sbatch $CLUSTER_USER@$CLUSTER_HOST:/fsx/$USER/vla/scripts/ +``` + +## Usage + +### Option A: Run the MCP Server (for any MCP client) + +Start the MCP server as a long-running process. Any MCP-compatible client (Claude Desktop, Cursor, a custom agent, or AgentCore Runtime) can connect: + +```bash +python3 slurm_mcp_server.py +``` + +The server exposes 6 tools over stdio: +- `slurm_submit` — Submit a training job +- `slurm_status` — Check job state +- `slurm_logs` — Read stdout/stderr +- `slurm_cancel` — Cancel a job +- `slurm_info` — Cluster/partition info +- `slurm_metrics` — Parse training metrics from logs + +### Option B: Run the Autonomous Agent (full loop) + +The agent submits a job and monitors it through completion, handling anomalies: + +```bash +python3 vla_training_agent.py +``` + +This will: +1. Submit `finetune_openvla.sbatch` via sbatch +2. Poll every 30 seconds for status and metrics +3. Detect divergence/stall/NaN and recover (cancel + resubmit with adjusted LR) +4. Report final results + +Expected runtime: ~12–15 minutes for 500 steps on 1 node (8x H200). + +### Option C: Monitor an Existing Job + +If you already have a running/completed job: + +```bash +python3 vla_training_agent.py --monitor-job --check-interval 10 --max-checks 5 +``` + +## Configuration + +All settings are read from environment variables (or `.env` file): + +| Variable | Default | Description | +|----------|---------|-------------| +| `CLUSTER_HOST` | *(required)* | Cluster login node hostname | +| `CLUSTER_USER` | *(required)* | SSH username | +| `SSH_KEY_PATH` | `~/.ssh/id_rsa` | Path to SSH private key | +| `VLA_HOME` | `/fsx/$USER/vla` | Base workspace on FSx | +| `MAX_STEPS` | `500` | Training steps | +| `LEARNING_RATE` | `5e-4` | Initial learning rate | +| `BATCH_SIZE` | `16` | Per-device batch size | +| `LORA_RANK` | `32` | LoRA adapter rank | + +## Production Deployment with AgentCore + +In production, this same MCP toolset can be hosted on Amazon Bedrock AgentCore Runtime: + +- **AgentCore Gateway** exposes the Slurm MCP tools with authentication, rate limiting, and audit logging +- **AgentCore Memory** stores experiment history across sessions +- **AgentCore Policy** enforces cost guardrails (max GPU hours, budget caps) +- **AgentCore Observability** logs all agent decisions for review + +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. + +## Security Notes + +- SSH uses `StrictHostKeyChecking=no` for demo convenience. In production, pin host keys or use a jump host with managed credentials. +- Job IDs are validated (digits only) to prevent shell injection through the MCP tools. +- The `.env` file contains credentials — never commit it to version control. + +## References + +- [Amazon Bedrock AgentCore](https://docs.aws.amazon.com/bedrock/latest/userguide/agentcore.html) +- [Model Context Protocol (MCP)](https://modelcontextprotocol.io/) +- [OpenVLA](https://github.com/openvla/openvla) — Vision-Language-Action model +- [LIBERO](https://libero-project.github.io/) — Robotic manipulation benchmark +- [SageMaker HyperPod](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod.html) +- [SRE Agent (sibling sample)](../SRE-agent/) — Similar architecture for infrastructure operations diff --git a/02-use-cases/openvla-agentcore-orchestrator/openvla.Dockerfile b/02-use-cases/openvla-agentcore-orchestrator/openvla.Dockerfile new file mode 100644 index 000000000..9a5e139b0 --- /dev/null +++ b/02-use-cases/openvla-agentcore-orchestrator/openvla.Dockerfile @@ -0,0 +1,69 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# +# OpenVLA fine-tuning container for HyperPod Slurm (P5/P5en) +# Build: +# docker build -t openvla-finetune -f openvla.Dockerfile . +# Import for Pyxis/Enroot: +# enroot import -o /fsx/$USER/openvla-finetune.sqsh dockerd://openvla-finetune:latest + +# Base image: AWS HPC container with CUDA, EFA, NCCL, and aws-ofi-nccl pre-installed. +# Same base used by sibling test cases (nanoVLM). Provides: +# CUDA 12.8.1, NCCL 2.27.7, EFA 1.43.2 (libfabric), aws-ofi-nccl 1.16.3 +# Note: No published tag meets the CI floor (NCCL>=2.28, CUDA>=13.0) yet. +# PyTorch wheels bundle their own CUDA runtime (cu124), so training itself is unaffected. +FROM public.ecr.aws/hpc-cloud/nccl-tests:cuda12.8.1-efa1.43.2-ofiv1.16.3-ncclv2.27.7-1-testsv2.16.9 + +RUN apt-get update && apt-get install -y --no-install-recommends \ + git \ + git-lfs \ + nvtop \ + && rm -rf /var/lib/apt/lists/* + +# --------------------------------------------------------------------------- +# Python training dependencies — pinned to the exact versions validated on +# HyperPod P5en (8x H200), job 5631, ~10 min for 500 LoRA steps. +# --------------------------------------------------------------------------- +RUN pip install --no-cache-dir \ + torch==2.6.0 \ + torchvision==0.21.0 \ + transformers==4.44.2 \ + peft==0.13.2 \ + accelerate==1.2.1 \ + "datasets>=2.14.0,<3.0.0" \ + tensorflow-datasets==4.9.3 \ + "tensorflow>=2.15.0,<2.18.0" \ + "tensorflow-graphics==2021.12.3" \ + "huggingface-hub>=0.20.0,<1.0.0" \ + "wandb>=0.16.0,<1.0.0" \ + "pillow>=10.0.0,<11.0.0" \ + "scipy>=1.11.0,<2.0.0" \ + "einops>=0.7.0,<1.0.0" \ + timm==0.9.10 \ + draccus==0.8.0 \ + sentencepiece==0.1.99 \ + tokenizers==0.19.1 + +# dlimp (data loading for RLDS) — no-deps to avoid pulling conflicting versions +RUN pip install --no-cache-dir --no-deps \ + git+https://github.com/kvablack/dlimp.git@5edaa4691567873d495633f2708982b42edf1972 + +# dlimp's fork used by OpenVLA for RLDS dataset loading +RUN pip install --no-cache-dir --no-deps \ + git+https://github.com/moojink/dlimp_openvla.git@040105d256bd28866cc6620621a3d5f7b6b91b46 + +# --------------------------------------------------------------------------- +# Clone OpenVLA and install in editable mode WITHOUT dependencies. +# This ensures our explicit pins above are not overwritten by OpenVLA's +# pyproject.toml (which hard-pins torch==2.2.0, transformers==4.40.1, etc.). +# --------------------------------------------------------------------------- +ARG OPENVLA_COMMIT=c8f03f48af692657d3060c19588038c7220e9af9 +RUN git clone https://github.com/openvla/openvla.git /openvla \ + && cd /openvla && git checkout ${OPENVLA_COMMIT} + +RUN cd /openvla && pip install --no-cache-dir --no-deps -e . + +# Symlink python for convenience +RUN ln -sf /usr/bin/python3 /usr/bin/python + +WORKDIR /openvla diff --git a/02-use-cases/openvla-agentcore-orchestrator/requirements.txt b/02-use-cases/openvla-agentcore-orchestrator/requirements.txt new file mode 100644 index 000000000..c9023384b --- /dev/null +++ b/02-use-cases/openvla-agentcore-orchestrator/requirements.txt @@ -0,0 +1,6 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# +# MCP server dependency (for slurm_mcp_server.py) +# The training agent (vla_training_agent.py) uses only stdlib. +mcp>=1.0.0,<2.0.0 diff --git a/02-use-cases/openvla-agentcore-orchestrator/slurm/finetune_openvla.sbatch b/02-use-cases/openvla-agentcore-orchestrator/slurm/finetune_openvla.sbatch new file mode 100644 index 000000000..f2e42329d --- /dev/null +++ b/02-use-cases/openvla-agentcore-orchestrator/slurm/finetune_openvla.sbatch @@ -0,0 +1,138 @@ +#!/bin/bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# +# OpenVLA LoRA Fine-Tuning on LIBERO (1 node, 8x GPUs) +# Tested on: HyperPod P5en (8x H200), ~10 min for 500 steps +# +# Prerequisites: +# 1. Build container: docker build -t openvla-finetune -f openvla.Dockerfile . +# 2. Import: enroot import -o $VLA_HOME/openvla-finetune.sqsh dockerd://openvla-finetune:latest +# 3. Download model + data (see README) +# +# Usage: +# export VLA_HOME=/fsx/$USER/vla +# cd $VLA_HOME # logs are written to $VLA_HOME/logs/ relative to cwd +# sbatch finetune_openvla.sbatch + +#SBATCH --job-name=openvla-finetune +#SBATCH --partition=p5en +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --mem=0 +#SBATCH --time=01:55:00 +#SBATCH --exclusive +#SBATCH --output=logs/%j.out +#SBATCH --error=logs/%j.err + +set -euo pipefail + +# --------------------------------------------------------------------------- +# Configuration — set VLA_HOME or edit paths below +# --------------------------------------------------------------------------- +VLA_HOME="${VLA_HOME:-/fsx/$USER/vla}" + +MODEL_DIR="${VLA_HOME}/models/openvla-7b" +DATA_DIR="${VLA_HOME}/data/libero_rlds" +CKPT_DIR="${VLA_HOME}/checkpoints/run_${SLURM_JOB_ID}" +CONTAINER_IMAGE="${VLA_HOME}/openvla-finetune.sqsh" + +# Training hyperparameters (override via environment before sbatch) +DATASET_NAME="${DATASET_NAME:-libero_10_no_noops}" +MAX_STEPS="${MAX_STEPS:-500}" +SAVE_STEPS="${SAVE_STEPS:-500}" +LEARNING_RATE="${LEARNING_RATE:-5e-4}" +BATCH_SIZE="${BATCH_SIZE:-16}" +LORA_RANK="${LORA_RANK:-32}" +WANDB_MODE="${WANDB_MODE:-disabled}" + +GPUS_PER_NODE=8 + +# --------------------------------------------------------------------------- +# Environment (EFA + NCCL for multi-node readiness) +# --------------------------------------------------------------------------- +export NCCL_DEBUG=INFO +export FI_PROVIDER=efa +export FI_EFA_SET_CUDA_SYNC_MEMOPS=0 +export NCCL_SOCKET_IFNAME=^lo,docker0 +export TORCHDYNAMO_DISABLE=1 + +# --------------------------------------------------------------------------- +# Container mounts +# --------------------------------------------------------------------------- +declare -a SRUN_ARGS=( + --container-image "${CONTAINER_IMAGE}" + --container-mounts "/fsx:/fsx" +) + +# --------------------------------------------------------------------------- +# Torchrun args +# For multi-node, replace --standalone with: +# --nnodes=$SLURM_JOB_NUM_NODES --rdzv_id=$SLURM_JOB_ID +# --rdzv_backend=c10d --rdzv_endpoint=$SLURMD_NODENAME:29500 +# --------------------------------------------------------------------------- +declare -a TORCHRUN_ARGS=( + --standalone + --nnodes=1 + --nproc_per_node=${GPUS_PER_NODE} +) + +# --------------------------------------------------------------------------- +# Training args +# --------------------------------------------------------------------------- +declare -a TRAINING_ARGS=( + --vla_path "${MODEL_DIR}" + --data_root_dir "${DATA_DIR}" + --dataset_name "${DATASET_NAME}" + --run_root_dir "${CKPT_DIR}" + --adapter_tmp_dir "${CKPT_DIR}/adapter_tmp" + --lora_rank "${LORA_RANK}" + --batch_size "${BATCH_SIZE}" + --grad_accumulation_steps 1 + --learning_rate "${LEARNING_RATE}" + --max_steps "${MAX_STEPS}" + --save_steps "${SAVE_STEPS}" + --image_aug true + --wandb_project openvla-finetune +) + +# Only pass --wandb_entity if set (avoids passing empty string) +if [ -n "${WANDB_ENTITY:-}" ]; then + TRAINING_ARGS+=(--wandb_entity "${WANDB_ENTITY}") +fi + +# --------------------------------------------------------------------------- +# Create output dirs +# --------------------------------------------------------------------------- +mkdir -p "${CKPT_DIR}" "logs" + +echo "=========================================" +echo "Job ID: ${SLURM_JOB_ID}" +echo "Container: ${CONTAINER_IMAGE}" +echo "Model: ${MODEL_DIR}" +echo "Data: ${DATA_DIR}" +echo "Output: ${CKPT_DIR}" +echo "=========================================" + +# --------------------------------------------------------------------------- +# HyperPod auto-resume +# --------------------------------------------------------------------------- +declare -a AUTO_RESUME_ARGS=() +if [ -d "/opt/sagemaker_cluster" ]; then + echo "Detected HyperPod cluster — enabling auto-resume" + AUTO_RESUME_ARGS=("--auto-resume=1") +fi + +# --------------------------------------------------------------------------- +# Launch +# --------------------------------------------------------------------------- +export WANDB_MODE +srun "${AUTO_RESUME_ARGS[@]}" -l "${SRUN_ARGS[@]}" \ + torchrun "${TORCHRUN_ARGS[@]}" \ + /openvla/vla-scripts/finetune.py "${TRAINING_ARGS[@]}" + +echo "=========================================" +echo "Fine-tuning complete!" +echo "Checkpoints at: ${CKPT_DIR}" +echo "=========================================" diff --git a/02-use-cases/openvla-agentcore-orchestrator/slurm_mcp_server.py b/02-use-cases/openvla-agentcore-orchestrator/slurm_mcp_server.py new file mode 100644 index 000000000..d26bdd515 --- /dev/null +++ b/02-use-cases/openvla-agentcore-orchestrator/slurm_mcp_server.py @@ -0,0 +1,438 @@ +#!/usr/bin/env python3 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +""" +Slurm MCP Server — Exposes Slurm cluster operations as MCP tools. + +An LLM agent uses these tools to submit, monitor, and manage training jobs +on a HyperPod Slurm cluster via SSH. + +Tools: + - slurm_submit: Submit a batch job (sbatch) + - slurm_status: Get job status (squeue/sacct) + - slurm_logs: Read job stdout/stderr + - slurm_cancel: Cancel a running job (scancel) + - slurm_info: Get cluster/partition info (sinfo) + - slurm_metrics: Parse training metrics from log files + +Configuration: + Set environment variables (or use a .env file): + CLUSTER_HOST, CLUSTER_USER, SSH_KEY_PATH, VLA_HOME + +Requires: Python 3.10+ (mcp library dependency) +""" + +import asyncio +import json +import os +import re +from dataclasses import dataclass + +from mcp.server import Server +from mcp.server.stdio import stdio_server +from mcp.types import Tool, TextContent + +# ============================================================================= +# Configuration (from environment variables) +# ============================================================================= + + +def _load_dotenv(): + """Load .env file if present (minimal implementation, no dependency).""" + env_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env") + if os.path.exists(env_path): + with open(env_path) as f: + for line in f: + line = line.strip() + if line and not line.startswith("#") and "=" in line: + key, _, value = line.partition("=") + os.environ.setdefault(key.strip(), value.strip()) + + +_load_dotenv() + + +@dataclass +class ClusterConfig: + host: str = os.environ.get("CLUSTER_HOST", "") + user: str = os.environ.get("CLUSTER_USER", "") + ssh_key: str = os.environ.get("SSH_KEY_PATH", "~/.ssh/id_rsa") + slurm_bin: str = os.environ.get("SLURM_BIN", "/opt/slurm/bin") + work_dir: str = os.environ.get("VLA_HOME", "") + + def validate(self): + missing = [] + if not self.host: + missing.append("CLUSTER_HOST") + if not self.user: + missing.append("CLUSTER_USER") + if not self.work_dir: + missing.append("VLA_HOME") + if missing: + raise EnvironmentError( + f"Missing required environment variables: {', '.join(missing)}. Set them in .env or export them." + ) + + +config = ClusterConfig() + + +# ============================================================================= +# Input Validation +# ============================================================================= + + +def _validate_job_id(job_id: str) -> str | None: + """Validate job_id is a numeric Slurm job ID. Returns error message or None.""" + if not re.fullmatch(r"\d+", str(job_id)): + return f"ERROR: job_id must be numeric, got: {job_id!r}" + return None + + +# ============================================================================= +# SSH Helper +# ============================================================================= + + +async def ssh_exec(cmd: str, timeout: int = 30) -> tuple[str, str, int]: + """Execute a command on the cluster via SSH. + + NOTE: StrictHostKeyChecking=no is used for convenience in dev/test. + For production, use known_hosts verification. + """ + ssh_key = os.path.expanduser(config.ssh_key) + ssh_cmd = [ + "ssh", + "-i", + ssh_key, + "-o", + "StrictHostKeyChecking=no", + "-o", + f"ConnectTimeout={timeout}", + f"{config.user}@{config.host}", + f"export PATH={config.slurm_bin}:$PATH && {cmd}", + ] + proc = await asyncio.create_subprocess_exec( + *ssh_cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE + ) + stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout + 10) + return stdout.decode(), stderr.decode(), proc.returncode + + +# ============================================================================= +# MCP Server +# ============================================================================= + +app = Server("slurm-mcp") + + +@app.list_tools() +async def list_tools(): + return [ + Tool( + name="slurm_submit", + description="Submit a Slurm batch job. Returns the job ID.", + inputSchema={ + "type": "object", + "properties": { + "script_path": { + "type": "string", + "description": "Absolute path to the .sbatch script on the cluster", + }, + "overrides": { + "type": "object", + "description": 'Optional sbatch overrides (e.g. {"time": "01:00:00", "job-name": "test"})', + "additionalProperties": {"type": "string"}, + }, + }, + "required": ["script_path"], + }, + ), + Tool( + name="slurm_status", + description="Get status of jobs. Returns job ID, state, time, node, reason.", + inputSchema={ + "type": "object", + "properties": { + "job_id": { + "type": "string", + "description": "Specific job ID to check (optional — defaults to all user jobs)", + } + }, + }, + ), + Tool( + name="slurm_logs", + description="Read stdout/stderr from a training job. Returns last N lines.", + inputSchema={ + "type": "object", + "properties": { + "job_id": {"type": "string", "description": "The Slurm job ID"}, + "stream": { + "type": "string", + "enum": ["stdout", "stderr", "both"], + "description": "Which log stream to read (default: both)", + }, + "tail_lines": {"type": "integer", "description": "Number of lines from the end (default: 50)"}, + }, + "required": ["job_id"], + }, + ), + Tool( + name="slurm_cancel", + description="Cancel a running or pending Slurm job.", + inputSchema={ + "type": "object", + "properties": { + "job_id": {"type": "string", "description": "The Slurm job ID to cancel"}, + "reason": {"type": "string", "description": "Reason for cancellation (logged)"}, + }, + "required": ["job_id"], + }, + ), + Tool( + name="slurm_info", + description="Get cluster partition info — available nodes, GPUs, state.", + inputSchema={ + "type": "object", + "properties": { + "partition": {"type": "string", "description": "Partition name (optional — defaults to all)"} + }, + }, + ), + Tool( + name="slurm_metrics", + description="Parse training metrics from job logs. Extracts loss, learning rate, steps, GPU utilization.", + inputSchema={ + "type": "object", + "properties": {"job_id": {"type": "string", "description": "The Slurm job ID to parse metrics from"}}, + "required": ["job_id"], + }, + ), + ] + + +@app.call_tool() +async def call_tool(name: str, arguments: dict): + config.validate() + + if name == "slurm_submit": + return await _slurm_submit(arguments) + elif name == "slurm_status": + return await _slurm_status(arguments) + elif name == "slurm_logs": + return await _slurm_logs(arguments) + elif name == "slurm_cancel": + return await _slurm_cancel(arguments) + elif name == "slurm_info": + return await _slurm_info(arguments) + elif name == "slurm_metrics": + return await _slurm_metrics(arguments) + else: + return [TextContent(type="text", text=f"Unknown tool: {name}")] + + +# ============================================================================= +# Tool Implementations +# ============================================================================= + + +async def _slurm_submit(args: dict): + script = args["script_path"] + overrides = args.get("overrides", {}) + + override_flags = " ".join(f"--{k}={v}" for k, v in overrides.items()) + cmd = f"sbatch {override_flags} {script}".strip() + + stdout, stderr, rc = await ssh_exec(cmd) + if rc != 0: + return [TextContent(type="text", text=f"ERROR: sbatch failed\n{stderr}")] + + match = re.search(r"Submitted batch job (\d+)", stdout) + job_id = match.group(1) if match else "unknown" + + return [ + TextContent( + type="text", + text=json.dumps( + {"status": "submitted", "job_id": job_id, "script": script, "overrides": overrides}, indent=2 + ), + ) + ] + + +async def _slurm_status(args: dict): + job_id = args.get("job_id", "") + + if job_id: + err = _validate_job_id(job_id) + if err: + return [TextContent(type="text", text=err)] + cmd = f"squeue -j {job_id} --format='%i|%j|%T|%M|%N|%r' --noheader" + else: + cmd = f"squeue -u {config.user} --format='%i|%j|%T|%M|%N|%r' --noheader" + + stdout, stderr, rc = await ssh_exec(cmd) + + if not stdout.strip() or rc != 0: + if not job_id: + return [TextContent(type="text", text="No jobs found")] + # Job may have completed — check sacct + cmd2 = f"sacct -j {job_id} --format=JobID,JobName,State,Elapsed,ExitCode --noheader -P" + stdout2, _, _ = await ssh_exec(cmd2) + if stdout2.strip(): + return [TextContent(type="text", text=f"Job completed:\n{stdout2}")] + return [TextContent(type="text", text="No jobs found")] + + jobs = [] + for line in stdout.strip().split("\n"): + parts = line.split("|") + if len(parts) >= 6: + jobs.append( + { + "job_id": parts[0].strip(), + "name": parts[1].strip(), + "state": parts[2].strip(), + "time": parts[3].strip(), + "node": parts[4].strip(), + "reason": parts[5].strip(), + } + ) + + return [TextContent(type="text", text=json.dumps(jobs, indent=2))] + + +async def _slurm_logs(args: dict): + job_id = args["job_id"] + err = _validate_job_id(job_id) + if err: + return [TextContent(type="text", text=err)] + stream = args.get("stream", "both") + tail_lines = args.get("tail_lines", 50) + + log_dir = f"{config.work_dir}/logs" + result = {} + + if stream in ("stdout", "both"): + cmd = f"tail -n {tail_lines} {log_dir}/finetune_{job_id}.out 2>/dev/null || echo 'NO_FILE'" + stdout, _, _ = await ssh_exec(cmd) + result["stdout"] = stdout + + if stream in ("stderr", "both"): + cmd = f"tail -n {tail_lines} {log_dir}/finetune_{job_id}.err 2>/dev/null || echo 'NO_FILE'" + stdout, _, _ = await ssh_exec(cmd) + result["stderr"] = stdout + + return [TextContent(type="text", text=json.dumps(result, indent=2))] + + +async def _slurm_cancel(args: dict): + job_id = args["job_id"] + err = _validate_job_id(job_id) + if err: + return [TextContent(type="text", text=err)] + reason = args.get("reason", "cancelled by agent") + + cmd = f"scancel {job_id}" + stdout, stderr, rc = await ssh_exec(cmd) + + if rc != 0: + return [TextContent(type="text", text=f"ERROR: scancel failed\n{stderr}")] + + return [ + TextContent(type="text", text=json.dumps({"status": "cancelled", "job_id": job_id, "reason": reason}, indent=2)) + ] + + +async def _slurm_info(args: dict): + partition = args.get("partition", "") + + if partition: + cmd = f"sinfo -p {partition} --format='%P|%a|%D|%T|%G|%l' --noheader" + else: + cmd = "sinfo --format='%P|%a|%D|%T|%G|%l' --noheader" + + stdout, stderr, rc = await ssh_exec(cmd) + + # Also get GPU utilization + cmd2 = "squeue --format='%i|%u|%T|%b' --noheader" + stdout2, _, _ = await ssh_exec(cmd2) + + return [ + TextContent( + type="text", text=json.dumps({"partitions": stdout.strip(), "running_jobs": stdout2.strip()}, indent=2) + ) + ] + + +async def _slurm_metrics(args: dict): + job_id = args["job_id"] + err = _validate_job_id(job_id) + if err: + return [TextContent(type="text", text=err)] + log_file = f"{config.work_dir}/logs/finetune_{job_id}.out" + + cmd = f"""python3 -c " +import re, json + +metrics = {{'job_id': '{job_id}', 'steps': [], 'latest': {{}}}} + +try: + with open('{log_file}') as f: + lines = f.readlines() + + for line in lines: + # Match [METRICS] structured lines first + if '[METRICS]' in line: + step_m = re.search(r'step=(\\d+)', line) + loss_m = re.search(r'loss=([0-9.]+)', line) + if step_m and loss_m: + entry = {{'step': int(step_m.group(1)), 'loss': float(loss_m.group(1))}} + metrics['steps'].append(entry) + continue + + # Fallback: generic loss/step patterns + loss_match = re.search(r'loss[=:\\s]+([0-9.]+)', line, re.IGNORECASE) + step_match = re.search(r'step[=:\\s]+(\\d+)', line, re.IGNORECASE) + lr_match = re.search(r'lr[=:\\s]+([0-9.e-]+)', line, re.IGNORECASE) + + if loss_match and step_match: + entry = {{'step': int(step_match.group(1)), 'loss': float(loss_match.group(1))}} + if lr_match: + entry['lr'] = float(lr_match.group(1)) + metrics['steps'].append(entry) + + if metrics['steps']: + metrics['latest'] = metrics['steps'][-1] + metrics['total_steps_logged'] = len(metrics['steps']) + first = metrics['steps'][0]['loss'] + last = metrics['steps'][-1]['loss'] + metrics['loss_trend'] = 'decreasing' if last < first else 'increasing' if last > first else 'flat' + +except FileNotFoundError: + metrics['error'] = 'Log file not found (job may still be starting)' + +print(json.dumps(metrics, indent=2)) +" +""" + stdout, stderr, rc = await ssh_exec(cmd, timeout=15) + + if rc != 0: + return [TextContent(type="text", text=f"ERROR parsing metrics: {stderr}")] + + return [TextContent(type="text", text=stdout)] + + +# ============================================================================= +# Main +# ============================================================================= + + +async def main(): + config.validate() + async with stdio_server() as (read_stream, write_stream): + await app.run(read_stream, write_stream, app.create_initialization_options()) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/02-use-cases/openvla-agentcore-orchestrator/vla_training_agent.py b/02-use-cases/openvla-agentcore-orchestrator/vla_training_agent.py new file mode 100644 index 000000000..5fc6b68a7 --- /dev/null +++ b/02-use-cases/openvla-agentcore-orchestrator/vla_training_agent.py @@ -0,0 +1,509 @@ +#!/usr/bin/env python3 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +""" +VLA Training Agent — Autonomous agent that orchestrates VLA fine-tuning via Slurm. + +This agent demonstrates how Amazon Bedrock AgentCore can manage ML training: + 1. Submits a training job to a HyperPod Slurm cluster + 2. Monitors loss curve periodically + 3. Detects issues (divergence, stalling, NaN) + 4. Takes corrective action (adjust LR, restart with different params) + 5. Reports results + +In production, this would run on AgentCore Runtime with: + - Gateway providing the Slurm MCP tools + - Memory storing experiment history + - Policy enforcing cost guardrails + - Observability logging all decisions + +For local testing, it SSHs directly to the cluster. + +Configuration: + Set environment variables (or use a .env file): + CLUSTER_HOST, CLUSTER_USER, SSH_KEY_PATH, VLA_HOME, MAX_STEPS +""" + +import json +import math +import os +import re +import subprocess +import time +from dataclasses import dataclass, field +from typing import Optional + + +# ============================================================================= +# Configuration (from environment) +# ============================================================================= + + +def _load_dotenv(): + """Load .env file if present.""" + env_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env") + if os.path.exists(env_path): + with open(env_path) as f: + for line in f: + line = line.strip() + if line and not line.startswith("#") and "=" in line: + key, _, value = line.partition("=") + os.environ.setdefault(key.strip(), value.strip()) + + +_load_dotenv() + + +def _get_env(key: str, default: str = "") -> str: + """Get environment variable with validation.""" + return os.environ.get(key, default) + + +# Cluster connection +CLUSTER_HOST = _get_env("CLUSTER_HOST") +CLUSTER_USER = _get_env("CLUSTER_USER") +SSH_KEY_PATH = os.path.expanduser(_get_env("SSH_KEY_PATH", "~/.ssh/id_rsa")) +VLA_HOME = _get_env("VLA_HOME") +MAX_STEPS = int(_get_env("MAX_STEPS", "500")) + + +@dataclass +class TrainingConfig: + """What the agent knows about the training setup.""" + + model_path: str = f"{VLA_HOME}/models/openvla-7b" + data_dir: str = f"{VLA_HOME}/data/libero_rlds" + dataset_name: str = _get_env("DATASET_NAME", "libero_10_no_noops") + script_path: str = f"{VLA_HOME}/scripts/finetune_openvla.sbatch" + max_steps: int = MAX_STEPS + + # Guardrails (would come from AgentCore Policy in production) + max_gpu_hours: float = 16.0 + max_retries: int = 3 + loss_divergence_threshold: float = 10.0 + loss_stall_patience: int = 5 + + +@dataclass +class ExperimentState: + """Agent's working memory (would be AgentCore Memory in production).""" + + current_job_id: Optional[str] = None + run_history: list = field(default_factory=list) + loss_history: list = field(default_factory=list) + decisions: list = field(default_factory=list) + retries: int = 0 + status: str = "idle" + + +# ============================================================================= +# Slurm Interface (SSH — in production, replaced by AgentCore Gateway + MCP) +# ============================================================================= + + +def _validate_config(): + """Ensure required config is set.""" + missing = [] + if not CLUSTER_HOST: + missing.append("CLUSTER_HOST") + if not CLUSTER_USER: + missing.append("CLUSTER_USER") + if not VLA_HOME: + missing.append("VLA_HOME") + if missing: + raise EnvironmentError( + f"Missing required environment variables: {', '.join(missing)}. " + f"Copy .env.example to .env and fill in your values." + ) + + +# NOTE: StrictHostKeyChecking=no is used for convenience in dev/test. +# For production, use known_hosts verification. +SSH_CMD_PREFIX = [ + "ssh", + "-i", + SSH_KEY_PATH, + "-o", + "StrictHostKeyChecking=no", + "-o", + "ConnectTimeout=10", + f"{CLUSTER_USER}@{CLUSTER_HOST}", +] + + +def ssh_exec(cmd: str) -> str: + """Execute command on cluster via SSH.""" + full_cmd = SSH_CMD_PREFIX + [f"export PATH=/opt/slurm/bin:$PATH && {cmd}"] + result = subprocess.run(full_cmd, capture_output=True, text=True, timeout=30) + return result.stdout + result.stderr + + +def submit_job(script_path: str, overrides: dict = None, env_vars: dict = None) -> str: + """Submit a Slurm job. Returns job ID. + + Args: + script_path: Path to the .sbatch script on the cluster. + overrides: sbatch option overrides (e.g. {"time": "01:00:00"}). + env_vars: Environment variables to export for the job (e.g. {"LEARNING_RATE": "1e-4"}). + """ + override_str = " ".join(f"--{k}={v}" for k, v in (overrides or {}).items()) + export_str = "" + if env_vars: + pairs = ",".join(f"{k}={v}" for k, v in env_vars.items()) + export_str = f"--export=ALL,{pairs}" + output = ssh_exec(f"sbatch {override_str} {export_str} {script_path}".strip()) + match = re.search(r"Submitted batch job (\d+)", output) + return match.group(1) if match else None + + +def get_job_status(job_id: str) -> dict: + """Get job state via squeue, falling back to sacct for completed jobs.""" + output = ssh_exec(f"squeue -j {job_id} --format='%T' --noheader") + state = output.strip() + if not state or "error" in state.lower() or "invalid" in state.lower(): + output = ssh_exec(f"sacct -j {job_id} --format=State --noheader -P") + state = output.strip().split("\n")[0] if output.strip() else "UNKNOWN" + return {"job_id": job_id, "state": state} + + +def get_training_metrics(job_id: str) -> dict: + """Parse training progress from log files. + + Sources (checked in priority order): + 1. [METRICS] lines in stdout — structured loss/acc/l1 per step + 2. tqdm progress bars in stderr — step counts and throughput + 3. Checkpoint existence — final artifact verification + """ + metrics = {"job_id": job_id, "steps": [], "losses": []} + max_steps = MAX_STEPS + + out_file = f"{VLA_HOME}/logs/finetune_{job_id}.out" + err_file = f"{VLA_HOME}/logs/finetune_{job_id}.err" + + # --- Source 1: Structured [METRICS] lines from stdout --- + cmd_metrics = f"""grep '\\[METRICS\\]' {out_file} 2>/dev/null""" + metrics_output = ssh_exec(cmd_metrics) + + for line in metrics_output.strip().split("\n"): + if "[METRICS]" not in line: + continue + step_m = re.search(r"step=(\d+)", line) + loss_m = re.search(r"loss=([0-9.]+)", line) + acc_m = re.search(r"acc=([0-9.]+)", line) + l1_m = re.search(r"l1=([0-9.]+)", line) + if step_m and loss_m: + entry = {"step": int(step_m.group(1)), "loss": float(loss_m.group(1))} + if acc_m: + entry["acc"] = float(acc_m.group(1)) + if l1_m: + entry["l1"] = float(l1_m.group(1)) + metrics["steps"].append(entry) + metrics["losses"].append(entry["loss"]) + + # --- Source 2: tqdm progress from stderr (fallback) --- + if not metrics["steps"]: + cmd = f"""grep -oP '\\d+/{max_steps} \\[\\d+:\\d+' {err_file} 2>/dev/null | sort -t/ -k1 -n -u | tail -20""" + output = ssh_exec(cmd) + for line in output.strip().split("\n"): + step_match = re.search(rf"(\d+)/{max_steps}", line) + if step_match: + step = int(step_match.group(1)) + metrics["steps"].append({"step": step}) + + # Throughput from last tqdm line + cmd2 = f"""grep -oP '[0-9.]+it/s' {err_file} 2>/dev/null | tail -1""" + throughput = ssh_exec(cmd2).strip() + if throughput: + metrics["throughput"] = throughput + + # --- Source 3: Status messages and checkpoint --- + cmd3 = f"""grep -E 'Saving|Max step|Fine-tuning complete|trainable params' {out_file} 2>/dev/null | sort -u | tail -5""" + status_output = ssh_exec(cmd3) + if status_output.strip(): + metrics["status_messages"] = [line.strip()[:80] for line in status_output.strip().split("\n") if line.strip()] + + ckpt_dir = f"{VLA_HOME}/checkpoints/run_{job_id}" + cmd4 = f"""ls -lh {ckpt_dir}/*/*.safetensors 2>/dev/null | wc -l; du -sh {ckpt_dir} 2>/dev/null""" + ckpt_output = ssh_exec(cmd4) + parts = ckpt_output.strip().split("\n") + if parts: + metrics["checkpoint_shards"] = parts[0].strip() + if len(parts) > 1: + metrics["checkpoint_size"] = parts[1].strip().split()[0] if parts[1].strip() else "unknown" + + # --- Summary --- + if metrics["steps"]: + metrics["latest_step"] = metrics["steps"][-1]["step"] + metrics["total_steps"] = max_steps + metrics["progress_pct"] = f"{metrics['latest_step'] / max_steps * 100:.0f}%" + + if metrics["losses"]: + metrics["latest_loss"] = metrics["losses"][-1] + metrics["min_loss"] = min(metrics["losses"]) + metrics["max_loss"] = max(metrics["losses"]) + + return metrics + + +def cancel_job(job_id: str) -> bool: + """Cancel a job.""" + output = ssh_exec(f"scancel {job_id}") + return "error" not in output.lower() + + +# ============================================================================= +# Agent Logic +# ============================================================================= + + +class VLATrainingAgent: + """ + Agent that manages VLA training autonomously. + + Decision loop: + 1. Submit job if none running + 2. Monitor metrics every N seconds + 3. Detect anomalies (divergence, NaN, stall) + 4. Take corrective action + 5. Report final results + """ + + def __init__(self, config: TrainingConfig = None): + self.config = config or TrainingConfig() + self.state = ExperimentState() + self.check_interval = 30 + + def log_decision(self, action: str, reason: str, details: dict = None): + """Record agent decisions (AgentCore Observability in production).""" + entry = { + "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), + "action": action, + "reason": reason, + "details": details or {}, + } + self.state.decisions.append(entry) + print(f"[AGENT] {action}: {reason}") + + def submit_training(self, overrides: dict = None, env_vars: dict = None) -> str: + """Submit a training job.""" + self.log_decision("SUBMIT", "Starting training run", {"overrides": overrides, "env_vars": env_vars}) + job_id = submit_job(self.config.script_path, overrides, env_vars) + + if job_id: + self.state.current_job_id = job_id + self.state.status = "training" + self.state.run_history.append({"job_id": job_id, "overrides": overrides}) + self.log_decision("SUBMITTED", f"Job {job_id} queued") + else: + self.log_decision("ERROR", "Failed to submit job") + self.state.status = "failed" + + return job_id + + def check_health(self) -> str: + """Check training health. Returns: healthy, diverged, stalled, failed, complete.""" + if not self.state.current_job_id: + return "no_job" + + status = get_job_status(self.state.current_job_id) + job_state = status["state"] + + if "COMPLETED" in job_state: + self.state.status = "complete" + return "complete" + + if "FAILED" in job_state or "CANCELLED" in job_state: + return "failed" + + if "PENDING" in job_state: + return "pending" + + if "RUNNING" not in job_state: + return "unknown" + + # Parse metrics for anomaly detection + metrics = get_training_metrics(self.state.current_job_id) + + if not metrics.get("steps"): + return "starting" + + latest_loss = metrics.get("latest_loss") + if latest_loss is not None: + self.state.loss_history.append(latest_loss) + + if latest_loss > self.config.loss_divergence_threshold: + self.log_decision("DETECT", f"Loss diverged: {latest_loss:.4f}", metrics) + return "diverged" + + if math.isnan(latest_loss): + self.log_decision("DETECT", "NaN loss detected", metrics) + return "diverged" + + if len(self.state.loss_history) >= self.config.loss_stall_patience: + recent = self.state.loss_history[-self.config.loss_stall_patience :] + if all(abs(recent[i] - recent[i - 1]) < 0.001 for i in range(1, len(recent))): + self.log_decision("DETECT", f"Loss stalled at {latest_loss:.4f}", metrics) + return "stalled" + + return "healthy" + + def recover(self, issue: str): + """Take corrective action based on the issue detected.""" + self.state.retries += 1 + + if self.state.retries > self.config.max_retries: + self.log_decision("ABORT", f"Max retries ({self.config.max_retries}) exceeded") + self.state.status = "failed" + return + + if self.state.current_job_id: + cancel_job(self.state.current_job_id) + self.log_decision("CANCEL", f"Cancelled job {self.state.current_job_id}") + + if issue == "diverged": + current_lr = 5e-4 / (5**self.state.retries) + self.log_decision("RECOVER", f"Reducing LR to {current_lr:.2e}") + self.submit_training( + overrides={"comment": f"retry_{self.state.retries}_lr_{current_lr}"}, + env_vars={"LEARNING_RATE": f"{current_lr:.2e}"}, + ) + + elif issue == "stalled": + current_lr = 5e-4 / (2**self.state.retries) + self.log_decision("RECOVER", f"Loss stalled — restarting with LR {current_lr:.2e}") + self.submit_training( + overrides={"comment": f"retry_{self.state.retries}_stall_fix"}, + env_vars={"LEARNING_RATE": f"{current_lr:.2e}"}, + ) + + elif issue == "failed": + self.log_decision("RECOVER", "Job failed — retrying") + self.submit_training({"comment": f"retry_{self.state.retries}_after_failure"}) + + def run(self, max_checks: int = 100): + """Main agent loop.""" + print("=" * 60) + print("VLA Training Agent — Starting") + print(f"Model: {self.config.model_path}") + print(f"Dataset: {self.config.dataset_name}") + print("=" * 60) + + self.submit_training() + + if self.state.status == "failed": + return self.report() + + for check_num in range(max_checks): + time.sleep(self.check_interval) + + health = self.check_health() + print(f"[CHECK {check_num + 1}] Health: {health} | Job: {self.state.current_job_id}") + + if health == "complete": + self.log_decision("COMPLETE", "Training finished successfully") + break + + elif health in ("diverged", "stalled", "failed"): + self.recover(health) + if self.state.status == "failed": + break + + elif health == "pending": + print(" Job still pending...") + + elif health == "starting": + print(" Job initializing (no metrics yet)...") + + return self.report() + + def report(self) -> dict: + """Generate final experiment report.""" + report = { + "status": self.state.status, + "total_runs": len(self.state.run_history), + "total_retries": self.state.retries, + "decisions": self.state.decisions, + "final_loss": self.state.loss_history[-1] if self.state.loss_history else None, + "loss_history_summary": { + "first": self.state.loss_history[0] if self.state.loss_history else None, + "last": self.state.loss_history[-1] if self.state.loss_history else None, + "min": min(self.state.loss_history) if self.state.loss_history else None, + "count": len(self.state.loss_history), + }, + } + + print("\n" + "=" * 60) + print("EXPERIMENT REPORT") + print("=" * 60) + print(json.dumps(report, indent=2, default=str)) + + return report + + +# ============================================================================= +# Entry point +# ============================================================================= + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="VLA Training Agent") + parser.add_argument("--check-interval", type=int, default=30, help="Seconds between health checks") + parser.add_argument("--max-checks", type=int, default=100, help="Maximum monitoring iterations") + parser.add_argument("--submit-only", action="store_true", help="Just submit, don't monitor") + parser.add_argument("--monitor-job", type=str, help="Monitor existing job ID instead of submitting") + args = parser.parse_args() + + _validate_config() + + agent = VLATrainingAgent() + agent.check_interval = args.check_interval + + if args.monitor_job: + agent.state.current_job_id = args.monitor_job + agent.state.status = "training" + print("=" * 60) + print(f"VLA Training Agent — Monitoring Job {args.monitor_job}") + print("=" * 60) + for i in range(args.max_checks): + time.sleep(agent.check_interval) + health = agent.check_health() + print(f"[CHECK {i + 1}] Health: {health} | Job: {agent.state.current_job_id}") + if health in ("complete", "failed"): + break + elif health in ("diverged", "stalled"): + agent.recover(health) + if agent.state.status == "failed": + break + # Final metrics + metrics = get_training_metrics(agent.state.current_job_id) + if metrics.get("steps"): + summary_parts = [f"Parsed {len(metrics['steps'])} unique steps"] + if metrics.get("throughput"): + summary_parts.append(f"Throughput: {metrics['throughput']}") + if metrics.get("checkpoint_size"): + summary_parts.append(f"Checkpoint: {metrics['checkpoint_size']}") + if metrics.get("progress_pct"): + summary_parts.append(f"Progress: {metrics['progress_pct']}") + agent.log_decision("METRICS", " | ".join(summary_parts)) + if metrics.get("losses"): + agent.log_decision( + "LOSS_CURVE", + f"Steps: {len(metrics['losses'])} | " + f"First: {metrics['losses'][0]:.4f} | " + f"Last: {metrics['losses'][-1]:.4f} | " + f"Min: {min(metrics['losses']):.4f} | " + f"Max: {max(metrics['losses']):.4f}", + ) + last_5 = metrics["losses"][-5:] + trend = " -> ".join(f"{val:.4f}" for val in last_5) + agent.log_decision("LOSS_TREND", trend) + if metrics.get("status_messages"): + for msg in metrics["status_messages"][-3:]: + agent.log_decision("STATUS", msg) + agent.report() + elif args.submit_only: + agent.submit_training() + print(f"Job submitted: {agent.state.current_job_id}") + else: + agent.run(max_checks=args.max_checks)