Skip to content

feat: auto-scale CPU nodes (2-30) instead of fixed 60 #37

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wdvr wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: auto-scale CPU nodes (2-30) instead of fixed 60 #37

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 21 additions & 6 deletions cli-tools/gpu-dev-cli/gpu_dev_cli/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand All @@ -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,
Expand Down Expand Up @@ -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:
Expand All @@ -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]"
Expand All @@ -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,
Expand Down
1 change: 1 addition & 0 deletions cli-tools/gpu-dev-cli/gpu_dev_cli/reservations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)),
Expand Down
11 changes: 10 additions & 1 deletion terraform-gpu-devservers/availability.tf
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = "*"
},
Expand Down
11 changes: 9 additions & 2 deletions terraform-gpu-devservers/eks.tf
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]
}
}


Expand Down
154 changes: 144 additions & 10 deletions terraform-gpu-devservers/lambda/availability_updater/index.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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)

Expand All @@ -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()),
}
)

Expand All @@ -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:
Expand Down Expand Up @@ -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,
)
16 changes: 12 additions & 4 deletions terraform-gpu-devservers/main.tf
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,15 +95,19 @@ 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"
}
"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"
Expand Down Expand Up @@ -221,15 +225,19 @@ 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"
}
"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"
Expand Down