Deploy Stack

# Exclude infrastructure and build artifacts
infra/
cdk.out/
__pycache__/
*.pyc
*.pyo
*.pyd
.Python
env/
venv/
.venv
pip-log.txt
pip-delete-this-directory.txt
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.log
.git/
.mypy_cache/
.pytest_cache/
.hypothesis/
**/.DS_Store

#!/bin/bash

# Build script for GR00T Fine-tuning Docker image

set -Eeuo pipefail

echo "=========================================="
echo "Building GR00T Fine-tuning Docker Image"
echo "=========================================="

# Default values
IMAGE_NAME="gr00t-finetune"
TAG="latest"
DOCKERFILE="Dockerfile"
PUSH_IMAGE=false
TEST_IMAGE=false
USE_STABLE=true

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -t|--tag)
            TAG="$2"
            shift 2
            ;;
        -n|--name)
            IMAGE_NAME="$2"
            shift 2
            ;;
        --latest)
            USE_STABLE=false
            shift
            ;;
        --push)
            PUSH_IMAGE=true
            shift
            ;;
        --test)
            TEST_IMAGE=true
            shift
            ;;
        -h|--help)
            echo "Usage: $0 [OPTIONS]"
            echo "Options:"
            echo "  -t, --tag TAG         Tag for the fine-tuning image (default: latest)"
            echo "  -n, --name NAME       Name for the fine-tuning image (default: gr00t-finetune)"
            echo "  --latest              Use latest GR00T from main branch (default: stable commit)"
            echo "  --push                Push image to registry after building"
            echo "  --test                Run basic tests after building"
            echo "  -h, --help            Show this help message"
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

FULL_IMAGE_NAME="${IMAGE_NAME}:${TAG}"

echo "Building image: ${FULL_IMAGE_NAME}"
echo "Using Dockerfile: ${DOCKERFILE}"

# Display GR00T version selection
if [[ "${USE_STABLE}" == "true" ]]; then
    echo "GR00T version: STABLE (tested commit from Sep 4, 2025) [default]"
else
    echo "GR00T version: LATEST (main branch, may have breaking changes)"
fi

# Build the fine-tuning image directly from the combined Dockerfile
echo "Building fine-tuning image..."
docker build \
    --build-arg USE_STABLE=${USE_STABLE} \
    -f ${DOCKERFILE} \
    -t ${FULL_IMAGE_NAME} \
    .

echo "Image built successfully: ${FULL_IMAGE_NAME}"

# Run basic tests if requested
if [[ "${TEST_IMAGE}" == "true" ]]; then
    echo "=========================================="
    echo "Running Basic Tests"
    echo "=========================================="
    
    # Test 1: Check if the image runs without errors (dry run)
    echo "Test 1: Checking if image starts correctly..."
    docker run --rm \
        -e HF_TOKEN="dummy" \
        -e HF_DATASET_ID="dummy/dummy" \
        -e HF_MODEL_REPO_ID="dummy/dummy" \
        --entrypoint /bin/bash \
        ${FULL_IMAGE_NAME} \
        -c "echo 'Image starts correctly' && python -c 'import sys; print(f\"Python version: {sys.version}\")' && which huggingface-cli"
    
    # Test 2: Check if finetune script can be imported
    echo "Test 2: Checking if finetune script imports correctly..."
    docker run --rm \
        --entrypoint /bin/bash \
        ${FULL_IMAGE_NAME} \
        -c "cd /workspace && python -c 'from scripts.finetune_gr00t import FinetuneWorkflow; print(\"Finetune script imports successfully\")'"
    
    # Test 3: Check if training script exists and can show help
    echo "Test 3: Checking if training script is accessible..."
    docker run --rm \
        --entrypoint /bin/bash \
        ${FULL_IMAGE_NAME} \
        -c "cd /workspace && python scripts/gr00t_finetune.py --help" | head -10
    
    echo "All tests passed!"
fi

# Push to registry if requested
if [[ "${PUSH_IMAGE}" == "true" ]]; then
    echo "=========================================="
    echo "Pushing to Registry"
    echo "=========================================="
    
    if [[ -z "${DOCKER_REGISTRY}" ]]; then
        echo "Warning: DOCKER_REGISTRY environment variable not set."
        echo "Assuming you want to push to Docker Hub or have already tagged appropriately."
    else
        # Re-tag with registry prefix
        REGISTRY_IMAGE="${DOCKER_REGISTRY}/${FULL_IMAGE_NAME}"
        docker tag ${FULL_IMAGE_NAME} ${REGISTRY_IMAGE}
        FULL_IMAGE_NAME=${REGISTRY_IMAGE}
    fi
    
    echo "Pushing image: ${FULL_IMAGE_NAME}"
    docker push ${FULL_IMAGE_NAME}
    echo "Image pushed successfully!"
fi

echo "=========================================="
echo "Build Complete!"
echo "=========================================="
echo "Image: ${FULL_IMAGE_NAME}"
if [[ "${USE_STABLE}" == "true" ]]; then
    echo "GR00T Version: STABLE (tested commit)"
else
    echo "GR00T Version: LATEST (main branch)"
fi
echo ""
echo "To run locally, create a local directory to simulate EFS mount:"
echo "mkdir -p ~/mock-efs/gr00t/checkpoints"
echo "Then run with a small number of steps for testing:"
echo "docker run --gpus all --network host \\"
echo "  -e MAX_STEPS=100 -e SAVE_STEPS=100 \\"
echo "  -v ~/mock-efs:/mnt/efs \\"
echo "  ${FULL_IMAGE_NAME}"
echo ""
echo "To rebuild with latest GR00T version from main branch:"
echo "  ./build_container.sh --latest" 

# Alternative Dockerfile from the GR00T repo: https://github.com/NVIDIA/Isaac-GR00T/blob/main/Dockerfile
# Optimized Dockerfile for Isaac-GR00T using UV package manager
FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04

# Build argument to control GR00T version (stable vs latest)
ARG USE_STABLE=true
# https://github.com/NVIDIA/Isaac-GR00T/releases/tag/n1.5-release
ARG STABLE_COMMIT=4af2b622892f7dcb5aae5a3fb70bcb02dc217b96

# Setting the frontend to be non-interactive
# This is to avoid any user input required during the installation of packages
ENV DEBIAN_FRONTEND=noninteractive

# System dependencies - consolidated for better layer caching
RUN apt-get update && apt-get install -y --no-install-recommends \
    # Core utilities
    wget curl ca-certificates unzip \
    # Git and version control
    git git-lfs \
    # Build essentials
    build-essential cmake \
    # Media processing
    ffmpeg \
    # OpenCV dependencies
    libopencv-dev libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender-dev \
    # Python development
    python3.10 python3.10-dev python3.10-distutils python3-pip \
    # Utilities for debugging
    vim less htop \
    && rm -rf /var/lib/apt/lists/*

# Install AWS CLI v2 (official AWS method)
RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" && \
    unzip awscliv2.zip && \
    ./aws/install && \
    rm -rf awscliv2.zip aws && \
    rm -rf /usr/local/aws-cli/v2/*/dist/awscli/examples

# Set Python 3.10 as default
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1 && \
    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1

# Install UV package manager (much faster than pip/conda)
COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv

# Configure UV to use system Python
ENV UV_SYSTEM_PYTHON=1
ENV PYTHONPATH=/workspace

# ============================================================================
# Clone Isaac-GR00T repository
# This layer changes when using --latest
# ============================================================================
RUN git clone https://github.com/NVIDIA/Isaac-GR00T.git /workspace && \
    cd /workspace && \
    if [ "${USE_STABLE}" = "true" ]; then \
    echo "Using stable commit: ${STABLE_COMMIT} (default)"; \
    git checkout ${STABLE_COMMIT}; \
    else \
    echo "Using latest version from main branch"; \
    fi && \
    echo "GR00T version info:" && \
    git log -1 --format="%H %ai %s"

# Set working directory
WORKDIR /workspace

# Upgrade pip and setuptools using UV
RUN uv pip install --upgrade pip setuptools wheel

# Install GR00T base dependencies using UV (faster resolution and installation)
RUN uv pip install --no-cache -e .[base]

# Install flash-attention separately (requires build isolation disabled)
RUN pip install --no-build-isolation flash-attn==2.7.1.post4

# Install additional utilities
RUN uv pip install --no-cache notebook gpustat wandb

# Install HuggingFace CLI and additional dependencies if necessary
# RUN pip install huggingface_hub[cli] datasets

# Copy the workflow scripts
COPY finetune_gr00t.py /workspace/scripts/
COPY run_finetune_workflow.sh /workspace/scripts/
RUN chmod +x /workspace/scripts/run_finetune_workflow.sh

