launch

#!/usr/bin/env bash
set -Eeuo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LAUNCH_D="${SCRIPT_DIR}/launch.d"
GENERIC_CONFIG="${LAUNCH_D}/generic.yaml"
INSTANCE_TYPES_D="${LAUNCH_D}/instance_types"
DRY_RUN="${DRY_RUN:-0}"

usage() {
  cat >&2 <<'EOF'
Usage:
  ./launch [--dry-run] <profile> <repo> <train_script> <train_config> [git_ref]

  profile     GPU / node profile: gpus | gpul | gpum | a100 | a100-80
  repo        Git URL, git@github.com:org/repo.git, or org/repo (GitHub shorthand)
  train_script Path inside the repo to the training entrypoint (e.g. train/act/train.sh)
  train_config Path inside the repo to the job YAML (for train/act/train.sh this should be
               train/act/configs/job.yaml — dataset + policy). For train/act/train.sh, launch
               auto-selects a device YAML from profile (gpus/gpul/gpum/a100/a100-80) unless overridden.
  git_ref     Optional branch or tag (default: main)

Environment overrides:
  launch.d/generic.yaml controls namespace, image, PVCs, shared memory, and trainer env.
  launch.d/instance_types/<profile>.yaml controls node selectors and resources.
  NAMESPACE, SERVICE_ACCOUNT_NAME, LAUNCH_IMAGE
  DATA_PVC, CHECKPOINT_PVC, CACHE_PVC, S3_DATA_PVC
  GIT_CLONE_DEPTH, SHARED_MEMORY_SIZE
  LAUNCH_MEM_REQUEST, LAUNCH_MEM_LIMIT, LAUNCH_CPU_REQUEST, LAUNCH_CPU_LIMIT
  LAUNCH_GPU_REQUEST, LAUNCH_GPU_LIMIT   (override profile GPU counts)
  LAUNCH_DEVICE_CONFIG_REL               Optional repo-relative path to force DEVICE_CONFIG
                                         (e.g. train/act/configs/aws_gpul.yaml)
  DRY_RUN=1 or --dry-run                 Print manifest only (no kubectl)
  LAUNCHER_ID                            Override detected launcher identity

Examples (lerobot-edit-scripts ACT — job YAML + profile-mapped device file):
  ./lib/update_volumes.sh
  ./launch gpul ETHRoboticsClub/lerobot-edit-scripts \\
    train/act/train.sh train/act/configs/job.yaml
  ./launch a100 ETHRoboticsClub/lerobot-edit-scripts \\
    train/act/train.sh train/act/configs/job.yaml

Follow logs:
  kubectl -n robot-learning logs -f job/<trainjob-pod-name> -c node
  (pod name: kubectl -n robot-learning get pods -l job-name=<TrainJob.metadata.name> ...)

Kill launched jobs:
  ./kill <TrainJob/Pod/ConfigMap name copied from k9s>
  ./kill robot-learning
EOF
}

log() {
  printf '[launch] %s\n' "$*" >&2
}

die() {
  log "$*"
  exit 1
}

yaml_sq() {
  local s="${1:-}"
  s="${s//\'/\'\'}"
  printf "'%s'" "$s"
}

yaml_value() {
  local file="$1"
  local wanted_key="$2"
  awk -v wanted_key="${wanted_key}" '
    index($0, ":") {
      key = substr($0, 1, index($0, ":") - 1)
      value = substr($0, index($0, ":") + 1)
      sub(/^[[:space:]]*/, "", key)
      sub(/[[:space:]]*$/, "", key)
      if (key != wanted_key) {
        next
      }
      sub(/^[[:space:]]*/, "", value)
      sub(/[[:space:]]*$/, "", value)
      print value
      exit
    }
  ' "${file}"
}

