[RHOAIENG-9004] Adjust existing test and workflow for GPU testing

Author: sutaakar

.github/workflows/e2e_tests.yaml

@@ -27,7 +27,7 @@ concurrency:
 jobs:
   kubernetes-e2e:
 
-    runs-on: ubuntu-20.04-4core
+    runs-on: ubuntu-20.04-4core-gpu
 
     steps:
       - name: Checkout code
@@ -52,33 +52,41 @@ jobs:
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
 
+      - name: Setup NVidia GPU environment for KinD
+        uses: ./common/github-actions/nvidia-gpu-setup
+
       - name: Setup and start KinD cluster
         uses: ./common/github-actions/kind
 
+      - name: Install NVidia GPU operator for KinD
+        uses: ./common/github-actions/nvidia-gpu-operator
+
       - name: Deploy CodeFlare stack
         id: deploy
         run: |
           echo Setting up CodeFlare stack
           make setup-e2e
 
           echo Deploying CodeFlare operator
-          IMG="${REGISTRY_ADDRESS}"/codeflare-operator
-          make image-push -e IMG="${IMG}"
+          IMG=localhost/codeflare-operator:test
+          make image-build -e IMG="${IMG}"
+          podman save -o cfo.tar ${IMG}
+          kind load image-archive cfo.tar --name cluster --verbosity 1000
           make deploy -e IMG="${IMG}" -e ENV="e2e"
           kubectl wait --timeout=120s --for=condition=Available=true deployment -n openshift-operators codeflare-operator-manager
 
       - name: Run e2e tests
         run: |
-          export CODEFLARE_TEST_TIMEOUT_SHORT=1m
-          export CODEFLARE_TEST_TIMEOUT_MEDIUM=5m
-          export CODEFLARE_TEST_TIMEOUT_LONG=10m
+          export CODEFLARE_TEST_TIMEOUT_SHORT=3m
+          export CODEFLARE_TEST_TIMEOUT_MEDIUM=10m
+          export CODEFLARE_TEST_TIMEOUT_LONG=20m
           export CODEFLARE_TEST_TIMEOUT_GPU_PROVISIONING=30m
 
           export CODEFLARE_TEST_OUTPUT_DIR=${{ env.TEMP_DIR }}
           echo "CODEFLARE_TEST_OUTPUT_DIR=${CODEFLARE_TEST_OUTPUT_DIR}" >> $GITHUB_ENV
 
           set -euo pipefail
-          go test -timeout 30m -v ./test/e2e -json 2>&1 | tee ${CODEFLARE_TEST_OUTPUT_DIR}/gotest.log | gotestfmt
+          go test -timeout 60m -v ./test/e2e -json 2>&1 | tee ${CODEFLARE_TEST_OUTPUT_DIR}/gotest.log | gotestfmt
 
       - name: Print CodeFlare operator logs
         if: always() && steps.deploy.outcome == 'success'

Dockerfile

@@ -2,7 +2,7 @@
 
 # BEGIN -- workaround lack of go-toolset for golang 1.22
 
-ARG GOLANG_IMAGE=golang:1.22
+ARG GOLANG_IMAGE=docker.io/library/golang:1.22
 
 ARG GOARCH=amd64
 

config/e2e/patch_resources.yaml

@@ -1,2 +1,5 @@
 - op: remove
   path: /spec/template/spec/containers/0/resources
+- op: replace
+  path: /spec/template/spec/containers/0/imagePullPolicy
+  value: IfNotPresent

test/e2e/mnist.py

@@ -20,7 +20,7 @@
 from pytorch_lightning.callbacks.progress import TQDMProgressBar
 from torch import nn
 from torch.nn import functional as F
-from torch.utils.data import DataLoader, random_split, RandomSampler
+from torch.utils.data import DataLoader, random_split
 from torchmetrics import Accuracy
 from torchvision import transforms
 from torchvision.datasets import MNIST
@@ -36,6 +36,9 @@
 print("MNIST_DATASET_URL: is ", os.getenv("MNIST_DATASET_URL"))
 MNIST_DATASET_URL = os.getenv("MNIST_DATASET_URL")
 
+print("ACCELERATOR: is ", os.getenv("ACCELERATOR"))
+ACCELERATOR = os.getenv("ACCELERATOR")
+
 class LitMNIST(LightningModule):
     def __init__(self, data_dir=PATH_DATASETS, hidden_size=64, learning_rate=2e-4):
 
@@ -158,7 +161,7 @@ def setup(self, stage=None):
             )
 
     def train_dataloader(self):
-        return DataLoader(self.mnist_train, batch_size=BATCH_SIZE, sampler=RandomSampler(self.mnist_train, num_samples=1000))
+        return DataLoader(self.mnist_train, batch_size=BATCH_SIZE)
 
     def val_dataloader(self):
         return DataLoader(self.mnist_val, batch_size=BATCH_SIZE)
