feature: SDK changes for TRCOMP support (aws#3714)

Author: BruceChang9688

src/sagemaker/huggingface/training_compiler/config.py

@@ -102,6 +102,17 @@ def validate(cls, estimator):
 
         super(TrainingCompilerConfig, cls).validate(estimator)
 
+        if estimator.pytorch_version:
+            if (Version(estimator.pytorch_version) in SpecifierSet("< 1.9")) or (
+                Version(estimator.pytorch_version) in SpecifierSet("> 1.11")
+            ):
+                error_helper_string = (
+                    "SageMaker Training Compiler is only supported "
+                    "with HuggingFace PyTorch 1.9-1.11. "
+                    "Received pytorch_version={} which is unsupported."
+                )
+                raise ValueError(error_helper_string.format(estimator.pytorch_version))
+
         if estimator.image_uri:
             error_helper_string = (
                 "Overriding the image URI is currently not supported "

src/sagemaker/tensorflow/training_compiler/config.py

@@ -26,6 +26,7 @@ class TrainingCompilerConfig(BaseConfig):
 
     SUPPORTED_INSTANCE_CLASS_PREFIXES = ["p3", "p3dn", "g4dn", "p4d", "g5"]
     MIN_SUPPORTED_VERSION = "2.9"
+    MAX_SUPPORTED_VERSION = "2.11"
 
     def __init__(self, enabled=True, debug=False):
         """This class initializes a ``TrainingCompilerConfig`` instance.
@@ -91,15 +92,17 @@ def validate(cls, estimator):
         super(TrainingCompilerConfig, cls).validate(estimator)
 
         if estimator.framework_version:
-            if Version(estimator.framework_version) in SpecifierSet(
-                f"< {cls.MIN_SUPPORTED_VERSION}"
+            if Version(estimator.framework_version) not in SpecifierSet(
+                f">= {cls.MIN_SUPPORTED_VERSION}", f"<= {cls.MAX_SUPPORTED_VERSION}"
             ):
                 error_helper_string = (
                     "SageMaker Training Compiler only supports TensorFlow version "
-                    ">= {} but received {}"
+                    "between {} to {} but received {}"
                 )
                 error_helper_string = error_helper_string.format(
-                    cls.MIN_SUPPORTED_VERSION, estimator.framework_version
+                    cls.MIN_SUPPORTED_VERSION,
+                    cls.MAX_SUPPORTED_VERSION,
+                    estimator.framework_version,
                 )
                 raise ValueError(error_helper_string)
 

tests/unit/sagemaker/training_compiler/test_huggingface_pytorch_compiler.py

@@ -218,7 +218,7 @@ def test_unsupported_gpu_instance(
         ).fit()
 
 
-def test_unsupported_framework_version(
+def test_unsupported_framework_version_min(
     huggingface_training_compiler_version,
 ):
     with pytest.raises(ValueError):
@@ -229,9 +229,24 @@ def test_unsupported_framework_version(
             instance_count=INSTANCE_COUNT,
             instance_type=INSTANCE_TYPE,
             transformers_version=huggingface_training_compiler_version,
-            pytorch_version=".".join(
-                ["99"] * len(huggingface_training_compiler_version.split("."))
-            ),
+            pytorch_version="1.8",
+            enable_sagemaker_metrics=False,
+            compiler_config=TrainingCompilerConfig(),
+        ).fit()
+
+
+def test_unsupported_framework_version_max(
+    huggingface_training_compiler_version,
+):
+    with pytest.raises(ValueError):
+        HuggingFace(
+            py_version="py38",
+            entry_point=SCRIPT_PATH,
+            role=ROLE,
+            instance_count=INSTANCE_COUNT,
+            instance_type=INSTANCE_TYPE,
+            transformers_version=huggingface_training_compiler_version,
+            pytorch_version="1.12",
             enable_sagemaker_metrics=False,
             compiler_config=TrainingCompilerConfig(),
         ).fit()

tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py

@@ -189,7 +189,7 @@ def test_gpu_instance(
                 compiler_config=TrainingCompilerConfig(),
             ).fit()
 
-    def test_framework_version(self, tensorflow_training_py_version):
+    def test_framework_version_min(self, tensorflow_training_py_version):
         with pytest.raises(ValueError):
             TensorFlow(
                 py_version=tensorflow_training_py_version,
@@ -202,6 +202,19 @@ def test_framework_version(self, tensorflow_training_py_version):
                 compiler_config=TrainingCompilerConfig(),
             ).fit()
 
+    def test_framework_version_max(self, tensorflow_training_py_version):
+        with pytest.raises(ValueError):
+            TensorFlow(
+                py_version=tensorflow_training_py_version,
+                entry_point=SCRIPT_PATH,
+                role=ROLE,
+                instance_count=INSTANCE_COUNT,
+                instance_type=INSTANCE_TYPE,
+                framework_version="2.12",
+                enable_sagemaker_metrics=False,
+                compiler_config=TrainingCompilerConfig(),
+            ).fit()
+
     def test_mwms(self, tensorflow_training_version, tensorflow_training_py_version):
         with pytest.raises(ValueError):
             TensorFlow(