change: Add PipelineVariable annotation to Amazon estimators

Author: Dewen Qi

src/sagemaker/amazon/amazon_estimator.py

@@ -16,7 +16,7 @@
 import json
 import logging
 import tempfile
-from typing import Union
+from typing import Union, Optional, Dict
 
 from six.moves.urllib.parse import urlparse
 
@@ -30,6 +30,7 @@
 from sagemaker.utils import sagemaker_timestamp
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.pipeline_context import runnable_by_pipeline
+from sagemaker.workflow import is_pipeline_variable
 
 logger = logging.getLogger(__name__)
 
@@ -42,16 +43,16 @@ class AmazonAlgorithmEstimatorBase(EstimatorBase):
 
     feature_dim = hp("feature_dim", validation.gt(0), data_type=int)
     mini_batch_size = hp("mini_batch_size", validation.gt(0), data_type=int)
-    repo_name = None
-    repo_version = None
+    repo_name: Optional[str] = None
+    repo_version: Optional[str] = None
 
     def __init__(
         self,
-        role,
-        instance_count=None,
-        instance_type=None,
-        data_location=None,
-        enable_network_isolation=False,
+        role: str,
+        instance_count: Optional[Union[int]] = None,
+        instance_type: Optional[Union[str, PipelineVariable]] = None,
+        data_location: Optional[str] = None,
+        enable_network_isolation: Union[bool, PipelineVariable] = False,
         **kwargs
     ):
         """Initialize an AmazonAlgorithmEstimatorBase.
@@ -115,6 +116,11 @@ def data_location(self):
     @data_location.setter
     def data_location(self, data_location):
         """Placeholder docstring"""
+        if is_pipeline_variable(data_location):
+            raise ValueError(
+                "data_location argument has to be an integer " + "rather than a pipeline variable"
+            )
+
         if not data_location.startswith("s3://"):
             raise ValueError(
                 'Expecting an S3 URL beginning with "s3://". Got "{}"'.format(data_location)
@@ -198,12 +204,12 @@ def _prepare_for_training(self, records, mini_batch_size=None, job_name=None):
     @runnable_by_pipeline
     def fit(
         self,
-        records,
-        mini_batch_size=None,
-        wait=True,
-        logs=True,
-        job_name=None,
-        experiment_config=None,
+        records: "RecordSet",
+        mini_batch_size: Optional[int] = None,
+        wait: bool = True,
+        logs: bool = True,
+        job_name: Optional[str] = None,
+        experiment_config: Optional[Dict[str, str]] = None,
     ):
         """Fit this Estimator on serialized Record objects, stored in S3.
 

src/sagemaker/amazon/factorization_machines.py