# Set environment variables with defaults
ENV DATASET_LOCAL_DIR="/workspace/train"
ENV OUTPUT_DIR="/workspace/checkpoints"

# If there is issue with the latest version, use tested checkpoint as of 09 July 2025 by commenting out the following two lines
# RUN hf download nvidia/GR00T-N1.5-3B --revision 869830fc749c35f34771aa5209f923ac57e4564e --local-dir ./GR00T-N1.5-3B
# ENV BASE_MODEL_PATH="./GR00T-N1.5-3B"

# Create directories using environment variables
RUN mkdir -p ${DATASET_LOCAL_DIR} ${OUTPUT_DIR}

# Setting the Entrypoint and Command to /bin/bash whenever executed
ENTRYPOINT ["/bin/bash"]
# Default command to run the workflow, but can be overridden
CMD ["/workspace/scripts/run_finetune_workflow.sh"]
# CMD ["jupyter", "notebook", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root"]

# Example Environment File for GR00T Fine-tuning Workflow
# Copy this file to .env and fill in your values

# ==========================================
# AUTHENTICATION (optional)
# ==========================================
# HF_TOKEN is only required if pulling dataset from HF or uploading model to HF
HF_TOKEN=<your_huggingface_token_here>

# ==========================================
# DATASET SOURCE SELECTION (all optional)
# The entrypoint resolves in this order: local -> s3 -> hf -> sample
# ==========================================
# DATASET_LOCAL_DIR: Use a local dataset path (e.g., mounted volume)
DATASET_LOCAL_DIR=/workspace/train
# DATASET_S3_URI: Sync dataset from S3 (e.g., s3://my-bucket/datasets/so100)
DATASET_S3_URI=
# HF_DATASET_ID: Download dataset from Hugging Face (e.g., lerobot/datasets)
HF_DATASET_ID=

# ==========================================
# UPLOAD CONFIGURATION (optional)
# UPLOAD_TARGET: hf | s3 | none
# ==========================================
# UPLOAD_TARGET: Set to 'hf' for Hugging Face, 's3' for S3, or 'none' to skip
UPLOAD_TARGET=none
# HF_MODEL_REPO_ID: Required if UPLOAD_TARGET=hf (e.g., your-username/your-model)
HF_MODEL_REPO_ID=
# S3_UPLOAD_URI: Required if UPLOAD_TARGET=s3 (e.g., s3://my-bucket/models/)
S3_UPLOAD_URI=

# ==========================================
# OUTPUT DIRECTORIES (optional)
# By default, the script writes to /mnt/efs/gr00t/checkpoints if mounted
# ==========================================
# OUTPUT_DIR: Override default checkpoint directory
OUTPUT_DIR=

# ==========================================
# BASIC TRAINING PARAMETERS
# ==========================================
MAX_STEPS=6000
SAVE_STEPS=2000
NUM_GPUS=1
BATCH_SIZE=32
LEARNING_RATE=1e-4
WEIGHT_DECAY=1e-5
WARMUP_RATIO=0.05

# ==========================================
# MODEL AND DATA CONFIGURATION
# ==========================================
BASE_MODEL_PATH=nvidia/GR00T-N1.5-3B
DATA_CONFIG=so100_dualcam
VIDEO_BACKEND=torchvision_av
EMBODIMENT_TAG=new_embodiment

# Valid DATA_CONFIG options:
# - so100_dualcam
# - fourier_gr1_arms_only
# - fourier_gr1_arms_waist
# - agibot_genie1_dualcam
# - oxe_droid_single_cam

# Valid EMBODIMENT_TAG options:
# - new_embodiment
# - gr1
# - oxe_droid
# - agibot_genie1

# ==========================================
# FINE-TUNING CONFIGURATION
# ==========================================
TUNE_LLM=false
TUNE_VISUAL=false
TUNE_PROJECTOR=true
TUNE_DIFFUSION_MODEL=true

# ==========================================
# LORA CONFIGURATION
# ==========================================
LORA_RANK=0
LORA_ALPHA=16
LORA_DROPOUT=0.1
LORA_FULL_MODEL=false

# ==========================================
# DATASET BALANCING
# ==========================================
BALANCE_DATASET_WEIGHTS=true
BALANCE_TRAJECTORY_WEIGHTS=true

# ==========================================
# PERFORMANCE AND SYSTEM
# ==========================================
DATALOADER_NUM_WORKERS=8
DATALOADER_PREFETCH_FACTOR=4

# ==========================================
# WORKFLOW CONTROL
# ==========================================
RESUME=false
CLEANUP_DATASET=false
CLEANUP_CHECKPOINTS=false

# ==========================================
# LOGGING AND MONITORING
# ==========================================
REPORT_TO=tensorboard

# Weights & Biases Configuration (set REPORT_TO=wandb to use)
WANDB_API_KEY=your_wandb_api_key
WANDB_PROJECT=gr00t-finetune
WANDB_ENTITY=your_wandb_entity

# ==========================================
# CONFIGURATION EXAMPLES
# ==========================================

# Example 1: Use Local Dataset and Upload to S3
# DATASET_LOCAL_DIR=/workspace/train
# UPLOAD_TARGET=s3
# S3_UPLOAD_URI=s3://my-bucket/models/gr00t-finetuned/

# Example 2: Download from HF and Upload to HF
# HF_DATASET_ID=lerobot/datasets
# UPLOAD_TARGET=hf
# HF_MODEL_REPO_ID=your-username/gr00t-so100-finetuned
# HF_TOKEN=your_actual_hf_token

# Example 3: Use Sample Dataset (Default) and Skip Upload
# UPLOAD_TARGET=none
# (No dataset or upload variables needed)

# Example 4: Use S3 Dataset and Upload to S3
# DATASET_S3_URI=s3://my-bucket/datasets/so100
# UPLOAD_TARGET=s3
# S3_UPLOAD_URI=s3://my-bucket/models/gr00t-finetuned/

# Example 5: LoRA Fine-tuning
# LORA_RANK=32
# LORA_ALPHA=64
# TUNE_PROJECTOR=false
# TUNE_DIFFUSION_MODEL=false
# TUNE_LLM=false
# TUNE_VISUAL=false

# Example 6: Full Fine-tuning
# TUNE_LLM=true
# TUNE_VISUAL=true
# TUNE_PROJECTOR=true
# TUNE_DIFFUSION_MODEL=true
# LORA_RANK=0

# Example 7: Projector + Diffusion Only (Recommended)
# TUNE_LLM=false
# TUNE_VISUAL=false
# TUNE_PROJECTOR=true
# TUNE_DIFFUSION_MODEL=true
# LORA_RANK=0

# Example 8: Resume Training from Checkpoint
# RESUME=true
# OUTPUT_DIR=/workspace/checkpoints  # Should contain previous checkpoints

# Example 9: WandB Logging
# REPORT_TO=wandb
# WANDB_API_KEY=your_actual_wandb_api_key
# WANDB_PROJECT=my-gr00t-project
# WANDB_ENTITY=my-wandb-username 

#!/usr/bin/env python3
"""
Workflow script for fine-tuning GR00T models with configurable parameters.
This script orchestrates the model-specific portion of the workflow:
1. Validate dataset path prepared by the entrypoint shell script
2. Run training
"""

import os
import sys
import logging
import json
from pathlib import Path
import torch
from transformers import TrainingArguments
from torch.distributed.run import main as torchrun

# Import GR00T training and data utilities
from gr00t.data.dataset import LeRobotMixtureDataset, LeRobotSingleDataset
from gr00t.data.schema import EmbodimentTag
from gr00t.experiment.data_config import load_data_config
from gr00t.experiment.runner import TrainRunner
from gr00t.model.gr00t_n1 import GR00T_N1_5
from gr00t.utils.peft import get_lora_model

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler("/workspace/finetune_gr00t.log"),
    ],
)
logger = logging.getLogger(__name__)


