change: Implement test mechanism for Pipeline variables

annotations for processors + estimators / test mechanism change

annotations for processors + estimators / test mechanism change

remove debug print

reformatting

update TM and resolve all AIs and all untested subclasses

Add ppl var annotation to all composite object for training

Author: Dewen Qi

src/sagemaker/amazon/kmeans.py

@@ -27,8 +27,6 @@
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 from sagemaker.workflow.entities import PipelineVariable
 
-from sagemaker.workflow.entities import PipelineVariable
-
 
 class KMeans(AmazonAlgorithmEstimatorBase):
     """An unsupervised learning algorithm that attempts to find discrete groupings within data.
@@ -76,7 +74,7 @@ def __init__(
         half_life_time_size: Optional[int] = None,
         epochs: Optional[int] = None,
         center_factor: Optional[int] = None,
-        eval_metrics: Optional[List[Union[str, PipelineVariable]]] = None,
+        eval_metrics: Optional[List[str]] = None,
         **kwargs
     ):
         """A k-means clustering class :class:`~sagemaker.amazon.AmazonAlgorithmEstimatorBase`.

src/sagemaker/amazon/linear_learner.py

@@ -13,6 +13,7 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
+import logging
 from typing import Union, Optional
 
 from sagemaker import image_uris
@@ -28,6 +29,8 @@
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow import is_pipeline_variable
 
+logger = logging.getLogger(__name__)
+
 
 class LinearLearner(AmazonAlgorithmEstimatorBase):
     """A supervised learning algorithms used for solving classification or regression problems.
@@ -437,9 +440,14 @@ def _prepare_for_training(self, records, mini_batch_size=None, job_name=None):
         # mini_batch_size can't be greater than number of records or training job fails
         if not mini_batch_size:
             if is_pipeline_variable(self.instance_count):
-                raise ValueError(
-                    "instance_count can not be a pipeline variable when mini_batch_size is not given."
+                logger.warning(
+                    "mini_batch_size is not given in .fit() and instance_count is a "
+                    "pipeline variable (%s) which is only parsed in execution time. "
+                    "Thus setting mini_batch_size to 1, as it can't be greater than "
+                    "number of records per instance_count, otherwise the training job fails.",
+                    type(self.instance_count),
                 )
+                mini_batch_size = 1
             else:
                 mini_batch_size = min(
                     self.DEFAULT_MINI_BATCH_SIZE, max(1, int(num_records / self.instance_count))

src/sagemaker/debugger/debugger.py

@@ -24,7 +24,6 @@
 
 from abc import ABC
 
-from sagemaker.workflow.entities import PipelineVariable
 from typing import Union, Optional, List, Dict
 
 import attr
@@ -33,6 +32,7 @@
 
 from sagemaker import image_uris
 from sagemaker.utils import build_dict
+from sagemaker.workflow.entities import PipelineVariable
 
 framework_name = "debugger"
 DEBUGGER_FLAG = "USE_SMDEBUG"

src/sagemaker/estimator.py

@@ -54,7 +54,6 @@
     get_jumpstart_base_name_if_jumpstart_model,
     update_inference_tags_with_jumpstart_training_tags,
 )
-from sagemaker.debugger import RuleBase
 from sagemaker.local import LocalSession
 from sagemaker.model import (
     CONTAINER_LOG_LEVEL_PARAM_NAME,
@@ -2503,7 +2502,7 @@ class Framework(EstimatorBase):
     def __init__(
         self,
         entry_point: str,
-        source_dir: Optional[Union[str, PipelineVariable]]= None,
+        source_dir: Optional[Union[str, PipelineVariable]] = None,
         hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
         container_log_level: Union[int, PipelineVariable] = logging.INFO,
         code_location: Optional[str] = None,

src/sagemaker/huggingface/estimator.py

@@ -28,7 +28,7 @@
 from sagemaker.huggingface.model import HuggingFaceModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
-from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig  # TODO: add annotations
+from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig
 from sagemaker.workflow.entities import PipelineVariable
 
 logger = logging.getLogger("sagemaker")
@@ -42,11 +42,11 @@ class HuggingFace(Framework):
     def __init__(
         self,
         py_version: str,
-        entry_point: str,
-        transformers_version: Optional[int] = None,
-        tensorflow_version: Optional[int] = None,
-        pytorch_version: Optional[int] = None,
-        source_dir: Optional[str] = None,
+        entry_point: Union[str, PipelineVariable],
+        transformers_version: Optional[str] = None,
+        tensorflow_version: Optional[str] = None,
+        pytorch_version: Optional[str] = None,
+        source_dir: Optional[Union[str, PipelineVariable]] = None,
         hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
         image_uri: Optional[Union[str, PipelineVariable]] = None,
         distribution: Optional[Dict] = None,
@@ -203,14 +203,8 @@ def __init__(
                     f"Instead got {type(compiler_config)}"
                 )
                 raise ValueError(error_string)
-
-            compiler_config.validate(
-                image_uri=image_uri,
-                instance_type=instance_type,
-                distribution=distribution,
-            )
-
-        self.distribution = distribution or {}
+            if compiler_config:
+                compiler_config.validate(self)
         self.compiler_config = compiler_config
 
     def _validate_args(self, image_uri):

src/sagemaker/inputs.py

@@ -14,8 +14,11 @@
 from __future__ import absolute_import, print_function
 
 from typing import Union, Optional, List
+
 import attr
 
+from sagemaker.workflow.entities import PipelineVariable
+
 FILE_SYSTEM_TYPES = ["FSxLustre", "EFS"]
 FILE_SYSTEM_ACCESS_MODES = ["ro", "rw"]
 

src/sagemaker/processing.py

@@ -40,7 +40,6 @@
 from sagemaker.dataset_definition.inputs import S3Input, DatasetDefinition
 from sagemaker.apiutils._base_types import ApiObject
 from sagemaker.s3 import S3Uploader
-from sagemaker.workflow.entities import PipelineVariable
 
 
 logger = logging.getLogger(__name__)

src/sagemaker/sklearn/estimator.py

@@ -156,12 +156,6 @@ 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/tensorflow/training_compiler/config.py

@@ -13,10 +13,13 @@
 """Configuration for the SageMaker Training Compiler."""
 from __future__ import absolute_import
 import logging
+from typing import Union
+
 from packaging.specifiers import SpecifierSet
 from packaging.version import Version
 
 from sagemaker.training_compiler.config import TrainingCompilerConfig as BaseConfig
+from sagemaker.workflow.entities import PipelineVariable
 
 logger = logging.getLogger(__name__)
 
@@ -29,8 +32,8 @@ class TrainingCompilerConfig(BaseConfig):
 
     def __init__(
         self,
-        enabled=True,
-        debug=False,
+        enabled: Union[bool, PipelineVariable] = True,
+        debug: Union[bool, PipelineVariable] = False,
     ):
         """This class initializes a ``TrainingCompilerConfig`` instance.
 