yaml_map_section() {
  local file="$1"
  local section="$2"
  awk -F: -v section="${section}" '
    $0 == section ":" {
      in_section = 1
      next
    }
    in_section && /^[^[:space:]]/ {
      exit
    }
    in_section && /^[[:space:]]+[A-Za-z0-9_]+:/ {
      key = $1
      value = substr($0, index($0, ":") + 1)
      sub(/^[[:space:]]*/, "", key)
      sub(/[[:space:]]*$/, "", key)
      sub(/^[[:space:]]*/, "", value)
      sub(/[[:space:]]*$/, "", value)
      print key "\t" value
    }
  ' "${file}"
}

load_generic_config() {
  [[ -f "${GENERIC_CONFIG}" ]] || die "Missing generic launch config: ${GENERIC_CONFIG}"

  local bootstrap_script
  NAMESPACE="${NAMESPACE:-$(yaml_value "${GENERIC_CONFIG}" namespace)}"
  SERVICE_ACCOUNT_NAME="${SERVICE_ACCOUNT_NAME:-$(yaml_value "${GENERIC_CONFIG}" service_account_name)}"
  IMAGE="${LAUNCH_IMAGE:-$(yaml_value "${GENERIC_CONFIG}" image)}"
  DATA_PVC="${DATA_PVC:-$(yaml_value "${GENERIC_CONFIG}" data_pvc)}"
  CHECKPOINT_PVC="${CHECKPOINT_PVC:-$(yaml_value "${GENERIC_CONFIG}" checkpoint_pvc)}"
  CACHE_PVC="${CACHE_PVC:-$(yaml_value "${GENERIC_CONFIG}" cache_pvc)}"
  S3_DATA_PVC="${S3_DATA_PVC:-$(yaml_value "${GENERIC_CONFIG}" s3_data_pvc)}"
  GIT_CLONE_DEPTH="${GIT_CLONE_DEPTH:-$(yaml_value "${GENERIC_CONFIG}" git_clone_depth)}"
  DEFAULT_SHARED_MEMORY_SIZE="$(yaml_value "${GENERIC_CONFIG}" default_shared_memory_size)"
  BOOTSTRAP_DIRS="${BOOTSTRAP_DIRS:-$(yaml_value "${GENERIC_CONFIG}" bootstrap_dirs)}"
  BOOTSTRAP_CHOWN_DIRS="${BOOTSTRAP_CHOWN_DIRS:-$(yaml_value "${GENERIC_CONFIG}" bootstrap_chown_dirs)}"

  bootstrap_script="$(yaml_value "${GENERIC_CONFIG}" bootstrap_script)"
  BOOTSTRAP_LOCAL="${BOOTSTRAP_LOCAL:-${LAUNCH_D}/${bootstrap_script}}"

  local required_var
  for required_var in NAMESPACE SERVICE_ACCOUNT_NAME IMAGE DATA_PVC CHECKPOINT_PVC CACHE_PVC S3_DATA_PVC GIT_CLONE_DEPTH DEFAULT_SHARED_MEMORY_SIZE BOOTSTRAP_DIRS BOOTSTRAP_CHOWN_DIRS BOOTSTRAP_LOCAL; do
    if [[ -z "${!required_var:-}" ]]; then
      die "Generic config is missing ${required_var} in ${GENERIC_CONFIG}"
    fi
  done
}

detect_launcher_identity() {
  if [[ -n "${LAUNCHER_ID:-}" ]]; then
    printf '%s\n' "${LAUNCHER_ID}"
    return 0
  fi

  if command -v aws >/dev/null 2>&1; then
    local aws_arn
    aws_arn="$(aws sts get-caller-identity --query Arn --output text 2>/dev/null || true)"
    if [[ -n "${aws_arn}" && "${aws_arn}" != "None" ]]; then
      printf 'aws:%s\n' "${aws_arn}"
      return 0
    fi
  fi

  printf 'local:%s@%s\n' "${USER:-unknown}" "$(hostname -s 2>/dev/null || hostname)"
}

launcher_hash() {
  printf '%s' "$1" | shasum -a 256 | awk '{print substr($1, 1, 16)}'
}

