gpu-dev 0.5.9__tar.gz → 0.5.12__tar.gz

{gpu_dev-0.5.9 → gpu_dev-0.5.12}/.github/workflows/no-gitlinks.yml

@@ -14,7 +14,7 @@ jobs:
         uses: actions/checkout@v4
       - name: Ensure no gitlinks are tracked
         run: |
-          gitlinks=$(git ls-files -s | awk "$1 == 160000 {print}")
+          gitlinks=$(git ls-files -s | awk '$1 == 160000 {print}')
           if [ -n "$gitlinks" ]; then
             echo "Unexpected gitlinks found:"
             echo "$gitlinks"

{gpu_dev-0.5.9 → gpu_dev-0.5.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpu-dev
-Version: 0.5.9
+Version: 0.5.12
 Summary: CLI tool for PyTorch GPU developer server reservations
 Author: PyTorch Team
 Requires-Python: >=3.10

{gpu_dev-0.5.9 → gpu_dev-0.5.12}/cli-tools/gpu-dev-cli/gpu_dev.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpu-dev
-Version: 0.5.9
+Version: 0.5.12
 Summary: CLI tool for PyTorch GPU developer server reservations
 Author: PyTorch Team
 Requires-Python: >=3.10

{gpu_dev-0.5.9 → gpu_dev-0.5.12}/cli-tools/gpu-dev-cli/gpu_dev.egg-info/SOURCES.txt

@@ -101,6 +101,7 @@ terraform-gpu-devservers/migrations/check_snapshots.py
 terraform-gpu-devservers/migrations/migrate_disks_to_named.py
 terraform-gpu-devservers/migrations/run_backfill.sh
 terraform-gpu-devservers/scripts/CLEANUP_GUIDE.md
+terraform-gpu-devservers/scripts/b200-mig-setup.sh
 terraform-gpu-devservers/scripts/detect_empty_volumes.sh
 terraform-gpu-devservers/scripts/ec2_avail_probe.sh
 terraform-gpu-devservers/scripts/inspect_user_data.sh

{gpu_dev-0.5.9 → gpu_dev-0.5.12}/cli-tools/gpu-dev-cli/gpu_dev_cli/cli.py

@@ -495,9 +495,9 @@ def main(ctx: click.Context) -> None:
     "--gpu-type",
     "-t",
     type=click.Choice(
-        ["b200", "h200", "h100", "h100-mig-1g", "h100-mig-2g", "h100-mig-3g", "a100", "rtxpro6000", "a10g", "t4", "l4", "t4-small", "cpu-arm", "cpu-x86"], case_sensitive=False
+        ["b200", "b200-mig-1g", "b200-mig-2g", "b200-mig-3g", "h200", "h100", "h100-mig-1g", "h100-mig-2g", "h100-mig-3g", "a100", "rtxpro6000", "a10g", "t4", "l4", "t4-small", "cpu-arm", "cpu-x86"], case_sensitive=False
     ),
-    help="GPU type to reserve. Full GPUs: b200, h200, h100, a100, rtxpro6000, a10g, t4, l4, t4-small. H100 MIG slices (partial GPU on a single shared node): h100-mig-1g (10 GB / 1/7 H100 compute), h100-mig-2g (20 GB / 2/7 H100), h100-mig-3g (40 GB / 3/7 H100). CPU only: cpu-arm, cpu-x86.",
+    help="GPU type to reserve. Full GPUs: b200, h200, h100, a100, rtxpro6000, a10g, t4, l4, t4-small. H100 MIG slices: h100-mig-1g (10 GB), h100-mig-2g (20 GB), h100-mig-3g (40 GB). B200 MIG slices (on the mixed B200 node): b200-mig-1g (23 GB), b200-mig-2g (45 GB), b200-mig-3g (90 GB). CPU: cpu-arm, cpu-x86.",
 )
 @click.option(
     "--hours",
@@ -656,6 +656,9 @@ def reserve(
             "h100-mig-1g": {"max_gpus": 16, "instance_type": "p5.48xlarge"},
             "h100-mig-2g": {"max_gpus": 8, "instance_type": "p5.48xlarge"},
             "h100-mig-3g": {"max_gpus": 8, "instance_type": "p5.48xlarge"},
+            "b200-mig-1g": {"max_gpus": 4, "instance_type": "p6-b200.48xlarge"},
+            "b200-mig-2g": {"max_gpus": 2, "instance_type": "p6-b200.48xlarge"},
+            "b200-mig-3g": {"max_gpus": 2, "instance_type": "p6-b200.48xlarge"},
             "h200": {"max_gpus": 8, "instance_type": "p5e.48xlarge"},
             "b200": {"max_gpus": 8, "instance_type": "p6-b200.48xlarge"},
             "cpu-arm": {"max_gpus": 0, "instance_type": "c7g.4xlarge"},
@@ -2454,6 +2457,9 @@ def _show_availability() -> None:
                 "h100-mig-1g": "Hopper (sm90, MIG 10GB)",
                 "h100-mig-2g": "Hopper (sm90, MIG 20GB)",
                 "h100-mig-3g": "Hopper (sm90, MIG 40GB)",
+                "b200-mig-1g": "Blackwell (sm100, MIG 23GB)",
+                "b200-mig-2g": "Blackwell (sm100, MIG 45GB)",
+                "b200-mig-3g": "Blackwell (sm100, MIG 90GB)",
                 "t4": "Turing (sm75)",
                 "cpu-x86": "CPU (x86_64)",
                 "cpu-arm": "CPU (arm64)",
@@ -2462,6 +2468,9 @@ def _show_availability() -> None:
             # Sort order: newest GPU architectures first, then CPUs at the bottom
             arch_priority = {
                 "Blackwell (sm100)": 0,
+                "Blackwell (sm100, MIG 90GB)": 0,
+                "Blackwell (sm100, MIG 45GB)": 0,
+                "Blackwell (sm100, MIG 23GB)": 0,
                 "Blackwell (sm120)": 0,
                 "Hopper (sm90)": 1,
                 "Hopper (sm90, MIG 40GB)": 1,
@@ -2609,6 +2618,9 @@ def _show_availability_watch(interval: int) -> None:
                             "h100-mig-1g": "Hopper (sm90, MIG 10GB)",
                             "h100-mig-2g": "Hopper (sm90, MIG 20GB)",
                             "h100-mig-3g": "Hopper (sm90, MIG 40GB)",
+                            "b200-mig-1g": "Blackwell (sm100, MIG 23GB)",
+                            "b200-mig-2g": "Blackwell (sm100, MIG 45GB)",
+                            "b200-mig-3g": "Blackwell (sm100, MIG 90GB)",
                             "t4": "Turing (sm75)",
                             "cpu-x86": "CPU (x86_64)",
                             "cpu-arm": "CPU (arm64)",
@@ -2617,6 +2629,9 @@ def _show_availability_watch(interval: int) -> None:
                         # Sort order: newest GPU architectures first, then CPUs at the bottom
                         arch_priority = {
                             "Blackwell (sm100)": 0,
+                            "Blackwell (sm100, MIG 90GB)": 0,
+                            "Blackwell (sm100, MIG 45GB)": 0,
+                            "Blackwell (sm100, MIG 23GB)": 0,
                             "Blackwell (sm120)": 0,
                             "Hopper (sm90)": 1,
                             "Hopper (sm90, MIG 40GB)": 1,

{gpu_dev-0.5.9 → gpu_dev-0.5.12}/cli-tools/gpu-dev-cli/gpu_dev_cli/interactive.py

@@ -64,17 +64,25 @@ def select_gpu_type_interactive(
         if "-mig-" not in gt
     }
 
-    # Aggregate MIG slice availability so we can hint it on the h100 row of this picker.
-    mig_total_available = sum(
-        int(info.get("available", 0))
-        for gt, info in (availability_info or {}).items()
-        if gt.startswith("h100-mig-")
-    )
-    mig_total_capacity = sum(
-        int(info.get("total", 0))
-        for gt, info in (availability_info or {}).items()
-        if gt.startswith("h100-mig-")
-    )
+    # Aggregate MIG slice availability per parent type, hinted on the h100/b200 rows.
+    def _mig_aggregates(parent: str):
+        avail = sum(
+            int(info.get("available", 0))
+            for gt, info in (availability_info or {}).items()
+            if gt.startswith(f"{parent}-mig-")
+        )
+        cap = sum(
+            int(info.get("total", 0))
+            for gt, info in (availability_info or {}).items()
+            if gt.startswith(f"{parent}-mig-")
+        )
+        return avail, cap
+
+    h100_mig_avail, h100_mig_capacity = _mig_aggregates("h100")
+    b200_mig_avail, b200_mig_capacity = _mig_aggregates("b200")
+    # Backwards-compat aliases for the existing h100 row code below.
+    mig_total_available = h100_mig_avail
+    mig_total_capacity = h100_mig_capacity
 
     # Display availability table first
     console.print("\n[cyan]🖥️  GPU Availability:[/cyan]")
@@ -146,6 +154,8 @@ def select_gpu_type_interactive(
                 choice_label += f" - {queue_length} in queue"
             if gpu_type == "h100" and mig_total_capacity > 0:
                 choice_label += f" — also {mig_total_available}/{mig_total_capacity} MIG slices"
+            elif gpu_type == "b200" and b200_mig_capacity > 0:
+                choice_label += f" — also {b200_mig_avail}/{b200_mig_capacity} MIG slices"
 
             choices.append(questionary.Choice(title=choice_label, value=gpu_type))
 
@@ -223,27 +233,31 @@ def select_gpu_count_interactive(
     parent_size_etas = parent_info.get("size_etas", {}) or {}
     _now_ts = int(_time.time())
 
-    # MIG slice submenu: only for h100. Each tuple is (target_gpu_type, gpu_count, gb_label).
+    # MIG slice submenu: h100 (16+8+8 slices/node) or b200 (4+2+2 slices/node).
     mig_options = []
-    if gpu_type == "h100":
-        # Map to internal SKUs; the count menu surfaces 1/2/4 of each slice size.
-        mig_specs = [
-            ("h100-mig-1g", "10GB"),
-            ("h100-mig-2g", "20GB"),
-            ("h100-mig-3g", "40GB"),
-        ]
-        for sku, gb in mig_specs:
-            slice_max = {"h100-mig-1g": 16, "h100-mig-2g": 8, "h100-mig-3g": 8}[sku]
-            free = None
-            if availability_info and sku in availability_info:
-                free = availability_info[sku].get("available", 0)
-            for n in [1, 2, 4]:
-                if n > slice_max:
-                    continue
-                noun = "slice" if n == 1 else "slices"
-                avail_suffix = f"  [{free} free]" if free is not None else ""
-                label = f"{n} × {gb} {noun}{avail_suffix}"
-                mig_options.append((sku, n, label))
+    mig_spec_map = {
+        "h100": [
+            ("h100-mig-1g", "10GB", 16),
+            ("h100-mig-2g", "20GB", 8),
+            ("h100-mig-3g", "40GB", 8),
+        ],
+        "b200": [
+            ("b200-mig-1g", "23GB", 4),
+            ("b200-mig-2g", "45GB", 2),
+            ("b200-mig-3g", "90GB", 2),
+        ],
+    }
+    for sku, gb, slice_max in mig_spec_map.get(gpu_type, []):
+        free = None
+        if availability_info and sku in availability_info:
+            free = availability_info[sku].get("available", 0)
+        for n in [1, 2, 4]:
+            if n > slice_max:
+                continue
+            noun = "slice" if n == 1 else "slices"
+            avail_suffix = f"  [{free} free]" if free is not None else ""
+            label = f"{n} × {gb} {noun}{avail_suffix}"
+            mig_options.append((sku, n, label))
 
     # Filter single-node by actual max for this GPU type
     valid_counts = [count for count in valid_counts if count <= max_gpus]

{gpu_dev-0.5.9 → gpu_dev-0.5.12}/cli-tools/gpu-dev-cli/gpu_dev_cli/reservations.py

@@ -543,6 +543,9 @@ class ReservationManager:
                 "h100-mig-1g": {"max_gpus": 16},
                 "h100-mig-2g": {"max_gpus": 8},
                 "h100-mig-3g": {"max_gpus": 8},
+                "b200-mig-1g": {"max_gpus": 4},
+                "b200-mig-2g": {"max_gpus": 2},
+                "b200-mig-3g": {"max_gpus": 2},
                 "h200": {"max_gpus": 8},
                 "b200": {"max_gpus": 8},
             }

{gpu_dev-0.5.9 → gpu_dev-0.5.12}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "gpu-dev"
-version = "0.5.9"
+version = "0.5.12"
 description = "CLI tool for PyTorch GPU developer server reservations"
 authors = [{name = "PyTorch Team"}]
 readme = "cli-tools/gpu-dev-cli/README.md"

{gpu_dev-0.5.9 → gpu_dev-0.5.12}/terraform-gpu-devservers/lambda/reservation_processor/index.py

@@ -71,6 +71,10 @@ GPU_CONFIG = {
     "h100-mig-1g": {"instance_type": "p5.48xlarge", "max_gpus": 16, "cpus": 192, "memory_gb": 2048, "efa_count": 0, "k8s_resource": "nvidia.com/mig-1g.10gb", "node_gpu_type": "h100"},
     "h100-mig-2g": {"instance_type": "p5.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048, "efa_count": 0, "k8s_resource": "nvidia.com/mig-2g.20gb", "node_gpu_type": "h100"},
     "h100-mig-3g": {"instance_type": "p5.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048, "efa_count": 0, "k8s_resource": "nvidia.com/mig-3g.40gb", "node_gpu_type": "h100"},
+    # B200 MIG slices on the b200-6full-2mig-balanced node (6 full GPUs + 2 partitioned per node).
+    "b200-mig-1g": {"instance_type": "p6-b200.48xlarge", "max_gpus": 4, "cpus": 192, "memory_gb": 2048, "efa_count": 0, "k8s_resource": "nvidia.com/mig-1g.23gb", "node_gpu_type": "b200"},
+    "b200-mig-2g": {"instance_type": "p6-b200.48xlarge", "max_gpus": 2, "cpus": 192, "memory_gb": 2048, "efa_count": 0, "k8s_resource": "nvidia.com/mig-2g.45gb", "node_gpu_type": "b200"},
+    "b200-mig-3g": {"instance_type": "p6-b200.48xlarge", "max_gpus": 2, "cpus": 192, "memory_gb": 2048, "efa_count": 0, "k8s_resource": "nvidia.com/mig-3g.90gb", "node_gpu_type": "b200"},
     "t4-small": {"instance_type": "g4dn.2xlarge", "max_gpus": 1, "cpus": 8, "memory_gb": 32, "efa_count": 0},
     "g5g": {"instance_type": "g5g.2xlarge", "max_gpus": 2, "cpus": 8, "memory_gb": 32, "efa_count": 0},
     "a100": {"instance_type": "p4d.24xlarge", "max_gpus": 8, "cpus": 96, "memory_gb": 1152, "efa_count": 4},
@@ -323,12 +327,12 @@ def get_target_az_for_reservation(gpu_type, gpus_requested):
             return target_az
 
         logger.warning(f"No ready/schedulable {gpu_type} nodes found in cluster")
-        return None
+        return None, None
 
     except Exception as e:
         logger.error(f"Error determining target AZ for {gpu_type}: {str(e)}")
-        # Fallback to primary AZ if detection fails
-        return PRIMARY_AVAILABILITY_ZONE
+        # Fallback to primary AZ if detection fails (no node hint — let k8s pick).
+        return PRIMARY_AVAILABILITY_ZONE, None
 
 
 def check_for_multiple_volumes(user_id):
@@ -2167,7 +2171,8 @@ def validate_reservation_request(request: dict[str, Any]) -> tuple[bool, str]:
     # Validate GPU type
     valid_gpu_types = ["t4", "l4", "a10g", "rtxpro6000", "t4-small", "a100",
                        "h100", "h100-mig-1g", "h100-mig-2g", "h100-mig-3g",
-                       "h200", "b200", "cpu-arm", "cpu-x86"]
+                       "h200", "b200", "b200-mig-1g", "b200-mig-2g", "b200-mig-3g",
+                       "cpu-arm", "cpu-x86"]
     if gpu_type not in valid_gpu_types:
         error_msg = f"Invalid GPU type: {gpu_type}. Must be one of: {', '.join(valid_gpu_types)}"
         logger.error(error_msg)
@@ -2408,6 +2413,9 @@ def update_gpu_availability_table(
             "h100-mig-1g": {"gpus_per_instance": 16},
             "h100-mig-2g": {"gpus_per_instance": 8},
             "h100-mig-3g": {"gpus_per_instance": 8},
+            "b200-mig-1g": {"gpus_per_instance": 4},
+            "b200-mig-2g": {"gpus_per_instance": 2},
+            "b200-mig-3g": {"gpus_per_instance": 2},
             "h200": {"gpus_per_instance": 8},
             "b200": {"gpus_per_instance": 8},
         }

{gpu_dev-0.5.9 → gpu_dev-0.5.12}/terraform-gpu-devservers/lambda.tf

@@ -180,8 +180,8 @@ resource "aws_lambda_function" "reservation_processor" {
       HOSTED_ZONE_ID                     = local.effective_domain_name != "" ? local.hosted_zone_id : ""
       SSH_DOMAIN_MAPPINGS_TABLE          = local.effective_domain_name != "" ? aws_dynamodb_table.ssh_domain_mappings.name : ""
       SSL_CERTIFICATE_ARN                = local.effective_domain_name != "" ? aws_acm_certificate.wildcard[0].arn : ""
-      LAMBDA_VERSION                     = "0.5.9"
-      MIN_CLI_VERSION                    = "0.5.5"
+      LAMBDA_VERSION                     = "0.5.12"
+      MIN_CLI_VERSION                    = "0.5.9"
       DISK_CONTENTS_BUCKET               = aws_s3_bucket.disk_contents.bucket
       OPERATIONS_TABLE                   = aws_dynamodb_table.operations.name
     }, local.alb_env_vars)

{gpu_dev-0.5.9 → gpu_dev-0.5.12}/terraform-gpu-devservers/main.tf

@@ -255,6 +255,46 @@ locals {
           k8s_resource        = "nvidia.com/mig-3g.40gb"
           node_gpu_type       = "h100"
         }
+        # B200 MIG slices — virtual SKUs backed by ONE B200 node labelled with the custom
+        # mig_profile "b200-6full-2mig-balanced": GPUs 0-5 stay as full B200 (still reservable
+        # via --gpu-type b200), GPUs 6-7 get partitioned per-GPU into 2x1g.23gb + 1x2g.45gb +
+        # 1x3g.90gb. Per node: 6 full + 4 small + 2 medium + 2 large slices.
+        "b200-mig-1g" = {
+          instance_type       = null
+          instance_types      = null
+          instance_count      = 0
+          gpus_per_instance   = 4 # 2 partitioned GPUs * 2 slices each
+          use_placement_group = false
+          architecture        = "x86_64"
+          efa_network_cards   = 0
+          virtual             = true
+          k8s_resource        = "nvidia.com/mig-1g.23gb"
+          node_gpu_type       = "b200"
+        }
+        "b200-mig-2g" = {
+          instance_type       = null
+          instance_types      = null
+          instance_count      = 0
+          gpus_per_instance   = 2 # 2 partitioned GPUs * 1 slice each
+          use_placement_group = false
+          architecture        = "x86_64"
+          efa_network_cards   = 0
+          virtual             = true
+          k8s_resource        = "nvidia.com/mig-2g.45gb"
+          node_gpu_type       = "b200"
+        }
+        "b200-mig-3g" = {
+          instance_type       = null
+          instance_types      = null
+          instance_count      = 0
+          gpus_per_instance   = 2 # 2 partitioned GPUs * 1 slice each
+          use_placement_group = false
+          architecture        = "x86_64"
+          efa_network_cards   = 0
+          virtual             = true
+          k8s_resource        = "nvidia.com/mig-3g.90gb"
+          node_gpu_type       = "b200"
+        }
         "cpu-arm" = {
           instance_type       = "c7g.8xlarge"
           instance_types      = null

gpu_dev-0.5.12/terraform-gpu-devservers/scripts/b200-mig-setup.sh

@@ -0,0 +1,75 @@
+#!/bin/bash
+# Post-deploy setup for B200 MIG split (6 full + 2 partitioned per node).
+# Run ONCE after PR #77 is merged + tf applied + the new docker/lambda is live.
+
+set -e
+
+NS=gpu-operator
+CM=default-mig-parted-config
+PROFILE_NAME=b200-6full-2mig-balanced
+
+echo "=== Checking current MIG profile in ConfigMap ==="
+if kubectl -n "$NS" get configmap "$CM" -o jsonpath='{.data.config\.yaml}' | grep -q "$PROFILE_NAME:"; then
+    echo "Profile $PROFILE_NAME already present — skipping ConfigMap edit"
+else
+    echo "Profile $PROFILE_NAME missing. Patching ConfigMap..."
+
+    # Save current ConfigMap content
+    kubectl -n "$NS" get configmap "$CM" -o yaml > /tmp/mig-config-backup.yaml
+    echo "Backup saved to /tmp/mig-config-backup.yaml"
+
+    # Append our profile under mig-configs:
+    # NOTE: this is a sed-driven append. ClusterPolicy's controller MAY revert this if it
+    # reconciles. If you see the profile disappear, re-run this script. If it keeps reverting,
+    # we'll need to fork the ConfigMap (next iteration).
+    kubectl -n "$NS" get configmap "$CM" -o jsonpath='{.data.config\.yaml}' > /tmp/mig-config.yaml
+
+    cat >> /tmp/mig-config.yaml <<'EOF'
+
+  # Mixed B200 split: GPUs 0-5 stay full (reservable as --gpu-type b200), GPUs 6-7 partitioned.
+  # Per partitioned GPU: 2x 1g.23gb + 1x 2g.45gb + 1x 3g.90gb. Per node: 6 full + 4 small + 2 medium + 2 large.
+  b200-6full-2mig-balanced:
+    - device-filter: ["0x290110DE"]
+      devices: [0, 1, 2, 3, 4, 5]
+      mig-enabled: false
+    - device-filter: ["0x290110DE"]
+      devices: [6, 7]
+      mig-enabled: true
+      mig-devices:
+        "1g.23gb": 2
+        "2g.45gb": 1
+        "3g.90gb": 1
+EOF
+
+    # Re-encode and patch
+    kubectl -n "$NS" create configmap "$CM" --from-file=config.yaml=/tmp/mig-config.yaml --dry-run=client -o yaml \
+        | kubectl -n "$NS" patch configmap "$CM" --patch-file=/dev/stdin
+    echo "ConfigMap patched."
+fi
+
+echo
+echo "=== Picking a B200 node to label ==="
+NODE=$(kubectl get nodes -l GpuType=b200 -o jsonpath='{.items[0].metadata.name}')
+if [ -z "$NODE" ]; then
+    echo "No B200 nodes found. Exiting."
+    exit 1
+fi
+echo "Will label: $NODE"
+read -p "Proceed? (y/N): " CONFIRM
+if [ "$CONFIRM" != "y" ]; then
+    echo "Aborted."
+    exit 0
+fi
+
+kubectl label node "$NODE" "nvidia.com/mig.config=$PROFILE_NAME" --overwrite
+echo "Node labelled. nvidia-mig-manager will partition GPUs 6-7 (drains existing pods if any)."
+echo
+echo "Watch progress with:"
+echo "  kubectl logs -n gpu-operator -l app=nvidia-mig-manager -f"
+echo "  kubectl get node $NODE -o jsonpath='{.status.allocatable}' | jq ."
+echo
+echo "After ~2-5 min, allocatable should show:"
+echo "  nvidia.com/gpu:           6"
+echo "  nvidia.com/mig-1g.23gb:   4"
+echo "  nvidia.com/mig-2g.45gb:   2"
+echo "  nvidia.com/mig-3g.90gb:   2"