change: Add Pipeline annotation in model base class and tensorflow estimator

Model annotate update

Author: Dewen Qi

src/sagemaker/drift_check_baselines.py

@@ -13,21 +13,25 @@
 """This file contains code related to drift check baselines"""
 from __future__ import absolute_import
 
+from typing import Optional
+
+from sagemaker.model_metrics import MetricsSource, FileSource
+
 
 class DriftCheckBaselines(object):
     """Accepts drift check baselines parameters for conversion to request dict."""
 
     def __init__(
         self,
-        model_statistics=None,
-        model_constraints=None,
-        model_data_statistics=None,
-        model_data_constraints=None,
-        bias_config_file=None,
-        bias_pre_training_constraints=None,
-        bias_post_training_constraints=None,
-        explainability_constraints=None,
-        explainability_config_file=None,
+        model_statistics: Optional[MetricsSource] = None,
+        model_constraints: Optional[MetricsSource] = None,
+        model_data_statistics: Optional[MetricsSource] = None,
+        model_data_constraints: Optional[MetricsSource] = None,
+        bias_config_file: Optional[FileSource] = None,
+        bias_pre_training_constraints: Optional[MetricsSource] = None,
+        bias_post_training_constraints: Optional[MetricsSource] = None,
+        explainability_constraints: Optional[MetricsSource] = None,
+        explainability_config_file: Optional[FileSource] = None,
     ):
         """Initialize a ``DriftCheckBaselines`` instance and turn parameters into dict.
 

src/sagemaker/estimator.py

@@ -18,7 +18,7 @@
 import os
 import uuid
 from abc import ABCMeta, abstractmethod
-from typing import Any, Dict
+from typing import Any, Dict, Optional, Union, List
 
 from six import string_types, with_metaclass
 from six.moves.urllib.parse import urlparse
@@ -36,6 +36,7 @@
     TensorBoardOutputConfig,
     get_default_profiler_rule,
     get_rule_container_image_uri,
+    RuleBase,
 )
 from sagemaker.deprecations import removed_function, removed_kwargs, renamed_kwargs
 from sagemaker.fw_utils import (
@@ -75,6 +76,7 @@
     name_from_base,
 )
 from sagemaker.workflow import is_pipeline_variable
+from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.pipeline_context import (
     PipelineSession,
     runnable_by_pipeline,
@@ -105,44 +107,44 @@ class EstimatorBase(with_metaclass(ABCMeta, object)):  # pylint: disable=too-man
 
     def __init__(
         self,
-        role,
-        instance_count=None,
-        instance_type=None,
-        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,
-        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,
-        checkpoint_s3_uri=None,
-        checkpoint_local_path=None,
-        rules=None,
-        debugger_hook_config=None,
-        tensorboard_output_config=None,
-        enable_sagemaker_metrics=None,
-        enable_network_isolation=False,
-        profiler_config=None,
-        disable_profiler=False,
-        environment=None,
-        max_retry_attempts=None,
-        source_dir=None,
-        git_config=None,
-        hyperparameters=None,
-        container_log_level=logging.INFO,
-        code_location=None,
-        entry_point=None,
-        dependencies=None,
+        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,
+        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: Optional[Union[int, PipelineVariable]] = None,
+        checkpoint_s3_uri: Optional[Union[str, PipelineVariable]] = None,
+        checkpoint_local_path: Optional[Union[str, PipelineVariable]] = None,
+        rules: Optional[List[RuleBase]] = None,
+        debugger_hook_config: Optional[Union[bool, DebuggerHookConfig]] = None,
+        tensorboard_output_config: Optional[TensorBoardOutputConfig] = None,
+        enable_sagemaker_metrics: Optional[Union[bool, PipelineVariable]] = None,
+        enable_network_isolation: Union[bool, PipelineVariable] = False,
+        profiler_config: Optional[ProfilerConfig] = None,
+        disable_profiler: bool = False,
+        environment: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        max_retry_attempts: Optional[Union[int, PipelineVariable]] = None,
+        source_dir: Optional[str] = None,
+        git_config: Optional[Dict[str, str]] = None,
+        hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        container_log_level: Union[int, PipelineVariable] = logging.INFO,
+        code_location: Optional[str] = None,
+        entry_point: Optional[str] = None,
+        dependencies: Optional[List[Union[str]]] = None,
         **kwargs,
     ):
         """Initialize an ``EstimatorBase`` instance.
