workflow:
  name: lerobot-training
  timeout:
    exec_timeout: 12h
  resources:
    default:
      gpu: {{ num_gpus }}
      cpu: 8
      memory: 48Gi
      storage: 80Gi
      platform: "{{ platform }}"
  tasks:
    - name: lerobot-train
      image: "{{ image }}"
      command:
        - /bin/bash
        - /tmp/entry.sh
      inputs:
        - url: "{{ code_url }}"
          regex: ".*"
      credentials:
        huggingface:
          HF_TOKEN: hf_token
      environment:
        GIT_PYTHON_REFRESH: "quiet"
        NVIDIA_DRIVER_CAPABILITIES: "all"
        DATASET_REPO_ID: "{{ dataset_repo_id }}"
        DATASET_ROOT: "{{ dataset_root }}"
        # Single-quoted: {{ blob_urls }} expands to a JSON array literal whose
        # double quotes would otherwise break the YAML string.
        BLOB_URLS: '{{ blob_urls }}'
        POLICY_TYPE: "{{ policy_type }}"
        OUTPUT_DIR: "{{ output_dir }}"
        JOB_NAME: "{{ job_name }}"
        POLICY_REPO_ID: "{{ policy_repo_id }}"
        TRAINING_STEPS: "{{ training_steps }}"
        BATCH_SIZE: "{{ batch_size }}"
        LEARNING_RATE: "{{ learning_rate }}"
        LR_WARMUP_STEPS: "{{ lr_warmup_steps }}"
        EVAL_FREQ: "{{ eval_freq }}"
        SAVE_FREQ: "{{ save_freq }}"
        LOG_FREQ: "{{ log_freq }}"
        VAL_SPLIT: "{{ val_split }}"
        SYSTEM_METRICS: "{{ system_metrics }}"
        MIXED_PRECISION: "{{ mixed_precision }}"
        LEROBOT_VERSION: "{{ lerobot_version }}"
        EXPERIMENT_NAME: "{{ experiment_name }}"
        REGISTER_CHECKPOINT: "{{ register_checkpoint }}"
        AZURE_SUBSCRIPTION_ID: "{{ azure_subscription_id }}"
        AZURE_RESOURCE_GROUP: "{{ azure_resource_group }}"
        AZUREML_WORKSPACE_NAME: "{{ azure_workspace_name }}"
        PAYLOAD_ROOT: "{{ payload_root }}"
        AZURE_AUTHORITY_HOST: "{{ azure_authority_host }}"
        MLFLOW_TRACKING_TOKEN_REFRESH_RETRIES: "{{ mlflow_token_refresh_retries }}"
        MLFLOW_HTTP_REQUEST_TIMEOUT: "{{ mlflow_http_request_timeout }}"
      files:
        - path: /tmp/entry.sh
          contents: |
            #!/bin/bash
            set -euo pipefail

            echo "=== LeRobot Training Workflow ==="
            echo "Dataset: ${DATASET_REPO_ID}"
            echo "Policy Type: ${POLICY_TYPE}"
            echo "Job Name: ${JOB_NAME}"
            echo "Output Dir: ${OUTPUT_DIR}"
            echo "Logging: Azure ML MLflow"
            echo "Val Split: ${VAL_SPLIT:-0.1}"
            echo "System Metrics: ${SYSTEM_METRICS:-true}"
            echo "Mixed Precision: ${MIXED_PRECISION:-no}"

            # Install system dependencies
            echo "Installing system dependencies..."
            export DEBIAN_FRONTEND=noninteractive
            apt-get update -qq -o Acquire::Retries=3
            APT_PACKAGES=(
              ffmpeg
              libgl1
              build-essential
              gcc
              unzip
              python3-dev
            )
            if apt-cache show libglib2.0-0t64 >/dev/null 2>&1; then
              APT_PACKAGES+=(libglib2.0-0t64)
            else
              APT_PACKAGES+=(libglib2.0-0)
            fi
            apt-get install -y -qq --no-install-recommends "${APT_PACKAGES[@]}"

            # Install UV package manager. --break-system-packages bypasses
            # PEP 668 (externally-managed-environment), which is safe in this
            # ephemeral container and portable across Debian, conda, and
            # PyTorch base images regardless of where the marker file lives.
            echo "Installing UV package manager..."
            pip install --quiet --break-system-packages uv==0.7.12

            echo "Creating Python 3.12 environment..."
            uv python install 3.12
            uv venv --python 3.12 /opt/lerobot-venv
            # shellcheck disable=SC1091
            source /opt/lerobot-venv/bin/activate

            # Unpack training scripts delivered via the OSMO url: input. The code
            # archive arrives as a downloaded object under the input mount rather
            # than an env var: a single env string is capped at 128 KiB
            # (MAX_ARG_STRLEN) and the archive (which ships the lerobot uv.lock
            # for `uv export`) exceeds that, failing the container execve E2BIG.
            PAYLOAD_ROOT="${PAYLOAD_ROOT:-/workspace/lerobot_payload}"
            mkdir -p "${PAYLOAD_ROOT}"
            ARCHIVE_PATH="$(find "{{input:0}}" -maxdepth 2 -name '*.zip' | head -n1)"
            if [[ -z "${ARCHIVE_PATH}" || ! -s "${ARCHIVE_PATH}" ]]; then
              echo "ERROR: no code archive found under the input mount; training payload is required." >&2
              ls -laR "{{input:0}}" || true
              exit 1
            fi
            unzip -oq "${ARCHIVE_PATH}" -d "${PAYLOAD_ROOT}"
            export PYTHONPATH="${PAYLOAD_ROOT}:${PYTHONPATH:-}"
            echo "Training scripts unpacked to ${PAYLOAD_ROOT} from ${ARCHIVE_PATH}"

            # Decide blob-vs-HuggingFace via the canonical Python helper so
            # whitespace, pretty-printed JSON, and [""] / [null] payloads agree
            # with download_dataset.prepare_dataset() instead of routing through
            # the blob branch and crashing.
            if python -c 'from training.il.scripts.lerobot._env import has_blob_urls; raise SystemExit(0 if has_blob_urls() else 1)'; then
              BLOB_DATASOURCE=1
              echo "Data Source: Azure Blob URLs"
            else
              BLOB_DATASOURCE=0
              echo "Data Source: HuggingFace Hub"
            fi

            # Install runtime dependencies from a build-time export of the lock
            LEROBOT_PROJECT="${PAYLOAD_ROOT}/training/il/lerobot"
            if [[ ! -f "${LEROBOT_PROJECT}/uv.lock" ]]; then
              echo "ERROR: LeRobot lockfile not found at ${LEROBOT_PROJECT}/uv.lock" >&2
              exit 1
            fi
            uv export --frozen --no-hashes --no-emit-project --project "${LEROBOT_PROJECT}" \
              | uv pip install --no-cache-dir --no-deps -r -

            # Build training command arguments
            TRAIN_ARGS=(
              "--policy.push_to_hub=false"
              "--wandb.enable=false"
              "--dataset.video_backend=pyav"
            )

            # Download dataset from Azure Blob Storage when configured
            if [[ "${BLOB_DATASOURCE}" == "1" ]]; then
              echo "Downloading dataset from Azure Blob Storage..."
              python -m training.il.scripts.lerobot.download_dataset

              FULL_DATASET_PATH="${DATASET_ROOT}/${DATASET_REPO_ID}"
              echo "Dataset downloaded to: ${FULL_DATASET_PATH}"
              TRAIN_ARGS+=(
                "--dataset.root=${FULL_DATASET_PATH}"
                "--dataset.use_imagenet_stats=false"
              )
            fi

            # Run training via Python orchestrator
            echo "Starting LeRobot training..."
            python -m training.il.scripts.lerobot.train "${TRAIN_ARGS[@]}"

            echo "=== Training Complete ==="
            ls -la "${OUTPUT_DIR}/" 2>/dev/null || true

