Allow same compute resource name in different queues

Signed-off-by: chenwany <[email protected]>

Author: chenwany

src/slurm_plugin/instance_manager.py

@@ -216,7 +216,7 @@ def _parse_requested_instances(self, node_list):
         for node in node_list:
             try:
                 queue_name, node_type, instance_name = parse_nodename(node)
-                instance_type = self._instance_name_type_mapping[instance_name]
+                instance_type = self._instance_name_type_mapping[queue_name][instance_name]
                 instances_to_launch[queue_name][instance_type].append(node)
             except (InvalidNodenameError, KeyError):
                 logger.warning("Discarding NodeName with invalid format: %s", node)

tests/slurm_plugin/slurm_resources/test_instance_manager.py

@@ -1050,11 +1050,11 @@ def test_parse_requested_instances(
     ):
         # Mock instance name/type mapping
         instance_name_type_mapping = {
-            "c5xlarge": "c5.xlarge",
-            "g34xlarge": "g3.4xlarge",
-            "g38xlarge": "g3.8xlarge",
-            "u6tb1metal": "u-6tb1.metal",
-            "i3enmetal2tb": "i3en.metal-2tb",
+            "queue1": {"c5xlarge": "c5.xlarge"},
+            "queue2": {"g34xlarge": "g3.4xlarge", "g38xlarge": "g3.8xlarge", "u6tb1metal": "u-6tb1.metal"},
+            "queuename-with-dash-and_underscore": {
+                "i3enmetal2tb": "i3en.metal-2tb",
+            },
         }
         instance_manager._instance_name_type_mapping = instance_name_type_mapping
 

tests/slurm_plugin/test_clustermgtd.py

@@ -1538,7 +1538,9 @@ def test_manage_cluster_boto3(
     # patch boto3 call
     boto3_stubber("ec2", mocked_boto3_request)
     mocker.patch("slurm_plugin.clustermgtd.datetime").now.return_value = datetime(2020, 1, 2, 0, 0, 0)
-    mocker.patch("slurm_plugin.clustermgtd.retrieve_instance_type_mapping", return_value={"c5xlarge": "c5.xlarge"})
+    mocker.patch(
+        "slurm_plugin.clustermgtd.retrieve_instance_type_mapping", return_value={"queue": {"c5xlarge": "c5.xlarge"}}
+    )
     sync_config = ClustermgtdConfig(test_datadir / config_file)
     cluster_manager = ClusterManager(sync_config)
     dynamodb_table_mock = mocker.patch.object(cluster_manager._compute_fleet_status_manager, "_table")

tests/slurm_plugin/test_resume.py

@@ -256,7 +256,7 @@ def test_resume_launch(
         cluster_name="hit",
         head_node_private_ip="some_ip",
         head_node_hostname="some_hostname",
-        instance_name_type_mapping={"c5xlarge": "c5.xlarge", "c52xlarge": "c5.2xlarge"},
+        instance_name_type_mapping={"queue1": {"c5xlarge": "c5.xlarge"}, "queue2": {"c52xlarge": "c5.2xlarge"}},
         clustermgtd_heartbeat_file_path="some_path",
         clustermgtd_timeout=600,
         boto3_config=botocore.config.Config(),