class FinetuneWorkflow:
    """Main workflow class for fine-tuning GR00T models."""

    def __init__(self):
        """Initialize the workflow with environment variables."""
        # Dataset directory
        self.dataset_local_dir = os.getenv("DATASET_LOCAL_DIR")

        # Output directories (prefer EFS paths if provided by the Batch Job Definition)
        self.output_dir = os.getenv("OUTPUT_DIR", "/workspace/checkpoints")

        # Training parameters
        self.max_steps = int(os.getenv("MAX_STEPS", "6000"))
        self.save_steps = int(os.getenv("SAVE_STEPS", "2000"))
        self.num_gpus = int(os.getenv("NUM_GPUS", "1"))
        self.data_config = os.getenv("DATA_CONFIG", "so100_dualcam")
        self.video_backend = os.getenv("VIDEO_BACKEND", "torchvision_av")
        self.batch_size = int(os.getenv("BATCH_SIZE", "32"))
        self.learning_rate = float(os.getenv("LEARNING_RATE", "1e-4"))
        self.base_model_path = os.getenv("BASE_MODEL_PATH", "nvidia/GR00T-N1.5-3B")
        self.embodiment_tag = os.getenv("EMBODIMENT_TAG", "new_embodiment")
        self.report_to = os.getenv("REPORT_TO", "tensorboard")

        # Optional parameters
        self.tune_llm = os.getenv("TUNE_LLM", "false").lower() == "true"
        self.tune_visual = os.getenv("TUNE_VISUAL", "false").lower() == "true"
        self.tune_projector = os.getenv("TUNE_PROJECTOR", "true").lower() == "true"
        self.tune_diffusion_model = (
            os.getenv("TUNE_DIFFUSION_MODEL", "true").lower() == "true"
        )
        self.lora_rank = int(os.getenv("LORA_RANK", "0"))
        self.lora_alpha = int(os.getenv("LORA_ALPHA", "16"))
        self.lora_dropout = float(os.getenv("LORA_DROPOUT", "0.1"))
        self.weight_decay = float(os.getenv("WEIGHT_DECAY", "1e-5"))
        self.warmup_ratio = float(os.getenv("WARMUP_RATIO", "0.05"))
        self.dataloader_num_workers = int(os.getenv("DATALOADER_NUM_WORKERS", "8"))
        self.dataloader_prefetch_factor = int(
            os.getenv("DATALOADER_PREFETCH_FACTOR", "4")
        )
        self.balance_dataset_weights = (
            os.getenv("BALANCE_DATASET_WEIGHTS", "true").lower() == "true"
        )
        self.balance_trajectory_weights = (
            os.getenv("BALANCE_TRAJECTORY_WEIGHTS", "true").lower() == "true"
        )
        self.lora_full_model = os.getenv("LORA_FULL_MODEL", "false").lower() == "true"
        self.resume = os.getenv("RESUME", "false").lower() == "true"

        # Validate required parameters
        self._validate_parameters()

    def _validate_parameters(self):
        """Validate required environment variables."""
        required_params = {
            "DATASET_LOCAL_DIR": self.dataset_local_dir,
        }

        missing_params = [
            param for param, value in required_params.items() if not value
        ]
        if missing_params:
            raise ValueError(
                f"Missing required environment variables: {', '.join(missing_params)}"
            )

        # Validate data_config
        valid_data_configs = [
            "so100_dualcam",
            "fourier_gr1_arms_only",
            "fourier_gr1_arms_waist",
            "agibot_genie1_dualcam",
            "oxe_droid_single_cam",
        ]
        if self.data_config not in valid_data_configs:
            logger.warning(f"Data config '{self.data_config}' may not be supported")

        # Validate embodiment_tag
        valid_embodiment_tags = ["new_embodiment", "gr1", "oxe_droid", "agibot_genie1"]
        if self.embodiment_tag not in valid_embodiment_tags:
            logger.warning(
                f"Embodiment tag '{self.embodiment_tag}' may not be supported"
            )

        logger.info("All required parameters validated successfully")

    def validate_dataset(self):
        """
        Ensure a local dataset directory is ready.
        """
        logger.info("Validating dataset...")

        # 1) Use explicit local directory if provided and non-empty
        if "DATASET_LOCAL_DIR" in os.environ and os.path.isdir(self.dataset_local_dir):
            if os.listdir(self.dataset_local_dir):
                logger.info(f"Using dataset directory: {self.dataset_local_dir}")

                meta_dir = os.path.join(self.dataset_local_dir, "meta")
                modality_json_path = os.path.join(meta_dir, "modality.json")
                if self.data_config == "so100_dualcam" and not os.path.isfile(
                    modality_json_path
                ):
                    os.makedirs(meta_dir, exist_ok=True)
                    modality_content = {
                        "state": {
                            "single_arm": {"start": 0, "end": 5},
                            "gripper": {"start": 5, "end": 6},
                        },
                        "action": {
                            "single_arm": {"start": 0, "end": 5},
                            "gripper": {"start": 5, "end": 6},
                        },
                        "video": {
                            "wrist": {"original_key": "observation.images.wrist"},
                            "front": {"original_key": "observation.images.front"},
                        },
                        "annotation": {
                            "human.task_description": {"original_key": "task_index"}
                        },
                    }
                    with open(modality_json_path, "w") as f:
                        json.dump(modality_content, f, indent=4)
                    logger.info(
                        f"Created missing modality.json at {modality_json_path} for so100_dualcam"
                    )
                return
            else:
                logger.warning(
                    f"DATASET_LOCAL_DIR is provided but empty: {self.dataset_local_dir}"
                )

        # Past this point, the dataset should already exist (shell prepared it)
        raise RuntimeError(
            "Dataset directory not prepared. Ensure entrypoint script resolved and downloaded the dataset."
        )

    def _train_once(self):
        """Run the fine-tuning steps in-process (ported from gr00t_finetune.py)."""
        logger.info("Starting training...")

        # ------------ step 1: load dataset ------------
        embodiment_tag = EmbodimentTag(self.embodiment_tag)

        # 1.1 modality configs and transforms
        data_config_cls = load_data_config(self.data_config)
        modality_configs = data_config_cls.modality_config()
        transforms = data_config_cls.transform()

        # 1.2 data loader: we will use either single dataset or mixture dataset
        dataset_path = [self.dataset_local_dir]
        if len(dataset_path) == 1:
            train_dataset = LeRobotSingleDataset(
                dataset_path=os.path.abspath(dataset_path[0]),
                modality_configs=modality_configs,
                transforms=transforms,
                embodiment_tag=embodiment_tag,  # This will override the dataset's embodiment tag to "new_embodiment"
                video_backend=self.video_backend,
            )
        else:
            single_datasets = []
            for p in dataset_path:
                assert os.path.exists(p), f"Dataset path {p} does not exist"
                # We use the same transforms, modality configs, and embodiment tag for all datasets here,
                # in reality, you can use dataset from different modalities and embodiment tags
                dataset = LeRobotSingleDataset(
                    dataset_path=p,
                    modality_configs=modality_configs,
                    transforms=transforms,
                    embodiment_tag=embodiment_tag,
                    video_backend=self.video_backend,
                )
                single_datasets.append(dataset)

            train_dataset = LeRobotMixtureDataset(
                data_mixture=[
                    (dataset, 1.0)
                    for dataset in single_datasets  # we will use equal weights for all datasets
                ],
                mode="train",
                balance_dataset_weights=self.balance_dataset_weights,
                balance_trajectory_weights=self.balance_trajectory_weights,
                seed=42,
                metadata_config={
                    "percentile_mixing_method": "weighted_average",
                },
            )
            print(f"Loaded {len(single_datasets)} datasets, with {dataset_path} ")

        # ------------ step 2: load model ------------
        # First, get the data config to determine action horizon
        data_action_horizon = len(data_config_cls.action_indices)

        # Load model
        model = GR00T_N1_5.from_pretrained(
            pretrained_model_name_or_path=self.base_model_path,
            tune_llm=self.tune_llm,  # backbone's LLM
            tune_visual=self.tune_visual,  # backbone's vision tower
            tune_projector=self.tune_projector,  # action head's projector
            tune_diffusion_model=self.tune_diffusion_model,  # action head's DiT
        )

        # Update action_horizon to match data config
        # Need to recreate action head with correct config since it was initialized with old config
        if data_action_horizon != model.action_head.config.action_horizon:
            print(
                f"Recreating action head with action_horizon {data_action_horizon} (was {model.action_head.config.action_horizon})"
            )

            # Update the action head config
            new_action_head_config = model.action_head.config
            new_action_head_config.action_horizon = data_action_horizon

            # Import the FlowmatchingActionHead class
            from gr00t.model.action_head.flow_matching_action_head import (
                FlowmatchingActionHead,
            )

            # Create new action head with updated config
            new_action_head = FlowmatchingActionHead(new_action_head_config)

            # Copy the weights from the old action head to the new one
            new_action_head.load_state_dict(
                model.action_head.state_dict(), strict=False
            )

            # Replace the action head
            model.action_head = new_action_head

            # Update model config AND the action_head_cfg dictionary that gets saved
            model.config.action_horizon = data_action_horizon
            model.action_horizon = data_action_horizon
            model.config.action_head_cfg["action_horizon"] = data_action_horizon

            # Set trainable parameters for the new action head
            model.action_head.set_trainable_parameters(
                tune_projector=self.tune_projector,
                tune_diffusion_model=self.tune_diffusion_model,
            )

        # Set the model's compute_dtype to bfloat16
        model.compute_dtype = "bfloat16"
        model.config.compute_dtype = "bfloat16"

        if self.lora_rank > 0:
            model = get_lora_model(
                model,
                rank=self.lora_rank,
                lora_alpha=self.lora_alpha,
                lora_dropout=self.lora_dropout,
                action_head_only=not self.lora_full_model,
            )

        # 2.1 modify training args
        training_args = TrainingArguments(
            output_dir=self.output_dir,
            run_name=None,
            remove_unused_columns=False,
            deepspeed="",
            gradient_checkpointing=False,
            bf16=True,
            tf32=True,
            per_device_train_batch_size=self.batch_size,
            gradient_accumulation_steps=1,
            dataloader_num_workers=self.dataloader_num_workers,
            dataloader_pin_memory=False,
            dataloader_prefetch_factor=self.dataloader_prefetch_factor,
            dataloader_persistent_workers=self.dataloader_num_workers > 0,
            optim="adamw_torch",
            adam_beta1=0.95,
            adam_beta2=0.999,
            adam_epsilon=1e-8,
            learning_rate=self.learning_rate,
            weight_decay=self.weight_decay,
            warmup_ratio=self.warmup_ratio,
            lr_scheduler_type="cosine",
            logging_steps=10.0,
            num_train_epochs=300,
            max_steps=self.max_steps,
            save_strategy="steps",
            save_steps=self.save_steps,
            # evaluation_strategy="no",
            save_total_limit=5,
            report_to=self.report_to,
            seed=42,
            do_eval=False,
            ddp_find_unused_parameters=False,
            ddp_bucket_cap_mb=100,
            torch_compile_mode=None,
        )

        # 2.2 run experiment
        experiment = TrainRunner(
            train_dataset=train_dataset,
            model=model,
            training_args=training_args,
            resume_from_checkpoint=self.resume,
        )

        # 2.3 run experiment
        experiment.train()

        logger.info("Training completed successfully")

    def run_training(self):
        """Run the training with optional multi-GPU support (torchrun)."""
        # Create output directories (checkpoints and tensorboard logs)
        os.makedirs(self.output_dir, exist_ok=True)

        available_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 1

        # Validate GPU configuration
        assert (
            self.num_gpus <= available_gpus
        ), f"Number of GPUs requested ({self.num_gpus}) is greater than the available GPUs ({available_gpus})"
        assert self.num_gpus > 0, "Number of GPUs must be greater than 0"
        print(f"Using {self.num_gpus} GPUs")

        if self.num_gpus == 1:
            # Single GPU mode - set CUDA_VISIBLE_DEVICES=0
            os.environ["CUDA_VISIBLE_DEVICES"] = "0"
            # Run training in-process
            self._train_once()
        else:
            # Multi-GPU mode - use torchrun to re-invoke this script with multiple processes
            if os.environ.get("IS_TORCHRUN", "0") == "1":
                # We are already inside a torchrun worker
                self._train_once()
            else:
                script_path = Path(__file__).absolute()
                # Remove any existing CUDA_VISIBLE_DEVICES from environment
                if "CUDA_VISIBLE_DEVICES" in os.environ:
                    del os.environ["CUDA_VISIBLE_DEVICES"]

                # Build torchrun args (call torch.distributed.run main directly)
                args = [
                    "--standalone",
                    f"--nproc_per_node={self.num_gpus}",
                    "--nnodes=1",
                    str(script_path),
                ]

                print("Running torchrun with args: ", args)
                os.environ["IS_TORCHRUN"] = "1"
                try:
                    torchrun(args=args)
                except SystemExit as e:
                    code = e.code if isinstance(e.code, int) else 0
                    sys.exit(code)

    def run_workflow(self):
        """Run the complete workflow."""
        try:
            logger.info("Starting GR00T fine-tuning...")

            # Step 1: Validate dataset was prepared by entrypoint script
            self.validate_dataset()

            # Step 2: Run training
            self.run_training()

            logger.info("GR00T fine-tuning completed successfully!")

        except Exception as e:
            logger.error(f"Workflow failed: {str(e)}")
            sys.exit(1)


