change: Add PipelineVariable annotation to framework estimators

Author: Dewen Qi

src/sagemaker/algorithm.py

@@ -13,15 +13,22 @@
 """Test docstring"""
 from __future__ import absolute_import
 
+from typing import Optional, Union, Dict, List
+
 import sagemaker
 import sagemaker.parameter
 from sagemaker import vpc_utils
 from sagemaker.deserializers import BytesDeserializer
 from sagemaker.deprecations import removed_kwargs
 from sagemaker.estimator import EstimatorBase
+from sagemaker.inputs import TrainingInput, FileSystemInput
 from sagemaker.serializers import IdentitySerializer
 from sagemaker.transformer import Transformer
 from sagemaker.predictor import Predictor
+from sagemaker.session import Session
+from sagemaker.workflow.entities import PipelineVariable
+
+from sagemaker.workflow import is_pipeline_variable
 
 
 class AlgorithmEstimator(EstimatorBase):
@@ -37,28 +44,28 @@ class AlgorithmEstimator(EstimatorBase):
 
     def __init__(
         self,
-        algorithm_arn,
-        role,
-        instance_count,
-        instance_type,
-        volume_size=30,
-        volume_kms_key=None,
-        max_run=24 * 60 * 60,
-        input_mode="File",
-        output_path=None,
-        output_kms_key=None,
-        base_job_name=None,
-        sagemaker_session=None,
-        hyperparameters=None,
-        tags=None,
-        subnets=None,
-        security_group_ids=None,
-        model_uri=None,
-        model_channel_name="model",
-        metric_definitions=None,
-        encrypt_inter_container_traffic=False,
-        use_spot_instances=False,
-        max_wait=None,
+        algorithm_arn: str,
+        role: str,
+        instance_count: Optional[Union[int, PipelineVariable]] = None,
+        instance_type: Optional[Union[str, PipelineVariable]] = None,
+        volume_size: Union[int, PipelineVariable] = 30,
+        volume_kms_key: Optional[Union[str, PipelineVariable]] = None,
+        max_run: Union[int, PipelineVariable] = 24 * 60 * 60,
+        input_mode: Union[str, PipelineVariable] = "File",
+        output_path: Optional[Union[str, PipelineVariable]] = None,
+        output_kms_key: Optional[Union[str, PipelineVariable]] = None,
+        base_job_name: Optional[str] = None,
+        sagemaker_session: Optional[Session] = None,
+        hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None,
+        subnets: Optional[List[Union[str, PipelineVariable]]] = None,
+        security_group_ids: Optional[List[Union[str, PipelineVariable]]] = None,
+        model_uri: Optional[str] = None,
+        model_channel_name: Union[str, PipelineVariable] = "model",
+        metric_definitions: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None,
+        encrypt_inter_container_traffic: Union[bool, PipelineVariable] = False,
+        use_spot_instances: Union[bool, PipelineVariable] = False,
+        max_wait: Union[int, PipelineVariable] = None,
         **kwargs  # pylint: disable=W0613
     ):
         """Initialize an ``AlgorithmEstimator`` instance.
@@ -186,22 +193,25 @@ def validate_train_spec(self):
         # Check that the input mode provided is compatible with the training input modes for the
         # algorithm.
         input_modes = self._algorithm_training_input_modes(train_spec["TrainingChannels"])
-        if self.input_mode not in input_modes:
+        if not is_pipeline_variable(self.input_mode) and self.input_mode not in input_modes:
             raise ValueError(
                 "Invalid input mode: %s. %s only supports: %s"
                 % (self.input_mode, algorithm_name, input_modes)
             )
 
         # Check that the training instance type is compatible with the algorithm.
         supported_instances = train_spec["SupportedTrainingInstanceTypes"]
-        if self.instance_type not in supported_instances:
+        if (
+            not is_pipeline_variable(self.instance_type)
+            and self.instance_type not in supported_instances
+        ):
             raise ValueError(
                 "Invalid instance_type: %s. %s supports the following instance types: %s"
                 % (self.instance_type, algorithm_name, supported_instances)
             )
 
         # Verify if distributed training is supported by the algorithm
-        if (
+        if not is_pipeline_variable(self.instance_count) and (
             self.instance_count > 1
             and "SupportsDistributedTraining" in train_spec
             and not train_spec["SupportsDistributedTraining"]
@@ -414,12 +424,18 @@ def _prepare_for_training(self, job_name=None):
 
         super(AlgorithmEstimator, self)._prepare_for_training(job_name)
 
-    def fit(self, inputs=None, wait=True, logs=True, job_name=None):
+    def fit(
+        self,
+        inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None,
+        wait: bool = True,
+        logs: bool = True,
+        job_name: Optional[str] = None,
+    ):
         """Placeholder docstring"""
         if inputs:
             self._validate_input_channels(inputs)
 
-        super(AlgorithmEstimator, self).fit(inputs, wait, logs, job_name)
+        return super(AlgorithmEstimator, self).fit(inputs, wait, logs, job_name)
 
     def _validate_input_channels(self, channels):
         """Placeholder docstring"""

src/sagemaker/chainer/estimator.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 import logging
-from typing import Union, Optional
+from typing import Union, Optional, Dict
 
 from sagemaker.estimator import Framework, EstimatorBase
 from sagemaker.fw_utils import (
@@ -34,26 +34,26 @@
 class Chainer(Framework):
     """Handle end-to-end training and deployment of custom Chainer code."""
 
-    _framework_name = "chainer"
+    _framework_name: str = "chainer"
 
     # Hyperparameters
-    _use_mpi = "sagemaker_use_mpi"
-    _num_processes = "sagemaker_num_processes"
-    _process_slots_per_host = "sagemaker_process_slots_per_host"
-    _additional_mpi_options = "sagemaker_additional_mpi_options"
+    _use_mpi: str = "sagemaker_use_mpi"
+    _num_processes: str = "sagemaker_num_processes"
+    _process_slots_per_host: str = "sagemaker_process_slots_per_host"
+    _additional_mpi_options: str = "sagemaker_additional_mpi_options"
 
     def __init__(
         self,
         entry_point: Union[str, PipelineVariable],
-        use_mpi=None,
-        num_processes=None,
-        process_slots_per_host=None,
-        additional_mpi_options=None,
+        use_mpi: Optional[Union[bool, PipelineVariable]] = None,
+        num_processes: Optional[Union[int, PipelineVariable]] = None,
+        process_slots_per_host: Optional[Union[int, PipelineVariable]] = None,
+        additional_mpi_options: Optional[Union[str, PipelineVariable]] = None,
         source_dir: Optional[Union[str, PipelineVariable]] = None,
-        hyperparameters=None,
-        framework_version=None,
-        py_version=None,
-        image_uri=None,
+        hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        framework_version: Optional[str] = None,
+        py_version: Optional[str] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
         **kwargs
     ):
         """This ``Estimator`` executes an Chainer script in a managed execution environment.

