testing Trn1 with Capacity blocks

Author: Himani Anil Deshpande

tests/integration-tests/tests/trainium/test_trainium.py

@@ -42,7 +42,7 @@ def test_trainium(
 
 
 def _test_allreduce_single_node(test_datadir, remote_command_executor, scheduler_commands):
-    result = scheduler_commands.submit_script(str(test_datadir / "neuron-allreduce.sh"), partition="queue-trn2")
+    result = scheduler_commands.submit_script(str(test_datadir / "neuron-allreduce.sh"), partition="queue-trn32")
     job_id = scheduler_commands.assert_job_submitted(result.stdout)
     scheduler_commands.wait_job_completed(job_id)
     scheduler_commands.assert_job_succeeded(job_id)

tests/integration-tests/tests/trainium/test_trainium/test_trainium/pcluster.config.yaml

@@ -18,13 +18,16 @@ Scheduling:
   Scheduler: slurm
   SlurmQueues:
     - Name: queue-trn32
+      CapacityType: CAPACITY_BLOCK
       ComputeResources:
         - Name: compute-resource-trn32
-          Instances:
-            - InstanceType: {{instance}}
+          InstanceType: {{instance}}
           MinCount: 2
+          MaxCount: 2
           Efa:
             Enabled: true
+          CapacityReservationTarget:
+            CapacityReservationId: cr-05b0c099ce2534ce3
       Networking:
         SubnetIds:
           - {{ private_subnet_id }}
@@ -42,24 +45,24 @@ Scheduling:
           - BucketName: {{ bucket_name }}
           # Needed to download neuronx packages and neff file --> FIXME to be removed once packages are public available
           - BucketName: aws-parallelcluster-beta
-    - Name: queue-trn2
-      ComputeResources:
-        - Name: compute-resource-trn2
-          Instances:
-            - InstanceType: trn1.2xlarge
-          MinCount: 0  # TODO change to 1 once allreduce test is passing
-      Networking:
-        SubnetIds:
-          - {{ private_subnet_id }}
-      CustomActions:
-        OnNodeConfigured:
-          Script: s3://{{ bucket_name }}/neuron-installation.sh
-      Iam:
-        # Policy to access to Trainium beta repository info
-        AdditionalIamPolicies:
-          - Policy: arn:aws:iam::447714826191:policy/TrainiumPreviewPolicy
-        S3Access:
-          # Needed to download post install script
-          - BucketName: {{ bucket_name }}
-          # Needed to download neuronx packages and neff file --> FIXME to be removed once packages are public available
-          - BucketName: aws-parallelcluster-beta
+#    - Name: queue-trn2
+#      ComputeResources:
+#        - Name: compute-resource-trn2
+#          Instances:
+#            - InstanceType: trn1.2xlarge
+#          MinCount: 0  # TODO change to 1 once allreduce test is passing
+#      Networking:
+#        SubnetIds:
+#          - {{ private_subnet_id }}
+#      CustomActions:
+#        OnNodeConfigured:
+#          Script: s3://{{ bucket_name }}/neuron-installation.sh
+#      Iam:
+#        # Policy to access to Trainium beta repository info
+#        AdditionalIamPolicies:
+#          - Policy: arn:aws:iam::447714826191:policy/TrainiumPreviewPolicy
+#        S3Access:
+#          # Needed to download post install script
+#          - BucketName: {{ bucket_name }}
+#          # Needed to download neuronx packages and neff file --> FIXME to be removed once packages are public available
+#          - BucketName: aws-parallelcluster-beta