gpu_capacity/config.py

import os

# One representative (smallest) size per GPU family, covered by the background
# sweep. Any other size can be probed on demand via the API or dashboard.
# Kept to small sizes on purpose: a successful probe holds the reservation
# ~2 s billed at the instance's on-demand rate.
INSTANCE_TYPES = [
    "g4ad.xlarge",
    "g4dn.xlarge",
    "g5.xlarge",
    "g5g.xlarge",
    "g6.xlarge",
    "g6e.xlarge",
    "g7e.2xlarge",
    "gr6.4xlarge",
    "p3.2xlarge",
    "p4d.24xlarge",
    "p4de.24xlarge",
    "p5.48xlarge",
    "p5e.48xlarge",
    "p5en.48xlarge",
    "p6-b200.48xlarge",
    "p6-b300.48xlarge",
]

REGIONS = [
    "us-east-1",
    "us-east-2",
    "us-west-2",
    "eu-west-1",
    "eu-west-2",
    "eu-west-3",
    "eu-north-1",
    "eu-central-1",
]

# Tag put on every probe reservation, used by the startup cleanup sweep.
PROBE_TAG_KEY = "purpose"
PROBE_TAG_VALUE = "capacity-probe"

# Dead-man switch: probe reservations auto-expire even if we crash mid-probe.
PROBE_END_DATE_MINUTES = 30

# Cached probe results older than this are flagged stale in the UI.
STALE_AFTER_SECONDS = int(os.environ.get("STALE_AFTER_SECONDS", "900"))

# Optional background refresh of all types/regions. 0 = disabled (probe on
# demand only). Each full sweep creates+cancels one reservation per offered AZ.
AUTO_REFRESH_MINUTES = int(os.environ.get("AUTO_REFRESH_MINUTES", "0"))

# On the Space this points into the mounted HF bucket (rtrm/gpu-capacity-storage
# on /data) so results survive restarts.
STATE_FILE = os.environ.get(
    "STATE_FILE", os.path.join(os.path.dirname(__file__), "..", "state.json")
)

# Probe history, one JSONL file per day next to the state file.
HISTORY_DIR = os.environ.get(
    "HISTORY_DIR", os.path.join(os.path.dirname(STATE_FILE), "history")
)
HISTORY_LOAD_DAYS = int(os.environ.get("HISTORY_LOAD_DAYS", "90"))

MAX_PROBE_WORKERS = 16