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

dlc_developer_config.toml

@@ -34,11 +34,11 @@ deep_canary_mode = false
 [build]
 # Add in frameworks you would like to build. By default, builds are disabled unless you specify building an image.
 # available frameworks - ["autogluon", "huggingface_tensorflow", "huggingface_pytorch", "huggingface_tensorflow_trcomp", "huggingface_pytorch_trcomp", "pytorch_trcomp", "tensorflow", "mxnet", "pytorch", "stabilityai_pytorch"]
-build_frameworks = []
+build_frameworks = ["pytorch"]
 
 # By default we build both training and inference containers. Set true/false values to determine which to build.
 build_training = true
-build_inference = true
+build_inference = false
 
 # Set do_build to "false" to skip builds and test the latest image built by this PR
 # Note: at least one build is required to set do_build to "false"
@@ -57,8 +57,8 @@ notify_test_failures = false
 sanity_tests = true
   safety_check_test = false
   ecr_scan_allowlist_feature = false
-ecs_tests = true
-eks_tests = true
+ecs_tests = false
+eks_tests = false
 ec2_tests = true
 # Set it to true if you are preparing a Benchmark related PR
 ec2_benchmark_tests = false
@@ -67,7 +67,7 @@ ec2_benchmark_tests = false
 ### default. If false, these types of tests will be skipped while other tests will run as usual.
 ### These tests are run in EC2 test jobs, so ec2_tests must be true if ec2_tests_on_heavy_instances is true.
 ### Off by default (set to false)
-ec2_tests_on_heavy_instances = false
+ec2_tests_on_heavy_instances = true
 
 ### SM specific tests
 ### Off by default
@@ -102,7 +102,7 @@ use_scheduler = false
 
 # Standard Framework Training
 dlc-pr-mxnet-training = ""
-dlc-pr-pytorch-training = ""
+dlc-pr-pytorch-training = "pytorch/training/buildspec-2-2-ec2.yml"
 dlc-pr-tensorflow-2-training = ""
 dlc-pr-autogluon-training = ""
 

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/training/test_performance_pytorch_training.py

@@ -7,6 +7,7 @@
     CONTAINER_TESTS_PREFIX,
     PT_GPU_PY3_BENCHMARK_IMAGENET_AMI_US_WEST_2,
     UBUNTU_18_HPU_DLAMI_US_WEST_2,
+    UBUNTU_20_BASE_PROPRIETARY_DLAMI_US_WEST_2,
     DEFAULT_REGION,
     get_framework_and_version_from_tag,
     is_pr_context,
@@ -54,6 +55,9 @@
 
 @pytest.mark.model("resnet50")
 @pytest.mark.parametrize("ec2_instance_type", [PT_EC2_GPU_SYNTHETIC_INSTANCE_TYPE], indirect=True)
+@pytest.mark.parametrize(
+    "ec2_instance_ami", [UBUNTU_20_BASE_PROPRIETARY_DLAMI_US_WEST_2], indirect=True
+)
 @pytest.mark.team("conda")
 def test_performance_pytorch_gpu_synthetic(
     pytorch_training, ec2_connection, gpu_only, py3_only, ec2_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)
     )
 
 
@@ -564,9 +553,14 @@ def ec2_instance(
         )
         if ec2_instance_ami != PT_GPU_PY3_BENCHMARK_IMAGENET_AMI_US_EAST_1:
             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, "us-east-1", linux_dist="AML2"
+                )
+                if ec2_instance_ami
+                == test_utils.get_instance_type_base_dlami(
+                    ec2_instance_type, "us-west-2", linux_dist="AML2"
+                )
+                else test_utils.get_instance_type_base_dlami(ec2_instance_type, "us-east-1")
             )
 
     ec2_key_name = f"{ec2_key_name}-{str(uuid.uuid4())}"

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

@@ -7,13 +7,15 @@
 
 from test.test_utils import (
     CONTAINER_TESTS_PREFIX,
+    UBUNTU_20_BASE_OSS_DLAMI_US_WEST_2,
     get_framework_and_version_from_tag,
     get_cuda_version_from_tag,
 )
 from test.test_utils.ec2 import (
     execute_ec2_training_test,
     get_ec2_instance_type,
     get_efa_ec2_instance_type,
+    get_efa_ec2_instance_ami,
 )
 
 # Test functions
@@ -44,6 +46,15 @@
     default="g4dn.12xlarge", filter_function=ec2_utils.filter_non_g3_instance_type
 )
 
