[Subnet Prioritization] Add EnableSingleAvailabilityZone to Fleet Config

CHANGELOG.md

@@ -11,6 +11,7 @@ This file is used to list changes made in each version of the AWS ParallelCluste
 
 **CHANGES**
 - Ubuntu 20.04 is no longer supported.
+- Add SingleAvailabilityZone parameter to fleet_config.json
 
 3.13.1
 ------

cookbooks/aws-parallelcluster-slurm/files/default/head_node_slurm/slurm/pcluster_fleet_config_generator.py

@@ -99,6 +99,9 @@ def generate_fleet_config_file(output_file: str, input_file: str):
                     queue_capacity_reservation=queue_capacity_reservation,
                     queue_capacity_type=queue_capacity_type,
                     queue_subnets=queue_config["Networking"]["SubnetIds"],
+                    queue_single_availability_zone=queue_config.get("Networking", {}).get(
+                        "EnableSingleAvailabilityZone", None
+                    ),
                 )
                 fleet_config[queue_name][compute_resource_name] = config_for_fleet
 
@@ -125,6 +128,7 @@ def _generate_compute_resource_fleet_config(
     queue_capacity_reservation: str,
     queue_capacity_type: str,
     queue_subnets: List,
+    queue_single_availability_zone: bool,
 ):
     """
     Generate compute resource config to add in the fleet-config.json, overriding values from the queue.
@@ -156,7 +160,10 @@ def _generate_compute_resource_fleet_config(
                 {
                     "Api": "create-fleet",
                     "Instances": copy.deepcopy(compute_resource_config["Instances"]),
-                    "Networking": {"SubnetIds": queue_subnets},
+                    "Networking": {
+                        "SubnetIds": queue_subnets,
+                        "SingleAvailabilityZone": queue_single_availability_zone,
+                    },
                 }
             )
             allocation_strategy = compute_resource_config.get("AllocationStrategy", queue_allocation_strategy)

test/unit/slurm/test_fleet_config_generator.py

@@ -62,6 +62,22 @@
             CriticalError,
             "Unable to find key 'Name' in the configuration file. Queue: q1",
         ),
+        (
+            {
+                "Scheduling": {
+                    "SlurmQueues": [
+                        {
+                            "Name": "q1",
+                            "CapacityType": "SPOT",
+                            "ComputeResources": [{"Instances": []}],
+                            "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},
+                        }
+                    ]
+                }
+            },
+            CriticalError,
+            "Unable to find key 'Name' in the configuration file. Queue: q1",
+        ),
         (
             {
                 "Scheduling": {
@@ -70,7 +86,7 @@
                             "Name": "q1",
                             "CapacityType": "ONDEMAND",
                             "ComputeResources": [{"Name": "cr1", "Instances": []}],
-                            "Networking": {"SubnetIds": ["123"]},
+                            "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},
                         }
                     ]
                 }
@@ -89,7 +105,7 @@
                                 {"Name": "cr1", "Instances": [{"InstanceType": "test"}]},
                                 {"Name": "cr2", "InstanceType": "test"},
                             ],
-                            "Networking": {"SubnetIds": ["123"]},
+                            "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},
                         }
                     ]
                 }
@@ -108,7 +124,7 @@
                                 {"Name": "cr1", "Instances": [{"InstanceType": "test"}, {"InstanceType": "test-2"}]},
                                 {"Name": "cr2", "InstanceType": "test"},
                             ],
-                            "Networking": {"SubnetIds": ["123", "456", "789"]},
+                            "Networking": {"SubnetIds": ["123", "456", "789"], "EnableSingleAvailabilityZone": None},
                         }
                     ]
                 }
@@ -131,7 +147,7 @@
                                 },
                                 {"Name": "cr2", "InstanceType": "test", "SpotPrice": "10"},
                             ],
-                            "Networking": {"SubnetIds": ["123", "456", "789"]},
+                            "Networking": {"SubnetIds": ["123", "456", "789"], "EnableSingleAvailabilityZone": None},
                         }
                     ]
                 }
@@ -147,7 +163,7 @@
                             "Name": "q1",
                             "CapacityType": "SPOT",
                             "ComputeResources": [{"Name": "cr1", "Instances": [{"InstanceType": "test"}]}],
-                            "Networking": {"SubnetIds": ["123"]},
+                            "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},
                         }
                     ]
                 }
@@ -165,7 +181,7 @@
                             "ComputeResources": [
                                 {"Name": "cr1", "Instances": [{"InstanceType": "test"}], "SpotPrice": 10}
                             ],
-                            "Networking": {"SubnetIds": ["123"]},
+                            "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},
                         }
                     ]
                 }
@@ -208,6 +224,24 @@
             CriticalError,
             "Unable to find key 'SubnetIds' in the configuration file. Queue: q1",
         ),
+        (
+            {
+                "Scheduling": {
+                    "SlurmQueues": [
+                        {
+                            "Name": "q1",
+                            "CapacityType": "SPOT",
+                            "ComputeResources": [
+                                {"Name": "cr1", "Instances": [{"InstanceType": "test"}], "SpotPrice": 10}
+                            ],
+                            "Networking": {"SubnetIds": ["123"]},
+                        }
+                    ]
+                }
+            },
+            None,
+            None,
+        ),
         (
             {
                 "Scheduling": {
@@ -231,7 +265,7 @@
                                     },
                                 },
                             ],
-                            "Networking": {"SubnetIds": ["123"]},
+                            "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},
                         }
                     ]
                 }

test/unit/slurm/test_fleet_config_generator/test_generate_fleet_config_file/expected_outputs/fleet-config.json

@@ -27,7 +27,8 @@
             "Networking": {
                 "SubnetIds": [
                     "subnet-0230367ab0e5123a4"
-                ]
+                ],
+                "SingleAvailabilityZone": true
             },
             "AllocationStrategy": "lowest-price"
         }
@@ -57,7 +58,8 @@
                 "SubnetIds": [
                     "subnet-0230367ab0e5123a4",
                     "subnet-0b903123096649662"
-                ]
+                ],
+                "SingleAvailabilityZone": false
             },
             "AllocationStrategy": "lowest-price"
         }
@@ -84,7 +86,8 @@
             "Networking": {
                 "SubnetIds": [
                     "subnet-0230367ab0e5123a4"
-                ]
+                ],
+                "SingleAvailabilityZone": null
             },
             "AllocationStrategy": "capacity-optimized",
             "MaxPrice": 10
@@ -100,7 +103,8 @@
             "Networking": {
                 "SubnetIds": [
                     "subnet-0230367ab0e5123a4"
-                ]
+                ],
+                "SingleAvailabilityZone": null
             },
             "AllocationStrategy": "capacity-optimized"
         }
@@ -131,7 +135,8 @@
             "Networking": {
                 "SubnetIds": [
                     "subnet-0230367ab0e5123a4"
-                ]
+                ],
+                "SingleAvailabilityZone": null
             }
         }
     }

test/unit/slurm/test_fleet_config_generator/test_generate_fleet_config_file/sample_input.yaml

@@ -62,6 +62,7 @@ Scheduling:
       SecurityGroups: null
       SubnetIds:
       - subnet-0230367ab0e5123a4
+      EnableSingleAvailabilityZone: true
     # queue ondemand without capacity reservations and with multiple subnets
   - AllocationStrategy: lowest-price
     CapacityReservationTarget: null
@@ -121,6 +122,7 @@ Scheduling:
       SubnetIds:
       - subnet-0230367ab0e5123a4
       - subnet-0b903123096649662
+      EnableSingleAvailabilityZone: false
     # queue spot
   - AllocationStrategy: capacity-optimized
     CapacityReservationTarget: null
@@ -191,6 +193,7 @@ Scheduling:
       SecurityGroups: null
       SubnetIds:
       - subnet-0230367ab0e5123a4
+      EnableSingleAvailabilityZone: null
     # queue for capacity-block
   - CapacityReservationTarget:
       CapacityReservationId: cr-987654