SDK changes for TRCOMP support

Author: BruceChang9688

src/sagemaker/huggingface/training_compiler/config.py

@@ -102,6 +102,16 @@ 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 = (
+                    "Training Compiler is only valid between HuggingFace PyTorch 1.9-1.11 "
+                    "for SageMaker Training Compiler."
+                    " 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.
@@ -102,6 +103,17 @@ def validate(cls, estimator):
                     cls.MIN_SUPPORTED_VERSION, estimator.framework_version
                 )
                 raise ValueError(error_helper_string)
+            elif Version(estimator.framework_version) in SpecifierSet(
+                f"> {cls.MAX_SUPPORTED_VERSION}"
+            ):
+                error_helper_string = (
+                    "SageMaker Training Compiler only supports TensorFlow version "
+                    "<= {} but received {}"
+                )
+                error_helper_string = error_helper_string.format(
+                    cls.MAX_SUPPORTED_VERSION, estimator.framework_version
+                )
+                raise ValueError(error_helper_string)
 
         if estimator.distribution and "multi_worker_mirrored_strategy" in estimator.distribution:
             mwms_enabled = estimator.distribution.get("multi_worker_mirrored_strategy").get(