diff --git a/cli-tools/gpu-dev-cli/gpu_dev_cli/cli.py b/cli-tools/gpu-dev-cli/gpu_dev_cli/cli.py index 1fe2a4fe..642464a8 100644 --- a/cli-tools/gpu-dev-cli/gpu_dev_cli/cli.py +++ b/cli-tools/gpu-dev-cli/gpu_dev_cli/cli.py @@ -2441,14 +2441,17 @@ def _show_availability() -> None: available = info.get("available", 0) max_reservable = info.get("max_reservable", 0) total = info.get("total", 0) + scalable_total = info.get("scalable_total", 0) full_nodes_available = info.get("full_nodes_available", 0) gpus_per_instance = info.get("gpus_per_instance", 0) queue_length = info.get("queue_length", 0) est_wait = info.get("estimated_wait_minutes", 0) - # Format wait time + # Format wait time — for autoscaled CPU types, show scale-up estimate if available > 0: wait_display = "Available now" + elif scalable_total > 0 and available == 0: + wait_display = "~3min (scaling up)" elif est_wait == 0: wait_display = "Unknown" elif est_wait < 60: @@ -2461,10 +2464,13 @@ def _show_availability() -> None: else: wait_display = f"{hours}h {minutes}min" + # Show total as "current / scalable" for autoscaled types + if scalable_total > 0: + total_display = f"{total} / {scalable_total}" + else: + total_display = str(total) + # Color code availability based on full nodes available - # Red: 0 GPUs available - # Yellow: Some GPUs available but no full node - # Green: At least one full node available if available == 0: available_display = f"[red]{available}[/red]" elif full_nodes_available > 0: @@ -2476,7 +2482,7 @@ def _show_availability() -> None: gpu_type.upper(), available_display, str(max_reservable), - str(total), + total_display, str(queue_length), arch, wait_display, @@ -2583,12 +2589,15 @@ def _show_availability_watch(interval: int) -> None: last_arch = arch available = info.get("available", 0) total = info.get("total", 0) + scalable_total = info.get("scalable_total", 0) queue_length = info.get("queue_length", 0) est_wait = info.get("estimated_wait_minutes", 0) # Format wait time if available > 0: wait_display = "Available now" + elif scalable_total > 0 and available == 0: + wait_display = "~3min (scaling up)" elif est_wait == 0: wait_display = "Unknown" elif est_wait < 60: @@ -2601,6 +2610,12 @@ def _show_availability_watch(interval: int) -> None: else: wait_display = f"{hours}h {minutes}min" + # Show total as "current / scalable" for autoscaled types + if scalable_total > 0: + total_display = f"{total} / {scalable_total}" + else: + total_display = str(total) + # Color code availability if available > 0: available_display = f"[green]{available}[/green]" @@ -2610,7 +2625,7 @@ def _show_availability_watch(interval: int) -> None: table.add_row( gpu_type.upper(), available_display, - str(total), + total_display, str(queue_length), arch, wait_display, diff --git a/cli-tools/gpu-dev-cli/gpu_dev_cli/reservations.py b/cli-tools/gpu-dev-cli/gpu_dev_cli/reservations.py index f2d4866b..00bcd1d4 100644 --- a/cli-tools/gpu-dev-cli/gpu_dev_cli/reservations.py +++ b/cli-tools/gpu-dev-cli/gpu_dev_cli/reservations.py @@ -959,6 +959,7 @@ def get_gpu_availability_by_type(self) -> Optional[Dict[str, Dict[str, Any]]]: availability_info[gpu_type] = { "available": int(item.get("available_gpus", 0)), "total": int(item.get("total_gpus", 0)), + "scalable_total": int(item.get("scalable_total", 0)), "max_reservable": int(item.get("max_reservable", 0)), "full_nodes_available": int(item.get("full_nodes_available", 0)), "gpus_per_instance": int(item.get("gpus_per_instance", 0)), diff --git a/terraform-gpu-devservers/availability.tf b/terraform-gpu-devservers/availability.tf index 738912f8..6361622b 100644 --- a/terraform-gpu-devservers/availability.tf +++ b/terraform-gpu-devservers/availability.tf @@ -106,7 +106,16 @@ resource "aws_iam_role_policy" "availability_updater_policy" { { Effect = "Allow" Action = [ - "autoscaling:DescribeAutoScalingGroups" + "autoscaling:DescribeAutoScalingGroups", + "autoscaling:SetDesiredCapacity", + "autoscaling:SetInstanceProtection" + ] + Resource = "*" + }, + { + Effect = "Allow" + Action = [ + "ec2:DescribeInstances" ] Resource = "*" }, diff --git a/terraform-gpu-devservers/eks.tf b/terraform-gpu-devservers/eks.tf index 5808cfeb..0c53f57c 100644 --- a/terraform-gpu-devservers/eks.tf +++ b/terraform-gpu-devservers/eks.tf @@ -229,8 +229,9 @@ resource "aws_autoscaling_group" "gpu_dev_nodes" { health_check_grace_period = 300 # Use dynamic instance count from capacity reservation - min_size = each.value.instance_count - max_size = each.value.instance_count + # CPU types with min/max_instance_count support autoscaling (Lambda adjusts desired_capacity) + min_size = try(each.value.gpu_config.min_instance_count, each.value.instance_count) + max_size = try(each.value.gpu_config.max_instance_count, each.value.instance_count) desired_capacity = each.value.instance_count # Don't wait for instances to become healthy - prevents Terraform failures when AWS can't place instances @@ -305,6 +306,12 @@ resource "aws_autoscaling_group" "gpu_dev_nodes" { value = each.value.capacity_reservation_id != null ? each.value.capacity_reservation_id : "none" propagate_at_launch = true } + + # Prevent terraform from resetting desired_capacity on apply — Lambda manages it for autoscaled CPU ASGs. + # Safe for GPU ASGs too since their min=max=desired so desired_capacity never drifts. + lifecycle { + ignore_changes = [desired_capacity] + } } diff --git a/terraform-gpu-devservers/lambda/availability_updater/index.py b/terraform-gpu-devservers/lambda/availability_updater/index.py index 2b4605ae..12720381 100644 --- a/terraform-gpu-devservers/lambda/availability_updater/index.py +++ b/terraform-gpu-devservers/lambda/availability_updater/index.py @@ -5,8 +5,10 @@ import json import logging +import math import os -from typing import Dict, Any +import time +from typing import Dict, Any, List import boto3 @@ -124,21 +126,24 @@ def update_gpu_availability(gpu_type: str, k8s_client=None) -> None: logger.info( f"CPU ASG calculation: {running_instances} instances * {max_users_per_node} slots = {total_gpus} total slots") + # Track per-node pod counts for autoscaling decisions + node_pod_counts = {} # node_name -> gpu_dev_pod_count + # Check actual pod usage on CPU nodes if k8s_client is not None: try: logger.info(f"Checking CPU node availability for {gpu_type}") - # Count available slots by checking pod count on each node + from kubernetes import client v1 = client.CoreV1Api(k8s_client) nodes = v1.list_node(label_selector=f"GpuType={gpu_type}") total_available_slots = 0 for node in nodes.items: if is_node_ready_and_schedulable(node): - # Count gpu-dev pods on this node pods = v1.list_pod_for_all_namespaces(field_selector=f"spec.nodeName={node.metadata.name}") gpu_dev_pods = [p for p in pods.items if p.metadata.name.startswith('gpu-dev-')] used_slots = len(gpu_dev_pods) + node_pod_counts[node.metadata.name] = used_slots available_slots = max(0, max_users_per_node - used_slots) total_available_slots += available_slots @@ -149,6 +154,28 @@ def update_gpu_availability(gpu_type: str, k8s_client=None) -> None: available_gpus = total_gpus else: available_gpus = total_gpus + + # --- CPU autoscaling logic --- + asg_min_size = matching_asgs[0]["MinSize"] + asg_max_size = matching_asgs[0]["MaxSize"] + current_desired = matching_asgs[0]["DesiredCapacity"] + asg_name_for_scaling = matching_asgs[0]["AutoScalingGroupName"] + + # Only autoscale if min != max (autoscaling is enabled for this ASG) + if asg_min_size < asg_max_size: + try: + autoscale_cpu_asg( + asg_name=asg_name_for_scaling, + current_desired=current_desired, + asg_min_size=asg_min_size, + asg_max_size=asg_max_size, + total_available_slots=available_gpus, + max_users_per_node=max_users_per_node, + node_pod_counts=node_pod_counts, + matching_asgs=matching_asgs, + ) + except Exception as scale_error: + logger.error(f"CPU autoscaling failed for {gpu_type}: {scale_error}") else: # GPU nodes - use existing logic total_gpus = running_instances * gpus_per_instance @@ -223,6 +250,13 @@ def update_gpu_availability(gpu_type: str, k8s_client=None) -> None: full_nodes_available = available_gpus # Each "GPU" represents one CPU node slot max_reservable = 1 if available_gpus > 0 else 0 # Max 1 CPU node per reservation + # For CPU types with autoscaling, report scalable total based on max ASG size + scalable_total = 0 + if is_cpu_type and matching_asgs: + asg_max = matching_asgs[0].get("MaxSize", 0) + if asg_max > matching_asgs[0].get("MinSize", 0): + scalable_total = asg_max * max_users_per_node + # Update DynamoDB table table = dynamodb.Table(AVAILABILITY_TABLE) @@ -231,15 +265,13 @@ def update_gpu_availability(gpu_type: str, k8s_client=None) -> None: "gpu_type": gpu_type, "total_gpus": total_gpus, "available_gpus": available_gpus, + "scalable_total": scalable_total, "max_reservable": max_reservable, "full_nodes_available": full_nodes_available, "running_instances": running_instances, "desired_capacity": desired_capacity, "gpus_per_instance": gpus_per_instance, - "last_updated": context.aws_request_id - if "context" in locals() - else "unknown", - "last_updated_timestamp": int(time.time()) if "time" in dir() else 0, + "last_updated_timestamp": int(time.time()), } ) @@ -252,9 +284,6 @@ def update_gpu_availability(gpu_type: str, k8s_client=None) -> None: raise -import time - - def check_schedulable_gpus_for_type(k8s_client, gpu_type: str) -> int: """Check how many GPUs of a specific type are schedulable (available for new pods)""" try: @@ -365,3 +394,108 @@ def get_available_gpus_on_node(v1_api, node) -> int: f"Error getting available GPUs on node {node.metadata.name}: {str(e)}" ) return 0 + + +# --- CPU Autoscaling --- + +MIN_SPARE_SLOTS = 2 # Minimum spare CPU slots to keep available +SCALE_DOWN_HYSTERESIS = 3 # Extra spare slots above MIN before scaling down (one full node worth) + + +def autoscale_cpu_asg( + asg_name: str, + current_desired: int, + asg_min_size: int, + asg_max_size: int, + total_available_slots: int, + max_users_per_node: int, + node_pod_counts: Dict[str, int], + matching_asgs: List[Dict], +) -> None: + """Scale CPU ASG up/down based on spare slot availability""" + logger.info( + f"Autoscale check: asg={asg_name} desired={current_desired} " + f"available_slots={total_available_slots} min={asg_min_size} max={asg_max_size}" + ) + + # Scale UP: fewer spare slots than minimum buffer + if total_available_slots < MIN_SPARE_SLOTS: + slots_needed = MIN_SPARE_SLOTS - total_available_slots + nodes_to_add = math.ceil(slots_needed / max_users_per_node) + new_desired = min(current_desired + nodes_to_add, asg_max_size) + if new_desired > current_desired: + logger.info(f"Scaling UP {asg_name}: {current_desired} -> {new_desired} (need {slots_needed} more slots)") + autoscaling.set_desired_capacity( + AutoScalingGroupName=asg_name, + DesiredCapacity=new_desired, + ) + return + + # Scale DOWN: more spare slots than needed (with hysteresis to avoid flapping) + scale_down_threshold = MIN_SPARE_SLOTS + max_users_per_node + SCALE_DOWN_HYSTERESIS + if total_available_slots > scale_down_threshold and current_desired > asg_min_size: + # Protect nodes that have active gpu-dev pods, unprotect empty ones + _update_instance_protection(matching_asgs, node_pod_counts) + + excess_slots = total_available_slots - MIN_SPARE_SLOTS + nodes_to_remove = excess_slots // max_users_per_node + new_desired = max(current_desired - nodes_to_remove, asg_min_size) + if new_desired < current_desired: + logger.info(f"Scaling DOWN {asg_name}: {current_desired} -> {new_desired} ({excess_slots} excess slots)") + autoscaling.set_desired_capacity( + AutoScalingGroupName=asg_name, + DesiredCapacity=new_desired, + ) + return + + logger.info(f"No scaling action needed for {asg_name}") + + +def _update_instance_protection(matching_asgs: List[Dict], node_pod_counts: Dict[str, int]) -> None: + """Set instance protection on nodes with active pods, remove from empty nodes""" + for asg in matching_asgs: + asg_name = asg["AutoScalingGroupName"] + in_service_instances = [ + inst for inst in asg["Instances"] + if inst["LifecycleState"] == "InService" + ] + + if not in_service_instances: + continue + + # Build instance_id -> node_name mapping via EC2 private DNS + ec2 = boto3.client("ec2") + instance_ids = [inst["InstanceId"] for inst in in_service_instances] + ec2_response = ec2.describe_instances(InstanceIds=instance_ids) + + instance_node_map = {} + for reservation in ec2_response.get("Reservations", []): + for instance in reservation.get("Instances", []): + # K8s node name is the EC2 private DNS name + private_dns = instance.get("PrivateDnsName", "") + instance_node_map[instance["InstanceId"]] = private_dns + + protect_ids = [] + unprotect_ids = [] + for instance_id, node_name in instance_node_map.items(): + pod_count = node_pod_counts.get(node_name, 0) + if pod_count > 0: + protect_ids.append(instance_id) + else: + unprotect_ids.append(instance_id) + + if protect_ids: + logger.info(f"Setting instance protection on {len(protect_ids)} instances with active pods") + autoscaling.set_instance_protection( + InstanceIds=protect_ids, + AutoScalingGroupName=asg_name, + ProtectedFromScaleIn=True, + ) + + if unprotect_ids: + logger.info(f"Removing instance protection from {len(unprotect_ids)} empty instances") + autoscaling.set_instance_protection( + InstanceIds=unprotect_ids, + AutoScalingGroupName=asg_name, + ProtectedFromScaleIn=False, + ) diff --git a/terraform-gpu-devservers/main.tf b/terraform-gpu-devservers/main.tf index 635a9d07..1f24895d 100644 --- a/terraform-gpu-devservers/main.tf +++ b/terraform-gpu-devservers/main.tf @@ -95,7 +95,9 @@ locals { "cpu-arm" = { instance_type = "c7g.4xlarge" instance_types = null - instance_count = 30 + instance_count = 2 + min_instance_count = 2 + max_instance_count = 30 gpus_per_instance = 0 use_placement_group = false architecture = "arm64" @@ -103,7 +105,9 @@ locals { "cpu-x86" = { instance_type = "c7i.4xlarge" instance_types = null - instance_count = 30 + instance_count = 2 + min_instance_count = 2 + max_instance_count = 30 gpus_per_instance = 0 use_placement_group = false architecture = "x86_64" @@ -221,7 +225,9 @@ locals { "cpu-arm" = { instance_type = "c7g.8xlarge" instance_types = null - instance_count = 30 + instance_count = 2 + min_instance_count = 2 + max_instance_count = 30 gpus_per_instance = 0 use_placement_group = false architecture = "arm64" @@ -229,7 +235,9 @@ locals { "cpu-x86" = { instance_type = "c7i.8xlarge" instance_types = null - instance_count = 30 + instance_count = 2 + min_instance_count = 2 + max_instance_count = 30 gpus_per_instance = 0 use_placement_group = false architecture = "x86_64"