.ci/aws: Add trainium tests to CI

Signed-off-by: Seth Zegelstein <[email protected]>
a-szegel · Aug 12, 2024 · 7c44533 · 7c44533
1 parent ece41f0
commit 7c44533
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/.ci/aws/Jenkinsfile b/.ci/aws/Jenkinsfile
@@ -222,6 +222,14 @@ pipeline {
                     def p5_region = "af-south-1"
                     def p5_odcr = "cr-02eb632dcd8175139"
                     def p4_p5_addl_args = "${base_args} --use-prebuilt-ami-with-efa-installer true"
+                    def trn1_lock_label = "trn1-1-4node"
+                    def trn1_region = "us-east-2"
+                    def trn1_odcr = "cr-0e9366fb7fa2772f1"
+                    def trn1_addl_args = "${base_args} --odcr-placement-group-name trn1-placement-group --test-list test_nccom_test"
+                    def trn1n_lock_label = "trn1n-1-4node"
+                    def trn1n_region = "us-east-1"
+                    def trn1n_odcr = "cr-07342cf6439332dce"
+                    def trn1n_addl_args = "${base_args} --odcr-placement-group-name trn1n-placement-group --test-list test_nccom_test"
 
                     // p3dn tests
                     stages["4_p3dn_al2"] = get_test_stage_with_lock("4_p3dn_al2", env.BUILD_TAG, "alinux2", "p3dn.24xlarge", p3dn_region, p3dn_lock_label, num_instances, p3dn_odcr, p3dn_al2_addl_args)
@@ -238,6 +246,12 @@ pipeline {
                     stages["4_p5_ubuntu2004"] = get_test_stage_with_lock("4_p5_ubuntu2004", env.BUILD_TAG, "ubuntu2004", "p5.48xlarge", p5_region, p5_lock_label, num_instances, p5_odcr, p4_p5_addl_args)
                     stages["4_p5_ubuntu2204"] = get_test_stage_with_lock("4_p5_ubuntu2204", env.BUILD_TAG, "ubuntu2204", "p5.48xlarge", p5_region, p5_lock_label, num_instances, p5_odcr, p4_p5_addl_args)
 
+                    // trn1 tests
+                    stages["4_trn1_ubuntu2004"] = get_test_stage_with_lock("4_trn1_ubuntu2004", env.BUILD_TAG, "ubuntu2004", "trn1.32xlarge", trn1_region, trn1_lock_label, num_instances, trn1_odcr, trn1_addl_args)
+
+                    // trn1n tests
+                    stages["4_trn1n_ubuntu2004"] = get_test_stage_with_lock("4_trn1n_ubuntu2004", env.BUILD_TAG, "ubuntu2004", "trn1n.32xlarge", trn1n_region, trn1n_lock_label, num_instances, trn1n_odcr, trn1n_addl_args)
+
                     parallel stages
                 }
             }