infer.yaml

workflow:
  name: isaaclab-inference
  timeout:
    exec_timeout: 2h
  resources:
    default:
      gpu: 1
      cpu: 4
      memory: 32Gi
      storage: 40Gi
  tasks:
    - name: isaac-inference
      image: "{{ image }}"
      command:
        - /bin/bash
        - /tmp/entry.sh
      environment:
        TASK: "{{ task }}"
        NUM_ENVS: "{{ num_envs }}"
        MAX_STEPS: "{{ max_steps }}"
        VIDEO_LENGTH: "{{ video_length }}"
        ACCEPT_EULA: "Y"
        PRIVACY_CONSENT: "Y"
        ENCODED_ARCHIVE: "{{ encoded_archive }}"
        PAYLOAD_ROOT: "{{ payload_root }}"
        CHECKPOINT_URI: "{{ checkpoint_uri }}"
        INFERENCE_FORMAT: "{{ inference_format }}"
        AZURE_SUBSCRIPTION_ID: "{{ azure_subscription_id }}"
        AZURE_RESOURCE_GROUP: "{{ azure_resource_group }}"
        AZUREML_WORKSPACE_NAME: "{{ azure_workspace_name }}"
        AZURE_STORAGE_ACCOUNT_NAME: "{{ azure_storage_account_name }}"
        AZURE_STORAGE_CONTAINER_NAME: "{{ azure_storage_container_name }}"
        AZURE_AUTHORITY_HOST: "{{ azure_authority_host }}"
        PYTHON: "{{ python }}"
      files:
        - path: /tmp/entry.sh
          contents: |
            #!/bin/bash
            set -euo pipefail

            if [[ -z "${ENCODED_ARCHIVE:-}" ]]; then
              echo "encoded_archive is required" >&2
              exit 1
            fi

            if [[ -z "${CHECKPOINT_URI:-}" ]]; then
              echo "checkpoint_uri is required" >&2
              exit 1
            fi

            ARCHIVE_PATH="${TEMP_ARCHIVE:-/tmp/isaac_payload.zip}"
            PAYLOAD_ROOT="${PAYLOAD_ROOT:-/workspace/isaac_payload}"
            INFERENCE_ROOT="${INFERENCE_ROOT:-${PAYLOAD_ROOT}/src/inference}"

            mkdir -p "${PAYLOAD_ROOT}"
            printf '%s' "${ENCODED_ARCHIVE}" | base64 -d > "${ARCHIVE_PATH}"
            unzip -oq "${ARCHIVE_PATH}" -d "${PAYLOAD_ROOT}"

            infer_args=(
              --task "${TASK:-}"
              --num-envs "${NUM_ENVS:-4}"
              --max-steps "${MAX_STEPS:-500}"
              --video-length "${VIDEO_LENGTH:-200}"
              --inference-format "${INFERENCE_FORMAT:-both}"
              --checkpoint-uri "${CHECKPOINT_URI:-}"
              --headless
            )

            bash "${INFERENCE_ROOT}/scripts/infer.sh" "${infer_args[@]}"

default-values:
  image: nvcr.io/nvidia/isaac-lab:2.2.0
  task: Isaac-Ant-v0
  num_envs: "4"
  max_steps: "500"
  video_length: "200"
  encoded_archive: ""
  payload_root: /workspace/isaac_payload
  checkpoint_uri: ""
  inference_format: "both"
  azure_subscription_id: ""
  azure_resource_group: ""
  azure_workspace_name: ""
  azure_storage_account_name: ""
  azure_storage_container_name: "isaaclab-training-logs"
  azure_authority_host: https://login.microsoftonline.com
  python: "/workspace/isaaclab/isaaclab.sh -p"