Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

scripts/cloud/sweep-launch.sh

@@ -0,0 +1,61 @@
+#!/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"