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: ["bash", "-euc"]
      args:
        - |
          base64 -d /tmp/entry.sh.b64 > /tmp/entry.sh
          [ -s /tmp/entry.sh ] || { echo "ERROR: /tmp/entry.sh is empty — submit via training/il/scripts/submit-osmo-lerobot-training.sh" >&2; exit 1; }
          exec bash /tmp/entry.sh
      inputs:
        - url: "{{ code_url }}"
          regex: ".*"
{% if use_huggingface_credential == "true" %}
      credentials:
        huggingface:
          HF_TOKEN: hf_token
{% endif %}
      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 }}"
        OSMO_INPUT_0: "{{input:0}}"
      files:
        - path: /tmp/entry.sh.b64
          contents: "{{ entry_script_b64 }}"

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: ""
  entry_script_b64: ""
  payload_root: /workspace/lerobot_payload
  dataset_repo_id: ""
  dataset_root: /workspace/data
  blob_urls: "[]"
  use_huggingface_credential: "true"
  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"