feat(nvidia): build pytorch to get older cuda compute capabilities and setup arm64 support

Author: ndbaker1

.github/workflows/ci.yaml

@@ -35,12 +35,16 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v3
-    - run: docker build --file test/images/nvidia-training/Dockerfile test/images/nvidia-training
+    - run: |
+        docker build --file test/images/nvidia-training/Dockerfile test/images/nvidia-training \
+          --build-arg PYTORCH_BUILD_ENV="MAX_JOBS=8 BUILD_TEST=0 USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 USE_DISTRIBUTED=0"
   build-image-nvidia-inference:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v3
-    - run: docker build --file test/images/nvidia-inference/Dockerfile test/images/nvidia-inference
+    - run: |
+        docker build --file test/images/nvidia-inference/Dockerfile test/images/nvidia-inference \
+          --build-arg PYTORCH_BUILD_ENV="MAX_JOBS=8 BUILD_TEST=0 USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 USE_DISTRIBUTED=0"
   build-image-neuron-training:
     runs-on: ubuntu-latest
     steps:

internal/deployers/eksapi/kubeconfig.go

@@ -2,9 +2,9 @@ package eksapi
 
 import (
 	"bytes"
+	"fmt"
 	"os"
 	"text/template"
-	"fmt"
 
 	"k8s.io/klog"
 )

test/cases/nvidia-training/bert_training_test.go

@@ -69,9 +69,11 @@ func TestBertTraining(t *testing.T) {
 				ObjectMeta: metav1.ObjectMeta{Name: "bert-training-launcher", Namespace: "default"},
 			}
 			err := wait.For(fwext.NewConditionExtension(cfg.Client().Resources()).JobSucceeded(job),
-				wait.WithTimeout(time.Minute*20))
+				wait.WithTimeout(time.Minute*20),
+				wait.WithContext(ctx),
+			)
 			if err != nil {
-				t.Fatal(err)
+				t.Error(err)
 			}
 
 			err = printJobLogs(ctx, cfg, "default", "bert-training-launcher")

test/cases/nvidia-training/main_test.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"os/signal"
 	"slices"
 	"testing"
 	"time"
