fix: add version length mismatch validation for HuggingFace

Author: Rui Wang Napieralski

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

@@ -193,7 +193,10 @@ def pytorch_inference_py_version(pytorch_inference_version, request):
 @pytest.fixture(scope="module")
 def huggingface_pytorch_version(huggingface_training_version):
     if Version(huggingface_training_version) <= Version("4.4.2"):
-        return "1.6.0"
+        if len(huggingface_training_version.split(".")) == 3:
+            return "1.6.0"
+        else:
+            return "1.6"
     else:
         pytest.skip("Skipping Huggingface version.")
 

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)