Fix all type hint and docstrings for callable (#5035)

* Fix all type hint and docstrings for callable

* Fix codestyle

---------

Co-authored-by: pintaoz <[email protected]>

Author: pintaoz-aws

src/sagemaker/amazon/hyperparameter.py

@@ -28,7 +28,7 @@ def __init__(self, name, validate=lambda _: True, validation_message="", data_ty
         """Args:
 
         name (str): The name of this hyperparameter validate
-        (callable[object]->[bool]): A validation function or list of validation
+        (Callable[object]->[bool]): A validation function or list of validation
         functions.
 
             Each function validates an object and returns False if the object

src/sagemaker/amazon/ipinsights.py

@@ -209,7 +209,7 @@ def __init__(
                 chain.
             serializer (sagemaker.serializers.BaseSerializer): Optional. Default
                 serializes input data to text/csv.
-            deserializer (callable): Optional. Default parses JSON responses
+            deserializer (Callable): Optional. Default parses JSON responses
                 using ``json.load(...)``.
             component_name (str): Optional. Name of the Amazon SageMaker inference
                 component corresponding the predictor.

src/sagemaker/automl/automl.py

@@ -478,7 +478,7 @@ def create_model(
                 training cluster for distributed training. Default: False
             model_kms_key (str): KMS key ARN used to encrypt the repacked
                 model archive file if the model is repacked
-            predictor_cls (callable[string, sagemaker.session.Session]): A
+            Callable[[string, sagemaker.session.Session], Any]: A
                 function to call to create a predictor (default: None). If
                 specified, ``deploy()``  returns the result of invoking this
                 function on the created endpoint name.
@@ -591,7 +591,7 @@ def deploy(
                 training cluster for distributed training. Default: False
             model_kms_key (str): KMS key ARN used to encrypt the repacked
                 model archive file if the model is repacked
-            predictor_cls (callable[string, sagemaker.session.Session]): A
+            predictor_cls (Callable[[string, sagemaker.session.Session], Any]): A
                 function to call to create a predictor (default: None). If
                 specified, ``deploy()``  returns the result of invoking this
                 function on the created endpoint name.
@@ -609,7 +609,7 @@ def deploy(
                 https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-inference-code.html#your-algorithms-inference-algo-ping-requests
 
         Returns:
-            callable[string, sagemaker.session.Session] or ``None``:
+            Optional[Callable[[string, sagemaker.session.Session], Any]]:
                 If ``predictor_cls`` is specified, the invocation of ``self.predictor_cls`` on
                 the created endpoint name. Otherwise, ``None``.
         """

src/sagemaker/automl/automlv2.py

@@ -1022,7 +1022,7 @@ def create_model(
                 training cluster for distributed training. Default: False
             model_kms_key (str): KMS key ARN used to encrypt the repacked
                 model archive file if the model is repacked
-            predictor_cls (callable[string, sagemaker.session.Session]): A
+            predictor_cls (Callable[[string, sagemaker.session.Session], Any]): A
                 function to call to create a predictor (default: None). If
                 specified, ``deploy()``  returns the result of invoking this
                 function on the created endpoint name.
@@ -1130,7 +1130,7 @@ def deploy(
                 training cluster for distributed training. Default: False
             model_kms_key (str): KMS key ARN used to encrypt the repacked
                 model archive file if the model is repacked
-            predictor_cls (callable[string, sagemaker.session.Session]): A
+            predictor_cls (Callable[[string, sagemaker.session.Session], Any]): A
                 function to call to create a predictor (default: None). If
                 specified, ``deploy()``  returns the result of invoking this
                 function on the created endpoint name.
@@ -1148,7 +1148,7 @@ def deploy(
                 https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-inference-code.html#your-algorithms-inference-algo-ping-requests
 
         Returns:
-            callable[string, sagemaker.session.Session] or ``None``:
+            Optional[Callable[[string, sagemaker.session.Session], Any]]:
                 If ``predictor_cls`` is specified, the invocation of ``self.predictor_cls`` on
                 the created endpoint name. Otherwise, ``None``.
         """

src/sagemaker/chainer/model.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 import logging
-from typing import Optional, Union, List, Dict
+from typing import Callable, Optional, Union, List, Dict
 
 import sagemaker
 from sagemaker import image_uris, ModelMetrics
@@ -96,7 +96,7 @@ def __init__(
         image_uri: Optional[Union[str, PipelineVariable]] = None,
         framework_version: Optional[str] = None,
         py_version: Optional[str] = None,
-        predictor_cls: callable = ChainerPredictor,
+        predictor_cls: Optional[Callable] = ChainerPredictor,
         model_server_workers: Optional[Union[int, PipelineVariable]] = None,
         **kwargs,
     ):
@@ -125,7 +125,7 @@ def __init__(
             py_version (str): Python version you want to use for executing your
                 model training code. Defaults to ``None``. Required unless
                 ``image_uri`` is provided.
-            predictor_cls (callable[str, sagemaker.session.Session]): A function
+            predictor_cls (Callable[[string, sagemaker.session.Session], Any]): A function
                 to call to create a predictor with an endpoint name and
                 SageMaker ``Session``. If specified, ``deploy()`` returns the
                 result of invoking this function on the created endpoint name.

src/sagemaker/djl_inference/model.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 import logging
-from typing import Optional, Dict, Any
+from typing import Callable, Optional, Dict, Any
 
 from sagemaker import image_uris
 from sagemaker.model import Model
@@ -54,7 +54,7 @@ def __init__(
         parallel_loading: bool = False,
         model_loading_timeout: Optional[int] = None,
         prediction_timeout: Optional[int] = None,
-        predictor_cls: callable = DJLPredictor,
+        predictor_cls: Optional[Callable] = DJLPredictor,
         huggingface_hub_token: Optional[str] = None,
         **kwargs,
     ):
@@ -97,10 +97,10 @@ def __init__(
                 None. If not provided, the default is 240 seconds.
             prediction_timeout (int): The worker predict call (handler) timeout in seconds.
                 Defaults to None. If not provided, the default is 120 seconds.
-            predictor_cls (callable[str, sagemaker.session.Session]): A function to call to create a
-                predictor with an endpoint name and SageMaker ``Session``. If specified,
-                ``deploy()`` returns
-                the result of invoking this function on the created endpoint name.
+            predictor_cls (Callable[[string, sagemaker.session.Session], Any]): A function to call
+                to create a predictor with an endpoint name and SageMaker ``Session``. If
+                specified, ``deploy()`` returns the result of invoking this function on the created
+                endpoint name.
             huggingface_hub_token (str): The HuggingFace Hub token to use for downloading the model
                 artifacts for a model stored on the huggingface hub.
                 Defaults to None. If not provided, the token must be specified in the

src/sagemaker/huggingface/model.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 import logging
-from typing import Optional, Union, List, Dict
+from typing import Callable, Optional, Union, List, Dict
 
 import sagemaker
 from sagemaker import image_uris, ModelMetrics
@@ -123,7 +123,7 @@ def __init__(
         pytorch_version: Optional[str] = None,
         py_version: Optional[str] = None,
         image_uri: Optional[Union[str, PipelineVariable]] = None,
-        predictor_cls: callable = HuggingFacePredictor,
+        predictor_cls: Optional[Callable] = HuggingFacePredictor,
         model_server_workers: Optional[Union[int, PipelineVariable]] = None,
         **kwargs,
     ):
@@ -158,7 +158,7 @@ def __init__(
                 If not specified, a default image for PyTorch will be used. If ``framework_version``
                 or ``py_version`` are ``None``, then ``image_uri`` is required. If
                 also ``None``, then a ``ValueError`` will be raised.
-            predictor_cls (callable[str, sagemaker.session.Session]): A function
+            predictor_cls (Callable[[string, sagemaker.session.Session], Any]): A function
                 to call to create a predictor with an endpoint name and
                 SageMaker ``Session``. If specified, ``deploy()`` returns the
                 result of invoking this function on the created endpoint name.
@@ -304,7 +304,7 @@ def deploy(
                 - If a wrong type of object is provided as serverless inference config or async
                     inference config
         Returns:
-            callable[string, sagemaker.session.Session] or None: Invocation of
+            Optional[Callable[[string, sagemaker.session.Session], Any]]: Invocation of
                 ``self.predictor_cls`` on the created endpoint name, if ``self.predictor_cls``
                 is not None. Otherwise, return None.
         """

src/sagemaker/jumpstart/estimator.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 
-from typing import Dict, List, Optional, Union
+from typing import Callable, Dict, List, Optional, Union
 from sagemaker import session
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
 from sagemaker.base_deserializers import BaseDeserializer
@@ -817,7 +817,7 @@ def deploy(
         explainer_config: Optional[ExplainerConfig] = None,
         image_uri: Optional[Union[str, PipelineVariable]] = None,
         role: Optional[str] = None,
-        predictor_cls: Optional[callable] = None,
+        predictor_cls: Optional[Callable] = None,
         env: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
         model_name: Optional[str] = None,
         vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = None,
@@ -918,7 +918,7 @@ def deploy(
                 It can be null if this is being used to create a Model to pass
                 to a ``PipelineModel`` which has its own Role field. (Default:
                 None).
-            predictor_cls (Optional[callable[string, sagemaker.session.Session]]): A
+            predictor_cls (Optional[Callable[[string, sagemaker.session.Session], Any]]): A
                 function to call to create a predictor (Default: None). If not
                 None, ``deploy`` will return the result of invoking this
                 function on the created endpoint name. (Default: None).

src/sagemaker/jumpstart/factory/estimator.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 
-from typing import Dict, List, Optional, Union
+from typing import Callable, Dict, List, Optional, Union
 from sagemaker import (
     environment_variables,
     hyperparameters as hyperparameters_utils,
@@ -330,7 +330,7 @@ def get_deploy_kwargs(
     explainer_config: Optional[ExplainerConfig] = None,
     image_uri: Optional[Union[str, PipelineVariable]] = None,
     role: Optional[str] = None,
-    predictor_cls: Optional[callable] = None,
+    predictor_cls: Optional[Callable] = None,
     env: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
     vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = None,
     sagemaker_session: Optional[Session] = None,

src/sagemaker/jumpstart/factory/model.py

@@ -15,7 +15,7 @@
 import json
 
 
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 from sagemaker_core.shapes import ModelAccessConfig
 from sagemaker import environment_variables, image_uris, instance_types, model_uris, script_uris
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
@@ -855,7 +855,7 @@ def get_init_kwargs(
     image_uri: Optional[Union[str, PipelineVariable]] = None,
     model_data: Optional[Union[str, PipelineVariable, dict]] = None,
     role: Optional[str] = None,
-    predictor_cls: Optional[callable] = 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,