def main():
    """Main entry point."""
    workflow = FinetuneWorkflow()
    workflow.run_workflow()


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
import sys
import os

# Add workstation module to Python path for import
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../.."))

from aws_cdk import App, Environment, Stack
from batch_stack import BatchStack
from workstation import DcvWorkstation, DcvWorkstationProps

app = App()

# Get environment variables with defaults
env = Environment(
    account=os.getenv("CDK_DEFAULT_ACCOUNT"),
    region=os.getenv("CDK_DEFAULT_REGION", "us-west-2"),
)

# If context values are provided, it will import the existing VPC/EFS directly.
ctx_vpc_id = app.node.try_get_context("vpc_id") or os.getenv("VPC_ID")
ctx_efs_id = app.node.try_get_context("efs_id") or os.getenv("EFS_ID")
ctx_efs_sg_id = app.node.try_get_context("efs_sg_id") or os.getenv("EFS_SG_ID")
ctx_ecr_image_uri = app.node.try_get_context("ecr_image_uri") or os.getenv(
    "ECR_IMAGE_URI"
)
ctx_dataset_bucket = app.node.try_get_context("dataset_bucket") or os.getenv(
    "DATASET_BUCKET"
)
ctx_s3_upload_uri = app.node.try_get_context("s3_upload_uri") or os.getenv(
    "S3_UPLOAD_URI"
)

batch_stack = BatchStack(
    app,
    "IsaacGr00tBatchStack",
    env=env,
    vpc_id=ctx_vpc_id,  # Optional: existing VPC
    efs_id=ctx_efs_id,  # Optional: existing EFS
    efs_sg_id=ctx_efs_sg_id,  # Optional: existing EFS SG
    ecr_image_uri=ctx_ecr_image_uri,  # Optional: pre-built ECR image
    dataset_bucket=ctx_dataset_bucket,  # Optional: S3 bucket for dataset access
    s3_upload_uri=ctx_s3_upload_uri,  # Optional: S3 URI for checkpoint uploads
)


class IsaacLabDcvStack(Stack):
    """DCV stack for gr00t visualization, integrated with BatchStack.

    Consumes the standalone workstation module from workstation/ and shares VPC/EFS
    with the gr00t BatchStack.
    """

    def __init__(
        self,
        scope,
        construct_id: str,
        batch_stack: BatchStack,
        **kwargs,
    ) -> None:
        super().__init__(scope, construct_id, **kwargs)

        # Configure DCV workstation with gr00t-specific settings
        # isaac_sim_version="5.1.0", isaac_lab_version="v2.3.0" for both N1.5 and N1.6
        # public_dcv_access: True (default) opens ports 8443/6006/8080 publicly.
        # Deploy with --context public_dcv_access=false for SSM-only access.
        public_dcv = self.node.try_get_context("public_dcv_access")
        public_dcv_access = public_dcv != "false"  # Default True unless explicitly "false"

        props = DcvWorkstationProps(
            vpc=batch_stack.vpc,              # Share VPC with Batch
            efs_id=batch_stack.efs_id,        # Share EFS with Batch
            efs_sg_id=batch_stack.efs_sg_id,  # Share security group
            # Preferred: g6.4xlarge, g6.2xlarge, g5.2xlarge or larger.
            # If capacity errors occur, try changing availability_zone between us-west-2a/b/c/d.
            instance_type="g6.4xlarge",       # Fall back to g6/g5.2xlarge if capacity unavailable.
            isaac_sim_version="5.1.0",        # Latest version with valid NGC container
            isaac_lab_version="v2.3.0",       # Matches isaac-lab:2.3.0 on NGC
            leisaac_enabled=True,             # Required for gr00t
            public_dcv_access=public_dcv_access,
        )

        self.dcv_workstation = DcvWorkstation(self, "DCV", props)


# DCV stack consumes shared infrastructure from BatchStack
dcv_stack = IsaacLabDcvStack(
    app,
    "IsaacLabDcvStack",  # ✅ Same stack ID as before
    env=env,
    batch_stack=batch_stack,
)

app.synth()

import os
from aws_cdk import (
    aws_ec2 as ec2,
    aws_batch as batch,
    aws_ecr as ecr,
    aws_iam as iam,
    aws_efs as efs,
    aws_ecs as ecs,
    aws_s3 as s3,
    Stack,
    CfnOutput,
    Duration,
    Size,
    RemovalPolicy,
)
from constructs import Construct
from codebuild_stack import CodeBuildStack


