change: Add PipelineVariable annotation in amazon models (#3187)

Co-authored-by: Dewen Qi <[email protected]>

Author: qidewenwhen

src/sagemaker/amazon/factorization_machines.py

@@ -13,6 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
+from typing import Union, Optional
+
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
@@ -21,7 +23,9 @@
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class FactorizationMachines(AmazonAlgorithmEstimatorBase):
@@ -319,7 +323,13 @@ class FactorizationMachinesModel(Model):
     returns :class:`FactorizationMachinesPredictor`.
     """
 
-    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
+    def __init__(
+        self,
+        model_data: Union[str, PipelineVariable],
+        role: str,
+        sagemaker_session: Optional[Session] = None,
+        **kwargs
+    ):
         """Initialization for FactorizationMachinesModel class.
 
         Args:
@@ -343,6 +353,8 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             sagemaker_session.boto_region_name,
             version=FactorizationMachines.repo_version,
         )
+        pop_out_unused_kwarg("predictor_cls", kwargs, FactorizationMachinesPredictor.__name__)
+        pop_out_unused_kwarg("image_uri", kwargs, image_uri)
         super(FactorizationMachinesModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/amazon/ipinsights.py

@@ -13,6 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
+from typing import Union, Optional
+
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
@@ -22,7 +24,9 @@
 from sagemaker.model import Model
 from sagemaker.serializers import CSVSerializer
 from sagemaker.session import Session
+from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class IPInsights(AmazonAlgorithmEstimatorBase):
@@ -222,7 +226,13 @@ class IPInsightsModel(Model):
     Predictor that calculates anomaly scores for data points.
     """
 
-    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
+    def __init__(
+        self,
+        model_data: Union[str, PipelineVariable],
+        role: str,
+        sagemaker_session: Optional[Session] = None,
+        **kwargs
+    ):
         """Creates object to get insights on S3 model data.
 
         Args:
@@ -246,6 +256,8 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             sagemaker_session.boto_region_name,
             version=IPInsights.repo_version,
         )
+        pop_out_unused_kwarg("predictor_cls", kwargs, IPInsightsPredictor.__name__)
+        pop_out_unused_kwarg("image_uri", kwargs, image_uri)
         super(IPInsightsModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/amazon/kmeans.py

@@ -13,6 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
+from typing import Union, Optional
+
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
@@ -21,7 +23,9 @@
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class KMeans(AmazonAlgorithmEstimatorBase):
@@ -246,7 +250,13 @@ class KMeansModel(Model):
     Predictor to performs k-means cluster assignment.
     """
 
-    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
+    def __init__(
+        self,
+        model_data: Union[str, PipelineVariable],
+        role: str,
+        sagemaker_session: Optional[Session] = None,
+        **kwargs
+    ):
         """Initialization for KMeansModel class.
 
         Args:
@@ -270,6 +280,8 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             sagemaker_session.boto_region_name,
             version=KMeans.repo_version,
         )
+        pop_out_unused_kwarg("predictor_cls", kwargs, KMeansPredictor.__name__)
+        pop_out_unused_kwarg("image_uri", kwargs, image_uri)
         super(KMeansModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/amazon/knn.py

@@ -13,6 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
+from typing import Union, Optional
+
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
@@ -21,7 +23,9 @@
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class KNN(AmazonAlgorithmEstimatorBase):
@@ -238,7 +242,13 @@ class KNNModel(Model):
     and returns :class:`KNNPredictor`.
     """
 
-    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
+    def __init__(
+        self,
+        model_data: Union[str, PipelineVariable],
+        role: str,
+        sagemaker_session: Optional[Session] = None,
+        **kwargs
+    ):
         """Function to initialize KNNModel.
 
         Args:
@@ -262,6 +272,8 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             sagemaker_session.boto_region_name,
             version=KNN.repo_version,
         )
+        pop_out_unused_kwarg("predictor_cls", kwargs, KNNPredictor.__name__)
+        pop_out_unused_kwarg("image_uri", kwargs, image_uri)
         super(KNNModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/amazon/lda.py

@@ -13,6 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
+from typing import Union, Optional
+
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
@@ -21,7 +23,9 @@
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class LDA(AmazonAlgorithmEstimatorBase):
@@ -220,7 +224,13 @@ class LDAModel(Model):
     Predictor that transforms vectors to a lower-dimensional representation.
     """
 
-    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
+    def __init__(
+        self,
+        model_data: Union[str, PipelineVariable],
+        role: str,
+        sagemaker_session: Optional[Session] = None,
+        **kwargs
+    ):
         """Initialization for LDAModel class.
 
         Args:
@@ -244,6 +254,8 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             sagemaker_session.boto_region_name,
             version=LDA.repo_version,
         )
