feat(hyperpod-eks): Local Zone deployment support for Terraform modules

1.architectures/7.sagemaker-hyperpod-eks/terraform-modules/README.md

@@ -491,6 +491,105 @@ After deployment, verify connectivity to AWS services:
 python3 tools/verify-aws-connectivity.py
 ```
 
+---
+
+### Local Zone Deployment
+
+You can place HyperPod worker instance groups in an [AWS Local Zone](https://docs.aws.amazon.com/local-zones/latest/ug/what-is-aws-local-zones.html) to run compute closer to a specific metro area. The EKS control plane stays in the parent Region: EKS cannot create control-plane ENIs in a Local Zone, so only the HyperPod worker subnet lives in the Local Zone while the control-plane subnets remain in standard parent Availability Zones.
+
+Local Zone support is opt-in and additive. All of the variables below default to standard-AZ behavior when unset, so existing deployments are unaffected.
+
+> **Note:** Not every Local Zone is supported by HyperPod, and there is no API to enumerate the supported zones. Confirm your target zone with AWS before deploying. The instance type must also be both offered in the Local Zone and present in HyperPod's `ClusterInstanceType` enum.
+
+#### Local Zone Variables
+
+| Variable | Usage |
+|----------|-------|
+| `private_subnet_availability_zone_ids` | Pins the HyperPod private (worker) subnets to explicit Availability Zone IDs, 1:1 with `private_subnet_cidrs`. This bypasses the `opt-in-status = "opt-in-not-required"` AZ-discovery filter, which excludes opt-in Local Zones. Default `[]` = discover standard AZs automatically. |
+| `local_zone_egress_zone_ids` | List of Local Zone AZ IDs that should get a Local-Zone-local NAT gateway. Default `[]` = worker subnets route through the regional NAT gateway. |
+| `local_zone_public_subnet_cidrs` | Local Zone public subnet CIDRs, 1:1 with `local_zone_egress_zone_ids`. Typically carved from the VPC primary CIDR (secondary CIDRs are usually consumed by the worker subnet). |
+| `local_zone_network_border_groups` | `NetworkBorderGroup` names for the Local Zone NAT Elastic IPs, 1:1 with `local_zone_egress_zone_ids`. Required: a plain VPC-scoped EIP cannot attach to a NAT gateway in a Local Zone subnet. The border group is the Local Zone name minus the trailing zone letter (e.g. `us-west-2-phx-2a` -> `us-west-2-phx-2`). |
+
+#### Local Zone egress (NAT placement)
+
+By default the `vpc` module creates a single regional NAT gateway in a standard-AZ public subnet. A worker subnet in a Local Zone routes `0.0.0.0/0` to that regional NAT, so egress traffic hairpins back to the parent Region and pays an added round trip per packet.
+
+Setting the three `local_zone_*` variables creates one Local-Zone-local NAT gateway per listed zone (with a border-group-scoped EIP) and routes matching worker subnets to it via the `vpc` module's `nat_gateway_ids_by_zone_id` output. Keeping egress in-zone significantly improves first-hop latency and internet throughput for Local Zone workers. Unmapped AZs continue to use the regional NAT. This has been validated with an end-to-end Local Zone HyperPod deployment.
+
+#### Example `custom.tfvars`
+
+```hcl
+resource_name_prefix = "hp-eks"
+aws_region           = "us-west-2"
+
+# VPC
+create_vpc_module    = true
+vpc_cidr             = "10.192.0.0/16"
+public_subnet_1_cidr = "10.192.10.0/24"
+public_subnet_2_cidr = "10.192.11.0/24"
+
+# Private (worker) subnet pinned to the Local Zone AZ ID, 1:1 with the CIDR.
+# This bypasses the opt-in-not-required discovery filter that excludes Local Zones.
+create_private_subnet_module         = true
+private_subnet_cidrs                 = ["10.1.0.0/16"]
+private_subnet_availability_zone_ids = ["usw2-phx2-az1"]
+
+# Optional: Local-Zone-local NAT gateway (all three lists non-empty and 1:1).
+# Uncomment to keep worker egress in-zone instead of hairpinning to the Region.
+# local_zone_egress_zone_ids       = ["usw2-phx2-az1"]
+# local_zone_public_subnet_cidrs   = ["10.192.20.0/24"]
+# local_zone_network_border_groups = ["us-west-2-phx-2"]
+
+# EKS control-plane subnets stay in parent AZs (cannot live in a Local Zone).
+create_eks_module         = true
+create_eks_subnets        = true
+eks_private_subnet_1_cidr = "10.192.7.0/28"
+eks_private_subnet_2_cidr = "10.192.8.0/28"
+
+# FSx placement: default co-locates FSx with the instance group's subnet (in-Local-Zone).
+create_fsx_module         = true
+create_new_fsx_filesystem = true
+fsx_storage_capacity      = 1200
+fsx_throughput            = 250
+# fsx_availability_zone_id = ""  # set to a parent-AZ ID for a cross-zone mount
+                                 # if the Local Zone does not offer FSx (or the tier).
+
+instance_groups = [
+  {
+    name                      = "instance-group-1"
+    instance_type             = "ml.c6i.2xlarge"
+    instance_count            = 1
+    availability_zone_id      = "usw2-phx2-az1" # land workers in the Local Zone
+    ebs_volume_size_in_gb     = 100
+    threads_per_core          = 2
+    enable_stress_check       = false
+    enable_connectivity_check = false
+    lifecycle_script          = "on_create.sh"
+  }
+]
+```
+
+#### FSx for Lustre in a Local Zone
+
+FSx placement is already configurable through `fsx_availability_zone_id` (see the [FSx for Lustre Module](#fsx-for-lustre-module) section). When empty (default), FSx is created in the first instance group's subnet, which co-locates it with compute in the Local Zone. FSx for Lustre availability and per-tier support vary by Local Zone; if your target zone does not offer FSx (or the tier you need), set `fsx_availability_zone_id` to a parent-AZ ID for a cross-zone mount, or set `create_new_fsx_filesystem = false`.
+
+#### Prerequisite: opt in to the Local Zone
+
+The target Local Zone must be opted in before you deploy (a not-yet-opted-in zone makes the private subnet fail to create):
+
+```bash
+# Look up your Local Zone's AZ ID and parent zone
+aws ec2 describe-availability-zones --all-availability-zones \
+  --query "AvailabilityZones[?ZoneType=='local-zone'].[ZoneName,ZoneId,ParentZoneName]" \
+  --output table
+
+# Opt in (opt-in is asynchronous - verify it reports opted-in before deploying)
+aws ec2 modify-availability-zone-group \
+  --group-name us-west-2-phx-2a --opt-in-status opted-in
+```
+
+For a complete, ready-to-run Local Zone example, see the [HyperPod Local Zone quickstart](https://github.com/aravneelaws/hyperpod-local-zone-quickstart/tree/main/terraform/eks). That repository ships only a `local-zone.tfvars` file and applies it against this reference stack (no forked Terraform), so the variable file lives there while the modules live here.
+
 ---
 ### Enabling Optional Addons 
 Set the following parameters to `true` in your `custom.tfvars` file to enable optional addons for your HyperPod cluster (e.g. `create_task_governance_module = true`):

1.architectures/7.sagemaker-hyperpod-eks/terraform-modules/hyperpod-eks-tf/main.tf

@@ -29,18 +29,21 @@ locals {
   ]
   is_amp_allowed = contains(local.amp_allowed_regions, var.aws_region)
 
-  vpc_id                   = var.create_vpc_module ? module.vpc[0].vpc_id : var.existing_vpc_id
-  private_subnet_ids       = var.create_private_subnet_module ? module.private_subnet[0].private_subnet_ids : var.existing_private_subnet_ids
-  security_group_id        = var.create_security_group_module ? module.security_group[0].security_group_id : var.existing_security_group_id
-  eks_cluster_name         = var.create_eks_module ? module.eks_cluster[0].eks_cluster_name : var.existing_eks_cluster_name
-  eks_cluster_arn          = var.create_eks_module ? module.eks_cluster[0].eks_cluster_arn : data.aws_eks_cluster.existing_eks_cluster[0].arn
-  sagemaker_iam_role_name  = var.create_sagemaker_iam_role_module ? module.sagemaker_iam_role[0].sagemaker_iam_role_name : var.existing_sagemaker_iam_role_name
-  create_hyperpod_module   = var.create_hyperpod_module && !(var.create_eks_module && !var.create_helm_chart_module)
-  karpenter_role_arn       = var.create_sagemaker_iam_role_module && length(module.sagemaker_iam_role[0].karpenter_role_arn) > 0 ? module.sagemaker_iam_role[0].karpenter_role_arn[0] : null
-  nat_gateway_id           = var.create_vpc_module ? module.vpc[0].nat_gateway_1_id : var.existing_nat_gateway_id
-  private_route_table_ids  = var.create_private_subnet_module ? module.private_subnet[0].private_route_table_ids : var.existing_private_route_table_ids
-  eks_private_subnet_cidrs = [var.eks_private_subnet_1_cidr, var.eks_private_subnet_2_cidr]
-  enable_guardduty_cleanup = var.enable_guardduty_cleanup && (var.create_vpc_module || var.create_private_subnet_module || var.create_eks_module)
+  vpc_id                  = var.create_vpc_module ? module.vpc[0].vpc_id : var.existing_vpc_id
+  private_subnet_ids      = var.create_private_subnet_module ? module.private_subnet[0].private_subnet_ids : var.existing_private_subnet_ids
+  security_group_id       = var.create_security_group_module ? module.security_group[0].security_group_id : var.existing_security_group_id
+  eks_cluster_name        = var.create_eks_module ? module.eks_cluster[0].eks_cluster_name : var.existing_eks_cluster_name
+  eks_cluster_arn         = var.create_eks_module ? module.eks_cluster[0].eks_cluster_arn : data.aws_eks_cluster.existing_eks_cluster[0].arn
+  sagemaker_iam_role_name = var.create_sagemaker_iam_role_module ? module.sagemaker_iam_role[0].sagemaker_iam_role_name : var.existing_sagemaker_iam_role_name
+  create_hyperpod_module  = var.create_hyperpod_module && !(var.create_eks_module && !var.create_helm_chart_module)
+  karpenter_role_arn      = var.create_sagemaker_iam_role_module && length(module.sagemaker_iam_role[0].karpenter_role_arn) > 0 ? module.sagemaker_iam_role[0].karpenter_role_arn[0] : null
+  nat_gateway_id          = var.create_vpc_module ? module.vpc[0].nat_gateway_1_id : var.existing_nat_gateway_id
+  # Per-AZ NAT map: only populated when this module manages the VPC. When
+  # BYO-VPC, callers currently have to set up LZ-local NATs themselves.
+  nat_gateway_ids_by_zone_id = var.create_vpc_module ? module.vpc[0].nat_gateway_ids_by_zone_id : {}
+  private_route_table_ids    = var.create_private_subnet_module ? module.private_subnet[0].private_route_table_ids : var.existing_private_route_table_ids
+  eks_private_subnet_cidrs   = [var.eks_private_subnet_1_cidr, var.eks_private_subnet_2_cidr]
+  enable_guardduty_cleanup   = var.enable_guardduty_cleanup && (var.create_vpc_module || var.create_private_subnet_module || var.create_eks_module)
 
   # Cilium CNI
   skip_vpc_cni  = var.enable_cilium && var.cilium_mode != "chaining"
@@ -68,28 +71,40 @@ locals {
   create_hyperpod_training_operator_module  = !local.rig_mode && var.create_hyperpod_training_operator_module
   create_observability_module               = !local.rig_mode && var.create_observability_module
   create_hyperpod_inference_operator_module = !local.rig_mode && var.create_hyperpod_inference_operator_module
+
+  # FSx subnet: by default the first instance group's subnet (primary_subnet_id).
+  # If fsx_availability_zone_id is set (e.g. the instance group is in a Local Zone
+  # where FSx is not offered), place FSx in the private subnet matching that AZ ID.
+  fsx_subnet_id = local.create_fsx_module ? (
+    var.fsx_availability_zone_id != "" ? local.az_to_subnet_map[var.fsx_availability_zone_id] : module.hyperpod_cluster[0].primary_subnet_id
+  ) : null
 }
 
 module "vpc" {
   count  = var.create_vpc_module ? 1 : 0
   source = "./modules/vpc"
 
-  resource_name_prefix = var.resource_name_prefix
-  vpc_cidr             = var.vpc_cidr
-  public_subnet_1_cidr = var.public_subnet_1_cidr
-  public_subnet_2_cidr = var.public_subnet_2_cidr
-  closed_network       = var.closed_network
+  resource_name_prefix             = var.resource_name_prefix
+  vpc_cidr                         = var.vpc_cidr
+  public_subnet_1_cidr             = var.public_subnet_1_cidr
+  public_subnet_2_cidr             = var.public_subnet_2_cidr
+  closed_network                   = var.closed_network
+  local_zone_egress_zone_ids       = var.local_zone_egress_zone_ids
+  local_zone_public_subnet_cidrs   = var.local_zone_public_subnet_cidrs
+  local_zone_network_border_groups = var.local_zone_network_border_groups
 }
 
 module "private_subnet" {
   count  = var.create_private_subnet_module ? 1 : 0
   source = "./modules/private_subnet"
 
-  resource_name_prefix = var.resource_name_prefix
-  vpc_id               = local.vpc_id
-  private_subnet_cidrs = var.private_subnet_cidrs
-  nat_gateway_id       = local.nat_gateway_id
-  closed_network       = var.closed_network
+  resource_name_prefix       = var.resource_name_prefix
+  vpc_id                     = local.vpc_id
+  private_subnet_cidrs       = var.private_subnet_cidrs
+  availability_zone_ids      = var.private_subnet_availability_zone_ids
+  nat_gateway_id             = local.nat_gateway_id
+  nat_gateway_ids_by_zone_id = local.nat_gateway_ids_by_zone_id
+  closed_network             = var.closed_network
 }
 
 module "security_group" {
@@ -282,7 +297,7 @@ module "fsx_lustre" {
 
   resource_name_prefix       = var.resource_name_prefix
   eks_cluster_name           = local.eks_cluster_name
-  subnet_id                  = module.hyperpod_cluster[0].primary_subnet_id
+  subnet_id                  = local.fsx_subnet_id
   security_group_id          = local.security_group_id
   create_new_filesystem      = var.create_new_fsx_filesystem
   storage_capacity           = var.fsx_storage_capacity

1.architectures/7.sagemaker-hyperpod-eks/terraform-modules/hyperpod-eks-tf/modules/private_subnet/main.tf

@@ -6,9 +6,18 @@ data "aws_availability_zones" "available" {
   }
 }
 
-# Set subnet count to the lesser of either the number of CIDRs provided or number of AZs available in the region
+# Subnet placement:
+#  - By default, subnets are spread across discovered standard AZs (opt-in-not-required),
+#    with count = min(#CIDRs, #AZs).
+#  - When availability_zone_ids is set, subnets are placed in those exact AZ IDs
+#    (1:1 with private_subnet_cidrs). This allows opt-in zones such as Local Zones,
+#    which the discovery filter deliberately excludes.
 locals {
-  subnet_count = min(length(var.private_subnet_cidrs), length(data.aws_availability_zones.available.names))
+  use_explicit_azs = length(var.availability_zone_ids) > 0
+  subnet_count = local.use_explicit_azs ? length(var.private_subnet_cidrs) : min(
+    length(var.private_subnet_cidrs), length(data.aws_availability_zones.available.names)
+  )
+  subnet_zone_ids = local.use_explicit_azs ? var.availability_zone_ids : data.aws_availability_zones.available.zone_ids
 }
 
 resource "aws_vpc_ipv4_cidr_block_association" "additional_cidr" {
@@ -21,7 +30,7 @@ resource "aws_subnet" "private" {
   count                = local.subnet_count
   vpc_id               = var.vpc_id
   cidr_block           = var.private_subnet_cidrs[count.index]
-  availability_zone_id = data.aws_availability_zones.available.zone_ids[count.index]
+  availability_zone_id = local.subnet_zone_ids[count.index]
 
   # Ensure the subnet is created after the CIDR block is associated
   depends_on = [aws_vpc_ipv4_cidr_block_association.additional_cidr]
@@ -52,10 +61,17 @@ resource "aws_route_table_association" "private" {
   route_table_id = aws_route_table.private[count.index].id
 }
 
-# NAT Gateway route - only created if NOT closed network
+# NAT Gateway route - only created if NOT closed network.
+# Per-AZ NAT selection: if this subnet's AZ ID has an entry in
+# nat_gateway_ids_by_zone_id (used for Local Zone-local NATs), use it;
+# otherwise fall back to the regional var.nat_gateway_id.
 resource "aws_route" "nat_gateway" {
   count                  = var.closed_network ? 0 : local.subnet_count
   route_table_id         = aws_route_table.private[count.index].id
   destination_cidr_block = "0.0.0.0/0"
-  nat_gateway_id         = var.nat_gateway_id
+  nat_gateway_id = lookup(
+    var.nat_gateway_ids_by_zone_id,
+    local.subnet_zone_ids[count.index],
+    var.nat_gateway_id
+  )
 }

1.architectures/7.sagemaker-hyperpod-eks/terraform-modules/hyperpod-eks-tf/modules/private_subnet/variables.tf

@@ -14,11 +14,40 @@ variable "private_subnet_cidrs" {
   default     = ["10.1.0.0/16", "10.2.0.0/16", "10.3.0.0/16", "10.4.0.0/16"]
 }
 
+variable "availability_zone_ids" {
+  description = <<-EOT
+    Optional list of Availability Zone IDs to place subnets in, 1:1 with
+    private_subnet_cidrs. When empty (default), AZs are discovered automatically
+    (standard opt-in-not-required zones only). Set this to target specific zones,
+    including opt-in zones such as Local Zones (e.g. usw2-phx2-az1) that the
+    discovery filter excludes.
+  EOT
+  type        = list(string)
+  default     = []
+
+  validation {
+    condition     = length(var.availability_zone_ids) == 0 || length(var.availability_zone_ids) == length(var.private_subnet_cidrs)
+    error_message = "When set, availability_zone_ids must have the same length as private_subnet_cidrs."
+  }
+}
+
 variable "nat_gateway_id" {
   description = "The Id of a NAT Gateway to route internet bound traffic (ignored if closed_network is true)"
   type        = string
 }
 
+variable "nat_gateway_ids_by_zone_id" {
+  description = <<-EOT
+    Optional per-AZ-ID NAT gateway override. When a private subnet's AZ ID is
+    a key in this map, its default route uses the mapped NAT gateway instead
+    of var.nat_gateway_id. Used to route Local Zone subnets to an LZ-local
+    NAT gateway. Falls back to var.nat_gateway_id for unmapped AZs.
+    Default: empty map (all subnets use var.nat_gateway_id).
+  EOT
+  type        = map(string)
+  default     = {}
+}
+
 variable "closed_network" {
   description = "Whether to deploy in closed network mode (no NAT gateway routes)"
   type        = bool