class BatchStack(Stack):
    def __init__(
        self,
        scope: Construct,
        construct_id: str,
        vpc_id: str = None,
        efs_id: str = None,
        efs_sg_id: str = None,
        ecr_image_uri: str = None,
        dataset_bucket: str = None,
        s3_upload_uri: str = None,
        **kwargs,
    ) -> None:
        """
        CDK stack for the AWS Batch resources used by the GR00T fine-tuning workflow.

        This file is intentionally structured to mirror the step-by-step flow in the
        blog (Draft Code Walkthrough) so it's easy to follow and customize:

        - 2.1 Create VPC and EFS
        - 2.2 Build the fine-tuning container and push to ECR
        - 2.3 Create Launch Template
        - 2.4 Create Compute Environment
        - 2.5 Create Job Queue and Job Definition

        Args:
            vpc_id: Existing VPC ID to reuse (optional)
            efs_id: Existing EFS file system ID to reuse (optional)
            efs_sg_id: Existing EFS security group ID (required if efs_id is provided)
            ecr_image_uri: Existing ECR image URI (e.g. 123456789012.dkr.ecr.us-west-2.amazonaws.com/gr00t-finetune:latest).
                          If not provided, builds from local Dockerfile.
            dataset_bucket: S3 bucket name for dataset read-only access (optional)
            s3_upload_uri: S3 URI for checkpoint uploads (e.g., s3://bucket/path). If not provided, creates a new bucket.
        """
        super().__init__(scope, construct_id, **kwargs)

        # ==============================================================
        # region 2.1 Create VPC and EFS
        # ==============================================================
        # Create or reference VPC. If you already have a VPC, pass its ID; otherwise we
        # create a VPC with one NAT gateway (to match "VPC and more" in the console flow).
        if vpc_id:
            vpc = ec2.Vpc.from_lookup(self, "BatchVPC", vpc_id=vpc_id)
        else:
            vpc = ec2.Vpc(
                self,
                "BatchVPC",
                vpc_name="BatchVPC",
                max_azs=2,
                ip_addresses=ec2.IpAddresses.cidr("10.0.0.0/16"),
                nat_gateways=1,
                subnet_configuration=[
                    # We keep a small public subnet for NAT/Egress. Jobs run in private subnets.
                    ec2.SubnetConfiguration(
                        name="Public",
                        subnet_type=ec2.SubnetType.PUBLIC,
                    ),
                    ec2.SubnetConfiguration(
                        name="Private",
                        subnet_type=ec2.SubnetType.PRIVATE_WITH_EGRESS,
                    ),
                ],
            )

        self.vpc = vpc

        # Create or import an EFS file system. This is where checkpoints and logs will be stored
        # so they persist across jobs and can be visualized from a DCV instance.
        if efs_id:
            efs_sg = ec2.SecurityGroup.from_security_group_id(
                self, "BatchEFSSecurityGroup", efs_sg_id, mutable=True
            )
            efs_fs = efs.FileSystem.from_file_system_attributes(
                self,
                "BatchEFS",
                file_system_id=efs_id,
                security_group=efs_sg,
            )
            # Expose attributes for cross-stack use
            self.efs_id = efs_id
            self.efs_sg_id = efs_sg_id
        else:
            efs_sg = ec2.SecurityGroup(
                self,
                "BatchEFSSecurityGroup",
                vpc=vpc,
                description="Security group for Batch instances and EFS",
            )
            efs_fs = efs.FileSystem(
                self,
                "BatchEFS",
                file_system_name="BatchEFS",
                vpc=vpc,
                security_group=efs_sg,
                performance_mode=efs.PerformanceMode.GENERAL_PURPOSE,
                throughput_mode=efs.ThroughputMode.BURSTING,
            )
            # Expose attributes for cross-stack use
            self.efs_id = efs_fs.file_system_id
            self.efs_sg_id = efs_sg.security_group_id

        # Add a self-referencing NFS rule so instances and EFS within the same SG can communicate.
        efs_sg.add_ingress_rule(
            peer=efs_sg,
            connection=ec2.Port.tcp(2049),
            description="Allow NFS within Batch EFS SG",
        )
        # endregion

        # ==============================================================
        # region 2.2 Build the fine-tuning container and push to ECR
        # ==============================================================
        # Container image selection strategy:
        # - Prefer an existing ECR image in the same account via ecr_image_uri
        # - Else automatically build via CodeBuild (works on any architecture)
        if ecr_image_uri:
            # Use provided ECR image
            # Expected format: <account>.dkr.ecr.<region>.amazonaws.com/<repo>:<tag>
            repo_and_tag = ecr_image_uri.split("/")[-1]
            if ":" in repo_and_tag:
                repo_name, tag = repo_and_tag.split(":", 1)
            else:
                repo_name, tag = repo_and_tag, "latest"
            repo = ecr.Repository.from_repository_name(
                self, "IsaacGr00tEcrRepo", repository_name=repo_name
            )
            container_image = ecs.ContainerImage.from_ecr_repository(
                repository=repo, tag=tag
            )
            codebuild_stack = None  # No CodeBuild needed
        else:
            # Automatically build container using CodeBuild
            # This works on any architecture (x86, ARM) since build happens in the cloud
            # build_target: "n15" for N1.5 Dockerfile (default), "n16" for N1.6, "n17" for N1.7
            ctx_build_target = self.node.try_get_context("build_target") or os.getenv("BUILD_TARGET", "n15")
            codebuild_stack = CodeBuildStack(
                self,
                "CodeBuild",
                ecr_repository_name="gr00t-finetune",
                use_stable=True,
                build_target=ctx_build_target,
            )
            # Use the built image
            container_image = ecs.ContainerImage.from_ecr_repository(
                repository=codebuild_stack.ecr_repository, tag="latest"
            )
            ecr_image_uri = codebuild_stack.image_uri

        # Store codebuild_stack for conditional outputs later
        self.codebuild_stack = codebuild_stack
        # endregion

        # ==============================================================
        # region 2.3 Create Launch Template
        # ==============================================================
        # Increase the Linux root volume size to 100 GiB for pulling docker containers.
        launch_template = ec2.LaunchTemplate(
            self,
            "BatchLaunchTemplate",
            launch_template_name="BatchLaunchTemplate",
            block_devices=[
                ec2.BlockDevice(
                    device_name="/dev/xvda",  # standard root device for Amazon Linux 2
                    volume=ec2.BlockDeviceVolume.ebs(
                        volume_size=100,
                        delete_on_termination=True,
                        volume_type=ec2.EbsDeviceVolumeType.GP3,
                    ),
                )
            ],
        )
        # endregion

        # ==============================================================
        # region 2.4 Create Compute Environment
        # ==============================================================
        # IAM role for Batch EC2 instances so they can pull images, mount EFS, access S3, etc.
        compute_env = batch.ManagedEc2EcsComputeEnvironment(
            self,
            "ComputeEnvironment",
            compute_environment_name="IsaacGr00tComputeEnvironment",
            vpc=vpc,
            vpc_subnets=ec2.SubnetSelection(
                subnet_type=ec2.SubnetType.PRIVATE_WITH_EGRESS
            ),  # Jobs run in private subnets
            launch_template=launch_template,
            security_groups=[efs_sg],
            instance_types=[
                # By default, limit to g6e family for cost savings
                # Single-GPU instances
                ec2.InstanceType("g6e.2xlarge"),
                ec2.InstanceType("g6e.4xlarge"),
                ec2.InstanceType("g6e.8xlarge"),
                # ec2.InstanceType("g6e.16xlarge"),
                # ec2.InstanceType("p5.4xlarge"),
                # Multi-GPU instances
                ec2.InstanceType("g6e.12xlarge"),  # 4 GPUs
                # ec2.InstanceType("g6e.24xlarge"),  # 4 GPUs
                ec2.InstanceType("g6e.48xlarge"),  # 8 GPUs
                # ec2.InstanceType("p4d.24xlarge"),  # 8 GPUs
                # ec2.InstanceType("p5.48xlarge"),  # 8 GPUs
            ],
            minv_cpus=0,
            maxv_cpus=192,
            instance_role=iam.Role(
                self,
                "BatchInstanceRole",
                assumed_by=iam.ServicePrincipal("ec2.amazonaws.com"),
                managed_policies=[
                    iam.ManagedPolicy.from_aws_managed_policy_name(
                        "service-role/AmazonEC2ContainerServiceforEC2Role"
                    ),
                    iam.ManagedPolicy.from_aws_managed_policy_name(
                        "AmazonSSMManagedInstanceCore"
                    ),
                ],
            ),
            # Uncomment for cost-optimized runs on Spot
            # spot=True,
            # spot_bid_percentage=70,
        )

        # No explicit allow needed beyond the SG rule above since EFS and instances share the SG.
        # endregion

        # ==============================================================
        # region 2.5 Create Job Queue and Job Definition
        # ==============================================================
        job_queue = batch.JobQueue(
            self,
            "JobQueue",
            job_queue_name="IsaacGr00tJobQueue",
            compute_environments=[
                batch.OrderedComputeEnvironment(
                    compute_environment=compute_env, order=1
                )
            ],
            priority=1,
        )

        # Job role for the container tasks (access to S3 during training/upload)
        job_role = iam.Role(
            self,
            "JobRole",
            assumed_by=iam.ServicePrincipal("ecs-tasks.amazonaws.com"),
            managed_policies=[],
        )

        # Separate dataset bucket (read-only) from checkpoint upload bucket (read/write).
        # 1) If dataset_bucket is provided, allow read-only on that bucket.
        # 2) If s3_upload_uri is provided (e.g., s3://bucket/path), allow read/write to its bucket/prefix.
        # 3) If s3_upload_uri is not provided, create a new checkpoint bucket and derive s3_upload_uri.
        # 4) If neither dataset nor upload buckets are specified, fall back to S3 read-only (useful for dataset downloads).
        if dataset_bucket:
            job_role.add_to_policy(
                iam.PolicyStatement(
                    actions=["s3:ListBucket"],
                    resources=[f"arn:aws:s3:::{dataset_bucket}"],
                )
            )
            job_role.add_to_policy(
                iam.PolicyStatement(
                    actions=["s3:GetObject"],
                    resources=[f"arn:aws:s3:::{dataset_bucket}/*"],
                )
            )

        # Resolve or create checkpoint upload bucket/URI
        checkpoint_bucket = None
        original_s3_upload_uri = s3_upload_uri  # Track original state
        if not s3_upload_uri:
            # Create a new S3 bucket for checkpoints
            checkpoint_bucket = s3.Bucket(
                self,
                "IsaacGr00tCheckpointBucket",
                block_public_access=s3.BlockPublicAccess.BLOCK_ALL,
                encryption=s3.BucketEncryption.S3_MANAGED,
                enforce_ssl=True,
                versioned=True,
                removal_policy=RemovalPolicy.RETAIN,
                auto_delete_objects=False,
            )
            # Derive s3_upload_uri from the created bucket
            s3_upload_uri = f"s3://{checkpoint_bucket.bucket_name}/gr00t/checkpoints"

        if checkpoint_bucket is not None:
            # Grant RW to the created checkpoint bucket
            checkpoint_bucket.grant_read_write(job_role)
        elif s3_upload_uri and s3_upload_uri.startswith("s3://"):
            remainder = s3_upload_uri[5:]
            if "/" in remainder:
                upload_bucket, upload_prefix = remainder.split("/", 1)
            else:
                upload_bucket, upload_prefix = remainder, ""
            job_role.add_to_policy(
                iam.PolicyStatement(
                    actions=["s3:ListBucket"],
                    resources=[f"arn:aws:s3:::{upload_bucket}"],
                )
            )
            object_resource = (
                f"arn:aws:s3:::{upload_bucket}/{upload_prefix}*"
                if upload_prefix
                else f"arn:aws:s3:::{upload_bucket}/*"
            )
            job_role.add_to_policy(
                iam.PolicyStatement(
                    actions=[
                        "s3:GetObject",
                        "s3:PutObject",
                        "s3:DeleteObject",
                        "s3:AbortMultipartUpload",
                        "s3:ListMultipartUploadParts",
                    ],
                    resources=[object_resource],
                )
            )

        # Only grant general S3 read-only access if no buckets were originally specified
        if (
            not dataset_bucket
            and not original_s3_upload_uri
            and checkpoint_bucket is None
        ):
            job_role.add_managed_policy(
                iam.ManagedPolicy.from_aws_managed_policy_name("AmazonS3ReadOnlyAccess")
            )

        # Mount EFS at /mnt/efs inside the container. The workflow script and training output
        # (checkpoints, tensorboard logs) are configured to write under this path.
        efs_volume = batch.EcsVolume.efs(
            name="BatchEFS", file_system=efs_fs, container_path="/mnt/efs"
        )

        # Prepare container environment with defaults and optional S3 settings
        container_environment = {
            # Optional default locations on EFS. You can override at submit time.
            "OUTPUT_DIR": "/mnt/efs/gr00t/checkpoints"
        }
        if s3_upload_uri:
            container_environment["UPLOAD_TARGET"] = "s3"
            container_environment["S3_UPLOAD_URI"] = s3_upload_uri

        job_def = batch.EcsJobDefinition(
            self,
            "IsaacGr00tJobDefinition",
            job_definition_name="IsaacGr00tJobDefinition",
            container=batch.EcsEc2ContainerDefinition(
                self,
                "IsaacGr00tContainer",
                image=container_image,
                memory=Size.gibibytes(64),
                cpu=8,
                gpu=1,
                job_role=job_role,
                environment=container_environment,
                volumes=[efs_volume],
                linux_parameters=batch.LinuxParameters(
                    self,
                    "IsaacGr00tLinuxParameters",
                    shared_memory_size=Size.gibibytes(64),
                ),
            ),
            timeout=Duration.hours(6),
        )
        # endregion

        # ==============================================================
        # region Outputs
        # ==============================================================
        CfnOutput(self, "VpcId", value=self.vpc.vpc_id)
        CfnOutput(self, "EFSFileSystemId", value=self.efs_id)
        CfnOutput(self, "EFSSecurityGroupId", value=self.efs_sg_id)

        # Additional outputs for convenience when submitting jobs via CLI/Console
        CfnOutput(
            self, "ComputeEnvironmentName", value=compute_env.compute_environment_name
        )
        CfnOutput(self, "JobQueueName", value=job_queue.job_queue_name)
        CfnOutput(self, "JobDefinitionName", value=job_def.job_definition_name)
        CfnOutput(self, "EcrImageUri", value=ecr_image_uri)
        if s3_upload_uri:
            CfnOutput(self, "CheckpointS3UploadUri", value=s3_upload_uri)

        # CodeBuild outputs (only when CodeBuild is used)
        if codebuild_stack:
            CfnOutput(
                self,
                "CodeBuildProjectName",
                value=codebuild_stack.build_project.project_name,
                description="CodeBuild project name for building the container",
            )
        # endregion

