test: Adding tests targeting trcomp support for PT 1.12

Author: Lokiiiiii

src/sagemaker/pytorch/__init__.py

@@ -16,3 +16,5 @@
 from sagemaker.pytorch.estimator import PyTorch  # noqa: F401
 from sagemaker.pytorch.model import PyTorchModel, PyTorchPredictor  # noqa: F401
 from sagemaker.pytorch.processing import PyTorchProcessor  # noqa: F401
+
+from sagemaker.pytorch.training_compiler.config import TrainingCompilerConfig  # noqa: F401

src/sagemaker/training_compiler/config.py

@@ -23,16 +23,12 @@ class TrainingCompilerConfig(object):
     """The SageMaker Training Compiler configuration class."""
 
     DEBUG_PATH = "/opt/ml/output/data/compiler/"
-    SUPPORTED_INSTANCE_CLASS_PREFIXES = ["p3", "g4dn", "p4d", "g5"]
+    SUPPORTED_INSTANCE_CLASS_PREFIXES = ["p3", "p3dn", "g4dn", "p4d", "g5"]
 
     HP_ENABLE_COMPILER = "sagemaker_training_compiler_enabled"
     HP_ENABLE_DEBUG = "sagemaker_training_compiler_debug_mode"
 
-    def __init__(
-        self,
-        enabled=True,
-        debug=False,
-    ):
+    def __init__(self, enabled=True, debug=False):
         """This class initializes a ``TrainingCompilerConfig`` instance.
 
         `Amazon SageMaker Training Compiler
@@ -118,10 +114,7 @@ def _to_hyperparameter_dict(self):
         return compiler_config_hyperparameters
 
     @classmethod
-    def validate(
-        cls,
-        estimator,
-    ):
+    def validate(cls, estimator):
         """Checks if SageMaker Training Compiler is configured correctly.
 
         Args:
@@ -138,19 +131,20 @@ def validate(
             warn_msg = (
                 "Estimator instance_type is a PipelineVariable (%s), "
                 "which has to be interpreted as one of the "
-                "[p3, g4dn, p4d, g5] classes in execution time."
+                "%s classes in execution time."
+            )
+            logger.warning(
+                warn_msg,
+                type(estimator.instance_type),
+                str(cls.SUPPORTED_INSTANCE_CLASS_PREFIXES).replace(",", ""),
             )
-            logger.warning(warn_msg, type(estimator.instance_type))
         elif estimator.instance_type:
             if "local" not in estimator.instance_type:
                 requested_instance_class = estimator.instance_type.split(".")[
                     1
                 ]  # Expecting ml.class.size
                 if not any(
-                    [
-                        requested_instance_class.startswith(i)
-                        for i in cls.SUPPORTED_INSTANCE_CLASS_PREFIXES
-                    ]
+                    [requested_instance_class == i for i in cls.SUPPORTED_INSTANCE_CLASS_PREFIXES]
                 ):
                     error_helper_string = (
                         "Unsupported Instance class {}."

tests/conftest.py

@@ -73,6 +73,7 @@
     "neo_pytorch",
     "neo_tensorflow",
     "pytorch",
+    "pytorch_training_compiler",
     "ray_pytorch",
     "ray_tensorflow",
     "sklearn",

tests/data/huggingface_byoc/requirements.txt

@@ -0,0 +1 @@
+transformers