Update DLAMI BASE AMI Logic to switch between OSS and Proprietary Nvidia Driver AMI

pytorch/training/docker/2.2/py3/cu121/Dockerfile.gpu

@@ -63,7 +63,7 @@ ENV TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
 ENV CUDNN_VERSION=8.9.2.26
 ENV NCCL_VERSION=2.19.4
 ENV EFA_VERSION=1.30.0
-ENV GDRCOPY_VERSION=2.3.1
+ENV GDRCOPY_VERSION=2.4.1
 
 ENV CMAKE_PREFIX_PATH="$(dirname $(which conda))/../"
 ENV OPEN_MPI_PATH=/opt/amazon/openmpi

test/dlc_tests/benchmark/ec2/pytorch/inference/test_performance_pytorch_inference.py

@@ -1,6 +1,7 @@
 import os
 import time
 import pytest
+
 from src.benchmark_metrics import (
     PYTORCH_INFERENCE_GPU_THRESHOLD,
     PYTORCH_INFERENCE_CPU_THRESHOLD,
@@ -10,6 +11,8 @@
     CONTAINER_TESTS_PREFIX,
     get_framework_and_version_from_tag,
     UL20_CPU_ARM64_US_WEST_2,
+    login_to_ecr_registry,
+    get_account_id_from_image_uri,
     LOGGER,
 )
 from test.test_utils.ec2 import (
@@ -85,7 +88,8 @@ def ec2_performance_pytorch_inference(
     repo_name, image_tag = image_uri.split("/")[-1].split(":")
 
     # Make sure we are logged into ECR so we can pull the image
-    ec2_connection.run(f"$(aws ecr get-login --no-include-email --region {region})", hide=True)
+    account_id = get_account_id_from_image_uri(image_uri)
+    login_to_ecr_registry(ec2_connection, account_id, region)
 
     ec2_connection.run(f"{docker_cmd} pull -q {image_uri} ")
 

test/dlc_tests/benchmark/ec2/pytorch/training/test_performance_pytorch_training.py

@@ -10,6 +10,8 @@
     DEFAULT_REGION,
     get_framework_and_version_from_tag,
     is_pr_context,
+    login_to_ecr_registry,
+    get_account_id_from_image_uri,
 )
 from test.test_utils.ec2 import (
     execute_ec2_training_performance_test,
@@ -142,7 +144,8 @@ def execute_pytorch_gpu_py3_imagenet_ec2_training_performance_test(
     container_name = f"{repo_name}-performance-{image_tag}-ec2"
 
     # Make sure we are logged into ECR so we can pull the image
-    connection.run(f"$(aws ecr get-login --no-include-email --region {region})", hide=True)
+    account_id = get_account_id_from_image_uri(ecr_uri)
+    login_to_ecr_registry(connection, account_id, region)
     # Do not add -q to docker pull as it leads to a hang for huge images like trcomp
     connection.run(f"nvidia-docker pull {ecr_uri}")
     timestamp = time.strftime("%Y-%m-%d-%H-%M-%S")

test/dlc_tests/benchmark/ec2/tensorflow/inference/test_performance_tensorflow_inference.py

@@ -14,6 +14,8 @@
     is_pr_context,
     is_tf_version,
     UL20_BENCHMARK_CPU_ARM64_US_WEST_2,
+    login_to_ecr_registry,
+    get_account_id_from_image_uri,
 )
 from test.test_utils.ec2 import (
     ec2_performance_upload_result_to_s3_and_validate,
@@ -83,7 +85,8 @@ def ec2_performance_tensorflow_inference(
     num_iterations = 500 if is_pr_context() or is_graviton else 1000
 
     # Make sure we are logged into ECR so we can pull the image
-    ec2_connection.run(f"$(aws ecr get-login --no-include-email --region {region})", hide=True)
+    account_id = get_account_id_from_image_uri(image_uri)
+    login_to_ecr_registry(ec2_connection, account_id, region)
     ec2_connection.run(f"{docker_cmd} pull -q {image_uri} ")
     if is_graviton:
         # TF training binary is used that is compatible for graviton instance type

test/dlc_tests/conftest.py

@@ -31,11 +31,7 @@
     is_nightly_context,
     DEFAULT_REGION,
     P3DN_REGION,
-    UBUNTU_20_BASE_DLAMI_US_EAST_1,
-    UBUNTU_20_BASE_DLAMI_US_WEST_2,
     PT_GPU_PY3_BENCHMARK_IMAGENET_AMI_US_EAST_1,
-    AML2_BASE_DLAMI_US_WEST_2,
-    AML2_BASE_DLAMI_US_EAST_1,
     KEYS_TO_DESTROY_FILE,
     are_efa_tests_disabled,
     get_repository_and_tag_from_image_uri,
@@ -330,18 +326,11 @@ def ec2_instance_role_name(request):
 
 
 @pytest.fixture(scope="function")
-def ec2_instance_ami(request, region):
+def ec2_instance_ami(request, region, ec2_instance_type):
     return (
         request.param
         if hasattr(request, "param")
-        else UBUNTU_20_BASE_DLAMI_US_EAST_1
-        if region == "us-east-1"
-        else UBUNTU_20_BASE_DLAMI_US_WEST_2
-        if region == "us-west-2"
-        else test_utils.get_ami_id_boto3(
-            region_name=region,
-            ami_name_pattern="Deep Learning Base GPU AMI (Ubuntu 20.04) ????????",
-        )
+        else test_utils.get_instance_type_base_dlami(ec2_instance_type, region)
     )
 
 
@@ -555,6 +544,8 @@ def ec2_instance(
 ):
     _validate_p4de_usage(request, ec2_instance_type)
     if ec2_instance_type == "p3dn.24xlarge":
+        # Keep track of initial region to get information about previous AMI
+        initial_region = region
         region = P3DN_REGION
         ec2_client = boto3.client(
             "ec2", region_name=region, config=Config(retries={"max_attempts": 10})
@@ -563,10 +554,16 @@ def ec2_instance(
             "ec2", region_name=region, config=Config(retries={"max_attempts": 10})
         )
         if ec2_instance_ami != PT_GPU_PY3_BENCHMARK_IMAGENET_AMI_US_EAST_1:
+            # Assign as AML2 if initial AMI is AML2, else use default
             ec2_instance_ami = (
-                AML2_BASE_DLAMI_US_EAST_1
-                if ec2_instance_ami == AML2_BASE_DLAMI_US_WEST_2
-                else UBUNTU_20_BASE_DLAMI_US_EAST_1
+                test_utils.get_instance_type_base_dlami(
+                    ec2_instance_type, region, linux_dist="AML2"
+                )
+                if ec2_instance_ami
+                == test_utils.get_instance_type_base_dlami(
+                    ec2_instance_type, initial_region, linux_dist="AML2"
+                )
+                else test_utils.get_instance_type_base_dlami(ec2_instance_type, region)
             )
 
     ec2_key_name = f"{ec2_key_name}-{str(uuid.uuid4())}"

test/dlc_tests/container_tests/bin/efa/testEFASanity

@@ -24,4 +24,4 @@ lsmod | grep ib_uverbs
 ibv_devinfo
 
 # check if gdr device is loaded
-grep -e '^1$' /sys/class/infiniband/**/device/gdr
+cat /sys/class/infiniband/**/device/p2p | grep 'NVIDIA'

test/dlc_tests/ec2/mxnet/inference/test_mxnet_inference.py

@@ -4,7 +4,12 @@
 import test.test_utils.ec2 as ec2_utils
 
 from test import test_utils
-from test.test_utils import CONTAINER_TESTS_PREFIX, get_framework_and_version_from_tag
+from test.test_utils import (
+    CONTAINER_TESTS_PREFIX,
+    get_framework_and_version_from_tag,
+    login_to_ecr_registry,
+    get_account_id_from_image_uri,
+)
 from test.test_utils.ec2 import (
     get_ec2_instance_type,
     execute_ec2_inference_test,
@@ -188,7 +193,8 @@ def run_ec2_mxnet_inference(
             f" {image_uri} {mms_inference_cmd}"
         )
     try:
-        ec2_connection.run(f"$(aws ecr get-login --no-include-email --region {region})", hide=True)
+        account_id = get_account_id_from_image_uri(image_uri)
+        login_to_ecr_registry(ec2_connection, account_id, region)
         LOGGER.info(docker_run_cmd)
         ec2_connection.run(docker_run_cmd, hide=True)
         if model_name == SQUEEZENET_MODEL:

test/dlc_tests/ec2/pytorch/inference/test_pytorch_inference.py

@@ -1,8 +1,6 @@
 import os
 import sys
-import time
 import logging
-from datetime import date, timedelta, datetime
 
 import pytest
 from packaging.version import Version
@@ -14,15 +12,15 @@
     CONTAINER_TESTS_PREFIX,
     get_framework_and_version_from_tag,
     get_inference_server_type,
-    get_cuda_version_from_tag,
+    login_to_ecr_registry,
+    get_account_id_from_image_uri,
 )
 from test.test_utils.ec2 import (
     get_ec2_instance_type,
     execute_ec2_inference_test,
     get_ec2_accelerator_type,
 )
 from test.dlc_tests.conftest import LOGGER
-import boto3
 
 LOGGER = logging.getLogger(__name__)
 LOGGER.addHandler(logging.StreamHandler(sys.stdout))
@@ -238,7 +236,8 @@ def ec2_pytorch_inference(image_uri, processor, ec2_connection, region):
             f" {image_uri} {inference_cmd}"
         )
     try:
-        ec2_connection.run(f"$(aws ecr get-login --no-include-email --region {region})", hide=True)
+        account_id = get_account_id_from_image_uri(image_uri)
+        login_to_ecr_registry(ec2_connection, account_id, region)
         LOGGER.info(docker_run_cmd)
         ec2_connection.run(docker_run_cmd, hide=True)
         server_type = get_inference_server_type(image_uri)

test/dlc_tests/ec2/pytorch/inference/test_pytorch_inference_inductor.py

@@ -4,10 +4,11 @@
 
 from test import test_utils
 from test.test_utils import (
-    CONTAINER_TESTS_PREFIX,
     get_framework_and_version_from_tag,
     get_inference_server_type,
     UL20_CPU_ARM64_US_WEST_2,
+    login_to_ecr_registry,
+    get_account_id_from_image_uri,
 )
 from test.test_utils.ec2 import (
     get_ec2_instance_type,
@@ -52,7 +53,7 @@ def test_ec2_pytorch_inference_cpu_compilation(pytorch_inference, ec2_connection
 @pytest.mark.parametrize("ec2_instance_type", PT_EC2_GRAVITON_INSTANCE_TYPES, indirect=True)
 @pytest.mark.parametrize("ec2_instance_ami", [UL20_CPU_ARM64_US_WEST_2], indirect=True)
 @pytest.mark.team("training-compiler")
-def test_ec2_pytorch_inference_cpu_compilation(
+def test_ec2_pytorch_inference_graviton_compilation(
     pytorch_inference_graviton, ec2_connection, region, cpu_only
 ):
     _, image_framework_version = get_framework_and_version_from_tag(pytorch_inference_graviton)
@@ -77,7 +78,8 @@ def ec2_pytorch_inference(image_uri, processor, ec2_connection, region):
         f" {image_uri} {inference_cmd}"
     )
     try:
-        ec2_connection.run(f"$(aws ecr get-login --no-include-email --region {region})", hide=True)
+        account_id = get_account_id_from_image_uri(image_uri)
+        login_to_ecr_registry(ec2_connection, account_id, region)
         LOGGER.info(docker_run_cmd)
         ec2_connection.run(docker_run_cmd, hide=True)
         server_type = get_inference_server_type(image_uri)

test/dlc_tests/ec2/pytorch/training/common_cases.py

@@ -9,6 +9,8 @@
     CONTAINER_TESTS_PREFIX,
     get_framework_and_version_from_tag,
     get_cuda_version_from_tag,
+    login_to_ecr_registry,
+    get_account_id_from_image_uri,
 )
 from test.test_utils.ec2 import (
     execute_ec2_training_test,
@@ -189,7 +191,8 @@ def pytorch_cudnn_match_gpu(pytorch_training, ec2_connection, region):
     PT 2.1 reintroduces a dependency on CUDNN to support NVDA TransformerEngine. This test is to ensure that torch CUDNN matches system CUDNN in the container.
     """
     container_name = "pt_cudnn_test"
-    ec2_connection.run(f"$(aws ecr get-login --no-include-email --region {region})", hide=True)
+    account_id = get_account_id_from_image_uri(pytorch_training)
+    login_to_ecr_registry(ec2_connection, account_id, region)
     ec2_connection.run(f"docker pull -q {pytorch_training}", hide=True)
     ec2_connection.run(
         f"nvidia-docker run --name {container_name} -itd {pytorch_training}", hide=True
@@ -221,6 +224,10 @@ def pytorch_cudnn_match_gpu(pytorch_training, ec2_connection, region):
     ), f"System CUDNN {system_cudnn} and torch cudnn {cudnn_from_torch} do not match. Please downgrade system CUDNN or recompile torch with correct CUDNN verson."
 
 
+def pytorch_curand_gpu(pytorch_training, ec2_connection):
+    execute_ec2_training_test(ec2_connection, pytorch_training, CURAND_CMD)
+
+
 def pytorch_linear_regression_cpu(pytorch_training, ec2_connection):
     execute_ec2_training_test(
         ec2_connection, pytorch_training, PT_REGRESSION_CMD, container_name="pt_reg"
@@ -240,10 +247,6 @@ def pytorch_telemetry_cpu(pytorch_training, ec2_connection):
     )
 
 
-def curand_gpu(training, ec2_connection):
-    execute_ec2_training_test(ec2_connection, training, CURAND_CMD)
-
-
 def pytorch_training_torchdata(pytorch_training, ec2_connection):
     _, image_framework_version = get_framework_and_version_from_tag(pytorch_training)
     # HACK including PT 1.13 in this condition because the Torchdata 0.5.0 tag includes old tests data

test/dlc_tests/ec2/pytorch/training/test_pytorch_training.py

@@ -12,6 +12,8 @@
     UBUNTU_18_HPU_DLAMI_US_WEST_2,
     get_framework_and_version_from_tag,
     get_cuda_version_from_tag,
+    login_to_ecr_registry,
+    get_account_id_from_image_uri,
 )
 from test.test_utils.ec2 import (
     execute_ec2_training_test,
@@ -747,7 +749,8 @@ def test_pytorch_cudnn_match_gpu(
     PT 2.1 reintroduces a dependency on CUDNN to support NVDA TransformerEngine. This test is to ensure that torch CUDNN matches system CUDNN in the container.
     """
     container_name = "pt_cudnn_test"
-    ec2_connection.run(f"$(aws ecr get-login --no-include-email --region {region})", hide=True)
+    account_id = get_account_id_from_image_uri(pytorch_training)
+    login_to_ecr_registry(ec2_connection, account_id, region)
     ec2_connection.run(f"docker pull -q {pytorch_training}", hide=True)
     ec2_connection.run(
         f"nvidia-docker run --name {container_name} -itd {pytorch_training}", hide=True

test/dlc_tests/ec2/pytorch/training/test_pytorch_training_2_2.py

@@ -34,6 +34,7 @@ def test_pytorch_2_2_gpu(
         (common_cases.nvapex, (pytorch_training, ec2_connection)),
         (common_cases.pytorch_training_torchaudio, (pytorch_training, ec2_connection)),
         (common_cases.pytorch_cudnn_match_gpu, (pytorch_training, ec2_connection, region)),
+        (common_cases.pytorch_curand_gpu, (pytorch_training, ec2_connection)),
     ]
 
     if "sagemaker" in pytorch_training:

test/dlc_tests/ec2/tensorflow/inference/test_tensorflow_inference.py

@@ -1,6 +1,5 @@
 import os
 import re
-import json
 from time import sleep
 import pytest
 
@@ -110,7 +109,8 @@ def test_ec2_tensorflow_inference_gpu_tensorrt(
     )
 
     try:
-        ec2_connection.run(f"$(aws ecr get-login --no-include-email --region {region})", hide=True)
+        account_id = test_utils.get_account_id_from_image_uri(tensorflow_inference)
+        test_utils.login_to_ecr_registry(ec2_connection, account_id, region)
         host_setup_for_tensorflow_inference(serving_folder_path, framework_version, ec2_connection)
         sleep(2)
 
@@ -268,7 +268,8 @@ def run_ec2_tensorflow_inference(
         if not is_neuron:
             train_mnist_model(serving_folder_path, ec2_connection)
             sleep(10)
-        ec2_connection.run(f"$(aws ecr get-login --no-include-email --region {region})", hide=True)
+        account_id = test_utils.get_account_id_from_image_uri(image_uri)
+        test_utils.login_to_ecr_registry(ec2_connection, account_id, region)
         ec2_connection.run(docker_run_cmd, hide=True)
         sleep(20)
         if is_neuron and str(framework_version).startswith(TENSORFLOW2_VERSION):

test/dlc_tests/ec2/test_gdrcopy.py

@@ -23,9 +23,6 @@
 @pytest.mark.team("conda")
 @pytest.mark.integration("gdrcopy")
 @pytest.mark.parametrize("ec2_instance_type,region", EC2_EFA_GPU_INSTANCE_TYPE_AND_REGION)
-@pytest.mark.parametrize(
-    "ec2_instance_ami", [test_utils.UBUNTU_20_BASE_DLAMI_US_WEST_2], indirect=True
-)
 @pytest.mark.skipif(
     is_pr_context() and not are_heavy_instance_ec2_tests_enabled(),
     reason="Skip GDRCopy test in PR context unless explicitly enabled",

test/dlc_tests/ec2/test_smdebug.py

@@ -1,5 +1,4 @@
 import os
-
 import pytest
 
 import test.test_utils as test_utils
@@ -13,6 +12,8 @@
     is_tf_version,
     get_framework_and_version_from_tag,
     is_nightly_context,
+    login_to_ecr_registry,
+    get_account_id_from_image_uri,
 )
 from test.test_utils.ec2 import get_ec2_instance_type
 
@@ -170,7 +171,8 @@ def run_smdebug_test(
     shm_setting = " --shm-size=1g " if ec2_instance_type in large_shm_instance_types else " "
     framework = get_framework_from_image_uri(image_uri)
     container_test_local_dir = os.path.join("$HOME", "container_tests")
-    ec2_connection.run(f"$(aws ecr get-login --no-include-email --region {region})", hide=True)
+    account_id = get_account_id_from_image_uri(image_uri)
+    login_to_ecr_registry(ec2_connection, account_id, region)
     # Do not add -q to docker pull as it leads to a hang for huge images like trcomp
     ec2_connection.run(f"docker pull {image_uri}")
 
@@ -209,7 +211,8 @@ def run_smprofiler_test(
     shm_setting = " --shm-size=1g " if ec2_instance_type in large_shm_instance_types else " "
     framework = get_framework_from_image_uri(image_uri)
     container_test_local_dir = os.path.join("$HOME", "container_tests")
-    ec2_connection.run(f"$(aws ecr get-login --no-include-email --region {region})", hide=True)
+    account_id = get_account_id_from_image_uri(image_uri)
+    login_to_ecr_registry(ec2_connection, account_id, region)
     # Do not add -q to docker pull as it leads to a hang for huge images like trcomp
     ec2_connection.run(f"docker pull {image_uri}")