version: 0.2

# Build specification for GR00T fine-tuning container
# This buildspec is used by AWS CodeBuild to build and push the Docker image to ECR

phases:
  pre_build:
    commands:
      - echo "Starting pre-build phase..."
      - echo "Build started on $(date)"
      - echo "Logging in to Amazon ECR..."

      # Get AWS account ID and region
      - export AWS_ACCOUNT_ID=$(aws sts get-caller-identity --query Account --output text)
      - "export AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-west-2}"
      - 'echo "AWS Account ID: $AWS_ACCOUNT_ID"'
      - 'echo "AWS Region: $AWS_DEFAULT_REGION"'

      # Login to ECR
      - aws ecr get-login-password --region $AWS_DEFAULT_REGION | docker login --username AWS --password-stdin $AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com

      # Set image variables
      - "export IMAGE_REPO_NAME=${ECR_REPOSITORY_NAME:-gr00t-finetune}"
      - "export IMAGE_TAG=${IMAGE_TAG:-latest}"
      - "export IMAGE_URI=$AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com/$IMAGE_REPO_NAME:$IMAGE_TAG"
      - 'echo "Image URI: $IMAGE_URI"'

      # Docker info
      - docker --version
      - docker info

  build:
    commands:
      - echo "Starting build phase..."
      - echo "Building Docker image using build_container.sh..."

      # Use the existing build_container.sh script to maintain consistency with Path 3
      # Select the build directory based on BUILD_TARGET (n16=N16/, n17=N17/; default is n15 root)
      - |
        if [ "${BUILD_TARGET}" = "n17" ]; then
          cd $CODEBUILD_SRC_DIR/N17
        elif [ "${BUILD_TARGET}" = "n16" ]; then
          cd $CODEBUILD_SRC_DIR/N16
        else
          cd $CODEBUILD_SRC_DIR
        fi
      - chmod +x build_container.sh

      # Build with the script, passing appropriate flags
      - 'if [ "${USE_STABLE}" = "true" ]; then ./build_container.sh --test -n $IMAGE_REPO_NAME -t $IMAGE_TAG; else ./build_container.sh --test --latest -n $IMAGE_REPO_NAME -t $IMAGE_TAG; fi'

      # Tag the image with build number for versioning
      - "docker tag $IMAGE_REPO_NAME:$IMAGE_TAG $AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com/$IMAGE_REPO_NAME:$CODEBUILD_BUILD_NUMBER"
      - "docker tag $IMAGE_REPO_NAME:$IMAGE_TAG $IMAGE_URI"

      - echo "Docker image built and tested successfully"
      - docker images

  post_build:
    commands:
      - echo "Starting post-build phase..."
      - echo "Pushing Docker image to ECR..."

      # Push the image with both tags (latest and build number)
      - "docker push $IMAGE_URI"
      - "docker push $AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com/$IMAGE_REPO_NAME:$CODEBUILD_BUILD_NUMBER"

      - echo "Docker image pushed successfully"
      - 'echo "Image URI: $IMAGE_URI"'
      - 'echo "Build number tag: $AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com/$IMAGE_REPO_NAME:$CODEBUILD_BUILD_NUMBER"'

      # Write image URI to file for easy retrieval
      - echo $IMAGE_URI > image_uri.txt
      - echo "Build completed on $(date)"