src/sagemaker/pytorch/estimator.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 import logging
-from typing import Union, Optional
+from typing import Union, Optional, Dict
 
 from packaging.version import Version
 
@@ -44,12 +44,12 @@ class PyTorch(Framework):
     def __init__(
         self,
         entry_point: Union[str, PipelineVariable],
-        framework_version=None,
-        py_version=None,
+        framework_version: Optional[str] = None,
+        py_version: Optional[str] = None,
         source_dir: Optional[Union[str, PipelineVariable]] = None,
-        hyperparameters=None,
-        image_uri=None,
-        distribution=None,
+        hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        distribution: Optional[Dict] = None,
         **kwargs
     ):
         """This ``Estimator`` executes a PyTorch script in a managed PyTorch execution environment.

src/sagemaker/rl/estimator.py

@@ -16,7 +16,7 @@
 import enum
 import logging
 import re
-from typing import Union, Optional
+from typing import Union, Optional, List, Dict
 
 from sagemaker import image_uris, fw_utils
 from sagemaker.estimator import Framework, EstimatorBase
@@ -77,13 +77,13 @@ class RLEstimator(Framework):
     def __init__(
         self,
         entry_point: Union[str, PipelineVariable],
-        toolkit=None,
-        toolkit_version=None,
-        framework=None,
+        toolkit: Optional[RLToolkit] = None,
+        toolkit_version: Optional[str] = None,
+        framework: Optional[Framework] = None,
         source_dir: Optional[Union[str, PipelineVariable]] = None,
-        hyperparameters=None,
-        image_uri=None,
-        metric_definitions=None,
+        hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        metric_definitions: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None,
         **kwargs
     ):
         """Creates an RLEstimator for managed Reinforcement Learning (RL).