+        pop_out_unused_kwarg("predictor_cls", kwargs, LDAPredictor.__name__)
+        pop_out_unused_kwarg("image_uri", kwargs, image_uri)
         super(LDAModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/amazon/linear_learner.py

@@ -13,6 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
+from typing import Union, Optional
+
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
@@ -21,7 +23,9 @@
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class LinearLearner(AmazonAlgorithmEstimatorBase):
@@ -481,7 +485,13 @@ class LinearLearnerModel(Model):
     :class:`LinearLearnerPredictor`
     """
 
-    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
+    def __init__(
+        self,
+        model_data: Union[str, PipelineVariable],
+        role: str,
+        sagemaker_session: Optional[Session] = None,
+        **kwargs
+    ):
         """Initialization for LinearLearnerModel.
 
         Args:
@@ -505,6 +515,8 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             sagemaker_session.boto_region_name,
             version=LinearLearner.repo_version,
         )
+        pop_out_unused_kwarg("predictor_cls", kwargs, LinearLearnerPredictor.__name__)
+        pop_out_unused_kwarg("image_uri", kwargs, image_uri)
         super(LinearLearnerModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/amazon/ntm.py

@@ -13,6 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
+from typing import Union, Optional
+
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
@@ -21,7 +23,9 @@
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class NTM(AmazonAlgorithmEstimatorBase):
@@ -249,7 +253,13 @@ class NTMModel(Model):
     Predictor that transforms vectors to a lower-dimensional representation.
     """
 
-    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
+    def __init__(
+        self,
+        model_data: Union[str, PipelineVariable],
+        role: str,
+        sagemaker_session: Optional[Session] = None,
+        **kwargs
+    ):
         """Initialization for NTMModel class.
 
         Args:
@@ -273,6 +283,8 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             sagemaker_session.boto_region_name,
             version=NTM.repo_version,
         )
+        pop_out_unused_kwarg("predictor_cls", kwargs, NTMPredictor.__name__)
+        pop_out_unused_kwarg("image_uri", kwargs, image_uri)
         super(NTMModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/amazon/object2vec.py

@@ -13,14 +13,18 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
+from typing import Union, Optional
+
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, le, isin
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
+from sagemaker.workflow.entities import PipelineVariable
 
 
 def _list_check_subset(valid_super_list):
@@ -344,7 +348,13 @@ class Object2VecModel(Model):
     Predictor that calculates anomaly scores for datapoints.
     """
 
-    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
+    def __init__(
+        self,
+        model_data: Union[str, PipelineVariable],
+        role: str,
+        sagemaker_session: Optional[Session] = None,
+        **kwargs
+    ):
         """Initialization for Object2VecModel class.
 
         Args:
@@ -368,6 +378,8 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             sagemaker_session.boto_region_name,
             version=Object2Vec.repo_version,
         )
+        pop_out_unused_kwarg("predictor_cls", kwargs, Predictor.__name__)
+        pop_out_unused_kwarg("image_uri", kwargs, image_uri)
         super(Object2VecModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/amazon/pca.py

@@ -13,6 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
+from typing import Union, Optional
+
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
@@ -21,7 +23,9 @@
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class PCA(AmazonAlgorithmEstimatorBase):
@@ -237,7 +241,13 @@ class PCAModel(Model):
     Predictor that transforms vectors to a lower-dimensional representation.
     """
 
-    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
+    def __init__(
+        self,
+        model_data: Union[str, PipelineVariable],
+        role: str,
+        sagemaker_session: Optional[Session] = None,
+        **kwargs
+    ):
         """Initialization for PCAModel.
 
         Args:
@@ -261,6 +271,8 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             sagemaker_session.boto_region_name,
             version=PCA.repo_version,
         )
+        pop_out_unused_kwarg("predictor_cls", kwargs, PCAPredictor.__name__)
+        pop_out_unused_kwarg("image_uri", kwargs, image_uri)
         super(PCAModel, self).__init__(
             image_uri,
             model_data,