go with model base and tf

Author: Dewen Qi

src/sagemaker/estimator.py

@@ -50,6 +50,7 @@
     validate_source_code_input_against_pipeline_variables,
 )
 from sagemaker.inputs import TrainingInput, FileSystemInput
+from sagemaker.instance_group import InstanceGroup
 from sagemaker.job import _Job
 from sagemaker.jumpstart.utils import (
     add_jumpstart_tags,
@@ -149,7 +150,7 @@ def __init__(
         code_location: Optional[str] = None,
         entry_point: Optional[Union[str, PipelineVariable]] = None,
         dependencies: Optional[List[Union[str]]] = None,
-        instance_groups: Optional[Dict[str, Union[str, int]]] = None,
+        instance_groups: Optional[List[InstanceGroup]] = None,
         **kwargs,
     ):
         """Initialize an ``EstimatorBase`` instance.
@@ -1580,6 +1581,8 @@ def _get_instance_type(self):
 
         for instance_group in self.instance_groups:
             instance_type = instance_group.instance_type
+            if is_pipeline_variable(instance_type):
+                continue
             match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
 
             if match:
@@ -2179,7 +2182,7 @@ def __init__(
         code_location: Optional[str] = None,
         entry_point: Optional[Union[str, PipelineVariable]] = None,
         dependencies: Optional[List[str]] = None,
-        instance_groups: Optional[Dict[str, Union[str, int]]] = None,
+        instance_groups: Optional[List[InstanceGroup]] = None,
         **kwargs,
     ):
         """Initialize an ``Estimator`` instance.

src/sagemaker/fw_utils.py

@@ -871,6 +871,14 @@ def validate_distribution_instance(sagemaker_session, distribution, instance_typ
         # Strategy modelparallel is not enabled
         return
 
+    if is_pipeline_variable(instance_type):
+        logger.warning(
+            "instance_type is a pipeline variable, which is only interpreted in "
+            "pipeline execution time. As modelparallel only runs on GPU-enabled "
+            "instances, in execution time, the specified instance type has to support GPU."
+        )
+        return
+
     instance_desc = sagemaker_session.boto_session.client("ec2").describe_instance_types(
         InstanceTypes=[f"{instance_type}"]
     )