32a241c28f
Some checks failed
nightly / Test against latest dependencies (py3.10) (push) Has been cancelled
nightly / Test against latest dependencies (py3.13) (push) Has been cancelled
tests / tests (3.13, locked) (push) Has been cancelled
tests / tests (3.13, unlocked) (push) Has been cancelled
tests / pyright (3.10) (push) Has been cancelled
tests / lint-format (push) Has been cancelled
tests / tests (3.10, locked) (push) Has been cancelled
tests / tests (3.11, locked) (push) Has been cancelled
tests / tests (3.12, locked) (push) Has been cancelled
tests / pyright (3.11) (push) Has been cancelled
tests / pyright (3.12) (push) Has been cancelled
tests / pyright (3.13) (push) Has been cancelled
tests / ty-check (3.10) (push) Has been cancelled
tests / ty-check (3.11) (push) Has been cancelled
tests / ty-check (3.12) (push) Has been cancelled
tests / ty-check (3.13) (push) Has been cancelled
tests / stubs (push) Has been cancelled
tests / smoke-test (push) Has been cancelled
Docker / check_paths (push) Has been cancelled
docs / build (push) Has been cancelled
Docker / build (push) Has been cancelled
Upstream: https://github.com/michaelgillett/mjlab Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047 Upstream-Branch: main
62 lines
1.7 KiB
Bash
62 lines
1.7 KiB
Bash
#!/usr/bin/env bash
|
|
# Launch a W&B sweep via SkyPilot.
|
|
#
|
|
# Provisions a single multi-GPU cluster and runs one sweep agent per GPU
|
|
# using SkyPilot's job queue. Each agent pulls hyperparameters from the
|
|
# W&B sweep controller independently.
|
|
#
|
|
# Usage:
|
|
# ./scripts/cloud/sweep-launch.sh [GPUS [CLOUD]]
|
|
#
|
|
# Examples:
|
|
# ./scripts/cloud/sweep-launch.sh A100:4 # 4 agents, default cloud
|
|
# ./scripts/cloud/sweep-launch.sh A100:8 gcp # 8 agents on GCP
|
|
# ./scripts/cloud/sweep-launch.sh A100:8 lambda # 8 agents on Lambda
|
|
# ./scripts/cloud/sweep-launch.sh # defaults to A100:4
|
|
|
|
set -euo pipefail
|
|
|
|
GPUS="${1:-A100:4}"
|
|
CLOUD="${2:-}"
|
|
GPU_TYPE="${GPUS%%:*}"
|
|
NUM_AGENTS="${GPUS##*:}"
|
|
CLUSTER_NAME="mjlab-sweep"
|
|
|
|
echo "Creating W&B sweep..."
|
|
SWEEP_ID=$(uv run wandb sweep scripts/cloud/sweep.yaml 2>&1 | grep "wandb agent" | awk '{print $NF}')
|
|
|
|
if [ -z "$SWEEP_ID" ]; then
|
|
echo "Failed to create sweep."
|
|
exit 1
|
|
fi
|
|
|
|
echo "Sweep created: $SWEEP_ID"
|
|
echo "Provisioning $GPUS cluster..."
|
|
|
|
# Provision the cluster and run setup (no run section in this YAML).
|
|
CLOUD_FLAG=${CLOUD:+--cloud "$CLOUD"}
|
|
sky launch scripts/cloud/sweep-cluster.yaml \
|
|
-c "$CLUSTER_NAME" \
|
|
--gpus "$GPUS" \
|
|
${CLOUD_FLAG} \
|
|
-y --retry-until-up
|
|
|
|
echo "Submitting $NUM_AGENTS agents to job queue..."
|
|
|
|
for i in $(seq 1 "$NUM_AGENTS"); do
|
|
echo " Agent $i/$NUM_AGENTS"
|
|
sky exec "$CLUSTER_NAME" \
|
|
--gpus "${GPU_TYPE}:1" \
|
|
--env "SWEEP_ID=$SWEEP_ID" \
|
|
-d \
|
|
scripts/cloud/sweep-agent.yaml
|
|
done
|
|
|
|
echo ""
|
|
echo "All agents launched. Monitor at:"
|
|
echo " sky queue $CLUSTER_NAME"
|
|
echo " sky logs $CLUSTER_NAME <JOB_ID>"
|
|
echo " W&B dashboard: https://wandb.ai/$SWEEP_ID"
|
|
echo ""
|
|
echo "When done: sky down $CLUSTER_NAME"
|