@@ -2730,15 +2732,25 @@ def _validate_and_set_debugger_configs(self):
         # Disable debugger if checkpointing is enabled by the customer
         if self.checkpoint_s3_uri and self.checkpoint_local_path and self.debugger_hook_config:
             if self._framework_name in {"mxnet", "pytorch", "tensorflow"}:
-                if self.instance_count > 1 or (
+                disable_debugger_hook_config = False
+                if not is_pipeline_variable(self.instance_count) and self.instance_count > 1:
+                    disable_debugger_hook_config = True
+                if (
                     hasattr(self, "distribution")
                     and self.distribution is not None  # pylint: disable=no-member
                 ):
+                    disable_debugger_hook_config = True
+                if disable_debugger_hook_config:
                     logger.info(
                         "SMDebug Does Not Currently Support \
                         Distributed Training Jobs With Checkpointing Enabled"
                     )
                     self.debugger_hook_config = False
+                elif is_pipeline_variable(self.instance_count):
+                    logger.warning(
+                        "SMDebug does not currently support when parameterized "
+                        "instance_count with value > 1 in pipeline execution."
+                    )
 
         if self.debugger_hook_config is False:
             if self.environment is None:

src/sagemaker/metadata_properties.py

@@ -13,16 +13,20 @@
 """This file contains code related to metadata properties."""
 from __future__ import absolute_import
 
+from typing import Optional, Union
+
+from sagemaker.workflow.entities import PipelineVariable
+
 
 class MetadataProperties(object):
     """Accepts metadata properties parameters for conversion to request dict."""
 
     def __init__(
         self,
-        commit_id=None,
-        repository=None,
-        generated_by=None,
-        project_id=None,
+        commit_id: Optional[Union[str, PipelineVariable]] = None,
+        repository: Optional[Union[str, PipelineVariable]] = None,
+        generated_by: Optional[Union[str, PipelineVariable]] = None,
+        project_id: Optional[Union[str, PipelineVariable]] = None,
     ):
         """Initialize a ``MetadataProperties`` instance and turn parameters into dict.
 

src/sagemaker/model.py

@@ -19,7 +19,7 @@
 import os
 import re
 import copy
-from typing import List, Dict
+from typing import List, Dict, Optional, Union
 
 import sagemaker
 from sagemaker import (
@@ -30,15 +30,20 @@
     utils,
     git_utils,
 )
+from sagemaker.session import Session
+from sagemaker.model_metrics import ModelMetrics
 from sagemaker.deprecations import removed_kwargs
+from sagemaker.drift_check_baselines import DriftCheckBaselines
+from sagemaker.metadata_properties import MetadataProperties
 from sagemaker.predictor import PredictorBase
 from sagemaker.serverless import ServerlessInferenceConfig
 from sagemaker.transformer import Transformer
 from sagemaker.jumpstart.utils import add_jumpstart_tags, get_jumpstart_base_name_if_jumpstart_model
-from sagemaker.utils import unique_name_from_base
+from sagemaker.utils import unique_name_from_base, to_string
 from sagemaker.async_inference import AsyncInferenceConfig
 from sagemaker.predictor_async import AsyncPredictor
 from sagemaker.workflow import is_pipeline_variable
+from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.pipeline_context import runnable_by_pipeline, PipelineSession
 
 LOGGER = logging.getLogger("sagemaker")
@@ -78,23 +83,23 @@ class Model(ModelBase):
 
     def __init__(
         self,
-        image_uri,
-        model_data=None,
-        role=None,
-        predictor_cls=None,
-        env=None,
-        name=None,
-        vpc_config=None,
-        sagemaker_session=None,
-        enable_network_isolation=False,
-        model_kms_key=None,
-        image_config=None,
-        source_dir=None,
-        code_location=None,
-        entry_point=None,
-        container_log_level=logging.INFO,
-        dependencies=None,
-        git_config=None,
+        image_uri: Union[str, PipelineVariable],
+        model_data: Optional[Union[str, PipelineVariable]] = None,
+        role: Optional[str] = None,
+        predictor_cls: Optional[callable] = None,
+        env: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        name: Optional[str] = None,
+        vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = None,
+        sagemaker_session: Optional[Session] = None,
+        enable_network_isolation: Union[bool, PipelineVariable] = False,
+        model_kms_key: Optional[str] = None,
+        image_config: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        source_dir: Optional[str] = None,
+        code_location: Optional[str] = None,
+        entry_point: Optional[str] = None,
+        container_log_level: Union[int, PipelineVariable] = logging.INFO,
+        dependencies: Optional[List[str]] = None,
+        git_config: Optional[Dict[str, str]] = None,
     ):
         """Initialize an SageMaker ``Model``.
 
@@ -294,22 +299,22 @@ def __init__(
     @runnable_by_pipeline
     def register(
         self,
-        content_types,
-        response_types,
-        inference_instances=None,
-        transform_instances=None,
-        model_package_name=None,
-        model_package_group_name=None,
-        image_uri=None,
-        model_metrics=None,
-        metadata_properties=None,
-        marketplace_cert=False,
-        approval_status=None,
-        description=None,
-        drift_check_baselines=None,
-        customer_metadata_properties=None,
-        validation_specification=None,
-        domain=None,
+        content_types: List[Union[str, PipelineVariable]],
+        response_types: List[Union[str, PipelineVariable]],
+        inference_instances: Optional[List[Union[str, PipelineVariable]]] = None,
+        transform_instances: Optional[List[Union[str, PipelineVariable]]] = None,
+        model_package_name: Optional[Union[str, PipelineVariable]] = None,
+        model_package_group_name: Optional[Union[str, PipelineVariable]] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        model_metrics: Optional[ModelMetrics] = None,
+        metadata_properties: Optional[MetadataProperties] = None,
+        marketplace_cert: bool = False,
+        approval_status: Optional[Union[str, PipelineVariable]] = None,
+        description: Optional[str] = None,
+        drift_check_baselines: Optional[DriftCheckBaselines] = None,
+        customer_metadata_properties: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        validation_specification: Optional[Union[str, PipelineVariable]] = None,
+        domain: Optional[Union[str, PipelineVariable]] = None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -385,10 +390,10 @@ def register(
     @runnable_by_pipeline
     def create(
         self,
-        instance_type: str = None,
-        accelerator_type: str = None,
-        serverless_inference_config: ServerlessInferenceConfig = None,
-        tags: List[Dict[str, str]] = None,
+        instance_type: Optional[str] = None,
+        accelerator_type: Optional[str] = None,
+        serverless_inference_config: Optional[ServerlessInferenceConfig] = None,
+        tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None,
     ):
         """Create a SageMaker Model Entity
 
@@ -570,7 +575,7 @@ def _script_mode_env_vars(self):
         return {
             SCRIPT_PARAM_NAME.upper(): script_name or str(),
             DIR_PARAM_NAME.upper(): dir_name or str(),
-            CONTAINER_LOG_LEVEL_PARAM_NAME.upper(): str(self.container_log_level),
+            CONTAINER_LOG_LEVEL_PARAM_NAME.upper(): to_string(self.container_log_level),
             SAGEMAKER_REGION_PARAM_NAME.upper(): self.sagemaker_session.boto_region_name,
         }
 
@@ -1239,19 +1244,19 @@ class FrameworkModel(Model):
 
     def __init__(
         self,
-        model_data,
-        image_uri,
-        role,
-        entry_point,
-        source_dir=None,
-        predictor_cls=None,
-        env=None,
-        name=None,
-        container_log_level=logging.INFO,
-        code_location=None,
-        sagemaker_session=None,
-        dependencies=None,
-        git_config=None,
+        model_data: Union[str, PipelineVariable],
+        image_uri: Union[str, PipelineVariable],
+        role: str,
+        entry_point: str,
+        source_dir: Optional[str] = None,
+        predictor_cls: Optional[callable] = None,
+        env: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        name: Optional[str] = None,
+        container_log_level: Union[int, PipelineVariable] = logging.INFO,
+        code_location: Optional[str] = None,
+        sagemaker_session: Optional[Session] = None,
+        dependencies: Optional[List[str]] = None,
+        git_config: Optional[Dict[str, str]] = None,
         **kwargs,
     ):
         """Initialize a ``FrameworkModel``.