# Artifacts to export (optional, for reference)
artifacts:
  files:
    - image_uri.txt
  name: gr00t-container-build-artifacts

{
    "app": "python3 app.py"
}

import os
from aws_cdk import (
    aws_codebuild as codebuild,
    aws_ecr as ecr,
    aws_iam as iam,
    aws_s3_assets as s3_assets,
    custom_resources as cr,
    CfnOutput,
    RemovalPolicy,
    Duration,
)
from constructs import Construct


class CodeBuildStack(Construct):
    def __init__(
        self,
        scope: Construct,
        construct_id: str,
        ecr_repository_name: str = "gr00t-finetune",
        use_stable: bool = True,
        build_target: str = "n15",
    ) -> None:
        """
        CDK construct for AWS CodeBuild project to build GR00T fine-tuning container.

        This construct creates:
        - ECR repository for storing container images
        - CodeBuild project with x86 compute for building containers
        - IAM roles and permissions
        - S3 bucket for source code (if using local source)

        Args:
            ecr_repository_name: Name for the ECR repository (default: gr00t-finetune)
            use_stable: Use stable GR00T commit vs latest (default: True)
            build_target: "n15" for N1.5 Dockerfile, "n16" for N16/, "n17" for N17/ (default: n15)
        """
        super().__init__(scope, construct_id)

        # ==============================================================
        # 1. ECR Repository
        # ==============================================================
        # Create ECR repository to store the built container images
        ecr_repo = ecr.Repository(
            self,
            "IsaacGr00tEcrRepository",
            repository_name=ecr_repository_name,
            removal_policy=RemovalPolicy.RETAIN,  # Keep images after stack deletion
            image_scan_on_push=True,  # Scan for vulnerabilities
            lifecycle_rules=[
                # Keep last 10 images, delete older ones
                ecr.LifecycleRule(
                    description="Keep last 10 images",
                    max_image_count=10,
                    rule_priority=1,
                )
            ],
        )

        # ==============================================================
        # 2. Source Code Asset
        # ==============================================================
        # Package the local source code and upload to a CDK managed S3 bucket
        # CodeBuild will download from S3 to build the container
        asset_path = os.path.abspath(
            os.path.join(os.path.dirname(__file__), "..")
        )  # training/gr00t/
        source_asset = s3_assets.Asset(
            self,
            "IsaacGr00tSourceAsset",
            path=asset_path,
            exclude=[
                ".git",
                ".gitignore",
                "*.pyc",
                "__pycache__",
                "cdk.context.json",
                ".venv",
                "venv",
                "*.egg-info",
                ".pytest_cache",
                ".mypy_cache",
                "cdk.out",
                "infra/.cdk.staging",
            ],
        )

        # ==============================================================
        # 3. CodeBuild Project
        # ==============================================================
        # Create CodeBuild project to build the Docker image
        build_project = codebuild.Project(
            self,
            "IsaacGr00tContainerBuild",
            project_name="IsaacGr00tContainerBuild",
            description="Build GR00T fine-tuning container and push to ECR",
            # Source: Use the S3 asset created above
            # Customize the source to use your own Git repository
            source=codebuild.Source.s3(
                bucket=source_asset.bucket,
                path=source_asset.s3_object_key,
            ),
            # Build environment
            environment=codebuild.BuildEnvironment(
                # Use x86_64 architecture (required for EC2 G6e, P4 and P5 instances)
                build_image=codebuild.LinuxBuildImage.STANDARD_7_0,  # Ubuntu 22.04, Docker 24
                compute_type=codebuild.ComputeType.LARGE,  # 8 vCPU, 15 GB RAM
                privileged=True,  # Required for Docker builds
            ),
            # Build specification (located in infra/ directory)
            build_spec=codebuild.BuildSpec.from_source_filename("infra/buildspec.yml"),
            # Environment variables
            environment_variables={
                "ECR_REPOSITORY_NAME": codebuild.BuildEnvironmentVariable(
                    value=ecr_repository_name
                ),
                "USE_STABLE": codebuild.BuildEnvironmentVariable(
                    value="true" if use_stable else "false"
                ),
                "IMAGE_TAG": codebuild.BuildEnvironmentVariable(value="latest"),
                "BUILD_TARGET": codebuild.BuildEnvironmentVariable(value=build_target),
            },
            # Timeout (building flash-attn takes time)
            timeout=Duration.hours(2),
            # Cache for faster rebuilds (optional)
            cache=codebuild.Cache.local(
                codebuild.LocalCacheMode.DOCKER_LAYER,
                codebuild.LocalCacheMode.CUSTOM,
            ),
        )

        # ==============================================================
        # 4. IAM Permissions
        # ==============================================================
        # Grant CodeBuild permissions to push to ECR
        ecr_repo.grant_pull_push(build_project.role)

        # Grant permissions to read source from S3
        source_asset.grant_read(build_project.role)

        # Add ECR authorization token permission (required for docker login)
        build_project.role.add_to_policy(
            iam.PolicyStatement(
                actions=["ecr:GetAuthorizationToken"],
                resources=["*"],
            )
        )

        # ==============================================================
        # 5. Auto-trigger Build on Stack Creation/Update
        # ==============================================================
        # Automatically trigger a CodeBuild build when the stack is created or updated
        # The build triggers when files in the asset path change because the physical
        # resource ID includes the asset hash, which changes when source files change
        trigger_build = cr.AwsCustomResource(
            self,
            "AutoTriggerBuild",
            # Use from_sdk_calls for simpler policy management
            # This automatically grants the necessary permissions for the SDK call
            policy=cr.AwsCustomResourcePolicy.from_sdk_calls(
                resources=cr.AwsCustomResourcePolicy.ANY_RESOURCE
            ),
            # Lambda timeout should be minimal - it just triggers the build, doesn't wait for completion
            timeout=Duration.minutes(5),
            # Trigger on both CREATE and UPDATE
            # The physical resource ID includes the asset's S3 object key hash, so it changes when files change
            on_create=cr.AwsSdkCall(
                service="CodeBuild",
                action="startBuild",
                parameters={
                    "projectName": build_project.project_name,
                },
                # Use S3 object key hash in physical resource ID so it changes when files change
                # The s3_object_key contains a hash that changes when asset contents change
                # This ensures builds trigger on every deploy when source files are modified
                physical_resource_id=cr.PhysicalResourceId.of(
                    f"{build_project.project_name}-build-{source_asset.s3_object_key}"
                ),
            ),
            # on_update: no-op — only trigger builds on CREATE (first deploy) or when source files
            # change (asset hash changes → physical resource ID changes → replacement → on_create fires).
            # Using batchGetProjects as a read-only no-op so repeated deploys don't re-trigger builds.
            on_update=cr.AwsSdkCall(
                service="CodeBuild",
                action="batchGetProjects",
                parameters={
                    "names": [build_project.project_name],
                },
                physical_resource_id=cr.PhysicalResourceId.of(
                    f"{build_project.project_name}-build-{source_asset.s3_object_key}"
                ),
            ),
            # Install latest AWS SDK in Lambda runtime for latest API support
            install_latest_aws_sdk=True,
        )
        # Ensure the build project exists before triggering
        trigger_build.node.add_dependency(build_project)

        # ==============================================================
        # 6. Outputs
        # ==============================================================
        CfnOutput(
            self,
            "CodeBuildProjectName",
            value=build_project.project_name,
            description="CodeBuild project name for building the container",
        )

        CfnOutput(
            self,
            "BuildCommand",
            value=f"aws codebuild start-build --project-name {build_project.project_name}",
            description="Command to trigger a container build",
        )

        CfnOutput(
            self,
            "EcrRepositoryUri",
            value=ecr_repo.repository_uri,
            description="ECR repository URI for the GR00T container image",
        )

        CfnOutput(
            self,
            "ImageUri",
            value=f"{ecr_repo.repository_uri}:latest",
            description="Full ECR image URI (use this for Batch stack deployment)",
        )

        # Store attributes for cross-stack references
        self.ecr_repository = ecr_repo
        self.build_project = build_project
        self.image_uri = f"{ecr_repo.repository_uri}:latest"