default-values:
  # Submission scripts set image from scripts/lib/common.sh via --set-string.
  # This value is the direct-workflow fallback; keep it in sync with DEFAULT_LEROBOT_TRAIN_IMAGE.
  image: pytorch/pytorch:2.11.0-cuda12.8-cudnn9-runtime@sha256:eee11b3b3872a8c838e35ef48f08b2d5def2080902c7f666831310ca1a0ef2be
  code_url: ""
  payload_root: /workspace/lerobot_payload
  dataset_repo_id: ""
  dataset_root: /workspace/data
  blob_urls: "[]"
  policy_type: act
  output_dir: /workspace/outputs/train
  job_name: lerobot-training
  policy_repo_id: ""
  training_steps: "100000"
  batch_size: "32"
  learning_rate: "1e-4"
  lr_warmup_steps: "1000"
  eval_freq: ""
  save_freq: "5000"
  log_freq: ""
  val_split: "0.1"
  system_metrics: "true"
  num_gpus: "1"
  mixed_precision: "no"
  platform: gpu_platform
  lerobot_version: ""
  experiment_name: ""
  register_checkpoint: ""
  azure_subscription_id: ""
  azure_resource_group: ""
  azure_workspace_name: ""
  azure_authority_host: https://login.microsoftonline.com
  mlflow_token_refresh_retries: "3"
  mlflow_http_request_timeout: "60"