@@ -37,83 +37,83 @@ class FactorizationMachines(AmazonAlgorithmEstimatorBase):
     sparse datasets economically.
     """
 
-    repo_name = "factorization-machines"
-    repo_version = 1
+    repo_name: str = "factorization-machines"
+    repo_version: int = 1
 
-    num_factors = hp("num_factors", gt(0), "An integer greater than zero", int)
-    predictor_type = hp(
+    num_factors: hp = hp("num_factors", gt(0), "An integer greater than zero", int)
+    predictor_type: hp = hp(
         "predictor_type",
         isin("binary_classifier", "regressor"),
         'Value "binary_classifier" or "regressor"',
         str,
     )
-    epochs = hp("epochs", gt(0), "An integer greater than 0", int)
-    clip_gradient = hp("clip_gradient", (), "A float value", float)
-    eps = hp("eps", (), "A float value", float)
-    rescale_grad = hp("rescale_grad", (), "A float value", float)
-    bias_lr = hp("bias_lr", ge(0), "A non-negative float", float)
-    linear_lr = hp("linear_lr", ge(0), "A non-negative float", float)
-    factors_lr = hp("factors_lr", ge(0), "A non-negative float", float)
-    bias_wd = hp("bias_wd", ge(0), "A non-negative float", float)
-    linear_wd = hp("linear_wd", ge(0), "A non-negative float", float)
-    factors_wd = hp("factors_wd", ge(0), "A non-negative float", float)
-    bias_init_method = hp(
+    epochs: hp = hp("epochs", gt(0), "An integer greater than 0", int)
+    clip_gradient: hp = hp("clip_gradient", (), "A float value", float)
+    eps: hp = hp("eps", (), "A float value", float)
+    rescale_grad: hp = hp("rescale_grad", (), "A float value", float)
+    bias_lr: hp = hp("bias_lr", ge(0), "A non-negative float", float)
+    linear_lr: hp = hp("linear_lr", ge(0), "A non-negative float", float)
+    factors_lr: hp = hp("factors_lr", ge(0), "A non-negative float", float)
+    bias_wd: hp = hp("bias_wd", ge(0), "A non-negative float", float)
+    linear_wd: hp = hp("linear_wd", ge(0), "A non-negative float", float)
+    factors_wd: hp = hp("factors_wd", ge(0), "A non-negative float", float)
+    bias_init_method: hp = hp(
         "bias_init_method",
         isin("normal", "uniform", "constant"),
         'Value "normal", "uniform" or "constant"',
         str,
     )
-    bias_init_scale = hp("bias_init_scale", ge(0), "A non-negative float", float)
-    bias_init_sigma = hp("bias_init_sigma", ge(0), "A non-negative float", float)
-    bias_init_value = hp("bias_init_value", (), "A float value", float)
-    linear_init_method = hp(
+    bias_init_scale: hp = hp("bias_init_scale", ge(0), "A non-negative float", float)
+    bias_init_sigma: hp = hp("bias_init_sigma", ge(0), "A non-negative float", float)
+    bias_init_value: hp = hp("bias_init_value", (), "A float value", float)
+    linear_init_method: hp = hp(
         "linear_init_method",
         isin("normal", "uniform", "constant"),
         'Value "normal", "uniform" or "constant"',
         str,
     )
-    linear_init_scale = hp("linear_init_scale", ge(0), "A non-negative float", float)
-    linear_init_sigma = hp("linear_init_sigma", ge(0), "A non-negative float", float)
-    linear_init_value = hp("linear_init_value", (), "A float value", float)
-    factors_init_method = hp(
+    linear_init_scale: hp = hp("linear_init_scale", ge(0), "A non-negative float", float)
+    linear_init_sigma: hp = hp("linear_init_sigma", ge(0), "A non-negative float", float)
+    linear_init_value: hp = hp("linear_init_value", (), "A float value", float)
+    factors_init_method: hp = hp(
         "factors_init_method",
         isin("normal", "uniform", "constant"),
         'Value "normal", "uniform" or "constant"',
         str,
     )
-    factors_init_scale = hp("factors_init_scale", ge(0), "A non-negative float", float)
-    factors_init_sigma = hp("factors_init_sigma", ge(0), "A non-negative float", float)
-    factors_init_value = hp("factors_init_value", (), "A float value", float)
+    factors_init_scale: hp = hp("factors_init_scale", ge(0), "A non-negative float", float)
+    factors_init_sigma: hp = hp("factors_init_sigma", ge(0), "A non-negative float", float)
+    factors_init_value: hp = hp("factors_init_value", (), "A float value", float)
 
     def __init__(
         self,
-        role,
-        instance_count=None,
-        instance_type=None,
-        num_factors=None,
-        predictor_type=None,
-        epochs=None,
-        clip_gradient=None,
-        eps=None,
-        rescale_grad=None,
-        bias_lr=None,
-        linear_lr=None,
-        factors_lr=None,
-        bias_wd=None,
-        linear_wd=None,
-        factors_wd=None,
-        bias_init_method=None,
-        bias_init_scale=None,
-        bias_init_sigma=None,
-        bias_init_value=None,
-        linear_init_method=None,
-        linear_init_scale=None,
-        linear_init_sigma=None,
-        linear_init_value=None,
-        factors_init_method=None,
-        factors_init_scale=None,
-        factors_init_sigma=None,
-        factors_init_value=None,
+        role: str,
+        instance_count: Optional[Union[int, PipelineVariable]] = None,
+        instance_type: Optional[Union[str, PipelineVariable]] = None,
+        num_factors: Optional[int] = None,
+        predictor_type: Optional[str] = None,
+        epochs: Optional[int] = None,
+        clip_gradient: Optional[float] = None,
+        eps: Optional[float] = None,
+        rescale_grad: Optional[float] = None,
+        bias_lr: Optional[float] = None,
+        linear_lr: Optional[float] = None,
+        factors_lr: Optional[float] = None,
+        bias_wd: Optional[float] = None,
+        linear_wd: Optional[float] = None,
+        factors_wd: Optional[float] = None,
+        bias_init_method: Optional[str] = None,
+        bias_init_scale: Optional[float] = None,
+        bias_init_sigma: Optional[float] = None,
+        bias_init_value: Optional[float] = None,
+        linear_init_method: Optional[str] = None,
+        linear_init_scale: Optional[float] = None,
+        linear_init_sigma: Optional[float] = None,
+        linear_init_value: Optional[float] = None,
+        factors_init_method: Optional[str] = None,
+        factors_init_scale: Optional[float] = None,
+        factors_init_sigma: Optional[float] = None,
+        factors_init_value: Optional[float] = None,
         **kwargs
     ):
         """Factorization Machines is :class:`Estimator` for general-purpose supervised learning.