@@ -176,13 +179,12 @@ def test_dataloader(self):
 
 # Initialize a trainer
 trainer = Trainer(
-    accelerator="auto",
+    accelerator=ACCELERATOR,
     # devices=1 if torch.cuda.is_available() else None,  # limiting got iPython runs
     max_epochs=3,
     callbacks=[TQDMProgressBar(refresh_rate=20)],
     num_nodes=int(os.environ.get("GROUP_WORLD_SIZE", 1)),
     devices=int(os.environ.get("LOCAL_WORLD_SIZE", 1)),
-    replace_sampler_ddp=False,
     strategy="ddp",
 )
 

test/e2e/mnist_pytorch_appwrapper_test.go

@@ -33,7 +33,6 @@ import (
 // Trains the MNIST dataset as a batch Job in an AppWrapper, and asserts successful completion of the training job.
 func TestMNISTPyTorchAppWrapper(t *testing.T) {
 	test := With(t)
-	test.T().Parallel()
 
 	// Create a namespace and localqueue in that namespace
 	namespace := test.NewTestNamespace()
@@ -85,6 +84,7 @@ func TestMNISTPyTorchAppWrapper(t *testing.T) {
 								{Name: "MNIST_DATASET_URL", Value: GetMnistDatasetURL()},
 								{Name: "PIP_INDEX_URL", Value: GetPipIndexURL()},
 								{Name: "PIP_TRUSTED_HOST", Value: GetPipTrustedHost()},
+								{Name: "ACCELERATOR", Value: "gpu"},
 							},
 							Command: []string{"/bin/sh", "-c", "pip install -r /test/requirements.txt && torchrun /test/mnist.py"},
 							VolumeMounts: []corev1.VolumeMount{

test/e2e/mnist_rayjob_raycluster_test.go

@@ -38,7 +38,6 @@ import (
 // directly managed by Kueue, and asserts successful completion of the training job.
 func TestMNISTRayJobRayCluster(t *testing.T) {
 	test := With(t)
-	test.T().Parallel()
 
 	// Create a namespace and localqueue in that namespace
 	namespace := test.NewTestNamespace()
@@ -91,7 +90,6 @@ func TestMNISTRayJobRayCluster(t *testing.T) {
 // Same as TestMNISTRayJobRayCluster, except the RayCluster is wrapped in an AppWrapper
 func TestMNISTRayJobRayClusterAppWrapper(t *testing.T) {
 	test := With(t)
-	test.T().Parallel()
 
 	// Create a namespace and localqueue in that namespace
 	namespace := test.NewTestNamespace()
@@ -282,11 +280,13 @@ func constructRayCluster(_ Test, namespace *corev1.Namespace, mnist *corev1.Conf
 									Resources: corev1.ResourceRequirements{
 										Requests: corev1.ResourceList{
 											corev1.ResourceCPU:    resource.MustParse("250m"),
-											corev1.ResourceMemory: resource.MustParse("256Mi"),
+											corev1.ResourceMemory: resource.MustParse("1G"),
+											"nvidia.com/gpu":      resource.MustParse("1"),
 										},
 										Limits: corev1.ResourceList{
 											corev1.ResourceCPU:    resource.MustParse("1"),
-											corev1.ResourceMemory: resource.MustParse("2G"),
+											corev1.ResourceMemory: resource.MustParse("4G"),
+											"nvidia.com/gpu":      resource.MustParse("1"),
 										},
 									},
 								},
@@ -320,6 +320,7 @@ func constructRayJob(_ Test, namespace *corev1.Namespace, rayCluster *rayv1.RayC
     MNIST_DATASET_URL: "` + GetMnistDatasetURL() + `"
     PIP_INDEX_URL: "` + GetPipIndexURL() + `"
     PIP_TRUSTED_HOST: "` + GetPipTrustedHost() + `"
+    ACCELERATOR: "gpu"
 `,
 			ClusterSelector: map[string]string{
 				RayJobDefaultClusterSelectorKey: rayCluster.Name,
@@ -336,6 +337,7 @@ func constructRayJob(_ Test, namespace *corev1.Namespace, rayCluster *rayv1.RayC
 					},
 				},
 			},
+			EntrypointNumGpus: 1,
 		},
 	}
 }

test/e2e/setup.sh

@@ -84,12 +84,14 @@ metadata:
 spec:
   namespaceSelector: {} # match all.
   resourceGroups:
-  - coveredResources: ["cpu","memory"]
+  - coveredResources: ["cpu","memory", "nvidia.com/gpu"]
     flavors:
     - name: "default-flavor"
       resources:
       - name: "cpu"
         nominalQuota: 4
       - name: "memory"
-        nominalQuota: "4G"
+        nominalQuota: "20G"
+      - name: "nvidia.com/gpu"
+        nominalQuota: "1"
 EOF