src/sagemaker/sklearn/estimator.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 import logging
-from typing import Union, Optional
+from typing import Union, Optional, Dict
 
 from sagemaker import image_uris
 from sagemaker.deprecations import renamed_kwargs
@@ -28,6 +28,7 @@
 from sagemaker.sklearn.model import SKLearnModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 from sagemaker.workflow.entities import PipelineVariable
+from sagemaker.workflow import is_pipeline_variable
 
 logger = logging.getLogger("sagemaker")
 
@@ -40,12 +41,12 @@ class SKLearn(Framework):
     def __init__(
         self,
         entry_point: Union[str, PipelineVariable],
-        framework_version=None,
-        py_version="py3",
+        framework_version: Optional[str] = None,
+        py_version: str = "py3",
         source_dir: Optional[Union[str, PipelineVariable]] = None,
-        hyperparameters=None,
-        image_uri=None,
-        image_uri_region=None,
+        hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        image_uri_region: Optional[str] = None,
         **kwargs
     ):
         """Creates a SKLearn Estimator for Scikit-learn environment.
@@ -128,11 +129,17 @@ def __init__(
         self.framework_version = framework_version
         self.py_version = py_version
 
+        if instance_type:
+            if is_pipeline_variable(instance_type):
+                raise ValueError("instance_count argument cannot be a pipeline variable")
         # SciKit-Learn does not support distributed training or training on GPU instance types.
         # Fail fast.
         _validate_not_gpu_instance_type(instance_type)
 
         if instance_count:
+            if is_pipeline_variable(instance_count):
+                raise ValueError("instance_count argument cannot be a pipeline variable")
+
             if instance_count != 1:
                 raise AttributeError(
                     "Scikit-Learn does not support distributed training. Please remove the "
@@ -148,6 +155,13 @@ def __init__(
         )
 
         if image_uri is None:
+
+            if is_pipeline_variable(instance_type):
+                raise ValueError(
+                    "instance_type argument cannot be a pipeline variable "
+                    "when image_uri is not given."
+                )
+
             self.image_uri = image_uris.retrieve(
                 SKLearn._framework_name,
                 image_uri_region or self.sagemaker_session.boto_region_name,

src/sagemaker/training_compiler/config.py

@@ -14,6 +14,8 @@
 from __future__ import absolute_import
 import logging
 
+from sagemaker.workflow import is_pipeline_variable
+
 logger = logging.getLogger(__name__)
 
 
@@ -132,7 +134,12 @@ def validate(
             ValueError: Raised if the requested configuration is not compatible
                         with SageMaker Training Compiler.
         """
-        if estimator.instance_type:
+        if is_pipeline_variable(estimator.instance_type):
+            logger.warning(
+                "Estimator instance_type is a PipelineVariable: %s, "
+                "which has to be one of the [p3, g4dn, p4d, g5] in execution time."
+            )
+        elif estimator.instance_type:
             if "local" not in estimator.instance_type:
                 requested_instance_class = estimator.instance_type.split(".")[
                     1

src/sagemaker/workflow/pipeline.py

@@ -486,7 +486,10 @@ def _generate_step_map(
     """Helper method to create a mapping from Step/Step Collection name to itself."""
     for step in steps:
         if step.name in step_map:
-            raise ValueError("Pipeline steps cannot have duplicate names.")
+            raise ValueError(
+                "Pipeline steps cannot have duplicate names. In addition, steps added to "
+                "the ConditionStep cannot be added to the Pipeline steps list."
+            )
         step_map[step.name] = step
         if isinstance(step, ConditionStep):
             _generate_step_map(step.if_steps + step.else_steps, step_map)

src/sagemaker/xgboost/estimator.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 import logging
-from typing import Union, Optional
+from typing import Union, Optional, Dict
 
 from sagemaker import image_uris
 from sagemaker.deprecations import renamed_kwargs
@@ -45,12 +45,12 @@ class XGBoost(Framework):
     def __init__(
         self,
         entry_point: Union[str, PipelineVariable],
-        framework_version,
+        framework_version: str,
         source_dir: Optional[Union[str, PipelineVariable]] = None,
-        hyperparameters=None,
-        py_version="py3",
-        image_uri=None,
-        image_uri_region=None,
+        hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        py_version: str = "py3",
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        image_uri_region: Optional[str] = None,
         **kwargs
     ):
         """An estimator that executes an XGBoost-based SageMaker Training Job.