src/sagemaker/amazon/ipinsights.py

@@ -36,45 +36,45 @@ class IPInsights(AmazonAlgorithmEstimatorBase):
     as user IDs or account numbers.
     """
 
-    repo_name = "ipinsights"
-    repo_version = 1
-    MINI_BATCH_SIZE = 10000
+    repo_name: str = "ipinsights"
+    repo_version: int = 1
+    MINI_BATCH_SIZE: int = 10000
 
-    num_entity_vectors = hp(
+    num_entity_vectors: hp = hp(
         "num_entity_vectors", (ge(1), le(250000000)), "An integer in [1, 250000000]", int
     )
-    vector_dim = hp("vector_dim", (ge(4), le(4096)), "An integer in [4, 4096]", int)
+    vector_dim: hp = hp("vector_dim", (ge(4), le(4096)), "An integer in [4, 4096]", int)
 
-    batch_metrics_publish_interval = hp(
+    batch_metrics_publish_interval: hp = hp(
         "batch_metrics_publish_interval", (ge(1)), "An integer greater than 0", int
     )
-    epochs = hp("epochs", (ge(1)), "An integer greater than 0", int)
-    learning_rate = hp("learning_rate", (ge(1e-6), le(10.0)), "A float in [1e-6, 10.0]", float)
-    num_ip_encoder_layers = hp(
+    epochs: hp = hp("epochs", (ge(1)), "An integer greater than 0", int)
+    learning_rate: hp = hp("learning_rate", (ge(1e-6), le(10.0)), "A float in [1e-6, 10.0]", float)
+    num_ip_encoder_layers: hp = hp(
         "num_ip_encoder_layers", (ge(0), le(100)), "An integer in [0, 100]", int
     )
-    random_negative_sampling_rate = hp(
+    random_negative_sampling_rate: hp = hp(
         "random_negative_sampling_rate", (ge(0), le(500)), "An integer in [0, 500]", int
     )
-    shuffled_negative_sampling_rate = hp(
+    shuffled_negative_sampling_rate: hp = hp(
         "shuffled_negative_sampling_rate", (ge(0), le(500)), "An integer in [0, 500]", int
     )
-    weight_decay = hp("weight_decay", (ge(0.0), le(10.0)), "A float in [0.0, 10.0]", float)
+    weight_decay: hp = hp("weight_decay", (ge(0.0), le(10.0)), "A float in [0.0, 10.0]", float)
 
     def __init__(
         self,
-        role,
-        instance_count=None,
-        instance_type=None,
-        num_entity_vectors=None,
-        vector_dim=None,
-        batch_metrics_publish_interval=None,
-        epochs=None,
-        learning_rate=None,
-        num_ip_encoder_layers=None,
-        random_negative_sampling_rate=None,
-        shuffled_negative_sampling_rate=None,
-        weight_decay=None,
+        role: str,
+        instance_count: Optional[Union[int, PipelineVariable]] = None,
+        instance_type: Optional[Union[str, PipelineVariable]] = None,
+        num_entity_vectors: Optional[int] = None,
+        vector_dim: Optional[int] = None,
+        batch_metrics_publish_interval: Optional[int] = None,
+        epochs: Optional[int] = None,
+        learning_rate: Optional[float] = None,
+        num_ip_encoder_layers: Optional[int] = None,
+        random_negative_sampling_rate: Optional[int] = None,
+        shuffled_negative_sampling_rate: Optional[int] = None,
+        weight_decay: Optional[float] = None,
         **kwargs
     ):
         """This estimator is for IP Insights.

