ci3/aws_request_instance at main · aztec-labs-eng/ci3 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
#!/usr/bin/env bash
[ -n "${BUILD_SYSTEM_DEBUG:-}" ] && set -x # conditionally trace
set -eu

name=$1
cpus=$2
arch=$3
state_dir=$4

cd $(dirname $0)

# Map a cpu count to its instance size suffix.
declare -A cpu_map
cpu_map=(
  [2]="large"
  [4]="xlarge"
  [8]="2xlarge"
  [16]="4xlarge"
  [32]="8xlarge"
  [48]="12xlarge"
  [64]="16xlarge"
  [96]="24xlarge"
  [128]="32xlarge"
  [192]="48xlarge"
  [0192]="metal"
)

# Spot pool preference by cpu count (lower number = tried first). 128 is the
# known-good build size, then 192, then 64 as a last resort before on-demand.
declare -A cpu_priority=([128]=1 [192]=2 [64]=3)

# Capacity-diverse spot pools per cpu count. Restricted to families with
# >=4 GB/vCPU (m/r) so the build's measured ~1.9 GB/vCPU peak keeps ~2x
# headroom; c-series (2 GB/vCPU) is opt-in via CI_SPOT_INCLUDE_C. All entries
# are confirmed offered as spot across every us-east-2 AZ.
function spot_pools_for {
  local cpu=$1
  case "$arch" in
    x86_64|amd64)
      case "$cpu" in
        128) echo "m6a.32xlarge m7a.32xlarge m6i.32xlarge r6a.32xlarge r7a.32xlarge" ;;
        192) echo "m6a.48xlarge m7a.48xlarge m7i.48xlarge r6a.48xlarge r7a.48xlarge" ;;
        64)  echo "m6a.16xlarge m7a.16xlarge m7i.16xlarge m6i.16xlarge m5.16xlarge r6a.16xlarge" ;;
        *)   local s=${cpu_map[$cpu]:?Unknown cpu count: $cpu}; echo "m6a.$s m7a.$s m6i.$s r6a.$s" ;;
      esac ;;
    arm64)
      case "$cpu" in
        64) echo "r7g.16xlarge r8g.16xlarge m7g.16xlarge m8g.16xlarge" ;;
        *)  local s=${cpu_map[$cpu]:?Unknown cpu count: $cpu}; echo "r7g.$s r8g.$s m7g.$s m8g.$s" ;;
      esac ;;
  esac
}

# Optional c-series pools (2 GB/vCPU). Gated behind CI_SPOT_INCLUDE_C until
# build-phase parallelism is throttled to fit the lower memory headroom.
function spot_pools_c_for {
  local cpu=$1
  case "$arch" in
    x86_64|amd64)
      case "$cpu" in
        128) echo "c6a.32xlarge c6i.32xlarge" ;;
        192) echo "c6a.48xlarge c7a.48xlarge" ;;
        64)  echo "c6a.16xlarge c6i.16xlarge" ;;
      esac ;;
  esac
}

if [ "$arch" == "x86_64" ] || [ "$arch" == "amd64" ]; then
  ami=${AMI:-$(cat ./aws/ami_id_amd64)}
elif [ "$arch" == "arm64" ]; then
  ami=${AMI:-$(cat ./aws/ami_id_arm64)}
else
  echo "Unsupported arch: $arch" >&2
  exit 1
fi

# Build the fleet override set as space-separated "instanceType:priority" tokens.
# AWS_INSTANCE pins a single pool (e.g. for debugging); otherwise diversify
# across every requested size.
overrides=""
if [ -n "${AWS_INSTANCE:-}" ]; then
  overrides="$AWS_INSTANCE:1"
else
  IFS=',' read -ra cpu_list <<< "$cpus"
  for cpu in "${cpu_list[@]}"; do
    prio=${cpu_priority[$cpu]:-9}
    types="$(spot_pools_for "$cpu")"
    [ "${CI_SPOT_INCLUDE_C:-0}" -eq 1 ] && types+=" $(spot_pools_c_for "$cpu")"
    for t in $types; do overrides+=" $t:$prio"; done
  done
fi

FLEET_OVERRIDES="${overrides# }" aws_request_instance_type "$name" "$ami" "$state_dir"