[Subnet Prioritization] Update unit test of generate_fleet_config_file

Allenz5 · Allenz5 · commit 116282d15c8b · 2025-06-16T12:53:16.000-04:00
Signed-off-by: Hanxuan Zhang &lt;hanxuanz@amazon.com&gt;
diff --git a/test/unit/slurm/test_fleet_config_generator.py b/test/unit/slurm/test_fleet_config_generator.py
@@ -46,6 +46,22 @@
             CriticalError,
             "Unable to find key 'Networking' in the configuration file. Queue: q1",
         ),
+        (
+                {
+                    "Scheduling": {
+                        "SlurmQueues": [
+                            {
+                                "Name": "q1",
+                                "CapacityType": "SPOT",
+                                "ComputeResources": [{"Instances": []}],
+                                "Networking": {"SubnetIds": ["123"]},
+                            }
+                        ]
+                    }
+                },
+                CriticalError,
+                "Unable to find key 'EnableSingleAvailabilityZone' in the configuration file. Queue: q1",
+        ),
         (
             {
                 "Scheduling": {
@@ -54,7 +70,7 @@
                             "Name": "q1",
                             "CapacityType": "SPOT",
                             "ComputeResources": [{"Instances": []}],
-                            "Networking": {"SubnetIds": ["123"]},
+                            "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},
                         }
                     ]
                 }
@@ -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"]},
+                            }
+                        ]
+                    }
+                },
+                CriticalError,
+                "Unable to find key 'EnableSingleAvailabilityZone' in the configuration file. Queue: q1",
+        ),
         (
             {
                 "Scheduling": {
@@ -231,7 +265,7 @@
                                     },
                                 },
                             ],
-                            "Networking": {"SubnetIds": ["123"]},
+                            "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},
                         }
                     ]
                 }
diff --git a/test/unit/slurm/test_fleet_config_generator/test_generate_fleet_config_file/expected_outputs/fleet-config.json b/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
             }
         }
     }
diff --git a/test/unit/slurm/test_fleet_config_generator/test_generate_fleet_config_file/sample_input.yaml b/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
@@ -250,6 +253,7 @@ Scheduling:
       SecurityGroups: null
       SubnetIds:
         - subnet-0230367ab0e5123a4
+      EnableSingleAvailabilityZone: null
   SlurmSettings:
     Dns:
       DisableManagedDns: false

Original file line number	Diff line number	Diff line change
`@@ -46,6 +46,22 @@`
`46`	`46`	`CriticalError,`
`47`	`47`	`"Unable to find key 'Networking' in the configuration file. Queue: q1",`
`48`	`48`	`),`
	`49`	`+ (`
	`50`	`+ {`
	`51`	`+ "Scheduling": {`
	`52`	`+ "SlurmQueues": [`
	`53`	`+ {`
	`54`	`+ "Name": "q1",`
	`55`	`+ "CapacityType": "SPOT",`
	`56`	`+ "ComputeResources": [{"Instances": []}],`
	`57`	`+ "Networking": {"SubnetIds": ["123"]},`
	`58`	`+ }`
	`59`	`+ ]`
	`60`	`+ }`
	`61`	`+ },`
	`62`	`+ CriticalError,`
	`63`	`+ "Unable to find key 'EnableSingleAvailabilityZone' in the configuration file. Queue: q1",`
	`64`	`+ ),`
`49`	`65`	`(`
`50`	`66`	`{`
`51`	`67`	`"Scheduling": {`
`@@ -54,7 +70,7 @@`
`54`	`70`	`"Name": "q1",`
`55`	`71`	`"CapacityType": "SPOT",`
`56`	`72`	`"ComputeResources": [{"Instances": []}],`
`57`		`- "Networking": {"SubnetIds": ["123"]},`
	`73`	`+ "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},`
`58`	`74`	`}`
`59`	`75`	`]`
`60`	`76`	`}`
`@@ -70,7 +86,7 @@`
`70`	`86`	`"Name": "q1",`
`71`	`87`	`"CapacityType": "ONDEMAND",`
`72`	`88`	`"ComputeResources": [{"Name": "cr1", "Instances": []}],`
`73`		`- "Networking": {"SubnetIds": ["123"]},`
	`89`	`+ "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},`
