Fixing logic errors in attach

Author: Lokiiiiii

src/sagemaker/fw_utils.py

@@ -493,7 +493,7 @@ def framework_name_from_image(image_uri):
     # We must support both the legacy and current image name format.
     name_pattern = re.compile(
         r"""^(?:sagemaker(?:-rl)?-)?
-        (tensorflow|mxnet|chainer|pytorch|scikit-learn|xgboost
+        (tensorflow|mxnet|chainer|pytorch|pytorch-trcomp|scikit-learn|xgboost
         |huggingface-tensorflow|huggingface-pytorch
         |huggingface-tensorflow-trcomp|huggingface-pytorch-trcomp)(?:-)?
         (scriptmode|training)?

src/sagemaker/pytorch/estimator.py

@@ -423,6 +423,7 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
         )
         image_uri = init_params.pop("image_uri")
         framework, py_version, tag, _ = framework_name_from_image(image_uri)
+        framework = framework.split("-")[0]
 
         if tag is None:
             framework_version = None

src/sagemaker/pytorch/training_compiler/config.py

@@ -26,7 +26,7 @@
 class TrainingCompilerConfig(BaseConfig):
     """The SageMaker Training Compiler configuration class."""
 
-    SUPPORTED_INSTANCE_CLASS_PREFIXES = ["p3", "g4dn", "p4d", "g5"]
+    SUPPORTED_INSTANCE_CLASS_PREFIXES = ["p3", "p3dn", "g4dn", "p4d", "g5"]
     SUPPORTED_INSTANCE_TYPES_WITH_EFA = [
         "ml.g4dn.8xlarge",
         "ml.g4dn.12xlarge",

tests/unit/sagemaker/training_compiler/test_pytorch_compiler.py

@@ -28,7 +28,7 @@
 from tests.unit.sagemaker.training_compiler import EC2_GPU_INSTANCE_CLASSES
 
 
-DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data")
+DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", "data")
 SCRIPT_PATH = os.path.join(DATA_DIR, "dummy_script.py")
 SERVING_SCRIPT_FILE = "another_dummy_script.py"
 MODEL_DATA = "s3://some/data.tar.gz"