From 212fba104b3d1477a0b86755f4f8e9bf47af896e Mon Sep 17 00:00:00 2001 From: pintaoz Date: Wed, 12 Feb 2025 14:46:39 -0800 Subject: [PATCH 1/2] Fix all type hint and docstrings for callable --- src/sagemaker/amazon/hyperparameter.py | 2 +- src/sagemaker/amazon/ipinsights.py | 2 +- src/sagemaker/automl/automl.py | 6 +- src/sagemaker/automl/automlv2.py | 6 +- src/sagemaker/chainer/model.py | 6 +- src/sagemaker/djl_inference/model.py | 6 +- src/sagemaker/huggingface/model.py | 8 +-- src/sagemaker/jumpstart/estimator.py | 6 +- src/sagemaker/jumpstart/factory/estimator.py | 4 +- src/sagemaker/jumpstart/factory/model.py | 4 +- src/sagemaker/jumpstart/model.py | 6 +- src/sagemaker/jumpstart/types.py | 6 +- src/sagemaker/model.py | 12 ++-- src/sagemaker/multidatamodel.py | 2 +- src/sagemaker/mxnet/model.py | 10 ++-- src/sagemaker/pipeline.py | 8 +-- src/sagemaker/pytorch/model.py | 6 +- src/sagemaker/remote_function/job.py | 2 +- src/sagemaker/serve/builder/schema_builder.py | 3 +- src/sagemaker/serve/utils/tuning.py | 3 +- src/sagemaker/sklearn/model.py | 6 +- src/sagemaker/tensorflow/model.py | 10 ++-- src/sagemaker/utils.py | 2 +- src/sagemaker/xgboost/model.py | 6 +- tests/integ/sagemaker/modules/train/test.py | 59 +++++++++++++++++++ 25 files changed, 126 insertions(+), 65 deletions(-) create mode 100644 tests/integ/sagemaker/modules/train/test.py diff --git a/src/sagemaker/amazon/hyperparameter.py b/src/sagemaker/amazon/hyperparameter.py index 856927cb13..b479f8a271 100644 --- a/src/sagemaker/amazon/hyperparameter.py +++ b/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 diff --git a/src/sagemaker/amazon/ipinsights.py b/src/sagemaker/amazon/ipinsights.py index 737d13dd44..bc8e1b5d86 100644 --- a/src/sagemaker/amazon/ipinsights.py +++ b/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. diff --git a/src/sagemaker/automl/automl.py b/src/sagemaker/automl/automl.py index bb4059c03a..e18d7ba2b9 100644 --- a/src/sagemaker/automl/automl.py +++ b/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``. """ diff --git a/src/sagemaker/automl/automlv2.py b/src/sagemaker/automl/automlv2.py index 0819e5384e..b071be3b24 100644 --- a/src/sagemaker/automl/automlv2.py +++ b/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``. """ diff --git a/src/sagemaker/chainer/model.py b/src/sagemaker/chainer/model.py index 806009b0f6..c2d2187b69 100644 --- a/src/sagemaker/chainer/model.py +++ b/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. diff --git a/src/sagemaker/djl_inference/model.py b/src/sagemaker/djl_inference/model.py index 8c724a6502..d27110ac90 100644 --- a/src/sagemaker/djl_inference/model.py +++ b/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,7 +97,7 @@ 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_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. diff --git a/src/sagemaker/huggingface/model.py b/src/sagemaker/huggingface/model.py index ea99be2fc0..05b981d21b 100644 --- a/src/sagemaker/huggingface/model.py +++ b/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. """ diff --git a/src/sagemaker/jumpstart/estimator.py b/src/sagemaker/jumpstart/estimator.py index def5121b9b..50f197c30e 100644 --- a/src/sagemaker/jumpstart/estimator.py +++ b/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). diff --git a/src/sagemaker/jumpstart/factory/estimator.py b/src/sagemaker/jumpstart/factory/estimator.py index e4020a39bd..2a54d9c4de 100644 --- a/src/sagemaker/jumpstart/factory/estimator.py +++ b/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, diff --git a/src/sagemaker/jumpstart/factory/model.py b/src/sagemaker/jumpstart/factory/model.py index 328e1e8227..4245c5ac91 100644 --- a/src/sagemaker/jumpstart/factory/model.py +++ b/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, diff --git a/src/sagemaker/jumpstart/model.py b/src/sagemaker/jumpstart/model.py index d26cf237f2..7dec3d78f9 100644 --- a/src/sagemaker/jumpstart/model.py +++ b/src/sagemaker/jumpstart/model.py @@ -14,7 +14,7 @@ from __future__ import absolute_import -from typing import Dict, List, Optional, Any, Union +from typing import Callable, Dict, List, Optional, Any, Union import pandas as pd from botocore.exceptions import ClientError @@ -95,7 +95,7 @@ def __init__( 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, @@ -149,7 +149,7 @@ def __init__( 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). diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py index f59e2eddf4..3dee2b3553 100644 --- a/src/sagemaker/jumpstart/types.py +++ b/src/sagemaker/jumpstart/types.py @@ -16,7 +16,7 @@ import re from copy import deepcopy from enum import Enum -from typing import Any, Dict, List, Optional, Set, Union +from typing import Any, Callable, Dict, List, Optional, Set, Union from sagemaker_core.shapes import ModelAccessConfig as CoreModelAccessConfig from sagemaker.model_card.model_card import ModelCard, ModelPackageModelCard from sagemaker.utils import ( @@ -2150,7 +2150,7 @@ def __init__( image_uri: Optional[Union[str, Any]] = None, model_data: Optional[Union[str, Any, dict]] = None, role: Optional[str] = None, - predictor_cls: Optional[callable] = None, + predictor_cls: Optional[Callable] = None, env: Optional[Dict[str, Union[str, Any]]] = None, name: Optional[str] = None, vpc_config: Optional[Dict[str, List[Union[str, Any]]]] = None, @@ -2698,7 +2698,7 @@ def __init__( explainer_config: Optional[Any] = None, image_uri: Optional[Union[str, Any]] = None, role: Optional[str] = None, - predictor_cls: Optional[callable] = None, + predictor_cls: Optional[Callable] = None, env: Optional[Dict[str, Union[str, Any]]] = None, model_name: Optional[str] = None, vpc_config: Optional[Dict[str, List[Union[str, Any]]]] = None, diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py index cdc1dfd898..5494bf5e22 100644 --- a/src/sagemaker/model.py +++ b/src/sagemaker/model.py @@ -20,7 +20,7 @@ import os import re import copy -from typing import List, Dict, Optional, Union, Any +from typing import Callable, List, Dict, Optional, Union, Any import sagemaker from sagemaker import ( @@ -154,7 +154,7 @@ def __init__( 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, @@ -186,7 +186,7 @@ def __init__( 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 (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 not None, ``deploy`` will return the result of invoking this function on the created endpoint name. @@ -1501,7 +1501,7 @@ def deploy( inference config or - If inference recommendation id is specified along with incompatible parameters Returns: - callable[string, sagemaker.session.Session] or None: Invocation of + Callable[[string, sagemaker.session.Session], Any] or None: Invocation of ``self.predictor_cls`` on the created endpoint name, if ``self.predictor_cls`` is not None. Otherwise, return None. """ @@ -1959,7 +1959,7 @@ def __init__( role: Optional[str] = None, entry_point: Optional[str] = None, source_dir: 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, container_log_level: Union[int, PipelineVariable] = logging.INFO, @@ -2012,7 +2012,7 @@ def __init__( >>> |----- test.py You can assign entry_point='inference.py', source_dir='src'. - 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 not None, ``deploy`` will return the result of invoking this function on the created endpoint name. diff --git a/src/sagemaker/multidatamodel.py b/src/sagemaker/multidatamodel.py index 9ed348c927..43a3588e6f 100644 --- a/src/sagemaker/multidatamodel.py +++ b/src/sagemaker/multidatamodel.py @@ -223,7 +223,7 @@ def deploy( Amazon SageMaker Model Monitoring. Default: None. 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. diff --git a/src/sagemaker/mxnet/model.py b/src/sagemaker/mxnet/model.py index 0dcd71741d..fa0c691d2d 100644 --- a/src/sagemaker/mxnet/model.py +++ b/src/sagemaker/mxnet/model.py @@ -14,7 +14,7 @@ from __future__ import absolute_import import logging -from typing import Union, Optional, List, Dict +from typing import Callable, Union, Optional, List, Dict import packaging.version @@ -68,9 +68,9 @@ def __init__( manages interactions with Amazon SageMaker APIs and any other AWS services needed. If not specified, the estimator creates one using the default AWS configuration chain. - serializer (callable): Optional. Default serializes input data to + serializer (Callable): Optional. Default serializes input data to json. Handles dicts, lists, and numpy arrays. - deserializer (callable): Optional. Default parses the response using + deserializer (Callable): Optional. Default parses the response using ``json.load(...)``. component_name (str): Optional. Name of the Amazon SageMaker inference component corresponding to the predictor. @@ -98,7 +98,7 @@ def __init__( framework_version: str = _LOWEST_MMS_VERSION, py_version: Optional[str] = None, image_uri: Optional[Union[str, PipelineVariable]] = None, - predictor_cls: callable = MXNetPredictor, + predictor_cls: Optional[Callable] = MXNetPredictor, model_server_workers: Optional[Union[int, PipelineVariable]] = None, **kwargs, ): @@ -127,7 +127,7 @@ def __init__( If ``framework_version`` or ``py_version`` are ``None``, then ``image_uri`` is required. If ``image_uri`` is 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. diff --git a/src/sagemaker/pipeline.py b/src/sagemaker/pipeline.py index 04fbc1cc93..1d1ece5965 100644 --- a/src/sagemaker/pipeline.py +++ b/src/sagemaker/pipeline.py @@ -13,7 +13,7 @@ """Placeholder docstring""" from __future__ import absolute_import -from typing import Optional, Dict, List, Union +from typing import Callable, Optional, Dict, List, Union import sagemaker from sagemaker import ModelMetrics, Model @@ -54,7 +54,7 @@ def __init__( self, models: List[Model], role: str = None, - predictor_cls: Optional[callable] = None, + predictor_cls: Optional[Callable] = None, name: Optional[str] = None, vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = None, sagemaker_session: Optional[Session] = None, @@ -75,7 +75,7 @@ def __init__( endpoints use this role to access training data and model artifacts. After the endpoint is created, the inference code might use the IAM role, if it needs to access an AWS resource. - 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 not None, ``deploy`` will return the result of invoking this function on the created endpoint name. @@ -230,7 +230,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: 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. """ diff --git a/src/sagemaker/pytorch/model.py b/src/sagemaker/pytorch/model.py index 329f9b83b5..958327ba08 100644 --- a/src/sagemaker/pytorch/model.py +++ b/src/sagemaker/pytorch/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 packaging.version @@ -99,7 +99,7 @@ def __init__( framework_version: str = "1.3", py_version: Optional[str] = None, image_uri: Optional[Union[str, PipelineVariable]] = None, - predictor_cls: callable = PyTorchPredictor, + predictor_cls: Optional[Callable] = PyTorchPredictor, model_server_workers: Optional[Union[int, PipelineVariable]] = None, **kwargs, ): @@ -128,7 +128,7 @@ def __init__( If ``framework_version`` or ``py_version`` are ``None``, then ``image_uri`` is required. If ``image_uri`` is 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. diff --git a/src/sagemaker/remote_function/job.py b/src/sagemaker/remote_function/job.py index f6c3a58ad6..52cb0ff04f 100644 --- a/src/sagemaker/remote_function/job.py +++ b/src/sagemaker/remote_function/job.py @@ -870,7 +870,7 @@ def compile( job_settings: _JobSettings, job_name: str, s3_base_uri: str, - func: callable, + func: Callable, func_args: tuple, func_kwargs: dict, run_info=None, diff --git a/src/sagemaker/serve/builder/schema_builder.py b/src/sagemaker/serve/builder/schema_builder.py index 3fd1816d0e..7f70e98747 100644 --- a/src/sagemaker/serve/builder/schema_builder.py +++ b/src/sagemaker/serve/builder/schema_builder.py @@ -4,6 +4,7 @@ import io import logging from pathlib import Path +from typing import Callable import numpy as np from pandas import DataFrame @@ -286,7 +287,7 @@ def _is_path_to_file(data: object) -> bool: def _validate_translations( - payload: object, serialize_callable: callable, deserialize_callable: callable + payload: object, serialize_callable: Callable, deserialize_callable: Callable ) -> None: """Placeholder docstring""" try: diff --git a/src/sagemaker/serve/utils/tuning.py b/src/sagemaker/serve/utils/tuning.py index b93c01b522..5a63cfe508 100644 --- a/src/sagemaker/serve/utils/tuning.py +++ b/src/sagemaker/serve/utils/tuning.py @@ -7,6 +7,7 @@ import collections from multiprocessing.pool import ThreadPool from math import ceil +from typing import Callable import pandas as pd from numpy import percentile, std from sagemaker.serve.model_server.djl_serving.utils import _tokens_from_chars, _tokens_from_words @@ -152,7 +153,7 @@ def _tokens_per_second(generated_text: str, max_token_length: int, latency: floa return min(est_tokens, max_token_length) / latency -def _timed_invoke(predict: callable, sample_input: object) -> tuple: +def _timed_invoke(predict: Callable, sample_input: object) -> tuple: """Placeholder docstring""" start_timer = perf_counter() response = predict(sample_input) diff --git a/src/sagemaker/sklearn/model.py b/src/sagemaker/sklearn/model.py index c3727b2fb5..a9b0e2e8f0 100644 --- a/src/sagemaker/sklearn/model.py +++ b/src/sagemaker/sklearn/model.py @@ -14,7 +14,7 @@ from __future__ import absolute_import import logging -from typing import Union, Optional, List, Dict +from typing import Callable, Union, Optional, List, Dict import sagemaker from sagemaker import image_uris, ModelMetrics @@ -92,7 +92,7 @@ def __init__( framework_version: Optional[str] = None, py_version: str = "py3", image_uri: Optional[Union[str, PipelineVariable]] = None, - predictor_cls: callable = SKLearnPredictor, + predictor_cls: Optional[Callable] = SKLearnPredictor, model_server_workers: Optional[Union[int, PipelineVariable]] = None, **kwargs, ): @@ -122,7 +122,7 @@ def __init__( If ``framework_version`` or ``py_version`` are ``None``, then ``image_uri`` is required. If ``image_uri`` is 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. diff --git a/src/sagemaker/tensorflow/model.py b/src/sagemaker/tensorflow/model.py index fe20994e20..c7f624114f 100644 --- a/src/sagemaker/tensorflow/model.py +++ b/src/sagemaker/tensorflow/model.py @@ -14,7 +14,7 @@ from __future__ import absolute_import import logging -from typing import Union, Optional, List, Dict +from typing import Callable, Union, Optional, List, Dict import sagemaker from sagemaker import image_uris, s3, ModelMetrics @@ -62,9 +62,9 @@ def __init__( manages interactions with Amazon SageMaker APIs and any other AWS services needed. If not specified, the estimator creates one using the default AWS configuration chain. - serializer (callable): Optional. Default serializes input data to + serializer (Callable): Optional. Default serializes input data to json. Handles dicts, lists, and numpy arrays. - deserializer (callable): Optional. Default parses the response using + deserializer (Callable): Optional. Default parses the response using ``json.load(...)``. model_name (str): Optional. The name of the SavedModel model that should handle the request. If not specified, the endpoint's @@ -146,7 +146,7 @@ def __init__( image_uri: Optional[Union[str, PipelineVariable]] = None, framework_version: Optional[str] = None, container_log_level: Optional[int] = None, - predictor_cls: callable = TensorFlowPredictor, + predictor_cls: Optional[Callable] = TensorFlowPredictor, **kwargs, ): """Initialize a Model. @@ -174,7 +174,7 @@ def __init__( container_log_level (int): Log level to use within the container (default: logging.ERROR). Valid values are defined in the Python logging module. - 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. diff --git a/src/sagemaker/utils.py b/src/sagemaker/utils.py index e8602de8d7..c575b1eeb6 100644 --- a/src/sagemaker/utils.py +++ b/src/sagemaker/utils.py @@ -726,7 +726,7 @@ def retry_with_backoff(callable_func, num_attempts=8, botocore_client_error_code """Retry with backoff until maximum attempts are reached Args: - callable_func (callable): The callable function to retry. + callable_func (Callable): The callable function to retry. num_attempts (int): The maximum number of attempts to retry.(Default: 8) botocore_client_error_code (str): The specific Botocore ClientError exception error code on which to retry on. diff --git a/src/sagemaker/xgboost/model.py b/src/sagemaker/xgboost/model.py index ea532b4c39..5744cc2422 100644 --- a/src/sagemaker/xgboost/model.py +++ b/src/sagemaker/xgboost/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 @@ -91,7 +91,7 @@ def __init__( framework_version: str = None, image_uri: Optional[Union[str, PipelineVariable]] = None, py_version: str = "py3", - predictor_cls: callable = XGBoostPredictor, + predictor_cls: Optional[Callable] = XGBoostPredictor, model_server_workers: Optional[Union[int, PipelineVariable]] = None, **kwargs, ): @@ -113,7 +113,7 @@ def __init__( (default: 'py3'). framework_version (str): XGBoost version you want to use for executing your model training code. - predictor_cls (callable[str, sagemaker.session.Session]): A function to call to create + 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. diff --git a/tests/integ/sagemaker/modules/train/test.py b/tests/integ/sagemaker/modules/train/test.py new file mode 100644 index 0000000000..9c87e978e5 --- /dev/null +++ b/tests/integ/sagemaker/modules/train/test.py @@ -0,0 +1,59 @@ +# import os +# import tempfile +# from tests.integ import DATA_DIR +# import tests.integ.lock as lock + +# from sagemaker.modules.configs import Compute, InputData, SourceCode +# from sagemaker.modules.distributed import Torchrun +# from sagemaker.modules.train.model_trainer import Mode, ModelTrainer +# import subprocess +# from sagemaker.modules import Session + +# DEFAULT_CPU_IMAGE = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310" +# CWD = os.getcwd() +# SOURCE_DIR = os.path.join(DATA_DIR, "modules/local_script") +# LOCK_PATH = os.path.join(tempfile.gettempdir(), "sagemaker_test_local_mode_lock") + +# source_code = SourceCode( +# source_dir=SOURCE_DIR, +# entry_script="local_training_script.py", +# ) + +# compute = Compute( +# instance_type="local_cpu", +# instance_count=1, +# ) + +# session = Session() +# bucket = session.default_bucket() +# session.upload_data( +# path=os.path.join(SOURCE_DIR, "data/train/"), +# bucket=bucket, +# key_prefix="data/train", +# ) +# session.upload_data( +# path=os.path.join(SOURCE_DIR, "data/test/"), +# bucket=bucket, +# key_prefix="data/test", +# ) + +# train_data = InputData(channel_name="train", data_source=f"s3://{bucket}/data/train/") + +# test_data = InputData(channel_name="test", data_source=f"s3://{bucket}/data/test/") + +# model_trainer = ModelTrainer( +# training_image=DEFAULT_CPU_IMAGE, +# source_code=source_code, +# compute=compute, +# input_data_config=[train_data, test_data], +# base_job_name="local_mode_single_container_local_data", +# training_mode=Mode.LOCAL_CONTAINER, +# remove_inputs_and_container_artifacts=False, +# ) + +# model_trainer.train() + +from sagemaker.local import LocalSession + +sagemaker_session = LocalSession() +sagemaker_session.config = {'local': {'local_code': True}} \ No newline at end of file From 495a0308962a7b1b02847337dfd50fc79079e6f2 Mon Sep 17 00:00:00 2001 From: pintaoz Date: Wed, 12 Feb 2025 15:10:47 -0800 Subject: [PATCH 2/2] Fix codestyle --- src/sagemaker/djl_inference/model.py | 8 +-- src/sagemaker/xgboost/model.py | 4 +- tests/integ/sagemaker/modules/train/test.py | 59 --------------------- 3 files changed, 6 insertions(+), 65 deletions(-) delete mode 100644 tests/integ/sagemaker/modules/train/test.py diff --git a/src/sagemaker/djl_inference/model.py b/src/sagemaker/djl_inference/model.py index d27110ac90..94db4efe29 100644 --- a/src/sagemaker/djl_inference/model.py +++ b/src/sagemaker/djl_inference/model.py @@ -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[[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. + 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 diff --git a/src/sagemaker/xgboost/model.py b/src/sagemaker/xgboost/model.py index 5744cc2422..f4797c79e7 100644 --- a/src/sagemaker/xgboost/model.py +++ b/src/sagemaker/xgboost/model.py @@ -113,8 +113,8 @@ def __init__( (default: 'py3'). framework_version (str): XGBoost version you want to use for executing your model training code. - predictor_cls (Callable[[string, sagemaker.session.Session], Any]): A function to call to create - a predictor with an endpoint name and SageMaker ``Session``. + 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. model_server_workers (int or PipelineVariable): Optional. The number of worker processes diff --git a/tests/integ/sagemaker/modules/train/test.py b/tests/integ/sagemaker/modules/train/test.py deleted file mode 100644 index 9c87e978e5..0000000000 --- a/tests/integ/sagemaker/modules/train/test.py +++ /dev/null @@ -1,59 +0,0 @@ -# import os -# import tempfile -# from tests.integ import DATA_DIR -# import tests.integ.lock as lock - -# from sagemaker.modules.configs import Compute, InputData, SourceCode -# from sagemaker.modules.distributed import Torchrun -# from sagemaker.modules.train.model_trainer import Mode, ModelTrainer -# import subprocess -# from sagemaker.modules import Session - -# DEFAULT_CPU_IMAGE = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310" -# CWD = os.getcwd() -# SOURCE_DIR = os.path.join(DATA_DIR, "modules/local_script") -# LOCK_PATH = os.path.join(tempfile.gettempdir(), "sagemaker_test_local_mode_lock") - -# source_code = SourceCode( -# source_dir=SOURCE_DIR, -# entry_script="local_training_script.py", -# ) - -# compute = Compute( -# instance_type="local_cpu", -# instance_count=1, -# ) - -# session = Session() -# bucket = session.default_bucket() -# session.upload_data( -# path=os.path.join(SOURCE_DIR, "data/train/"), -# bucket=bucket, -# key_prefix="data/train", -# ) -# session.upload_data( -# path=os.path.join(SOURCE_DIR, "data/test/"), -# bucket=bucket, -# key_prefix="data/test", -# ) - -# train_data = InputData(channel_name="train", data_source=f"s3://{bucket}/data/train/") - -# test_data = InputData(channel_name="test", data_source=f"s3://{bucket}/data/test/") - -# model_trainer = ModelTrainer( -# training_image=DEFAULT_CPU_IMAGE, -# source_code=source_code, -# compute=compute, -# input_data_config=[train_data, test_data], -# base_job_name="local_mode_single_container_local_data", -# training_mode=Mode.LOCAL_CONTAINER, -# remove_inputs_and_container_artifacts=False, -# ) - -# model_trainer.train() - -from sagemaker.local import LocalSession - -sagemaker_session = LocalSession() -sagemaker_session.config = {'local': {'local_code': True}} \ No newline at end of file