+# Instance AMI filters
+PT_EC2_CPU_INSTANCE_AMI = [UBUNTU_20_BASE_OSS_DLAMI_US_WEST_2]
+
+PT_EC2_GPU_INSTANCE_AMI = get_efa_ec2_instance_ami(PT_EC2_GPU_INSTANCE_TYPE_AND_REGION)
+
+PT_EC2_GPU_INDUCTOR_INSTANCE_AMI = get_efa_ec2_instance_ami(
+    PT_EC2_GPU_INDUCTOR_INSTANCE_TYPE_AND_REGION
+)
+
 
 def pytorch_standalone(pytorch_training, ec2_connection):
     execute_ec2_training_test(
@@ -219,6 +230,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"
@@ -236,7 +251,3 @@ def pytorch_telemetry_cpu(pytorch_training, ec2_connection):
     execute_ec2_training_test(
         ec2_connection, pytorch_training, PT_TELEMETRY_CMD, timeout=900, container_name="telemetry"
     )
-
-
-def curand_gpu(training, ec2_connection):
-    execute_ec2_training_test(ec2_connection, training, CURAND_CMD)

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

@@ -15,6 +15,7 @@
 @pytest.mark.parametrize(
     "ec2_instance_type, region", common_cases.PT_EC2_GPU_INSTANCE_TYPE_AND_REGION, indirect=True
 )
+@pytest.mark.parametrize("ec2_instance_ami", common_cases.PT_EC2_GPU_INSTANCE_AMI, indirect=True)
 def test_pytorch_2_2_gpu(
     pytorch_training___2__2, ec2_connection, region, gpu_only, ec2_instance_type
 ):
@@ -34,6 +35,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:
@@ -57,6 +59,9 @@ def test_pytorch_2_2_gpu(
     common_cases.PT_EC2_GPU_INDUCTOR_INSTANCE_TYPE_AND_REGION,
     indirect=True,
 )
+@pytest.mark.parametrize(
+    "ec2_instance_ami", common_cases.PT_EC2_GPU_INDUCTOR_INSTANCE_AMI, indirect=True
+)
 def test_pytorch_2_2_gpu_inductor(
     pytorch_training___2__2, ec2_connection, region, gpu_only, ec2_instance_type
 ):
@@ -81,6 +86,7 @@ def test_pytorch_2_2_gpu_inductor(
 @pytest.mark.model("N/A")
 @pytest.mark.team("conda")
 @pytest.mark.parametrize("ec2_instance_type", common_cases.PT_EC2_CPU_INSTANCE_TYPE, indirect=True)
+@pytest.mark.parametrize("ec2_instance_ami", common_cases.PT_EC2_CPU_INSTANCE_AMI, indirect=True)
 def test_pytorch_2_2_cpu(pytorch_training___2__2, ec2_connection, cpu_only):
     pytorch_training = pytorch_training___2__2
 

test/dlc_tests/ec2/test_curand.py

@@ -20,6 +20,9 @@
 @pytest.mark.model("N/A")
 @pytest.mark.team("frameworks")
 @pytest.mark.parametrize("ec2_instance_type", CURAND_EC2_SINGLE_GPU_INSTANCE_TYPE, indirect=True)
+@pytest.mark.parametrize(
+    "ec2_instance_ami", [test_utils.UBUNTU_20_BASE_PROPRIETARY_DLAMI_US_WEST_2], indirect=True
+)
 def test_curand_gpu(training, ec2_connection, gpu_only, ec2_instance_type):
     if test_utils.is_image_incompatible_with_instance_type(training, ec2_instance_type):
         pytest.skip(f"Image {training} is incompatible with instance type {ec2_instance_type}")

test/dlc_tests/ec2/test_gdrcopy.py

@@ -24,7 +24,7 @@
 @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
+    "ec2_instance_ami", [test_utils.UBUNTU_20_BASE_OSS_DLAMI_US_WEST_2], indirect=True
 )
 @pytest.mark.skipif(
     is_pr_context() and not are_heavy_instance_ec2_tests_enabled(),

test/test_utils/__init__.py

@@ -78,18 +78,38 @@ def get_ami_id_ssm(region_name, parameter_path):
     return ami_id
 
 
-# The Ubuntu 20.04 AMI which adds GDRCopy is used only for GDRCopy feature that is supported on PT1.13 and PT2.0
-UBUNTU_20_BASE_DLAMI_US_WEST_2 = get_ami_id_boto3(
-    region_name="us-west-2", ami_name_pattern="Deep Learning Base GPU AMI (Ubuntu 20.04) ????????"
+# DLAMI Base is split between OSS Nvidia Driver and Propietary Nvidia Driver. see https://docs.aws.amazon.com/dlami/latest/devguide/important-changes.html
+UBUNTU_20_BASE_OSS_DLAMI_US_WEST_2 = get_ami_id_boto3(
+    region_name="us-west-2",
+    ami_name_pattern="Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 20.04) ????????",
 )
-UBUNTU_20_BASE_DLAMI_US_EAST_1 = get_ami_id_boto3(
-    region_name="us-east-1", ami_name_pattern="Deep Learning Base GPU AMI (Ubuntu 20.04) ????????"
+UBUNTU_20_BASE_OSS_DLAMI_US_EAST_1 = get_ami_id_boto3(
+    region_name="us-east-1",
+    ami_name_pattern="Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 20.04) ????????",
 )
-AML2_BASE_DLAMI_US_WEST_2 = get_ami_id_boto3(
-    region_name="us-west-2", ami_name_pattern="Deep Learning Base AMI (Amazon Linux 2) Version ??.?"
+UBUNTU_20_BASE_PROPRIETARY_DLAMI_US_WEST_2 = get_ami_id_boto3(
+    region_name="us-west-2",
+    ami_name_pattern="Deep Learning Base Proprietary Nvidia Driver GPU AMI (Ubuntu 20.04) ????????",
 )
-AML2_BASE_DLAMI_US_EAST_1 = get_ami_id_boto3(
-    region_name="us-east-1", ami_name_pattern="Deep Learning Base AMI (Amazon Linux 2) Version ??.?"
+UBUNTU_20_BASE_PROPRIETARY_DLAMI_US_EAST_1 = get_ami_id_boto3(
+    region_name="us-east-1",
+    ami_name_pattern="Deep Learning Base Proprietary Nvidia Driver GPU AMI (Ubuntu 20.04) ????????",
+)
+AML2_BASE_OSS_DLAMI_US_WEST_2 = get_ami_id_boto3(
+    region_name="us-west-2",
+    ami_name_pattern="Deep Learning Base OSS Nvidia Driver AMI (Amazon Linux 2) Version ??.?",
+)
+AML2_BASE_OSS_DLAMI_US_EAST_1 = get_ami_id_boto3(
+    region_name="us-east-1",
+    ami_name_pattern="Deep Learning Base OSS Nvidia Driver AMI (Amazon Linux 2) Version ??.?",
+)
+AML2_BASE_PROPRIETARY_DLAMI_US_WEST_2 = get_ami_id_boto3(
+    region_name="us-west-2",
+    ami_name_pattern="Deep Learning Base Proprietary Nvidia Driver AMI (Amazon Linux 2) Version ??.?",
+)
+AML2_BASE_PROPRIETARY_DLAMI_US_EAST_1 = get_ami_id_boto3(
+    region_name="us-east-1",
+    ami_name_pattern="Deep Learning Base Proprietary Nvidia Driver AMI (Amazon Linux 2) Version ??.?",
 )
 # We use the following DLAMI for MXNet and TensorFlow tests as well, but this is ok since we use custom DLC Graviton containers on top. We just need an ARM base DLAMI.
 UL20_CPU_ARM64_US_WEST_2 = get_ami_id_boto3(
@@ -145,8 +165,10 @@ def get_ami_id_ssm(region_name, parameter_path):
 UBUNTU_18_HPU_DLAMI_US_WEST_2 = "ami-03cdcfc91a96a8f92"
 UBUNTU_18_HPU_DLAMI_US_EAST_1 = "ami-0d83d7487f322545a"
 UL_AMI_LIST = [
-    UBUNTU_20_BASE_DLAMI_US_WEST_2,
-    UBUNTU_20_BASE_DLAMI_US_EAST_1,
+    UBUNTU_20_BASE_OSS_DLAMI_US_WEST_2,
+    UBUNTU_20_BASE_OSS_DLAMI_US_EAST_1,
+    UBUNTU_20_BASE_PROPRIETARY_DLAMI_US_WEST_2,
+    UBUNTU_20_BASE_PROPRIETARY_DLAMI_US_EAST_1,
     UBUNTU_18_HPU_DLAMI_US_WEST_2,
     UBUNTU_18_HPU_DLAMI_US_EAST_1,
     PT_GPU_PY3_BENCHMARK_IMAGENET_AMI_US_EAST_1,
@@ -2370,3 +2392,93 @@ def get_image_spec_from_buildspec(image_uri, dlc_folder_path):
         raise ValueError(f"No corresponding entry found for {image_uri} in {buildspec_path}")
 
     return matched_image_spec
+
+
+def get_instance_type_base_dlami(instance_type, region, linux_dist="UBUNTU_20"):
+    """
+    Get Instance types based on EC2 instance, see https://docs.aws.amazon.com/dlami/latest/devguide/important-changes.html
+    OSS Nvidia Driver DLAMI supports the following: ["g4dn.xlarge",
+                                                     "g4dn.2xlarge",
+                                                     "g4dn.4xlarge",
+                                                     "g4dn.8xlarge",
+                                                     "g4dn.16xlarge",
+                                                     "g4dn.12xlarge",
+                                                     "g4dn.metal",
+                                                     "g4dn.xlarge",
+                                                     "g5.xlarge",
+                                                     "g5.2xlarge",
+                                                     "g5.4xlarge",
+                                                     "g5.8xlarge",
+                                                     "g5.16xlarge",
+                                                     "g5.12xlarge",
+                                                     "g5.24xlarge",
+                                                     "g5.48xlarge",
+                                                     "p4d.24xlarge",
+                                                     "p4de.24xlarge",
+                                                     "p5.48xlarge",]
+
+    Proprietary Nvidia Driver DLAMI supports the following: ["p3.2xlarge",
+                                                             "p3.8xlarge",
+                                                             "p3.16xlarge",
+                                                             "p3dn.24xlarge",
+                                                             "g3s.xlarge",
+                                                             "g3.4xlarge",
+                                                             "g3.8xlarge",
+                                                             "g3.16xlarge",]
+
+    Other instances will default to Proprietary Nvidia Driver DLAMI
+    """
+
+    base_proprietary_dlami_instances = [
+        "p3.2xlarge",
+        "p3.8xlarge",
+        "p3.16xlarge",
+        "p3dn.24xlarge",
+        "g3s.xlarge",
+        "g3.4xlarge",
+        "g3.8xlarge",
+        "g3.16xlarge",
+    ]
+
+    # set defaults
+    if linux_dist == "AML2":
+        oss_dlami_us_east_1 = AML2_BASE_OSS_DLAMI_US_EAST_1
+        oss_dlami_us_west_2 = AML2_BASE_OSS_DLAMI_US_WEST_2
+        oss_dlami_name_pattern = (
+            "Deep Learning Base OSS Nvidia Driver AMI (Amazon Linux 2) Version ??.?"
+        )
+
+        proprietary_dlami_us_east_1 = AML2_BASE_PROPRIETARY_DLAMI_US_EAST_1
+        proprietary_dlami_us_west_2 = AML2_BASE_PROPRIETARY_DLAMI_US_WEST_2
+        proprietary_dlami_name_pattern = (
+            "Deep Learning Base Proprietary Nvidia Driver AMI (Amazon Linux 2) Version ??.?"
+        )
+    else:
+        oss_dlami_us_east_1 = UBUNTU_20_BASE_OSS_DLAMI_US_EAST_1
+        oss_dlami_us_west_2 = UBUNTU_20_BASE_OSS_DLAMI_US_WEST_2
+        oss_dlami_name_pattern = (
+            "Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 20.04) ????????"
+        )
+
+        proprietary_dlami_us_east_1 = UBUNTU_20_BASE_OSS_DLAMI_US_EAST_1
+        proprietary_dlami_us_west_2 = UBUNTU_20_BASE_OSS_DLAMI_US_WEST_2
+        proprietary_dlami_name_pattern = (
+            "Deep Learning Base Proprietary Nvidia Driver GPU AMI (Ubuntu 20.04) ????????"
+        )
+
+    return (
+        proprietary_dlami_us_east_1
+        if region == "us-east-1" and instance_type in base_proprietary_dlami_instances
+        else proprietary_dlami_us_west_2
+        if region == "us-west-2" and instance_type in base_proprietary_dlami_instances
+        else get_ami_id_boto3(
+            region_name=region,
+            ami_name_pattern=proprietary_dlami_name_pattern,
+        )
+        if instance_type in base_proprietary_dlami_instances
+        else oss_dlami_us_east_1
+        if region == "us-east-1"
+        else oss_dlami_us_west_2
+        if region == "us-west-2"
+        else get_ami_id_boto3(region_name=region, ami_name_pattern=oss_dlami_name_pattern)
+    )