src/sagemaker/amazon/kmeans.py

@@ -13,7 +13,7 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-from typing import Union, Optional
+from typing import Union, Optional, List
 
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
@@ -36,43 +36,45 @@ class KMeans(AmazonAlgorithmEstimatorBase):
     the algorithm to use to determine similarity.
     """
 
-    repo_name = "kmeans"
-    repo_version = 1
+    repo_name: str = "kmeans"
+    repo_version: str = 1
 
-    k = hp("k", gt(1), "An integer greater-than 1", int)
-    init_method = hp("init_method", isin("random", "kmeans++"), 'One of "random", "kmeans++"', str)
-    max_iterations = hp("local_lloyd_max_iter", gt(0), "An integer greater-than 0", int)
-    tol = hp("local_lloyd_tol", (ge(0), le(1)), "An float in [0, 1]", float)
-    num_trials = hp("local_lloyd_num_trials", gt(0), "An integer greater-than 0", int)
-    local_init_method = hp(
+    k: hp = hp("k", gt(1), "An integer greater-than 1", int)
+    init_method: hp = hp(
+        "init_method", isin("random", "kmeans++"), 'One of "random", "kmeans++"', str
+    )
+    max_iterations: hp = hp("local_lloyd_max_iter", gt(0), "An integer greater-than 0", int)
+    tol: hp = hp("local_lloyd_tol", (ge(0), le(1)), "An float in [0, 1]", float)
+    num_trials: hp = hp("local_lloyd_num_trials", gt(0), "An integer greater-than 0", int)
+    local_init_method: hp = hp(
         "local_lloyd_init_method", isin("random", "kmeans++"), 'One of "random", "kmeans++"', str
     )
-    half_life_time_size = hp(
+    half_life_time_size: hp = hp(
         "half_life_time_size", ge(0), "An integer greater-than-or-equal-to 0", int
     )
-    epochs = hp("epochs", gt(0), "An integer greater-than 0", int)
-    center_factor = hp("extra_center_factor", gt(0), "An integer greater-than 0", int)
-    eval_metrics = hp(
+    epochs: hp = hp("epochs", gt(0), "An integer greater-than 0", int)
+    center_factor: hp = hp("extra_center_factor", gt(0), "An integer greater-than 0", int)
+    eval_metrics: hp = hp(
         name="eval_metrics",
         validation_message='A comma separated list of "msd" or "ssd"',
         data_type=list,
     )
 
     def __init__(
         self,
-        role,
-        instance_count=None,
-        instance_type=None,
-        k=None,
-        init_method=None,
-        max_iterations=None,
-        tol=None,
-        num_trials=None,
-        local_init_method=None,
-        half_life_time_size=None,
-        epochs=None,
-        center_factor=None,
-        eval_metrics=None,
+        role: str,
+        instance_count: Optional[Union[int, PipelineVariable]] = None,
+        instance_type: Optional[Union[str, PipelineVariable]] = None,
+        k: Optional[int] = None,
+        init_method: Optional[str] = None,
+        max_iterations: Optional[int] = None,
+        tol: Optional[float] = None,
+        num_trials: Optional[int] = None,
+        local_init_method: Optional[str] = None,
+        half_life_time_size: Optional[int] = None,
+        epochs: Optional[int] = None,
+        center_factor: Optional[int] = None,
+        eval_metrics: Optional[List[Union[str, PipelineVariable]]] = None,
         **kwargs
     ):
         """A k-means clustering class :class:`~sagemaker.amazon.AmazonAlgorithmEstimatorBase`.