safe_rel_path() {
  local p="$1"
  local name="$2"
  if [[ -z "${p}" || "${p}" == /* || "${p}" == ../* || "${p}" == */.. || "${p}" == */../* ]]; then
    die "${name} must be a relative repo path (no leading / or ..)."
  fi
  if [[ "${p}" != *[a-zA-Z]* ]]; then
    die "${name} looks invalid."
  fi
}

resolve_repo_url() {
  local raw="$1"
  case "${raw}" in
    https://* | http://* | git@* | ssh://*)
      printf '%s\n' "${raw}"
      ;;
    *@*:*)
      printf '%s\n' "${raw}"
      ;;
    */*)
      printf 'https://github.com/%s.git\n' "${raw}"
      ;;
    *)
      die "Could not interpret repo as a URL or org/repo: ${raw}"
      ;;
  esac
}

profile_apply() {
  local name="$1"
  local config="${INSTANCE_TYPES_D}/${name}.yaml"

  if [[ ! "${name}" =~ ^[a-z0-9][a-z0-9-]*$ || ! -f "${config}" ]]; then
    die "Unknown profile '${name}'. Run './launch --list-profiles' to see available profiles."
  fi

  INSTANCE_TYPE_CONFIG="${config}"
  NODEPOOL=""
  NODETIER=""
  GPU_REQUEST=""
  GPU_LIMIT=""
  CPU_REQUEST=""
  CPU_LIMIT=""
  MEM_REQUEST=""
  MEM_LIMIT=""
  INSTANCE_SHARED_MEMORY_SIZE=""
  MULTI_GPU_ENV=""
  TRAIN_ACT_DEVICE_CONFIG_DEFAULT=""

  local line
  local key
  local value
  while IFS= read -r line || [[ -n "${line}" ]]; do
    [[ -z "${line}" || "${line}" =~ ^[[:space:]]*# ]] && continue
    [[ "${line}" =~ ^[[:space:]] ]] && continue
    [[ "${line}" == *:* ]] || continue

    key="${line%%:*}"
    value="${line#*:}"
    key="$(printf '%s' "${key}" | tr -d '[:space:]')"
    value="$(printf '%s' "${value}" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')"

    case "${key}" in
      description | instance_type | instance_types | vcpus | memory | gpu_model | gpu_memory)
        ;;
      nodepool)
        NODEPOOL="${value}"
        ;;
      nodetier)
        NODETIER="${value}"
        ;;
      gpu_request)
        GPU_REQUEST="${value}"
        ;;
      gpu_limit)
        GPU_LIMIT="${value}"
        ;;
      cpu_request)
        CPU_REQUEST="${value}"
        ;;
      cpu_limit)
        CPU_LIMIT="${value}"
        ;;
      memory_request)
        MEM_REQUEST="${value}"
        ;;
      memory_limit)
        MEM_LIMIT="${value}"
        ;;
      shared_memory_size)
        INSTANCE_SHARED_MEMORY_SIZE="${value}"
        ;;
      multi_gpu_env)
        MULTI_GPU_ENV="${value}"
        ;;
      train_act_device_config)
        TRAIN_ACT_DEVICE_CONFIG_DEFAULT="${value}"
        ;;
      extra_trainer_env)
        ;;
      *)
        die "Unknown key '${key}' in ${config}"
        ;;
    esac
  done <"${config}"

  GPU_REQUEST="${LAUNCH_GPU_REQUEST:-${GPU_REQUEST}}"
  GPU_LIMIT="${LAUNCH_GPU_LIMIT:-${GPU_LIMIT}}"
  CPU_REQUEST="${LAUNCH_CPU_REQUEST:-${CPU_REQUEST}}"
  CPU_LIMIT="${LAUNCH_CPU_LIMIT:-${CPU_LIMIT}}"
  MEM_REQUEST="${LAUNCH_MEM_REQUEST:-${MEM_REQUEST}}"
  MEM_LIMIT="${LAUNCH_MEM_LIMIT:-${MEM_LIMIT}}"
  SHARED_MEMORY_SIZE="${SHARED_MEMORY_SIZE:-${INSTANCE_SHARED_MEMORY_SIZE:-${DEFAULT_SHARED_MEMORY_SIZE}}}"

  local required_var
  for required_var in NODEPOOL NODETIER GPU_REQUEST GPU_LIMIT CPU_REQUEST CPU_LIMIT MEM_REQUEST MEM_LIMIT SHARED_MEMORY_SIZE MULTI_GPU_ENV; do
    if [[ -z "${!required_var:-}" ]]; then
      die "Profile '${name}' is missing ${required_var} in ${config}"
    fi
  done
}