aws-cdk-lib>=2.0.0
constructs>=10.0.0

# Exclude infrastructure and build artifacts
infra/
cdk.out/
__pycache__/
*.pyc
*.pyo
*.pyd
.Python
env/
venv/
.venv
pip-log.txt
pip-delete-this-directory.txt
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.log
.git/
.mypy_cache/
.pytest_cache/
.hypothesis/
**/.DS_Store

#!/bin/bash

# Build script for GR00T Fine-tuning Docker image

set -Eeuo pipefail

echo "=========================================="
echo "Building GR00T Fine-tuning Docker Image"
echo "=========================================="

# Default values
IMAGE_NAME="gr00t-finetune"
TAG="latest"
DOCKERFILE="Dockerfile"
PUSH_IMAGE=false
TEST_IMAGE=false
USE_STABLE=true

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -t|--tag)
            TAG="$2"
            shift 2
            ;;
        -n|--name)
            IMAGE_NAME="$2"
            shift 2
            ;;
        --latest)
            USE_STABLE=false
            shift
            ;;
        --push)
            PUSH_IMAGE=true
            shift
            ;;
        --test)
            TEST_IMAGE=true
            shift
            ;;
        -h|--help)
            echo "Usage: $0 [OPTIONS]"
            echo "Options:"
            echo "  -t, --tag TAG         Tag for the fine-tuning image (default: latest)"
            echo "  -n, --name NAME       Name for the fine-tuning image (default: gr00t-finetune)"
            echo "  --latest              Use latest GR00T from main branch (default: stable commit)"
            echo "  --push                Push image to registry after building"
            echo "  --test                Run basic tests after building"
            echo "  -h, --help            Show this help message"
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

FULL_IMAGE_NAME="${IMAGE_NAME}:${TAG}"

echo "Building image: ${FULL_IMAGE_NAME}"
echo "Using Dockerfile: ${DOCKERFILE}"

# Display GR00T version selection
if [[ "${USE_STABLE}" == "true" ]]; then
    echo "GR00T version: STABLE (N1.6, tested commit) [default]"
else
    echo "GR00T version: LATEST (main branch, may have breaking changes)"
fi

# Build the fine-tuning image directly from the combined Dockerfile
echo "Building fine-tuning image..."
docker build \
    --build-arg USE_STABLE=${USE_STABLE} \
    -f ${DOCKERFILE} \
    -t ${FULL_IMAGE_NAME} \
    .

echo "Image built successfully: ${FULL_IMAGE_NAME}"

# Run basic tests if requested
if [[ "${TEST_IMAGE}" == "true" ]]; then
    echo "=========================================="
    echo "Running Basic Tests"
    echo "=========================================="
    
    # Test 1: Check if the image runs without errors (dry run)
    echo "Test 1: Checking if image starts correctly..."
    docker run --rm \
        -e HF_TOKEN="dummy" \
        -e HF_DATASET_ID="dummy/dummy" \
        -e HF_MODEL_REPO_ID="dummy/dummy" \
        --entrypoint /bin/bash \
        ${FULL_IMAGE_NAME} \
        -c "echo 'Image starts correctly' && python -c 'import sys; print(f\"Python version: {sys.version}\")' && which huggingface-cli"
    
    # Test 2: Check if finetune script can be imported
    echo "Test 2: Checking if finetune script imports correctly..."
    docker run --rm \
        --entrypoint /bin/bash \
        ${FULL_IMAGE_NAME} \
        -c "PYTHONPATH=/workspace/gr00t-repo:/workspace python -c 'import importlib.util; spec = importlib.util.spec_from_file_location(\"finetune\", \"/workspace/scripts/finetune_gr00t.py\"); mod = importlib.util.module_from_spec(spec); spec.loader.exec_module(mod); print(\"Finetune script imports successfully\")'"

    # Test 3: Check if modality config can be loaded
    echo "Test 3: Checking if modality config is accessible..."
    docker run --rm \
        --entrypoint /bin/bash \
        ${FULL_IMAGE_NAME} \
        -c "PYTHONPATH=/workspace/gr00t-repo:/workspace python -c 'import importlib.util; spec = importlib.util.spec_from_file_location(\"m\", \"/workspace/scripts/so101_modality_config.py\"); print(\"Modality config accessible\")'"
    
    echo "All tests passed!"
fi

# Push to registry if requested
if [[ "${PUSH_IMAGE}" == "true" ]]; then
    echo "=========================================="
    echo "Pushing to Registry"
    echo "=========================================="
    
    if [[ -z "${DOCKER_REGISTRY}" ]]; then
        echo "Warning: DOCKER_REGISTRY environment variable not set."
        echo "Assuming you want to push to Docker Hub or have already tagged appropriately."
    else
        # Re-tag with registry prefix
        REGISTRY_IMAGE="${DOCKER_REGISTRY}/${FULL_IMAGE_NAME}"
        docker tag ${FULL_IMAGE_NAME} ${REGISTRY_IMAGE}
        FULL_IMAGE_NAME=${REGISTRY_IMAGE}
    fi
    
    echo "Pushing image: ${FULL_IMAGE_NAME}"
    docker push ${FULL_IMAGE_NAME}
    echo "Image pushed successfully!"
fi

echo "=========================================="
echo "Build Complete!"
echo "=========================================="
echo "Image: ${FULL_IMAGE_NAME}"
if [[ "${USE_STABLE}" == "true" ]]; then
    echo "GR00T Version: STABLE (tested commit)"
else
    echo "GR00T Version: LATEST (main branch)"
fi
echo ""
echo "To run locally, create a local directory to simulate EFS mount:"
echo "mkdir -p ~/mock-efs/gr00t/checkpoints"
echo "Then run with a small number of steps for testing:"
echo "docker run --gpus all --network host \\"
echo "  -e MAX_STEPS=100 -e SAVE_STEPS=100 \\"
echo "  -v ~/mock-efs:/mnt/efs \\"
echo "  ${FULL_IMAGE_NAME}"
echo ""
echo "To rebuild with latest GR00T version from main branch:"
echo "  ./build_container.sh --latest" 

"""
Modality configuration for SO-ARM101 (6-DOF) robot.

Based on upstream examples/SO100/so100_config.py — the SO-ARM100/101
share the same 6-DOF joint structure (5 arm joints + 1 gripper).

This file is loaded by the training script via N1.6's
register_modality_config() mechanism.
"""

from gr00t.configs.data.embodiment_configs import register_modality_config
from gr00t.data.embodiment_tags import EmbodimentTag
from gr00t.data.types import (
    ActionConfig,
    ActionFormat,
    ActionRepresentation,
    ActionType,
    ModalityConfig,
)

so101_config = {
    "video": ModalityConfig(
        delta_indices=[0],
        modality_keys=["front", "wrist"],
    ),
    "state": ModalityConfig(
        delta_indices=[0],
        modality_keys=["single_arm", "gripper"],
    ),
    "action": ModalityConfig(
        delta_indices=list(range(0, 16)),
        modality_keys=["single_arm", "gripper"],
        action_configs=[
            ActionConfig(
                rep=ActionRepresentation.RELATIVE,
                type=ActionType.NON_EEF,
                format=ActionFormat.DEFAULT,
            ),
            ActionConfig(
                rep=ActionRepresentation.ABSOLUTE,
                type=ActionType.NON_EEF,
                format=ActionFormat.DEFAULT,
            ),
        ],
    ),
    "language": ModalityConfig(
        delta_indices=[0],
        modality_keys=["annotation.human.task_description"],
    ),
}

register_modality_config(so101_config, embodiment_tag=EmbodimentTag.NEW_EMBODIMENT)

# Exclude infrastructure and build artifacts
infra/
cdk.out/
__pycache__/
*.pyc
*.pyo
*.pyd
.Python
env/
venv/
.venv
pip-log.txt
pip-delete-this-directory.txt
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.log
.git/
.mypy_cache/
.pytest_cache/
.hypothesis/
**/.DS_Store