[SlurmTopo] Code linters

Author: Himani Anil Deshpande

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

@@ -10,10 +10,12 @@
 # This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied.
 # See the License for the specific language governing permissions and limitations under the License.
 import argparse
-import yaml
 import logging
-import traceback
 import os
+import traceback
+
+import yaml
+
 log = logging.getLogger()
 
 
@@ -43,12 +45,12 @@ def _load_cluster_config(input_file_path):
         return yaml.load(input_file, Loader=yaml.SafeLoader)
 
 
-def generate_topology_config_file(output_file: str, input_file: str, block_sizes: str):
+def generate_topology_config_file(output_file: str, input_file: str, block_sizes: str):  # noqa: C901
     """
     Generate Topology configuration file.
 
     Generate topology.conf
-    
+
     # This file is automatically generated by pcluster
     BlockName=block1 Nodes=queue-1-st-compute-resource-0-[1-9]  #### 9 nodes
     BlockName=block2 Nodes=queue-1-st-compute-resource-0-[1-18] #### 18 nodes
@@ -81,15 +83,28 @@ def generate_topology_config_file(output_file: str, input_file: str, block_sizes
                     else:
                         continue
 
-                    ### Check for if reservation is for NVLink and size matches min_block_size_list
-                    if compute_resource_config.get('InstanceType') == 'p6e-gb200.36xlarge':
+                    # Check for if reservation is for NVLink and size matches min_block_size_list
+                    if compute_resource_config.get("InstanceType") == "p6e-gb200.36xlarge":
                         if min_block_size_list == compute_min_count or max_block_size_list == compute_max_count:
                             block_count += 1
-                            ### Each Capacity Reservation ID is a Capacity Block and we associate each slurm block with a single capacity Block
-                            topology_config += "BlockName=Block" + str(block_count)+ "  Nodes=" + str(queue_name) + "-" + str(node_type) + "-" + str(compute_resource_name) + "-[1-" + str(compute_max_count) + "]\n"
-
-            topology_config += "BlockSizes="+ str(block_sizes)+"\n"
-        except(KeyError, AttributeError) as e:
+                            # Each Capacity Reservation ID is a Capacity Block,
+                            # we associate each slurm block with a single capacity Block
+                            topology_config += (
+                                "BlockName=Block"
+                                + str(block_count)
+                                + "  Nodes="
+                                + str(queue_name)
+                                + "-"
+                                + str(node_type)
+                                + "-"
+                                + str(compute_resource_name)
+                                + "-[1-"
+                                + str(compute_max_count)
+                                + "]\n"
+                            )
+
+            topology_config += "BlockSizes=" + str(block_sizes) + "\n"
+        except (KeyError, AttributeError) as e:
             if isinstance(e, KeyError):
                 message = f"Unable to find key {e} in the configuration file."
             else:

test/unit/slurm/test_topology_generator.py

@@ -9,8 +9,9 @@
 # OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
 import os
+
+import pytest
 from assertpy import assert_that
 from pcluster_topology_generator import (
     cleanup_topology_config_file,
@@ -23,27 +24,21 @@ def _assert_files_are_equal(file, expected_file):
         assert_that(f.read()).is_equal_to(exp_f.read())
 
 
-@pytest.mark.parametrize("file_name_suffix", [
-    "with_capacity_block",
-    "no_capacity_block"
-])
+@pytest.mark.parametrize("file_name_suffix", ["with_capacity_block", "no_capacity_block"])
 def test_generate_topology_config(test_datadir, tmpdir, file_name_suffix):
-    block_sizes = "9,18" if 'no' not in file_name_suffix else None
+    block_sizes = "9,18" if "no" not in file_name_suffix else None
     file_name = "sample_" + file_name_suffix + ".yaml"
     input_file_path = str(test_datadir / file_name)
     output_file_name = "topology_" + file_name_suffix + ".conf"
     output_file_path = f"{tmpdir}/{output_file_name}"
     generate_topology_config_file(output_file_path, input_file_path, block_sizes)
-    if 'no' in file_name_suffix:
+    if "no" in file_name_suffix:
         assert_that(os.path.isfile(output_file_path)).is_equal_to(False)
     else:
         _assert_files_are_equal(output_file_path, test_datadir / "expected_outputs" / output_file_name)
 
 
-@pytest.mark.parametrize("file_exists", [
-    True,
-    False
-])
+@pytest.mark.parametrize("file_exists", [True, False])
 def test_cleanup_topology_config_file(mocker, tmpdir, file_exists):
     topology_file_path = tmpdir / "topology.conf"
     mocker.patch("os.path.exists", return_value=file_exists)
@@ -53,4 +48,3 @@ def test_cleanup_topology_config_file(mocker, tmpdir, file_exists):
         mock_remove.assert_called_once_with(str(topology_file_path))
     else:
         mock_remove.assert_not_called()
-