fix: broken unit tests for trcomp

Author: Lokiiiiii

tests/conftest.py

@@ -265,14 +265,21 @@ def huggingface_training_compiler_tensorflow_version(huggingface_training_compil
 
 
 @pytest.fixture(scope="module")
-def huggingface_training_compiler_py_version(huggingface_training_compiler_tensorflow_version):
+def huggingface_training_compiler_tensorflow_py_version(
+    huggingface_training_compiler_tensorflow_version,
+):
     return (
         "py37"
         if Version(huggingface_training_compiler_tensorflow_version) < Version("2.6")
         else "py38"
     )
 
 
+@pytest.fixture(scope="module")
+def huggingface_training_compiler_pytorch_py_version(huggingface_training_compiler_pytorch_version):
+    return "py38"
+
+
 @pytest.fixture(scope="module")
 def huggingface_pytorch_latest_training_py_version(huggingface_training_pytorch_latest_version):
     return (
@@ -545,6 +552,11 @@ def _huggingface_base_fm_version(huggingface_version, base_fw, fixture_prefix):
             if len(original_version.split(".")) == 2:
                 base_fw_version = ".".join(base_fw_version.split(".")[:-1])
             versions.append(base_fw_version)
+    if not versions:
+        pytest.skip(
+            f"{fixture_prefix} version {huggingface_version} does not have an"
+            f" image URI configuration for {base_fw}"
+        )
     return sorted(versions, reverse=True)
 
 

tests/unit/sagemaker/training_compiler/test_huggingface_pytorch_compiler.py

@@ -18,6 +18,7 @@
 
 import pytest
 from mock import MagicMock, Mock, patch, ANY
+from packaging.version import Version
 
 from sagemaker import image_uris
 from sagemaker.huggingface import HuggingFace, TrainingCompilerConfig
@@ -96,7 +97,9 @@ def _get_full_gpu_image_uri(
     )
 
 
-def _create_train_job(version, base_framework_version, instance_type, training_compiler_config):
+def _create_train_job(
+    version, base_framework_version, instance_type, training_compiler_config, instance_count=1
+):
     return {
         "image_uri": _get_full_gpu_image_uri(
             version, base_framework_version, instance_type, training_compiler_config
@@ -118,7 +121,7 @@ def _create_train_job(version, base_framework_version, instance_type, training_c
         "output_config": {"S3OutputPath": "s3://{}/".format(BUCKET_NAME)},
         "resource_config": {
             "InstanceType": instance_type,
-            "InstanceCount": 1,
+            "InstanceCount": instance_count,
             "VolumeSizeInGB": 30,
         },
         "hyperparameters": {
@@ -276,6 +279,8 @@ def test_unsupported_instance_group(
     huggingface_training_compiler_version,
     huggingface_training_compiler_pytorch_version,
 ):
+    if Version(huggingface_training_compiler_pytorch_version) < Version("1.11"):
+        pytest.skip("This test is intended for PyTorch 1.11 and above")
     with pytest.raises(ValueError):
         HuggingFace(
             py_version="py38",
@@ -296,6 +301,8 @@ def test_unsupported_distribution(
     huggingface_training_compiler_version,
     huggingface_training_compiler_pytorch_version,
 ):
+    if Version(huggingface_training_compiler_pytorch_version) < Version("1.11"):
+        pytest.skip("This test is intended for PyTorch 1.11 and above")
     with pytest.raises(ValueError):
         HuggingFace(
             py_version="py38",
@@ -383,6 +390,7 @@ def test_pytorchxla_distribution(
         f"pytorch{huggingface_training_compiler_pytorch_version}",
         instance_type,
         compiler_config,
+        instance_count=2,
     )
     expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs
     expected_train_args["enable_sagemaker_metrics"] = False
@@ -642,7 +650,7 @@ def test_register_hf_pytorch_model_auto_infer_framework(
     sagemaker_session,
     huggingface_training_compiler_version,
     huggingface_training_compiler_pytorch_version,
-    huggingface_training_compiler_py_version,
+    huggingface_training_compiler_pytorch_py_version,
 ):
 
     model_package_group_name = "test-hf-tfs-register-model"
@@ -657,7 +665,7 @@ def test_register_hf_pytorch_model_auto_infer_framework(
         role=ROLE,
         transformers_version=huggingface_training_compiler_version,
         pytorch_version=huggingface_training_compiler_pytorch_version,
-        py_version=huggingface_training_compiler_py_version,
+        py_version=huggingface_training_compiler_pytorch_py_version,
         sagemaker_session=sagemaker_session,
     )
 

tests/unit/sagemaker/training_compiler/test_huggingface_tensorflow_compiler.py

@@ -158,18 +158,18 @@ def _create_train_job(
 def test_unsupported_BYOC(
     huggingface_training_compiler_version,
     huggingface_training_compiler_tensorflow_version,
-    huggingface_training_compiler_py_version,
+    huggingface_training_compiler_tensorflow_py_version,
 ):
     byoc = (
         f"1.dkr.ecr.us-east-1.amazonaws.com/huggingface-tensorflow-trcomp-training:"
         f"2.6.3-"
         f"transformers4.17.0-gpu-"
-        f"{huggingface_training_compiler_py_version}-cu112-ubuntu20.04"
+        f"{huggingface_training_compiler_tensorflow_py_version}-cu112-ubuntu20.04"
     )
     with pytest.raises(ValueError):
         HuggingFace(
             image_uri=byoc,
-            py_version=huggingface_training_compiler_py_version,
+            py_version=huggingface_training_compiler_tensorflow_py_version,
             entry_point=SCRIPT_PATH,
             role=ROLE,
             instance_count=INSTANCE_COUNT,
@@ -185,11 +185,11 @@ def test_unsupported_cpu_instance(
     cpu_instance_type,
     huggingface_training_compiler_version,
     huggingface_training_compiler_tensorflow_version,
-    huggingface_training_compiler_py_version,
+    huggingface_training_compiler_tensorflow_py_version,
 ):
     with pytest.raises(ValueError):
         HuggingFace(
-            py_version=huggingface_training_compiler_py_version,
+            py_version=huggingface_training_compiler_tensorflow_py_version,
             entry_point=SCRIPT_PATH,
             role=ROLE,
             instance_count=INSTANCE_COUNT,
@@ -206,11 +206,11 @@ def test_unsupported_gpu_instance(
     unsupported_gpu_instance_class,
     huggingface_training_compiler_version,
     huggingface_training_compiler_tensorflow_version,
-    huggingface_training_compiler_py_version,
+    huggingface_training_compiler_tensorflow_py_version,
 ):
     with pytest.raises(ValueError):
         HuggingFace(
-            py_version=huggingface_training_compiler_py_version,
+            py_version=huggingface_training_compiler_tensorflow_py_version,
             entry_point=SCRIPT_PATH,
             role=ROLE,
             instance_count=INSTANCE_COUNT,
@@ -224,11 +224,11 @@ def test_unsupported_gpu_instance(
 
 def test_unsupported_framework_version(
     huggingface_training_compiler_version,
-    huggingface_training_compiler_py_version,
+    huggingface_training_compiler_tensorflow_py_version,
 ):
     with pytest.raises(ValueError):
         HuggingFace(
-            py_version=huggingface_training_compiler_py_version,
+            py_version=huggingface_training_compiler_tensorflow_py_version,
             entry_point=SCRIPT_PATH,
             role=ROLE,
             instance_count=INSTANCE_COUNT,
@@ -244,11 +244,11 @@ def test_unsupported_framework_version(
 
 def test_unsupported_framework_mxnet(
     huggingface_training_compiler_version,
-    huggingface_training_compiler_py_version,
+    huggingface_training_compiler_tensorflow_py_version,
 ):
     with pytest.raises(ValueError):
         HuggingFace(
-            py_version=huggingface_training_compiler_py_version,
+            py_version=huggingface_training_compiler_tensorflow_py_version,
             entry_point=SCRIPT_PATH,
             role=ROLE,
             instance_count=INSTANCE_COUNT,
@@ -323,13 +323,13 @@ def test_default_compiler_config(
     huggingface_training_compiler_version,
     huggingface_training_compiler_tensorflow_version,
     instance_class,
-    huggingface_training_compiler_py_version,
+    huggingface_training_compiler_tensorflow_py_version,
 ):
     compiler_config = TrainingCompilerConfig()
     instance_type = f"ml.{instance_class}.xlarge"
 
     hf = HuggingFace(
-        py_version=huggingface_training_compiler_py_version,
+        py_version=huggingface_training_compiler_tensorflow_py_version,
         entry_point=SCRIPT_PATH,
         role=ROLE,
         sagemaker_session=sagemaker_session,
@@ -355,7 +355,7 @@ def test_default_compiler_config(
         f"tensorflow{huggingface_training_compiler_tensorflow_version}",
         instance_type,
         compiler_config,
-        huggingface_training_compiler_py_version,
+        huggingface_training_compiler_tensorflow_py_version,
     )
     expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs
     expected_train_args["enable_sagemaker_metrics"] = False
@@ -382,12 +382,12 @@ def test_debug_compiler_config(
     sagemaker_session,
     huggingface_training_compiler_version,
     huggingface_training_compiler_tensorflow_version,
-    huggingface_training_compiler_py_version,
+    huggingface_training_compiler_tensorflow_py_version,
 ):
     compiler_config = TrainingCompilerConfig(debug=True)
 
     hf = HuggingFace(
-        py_version=huggingface_training_compiler_py_version,
+        py_version=huggingface_training_compiler_tensorflow_py_version,
         entry_point=SCRIPT_PATH,
         role=ROLE,
         sagemaker_session=sagemaker_session,
@@ -413,7 +413,7 @@ def test_debug_compiler_config(
         f"tensorflow{huggingface_training_compiler_tensorflow_version}",
         INSTANCE_TYPE,
         compiler_config,
-        huggingface_training_compiler_py_version,
+        huggingface_training_compiler_tensorflow_py_version,
     )
     expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs
     expected_train_args["enable_sagemaker_metrics"] = False
@@ -440,12 +440,12 @@ def test_disable_compiler_config(
     sagemaker_session,
     huggingface_training_compiler_version,
     huggingface_training_compiler_tensorflow_version,
-    huggingface_training_compiler_py_version,
+    huggingface_training_compiler_tensorflow_py_version,
 ):
     compiler_config = TrainingCompilerConfig(enabled=False)
 
     hf = HuggingFace(
-        py_version=huggingface_training_compiler_py_version,
+        py_version=huggingface_training_compiler_tensorflow_py_version,
         entry_point=SCRIPT_PATH,
         role=ROLE,
         sagemaker_session=sagemaker_session,
@@ -471,7 +471,7 @@ def test_disable_compiler_config(
         f"tensorflow{huggingface_training_compiler_tensorflow_version}",
         INSTANCE_TYPE,
         compiler_config,
-        huggingface_training_compiler_py_version,
+        huggingface_training_compiler_tensorflow_py_version,
     )
     expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs
     expected_train_args["enable_sagemaker_metrics"] = False
@@ -492,13 +492,16 @@ def test_disable_compiler_config(
     ["compiler_enabled", "debug_enabled"], [(True, False), (True, True), (False, False)]
 )
 def test_attach(
-    sagemaker_session, compiler_enabled, debug_enabled, huggingface_training_compiler_py_version
+    sagemaker_session,
+    compiler_enabled,
+    debug_enabled,
+    huggingface_training_compiler_tensorflow_py_version,
 ):
     training_image = (
         f"1.dkr.ecr.us-east-1.amazonaws.com/huggingface-tensorflow-trcomp-training:"
         f"2.6.3-"
         f"transformers4.17.0-gpu-"
-        f"{huggingface_training_compiler_py_version}-cu112-ubuntu20.04"
+        f"{huggingface_training_compiler_tensorflow_py_version}-cu112-ubuntu20.04"
     )
     returned_job_description = {
         "AlgorithmSpecification": {"TrainingInputMode": "File", "TrainingImage": training_image},
@@ -532,7 +535,7 @@ def test_attach(
 
     estimator = HuggingFace.attach(training_job_name="trcomp", sagemaker_session=sagemaker_session)
     assert estimator.latest_training_job.job_name == "trcomp"
-    assert estimator.py_version == huggingface_training_compiler_py_version
+    assert estimator.py_version == huggingface_training_compiler_tensorflow_py_version
     assert estimator.framework_version == "4.17.0"
     assert estimator.tensorflow_version == "2.6.3"
     assert estimator.role == "arn:aws:iam::366:role/SageMakerRole"
@@ -557,7 +560,7 @@ def test_register_hf_tfs_model_auto_infer_framework(
     sagemaker_session,
     huggingface_training_compiler_version,
     huggingface_training_compiler_tensorflow_version,
-    huggingface_training_compiler_py_version,
+    huggingface_training_compiler_tensorflow_py_version,
 ):
 
     model_package_group_name = "test-hf-tfs-register-model"
@@ -572,7 +575,7 @@ def test_register_hf_tfs_model_auto_infer_framework(
         role=ROLE,
         transformers_version=huggingface_training_compiler_version,
         tensorflow_version=huggingface_training_compiler_tensorflow_version,
-        py_version=huggingface_training_compiler_py_version,
+        py_version=huggingface_training_compiler_tensorflow_py_version,
         sagemaker_session=sagemaker_session,
     )
 

tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py

@@ -41,7 +41,7 @@
 ROLE = "Dummy"
 REGION = "us-east-1"
 GPU = "ml.p3.2xlarge"
-SUPPORTED_GPU_INSTANCE_CLASSES = {"p3", "p3dn", "g4dn", "p4dn", "g5"}
+SUPPORTED_GPU_INSTANCE_CLASSES = {"p3", "p3dn", "g4dn", "p4d", "g5"}
 UNSUPPORTED_GPU_INSTANCE_CLASSES = EC2_GPU_INSTANCE_CLASSES - SUPPORTED_GPU_INSTANCE_CLASSES
 
 LIST_TAGS_RESULT = {"Tags": [{"Key": "TagtestKey", "Value": "TagtestValue"}]}