`74`	`90`	`}`
`75`	`91`	`]`
`76`	`92`	`}`
`@@ -89,7 +105,7 @@`
`89`	`105`	`{"Name": "cr1", "Instances": [{"InstanceType": "test"}]},`
`90`	`106`	`{"Name": "cr2", "InstanceType": "test"},`
`91`	`107`	`],`
`92`		`- "Networking": {"SubnetIds": ["123"]},`
	`108`	`+ "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},`
`93`	`109`	`}`
`94`	`110`	`]`
`95`	`111`	`}`
`@@ -108,7 +124,7 @@`
`108`	`124`	`{"Name": "cr1", "Instances": [{"InstanceType": "test"}, {"InstanceType": "test-2"}]},`
`109`	`125`	`{"Name": "cr2", "InstanceType": "test"},`
`110`	`126`	`],`
`111`		`- "Networking": {"SubnetIds": ["123", "456", "789"]},`
	`127`	`+ "Networking": {"SubnetIds": ["123", "456", "789"], "EnableSingleAvailabilityZone": None},`
`112`	`128`	`}`
`113`	`129`	`]`
`114`	`130`	`}`
`@@ -131,7 +147,7 @@`
`131`	`147`	`},`
`132`	`148`	`{"Name": "cr2", "InstanceType": "test", "SpotPrice": "10"},`
`133`	`149`	`],`
`134`		`- "Networking": {"SubnetIds": ["123", "456", "789"]},`
	`150`	`+ "Networking": {"SubnetIds": ["123", "456", "789"], "EnableSingleAvailabilityZone": None},`
`135`	`151`	`}`
`136`	`152`	`]`
`137`	`153`	`}`
`@@ -147,7 +163,7 @@`
`147`	`163`	`"Name": "q1",`
`148`	`164`	`"CapacityType": "SPOT",`
`149`	`165`	`"ComputeResources": [{"Name": "cr1", "Instances": [{"InstanceType": "test"}]}],`
`150`		`- "Networking": {"SubnetIds": ["123"]},`
	`166`	`+ "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},`
`151`	`167`	`}`
`152`	`168`	`]`
`153`	`169`	`}`
`@@ -165,7 +181,7 @@`
`165`	`181`	`"ComputeResources": [`
`166`	`182`	`{"Name": "cr1", "Instances": [{"InstanceType": "test"}], "SpotPrice": 10}`
`167`	`183`	`],`
`168`		`- "Networking": {"SubnetIds": ["123"]},`
	`184`	`+ "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},`
`169`	`185`	`}`
`170`	`186`	`]`
`171`	`187`	`}`
`@@ -208,6 +224,24 @@`
`208`	`224`	`CriticalError,`
`209`	`225`	`"Unable to find key 'SubnetIds' in the configuration file. Queue: q1",`
`210`	`226`	`),`
	`227`	`+ (`
	`228`	`+ {`
	`229`	`+ "Scheduling": {`
	`230`	`+ "SlurmQueues": [`
	`231`	`+ {`
	`232`	`+ "Name": "q1",`
	`233`	`+ "CapacityType": "SPOT",`
	`234`	`+ "ComputeResources": [`
	`235`	`+ {"Name": "cr1", "Instances": [{"InstanceType": "test"}], "SpotPrice": 10}`
	`236`	`+ ],`
	`237`	`+ "Networking": {"SubnetIds": ["123"]},`
	`238`	`+ }`
	`239`	`+ ]`
	`240`	`+ }`
	`241`	`+ },`
	`242`	`+ CriticalError,`
	`243`	`+ "Unable to find key 'EnableSingleAvailabilityZone' in the configuration file. Queue: q1",`
	`244`	`+ ),`
`211`	`245`	`(`
`212`	`246`	`{`
`213`	`247`	`"Scheduling": {`
`@@ -231,7 +265,7 @@`
`231`	`265`	`},`
`232`	`266`	`},`
`233`	`267`	`],`
`234`		`- "Networking": {"SubnetIds": ["123"]},`
	`268`	`+ "Networking": {"SubnetIds": ["123"], "EnableSingleAvailabilityZone": None},`
`235`	`269`	`}`
`236`	`270`	`]`
`237`	`271`	`}`