list_profiles() {
  cat >&2 <<EOF
Profiles from ${INSTANCE_TYPES_D}:
EOF

  local config
  local name
  local description
  for config in "${INSTANCE_TYPES_D}"/*.yaml; do
    [[ -f "${config}" ]] || continue
    name="$(basename "${config}" .yaml)"
    description="$(yaml_value "${config}" description)"
    printf '  %-8s %s\n' "${name}" "${description}" >&2
  done

  cat >&2 <<'EOF'
Override RAM/GPU if scheduling fails, e.g.:
  LAUNCH_MEM_REQUEST=256Gi LAUNCH_MEM_LIMIT=320Gi ./launch gpul ...
EOF
}

emit_trainer_env_yaml() {
  local key
  local value
  while IFS=$'\t' read -r key value; do
    [[ -n "${key}" ]] || continue
    printf '      - name: %s\n' "${key}"
    printf '        value: %s\n' "$(yaml_sq "${value}")"
  done < <(yaml_map_section "${GENERIC_CONFIG}" trainer_env)
}

emit_instance_env_yaml() {
  local key
  local value
  while IFS=$'\t' read -r key value; do
    [[ -n "${key}" ]] || continue
    printf '      - name: %s\n' "${key}"
    printf '        value: %s\n' "$(yaml_sq "${value}")"
  done < <(yaml_map_section "${INSTANCE_TYPE_CONFIG}" extra_trainer_env)
}

default_device_config_for_profile() {
  local train_script="$2"

  if [[ "${train_script}" != "train/act/train.sh" ]]; then
    return 0
  fi

  printf '%s\n' "${TRAIN_ACT_DEVICE_CONFIG_DEFAULT:-}"
}

emit_trainjob_yaml() {
  local job_name="$1"
  local cm_name="$2"
  local profile="$3"
  local train_git_url="$4"
  local train_git_ref="$5"
  local train_script_rel="$6"
  local train_config_rel="$7"

  local device_override_block=""
  if [[ -n "${TRAIN_DEVICE_CONFIG_REL:-}" ]]; then
    device_override_block="$(printf '%s\n' "      - name: TRAIN_DEVICE_CONFIG_REL" "        value: $(yaml_sq "${TRAIN_DEVICE_CONFIG_REL}")")"
  fi

  local trainer_env_block
  local instance_env_block
  trainer_env_block="$(emit_trainer_env_yaml)"
  instance_env_block="$(emit_instance_env_yaml)"

  cat <<EOF
apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
  name: ${job_name}
  namespace: ${NAMESPACE}
  labels:
    app: generic-train
    launch.ethrc.github.io/managed-by: launch
    launch.ethrc.github.io/job-name: ${job_name}
    launch.ethrc.github.io/profile: ${profile}
    launch.ethrc.github.io/launcher-hash: ${LAUNCHER_HASH}
  annotations:
    launch.ethrc.github.io/launcher: $(yaml_sq "${LAUNCHER_ID_RAW}")
    launch.ethrc.github.io/kill-mine: $(yaml_sq "./kill mine")
    launch.ethrc.github.io/kill-one: $(yaml_sq "./kill ${job_name}")
    launch.ethrc.github.io/kill-namespace: $(yaml_sq "./kill ${NAMESPACE}")
spec:
  runtimeRef:
    name: torch-distributed
  trainer:
    image: ${IMAGE}
    env:
      - name: TRAIN_GIT_URL
        value: $(yaml_sq "${train_git_url}")
      - name: TRAIN_GIT_REF
        value: $(yaml_sq "${train_git_ref}")
      - name: TRAIN_SCRIPT_REL
        value: $(yaml_sq "${train_script_rel}")
      - name: TRAIN_CONFIG_REL
        value: $(yaml_sq "${train_config_rel}")
      - name: GIT_CLONE_DEPTH
        value: $(yaml_sq "${GIT_CLONE_DEPTH}")
${trainer_env_block}
${device_override_block}
${instance_env_block}
    command:
      - /bin/bash
      - -lc
      - |
        set -euxo pipefail
        sudo mkdir -p ${BOOTSTRAP_DIRS}
        sudo chown -R "\$(id -u):\$(id -g)" ${BOOTSTRAP_CHOWN_DIRS}
        exec > >(tee -a /checkpoints/generic-train-bootstrap.log) 2>&1
        set +x
        export WANDB_API_KEY="\$(cat /var/run/secrets/wandb/WANDB_API_KEY)"
        export HF_TOKEN="\$(cat /var/run/secrets/hf/HF_TOKEN)"
        set -x
        exec /bin/bash /job-files/pod-bootstrap.sh
    numNodes: 1
    resourcesPerNode:
      requests:
        cpu: "${CPU_REQUEST}"
        memory: "${MEM_REQUEST}"
        nvidia.com/gpu: "${GPU_REQUEST}"
      limits:
        cpu: "${CPU_LIMIT}"
        memory: "${MEM_LIMIT}"
        nvidia.com/gpu: "${GPU_LIMIT}"
  podTemplateOverrides:
    - targetJobs:
        - name: node
      spec:
        serviceAccountName: ${SERVICE_ACCOUNT_NAME}
        nodeSelector:
          karpenter.sh/nodepool: ${NODEPOOL}
          node-tier: ${NODETIER}
        tolerations:
          - key: nvidia.com/gpu
            operator: Exists
            effect: NoSchedule
        volumes:
          - name: data
            persistentVolumeClaim:
              claimName: ${DATA_PVC}
          - name: checkpoints
            persistentVolumeClaim:
              claimName: ${CHECKPOINT_PVC}
          - name: s3-data
            persistentVolumeClaim:
              claimName: ${S3_DATA_PVC}
          - name: dshm
            emptyDir:
              medium: Memory
              sizeLimit: ${SHARED_MEMORY_SIZE}
          - name: wandb-secret
            secret:
              secretName: wandb-secret
          - name: hf-secret
            secret:
              secretName: hf-secret
          - name: job-files
            configMap:
              name: ${cm_name}
              defaultMode: 0755
        containers:
          - name: node
            volumeMounts:
              - name: data
                mountPath: /data
              - name: checkpoints
                mountPath: /checkpoints
              - name: s3-data
                mountPath: /s3
              - name: dshm
                mountPath: /dev/shm
              - name: wandb-secret
                mountPath: /var/run/secrets/wandb
                readOnly: true
              - name: hf-secret
                mountPath: /var/run/secrets/hf
                readOnly: true
              - name: job-files
                mountPath: /job-files
                readOnly: true
EOF
}

main() {
  if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
    usage
    exit 0
  fi
  if [[ "${1:-}" == "--list-profiles" ]]; then
    list_profiles
    exit 0
  fi

  if [[ "${1:-}" == "--dry-run" ]]; then
    DRY_RUN=1
    shift
  fi

  local profile="${1:-}"
  local repo="${2:-}"
  local train_script="${3:-}"
  local train_config="${4:-}"
  local git_ref="${5:-main}"

  if [[ -z "${profile}" || -z "${repo}" || -z "${train_script}" || -z "${train_config}" ]]; then
    usage
    exit 1
  fi

  load_generic_config
  [[ -f "${BOOTSTRAP_LOCAL}" ]] || die "Missing bootstrap script: ${BOOTSTRAP_LOCAL}"

  safe_rel_path "${train_script}" "train_script"
  safe_rel_path "${train_config}" "train_config"

  local train_git_url
  train_git_url="$(resolve_repo_url "${repo}")"

  LAUNCHER_ID_RAW="$(detect_launcher_identity)"
  LAUNCHER_HASH="$(launcher_hash "${LAUNCHER_ID_RAW}")"

  profile_apply "${profile}"

  TRAIN_DEVICE_CONFIG_REL="${LAUNCH_DEVICE_CONFIG_REL:-}"
  if [[ -z "${TRAIN_DEVICE_CONFIG_REL}" ]]; then
    TRAIN_DEVICE_CONFIG_REL="$(default_device_config_for_profile "${profile}" "${train_script}" || true)"
  fi
  if [[ -n "${TRAIN_DEVICE_CONFIG_REL}" ]]; then
    safe_rel_path "${TRAIN_DEVICE_CONFIG_REL}" "LAUNCH_DEVICE_CONFIG_REL"
    log "Device config (DEVICE_CONFIG): ${TRAIN_DEVICE_CONFIG_REL}"
  fi

  local run_id
  run_id="$(date +%Y%m%d-%H%M%S)"
  local job_name="generic-train-${run_id}"
  local cm_name="${job_name}-bootstrap"

  log "TrainJob: ${job_name}"
  log "Launcher: ${LAUNCHER_ID_RAW} (${LAUNCHER_HASH})"
  log "Profile: ${profile} (nodepool=${NODEPOOL}, gpus=${GPU_REQUEST}, cpu=${CPU_REQUEST}/${CPU_LIMIT}, mem=${MEM_REQUEST}/${MEM_LIMIT})"
  log "Repo: ${train_git_url} @ ${git_ref}"
  log "Script: ${train_script}"
  log "Config (JOB_CONFIG): ${train_config}"

  local manifest
  manifest="$(mktemp)"
  emit_trainjob_yaml "${job_name}" "${cm_name}" "${profile}" "${train_git_url}" "${git_ref}" "${train_script}" "${train_config}" >"${manifest}"

  if [[ "${DRY_RUN}" == 1 ]]; then
    cat "${manifest}"
    rm -f "${manifest}"
    log "Dry run: no resources applied."
    exit 0
  fi

  kubectl create configmap "${cm_name}" \
    -n "${NAMESPACE}" \
    --from-file=pod-bootstrap.sh="${BOOTSTRAP_LOCAL}" \
    --dry-run=client \
    -o yaml \
    | kubectl apply -f -

  kubectl -n "${NAMESPACE}" label configmap "${cm_name}" \
    app=generic-train \
    launch.ethrc.github.io/managed-by=launch \
    launch.ethrc.github.io/job-name="${job_name}" \
    launch.ethrc.github.io/launcher-hash="${LAUNCHER_HASH}" \
    --overwrite

  kubectl -n "${NAMESPACE}" annotate configmap "${cm_name}" \
    launch.ethrc.github.io/launcher="${LAUNCHER_ID_RAW}" \
    launch.ethrc.github.io/kill-mine="./kill mine" \
    --overwrite

  kubectl apply -f "${manifest}"
  rm -f "${manifest}"

  log "Applied. Tail logs:"
  log "  kubectl -n ${NAMESPACE} get pods | grep ${job_name}"
  log "  kubectl -n ${NAMESPACE} logs -f <pod-name> -c node"
  log "Kill this job:"
  log "  ./kill ${job_name}"
  log "  (or copy a TrainJob, Pod, or ConfigMap name from k9s and run ./kill <copied-name>)"
  log "Kill your ./launch jobs:"
  log "  ./kill mine"
  log "Kill all ./launch jobs in this namespace:"
  log "  ./kill ${NAMESPACE}"
}

main "$@"