fix: add version length mismatch validation for HuggingFace

src/sagemaker/huggingface/estimator.py

@@ -189,6 +189,7 @@ def _validate_args(self, image_uri):
         """Placeholder docstring"""
         if image_uri is not None:
             return
+
         if self.framework_version is None and image_uri is None:
             raise ValueError(
                 "transformers_version, and image_uri are both None. "
@@ -204,6 +205,17 @@ def _validate_args(self, image_uri):
                 "tensorflow_version and pytorch_version are both None. "
                 "Specify either tensorflow_version or pytorch_version."
             )
+        base_framework_version_len = (
+            len(self.tensorflow_version.split("."))
+            if self.tensorflow_version is not None
+            else len(self.pytorch_version.split("."))
+        )
+        transformers_version_len = len(self.framework_version.split("."))
+        if transformers_version_len != base_framework_version_len:
+            raise ValueError(
+                "Please use either full version or shortened version for both "
+                "transformers_version, tensorflow_version and pytorch_version."
+            )
 
     def hyperparameters(self):
         """Return hyperparameters used by your custom PyTorch code during model training."""

tests/conftest.py

@@ -190,12 +190,31 @@ def pytorch_inference_py_version(pytorch_inference_version, request):
         return "py3"
 
 
+def _huggingface_pytorch_version(huggingface_vesion):
+    config = image_uris.config_for_framework("huggingface")
+    training_config = config.get("training")
+    original_version = huggingface_vesion
+    if "version_aliases" in training_config:
+        huggingface_vesion = training_config.get("version_aliases").get(
+            huggingface_vesion, huggingface_vesion
+        )
+    version_config = training_config.get("versions").get(huggingface_vesion)
+    for key in list(version_config.keys()):
+        if key.startswith("pytorch"):
+            pt_version = key[7:]
+            if len(original_version.split(".")) == 2:
+                pt_version = ".".join(pt_version.split(".")[:-1])
+            return pt_version
+
+
 @pytest.fixture(scope="module")
 def huggingface_pytorch_version(huggingface_training_version):
-    if Version(huggingface_training_version) <= Version("4.4.2"):
-        return "1.6.0"
-    else:
-        pytest.skip("Skipping Huggingface version.")
+    return _huggingface_pytorch_version(huggingface_training_version)
+
+
+@pytest.fixture(scope="module")
+def huggingface_pytorch_latest_version(huggingface_training_latest_version):
+    return _huggingface_pytorch_version(huggingface_training_latest_version)
 
 
 @pytest.fixture(scope="module")

tests/integ/test_huggingface.py

@@ -30,7 +30,7 @@ def test_huggingface_training(
     sagemaker_session,
     gpu_instance_type,
     huggingface_training_latest_version,
-    huggingface_pytorch_version,
+    huggingface_pytorch_latest_version,
 ):
     with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, "huggingface")
@@ -40,7 +40,7 @@ def test_huggingface_training(
             entry_point="examples/text-classification/run_glue.py",
             role="SageMakerRole",
             transformers_version=huggingface_training_latest_version,
-            pytorch_version=huggingface_pytorch_version,
+            pytorch_version=huggingface_pytorch_latest_version,
             instance_count=1,
             instance_type=gpu_instance_type,
             hyperparameters={

tests/unit/sagemaker/huggingface/test_estimator.py

@@ -168,6 +168,59 @@ def _create_train_job(version, base_framework_version):
     }
 
 
+def test_huggingface_invalid_args():
+    with pytest.raises(ValueError) as error:
+        HuggingFace(
+            py_version="py36",
+            entry_point=SCRIPT_PATH,
+            role=ROLE,
+            instance_count=INSTANCE_COUNT,
+            instance_type=INSTANCE_TYPE,
+            transformers_version="4.2.1",
+            pytorch_version="1.6",
+            enable_sagemaker_metrics=False,
+        )
+    assert "use either full version or shortened version" in str(error)
+
+    with pytest.raises(ValueError) as error:
+        HuggingFace(
+            py_version="py36",
+            entry_point=SCRIPT_PATH,
+            role=ROLE,
+            instance_count=INSTANCE_COUNT,
+            instance_type=INSTANCE_TYPE,
+            pytorch_version="1.6",
+            enable_sagemaker_metrics=False,
+        )
+    assert "transformers_version, and image_uri are both None." in str(error)
+
+    with pytest.raises(ValueError) as error:
+        HuggingFace(
+            py_version="py36",
+            entry_point=SCRIPT_PATH,
+            role=ROLE,
+            instance_count=INSTANCE_COUNT,
+            instance_type=INSTANCE_TYPE,
+            transformers_version="4.2.1",
+            enable_sagemaker_metrics=False,
+        )
+    assert "tensorflow_version and pytorch_version are both None." in str(error)
+
+    with pytest.raises(ValueError) as error:
+        HuggingFace(
+            py_version="py36",
+            entry_point=SCRIPT_PATH,
+            role=ROLE,
+            instance_count=INSTANCE_COUNT,
+            instance_type=INSTANCE_TYPE,
+            transformers_version="4.2",
+            pytorch_version="1.6",
+            tensorflow_version="2.3",
+            enable_sagemaker_metrics=False,
+        )
+    assert "tensorflow_version and pytorch_version are both not None." in str(error)
+
+
 @patch("sagemaker.utils.repack_model", MagicMock())
 @patch("sagemaker.utils.create_tar_file", MagicMock())
 @patch("sagemaker.estimator.name_from_base", return_value=JOB_NAME)