@@ -37,7 +38,10 @@ func TestMain(m *testing.M) {
 	if err != nil {
 		log.Fatalf("failed to initialize test environment: %v", err)
 	}
-	testenv = env.NewWithConfig(cfg)
+
+	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
+	defer cancel()
+	testenv = env.NewWithConfig(cfg).WithContext(ctx)
 
 	manifests := [][]byte{
 		nvidiaDevicePluginManifest,
@@ -147,16 +151,11 @@ func checkNodeTypes(ctx context.Context, config *envconf.Config) (context.Contex
 		return ctx, fmt.Errorf("no nodes found in the cluster")
 	}
 
-	singleNodeType := true
 	for i := 1; i < len(nodes.Items); i++ {
 		if nodes.Items[i].Labels["node.kubernetes.io/instance-type"] != nodes.Items[i-1].Labels["node.kubernetes.io/instance-type"] {
-			singleNodeType = false
-			break
+			return ctx, fmt.Errorf("node types are not the same, all node types must be the same in the cluster")
 		}
 	}
-	if !singleNodeType {
-		return ctx, fmt.Errorf("node types are not the same, all node types must be the same in the cluster")
-	}
 
 	if *nodeType != "" {
 		count := 0

test/cases/nvidia/main_test.go

@@ -9,9 +9,9 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"os/signal"
 	"slices"
 	"testing"
-	"time"
 
 	fwext "github.com/aws/aws-k8s-tester/internal/e2e"
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -31,6 +31,7 @@ var (
 	installDevicePlugin    *bool
 	efaEnabled             *bool
 	nvidiaTestImage        *string
+	pytorchImage           *string
 	skipUnitTestSubcommand *string
 	nodeCount              int
 	gpuPerNode             int
@@ -99,15 +100,11 @@ func checkNodeTypes(ctx context.Context, config *envconf.Config) (context.Contex
 		return ctx, err
 	}
 
-	singleNodeType := true
 	for i := 1; i < len(nodes.Items)-1; i++ {
 		if nodes.Items[i].Labels["node.kubernetes.io/instance-type"] != nodes.Items[i-1].Labels["node.kubernetes.io/instance-type"] {
-			singleNodeType = false
+			return ctx, fmt.Errorf("Node types are not the same, all node types must be the same in the cluster")
 		}
 	}
-	if !singleNodeType {
-		return ctx, fmt.Errorf("Node types are not the same, all node types must be the same in the cluster")
-	}
 
 	if *nodeType != "" {
 		for _, v := range nodes.Items {
@@ -135,6 +132,7 @@ func checkNodeTypes(ctx context.Context, config *envconf.Config) (context.Contex
 func TestMain(m *testing.M) {
 	nodeType = flag.String("nodeType", "", "node type for the tests")
 	nvidiaTestImage = flag.String("nvidiaTestImage", "", "nccl test image for nccl tests")
+	pytorchImage = flag.String("pytorchImage", "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.1.0-gpu-py310-cu121-ubuntu20.04-ec2", "pytorch cuda image for single node tests")
 	efaEnabled = flag.Bool("efaEnabled", false, "enable efa tests")
 	installDevicePlugin = flag.Bool("installDevicePlugin", true, "install nvidia device plugin")
 	skipUnitTestSubcommand = flag.String("skipUnitTestSubcommand", "", "optional command to skip specified unit test, `-s test1|test2|...`")
@@ -143,7 +141,7 @@ func TestMain(m *testing.M) {
 		log.Fatalf("failed to initialize test environment: %v", err)
 	}
 	testenv = env.NewWithConfig(cfg)
-	ctx, cancel := context.WithTimeout(context.Background(), 55*time.Minute)
+	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
 	defer cancel()
 	testenv = testenv.WithContext(ctx)
 

test/cases/nvidia/manifests/job-hpc-benchmarks.yaml

@@ -35,11 +35,15 @@ spec:
         - hpl.sh 
         - --mem-affinity 
         - 0:0:0:0:1:1:1:1 
+        # --cpu-affinity needs to be tuned depending on the number of CPUs
+        # available on the instance type.
         - --cpu-affinity 
         - 0-13:14-27:28-41:42-55:56-69:70-83:84-97:98-111
         - --no-multinode 
         - --dat 
         - hpl-linux-x86_64/sample-dat/HPL-dgx-1N.dat
+        # TODO: the path differs for arm64
+        #- hpl-linux-aarch64-gpu/sample-dat/HPL-dgx-1N.dat
         volumeMounts:
         - mountPath: /dev/shm
           name: dshm

test/cases/nvidia/manifests/job-unit-test-single-node.yaml

@@ -17,8 +17,12 @@ spec:
         - /bin/bash
         - ./gpu_unit_tests/unit_test
         env:
-            - name: SKIP_TESTS_SUBCOMMAND
-              value: {{.SkipTestSubcommand}}
+          - name: SKIP_TESTS_SUBCOMMAND
+            value: {{.SkipTestSubcommand}}
+          # because we started building these from source, this is just a
+          # regular binary.
+          - name: DEMO_SUITE_DIR
+            value: /usr/bin
         imagePullPolicy: Always
         resources:
           limits:
@@ -29,4 +33,4 @@ spec:
             cpu: "1"
             memory: 1Gi
       restartPolicy: Never
-  backoffLimit: 4
+  backoffLimit: 4

test/cases/nvidia/manifests/mpi-job-pytorch-training-single-node.yaml

@@ -16,7 +16,7 @@ spec:
          spec:
            restartPolicy: OnFailure
            containers:
-           - image: 763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.1.0-gpu-py310-cu121-ubuntu20.04-ec2
+           - image: {{.PytorchTestImage}}
              name: gpu-test
              command:
               - mpirun
@@ -48,7 +48,7 @@ spec:
               - MXNET_CUDNN_AUTOTUNE_DEFAULT=0
               - python
               - -c
-              - import os; os.system("git clone https://github.com/pytorch/examples.git /pytorch-examples"); os.system("git -C pytorch-examples checkout 0f0c9131ca5c79d1332dce1f4c06fe942fbdc665"); os.system("python /pytorch-examples/mnist/main.py --epochs 1")
+              - import os; os.system("git clone https://github.com/pytorch/examples.git pytorch-examples"); os.system("git -C pytorch-examples checkout 0f0c9131ca5c79d1332dce1f4c06fe942fbdc665"); os.system("python pytorch-examples/mnist/main.py --epochs 1")
              resources:
                limits:
                  nvidia.com/gpu: 1

test/cases/nvidia/mpi_test.go

@@ -48,7 +48,15 @@ func TestMPIJobPytorchTraining(t *testing.T) {
 		WithLabel("hardware", "gpu").
 		Setup(func(ctx context.Context, t *testing.T, cfg *envconf.Config) context.Context {
 			t.Log("Applying single node manifest")
-			err := fwext.ApplyManifests(cfg.Client().RESTConfig(), mpiJobPytorchTrainingSingleNodeManifest)
+			renderedSingleNodeManifest, err := fwext.RenderManifests(mpiJobPytorchTrainingSingleNodeManifest, struct {
+				PytorchTestImage string
+			}{
+				PytorchTestImage: *pytorchImage,
+			})
+			if err != nil {
+				t.Fatal(err)
+			}
+			err = fwext.ApplyManifests(cfg.Client().RESTConfig(), renderedSingleNodeManifest)
 			if err != nil {
 				t.Fatal(err)
 			}

test/cases/nvidia/unit_test.go

@@ -75,13 +75,13 @@ func TestSingleNodeUnitTest(t *testing.T) {
 				ObjectMeta: metav1.ObjectMeta{Name: "unit-test-job", Namespace: "default"},
 			})
 			if err != nil {
-				t.Fatal(err)
+				t.Error(err)
 			}
 			t.Log("Test log for unit-test-job:")
 			t.Log(log)
 			err = fwext.DeleteManifests(cfg.Client().RESTConfig(), renderedJobUnitTestSingleNodeManifest)
 			if err != nil {
-				t.Fatal(err)
+				t.Error(err)
 			}
 			return ctx
 		}).
@@ -120,13 +120,13 @@ func TestSingleNodeUnitTest(t *testing.T) {
 				ObjectMeta: metav1.ObjectMeta{Name: "hpc-benckmarks-job", Namespace: "default"},
 			})
 			if err != nil {
-				t.Fatal(err)
+				t.Error(err)
 			}
 			t.Log("Test log for hpc-benckmarks-job:")
 			t.Log(log)
 			err = fwext.DeleteManifests(cfg.Client().RESTConfig(), renderedJobHpcBenchmarksSingleNodeManifest)
 			if err != nil {
-				t.Fatal(err)
+				t.Error(err)
 			}
 			return ctx
 		}).