src/sagemaker/utils.py

@@ -93,6 +93,7 @@ def unique_name_from_base(base, max_length=63):
 
 def base_name_from_image(image, default_base_name=None):
     """Extract the base name of the image to use as the 'algorithm name' for the job.
+
     Args:
         image (str): Image name.
         default_base_name (str): The default base name

src/sagemaker/workflow/__init__.py

@@ -34,6 +34,7 @@ def is_pipeline_variable(var: object) -> bool:
 
 def is_pipeline_parameter_string(var: object) -> bool:
     """Check if the variable is a pipeline parameter string
+
     Args:
         var (object): The variable to be verified.
     Returns:

src/sagemaker/workflow/step_collections.py

@@ -14,7 +14,6 @@
 from __future__ import absolute_import
 
 import warnings
-from sagemaker.deprecations import deprecated_class
 from typing import List, Union, Optional
 
 import attr

src/sagemaker/xgboost/processing.py

@@ -33,17 +33,17 @@ class XGBoostProcessor(FrameworkProcessor):
 
     def __init__(
         self,
-        framework_version: str,  # New arg
+        framework_version: str,
         role: str,
         instance_count: Union[int, PipelineVariable],
         instance_type: Union[str, PipelineVariable],
-        py_version: str = "py3",  # New kwarg
+        py_version: str = "py3",
         image_uri: Optional[Union[str, PipelineVariable]] = None,
         command: Optional[List[str]] = None,
         volume_size_in_gb: Union[int, PipelineVariable] = 30,
         volume_kms_key: Optional[Union[str, PipelineVariable]] = None,
         output_kms_key: Optional[Union[str, PipelineVariable]] = None,
-        code_location: Optional[str] = None,  # New arg
+        code_location: Optional[str] = None,
         max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None,
         base_job_name: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,