From 7767b22aa7d13c061f5a4080f16f86b02815d872 Mon Sep 17 00:00:00 2001 From: DemyCode Date: Fri, 23 Jun 2023 17:21:45 +0200 Subject: [PATCH 01/10] feat: adding sagemaker raw stubs --- pyrightconfig.stricter.json | 1 + stubs/sagemaker/METADATA.toml | 1 + stubs/sagemaker/sagemaker/__init__.pyi | 49 ++ stubs/sagemaker/sagemaker/_studio.pyi | 4 + stubs/sagemaker/sagemaker/accept_types.pyi | 16 + stubs/sagemaker/sagemaker/algorithm.pyi | 75 +++ stubs/sagemaker/sagemaker/amazon/__init__.pyi | 0 .../sagemaker/amazon/amazon_estimator.pyi | 84 +++ stubs/sagemaker/sagemaker/amazon/common.pyi | 24 + .../amazon/factorization_machines.pyi | 81 +++ .../sagemaker/amazon/hyperparameter.pyi | 14 + .../sagemaker/sagemaker/amazon/ipinsights.pyi | 52 ++ stubs/sagemaker/sagemaker/amazon/kmeans.pyi | 54 ++ stubs/sagemaker/sagemaker/amazon/knn.pyi | 51 ++ stubs/sagemaker/sagemaker/amazon/lda.pyi | 44 ++ .../sagemaker/amazon/linear_learner.pyi | 120 ++++ stubs/sagemaker/sagemaker/amazon/ntm.pyi | 57 ++ .../sagemaker/sagemaker/amazon/object2vec.pyi | 95 +++ stubs/sagemaker/sagemaker/amazon/pca.pyi | 42 ++ .../sagemaker/amazon/randomcutforest.pyi | 41 ++ .../sagemaker/sagemaker/amazon/record_pb2.pyi | 9 + .../sagemaker/sagemaker/amazon/validation.pyi | 6 + stubs/sagemaker/sagemaker/analytics.pyi | 64 ++ .../sagemaker/sagemaker/apiutils/__init__.pyi | 0 .../sagemaker/apiutils/_base_types.pyi | 16 + .../sagemaker/apiutils/_boto_functions.pyi | 4 + stubs/sagemaker/sagemaker/apiutils/_utils.pyi | 3 + .../sagemaker/async_inference/__init__.pyi | 3 + .../async_inference_config.pyi | 16 + .../async_inference_response.pyi | 8 + .../async_inference/waiter_config.pyi | 6 + stubs/sagemaker/sagemaker/automl/__init__.pyi | 0 stubs/sagemaker/sagemaker/automl/automl.pyi | 145 +++++ .../sagemaker/automl/candidate_estimator.pyi | 30 + .../sagemaker/base_deserializers.pyi | 49 ++ stubs/sagemaker/sagemaker/base_predictor.pyi | 71 +++ .../sagemaker/sagemaker/base_serializers.pyi | 47 ++ .../sagemaker/sagemaker/chainer/__init__.pyi | 2 + .../sagemaker/sagemaker/chainer/defaults.pyi | 1 + .../sagemaker/sagemaker/chainer/estimator.pyi | 40 ++ stubs/sagemaker/sagemaker/chainer/model.pyi | 69 +++ stubs/sagemaker/sagemaker/clarify.pyi | 281 +++++++++ stubs/sagemaker/sagemaker/cli/__init__.pyi | 0 .../sagemaker/cli/compatibility/__init__.pyi | 0 .../cli/compatibility/v2/__init__.pyi | 0 .../cli/compatibility/v2/ast_transformer.pyi | 15 + .../sagemaker/cli/compatibility/v2/files.pyi | 18 + .../compatibility/v2/modifiers/__init__.pyi | 13 + .../compatibility/v2/modifiers/airflow.pyi | 20 + .../v2/modifiers/deprecated_params.pyi | 9 + .../v2/modifiers/framework_version.pyi | 15 + .../compatibility/v2/modifiers/image_uris.pyi | 14 + .../compatibility/v2/modifiers/matching.pyi | 6 + .../compatibility/v2/modifiers/modifier.pyi | 9 + .../compatibility/v2/modifiers/parsing.pyi | 2 + .../compatibility/v2/modifiers/predictors.pyi | 14 + .../v2/modifiers/renamed_params.pyi | 76 +++ .../cli/compatibility/v2/modifiers/serde.pyi | 47 ++ .../v2/modifiers/tf_legacy_mode.pyi | 17 + .../cli/compatibility/v2/modifiers/tfs.pyi | 18 + .../v2/modifiers/training_input.pyi | 22 + .../v2/modifiers/training_params.pyi | 10 + .../compatibility/v2/sagemaker_upgrade_v2.pyi | 1 + .../sagemaker/cli/framework_upgrade.pyi | 14 + stubs/sagemaker/sagemaker/collection.pyi | 12 + stubs/sagemaker/sagemaker/config/__init__.pyi | 149 +++++ stubs/sagemaker/sagemaker/config/config.pyi | 10 + .../sagemaker/config/config_schema.pyi | 191 ++++++ .../sagemaker/config/config_utils.pyi | 1 + stubs/sagemaker/sagemaker/content_types.pyi | 21 + .../sagemaker/dataset_definition/__init__.pyi | 6 + .../sagemaker/dataset_definition/inputs.pyi | 51 ++ .../sagemaker/sagemaker/debugger/__init__.pyi | 23 + .../sagemaker/sagemaker/debugger/debugger.pyi | 129 ++++ .../sagemaker/debugger/framework_profile.pyi | 22 + .../sagemaker/debugger/metrics_config.pyi | 75 +++ .../sagemaker/debugger/profiler_config.pyi | 20 + .../sagemaker/debugger/profiler_constants.pyi | 17 + stubs/sagemaker/sagemaker/debugger/utils.pyi | 30 + stubs/sagemaker/sagemaker/deprecations.pyi | 19 + stubs/sagemaker/sagemaker/deserializers.pyi | 30 + .../sagemaker/djl_inference/__init__.pyi | 7 + .../sagemaker/djl_inference/defaults.pyi | 8 + .../sagemaker/djl_inference/model.pyi | 168 +++++ .../sagemaker/drift_check_baselines.pyi | 27 + .../sagemaker/environment_variables.pyi | 13 + stubs/sagemaker/sagemaker/estimator.pyi | 393 ++++++++++++ stubs/sagemaker/sagemaker/exceptions.pyi | 34 + .../sagemaker/experiments/__init__.pyi | 8 + .../sagemaker/experiments/_api_types.pyi | 85 +++ .../sagemaker/experiments/_environment.pyi | 28 + .../sagemaker/experiments/_helper.pyi | 47 ++ .../sagemaker/experiments/_metrics.pyi | 42 ++ .../sagemaker/experiments/_run_context.pyi | 9 + .../sagemaker/experiments/_utils.pyi | 13 + .../sagemaker/experiments/experiment.pyi | 29 + stubs/sagemaker/sagemaker/experiments/run.pyi | 81 +++ .../sagemaker/sagemaker/experiments/trial.pyi | 44 ++ .../sagemaker/experiments/trial_component.pyi | 60 ++ .../sagemaker/explainer/__init__.pyi | 8 + .../explainer/clarify_explainer_config.pyi | 68 ++ .../sagemaker/explainer/explainer_config.pyi | 8 + .../sagemaker/feature_store/__init__.pyi | 0 .../feature_store/dataset_builder.pyi | 128 ++++ .../feature_store/feature_definition.pyi | 28 + .../sagemaker/feature_store/feature_group.pyi | 117 ++++ .../feature_processor/__init__.pyi | 16 + .../sagemaker/feature_store/feature_store.pyi | 48 ++ .../sagemaker/feature_store/inputs.pyi | 135 ++++ stubs/sagemaker/sagemaker/fw_utils.pyi | 74 +++ stubs/sagemaker/sagemaker/git_utils.pyi | 3 + .../sagemaker/huggingface/__init__.pyi | 5 + .../sagemaker/huggingface/estimator.pyi | 46 ++ .../sagemaker/huggingface/llm_utils.pyi | 7 + .../sagemaker/sagemaker/huggingface/model.pyi | 99 +++ .../sagemaker/huggingface/processing.pyi | 35 ++ .../training_compiler/__init__.pyi | 0 .../huggingface/training_compiler/config.pyi | 14 + stubs/sagemaker/sagemaker/hyperparameters.pyi | 22 + stubs/sagemaker/sagemaker/image_uris.pyi | 46 ++ .../inference_recommender/__init__.pyi | 4 + .../inference_recommender_mixin.pyi | 36 ++ stubs/sagemaker/sagemaker/inputs.pyi | 83 +++ stubs/sagemaker/sagemaker/instance_group.pyi | 12 + stubs/sagemaker/sagemaker/instance_types.pyi | 21 + stubs/sagemaker/sagemaker/job.pyi | 18 + .../sagemaker/jumpstart/__init__.pyi | 0 .../sagemaker/jumpstart/accessors.pyi | 21 + .../jumpstart/artifacts/__init__.pyi | 0 .../artifacts/environment_variables.pyi | 0 .../jumpstart/artifacts/hyperparameters.pyi | 0 .../jumpstart/artifacts/image_uris.pyi | 0 .../artifacts/incremental_training.pyi | 0 .../jumpstart/artifacts/instance_types.pyi | 0 .../sagemaker/jumpstart/artifacts/kwargs.pyi | 0 .../artifacts/metric_definitions.pyi | 0 .../jumpstart/artifacts/model_uris.pyi | 0 .../jumpstart/artifacts/predictors.pyi | 0 .../sagemaker/jumpstart/artifacts/prepack.pyi | 0 .../jumpstart/artifacts/resource_names.pyi | 0 .../jumpstart/artifacts/script_uris.pyi | 0 stubs/sagemaker/sagemaker/jumpstart/cache.pyi | 30 + .../sagemaker/jumpstart/constants.pyi | 28 + stubs/sagemaker/sagemaker/jumpstart/enums.pyi | 55 ++ .../sagemaker/jumpstart/estimator.pyi | 127 ++++ .../sagemaker/jumpstart/exceptions.pyi | 26 + .../sagemaker/jumpstart/factory/__init__.pyi | 0 .../sagemaker/jumpstart/factory/estimator.pyi | 122 ++++ .../sagemaker/jumpstart/factory/model.pyi | 74 +++ .../sagemaker/sagemaker/jumpstart/filters.pyi | 87 +++ stubs/sagemaker/sagemaker/jumpstart/model.pyi | 70 +++ .../sagemaker/jumpstart/notebook_utils.pyi | 16 + .../sagemaker/jumpstart/parameters.pyi | 6 + stubs/sagemaker/sagemaker/jumpstart/types.pyi | 449 ++++++++++++++ stubs/sagemaker/sagemaker/jumpstart/utils.pyi | 53 ++ .../sagemaker/jumpstart/validators.pyi | 11 + stubs/sagemaker/sagemaker/lambda_helper.pyi | 41 ++ .../sagemaker/sagemaker/lineage/__init__.pyi | 0 .../sagemaker/lineage/_api_types.pyi | 70 +++ stubs/sagemaker/sagemaker/lineage/_utils.pyi | 1 + stubs/sagemaker/sagemaker/lineage/action.pyi | 62 ++ .../sagemaker/sagemaker/lineage/artifact.pyi | 73 +++ .../sagemaker/lineage/association.pyi | 35 ++ stubs/sagemaker/sagemaker/lineage/context.pyi | 67 ++ .../lineage/lineage_trial_component.pyi | 34 + stubs/sagemaker/sagemaker/lineage/query.pyi | 91 +++ .../sagemaker/lineage/visualizer.pyi | 18 + stubs/sagemaker/sagemaker/local/__init__.pyi | 6 + stubs/sagemaker/sagemaker/local/data.pyi | 51 ++ stubs/sagemaker/sagemaker/local/entities.pyi | 152 +++++ .../sagemaker/sagemaker/local/exceptions.pyi | 6 + stubs/sagemaker/sagemaker/local/image.pyi | 54 ++ .../sagemaker/local/local_session.pyi | 85 +++ stubs/sagemaker/sagemaker/local/pipeline.pyi | 49 ++ stubs/sagemaker/sagemaker/local/utils.pyi | 11 + stubs/sagemaker/sagemaker/logs.pyi | 20 + .../sagemaker/metadata_properties.pyi | 17 + .../sagemaker/metric_definitions.pyi | 12 + stubs/sagemaker/sagemaker/model.pyi | 211 +++++++ .../sagemaker/model_card/__init__.pyi | 25 + .../model_card/evaluation_metric_parsers.pyi | 24 + .../sagemaker/model_card/helpers.pyi | 65 ++ .../sagemaker/model_card/model_card.pyi | 272 ++++++++ .../model_card/schema_constraints.pyi | 44 ++ stubs/sagemaker/sagemaker/model_metrics.pyi | 47 ++ .../sagemaker/model_monitor/__init__.pyi | 32 + .../clarify_model_monitoring.pyi | 186 ++++++ .../cron_expression_generator.pyi | 7 + .../model_monitor/data_capture_config.pyi | 22 + .../data_quality_monitoring_config.pyi | 18 + .../model_monitor/dataset_format.pyi | 15 + .../model_monitor/model_monitoring.pyi | 382 ++++++++++++ .../model_monitor/monitoring_alert.pyi | 41 ++ .../model_monitor/monitoring_files.pyi | 79 +++ stubs/sagemaker/sagemaker/model_uris.pyi | 13 + stubs/sagemaker/sagemaker/multidatamodel.pyi | 50 ++ stubs/sagemaker/sagemaker/mxnet/__init__.pyi | 3 + stubs/sagemaker/sagemaker/mxnet/defaults.pyi | 1 + stubs/sagemaker/sagemaker/mxnet/estimator.pyi | 33 + stubs/sagemaker/sagemaker/mxnet/model.pyi | 69 +++ .../sagemaker/sagemaker/mxnet/processing.pyi | 30 + stubs/sagemaker/sagemaker/network.pyi | 17 + stubs/sagemaker/sagemaker/parameter.pyi | 40 ++ stubs/sagemaker/sagemaker/pipeline.pyi | 86 +++ stubs/sagemaker/sagemaker/predictor.pyi | 14 + stubs/sagemaker/sagemaker/predictor_async.pyi | 44 ++ stubs/sagemaker/sagemaker/processing.pyi | 230 +++++++ .../sagemaker/sagemaker/pytorch/__init__.pyi | 4 + .../sagemaker/sagemaker/pytorch/defaults.pyi | 1 + .../sagemaker/sagemaker/pytorch/estimator.pyi | 40 ++ stubs/sagemaker/sagemaker/pytorch/model.pyi | 69 +++ .../sagemaker/pytorch/processing.pyi | 30 + .../pytorch/training_compiler/__init__.pyi | 0 .../pytorch/training_compiler/config.pyi | 14 + .../sagemaker/remote_function/__init__.pyi | 1 + .../sagemaker/remote_function/client.pyi | 101 +++ .../sagemaker/remote_function/errors.pyi | 14 + .../remote_function/invoke_function.pyi | 3 + .../sagemaker/remote_function/job.pyi | 102 +++ .../remote_function/logging_config.pyi | 7 + .../runtime_environment/__init__.pyi | 0 .../bootstrap_runtime_environment.pyi | 12 + .../runtime_environment_manager.pyi | 19 + .../runtime_environment/spark_app.pyi | 0 .../remote_function/spark_config.pyi | 13 + stubs/sagemaker/sagemaker/rl/__init__.pyi | 6 + stubs/sagemaker/sagemaker/rl/estimator.pyi | 55 ++ stubs/sagemaker/sagemaker/s3.pyi | 32 + stubs/sagemaker/sagemaker/s3_utils.pyi | 10 + stubs/sagemaker/sagemaker/script_uris.pyi | 13 + stubs/sagemaker/sagemaker/serializers.pyi | 29 + .../sagemaker/serverless/__init__.pyi | 3 + .../sagemaker/sagemaker/serverless/model.pyi | 1 + .../sagemaker/serverless/predictor.pyi | 1 + .../serverless_inference_config.pyi | 10 + stubs/sagemaker/sagemaker/session.pyi | 585 ++++++++++++++++++ .../sagemaker/sagemaker/session_settings.pyi | 8 + .../sagemaker/sagemaker/sklearn/__init__.pyi | 3 + .../sagemaker/sagemaker/sklearn/defaults.pyi | 1 + .../sagemaker/sagemaker/sklearn/estimator.pyi | 33 + stubs/sagemaker/sagemaker/sklearn/model.pyi | 63 ++ .../sagemaker/sklearn/processing.pyi | 25 + stubs/sagemaker/sagemaker/spark/__init__.pyi | 1 + stubs/sagemaker/sagemaker/spark/defaults.pyi | 1 + .../sagemaker/sagemaker/spark/processing.pyi | 182 ++++++ .../sagemaker/sagemaker/sparkml/__init__.pyi | 1 + stubs/sagemaker/sagemaker/sparkml/model.pyi | 20 + .../sagemaker/tensorflow/__init__.pyi | 4 + .../sagemaker/tensorflow/defaults.pyi | 1 + .../sagemaker/tensorflow/estimator.pyi | 56 ++ .../sagemaker/sagemaker/tensorflow/model.pyi | 103 +++ .../sagemaker/tensorflow/processing.pyi | 30 + .../sagemaker/tensorflow/serving.pyi | 4 + .../tensorflow/training_compiler/__init__.pyi | 0 .../tensorflow/training_compiler/config.pyi | 13 + .../sagemaker/training_compiler/__init__.pyi | 0 .../sagemaker/training_compiler/config.pyi | 16 + stubs/sagemaker/sagemaker/transformer.pyi | 108 ++++ stubs/sagemaker/sagemaker/tuner.pyi | 223 +++++++ stubs/sagemaker/sagemaker/user_agent.pyi | 10 + .../sagemaker/utilities/__init__.pyi | 0 stubs/sagemaker/sagemaker/utilities/cache.pyi | 24 + .../sagemaker/utilities/search_expression.pyi | 47 ++ stubs/sagemaker/sagemaker/utils.pyi | 113 ++++ stubs/sagemaker/sagemaker/vpc_utils.pyi | 8 + .../sagemaker/sagemaker/workflow/__init__.pyi | 2 + .../sagemaker/workflow/_repack_model.pyi | 5 + stubs/sagemaker/sagemaker/workflow/_utils.pyi | 97 +++ .../sagemaker/sagemaker/workflow/airflow.pyi | 102 +++ .../sagemaker/workflow/automl_step.pyi | 28 + .../sagemaker/workflow/callback_step.pyi | 46 ++ .../sagemaker/workflow/check_job_config.pyi | 30 + .../sagemaker/workflow/clarify_check_step.pyi | 90 +++ .../sagemaker/workflow/condition_step.pyi | 33 + .../sagemaker/workflow/conditions.pyi | 92 +++ .../sagemaker/sagemaker/workflow/emr_step.pyi | 46 ++ .../sagemaker/sagemaker/workflow/entities.pyi | 30 + .../workflow/execution_variables.pyi | 21 + .../sagemaker/workflow/fail_step.pyi | 21 + .../sagemaker/workflow/functions.pyi | 28 + .../sagemaker/workflow/lambda_step.pyi | 47 ++ .../sagemaker/workflow/model_step.pyi | 25 + .../workflow/monitor_batch_transform_step.pyi | 26 + .../sagemaker/workflow/parallelism_config.pyi | 8 + .../sagemaker/workflow/parameters.pyi | 47 ++ .../sagemaker/sagemaker/workflow/pipeline.pyi | 91 +++ .../sagemaker/workflow/pipeline_context.pyi | 60 ++ .../workflow/pipeline_experiment_config.pyi | 26 + .../sagemaker/workflow/properties.pyi | 46 ++ .../sagemaker/workflow/quality_check_step.pyi | 84 +++ stubs/sagemaker/sagemaker/workflow/retry.pyi | 62 ++ .../workflow/selective_execution_config.pyi | 10 + .../sagemaker/workflow/step_collections.pyi | 92 +++ stubs/sagemaker/sagemaker/workflow/steps.pyi | 210 +++++++ .../sagemaker/workflow/utilities.pyi | 23 + .../sagemaker/sagemaker/wrangler/__init__.pyi | 0 .../sagemaker/wrangler/ingestion.pyi | 18 + .../sagemaker/wrangler/processing.pyi | 26 + .../sagemaker/sagemaker/xgboost/__init__.pyi | 4 + .../sagemaker/sagemaker/xgboost/defaults.pyi | 4 + .../sagemaker/sagemaker/xgboost/estimator.pyi | 35 ++ stubs/sagemaker/sagemaker/xgboost/model.pyi | 63 ++ .../sagemaker/xgboost/processing.pyi | 30 + stubs/sagemaker/sagemaker/xgboost/utils.pyi | 2 + 304 files changed, 13311 insertions(+) create mode 100644 stubs/sagemaker/METADATA.toml create mode 100644 stubs/sagemaker/sagemaker/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/_studio.pyi create mode 100644 stubs/sagemaker/sagemaker/accept_types.pyi create mode 100644 stubs/sagemaker/sagemaker/algorithm.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/common.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/hyperparameter.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/ipinsights.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/kmeans.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/knn.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/lda.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/linear_learner.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/ntm.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/object2vec.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/pca.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/record_pb2.pyi create mode 100644 stubs/sagemaker/sagemaker/amazon/validation.pyi create mode 100644 stubs/sagemaker/sagemaker/analytics.pyi create mode 100644 stubs/sagemaker/sagemaker/apiutils/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/apiutils/_base_types.pyi create mode 100644 stubs/sagemaker/sagemaker/apiutils/_boto_functions.pyi create mode 100644 stubs/sagemaker/sagemaker/apiutils/_utils.pyi create mode 100644 stubs/sagemaker/sagemaker/async_inference/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/async_inference/async_inference_config.pyi create mode 100644 stubs/sagemaker/sagemaker/async_inference/async_inference_response.pyi create mode 100644 stubs/sagemaker/sagemaker/async_inference/waiter_config.pyi create mode 100644 stubs/sagemaker/sagemaker/automl/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/automl/automl.pyi create mode 100644 stubs/sagemaker/sagemaker/automl/candidate_estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/base_deserializers.pyi create mode 100644 stubs/sagemaker/sagemaker/base_predictor.pyi create mode 100644 stubs/sagemaker/sagemaker/base_serializers.pyi create mode 100644 stubs/sagemaker/sagemaker/chainer/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/chainer/defaults.pyi create mode 100644 stubs/sagemaker/sagemaker/chainer/estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/chainer/model.pyi create mode 100644 stubs/sagemaker/sagemaker/clarify.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/ast_transformer.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/files.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/airflow.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/deprecated_params.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/framework_version.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/image_uris.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/matching.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/modifier.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/parsing.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/predictors.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/renamed_params.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/serde.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/tf_legacy_mode.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/tfs.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/training_input.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/training_params.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/compatibility/v2/sagemaker_upgrade_v2.pyi create mode 100644 stubs/sagemaker/sagemaker/cli/framework_upgrade.pyi create mode 100644 stubs/sagemaker/sagemaker/collection.pyi create mode 100644 stubs/sagemaker/sagemaker/config/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/config/config.pyi create mode 100644 stubs/sagemaker/sagemaker/config/config_schema.pyi create mode 100644 stubs/sagemaker/sagemaker/config/config_utils.pyi create mode 100644 stubs/sagemaker/sagemaker/content_types.pyi create mode 100644 stubs/sagemaker/sagemaker/dataset_definition/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/dataset_definition/inputs.pyi create mode 100644 stubs/sagemaker/sagemaker/debugger/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/debugger/debugger.pyi create mode 100644 stubs/sagemaker/sagemaker/debugger/framework_profile.pyi create mode 100644 stubs/sagemaker/sagemaker/debugger/metrics_config.pyi create mode 100644 stubs/sagemaker/sagemaker/debugger/profiler_config.pyi create mode 100644 stubs/sagemaker/sagemaker/debugger/profiler_constants.pyi create mode 100644 stubs/sagemaker/sagemaker/debugger/utils.pyi create mode 100644 stubs/sagemaker/sagemaker/deprecations.pyi create mode 100644 stubs/sagemaker/sagemaker/deserializers.pyi create mode 100644 stubs/sagemaker/sagemaker/djl_inference/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/djl_inference/defaults.pyi create mode 100644 stubs/sagemaker/sagemaker/djl_inference/model.pyi create mode 100644 stubs/sagemaker/sagemaker/drift_check_baselines.pyi create mode 100644 stubs/sagemaker/sagemaker/environment_variables.pyi create mode 100644 stubs/sagemaker/sagemaker/estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/exceptions.pyi create mode 100644 stubs/sagemaker/sagemaker/experiments/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/experiments/_api_types.pyi create mode 100644 stubs/sagemaker/sagemaker/experiments/_environment.pyi create mode 100644 stubs/sagemaker/sagemaker/experiments/_helper.pyi create mode 100644 stubs/sagemaker/sagemaker/experiments/_metrics.pyi create mode 100644 stubs/sagemaker/sagemaker/experiments/_run_context.pyi create mode 100644 stubs/sagemaker/sagemaker/experiments/_utils.pyi create mode 100644 stubs/sagemaker/sagemaker/experiments/experiment.pyi create mode 100644 stubs/sagemaker/sagemaker/experiments/run.pyi create mode 100644 stubs/sagemaker/sagemaker/experiments/trial.pyi create mode 100644 stubs/sagemaker/sagemaker/experiments/trial_component.pyi create mode 100644 stubs/sagemaker/sagemaker/explainer/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi create mode 100644 stubs/sagemaker/sagemaker/explainer/explainer_config.pyi create mode 100644 stubs/sagemaker/sagemaker/feature_store/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi create mode 100644 stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi create mode 100644 stubs/sagemaker/sagemaker/feature_store/feature_group.pyi create mode 100644 stubs/sagemaker/sagemaker/feature_store/feature_processor/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/feature_store/feature_store.pyi create mode 100644 stubs/sagemaker/sagemaker/feature_store/inputs.pyi create mode 100644 stubs/sagemaker/sagemaker/fw_utils.pyi create mode 100644 stubs/sagemaker/sagemaker/git_utils.pyi create mode 100644 stubs/sagemaker/sagemaker/huggingface/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/huggingface/estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi create mode 100644 stubs/sagemaker/sagemaker/huggingface/model.pyi create mode 100644 stubs/sagemaker/sagemaker/huggingface/processing.pyi create mode 100644 stubs/sagemaker/sagemaker/huggingface/training_compiler/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi create mode 100644 stubs/sagemaker/sagemaker/hyperparameters.pyi create mode 100644 stubs/sagemaker/sagemaker/image_uris.pyi create mode 100644 stubs/sagemaker/sagemaker/inference_recommender/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi create mode 100644 stubs/sagemaker/sagemaker/inputs.pyi create mode 100644 stubs/sagemaker/sagemaker/instance_group.pyi create mode 100644 stubs/sagemaker/sagemaker/instance_types.pyi create mode 100644 stubs/sagemaker/sagemaker/job.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/accessors.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/environment_variables.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/hyperparameters.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/image_uris.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/incremental_training.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/instance_types.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/kwargs.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/metric_definitions.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/model_uris.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/predictors.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/prepack.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/resource_names.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/artifacts/script_uris.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/cache.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/constants.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/enums.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/factory/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/filters.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/model.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/parameters.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/types.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/utils.pyi create mode 100644 stubs/sagemaker/sagemaker/jumpstart/validators.pyi create mode 100644 stubs/sagemaker/sagemaker/lambda_helper.pyi create mode 100644 stubs/sagemaker/sagemaker/lineage/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/lineage/_api_types.pyi create mode 100644 stubs/sagemaker/sagemaker/lineage/_utils.pyi create mode 100644 stubs/sagemaker/sagemaker/lineage/action.pyi create mode 100644 stubs/sagemaker/sagemaker/lineage/artifact.pyi create mode 100644 stubs/sagemaker/sagemaker/lineage/association.pyi create mode 100644 stubs/sagemaker/sagemaker/lineage/context.pyi create mode 100644 stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi create mode 100644 stubs/sagemaker/sagemaker/lineage/query.pyi create mode 100644 stubs/sagemaker/sagemaker/lineage/visualizer.pyi create mode 100644 stubs/sagemaker/sagemaker/local/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/local/data.pyi create mode 100644 stubs/sagemaker/sagemaker/local/entities.pyi create mode 100644 stubs/sagemaker/sagemaker/local/exceptions.pyi create mode 100644 stubs/sagemaker/sagemaker/local/image.pyi create mode 100644 stubs/sagemaker/sagemaker/local/local_session.pyi create mode 100644 stubs/sagemaker/sagemaker/local/pipeline.pyi create mode 100644 stubs/sagemaker/sagemaker/local/utils.pyi create mode 100644 stubs/sagemaker/sagemaker/logs.pyi create mode 100644 stubs/sagemaker/sagemaker/metadata_properties.pyi create mode 100644 stubs/sagemaker/sagemaker/metric_definitions.pyi create mode 100644 stubs/sagemaker/sagemaker/model.pyi create mode 100644 stubs/sagemaker/sagemaker/model_card/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/model_card/evaluation_metric_parsers.pyi create mode 100644 stubs/sagemaker/sagemaker/model_card/helpers.pyi create mode 100644 stubs/sagemaker/sagemaker/model_card/model_card.pyi create mode 100644 stubs/sagemaker/sagemaker/model_card/schema_constraints.pyi create mode 100644 stubs/sagemaker/sagemaker/model_metrics.pyi create mode 100644 stubs/sagemaker/sagemaker/model_monitor/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/model_monitor/clarify_model_monitoring.pyi create mode 100644 stubs/sagemaker/sagemaker/model_monitor/cron_expression_generator.pyi create mode 100644 stubs/sagemaker/sagemaker/model_monitor/data_capture_config.pyi create mode 100644 stubs/sagemaker/sagemaker/model_monitor/data_quality_monitoring_config.pyi create mode 100644 stubs/sagemaker/sagemaker/model_monitor/dataset_format.pyi create mode 100644 stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi create mode 100644 stubs/sagemaker/sagemaker/model_monitor/monitoring_alert.pyi create mode 100644 stubs/sagemaker/sagemaker/model_monitor/monitoring_files.pyi create mode 100644 stubs/sagemaker/sagemaker/model_uris.pyi create mode 100644 stubs/sagemaker/sagemaker/multidatamodel.pyi create mode 100644 stubs/sagemaker/sagemaker/mxnet/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/mxnet/defaults.pyi create mode 100644 stubs/sagemaker/sagemaker/mxnet/estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/mxnet/model.pyi create mode 100644 stubs/sagemaker/sagemaker/mxnet/processing.pyi create mode 100644 stubs/sagemaker/sagemaker/network.pyi create mode 100644 stubs/sagemaker/sagemaker/parameter.pyi create mode 100644 stubs/sagemaker/sagemaker/pipeline.pyi create mode 100644 stubs/sagemaker/sagemaker/predictor.pyi create mode 100644 stubs/sagemaker/sagemaker/predictor_async.pyi create mode 100644 stubs/sagemaker/sagemaker/processing.pyi create mode 100644 stubs/sagemaker/sagemaker/pytorch/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/pytorch/defaults.pyi create mode 100644 stubs/sagemaker/sagemaker/pytorch/estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/pytorch/model.pyi create mode 100644 stubs/sagemaker/sagemaker/pytorch/processing.pyi create mode 100644 stubs/sagemaker/sagemaker/pytorch/training_compiler/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi create mode 100644 stubs/sagemaker/sagemaker/remote_function/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/remote_function/client.pyi create mode 100644 stubs/sagemaker/sagemaker/remote_function/errors.pyi create mode 100644 stubs/sagemaker/sagemaker/remote_function/invoke_function.pyi create mode 100644 stubs/sagemaker/sagemaker/remote_function/job.pyi create mode 100644 stubs/sagemaker/sagemaker/remote_function/logging_config.pyi create mode 100644 stubs/sagemaker/sagemaker/remote_function/runtime_environment/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.pyi create mode 100644 stubs/sagemaker/sagemaker/remote_function/runtime_environment/runtime_environment_manager.pyi create mode 100644 stubs/sagemaker/sagemaker/remote_function/runtime_environment/spark_app.pyi create mode 100644 stubs/sagemaker/sagemaker/remote_function/spark_config.pyi create mode 100644 stubs/sagemaker/sagemaker/rl/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/rl/estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/s3.pyi create mode 100644 stubs/sagemaker/sagemaker/s3_utils.pyi create mode 100644 stubs/sagemaker/sagemaker/script_uris.pyi create mode 100644 stubs/sagemaker/sagemaker/serializers.pyi create mode 100644 stubs/sagemaker/sagemaker/serverless/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/serverless/model.pyi create mode 100644 stubs/sagemaker/sagemaker/serverless/predictor.pyi create mode 100644 stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi create mode 100644 stubs/sagemaker/sagemaker/session.pyi create mode 100644 stubs/sagemaker/sagemaker/session_settings.pyi create mode 100644 stubs/sagemaker/sagemaker/sklearn/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/sklearn/defaults.pyi create mode 100644 stubs/sagemaker/sagemaker/sklearn/estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/sklearn/model.pyi create mode 100644 stubs/sagemaker/sagemaker/sklearn/processing.pyi create mode 100644 stubs/sagemaker/sagemaker/spark/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/spark/defaults.pyi create mode 100644 stubs/sagemaker/sagemaker/spark/processing.pyi create mode 100644 stubs/sagemaker/sagemaker/sparkml/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/sparkml/model.pyi create mode 100644 stubs/sagemaker/sagemaker/tensorflow/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/tensorflow/defaults.pyi create mode 100644 stubs/sagemaker/sagemaker/tensorflow/estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/tensorflow/model.pyi create mode 100644 stubs/sagemaker/sagemaker/tensorflow/processing.pyi create mode 100644 stubs/sagemaker/sagemaker/tensorflow/serving.pyi create mode 100644 stubs/sagemaker/sagemaker/tensorflow/training_compiler/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/tensorflow/training_compiler/config.pyi create mode 100644 stubs/sagemaker/sagemaker/training_compiler/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/training_compiler/config.pyi create mode 100644 stubs/sagemaker/sagemaker/transformer.pyi create mode 100644 stubs/sagemaker/sagemaker/tuner.pyi create mode 100644 stubs/sagemaker/sagemaker/user_agent.pyi create mode 100644 stubs/sagemaker/sagemaker/utilities/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/utilities/cache.pyi create mode 100644 stubs/sagemaker/sagemaker/utilities/search_expression.pyi create mode 100644 stubs/sagemaker/sagemaker/utils.pyi create mode 100644 stubs/sagemaker/sagemaker/vpc_utils.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/_repack_model.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/_utils.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/airflow.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/automl_step.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/callback_step.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/check_job_config.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/condition_step.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/conditions.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/emr_step.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/entities.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/execution_variables.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/fail_step.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/functions.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/lambda_step.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/model_step.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/parallelism_config.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/parameters.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/pipeline.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/properties.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/retry.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/step_collections.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/steps.pyi create mode 100644 stubs/sagemaker/sagemaker/workflow/utilities.pyi create mode 100644 stubs/sagemaker/sagemaker/wrangler/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/wrangler/ingestion.pyi create mode 100644 stubs/sagemaker/sagemaker/wrangler/processing.pyi create mode 100644 stubs/sagemaker/sagemaker/xgboost/__init__.pyi create mode 100644 stubs/sagemaker/sagemaker/xgboost/defaults.pyi create mode 100644 stubs/sagemaker/sagemaker/xgboost/estimator.pyi create mode 100644 stubs/sagemaker/sagemaker/xgboost/model.pyi create mode 100644 stubs/sagemaker/sagemaker/xgboost/processing.pyi create mode 100644 stubs/sagemaker/sagemaker/xgboost/utils.pyi diff --git a/pyrightconfig.stricter.json b/pyrightconfig.stricter.json index 2fe331589d39..b6b27aae3c62 100644 --- a/pyrightconfig.stricter.json +++ b/pyrightconfig.stricter.json @@ -69,6 +69,7 @@ "stubs/PyYAML", "stubs/redis", "stubs/requests", + "stubs/sagemaker", "stubs/setuptools", "stubs/SQLAlchemy", "stubs/stripe", diff --git a/stubs/sagemaker/METADATA.toml b/stubs/sagemaker/METADATA.toml new file mode 100644 index 000000000000..4b4e7f8b507a --- /dev/null +++ b/stubs/sagemaker/METADATA.toml @@ -0,0 +1 @@ +version = "2.168.*" \ No newline at end of file diff --git a/stubs/sagemaker/sagemaker/__init__.pyi b/stubs/sagemaker/sagemaker/__init__.pyi new file mode 100644 index 000000000000..350dbc228d8c --- /dev/null +++ b/stubs/sagemaker/sagemaker/__init__.pyi @@ -0,0 +1,49 @@ +from sagemaker import estimator as estimator, parameter as parameter, tuner as tuner +from sagemaker.algorithm import AlgorithmEstimator as AlgorithmEstimator +from sagemaker.amazon.factorization_machines import ( + FactorizationMachines as FactorizationMachines, + FactorizationMachinesModel as FactorizationMachinesModel, + FactorizationMachinesPredictor as FactorizationMachinesPredictor, +) +from sagemaker.amazon.ipinsights import ( + IPInsights as IPInsights, + IPInsightsModel as IPInsightsModel, + IPInsightsPredictor as IPInsightsPredictor, +) +from sagemaker.amazon.kmeans import KMeans as KMeans, KMeansModel as KMeansModel, KMeansPredictor as KMeansPredictor +from sagemaker.amazon.knn import KNN as KNN, KNNModel as KNNModel, KNNPredictor as KNNPredictor +from sagemaker.amazon.lda import LDA as LDA, LDAModel as LDAModel, LDAPredictor as LDAPredictor +from sagemaker.amazon.linear_learner import ( + LinearLearner as LinearLearner, + LinearLearnerModel as LinearLearnerModel, + LinearLearnerPredictor as LinearLearnerPredictor, +) +from sagemaker.amazon.ntm import NTM as NTM, NTMModel as NTMModel, NTMPredictor as NTMPredictor +from sagemaker.amazon.object2vec import Object2Vec as Object2Vec, Object2VecModel as Object2VecModel +from sagemaker.amazon.pca import PCA as PCA, PCAModel as PCAModel, PCAPredictor as PCAPredictor +from sagemaker.amazon.randomcutforest import ( + RandomCutForest as RandomCutForest, + RandomCutForestModel as RandomCutForestModel, + RandomCutForestPredictor as RandomCutForestPredictor, +) +from sagemaker.analytics import ( + HyperparameterTuningJobAnalytics as HyperparameterTuningJobAnalytics, + TrainingJobAnalytics as TrainingJobAnalytics, +) +from sagemaker.automl.automl import AutoML as AutoML, AutoMLInput as AutoMLInput, AutoMLJob as AutoMLJob +from sagemaker.automl.candidate_estimator import CandidateEstimator as CandidateEstimator, CandidateStep as CandidateStep +from sagemaker.inputs import TrainingInput as TrainingInput +from sagemaker.local.local_session import LocalSession as LocalSession +from sagemaker.model import Model as Model, ModelPackage as ModelPackage +from sagemaker.model_metrics import FileSource as FileSource, MetricsSource as MetricsSource, ModelMetrics as ModelMetrics +from sagemaker.pipeline import PipelineModel as PipelineModel +from sagemaker.predictor import Predictor as Predictor +from sagemaker.processing import Processor as Processor, ScriptProcessor as ScriptProcessor +from sagemaker.session import ( + Session as Session, + container_def as container_def, + get_execution_role as get_execution_role, + get_model_package_args as get_model_package_args, + pipeline_container_def as pipeline_container_def, + production_variant as production_variant, +) diff --git a/stubs/sagemaker/sagemaker/_studio.pyi b/stubs/sagemaker/sagemaker/_studio.pyi new file mode 100644 index 000000000000..4596620a7650 --- /dev/null +++ b/stubs/sagemaker/sagemaker/_studio.pyi @@ -0,0 +1,4 @@ +from _typeshed import Incomplete + +STUDIO_PROJECT_CONFIG: str +logger: Incomplete diff --git a/stubs/sagemaker/sagemaker/accept_types.pyi b/stubs/sagemaker/sagemaker/accept_types.pyi new file mode 100644 index 000000000000..9f6a7a2b80fe --- /dev/null +++ b/stubs/sagemaker/sagemaker/accept_types.pyi @@ -0,0 +1,16 @@ +from typing import List, Optional + +def retrieve_options( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> List[str]: ... +def retrieve_default( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> str: ... diff --git a/stubs/sagemaker/sagemaker/algorithm.pyi b/stubs/sagemaker/sagemaker/algorithm.pyi new file mode 100644 index 000000000000..602bbf46bd23 --- /dev/null +++ b/stubs/sagemaker/sagemaker/algorithm.pyi @@ -0,0 +1,75 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker.estimator import EstimatorBase +from sagemaker.inputs import FileSystemInput, TrainingInput +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class AlgorithmEstimator(EstimatorBase): + algorithm_arn: Incomplete + algorithm_spec: Incomplete + hyperparameter_definitions: Incomplete + def __init__( + self, + algorithm_arn: str, + role: str = None, + 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] = 86400, + 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, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = 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, + **kwargs, + ) -> None: ... + def validate_train_spec(self) -> None: ... + def set_hyperparameters(self, **kwargs) -> None: ... + def hyperparameters(self): ... + def training_image_uri(self) -> None: ... + def enable_network_isolation(self): ... + def create_model( + self, + role: Incomplete | None = None, + predictor_cls: Incomplete | None = None, + serializer=..., + deserializer=..., + vpc_config_override="VPC_CONFIG_DEFAULT", + **kwargs, + ): ... + def transformer( + self, + instance_count, + instance_type, + strategy: Incomplete | None = None, + assemble_with: Incomplete | None = None, + output_path: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + accept: Incomplete | None = None, + env: Incomplete | None = None, + max_concurrent_transforms: Incomplete | None = None, + max_payload: Incomplete | None = None, + tags: Incomplete | None = None, + role: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + ): ... + def fit( + self, + inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None, + wait: bool = True, + logs: bool = True, + job_name: Optional[str] = None, + ): ... diff --git a/stubs/sagemaker/sagemaker/amazon/__init__.pyi b/stubs/sagemaker/sagemaker/amazon/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi b/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi new file mode 100644 index 000000000000..77fa3bcf2c54 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi @@ -0,0 +1,84 @@ +import abc +from _typeshed import Incomplete +from typing import Dict, Optional, Union + +from sagemaker.amazon.hyperparameter import Hyperparameter as hp +from sagemaker.estimator import EstimatorBase +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class AmazonAlgorithmEstimatorBase(EstimatorBase, metaclass=abc.ABCMeta): + feature_dim: hp + mini_batch_size: hp + repo_name: Optional[str] + repo_version: Optional[str] + DEFAULT_MINI_BATCH_SIZE: Optional[int] + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Optional[Union[int, PipelineVariable]] = None, + instance_type: Optional[Union[str, PipelineVariable]] = None, + data_location: Optional[str] = None, + enable_network_isolation: Union[bool, PipelineVariable] = False, + **kwargs, + ) -> None: ... + def training_image_uri(self): ... + def hyperparameters(self): ... + @property + def data_location(self): ... + @data_location.setter + def data_location(self, data_location: str): ... + def prepare_workflow_for_training( + self, records: Incomplete | None = None, mini_batch_size: Incomplete | None = None, job_name: Incomplete | None = None + ) -> None: ... + latest_training_job: Incomplete + def fit( + self, + 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, + ): ... + def record_set(self, train, labels: Incomplete | None = None, channel: str = "train", encrypt: bool = False): ... + +class RecordSet: + s3_data: Incomplete + feature_dim: Incomplete + num_records: Incomplete + s3_data_type: Incomplete + channel: Incomplete + def __init__( + self, + s3_data: Union[str, PipelineVariable], + num_records: int, + feature_dim: int, + s3_data_type: Union[str, PipelineVariable] = "ManifestFile", + channel: Union[str, PipelineVariable] = "train", + ) -> None: ... + def data_channel(self): ... + def records_s3_input(self): ... + +class FileSystemRecordSet: + file_system_input: Incomplete + feature_dim: Incomplete + num_records: Incomplete + channel: Incomplete + def __init__( + self, + file_system_id, + file_system_type, + directory_path, + num_records, + feature_dim, + file_system_access_mode: str = "ro", + channel: str = "train", + ) -> None: ... + def data_channel(self): ... + +def upload_numpy_to_s3_shards( + num_shards, s3, bucket, key_prefix, array, labels: Incomplete | None = None, encrypt: bool = False +): ... +def get_image_uri(region_name, repo_name, repo_version: str = "1"): ... diff --git a/stubs/sagemaker/sagemaker/amazon/common.pyi b/stubs/sagemaker/sagemaker/amazon/common.pyi new file mode 100644 index 000000000000..b754cda8b8d9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/common.pyi @@ -0,0 +1,24 @@ +from _typeshed import Incomplete +from collections.abc import Generator + +from sagemaker.deserializers import SimpleBaseDeserializer +from sagemaker.serializers import SimpleBaseSerializer + +class RecordSerializer(SimpleBaseSerializer): + def __init__(self, content_type: str = "application/x-recordio-protobuf") -> None: ... + def serialize(self, data): ... + +class RecordDeserializer(SimpleBaseDeserializer): + def __init__(self, accept: str = "application/x-recordio-protobuf") -> None: ... + def deserialize(self, data, content_type): ... + +def write_numpy_to_dense_tensor(file, array, labels: Incomplete | None = None) -> None: ... +def write_spmatrix_to_sparse_tensor(file, array, labels: Incomplete | None = None) -> None: ... +def read_records(file): ... + +padding: Incomplete + +def read_recordio(f) -> Generator[Incomplete, None, None]: ... + +numpy_to_record_serializer: Incomplete +record_deserializer: Incomplete diff --git a/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi b/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi new file mode 100644 index 000000000000..3d70a2034c8b --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi @@ -0,0 +1,81 @@ +from _typeshed import Incomplete +from typing import Optional, Union + +from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase +from sagemaker.amazon.hyperparameter import Hyperparameter as hp +from sagemaker.model import Model +from sagemaker.predictor import Predictor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class FactorizationMachines(AmazonAlgorithmEstimatorBase): + repo_name: str + repo_version: str + num_factors: hp + predictor_type: hp + epochs: hp + clip_gradient: hp + eps: hp + rescale_grad: hp + bias_lr: hp + linear_lr: hp + factors_lr: hp + bias_wd: hp + linear_wd: hp + factors_wd: hp + bias_init_method: hp + bias_init_scale: hp + bias_init_sigma: hp + bias_init_value: hp + linear_init_method: hp + linear_init_scale: hp + linear_init_sigma: hp + linear_init_value: hp + factors_init_method: hp + factors_init_scale: hp + factors_init_sigma: hp + factors_init_value: hp + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + 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, + ) -> None: ... + def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... + +class FactorizationMachinesPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class FactorizationMachinesModel(Model): + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/hyperparameter.pyi b/stubs/sagemaker/sagemaker/amazon/hyperparameter.pyi new file mode 100644 index 000000000000..5886f39c6e49 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/hyperparameter.pyi @@ -0,0 +1,14 @@ +from _typeshed import Incomplete + +class Hyperparameter: + validation: Incomplete + validation_message: Incomplete + name: Incomplete + data_type: Incomplete + def __init__(self, name, validate=..., validation_message: str = "", data_type=...) -> None: ... + def validate(self, value) -> None: ... + def __get__(self, obj, objtype): ... + def __set__(self, obj, value) -> None: ... + def __delete__(self, obj) -> None: ... + @staticmethod + def serialize_all(obj): ... diff --git a/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi b/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi new file mode 100644 index 000000000000..f893b5c38b13 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi @@ -0,0 +1,52 @@ +from _typeshed import Incomplete +from typing import Optional, Union + +from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase +from sagemaker.amazon.hyperparameter import Hyperparameter as hp +from sagemaker.model import Model +from sagemaker.predictor import Predictor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class IPInsights(AmazonAlgorithmEstimatorBase): + repo_name: str + repo_version: str + MINI_BATCH_SIZE: int + num_entity_vectors: hp + vector_dim: hp + batch_metrics_publish_interval: hp + epochs: hp + learning_rate: hp + num_ip_encoder_layers: hp + random_negative_sampling_rate: hp + shuffled_negative_sampling_rate: hp + weight_decay: hp + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + 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, + ) -> None: ... + def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... + +class IPInsightsPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class IPInsightsModel(Model): + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/kmeans.pyi b/stubs/sagemaker/sagemaker/amazon/kmeans.pyi new file mode 100644 index 000000000000..121cfa8088f8 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/kmeans.pyi @@ -0,0 +1,54 @@ +from _typeshed import Incomplete +from typing import List, Optional, Union + +from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase +from sagemaker.amazon.hyperparameter import Hyperparameter as hp +from sagemaker.model import Model +from sagemaker.predictor import Predictor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class KMeans(AmazonAlgorithmEstimatorBase): + repo_name: str + repo_version: str + k: hp + init_method: hp + max_iterations: hp + tol: hp + num_trials: hp + local_init_method: hp + half_life_time_size: hp + epochs: hp + center_factor: hp + eval_metrics: hp + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + 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, + ) -> None: ... + def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... + def hyperparameters(self): ... + +class KMeansPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class KMeansModel(Model): + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/knn.pyi b/stubs/sagemaker/sagemaker/amazon/knn.pyi new file mode 100644 index 000000000000..5db619fd0e0f --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/knn.pyi @@ -0,0 +1,51 @@ +from _typeshed import Incomplete +from typing import Optional, Union + +from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase +from sagemaker.amazon.hyperparameter import Hyperparameter as hp +from sagemaker.model import Model +from sagemaker.predictor import Predictor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class KNN(AmazonAlgorithmEstimatorBase): + repo_name: str + repo_version: str + k: hp + sample_size: hp + predictor_type: hp + dimension_reduction_target: hp + dimension_reduction_type: hp + index_metric: hp + index_type: hp + faiss_index_ivf_nlists: hp + faiss_index_pq_m: hp + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Optional[Union[int, PipelineVariable]] = None, + instance_type: Optional[Union[str, PipelineVariable]] = None, + k: Optional[int] = None, + sample_size: Optional[int] = None, + predictor_type: Optional[str] = None, + dimension_reduction_type: Optional[str] = None, + dimension_reduction_target: Optional[int] = None, + index_type: Optional[str] = None, + index_metric: Optional[str] = None, + faiss_index_ivf_nlists: Optional[str] = None, + faiss_index_pq_m: Optional[int] = None, + **kwargs, + ) -> None: ... + def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... + +class KNNPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class KNNModel(Model): + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/lda.pyi b/stubs/sagemaker/sagemaker/amazon/lda.pyi new file mode 100644 index 000000000000..49efbaf7a1db --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/lda.pyi @@ -0,0 +1,44 @@ +from _typeshed import Incomplete +from typing import Optional, Union + +from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase +from sagemaker.amazon.hyperparameter import Hyperparameter as hp +from sagemaker.model import Model +from sagemaker.predictor import Predictor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class LDA(AmazonAlgorithmEstimatorBase): + repo_name: str + repo_version: str + num_topics: hp + alpha0: hp + max_restarts: hp + max_iterations: hp + tol: hp + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + instance_type: Optional[Union[str, PipelineVariable]] = None, + num_topics: Optional[int] = None, + alpha0: Optional[float] = None, + max_restarts: Optional[int] = None, + max_iterations: Optional[int] = None, + tol: Optional[float] = None, + **kwargs, + ) -> None: ... + def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... + +class LDAPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class LDAModel(Model): + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi b/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi new file mode 100644 index 000000000000..867d6450a795 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi @@ -0,0 +1,120 @@ +from _typeshed import Incomplete +from typing import Optional, Union + +from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase +from sagemaker.amazon.hyperparameter import Hyperparameter as hp +from sagemaker.model import Model +from sagemaker.predictor import Predictor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class LinearLearner(AmazonAlgorithmEstimatorBase): + repo_name: str + repo_version: str + DEFAULT_MINI_BATCH_SIZE: int + binary_classifier_model_selection_criteria: hp + target_recall: hp + target_precision: hp + positive_example_weight_mult: hp + epochs: hp + predictor_type: hp + use_bias: hp + num_models: hp + num_calibration_samples: hp + init_method: hp + init_scale: hp + init_sigma: hp + init_bias: hp + optimizer: hp + loss: hp + wd: hp + l1: hp + momentum: hp + learning_rate: hp + beta_1: hp + beta_2: hp + bias_lr_mult: hp + bias_wd_mult: hp + use_lr_scheduler: hp + lr_scheduler_step: hp + lr_scheduler_factor: hp + lr_scheduler_minimum_lr: hp + normalize_data: hp + normalize_label: hp + unbias_data: hp + unbias_label: hp + num_point_for_scaler: hp + margin: hp + quantile: hp + loss_insensitivity: hp + huber_delta: hp + early_stopping_patience: hp + early_stopping_tolerance: hp + num_classes: hp + accuracy_top_k: hp + f_beta: hp + balance_multiclass_weights: hp + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Optional[Union[int, PipelineVariable]] = None, + instance_type: Optional[Union[str, PipelineVariable]] = None, + predictor_type: Optional[str] = None, + binary_classifier_model_selection_criteria: Optional[str] = None, + target_recall: Optional[float] = None, + target_precision: Optional[float] = None, + positive_example_weight_mult: Optional[float] = None, + epochs: Optional[int] = None, + use_bias: Optional[bool] = None, + num_models: Optional[int] = None, + num_calibration_samples: Optional[int] = None, + init_method: Optional[str] = None, + init_scale: Optional[float] = None, + init_sigma: Optional[float] = None, + init_bias: Optional[float] = None, + optimizer: Optional[str] = None, + loss: Optional[str] = None, + wd: Optional[float] = None, + l1: Optional[float] = None, + momentum: Optional[float] = None, + learning_rate: Optional[float] = None, + beta_1: Optional[float] = None, + beta_2: Optional[float] = None, + bias_lr_mult: Optional[float] = None, + bias_wd_mult: Optional[float] = None, + use_lr_scheduler: Optional[bool] = None, + lr_scheduler_step: Optional[int] = None, + lr_scheduler_factor: Optional[float] = None, + lr_scheduler_minimum_lr: Optional[float] = None, + normalize_data: Optional[bool] = None, + normalize_label: Optional[bool] = None, + unbias_data: Optional[bool] = None, + unbias_label: Optional[bool] = None, + num_point_for_scaler: Optional[int] = None, + margin: Optional[float] = None, + quantile: Optional[float] = None, + loss_insensitivity: Optional[float] = None, + huber_delta: Optional[float] = None, + early_stopping_patience: Optional[int] = None, + early_stopping_tolerance: Optional[float] = None, + num_classes: Optional[int] = None, + accuracy_top_k: Optional[int] = None, + f_beta: Optional[float] = None, + balance_multiclass_weights: Optional[bool] = None, + **kwargs, + ) -> None: ... + def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... + +class LinearLearnerPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class LinearLearnerModel(Model): + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/ntm.pyi b/stubs/sagemaker/sagemaker/amazon/ntm.pyi new file mode 100644 index 000000000000..4ca8a2ab3a5e --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/ntm.pyi @@ -0,0 +1,57 @@ +from _typeshed import Incomplete +from typing import List, Optional, Union + +from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase +from sagemaker.amazon.hyperparameter import Hyperparameter as hp +from sagemaker.model import Model +from sagemaker.predictor import Predictor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class NTM(AmazonAlgorithmEstimatorBase): + repo_name: str + repo_version: str + num_topics: hp + encoder_layers: hp + epochs: hp + encoder_layers_activation: hp + optimizer: hp + tolerance: hp + num_patience_epochs: hp + batch_norm: hp + rescale_gradient: hp + clip_gradient: hp + weight_decay: hp + learning_rate: hp + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Optional[Union[int, PipelineVariable]] = None, + instance_type: Optional[Union[str, PipelineVariable]] = None, + num_topics: Optional[int] = None, + encoder_layers: Optional[List] = None, + epochs: Optional[int] = None, + encoder_layers_activation: Optional[str] = None, + optimizer: Optional[str] = None, + tolerance: Optional[float] = None, + num_patience_epochs: Optional[int] = None, + batch_norm: Optional[bool] = None, + rescale_gradient: Optional[float] = None, + clip_gradient: Optional[float] = None, + weight_decay: Optional[float] = None, + learning_rate: Optional[float] = None, + **kwargs, + ) -> None: ... + def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... + +class NTMPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class NTMModel(Model): + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/object2vec.pyi b/stubs/sagemaker/sagemaker/amazon/object2vec.pyi new file mode 100644 index 000000000000..f537d79b6621 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/object2vec.pyi @@ -0,0 +1,95 @@ +from typing import Optional, Union + +from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase +from sagemaker.amazon.hyperparameter import Hyperparameter as hp +from sagemaker.model import Model +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class Object2Vec(AmazonAlgorithmEstimatorBase): + repo_name: str + repo_version: str + MINI_BATCH_SIZE: int + enc_dim: hp + mini_batch_size: hp + epochs: hp + early_stopping_patience: hp + early_stopping_tolerance: hp + dropout: hp + weight_decay: hp + bucket_width: hp + num_classes: hp + mlp_layers: hp + mlp_dim: hp + mlp_activation: hp + output_layer: hp + optimizer: hp + learning_rate: hp + negative_sampling_rate: hp + comparator_list: hp + tied_token_embedding_weight: hp + token_embedding_storage_type: hp + enc0_network: hp + enc1_network: hp + enc0_cnn_filter_width: hp + enc1_cnn_filter_width: hp + enc0_max_seq_len: hp + enc1_max_seq_len: hp + enc0_token_embedding_dim: hp + enc1_token_embedding_dim: hp + enc0_vocab_size: hp + enc1_vocab_size: hp + enc0_layers: hp + enc1_layers: hp + enc0_freeze_pretrained_embedding: hp + enc1_freeze_pretrained_embedding: hp + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Optional[Union[int, PipelineVariable]] = None, + instance_type: Optional[Union[str, PipelineVariable]] = None, + epochs: Optional[int] = None, + enc0_max_seq_len: Optional[int] = None, + enc0_vocab_size: Optional[int] = None, + enc_dim: Optional[int] = None, + mini_batch_size: Optional[int] = None, + early_stopping_patience: Optional[int] = None, + early_stopping_tolerance: Optional[float] = None, + dropout: Optional[float] = None, + weight_decay: Optional[float] = None, + bucket_width: Optional[int] = None, + num_classes: Optional[int] = None, + mlp_layers: Optional[int] = None, + mlp_dim: Optional[int] = None, + mlp_activation: Optional[str] = None, + output_layer: Optional[str] = None, + optimizer: Optional[str] = None, + learning_rate: Optional[float] = None, + negative_sampling_rate: Optional[int] = None, + comparator_list: Optional[str] = None, + tied_token_embedding_weight: Optional[bool] = None, + token_embedding_storage_type: Optional[str] = None, + enc0_network: Optional[str] = None, + enc1_network: Optional[str] = None, + enc0_cnn_filter_width: Optional[int] = None, + enc1_cnn_filter_width: Optional[int] = None, + enc1_max_seq_len: Optional[int] = None, + enc0_token_embedding_dim: Optional[int] = None, + enc1_token_embedding_dim: Optional[int] = None, + enc1_vocab_size: Optional[int] = None, + enc0_layers: Optional[int] = None, + enc1_layers: Optional[int] = None, + enc0_freeze_pretrained_embedding: Optional[bool] = None, + enc1_freeze_pretrained_embedding: Optional[bool] = None, + **kwargs, + ) -> None: ... + def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... + +class Object2VecModel(Model): + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/pca.pyi b/stubs/sagemaker/sagemaker/amazon/pca.pyi new file mode 100644 index 000000000000..af1a2fa8ce63 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/pca.pyi @@ -0,0 +1,42 @@ +from _typeshed import Incomplete +from typing import Optional, Union + +from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase +from sagemaker.amazon.hyperparameter import Hyperparameter as hp +from sagemaker.model import Model +from sagemaker.predictor import Predictor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class PCA(AmazonAlgorithmEstimatorBase): + repo_name: str + repo_version: str + DEFAULT_MINI_BATCH_SIZE: int + num_components: hp + algorithm_mode: hp + subtract_mean: hp + extra_components: hp + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Optional[Union[int, PipelineVariable]] = None, + instance_type: Optional[Union[str, PipelineVariable]] = None, + num_components: Optional[int] = None, + algorithm_mode: Optional[str] = None, + subtract_mean: Optional[bool] = None, + extra_components: Optional[int] = None, + **kwargs, + ) -> None: ... + def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... + +class PCAPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class PCAModel(Model): + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi b/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi new file mode 100644 index 000000000000..f646ff1eb145 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi @@ -0,0 +1,41 @@ +from _typeshed import Incomplete +from typing import List, Optional, Union + +from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase +from sagemaker.amazon.hyperparameter import Hyperparameter as hp +from sagemaker.model import Model +from sagemaker.predictor import Predictor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class RandomCutForest(AmazonAlgorithmEstimatorBase): + repo_name: str + repo_version: str + MINI_BATCH_SIZE: int + eval_metrics: hp + num_trees: hp + num_samples_per_tree: hp + feature_dim: hp + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Optional[Union[int, PipelineVariable]] = None, + instance_type: Optional[Union[str, PipelineVariable]] = None, + num_samples_per_tree: Optional[int] = None, + num_trees: Optional[int] = None, + eval_metrics: Optional[List] = None, + **kwargs, + ) -> None: ... + def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... + +class RandomCutForestPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class RandomCutForestModel(Model): + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/record_pb2.pyi b/stubs/sagemaker/sagemaker/amazon/record_pb2.pyi new file mode 100644 index 000000000000..41d320178f75 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/record_pb2.pyi @@ -0,0 +1,9 @@ +from _typeshed import Incomplete + +DESCRIPTOR: Incomplete +Float32Tensor: Incomplete +Float64Tensor: Incomplete +Int32Tensor: Incomplete +Bytes: Incomplete +Value: Incomplete +Record: Incomplete diff --git a/stubs/sagemaker/sagemaker/amazon/validation.pyi b/stubs/sagemaker/sagemaker/amazon/validation.pyi new file mode 100644 index 000000000000..a952c0e251b9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/validation.pyi @@ -0,0 +1,6 @@ +def gt(minimum): ... +def ge(minimum): ... +def lt(maximum): ... +def le(maximum): ... +def isin(*expected): ... +def istype(expected): ... diff --git a/stubs/sagemaker/sagemaker/analytics.pyi b/stubs/sagemaker/sagemaker/analytics.pyi new file mode 100644 index 000000000000..33505cc06a05 --- /dev/null +++ b/stubs/sagemaker/sagemaker/analytics.pyi @@ -0,0 +1,64 @@ +import abc +from _typeshed import Incomplete + +logger: Incomplete +METRICS_PERIOD_DEFAULT: int + +class AnalyticsMetricsBase(metaclass=abc.ABCMeta): + def __init__(self) -> None: ... + def export_csv(self, filename) -> None: ... + def dataframe(self, force_refresh: bool = False): ... + def clear_cache(self) -> None: ... + +class HyperparameterTuningJobAnalytics(AnalyticsMetricsBase): + def __init__(self, hyperparameter_tuning_job_name, sagemaker_session: Incomplete | None = None) -> None: ... + @property + def name(self): ... + def clear_cache(self) -> None: ... + @property + def tuning_ranges(self): ... + def description(self, force_refresh: bool = False): ... + def training_job_summaries(self, force_refresh: bool = False): ... + +class TrainingJobAnalytics(AnalyticsMetricsBase): + CLOUDWATCH_NAMESPACE: str + def __init__( + self, + training_job_name, + metric_names: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + start_time: Incomplete | None = None, + end_time: Incomplete | None = None, + period: Incomplete | None = None, + ) -> None: ... + @property + def name(self): ... + def clear_cache(self) -> None: ... + +class ArtifactAnalytics(AnalyticsMetricsBase): + def __init__( + self, + sort_by: Incomplete | None = None, + sort_order: Incomplete | None = None, + source_uri: Incomplete | None = None, + artifact_type: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + ) -> None: ... + +class ExperimentAnalytics(AnalyticsMetricsBase): + MAX_TRIAL_COMPONENTS: int + def __init__( + self, + experiment_name: Incomplete | None = None, + search_expression: Incomplete | None = None, + sort_by: Incomplete | None = None, + sort_order: Incomplete | None = None, + metric_names: Incomplete | None = None, + parameter_names: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + input_artifact_names: Incomplete | None = None, + output_artifact_names: Incomplete | None = None, + ) -> None: ... + @property + def name(self): ... + def clear_cache(self) -> None: ... diff --git a/stubs/sagemaker/sagemaker/apiutils/__init__.pyi b/stubs/sagemaker/sagemaker/apiutils/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/apiutils/_base_types.pyi b/stubs/sagemaker/sagemaker/apiutils/_base_types.pyi new file mode 100644 index 000000000000..b655044fb08b --- /dev/null +++ b/stubs/sagemaker/sagemaker/apiutils/_base_types.pyi @@ -0,0 +1,16 @@ +from _typeshed import Incomplete + +class ApiObject: + def __init__(self, **kwargs) -> None: ... + @classmethod + def from_boto(cls, boto_dict, **kwargs): ... + @classmethod + def to_boto(cls, obj): ... + def __eq__(self, other): ... + def __ne__(self, other): ... + def __hash__(self): ... + +class Record(ApiObject): + sagemaker_session: Incomplete + def __init__(self, sagemaker_session: Incomplete | None = None, **kwargs) -> None: ... + def with_boto(self, boto_dict): ... diff --git a/stubs/sagemaker/sagemaker/apiutils/_boto_functions.pyi b/stubs/sagemaker/sagemaker/apiutils/_boto_functions.pyi new file mode 100644 index 000000000000..7037ea90664b --- /dev/null +++ b/stubs/sagemaker/sagemaker/apiutils/_boto_functions.pyi @@ -0,0 +1,4 @@ +def to_camel_case(snake_case): ... +def to_snake_case(name): ... +def from_boto(boto_dict, boto_name_to_member_name, member_name_to_type): ... +def to_boto(member_vars, member_name_to_boto_name, member_name_to_type): ... diff --git a/stubs/sagemaker/sagemaker/apiutils/_utils.pyi b/stubs/sagemaker/sagemaker/apiutils/_utils.pyi new file mode 100644 index 000000000000..4d2a6fe76c95 --- /dev/null +++ b/stubs/sagemaker/sagemaker/apiutils/_utils.pyi @@ -0,0 +1,3 @@ +def suffix(): ... +def name(prefix): ... +def default_session(): ... diff --git a/stubs/sagemaker/sagemaker/async_inference/__init__.pyi b/stubs/sagemaker/sagemaker/async_inference/__init__.pyi new file mode 100644 index 000000000000..406ccb5a63b6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/async_inference/__init__.pyi @@ -0,0 +1,3 @@ +from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig as AsyncInferenceConfig +from sagemaker.async_inference.async_inference_response import AsyncInferenceResponse as AsyncInferenceResponse +from sagemaker.async_inference.waiter_config import WaiterConfig as WaiterConfig diff --git a/stubs/sagemaker/sagemaker/async_inference/async_inference_config.pyi b/stubs/sagemaker/sagemaker/async_inference/async_inference_config.pyi new file mode 100644 index 000000000000..c3a48b38e248 --- /dev/null +++ b/stubs/sagemaker/sagemaker/async_inference/async_inference_config.pyi @@ -0,0 +1,16 @@ +from _typeshed import Incomplete + +class AsyncInferenceConfig: + output_path: Incomplete + max_concurrent_invocations_per_instance: Incomplete + kms_key_id: Incomplete + notification_config: Incomplete + failure_path: Incomplete + def __init__( + self, + output_path: Incomplete | None = None, + max_concurrent_invocations_per_instance: Incomplete | None = None, + kms_key_id: Incomplete | None = None, + notification_config: Incomplete | None = None, + failure_path: Incomplete | None = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/async_inference/async_inference_response.pyi b/stubs/sagemaker/sagemaker/async_inference/async_inference_response.pyi new file mode 100644 index 000000000000..64d0127773c1 --- /dev/null +++ b/stubs/sagemaker/sagemaker/async_inference/async_inference_response.pyi @@ -0,0 +1,8 @@ +from _typeshed import Incomplete + +class AsyncInferenceResponse: + predictor_async: Incomplete + output_path: Incomplete + failure_path: Incomplete + def __init__(self, predictor_async, output_path, failure_path) -> None: ... + def get_result(self, waiter_config: Incomplete | None = None): ... diff --git a/stubs/sagemaker/sagemaker/async_inference/waiter_config.pyi b/stubs/sagemaker/sagemaker/async_inference/waiter_config.pyi new file mode 100644 index 000000000000..24fdba1fd434 --- /dev/null +++ b/stubs/sagemaker/sagemaker/async_inference/waiter_config.pyi @@ -0,0 +1,6 @@ +from _typeshed import Incomplete + +class WaiterConfig: + max_attempts: Incomplete + delay: Incomplete + def __init__(self, max_attempts: int = 60, delay: int = 15) -> None: ... diff --git a/stubs/sagemaker/sagemaker/automl/__init__.pyi b/stubs/sagemaker/sagemaker/automl/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/automl/automl.pyi b/stubs/sagemaker/sagemaker/automl/automl.pyi new file mode 100644 index 000000000000..aff51516981b --- /dev/null +++ b/stubs/sagemaker/sagemaker/automl/automl.pyi @@ -0,0 +1,145 @@ +import abc +from _typeshed import Incomplete +from typing import Dict, List, Optional + +from sagemaker.job import _Job +from sagemaker.session import Session + +logger: Incomplete + +class AutoMLInput: + inputs: Incomplete + target_attribute_name: Incomplete + compression: Incomplete + channel_type: Incomplete + content_type: Incomplete + s3_data_type: Incomplete + sample_weight_attribute_name: Incomplete + def __init__( + self, + inputs, + target_attribute_name, + compression: Incomplete | None = None, + channel_type: Incomplete | None = None, + content_type: Incomplete | None = None, + s3_data_type: Incomplete | None = None, + sample_weight_attribute_name: Incomplete | None = None, + ) -> None: ... + def to_request_dict(self): ... + +class AutoML: + output_path: Incomplete + base_job_name: Incomplete + compression_type: Incomplete + encrypt_inter_container_traffic: Incomplete + problem_type: Incomplete + max_candidate: Incomplete + max_runtime_per_training_job_in_seconds: Incomplete + total_job_runtime_in_seconds: Incomplete + target_attribute_name: Incomplete + job_objective: Incomplete + generate_candidate_definitions_only: Incomplete + tags: Incomplete + content_type: Incomplete + s3_data_type: Incomplete + feature_specification_s3_uri: Incomplete + validation_fraction: Incomplete + mode: Incomplete + auto_generate_endpoint_name: Incomplete + endpoint_name: Incomplete + current_job_name: Incomplete + sagemaker_session: Incomplete + vpc_config: Incomplete + volume_kms_key: Incomplete + output_kms_key: Incomplete + role: Incomplete + sample_weight_attribute_name: Incomplete + def __init__( + self, + role: Optional[str] = None, + target_attribute_name: str = None, + output_kms_key: Optional[str] = None, + output_path: Optional[str] = None, + base_job_name: Optional[str] = None, + compression_type: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + volume_kms_key: Optional[str] = None, + encrypt_inter_container_traffic: Optional[bool] = None, + vpc_config: Optional[Dict[str, List]] = None, + problem_type: Optional[str] = None, + max_candidates: Optional[int] = None, + max_runtime_per_training_job_in_seconds: Optional[int] = None, + total_job_runtime_in_seconds: Optional[int] = None, + job_objective: Optional[Dict[str, str]] = None, + generate_candidate_definitions_only: Optional[bool] = False, + tags: Optional[List[Dict[str, str]]] = None, + content_type: Optional[str] = None, + s3_data_type: Optional[str] = None, + feature_specification_s3_uri: Optional[str] = None, + validation_fraction: Optional[float] = None, + mode: Optional[str] = None, + auto_generate_endpoint_name: Optional[bool] = None, + endpoint_name: Optional[str] = None, + sample_weight_attribute_name: str = None, + ) -> None: ... + latest_auto_ml_job: Incomplete + def fit( + self, inputs: Incomplete | None = None, wait: bool = True, logs: bool = True, job_name: Incomplete | None = None + ) -> None: ... + @classmethod + def attach(cls, auto_ml_job_name, sagemaker_session: Incomplete | None = None): ... + def describe_auto_ml_job(self, job_name: Incomplete | None = None): ... + def best_candidate(self, job_name: Incomplete | None = None): ... + def list_candidates( + self, + job_name: Incomplete | None = None, + status_equals: Incomplete | None = None, + candidate_name: Incomplete | None = None, + candidate_arn: Incomplete | None = None, + sort_order: Incomplete | None = None, + sort_by: Incomplete | None = None, + max_results: Incomplete | None = None, + ): ... + def create_model( + self, + name, + sagemaker_session: Incomplete | None = None, + candidate: Incomplete | None = None, + vpc_config: Incomplete | None = None, + enable_network_isolation: bool = False, + model_kms_key: Incomplete | None = None, + predictor_cls: Incomplete | None = None, + inference_response_keys: Incomplete | None = None, + ): ... + def deploy( + self, + initial_instance_count, + instance_type, + serializer: Incomplete | None = None, + deserializer: Incomplete | None = None, + candidate: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + name: Incomplete | None = None, + endpoint_name: Incomplete | None = None, + tags: Incomplete | None = None, + wait: bool = True, + vpc_config: Incomplete | None = None, + enable_network_isolation: bool = False, + model_kms_key: Incomplete | None = None, + predictor_cls: Incomplete | None = None, + inference_response_keys: Incomplete | None = None, + volume_size: Incomplete | None = None, + model_data_download_timeout: Incomplete | None = None, + container_startup_health_check_timeout: Incomplete | None = None, + ): ... + @classmethod + def validate_and_update_inference_response(cls, inference_containers, inference_response_keys) -> None: ... + +class AutoMLJob(_Job, metaclass=abc.ABCMeta): + inputs: Incomplete + job_name: Incomplete + def __init__(self, sagemaker_session, job_name, inputs) -> None: ... + @classmethod + def start_new(cls, auto_ml, inputs): ... + def describe(self): ... + def wait(self, logs: bool = True) -> None: ... diff --git a/stubs/sagemaker/sagemaker/automl/candidate_estimator.pyi b/stubs/sagemaker/sagemaker/automl/candidate_estimator.pyi new file mode 100644 index 000000000000..e8660365e47e --- /dev/null +++ b/stubs/sagemaker/sagemaker/automl/candidate_estimator.pyi @@ -0,0 +1,30 @@ +from _typeshed import Incomplete + +class CandidateEstimator: + name: Incomplete + containers: Incomplete + steps: Incomplete + sagemaker_session: Incomplete + def __init__(self, candidate, sagemaker_session: Incomplete | None = None) -> None: ... + def get_steps(self): ... + def fit( + self, + inputs, + candidate_name: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + encrypt_inter_container_traffic: Incomplete | None = None, + vpc_config: Incomplete | None = None, + wait: bool = True, + logs: bool = True, + ) -> None: ... + +class CandidateStep: + def __init__(self, name, inputs, step_type, description) -> None: ... + @property + def name(self): ... + @property + def inputs(self): ... + @property + def type(self): ... + @property + def description(self): ... diff --git a/stubs/sagemaker/sagemaker/base_deserializers.pyi b/stubs/sagemaker/sagemaker/base_deserializers.pyi new file mode 100644 index 000000000000..4cd4ea9f4c3c --- /dev/null +++ b/stubs/sagemaker/sagemaker/base_deserializers.pyi @@ -0,0 +1,49 @@ +import abc +from _typeshed import Incomplete + +class BaseDeserializer(abc.ABC, metaclass=abc.ABCMeta): + @abc.abstractmethod + def deserialize(self, stream, content_type): ... + @property + @abc.abstractmethod + def ACCEPT(self): ... + +class SimpleBaseDeserializer: + accept: Incomplete + def __init__(self, accept: str = "*/*") -> None: ... + @property + def ACCEPT(self): ... + +class StringDeserializer(SimpleBaseDeserializer): + encoding: Incomplete + def __init__(self, encoding: str = "UTF-8", accept: str = "application/json") -> None: ... + def deserialize(self, stream, content_type): ... + +class BytesDeserializer(SimpleBaseDeserializer): + def deserialize(self, stream, content_type): ... + +class CSVDeserializer(SimpleBaseDeserializer): + encoding: Incomplete + def __init__(self, encoding: str = "utf-8", accept: str = "text/csv") -> None: ... + def deserialize(self, stream, content_type): ... + +class StreamDeserializer(SimpleBaseDeserializer): + def deserialize(self, stream, content_type): ... + +class NumpyDeserializer(SimpleBaseDeserializer): + dtype: Incomplete + allow_pickle: Incomplete + def __init__(self, dtype: Incomplete | None = None, accept: str = "application/x-npy", allow_pickle: bool = True) -> None: ... + def deserialize(self, stream, content_type): ... + +class JSONDeserializer(SimpleBaseDeserializer): + def __init__(self, accept: str = "application/json") -> None: ... + def deserialize(self, stream, content_type): ... + +class PandasDeserializer(SimpleBaseDeserializer): + def __init__(self, accept=("text/csv", "application/json")) -> None: ... + def deserialize(self, stream, content_type): ... + +class JSONLinesDeserializer(SimpleBaseDeserializer): + def __init__(self, accept: str = "application/jsonlines") -> None: ... + def deserialize(self, stream, content_type): ... diff --git a/stubs/sagemaker/sagemaker/base_predictor.pyi b/stubs/sagemaker/sagemaker/base_predictor.pyi new file mode 100644 index 000000000000..7d6dd88bb122 --- /dev/null +++ b/stubs/sagemaker/sagemaker/base_predictor.pyi @@ -0,0 +1,71 @@ +import abc +from _typeshed import Incomplete +from typing import Any, Tuple + +from sagemaker.deserializers import StreamDeserializer as StreamDeserializer, StringDeserializer as StringDeserializer + +class PredictorBase(abc.ABC, metaclass=abc.ABCMeta): + @abc.abstractmethod + def predict(self, *args, **kwargs) -> Any: ... + @abc.abstractmethod + def delete_predictor(self, *args, **kwargs) -> None: ... + @property + @abc.abstractmethod + def content_type(self) -> str: ... + @property + @abc.abstractmethod + def accept(self) -> Tuple[str]: ... + +class Predictor(PredictorBase): + endpoint_name: Incomplete + sagemaker_session: Incomplete + serializer: Incomplete + deserializer: Incomplete + def __init__( + self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=..., **kwargs + ) -> None: ... + def predict( + self, + data, + initial_args: Incomplete | None = None, + target_model: Incomplete | None = None, + target_variant: Incomplete | None = None, + inference_id: Incomplete | None = None, + ): ... + def update_endpoint( + self, + initial_instance_count: Incomplete | None = None, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + model_name: Incomplete | None = None, + tags: Incomplete | None = None, + kms_key: Incomplete | None = None, + data_capture_config_dict: Incomplete | None = None, + wait: bool = True, + ) -> None: ... + def delete_endpoint(self, delete_endpoint_config: bool = True) -> None: ... + delete_predictor = delete_endpoint + def delete_model(self) -> None: ... + def enable_data_capture(self) -> None: ... + def disable_data_capture(self) -> None: ... + def update_data_capture_config(self, data_capture_config) -> None: ... + def list_monitors(self): ... + def endpoint_context(self): ... + @property + def content_type(self): ... + @property + def accept(self): ... + @content_type.setter + def content_type(self, val: str): ... + @accept.setter + def accept(self, val: str): ... + @property + def endpoint(self): ... + +csv_serializer: Incomplete +json_serializer: Incomplete +npy_serializer: Incomplete +csv_deserializer: Incomplete +json_deserializer: Incomplete +numpy_deserializer: Incomplete +RealTimePredictor: Incomplete diff --git a/stubs/sagemaker/sagemaker/base_serializers.pyi b/stubs/sagemaker/sagemaker/base_serializers.pyi new file mode 100644 index 000000000000..14d27f9e67ef --- /dev/null +++ b/stubs/sagemaker/sagemaker/base_serializers.pyi @@ -0,0 +1,47 @@ +import abc +from _typeshed import Incomplete + +class BaseSerializer(abc.ABC, metaclass=abc.ABCMeta): + @abc.abstractmethod + def serialize(self, data): ... + @property + @abc.abstractmethod + def CONTENT_TYPE(self): ... + +class SimpleBaseSerializer: + content_type: Incomplete + def __init__(self, content_type: str = "application/json") -> None: ... + @property + def CONTENT_TYPE(self): ... + +class CSVSerializer(SimpleBaseSerializer): + def __init__(self, content_type: str = "text/csv") -> None: ... + def serialize(self, data): ... + +class NumpySerializer(SimpleBaseSerializer): + dtype: Incomplete + def __init__(self, dtype: Incomplete | None = None, content_type: str = "application/x-npy") -> None: ... + def serialize(self, data): ... + +class JSONSerializer(SimpleBaseSerializer): + def serialize(self, data): ... + +class IdentitySerializer(SimpleBaseSerializer): + def __init__(self, content_type: str = "application/octet-stream") -> None: ... + def serialize(self, data): ... + +class JSONLinesSerializer(SimpleBaseSerializer): + def __init__(self, content_type: str = "application/jsonlines") -> None: ... + def serialize(self, data): ... + +class SparseMatrixSerializer(SimpleBaseSerializer): + def __init__(self, content_type: str = "application/x-npz") -> None: ... + def serialize(self, data): ... + +class LibSVMSerializer(SimpleBaseSerializer): + def __init__(self, content_type: str = "text/libsvm") -> None: ... + def serialize(self, data): ... + +class DataSerializer(SimpleBaseSerializer): + def __init__(self, content_type: str = "file-path/raw-bytes") -> None: ... + def serialize(self, data): ... diff --git a/stubs/sagemaker/sagemaker/chainer/__init__.pyi b/stubs/sagemaker/sagemaker/chainer/__init__.pyi new file mode 100644 index 000000000000..39f78da6a288 --- /dev/null +++ b/stubs/sagemaker/sagemaker/chainer/__init__.pyi @@ -0,0 +1,2 @@ +from sagemaker.chainer.estimator import Chainer as Chainer +from sagemaker.chainer.model import ChainerModel as ChainerModel, ChainerPredictor as ChainerPredictor diff --git a/stubs/sagemaker/sagemaker/chainer/defaults.pyi b/stubs/sagemaker/sagemaker/chainer/defaults.pyi new file mode 100644 index 000000000000..0555246a1ab9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/chainer/defaults.pyi @@ -0,0 +1 @@ +LATEST_PY2_VERSION: str diff --git a/stubs/sagemaker/sagemaker/chainer/estimator.pyi b/stubs/sagemaker/sagemaker/chainer/estimator.pyi new file mode 100644 index 000000000000..9e1cd059a122 --- /dev/null +++ b/stubs/sagemaker/sagemaker/chainer/estimator.pyi @@ -0,0 +1,40 @@ +from _typeshed import Incomplete +from typing import Dict, Optional, Union + +from sagemaker.estimator import Framework +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class Chainer(Framework): + framework_version: Incomplete + py_version: Incomplete + use_mpi: Incomplete + num_processes: Incomplete + process_slots_per_host: Incomplete + additional_mpi_options: Incomplete + def __init__( + self, + entry_point: Union[str, PipelineVariable], + use_mpi: Optional[Union[bool, PipelineVariable]] = None, + num_processes: Optional[Union[int, PipelineVariable]] = None, + process_slots_per_host: Optional[Union[int, PipelineVariable]] = None, + additional_mpi_options: Optional[Union[str, PipelineVariable]] = None, + source_dir: Optional[Union[str, PipelineVariable]] = None, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + framework_version: Optional[str] = None, + py_version: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + **kwargs, + ) -> None: ... + def hyperparameters(self): ... + def create_model( + self, + model_server_workers: Incomplete | None = None, + role: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + entry_point: Incomplete | None = None, + source_dir: Incomplete | None = None, + dependencies: Incomplete | None = None, + **kwargs, + ): ... diff --git a/stubs/sagemaker/sagemaker/chainer/model.pyi b/stubs/sagemaker/sagemaker/chainer/model.pyi new file mode 100644 index 000000000000..9a03e7b2ddc1 --- /dev/null +++ b/stubs/sagemaker/sagemaker/chainer/model.pyi @@ -0,0 +1,69 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker import ModelMetrics +from sagemaker.drift_check_baselines import DriftCheckBaselines +from sagemaker.metadata_properties import MetadataProperties +from sagemaker.model import FrameworkModel +from sagemaker.predictor import Predictor +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class ChainerPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class ChainerModel(FrameworkModel): + framework_version: Incomplete + py_version: Incomplete + model_server_workers: Incomplete + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + entry_point: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + framework_version: Optional[str] = None, + py_version: Optional[str] = None, + predictor_cls: callable = ..., + model_server_workers: Optional[Union[int, PipelineVariable]] = None, + **kwargs, + ) -> None: ... + image_uri: Incomplete + def register( + self, + 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, + domain: Optional[Union[str, PipelineVariable]] = None, + sample_payload_url: Optional[Union[str, PipelineVariable]] = None, + task: Optional[Union[str, PipelineVariable]] = None, + framework: Optional[Union[str, PipelineVariable]] = None, + framework_version: Optional[Union[str, PipelineVariable]] = None, + nearest_model_name: Optional[Union[str, PipelineVariable]] = None, + data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + ): ... + def prepare_container_def( + self, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... + def serving_image_uri( + self, + region_name, + instance_type, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... diff --git a/stubs/sagemaker/sagemaker/clarify.pyi b/stubs/sagemaker/sagemaker/clarify.pyi new file mode 100644 index 000000000000..739411a48042 --- /dev/null +++ b/stubs/sagemaker/sagemaker/clarify.pyi @@ -0,0 +1,281 @@ +import abc +from _typeshed import Incomplete +from abc import ABC, abstractmethod +from enum import Enum +from typing import Any, Dict, List, Optional, Union + +from sagemaker.network import NetworkConfig +from sagemaker.processing import Processor +from sagemaker.session import Session + +logger: Incomplete +ENDPOINT_NAME_PREFIX_PATTERN: str +ANALYSIS_CONFIG_SCHEMA_V1_0: Incomplete + +class DatasetType(Enum): + TEXTCSV: str + JSONLINES: str + JSON: str + PARQUET: str + IMAGE: str + +class DataConfig: + s3_data_input_path: Incomplete + s3_output_path: Incomplete + s3_analysis_config_output_path: Incomplete + s3_data_distribution_type: str + s3_compression_type: Incomplete + label: Incomplete + headers: Incomplete + features: Incomplete + facet_dataset_uri: Incomplete + facet_headers: Incomplete + predicted_label_dataset_uri: Incomplete + predicted_label_headers: Incomplete + predicted_label: Incomplete + excluded_columns: Incomplete + analysis_config: Incomplete + def __init__( + self, + s3_data_input_path: str, + s3_output_path: str, + s3_analysis_config_output_path: Optional[str] = None, + label: Optional[str] = None, + headers: Optional[List[str]] = None, + features: Optional[str] = None, + dataset_type: str = "text/csv", + s3_compression_type: str = "None", + joinsource: Optional[Union[str, int]] = None, + facet_dataset_uri: Optional[str] = None, + facet_headers: Optional[List[str]] = None, + predicted_label_dataset_uri: Optional[str] = None, + predicted_label_headers: Optional[List[str]] = None, + predicted_label: Optional[Union[str, int]] = None, + excluded_columns: Optional[Union[List[int], List[str]]] = None, + ) -> None: ... + def get_config(self): ... + +class BiasConfig: + analysis_config: Incomplete + def __init__( + self, + label_values_or_threshold: Union[int, float, str], + facet_name: Union[str, int, List[str], List[int]], + facet_values_or_threshold: Optional[Union[int, float, str]] = None, + group_name: Optional[str] = None, + ) -> None: ... + def get_config(self): ... + +class ModelConfig: + predictor_config: Incomplete + def __init__( + self, + model_name: Optional[str] = None, + instance_count: Optional[int] = None, + instance_type: Optional[str] = None, + accept_type: Optional[str] = None, + content_type: Optional[str] = None, + content_template: Optional[str] = None, + record_template: Optional[str] = None, + custom_attributes: Optional[str] = None, + accelerator_type: Optional[str] = None, + endpoint_name_prefix: Optional[str] = None, + target_model: Optional[str] = None, + endpoint_name: Optional[str] = None, + ) -> None: ... + def get_predictor_config(self): ... + +class ModelPredictedLabelConfig: + label: Incomplete + probability: Incomplete + probability_threshold: Incomplete + label_headers: Incomplete + predictor_config: Incomplete + def __init__( + self, + label: Optional[Union[str, int]] = None, + probability: Optional[Union[str, int]] = None, + probability_threshold: Optional[float] = None, + label_headers: Optional[List[str]] = None, + ) -> None: ... + def get_predictor_config(self): ... + +class ExplainabilityConfig(ABC, metaclass=abc.ABCMeta): + @abstractmethod + def get_explainability_config(self): ... + +class PDPConfig(ExplainabilityConfig): + pdp_config: Incomplete + def __init__(self, features: Optional[List] = None, grid_resolution: int = 15, top_k_features: int = 10) -> None: ... + def get_explainability_config(self): ... + +class TextConfig: + text_config: Incomplete + def __init__(self, granularity: str, language: str) -> None: ... + def get_text_config(self): ... + +class ImageConfig: + image_config: Incomplete + def __init__( + self, + model_type: str, + num_segments: Optional[int] = None, + feature_extraction_method: Optional[str] = None, + segment_compactness: Optional[float] = None, + max_objects: Optional[int] = None, + iou_threshold: Optional[float] = None, + context: Optional[float] = None, + ) -> None: ... + def get_image_config(self): ... + +class SHAPConfig(ExplainabilityConfig): + shap_config: Incomplete + def __init__( + self, + baseline: Optional[Union[str, List, Dict]] = None, + num_samples: Optional[int] = None, + agg_method: Optional[str] = None, + use_logit: bool = False, + save_local_shap_values: bool = True, + seed: Optional[int] = None, + num_clusters: Optional[int] = None, + text_config: Optional[TextConfig] = None, + image_config: Optional[ImageConfig] = None, + ) -> None: ... + def get_explainability_config(self): ... + +class SageMakerClarifyProcessor(Processor): + job_name_prefix: Incomplete + skip_early_validation: Incomplete + def __init__( + self, + role: Optional[str] = None, + instance_count: int = None, + instance_type: str = None, + volume_size_in_gb: int = 30, + volume_kms_key: Optional[str] = None, + output_kms_key: Optional[str] = None, + max_runtime_in_seconds: Optional[int] = None, + sagemaker_session: Optional[Session] = None, + env: Optional[Dict[str, str]] = None, + tags: Optional[List[Dict[str, str]]] = None, + network_config: Optional[NetworkConfig] = None, + job_name_prefix: Optional[str] = None, + version: Optional[str] = None, + skip_early_validation: bool = False, + ) -> None: ... + def run(self, **_) -> None: ... + def run_pre_training_bias( + self, + data_config: DataConfig, + data_bias_config: BiasConfig, + methods: Union[str, List[str]] = "all", + wait: bool = True, + logs: bool = True, + job_name: Optional[str] = None, + kms_key: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + ): ... + def run_post_training_bias( + self, + data_config: DataConfig, + data_bias_config: BiasConfig, + model_config: Optional[ModelConfig] = None, + model_predicted_label_config: Optional[ModelPredictedLabelConfig] = None, + methods: Union[str, List[str]] = "all", + wait: bool = True, + logs: bool = True, + job_name: Optional[str] = None, + kms_key: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + ): ... + def run_bias( + self, + data_config: DataConfig, + bias_config: BiasConfig, + model_config: Optional[ModelConfig] = None, + model_predicted_label_config: Optional[ModelPredictedLabelConfig] = None, + pre_training_methods: Union[str, List[str]] = "all", + post_training_methods: Union[str, List[str]] = "all", + wait: bool = True, + logs: bool = True, + job_name: Optional[str] = None, + kms_key: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + ): ... + def run_explainability( + self, + data_config: DataConfig, + model_config: ModelConfig, + explainability_config: Union[ExplainabilityConfig, List], + model_scores: Optional[Union[int, str, ModelPredictedLabelConfig]] = None, + wait: bool = True, + logs: bool = True, + job_name: Optional[str] = None, + kms_key: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + ): ... + def run_bias_and_explainability( + self, + data_config: DataConfig, + model_config: ModelConfig, + explainability_config: Union[ExplainabilityConfig, List[ExplainabilityConfig]], + bias_config: BiasConfig, + pre_training_methods: Union[str, List[str]] = "all", + post_training_methods: Union[str, List[str]] = "all", + model_predicted_label_config: ModelPredictedLabelConfig = None, + wait: bool = True, + logs: bool = True, + job_name: Incomplete | None = None, + kms_key: Incomplete | None = None, + experiment_config: Incomplete | None = None, + ): ... + +class _AnalysisConfigGenerator: + @classmethod + def bias_and_explainability( + cls, + data_config: DataConfig, + model_config: ModelConfig, + model_predicted_label_config: ModelPredictedLabelConfig, + explainability_config: Union[ExplainabilityConfig, List[ExplainabilityConfig]], + bias_config: BiasConfig, + pre_training_methods: Union[str, List[str]] = "all", + post_training_methods: Union[str, List[str]] = "all", + ): ... + @classmethod + def explainability( + cls, + data_config: DataConfig, + model_config: ModelConfig, + model_predicted_label_config: ModelPredictedLabelConfig, + explainability_config: Union[ExplainabilityConfig, List[ExplainabilityConfig]], + ): ... + @classmethod + def bias_pre_training(cls, data_config: DataConfig, bias_config: BiasConfig, methods: Union[str, List[str]]): ... + @classmethod + def bias_post_training( + cls, + data_config: DataConfig, + bias_config: BiasConfig, + model_predicted_label_config: ModelPredictedLabelConfig, + methods: Union[str, List[str]], + model_config: ModelConfig, + ): ... + @classmethod + def bias( + cls, + data_config: DataConfig, + bias_config: BiasConfig, + model_config: ModelConfig, + model_predicted_label_config: ModelPredictedLabelConfig, + pre_training_methods: Union[str, List[str]] = "all", + post_training_methods: Union[str, List[str]] = "all", + ): ... + +class ProcessingOutputHandler: + class S3UploadMode(Enum): + CONTINUOUS: str + ENDOFJOB: str + @classmethod + def get_s3_upload_mode(cls, analysis_config: Dict[str, Any]) -> str: ... diff --git a/stubs/sagemaker/sagemaker/cli/__init__.pyi b/stubs/sagemaker/sagemaker/cli/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/__init__.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/__init__.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/ast_transformer.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/ast_transformer.pyi new file mode 100644 index 000000000000..4684aadcc2f0 --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/ast_transformer.pyi @@ -0,0 +1,15 @@ +import ast +from _typeshed import Incomplete + +FUNCTION_CALL_MODIFIERS: Incomplete +IMPORT_MODIFIERS: Incomplete +NAME_MODIFIERS: Incomplete +MODULE_MODIFIERS: Incomplete +IMPORT_FROM_MODIFIERS: Incomplete + +class ASTTransformer(ast.NodeTransformer): + def visit_Call(self, node): ... + def visit_Name(self, node): ... + def visit_Import(self, node): ... + def visit_Module(self, node): ... + def visit_ImportFrom(self, node): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/files.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/files.pyi new file mode 100644 index 000000000000..f744f79eaa09 --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/files.pyi @@ -0,0 +1,18 @@ +import abc +from _typeshed import Incomplete +from abc import abstractmethod + +LOGGER: Incomplete + +class FileUpdater(metaclass=abc.ABCMeta): + input_path: Incomplete + output_path: Incomplete + def __init__(self, input_path, output_path) -> None: ... + @abstractmethod + def update(self): ... + +class PyFileUpdater(FileUpdater): + def update(self) -> None: ... + +class JupyterNotebookFileUpdater(FileUpdater): + def update(self) -> None: ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/__init__.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/__init__.pyi new file mode 100644 index 000000000000..82f875003c99 --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/__init__.pyi @@ -0,0 +1,13 @@ +from sagemaker.cli.compatibility.v2.modifiers import ( + airflow as airflow, + deprecated_params as deprecated_params, + framework_version as framework_version, + image_uris as image_uris, + predictors as predictors, + renamed_params as renamed_params, + serde as serde, + tf_legacy_mode as tf_legacy_mode, + tfs as tfs, + training_input as training_input, + training_params as training_params, +) diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/airflow.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/airflow.pyi new file mode 100644 index 000000000000..e0b278f673fb --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/airflow.pyi @@ -0,0 +1,20 @@ +from _typeshed import Incomplete + +from sagemaker.cli.compatibility.v2.modifiers import renamed_params +from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier + +FUNCTION_NAMES: Incomplete +NAMESPACES: Incomplete +FUNCTIONS: Incomplete + +class ModelConfigArgModifier(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class ModelConfigImageURIRenamer(renamed_params.ParamRenamer): + @property + def calls_to_modify(self): ... + @property + def old_param_name(self): ... + @property + def new_param_name(self): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/deprecated_params.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/deprecated_params.pyi new file mode 100644 index 000000000000..49317ded4b89 --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/deprecated_params.pyi @@ -0,0 +1,9 @@ +from _typeshed import Incomplete + +from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier + +TF_NAMESPACES: Incomplete + +class TensorFlowScriptModeParameterRemover(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/framework_version.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/framework_version.pyi new file mode 100644 index 000000000000..c7124a19b486 --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/framework_version.pyi @@ -0,0 +1,15 @@ +from _typeshed import Incomplete + +from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier + +FRAMEWORK_ARG: str +IMAGE_ARG: str +PY_ARG: str +FRAMEWORK_DEFAULTS: Incomplete +FRAMEWORK_CLASSES: Incomplete +ESTIMATORS: Incomplete +MODELS: Incomplete + +class FrameworkVersionEnforcer(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/image_uris.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/image_uris.pyi new file mode 100644 index 000000000000..93256cbe352f --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/image_uris.pyi @@ -0,0 +1,14 @@ +from _typeshed import Incomplete + +from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier + +GET_IMAGE_URI_NAME: str +GET_IMAGE_URI_NAMESPACES: Incomplete + +class ImageURIRetrieveRefactor(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class ImageURIRetrieveImportFromRenamer(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/matching.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/matching.pyi new file mode 100644 index 000000000000..2a189dee95fe --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/matching.pyi @@ -0,0 +1,6 @@ +def matches_any(node, name_to_namespaces_dict): ... +def matches_name_or_namespaces(node, name, namespaces): ... +def matches_name(node, name): ... +def matches_attr(node, name): ... +def matches_namespace(node, namespace): ... +def has_arg(node, arg): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/modifier.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/modifier.pyi new file mode 100644 index 000000000000..daea66f0dcf9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/modifier.pyi @@ -0,0 +1,9 @@ +import abc +from abc import abstractmethod + +class Modifier(metaclass=abc.ABCMeta): + def check_and_modify_node(self, node): ... + @abstractmethod + def node_should_be_modified(self, node): ... + @abstractmethod + def modify_node(self, node): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/parsing.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/parsing.pyi new file mode 100644 index 000000000000..00d78b6f0f70 --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/parsing.pyi @@ -0,0 +1,2 @@ +def arg_from_keywords(node, arg): ... +def arg_value(node, arg): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/predictors.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/predictors.pyi new file mode 100644 index 000000000000..cc4d682a394f --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/predictors.pyi @@ -0,0 +1,14 @@ +from _typeshed import Incomplete + +from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier + +BASE_PREDICTOR: str +PREDICTORS: Incomplete + +class PredictorConstructorRefactor(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class PredictorImportFromRenamer(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/renamed_params.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/renamed_params.pyi new file mode 100644 index 000000000000..9047ecd184c7 --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/renamed_params.pyi @@ -0,0 +1,76 @@ +import abc +from abc import abstractmethod + +from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier + +class ParamRenamer(Modifier, metaclass=abc.ABCMeta): + @property + @abstractmethod + def calls_to_modify(self): ... + @property + @abstractmethod + def old_param_name(self): ... + @property + @abstractmethod + def new_param_name(self): ... + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class MethodParamRenamer(ParamRenamer, metaclass=abc.ABCMeta): + def node_should_be_modified(self, node): ... + +class DistributionParameterRenamer(ParamRenamer): + @property + def calls_to_modify(self): ... + @property + def old_param_name(self): ... + @property + def new_param_name(self): ... + +class S3SessionRenamer(MethodParamRenamer): + @property + def calls_to_modify(self): ... + @property + def old_param_name(self): ... + @property + def new_param_name(self): ... + +class EstimatorImageURIRenamer(ParamRenamer): + @property + def calls_to_modify(self): ... + @property + def old_param_name(self): ... + @property + def new_param_name(self): ... + +class ModelImageURIRenamer(ParamRenamer): + @property + def calls_to_modify(self): ... + @property + def old_param_name(self): ... + @property + def new_param_name(self): ... + +class EstimatorCreateModelImageURIRenamer(MethodParamRenamer): + @property + def calls_to_modify(self): ... + @property + def old_param_name(self): ... + @property + def new_param_name(self): ... + +class SessionCreateModelImageURIRenamer(MethodParamRenamer): + @property + def calls_to_modify(self): ... + @property + def old_param_name(self): ... + @property + def new_param_name(self): ... + +class SessionCreateEndpointImageURIRenamer(MethodParamRenamer): + @property + def calls_to_modify(self): ... + @property + def old_param_name(self): ... + @property + def new_param_name(self): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/serde.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/serde.pyi new file mode 100644 index 000000000000..26d19b46f596 --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/serde.pyi @@ -0,0 +1,47 @@ +from _typeshed import Incomplete + +from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier + +OLD_AMAZON_CLASS_NAMES: Incomplete +NEW_AMAZON_CLASS_NAMES: Incomplete +OLD_PREDICTOR_CLASS_NAMES: Incomplete +OLD_CLASS_NAME_TO_NAMESPACES: Incomplete +NEW_CLASS_NAME_TO_NAMESPACES: Incomplete +OLD_CLASS_NAME_TO_NEW_CLASS_NAME: Incomplete +OLD_OBJECT_NAME_TO_NEW_CLASS_NAME: Incomplete +NEW_CLASS_NAMES: Incomplete +OLD_CLASS_NAMES: Incomplete +OLD_OBJECT_NAMES: Incomplete + +class SerdeConstructorRenamer(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class SerdeKeywordRemover(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class SerdeObjectRenamer(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class SerdeImportFromPredictorRenamer(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class SerdeImportFromAmazonCommonRenamer(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class _ImportInserter(Modifier): + class_names: Incomplete + import_node: Incomplete + def __init__(self, class_names, import_node) -> None: ... + def node_should_be_modified(self, module): ... + def modify_node(self, module): ... + +class SerializerImportInserter(_ImportInserter): + def __init__(self) -> None: ... + +class DeserializerImportInserter(_ImportInserter): + def __init__(self) -> None: ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/tf_legacy_mode.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/tf_legacy_mode.pyi new file mode 100644 index 000000000000..9ff7340fa4de --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/tf_legacy_mode.pyi @@ -0,0 +1,17 @@ +from _typeshed import Incomplete + +from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier + +TF_NAMESPACES: Incomplete +LEGACY_MODE_PARAMETERS: Incomplete + +class TensorFlowLegacyModeConstructorUpgrader(Modifier): + def __init__(self) -> None: ... + @property + def region(self): ... + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class TensorBoardParameterRemover(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/tfs.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/tfs.pyi new file mode 100644 index 000000000000..253bb3090ece --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/tfs.pyi @@ -0,0 +1,18 @@ +import abc +from _typeshed import Incomplete + +from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier + +CLASS_NAMES: Incomplete +TFS_CLASSES: Incomplete + +class TensorFlowServingConstructorRenamer(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class TensorFlowServingImportFromRenamer(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class TensorFlowServingImportRenamer(Modifier, metaclass=abc.ABCMeta): + def check_and_modify_node(self, node): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/training_input.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/training_input.pyi new file mode 100644 index 000000000000..bf937de5bc8b --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/training_input.pyi @@ -0,0 +1,22 @@ +from _typeshed import Incomplete + +from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier + +S3_INPUT_NAME: str +S3_INPUT_NAMESPACES: Incomplete + +class TrainingInputConstructorRefactor(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class TrainingInputImportFromRenamer(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class ShuffleConfigModuleRenamer(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... + +class ShuffleConfigImportFromRenamer(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/training_params.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/training_params.pyi new file mode 100644 index 000000000000..376993028e59 --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/modifiers/training_params.pyi @@ -0,0 +1,10 @@ +from _typeshed import Incomplete + +from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier + +ESTIMATORS: Incomplete +PARAMS: Incomplete + +class TrainPrefixRemover(Modifier): + def node_should_be_modified(self, node): ... + def modify_node(self, node): ... diff --git a/stubs/sagemaker/sagemaker/cli/compatibility/v2/sagemaker_upgrade_v2.pyi b/stubs/sagemaker/sagemaker/cli/compatibility/v2/sagemaker_upgrade_v2.pyi new file mode 100644 index 000000000000..7e7363e797f3 --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/compatibility/v2/sagemaker_upgrade_v2.pyi @@ -0,0 +1 @@ +def main() -> None: ... diff --git a/stubs/sagemaker/sagemaker/cli/framework_upgrade.pyi b/stubs/sagemaker/sagemaker/cli/framework_upgrade.pyi new file mode 100644 index 000000000000..54bf7f665e37 --- /dev/null +++ b/stubs/sagemaker/sagemaker/cli/framework_upgrade.pyi @@ -0,0 +1,14 @@ +from _typeshed import Incomplete + +IMAGE_URI_CONFIG_DIR: Incomplete + +def get_latest_values(existing_content, scope: Incomplete | None = None): ... +def add_dlc_framework_version( + existing_content, short_version, full_version, scope, processors, py_versions, registries, repository +) -> None: ... +def add_algo_version( + existing_content, processors, scopes, full_version, py_versions, registries, repository, tag_prefix +) -> None: ... +def add_region(existing_content, region, account) -> None: ... +def add_version(existing_content, short_version, full_version, scope, processors, py_versions, tag_prefix) -> None: ... +def main() -> None: ... diff --git a/stubs/sagemaker/sagemaker/collection.pyi b/stubs/sagemaker/sagemaker/collection.pyi new file mode 100644 index 000000000000..447c831f8f74 --- /dev/null +++ b/stubs/sagemaker/sagemaker/collection.pyi @@ -0,0 +1,12 @@ +from _typeshed import Incomplete +from typing import List + +class Collection: + sagemaker_session: Incomplete + def __init__(self, sagemaker_session) -> None: ... + def create(self, collection_name: str, parent_collection_name: str = None): ... + def delete(self, collections: List[str]): ... + def add_model_groups(self, collection_name: str, model_groups: List[str]): ... + def remove_model_groups(self, collection_name: str, model_groups: List[str]): ... + def move_model_group(self, source_collection_name: str, model_group: str, destination_collection_name: str): ... + def list_collection(self, collection_name: str = None): ... diff --git a/stubs/sagemaker/sagemaker/config/__init__.pyi b/stubs/sagemaker/sagemaker/config/__init__.pyi new file mode 100644 index 000000000000..30c1276d9aa6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/config/__init__.pyi @@ -0,0 +1,149 @@ +from sagemaker.config.config import ( + load_sagemaker_config as load_sagemaker_config, + validate_sagemaker_config as validate_sagemaker_config, +) +from sagemaker.config.config_schema import ( + ASYNC_INFERENCE_CONFIG as ASYNC_INFERENCE_CONFIG, + ATHENA_DATASET_DEFINITION as ATHENA_DATASET_DEFINITION, + AUTO_ML_INTER_CONTAINER_ENCRYPTION_PATH as AUTO_ML_INTER_CONTAINER_ENCRYPTION_PATH, + AUTO_ML_JOB as AUTO_ML_JOB, + AUTO_ML_JOB_CONFIG as AUTO_ML_JOB_CONFIG, + AUTO_ML_JOB_CONFIG_PATH as AUTO_ML_JOB_CONFIG_PATH, + AUTO_ML_KMS_KEY_ID_PATH as AUTO_ML_KMS_KEY_ID_PATH, + AUTO_ML_OUTPUT_CONFIG_PATH as AUTO_ML_OUTPUT_CONFIG_PATH, + AUTO_ML_ROLE_ARN_PATH as AUTO_ML_ROLE_ARN_PATH, + AUTO_ML_VOLUME_KMS_KEY_ID_PATH as AUTO_ML_VOLUME_KMS_KEY_ID_PATH, + AUTO_ML_VPC_CONFIG_PATH as AUTO_ML_VPC_CONFIG_PATH, + CLUSTER_CONFIG as CLUSTER_CONFIG, + CLUSTER_ROLE_ARN as CLUSTER_ROLE_ARN, + COMPILATION_JOB as COMPILATION_JOB, + COMPILATION_JOB_OUTPUT_CONFIG_PATH as COMPILATION_JOB_OUTPUT_CONFIG_PATH, + COMPILATION_JOB_ROLE_ARN_PATH as COMPILATION_JOB_ROLE_ARN_PATH, + COMPILATION_JOB_VPC_CONFIG_PATH as COMPILATION_JOB_VPC_CONFIG_PATH, + CONTAINERS as CONTAINERS, + DATA_CAPTURE_CONFIG as DATA_CAPTURE_CONFIG, + DATASET_DEFINITION as DATASET_DEFINITION, + DEFAULT_S3_BUCKET as DEFAULT_S3_BUCKET, + DEFAULT_S3_OBJECT_KEY_PREFIX as DEFAULT_S3_OBJECT_KEY_PREFIX, + DISABLE_PROFILER as DISABLE_PROFILER, + EDGE_PACKAGING_JOB as EDGE_PACKAGING_JOB, + EDGE_PACKAGING_KMS_KEY_ID_PATH as EDGE_PACKAGING_KMS_KEY_ID_PATH, + EDGE_PACKAGING_OUTPUT_CONFIG_PATH as EDGE_PACKAGING_OUTPUT_CONFIG_PATH, + EDGE_PACKAGING_RESOURCE_KEY_PATH as EDGE_PACKAGING_RESOURCE_KEY_PATH, + EDGE_PACKAGING_ROLE_ARN_PATH as EDGE_PACKAGING_ROLE_ARN_PATH, + ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION as ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION, + ENABLE_NETWORK_ISOLATION as ENABLE_NETWORK_ISOLATION, + ENDPOINT as ENDPOINT, + ENDPOINT_CONFIG as ENDPOINT_CONFIG, + ENDPOINT_CONFIG_ASYNC_INFERENCE_PATH as ENDPOINT_CONFIG_ASYNC_INFERENCE_PATH, + ENDPOINT_CONFIG_ASYNC_KMS_KEY_ID_PATH as ENDPOINT_CONFIG_ASYNC_KMS_KEY_ID_PATH, + ENDPOINT_CONFIG_DATA_CAPTURE_KMS_KEY_ID_PATH as ENDPOINT_CONFIG_DATA_CAPTURE_KMS_KEY_ID_PATH, + ENDPOINT_CONFIG_DATA_CAPTURE_PATH as ENDPOINT_CONFIG_DATA_CAPTURE_PATH, + ENDPOINT_CONFIG_KMS_KEY_ID_PATH as ENDPOINT_CONFIG_KMS_KEY_ID_PATH, + ENDPOINT_CONFIG_PRODUCTION_VARIANTS_PATH as ENDPOINT_CONFIG_PRODUCTION_VARIANTS_PATH, + ENDPOINT_TAGS_PATH as ENDPOINT_TAGS_PATH, + ENVIRONMENT as ENVIRONMENT, + EXECUTION_ROLE_ARN as EXECUTION_ROLE_ARN, + FEATURE_GROUP as FEATURE_GROUP, + FEATURE_GROUP_OFFLINE_STORE_CONFIG_PATH as FEATURE_GROUP_OFFLINE_STORE_CONFIG_PATH, + FEATURE_GROUP_OFFLINE_STORE_KMS_KEY_ID_PATH as FEATURE_GROUP_OFFLINE_STORE_KMS_KEY_ID_PATH, + FEATURE_GROUP_ONLINE_STORE_CONFIG_PATH as FEATURE_GROUP_ONLINE_STORE_CONFIG_PATH, + FEATURE_GROUP_ONLINE_STORE_KMS_KEY_ID_PATH as FEATURE_GROUP_ONLINE_STORE_KMS_KEY_ID_PATH, + FEATURE_GROUP_ROLE_ARN_PATH as FEATURE_GROUP_ROLE_ARN_PATH, + INFERENCE_SPECIFICATION as INFERENCE_SPECIFICATION, + KEY as KEY, + KMS_KEY_ID as KMS_KEY_ID, + MODEL as MODEL, + MODEL_CONTAINERS_PATH as MODEL_CONTAINERS_PATH, + MODEL_ENABLE_NETWORK_ISOLATION_PATH as MODEL_ENABLE_NETWORK_ISOLATION_PATH, + MODEL_EXECUTION_ROLE_ARN_PATH as MODEL_EXECUTION_ROLE_ARN_PATH, + MODEL_PACKAGE as MODEL_PACKAGE, + MODEL_PACKAGE_INFERENCE_SPECIFICATION_CONTAINERS_PATH as MODEL_PACKAGE_INFERENCE_SPECIFICATION_CONTAINERS_PATH, + MODEL_PACKAGE_VALIDATION_PROFILES_PATH as MODEL_PACKAGE_VALIDATION_PROFILES_PATH, + MODEL_PACKAGE_VALIDATION_ROLE_PATH as MODEL_PACKAGE_VALIDATION_ROLE_PATH, + MODEL_PRIMARY_CONTAINER_ENVIRONMENT_PATH as MODEL_PRIMARY_CONTAINER_ENVIRONMENT_PATH, + MODEL_PRIMARY_CONTAINER_PATH as MODEL_PRIMARY_CONTAINER_PATH, + MODEL_VPC_CONFIG_PATH as MODEL_VPC_CONFIG_PATH, + MODULES as MODULES, + MONITORING_JOB_DEFINITION as MONITORING_JOB_DEFINITION, + MONITORING_JOB_ENVIRONMENT_PATH as MONITORING_JOB_ENVIRONMENT_PATH, + MONITORING_JOB_NETWORK_CONFIG_PATH as MONITORING_JOB_NETWORK_CONFIG_PATH, + MONITORING_JOB_OUTPUT_KMS_KEY_ID_PATH as MONITORING_JOB_OUTPUT_KMS_KEY_ID_PATH, + MONITORING_JOB_ROLE_ARN_PATH as MONITORING_JOB_ROLE_ARN_PATH, + MONITORING_JOB_VOLUME_KMS_KEY_ID_PATH as MONITORING_JOB_VOLUME_KMS_KEY_ID_PATH, + MONITORING_OUTPUT_CONFIG as MONITORING_OUTPUT_CONFIG, + MONITORING_RESOURCES as MONITORING_RESOURCES, + MONITORING_SCHEDULE as MONITORING_SCHEDULE, + MONITORING_SCHEDULE_CONFIG as MONITORING_SCHEDULE_CONFIG, + MONITORING_SCHEDULE_INTER_CONTAINER_ENCRYPTION_PATH as MONITORING_SCHEDULE_INTER_CONTAINER_ENCRYPTION_PATH, + NETWORK_CONFIG as NETWORK_CONFIG, + OFFLINE_STORE_CONFIG as OFFLINE_STORE_CONFIG, + ONLINE_STORE_CONFIG as ONLINE_STORE_CONFIG, + OUTPUT_CONFIG as OUTPUT_CONFIG, + OUTPUT_DATA_CONFIG as OUTPUT_DATA_CONFIG, + PIPELINE_ROLE_ARN_PATH as PIPELINE_ROLE_ARN_PATH, + PIPELINE_TAGS_PATH as PIPELINE_TAGS_PATH, + PRIMARY_CONTAINER as PRIMARY_CONTAINER, + PROCESSING_INPUTS as PROCESSING_INPUTS, + PROCESSING_JOB as PROCESSING_JOB, + PROCESSING_JOB_ENABLE_NETWORK_ISOLATION_PATH as PROCESSING_JOB_ENABLE_NETWORK_ISOLATION_PATH, + PROCESSING_JOB_ENVIRONMENT_PATH as PROCESSING_JOB_ENVIRONMENT_PATH, + PROCESSING_JOB_INPUTS_PATH as PROCESSING_JOB_INPUTS_PATH, + PROCESSING_JOB_INTER_CONTAINER_ENCRYPTION_PATH as PROCESSING_JOB_INTER_CONTAINER_ENCRYPTION_PATH, + PROCESSING_JOB_KMS_KEY_ID_PATH as PROCESSING_JOB_KMS_KEY_ID_PATH, + PROCESSING_JOB_NETWORK_CONFIG_PATH as PROCESSING_JOB_NETWORK_CONFIG_PATH, + PROCESSING_JOB_PROCESSING_RESOURCES_PATH as PROCESSING_JOB_PROCESSING_RESOURCES_PATH, + PROCESSING_JOB_ROLE_ARN_PATH as PROCESSING_JOB_ROLE_ARN_PATH, + PROCESSING_JOB_SECURITY_GROUP_IDS_PATH as PROCESSING_JOB_SECURITY_GROUP_IDS_PATH, + PROCESSING_JOB_SUBNETS_PATH as PROCESSING_JOB_SUBNETS_PATH, + PROCESSING_JOB_VOLUME_KMS_KEY_ID_PATH as PROCESSING_JOB_VOLUME_KMS_KEY_ID_PATH, + PROCESSING_OUTPUT_CONFIG as PROCESSING_OUTPUT_CONFIG, + PROCESSING_OUTPUT_CONFIG_PATH as PROCESSING_OUTPUT_CONFIG_PATH, + PROCESSING_RESOURCES as PROCESSING_RESOURCES, + PRODUCTION_VARIANTS as PRODUCTION_VARIANTS, + PROFILER_CONFIG as PROFILER_CONFIG, + PYTHON_SDK as PYTHON_SDK, + REDSHIFT_DATASET_DEFINITION as REDSHIFT_DATASET_DEFINITION, + RESOURCE_CONFIG as RESOURCE_CONFIG, + RESOURCE_KEY as RESOURCE_KEY, + ROLE_ARN as ROLE_ARN, + S3_STORAGE_CONFIG as S3_STORAGE_CONFIG, + SAGEMAKER as SAGEMAKER, + SCHEMA_VERSION as SCHEMA_VERSION, + SECURITY_CONFIG as SECURITY_CONFIG, + SECURITY_GROUP_IDS as SECURITY_GROUP_IDS, + SESSION as SESSION, + SESSION_DEFAULT_S3_BUCKET_PATH as SESSION_DEFAULT_S3_BUCKET_PATH, + SESSION_DEFAULT_S3_OBJECT_KEY_PREFIX_PATH as SESSION_DEFAULT_S3_OBJECT_KEY_PREFIX_PATH, + SUBNETS as SUBNETS, + TAGS as TAGS, + TRAINING_JOB as TRAINING_JOB, + TRAINING_JOB_DISABLE_PROFILER_PATH as TRAINING_JOB_DISABLE_PROFILER_PATH, + TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH as TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH, + TRAINING_JOB_ENVIRONMENT_PATH as TRAINING_JOB_ENVIRONMENT_PATH, + TRAINING_JOB_INTER_CONTAINER_ENCRYPTION_PATH as TRAINING_JOB_INTER_CONTAINER_ENCRYPTION_PATH, + TRAINING_JOB_KMS_KEY_ID_PATH as TRAINING_JOB_KMS_KEY_ID_PATH, + TRAINING_JOB_OUTPUT_DATA_CONFIG_PATH as TRAINING_JOB_OUTPUT_DATA_CONFIG_PATH, + TRAINING_JOB_PROFILE_CONFIG_PATH as TRAINING_JOB_PROFILE_CONFIG_PATH, + TRAINING_JOB_RESOURCE_CONFIG_PATH as TRAINING_JOB_RESOURCE_CONFIG_PATH, + TRAINING_JOB_ROLE_ARN_PATH as TRAINING_JOB_ROLE_ARN_PATH, + TRAINING_JOB_SECURITY_GROUP_IDS_PATH as TRAINING_JOB_SECURITY_GROUP_IDS_PATH, + TRAINING_JOB_SUBNETS_PATH as TRAINING_JOB_SUBNETS_PATH, + TRAINING_JOB_VOLUME_KMS_KEY_ID_PATH as TRAINING_JOB_VOLUME_KMS_KEY_ID_PATH, + TRAINING_JOB_VPC_CONFIG_PATH as TRAINING_JOB_VPC_CONFIG_PATH, + TRANSFORM_JOB as TRANSFORM_JOB, + TRANSFORM_JOB_DEFINITION as TRANSFORM_JOB_DEFINITION, + TRANSFORM_JOB_ENVIRONMENT_PATH as TRANSFORM_JOB_ENVIRONMENT_PATH, + TRANSFORM_JOB_KMS_KEY_ID_PATH as TRANSFORM_JOB_KMS_KEY_ID_PATH, + TRANSFORM_JOB_VOLUME_KMS_KEY_ID_PATH as TRANSFORM_JOB_VOLUME_KMS_KEY_ID_PATH, + TRANSFORM_OUTPUT as TRANSFORM_OUTPUT, + TRANSFORM_OUTPUT_KMS_KEY_ID_PATH as TRANSFORM_OUTPUT_KMS_KEY_ID_PATH, + TRANSFORM_RESOURCES as TRANSFORM_RESOURCES, + TRANSFORM_RESOURCES_VOLUME_KMS_KEY_ID_PATH as TRANSFORM_RESOURCES_VOLUME_KMS_KEY_ID_PATH, + VALIDATION_PROFILES as VALIDATION_PROFILES, + VALIDATION_ROLE as VALIDATION_ROLE, + VALIDATION_SPECIFICATION as VALIDATION_SPECIFICATION, + VALUE as VALUE, + VOLUME_KMS_KEY_ID as VOLUME_KMS_KEY_ID, + VPC_CONFIG as VPC_CONFIG, +) diff --git a/stubs/sagemaker/sagemaker/config/config.pyi b/stubs/sagemaker/sagemaker/config/config.pyi new file mode 100644 index 000000000000..44b25ad4d987 --- /dev/null +++ b/stubs/sagemaker/sagemaker/config/config.pyi @@ -0,0 +1,10 @@ +from _typeshed import Incomplete +from typing import List + +logger: Incomplete +ENV_VARIABLE_ADMIN_CONFIG_OVERRIDE: str +ENV_VARIABLE_USER_CONFIG_OVERRIDE: str +S3_PREFIX: str + +def load_sagemaker_config(additional_config_paths: List[str] = None, s3_resource: Incomplete | None = None) -> dict: ... +def validate_sagemaker_config(sagemaker_config: dict = None): ... diff --git a/stubs/sagemaker/sagemaker/config/config_schema.pyi b/stubs/sagemaker/sagemaker/config/config_schema.pyi new file mode 100644 index 000000000000..d6308d476dfe --- /dev/null +++ b/stubs/sagemaker/sagemaker/config/config_schema.pyi @@ -0,0 +1,191 @@ +from _typeshed import Incomplete + +SECURITY_GROUP_IDS: str +SUBNETS: str +ENABLE_NETWORK_ISOLATION: str +VOLUME_KMS_KEY_ID: str +KMS_KEY_ID: str +RESOURCE_KEY: str +ROLE_ARN: str +TAGS: str +KEY: str +VALUE: str +EXECUTION_ROLE_ARN: str +CLUSTER_ROLE_ARN: str +VPC_CONFIG: str +OUTPUT_DATA_CONFIG: str +AUTO_ML_JOB_CONFIG: str +ASYNC_INFERENCE_CONFIG: str +OUTPUT_CONFIG: str +PROCESSING_OUTPUT_CONFIG: str +CLUSTER_CONFIG: str +NETWORK_CONFIG: str +CORE_DUMP_CONFIG: str +DATA_CAPTURE_CONFIG: str +MONITORING_OUTPUT_CONFIG: str +RESOURCE_CONFIG: str +SCHEMA_VERSION: str +DATASET_DEFINITION: str +ATHENA_DATASET_DEFINITION: str +REDSHIFT_DATASET_DEFINITION: str +MONITORING_JOB_DEFINITION: str +SAGEMAKER: str +PYTHON_SDK: str +MODULES: str +REMOTE_FUNCTION: str +DEPENDENCIES: str +PRE_EXECUTION_SCRIPT: str +PRE_EXECUTION_COMMANDS: str +ENVIRONMENT_VARIABLES: str +IMAGE_URI: str +INCLUDE_LOCAL_WORKDIR: str +INSTANCE_TYPE: str +S3_KMS_KEY_ID: str +S3_ROOT_URI: str +JOB_CONDA_ENV: str +OFFLINE_STORE_CONFIG: str +ONLINE_STORE_CONFIG: str +S3_STORAGE_CONFIG: str +SECURITY_CONFIG: str +TRANSFORM_JOB_DEFINITION: str +MONITORING_SCHEDULE_CONFIG: str +MONITORING_RESOURCES: str +PROCESSING_RESOURCES: str +PRODUCTION_VARIANTS: str +TRANSFORM_OUTPUT: str +TRANSFORM_RESOURCES: str +VALIDATION_ROLE: str +VALIDATION_SPECIFICATION: str +VALIDATION_PROFILES: str +PROCESSING_INPUTS: str +FEATURE_GROUP: str +EDGE_PACKAGING_JOB: str +TRAINING_JOB: str +PROCESSING_JOB: str +MODEL_PACKAGE: str +MODEL: str +MONITORING_SCHEDULE: str +ENDPOINT_CONFIG: str +ENDPOINT: str +AUTO_ML_JOB: str +COMPILATION_JOB: str +CUSTOM_PARAMETERS: str +PIPELINE: str +TRANSFORM_JOB: str +PROPERTIES: str +PATTERN_PROPERTIES: str +TYPE: str +OBJECT: str +ADDITIONAL_PROPERTIES: str +ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION: str +SESSION: str +DEFAULT_S3_BUCKET: str +DEFAULT_S3_OBJECT_KEY_PREFIX: str +ENVIRONMENT: str +CONTAINERS: str +PRIMARY_CONTAINER: str +INFERENCE_SPECIFICATION: str +PROFILER_CONFIG: str +DISABLE_PROFILER: str +COMPILATION_JOB_VPC_CONFIG_PATH: Incomplete +COMPILATION_JOB_KMS_KEY_ID_PATH: Incomplete +COMPILATION_JOB_OUTPUT_CONFIG_PATH: Incomplete +COMPILATION_JOB_ROLE_ARN_PATH: Incomplete +TRAINING_JOB_ENVIRONMENT_PATH: Incomplete +TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH: Incomplete +TRAINING_JOB_KMS_KEY_ID_PATH: Incomplete +TRAINING_JOB_RESOURCE_CONFIG_PATH: Incomplete +TRAINING_JOB_OUTPUT_DATA_CONFIG_PATH: Incomplete +TRAINING_JOB_VOLUME_KMS_KEY_ID_PATH: Incomplete +TRAINING_JOB_ROLE_ARN_PATH: Incomplete +TRAINING_JOB_VPC_CONFIG_PATH: Incomplete +TRAINING_JOB_SECURITY_GROUP_IDS_PATH: Incomplete +TRAINING_JOB_SUBNETS_PATH: Incomplete +TRAINING_JOB_PROFILE_CONFIG_PATH: Incomplete +TRAINING_JOB_DISABLE_PROFILER_PATH: Incomplete +EDGE_PACKAGING_KMS_KEY_ID_PATH: Incomplete +EDGE_PACKAGING_OUTPUT_CONFIG_PATH: Incomplete +EDGE_PACKAGING_ROLE_ARN_PATH: Incomplete +EDGE_PACKAGING_RESOURCE_KEY_PATH: Incomplete +ENDPOINT_CONFIG_DATA_CAPTURE_KMS_KEY_ID_PATH: Incomplete +ENDPOINT_CONFIG_DATA_CAPTURE_PATH: Incomplete +ENDPOINT_CONFIG_ASYNC_INFERENCE_PATH: Incomplete +ENDPOINT_CONFIG_PRODUCTION_VARIANTS_PATH: Incomplete +ENDPOINT_CONFIG_ASYNC_KMS_KEY_ID_PATH: Incomplete +ENDPOINT_CONFIG_KMS_KEY_ID_PATH: Incomplete +ENDPOINT_TAGS_PATH: Incomplete +FEATURE_GROUP_ONLINE_STORE_CONFIG_PATH: Incomplete +FEATURE_GROUP_OFFLINE_STORE_CONFIG_PATH: Incomplete +FEATURE_GROUP_ROLE_ARN_PATH: Incomplete +FEATURE_GROUP_OFFLINE_STORE_KMS_KEY_ID_PATH: Incomplete +FEATURE_GROUP_ONLINE_STORE_KMS_KEY_ID_PATH: Incomplete +AUTO_ML_OUTPUT_CONFIG_PATH: Incomplete +AUTO_ML_KMS_KEY_ID_PATH: Incomplete +AUTO_ML_VOLUME_KMS_KEY_ID_PATH: Incomplete +AUTO_ML_ROLE_ARN_PATH: Incomplete +AUTO_ML_VPC_CONFIG_PATH: Incomplete +AUTO_ML_JOB_CONFIG_PATH: Incomplete +MONITORING_JOB_DEFINITION_PREFIX: Incomplete +MONITORING_JOB_ENVIRONMENT_PATH: Incomplete +MONITORING_JOB_OUTPUT_KMS_KEY_ID_PATH: Incomplete +MONITORING_JOB_VOLUME_KMS_KEY_ID_PATH: Incomplete +MONITORING_JOB_NETWORK_CONFIG_PATH: Incomplete +MONITORING_JOB_ENABLE_NETWORK_ISOLATION_PATH: Incomplete +MONITORING_JOB_VPC_CONFIG_PATH: Incomplete +MONITORING_JOB_SECURITY_GROUP_IDS_PATH: Incomplete +MONITORING_JOB_SUBNETS_PATH: Incomplete +MONITORING_JOB_ROLE_ARN_PATH: Incomplete +PIPELINE_ROLE_ARN_PATH: Incomplete +PIPELINE_TAGS_PATH: Incomplete +TRANSFORM_JOB_ENVIRONMENT_PATH: Incomplete +TRANSFORM_OUTPUT_KMS_KEY_ID_PATH: Incomplete +TRANSFORM_RESOURCES_VOLUME_KMS_KEY_ID_PATH: Incomplete +TRANSFORM_JOB_KMS_KEY_ID_PATH: Incomplete +TRANSFORM_JOB_VOLUME_KMS_KEY_ID_PATH: Incomplete +MODEL_CONTAINERS_PATH: Incomplete +MODEL_VPC_CONFIG_PATH: Incomplete +MODEL_ENABLE_NETWORK_ISOLATION_PATH: Incomplete +MODEL_EXECUTION_ROLE_ARN_PATH: Incomplete +MODEL_PRIMARY_CONTAINER_PATH: Incomplete +MODEL_PRIMARY_CONTAINER_ENVIRONMENT_PATH: Incomplete +PROCESSING_JOB_ENABLE_NETWORK_ISOLATION_PATH: Incomplete +PROCESSING_JOB_INPUTS_PATH: Incomplete +REDSHIFT_DATASET_DEFINITION_KMS_KEY_ID_PATH: Incomplete +ATHENA_DATASET_DEFINITION_KMS_KEY_ID_PATH: Incomplete +REDSHIFT_DATASET_DEFINITION_CLUSTER_ROLE_ARN_PATH: Incomplete +PROCESSING_JOB_NETWORK_CONFIG_PATH: Incomplete +PROCESSING_JOB_VPC_CONFIG_PATH: Incomplete +PROCESSING_JOB_SUBNETS_PATH: Incomplete +PROCESSING_JOB_SECURITY_GROUP_IDS_PATH: Incomplete +PROCESSING_OUTPUT_CONFIG_PATH: Incomplete +PROCESSING_JOB_KMS_KEY_ID_PATH: Incomplete +PROCESSING_JOB_PROCESSING_RESOURCES_PATH: Incomplete +PROCESSING_JOB_VOLUME_KMS_KEY_ID_PATH: Incomplete +PROCESSING_JOB_ROLE_ARN_PATH: Incomplete +MODEL_PACKAGE_INFERENCE_SPECIFICATION_CONTAINERS_PATH: Incomplete +MODEL_PACKAGE_VALIDATION_ROLE_PATH: Incomplete +MODEL_PACKAGE_VALIDATION_PROFILES_PATH: Incomplete +REMOTE_FUNCTION_DEPENDENCIES: Incomplete +REMOTE_FUNCTION_PRE_EXECUTION_COMMANDS: Incomplete +REMOTE_FUNCTION_PRE_EXECUTION_SCRIPT: Incomplete +REMOTE_FUNCTION_ENVIRONMENT_VARIABLES: Incomplete +REMOTE_FUNCTION_IMAGE_URI: Incomplete +REMOTE_FUNCTION_INCLUDE_LOCAL_WORKDIR: Incomplete +REMOTE_FUNCTION_INSTANCE_TYPE: Incomplete +REMOTE_FUNCTION_JOB_CONDA_ENV: Incomplete +REMOTE_FUNCTION_ROLE_ARN: Incomplete +REMOTE_FUNCTION_S3_KMS_KEY_ID: Incomplete +REMOTE_FUNCTION_S3_ROOT_URI: Incomplete +REMOTE_FUNCTION_TAGS: Incomplete +REMOTE_FUNCTION_VOLUME_KMS_KEY_ID: Incomplete +REMOTE_FUNCTION_VPC_CONFIG_SUBNETS: Incomplete +REMOTE_FUNCTION_VPC_CONFIG_SECURITY_GROUP_IDS: Incomplete +REMOTE_FUNCTION_ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION: Incomplete +MONITORING_SCHEDULE_INTER_CONTAINER_ENCRYPTION_PATH: Incomplete +AUTO_ML_INTER_CONTAINER_ENCRYPTION_PATH: Incomplete +PROCESSING_JOB_ENVIRONMENT_PATH: Incomplete +PROCESSING_JOB_INTER_CONTAINER_ENCRYPTION_PATH: Incomplete +TRAINING_JOB_INTER_CONTAINER_ENCRYPTION_PATH: Incomplete +SESSION_DEFAULT_S3_BUCKET_PATH: Incomplete +SESSION_DEFAULT_S3_OBJECT_KEY_PREFIX_PATH: Incomplete +SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA: Incomplete diff --git a/stubs/sagemaker/sagemaker/config/config_utils.pyi b/stubs/sagemaker/sagemaker/config/config_utils.pyi new file mode 100644 index 000000000000..60657d6174b1 --- /dev/null +++ b/stubs/sagemaker/sagemaker/config/config_utils.pyi @@ -0,0 +1 @@ +def get_sagemaker_config_logger(): ... diff --git a/stubs/sagemaker/sagemaker/content_types.pyi b/stubs/sagemaker/sagemaker/content_types.pyi new file mode 100644 index 000000000000..b8d3fac6a6f3 --- /dev/null +++ b/stubs/sagemaker/sagemaker/content_types.pyi @@ -0,0 +1,21 @@ +from typing import List, Optional + +def retrieve_options( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> List[str]: ... +def retrieve_default( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> str: ... + +CONTENT_TYPE_JSON: str +CONTENT_TYPE_CSV: str +CONTENT_TYPE_OCTET_STREAM: str +CONTENT_TYPE_NPY: str diff --git a/stubs/sagemaker/sagemaker/dataset_definition/__init__.pyi b/stubs/sagemaker/sagemaker/dataset_definition/__init__.pyi new file mode 100644 index 000000000000..e3d8a5c8fca8 --- /dev/null +++ b/stubs/sagemaker/sagemaker/dataset_definition/__init__.pyi @@ -0,0 +1,6 @@ +from sagemaker.dataset_definition.inputs import ( + AthenaDatasetDefinition as AthenaDatasetDefinition, + DatasetDefinition as DatasetDefinition, + RedshiftDatasetDefinition as RedshiftDatasetDefinition, + S3Input as S3Input, +) diff --git a/stubs/sagemaker/sagemaker/dataset_definition/inputs.pyi b/stubs/sagemaker/sagemaker/dataset_definition/inputs.pyi new file mode 100644 index 000000000000..f0cee2eb47e6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/dataset_definition/inputs.pyi @@ -0,0 +1,51 @@ +from _typeshed import Incomplete + +from sagemaker.apiutils._base_types import ApiObject + +class RedshiftDatasetDefinition(ApiObject): + def __init__( + self, + cluster_id: Incomplete | None = None, + database: Incomplete | None = None, + db_user: Incomplete | None = None, + query_string: Incomplete | None = None, + cluster_role_arn: Incomplete | None = None, + output_s3_uri: Incomplete | None = None, + kms_key_id: Incomplete | None = None, + output_format: Incomplete | None = None, + output_compression: Incomplete | None = None, + ) -> None: ... + +class AthenaDatasetDefinition(ApiObject): + def __init__( + self, + catalog: Incomplete | None = None, + database: Incomplete | None = None, + query_string: Incomplete | None = None, + output_s3_uri: Incomplete | None = None, + work_group: Incomplete | None = None, + kms_key_id: Incomplete | None = None, + output_format: Incomplete | None = None, + output_compression: Incomplete | None = None, + ) -> None: ... + +class DatasetDefinition(ApiObject): + def __init__( + self, + data_distribution_type: str = "ShardedByS3Key", + input_mode: str = "File", + local_path: Incomplete | None = None, + redshift_dataset_definition: Incomplete | None = None, + athena_dataset_definition: Incomplete | None = None, + ) -> None: ... + +class S3Input(ApiObject): + def __init__( + self, + s3_uri: Incomplete | None = None, + local_path: Incomplete | None = None, + s3_data_type: str = "S3Prefix", + s3_input_mode: str = "File", + s3_data_distribution_type: str = "FullyReplicated", + s3_compression_type: Incomplete | None = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/debugger/__init__.pyi b/stubs/sagemaker/sagemaker/debugger/__init__.pyi new file mode 100644 index 000000000000..65dbe384b852 --- /dev/null +++ b/stubs/sagemaker/sagemaker/debugger/__init__.pyi @@ -0,0 +1,23 @@ +from sagemaker.debugger.debugger import ( + DEBUGGER_FLAG as DEBUGGER_FLAG, + CollectionConfig as CollectionConfig, + DebuggerHookConfig as DebuggerHookConfig, + ProfilerRule as ProfilerRule, + Rule as Rule, + RuleBase as RuleBase, + TensorBoardOutputConfig as TensorBoardOutputConfig, + framework_name as framework_name, + get_default_profiler_rule as get_default_profiler_rule, + get_rule_container_image_uri as get_rule_container_image_uri, + rule_configs as rule_configs, +) +from sagemaker.debugger.framework_profile import FrameworkProfile as FrameworkProfile +from sagemaker.debugger.metrics_config import ( + DataloaderProfilingConfig as DataloaderProfilingConfig, + DetailedProfilingConfig as DetailedProfilingConfig, + HorovodProfilingConfig as HorovodProfilingConfig, + PythonProfilingConfig as PythonProfilingConfig, + SMDataParallelProfilingConfig as SMDataParallelProfilingConfig, +) +from sagemaker.debugger.profiler_config import ProfilerConfig as ProfilerConfig +from sagemaker.debugger.utils import PythonProfiler as PythonProfiler, cProfileTimer as cProfileTimer diff --git a/stubs/sagemaker/sagemaker/debugger/debugger.pyi b/stubs/sagemaker/sagemaker/debugger/debugger.pyi new file mode 100644 index 000000000000..ae44e407fff1 --- /dev/null +++ b/stubs/sagemaker/sagemaker/debugger/debugger.pyi @@ -0,0 +1,129 @@ +from _typeshed import Incomplete +from abc import ABC +from typing import Dict, List, Optional, Union + +from sagemaker.workflow.entities import PipelineVariable + +framework_name: str +DEBUGGER_FLAG: str + +def get_rule_container_image_uri(region): ... +def get_default_profiler_rule(): ... + +class RuleBase(ABC): + name: Incomplete + image_uri: Incomplete + instance_type: Incomplete + container_local_output_path: Incomplete + s3_output_path: Incomplete + volume_size_in_gb: Incomplete + rule_parameters: Incomplete + def __init__( + self, name, image_uri, instance_type, container_local_output_path, s3_output_path, volume_size_in_gb, rule_parameters + ) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class Rule(RuleBase): + collection_configs: Incomplete + actions: Incomplete + def __init__( + self, + name, + image_uri, + instance_type, + container_local_output_path, + s3_output_path, + volume_size_in_gb, + rule_parameters, + collections_to_save, + actions: Incomplete | None = None, + ) -> None: ... + @classmethod + def sagemaker( + cls, + base_config, + name: Incomplete | None = None, + container_local_output_path: Incomplete | None = None, + s3_output_path: Incomplete | None = None, + other_trials_s3_input_paths: Incomplete | None = None, + rule_parameters: Incomplete | None = None, + collections_to_save: Incomplete | None = None, + actions: Incomplete | None = None, + ): ... + @classmethod + def custom( + cls, + name: str, + image_uri: Union[str, PipelineVariable], + instance_type: Union[str, PipelineVariable], + volume_size_in_gb: Union[int, PipelineVariable], + source: Optional[str] = None, + rule_to_invoke: Optional[Union[str, PipelineVariable]] = None, + container_local_output_path: Optional[Union[str, PipelineVariable]] = None, + s3_output_path: Optional[Union[str, PipelineVariable]] = None, + other_trials_s3_input_paths: Optional[List[Union[str, PipelineVariable]]] = None, + rule_parameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + collections_to_save: Optional[List["CollectionConfig"]] = None, + actions: Incomplete | None = None, + ): ... + def prepare_actions(self, training_job_name) -> None: ... + def to_debugger_rule_config_dict(self): ... + +class ProfilerRule(RuleBase): + @classmethod + def sagemaker( + cls, + base_config, + name: Incomplete | None = None, + container_local_output_path: Incomplete | None = None, + s3_output_path: Incomplete | None = None, + ): ... + @classmethod + def custom( + cls, + name, + image_uri, + instance_type, + volume_size_in_gb, + source: Incomplete | None = None, + rule_to_invoke: Incomplete | None = None, + container_local_output_path: Incomplete | None = None, + s3_output_path: Incomplete | None = None, + rule_parameters: Incomplete | None = None, + ): ... + def to_profiler_rule_config_dict(self): ... + +class DebuggerHookConfig: + s3_output_path: Incomplete + container_local_output_path: Incomplete + hook_parameters: Incomplete + collection_configs: Incomplete + def __init__( + self, + s3_output_path: Optional[Union[str, PipelineVariable]] = None, + container_local_output_path: Optional[Union[str, PipelineVariable]] = None, + hook_parameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + collection_configs: Optional[List["CollectionConfig"]] = None, + ) -> None: ... + +class TensorBoardOutputConfig: + s3_output_path: Incomplete + container_local_output_path: Incomplete + def __init__( + self, + s3_output_path: Union[str, PipelineVariable], + container_local_output_path: Optional[Union[str, PipelineVariable]] = None, + ) -> None: ... + +class CollectionConfig: + name: Incomplete + parameters: Incomplete + def __init__( + self, name: Union[str, PipelineVariable], parameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None + ) -> None: ... + def __eq__(self, other): ... + def __ne__(self, other): ... + def __hash__(self): ... diff --git a/stubs/sagemaker/sagemaker/debugger/framework_profile.pyi b/stubs/sagemaker/sagemaker/debugger/framework_profile.pyi new file mode 100644 index 000000000000..220d7e3233a5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/debugger/framework_profile.pyi @@ -0,0 +1,22 @@ +from _typeshed import Incomplete + +ALL_METRIC_CONFIGS: Incomplete + +class FrameworkProfile: + profiling_parameters: Incomplete + def __init__( + self, + local_path="/opt/ml/output/profiler", + file_max_size=10485760, + file_close_interval=60, + file_open_fail_threshold=50, + detailed_profiling_config: Incomplete | None = None, + dataloader_profiling_config: Incomplete | None = None, + python_profiling_config: Incomplete | None = None, + horovod_profiling_config: Incomplete | None = None, + smdataparallel_profiling_config: Incomplete | None = None, + start_step: Incomplete | None = None, + num_steps: Incomplete | None = None, + start_unix_time: Incomplete | None = None, + duration: Incomplete | None = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/debugger/metrics_config.pyi b/stubs/sagemaker/sagemaker/debugger/metrics_config.pyi new file mode 100644 index 000000000000..03c54951df20 --- /dev/null +++ b/stubs/sagemaker/sagemaker/debugger/metrics_config.pyi @@ -0,0 +1,75 @@ +from _typeshed import Incomplete + +class StepRange: + start_step: Incomplete + num_steps: Incomplete + def __init__(self, start_step, num_steps) -> None: ... + def to_json(self): ... + +class TimeRange: + start_unix_time: Incomplete + duration: Incomplete + def __init__(self, start_unix_time, duration) -> None: ... + def to_json(self): ... + +class MetricsConfigBase: + name: Incomplete + range: Incomplete + def __init__(self, name, start_step, num_steps, start_unix_time, duration) -> None: ... + def to_json_string(self): ... + +class DetailedProfilingConfig(MetricsConfigBase): + def __init__( + self, + start_step: Incomplete | None = None, + num_steps: Incomplete | None = None, + start_unix_time: Incomplete | None = None, + duration: Incomplete | None = None, + profile_default_steps: bool = False, + ) -> None: ... + +class DataloaderProfilingConfig(MetricsConfigBase): + metrics_regex: Incomplete + def __init__( + self, + start_step: Incomplete | None = None, + num_steps: Incomplete | None = None, + start_unix_time: Incomplete | None = None, + duration: Incomplete | None = None, + profile_default_steps: bool = False, + metrics_regex: str = ".*", + ) -> None: ... + +class PythonProfilingConfig(MetricsConfigBase): + python_profiler: Incomplete + cprofile_timer: Incomplete + def __init__( + self, + start_step: Incomplete | None = None, + num_steps: Incomplete | None = None, + start_unix_time: Incomplete | None = None, + duration: Incomplete | None = None, + profile_default_steps: bool = False, + python_profiler=..., + cprofile_timer=..., + ) -> None: ... + +class HorovodProfilingConfig(MetricsConfigBase): + def __init__( + self, + start_step: Incomplete | None = None, + num_steps: Incomplete | None = None, + start_unix_time: Incomplete | None = None, + duration: Incomplete | None = None, + profile_default_steps: bool = False, + ) -> None: ... + +class SMDataParallelProfilingConfig(MetricsConfigBase): + def __init__( + self, + start_step: Incomplete | None = None, + num_steps: Incomplete | None = None, + start_unix_time: Incomplete | None = None, + duration: Incomplete | None = None, + profile_default_steps: bool = False, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi b/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi new file mode 100644 index 000000000000..f5e7cb8f2dec --- /dev/null +++ b/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi @@ -0,0 +1,20 @@ +from _typeshed import Incomplete +from typing import Optional, Union + +from sagemaker.debugger.framework_profile import FrameworkProfile +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class ProfilerConfig: + s3_output_path: Incomplete + system_monitor_interval_millis: Incomplete + framework_profile_params: Incomplete + disable_profiler: Incomplete + def __init__( + self, + s3_output_path: Optional[Union[str, PipelineVariable]] = None, + system_monitor_interval_millis: Optional[Union[int, PipelineVariable]] = None, + framework_profile_params: Optional[FrameworkProfile] = None, + disable_profiler: Optional[Union[str, PipelineVariable]] = False, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/debugger/profiler_constants.pyi b/stubs/sagemaker/sagemaker/debugger/profiler_constants.pyi new file mode 100644 index 000000000000..ae4cfe6ea21e --- /dev/null +++ b/stubs/sagemaker/sagemaker/debugger/profiler_constants.pyi @@ -0,0 +1,17 @@ +BASE_FOLDER_DEFAULT: str +MAX_FILE_SIZE_DEFAULT: int +CLOSE_FILE_INTERVAL_DEFAULT: int +FILE_OPEN_FAIL_THRESHOLD_DEFAULT: int +DETAILED_PROFILING_CONFIG_NAME: str +DATALOADER_PROFILING_CONFIG_NAME: str +PYTHON_PROFILING_CONFIG_NAME: str +HOROVOD_PROFILING_CONFIG_NAME: str +SMDATAPARALLEL_PROFILING_CONFIG_NAME: str +DETAILED_PROFILING_START_STEP_DEFAULT: int +DATALOADER_PROFILING_START_STEP_DEFAULT: int +PYTHON_PROFILING_START_STEP_DEFAULT: int +HOROVOD_PROFILING_START_STEP_DEFAULT: int +SMDATAPARALLEL_PROFILING_START_STEP_DEFAULT: int +PROFILING_NUM_STEPS_DEFAULT: int +START_STEP_DEFAULT: int +PYTHON_PROFILING_NUM_STEPS_DEFAULT: int diff --git a/stubs/sagemaker/sagemaker/debugger/utils.pyi b/stubs/sagemaker/sagemaker/debugger/utils.pyi new file mode 100644 index 000000000000..879b039a744b --- /dev/null +++ b/stubs/sagemaker/sagemaker/debugger/utils.pyi @@ -0,0 +1,30 @@ +from enum import Enum + +def convert_json_config_to_string(config): ... +def is_valid_unix_time(unix_time): ... +def is_valid_regex(regex): ... + +class ErrorMessages(Enum): + INVALID_LOCAL_PATH: str + INVALID_FILE_MAX_SIZE: str + INVALID_FILE_CLOSE_INTERVAL: str + INVALID_FILE_OPEN_FAIL_THRESHOLD: str + INVALID_PROFILE_DEFAULT_STEPS: str + INVALID_START_STEP: str + INVALID_NUM_STEPS: str + INVALID_START_UNIX_TIME: str + INVALID_DURATION: str + FOUND_BOTH_STEP_AND_TIME_FIELDS: str + INVALID_METRICS_REGEX: str + INVALID_PYTHON_PROFILER: str + INVALID_CPROFILE_TIMER: str + +class PythonProfiler(Enum): + CPROFILE: str + PYINSTRUMENT: str + +class cProfileTimer(Enum): + TOTAL_TIME: str + CPU_TIME: str + OFF_CPU_TIME: str + DEFAULT: str diff --git a/stubs/sagemaker/sagemaker/deprecations.pyi b/stubs/sagemaker/sagemaker/deprecations.pyi new file mode 100644 index 000000000000..455a12da8356 --- /dev/null +++ b/stubs/sagemaker/sagemaker/deprecations.pyi @@ -0,0 +1,19 @@ +from _typeshed import Incomplete + +logger: Incomplete +V2_URL: str + +def removed_warning(phrase, sdk_version: Incomplete | None = None) -> None: ... +def renamed_warning(phrase) -> None: ... +def deprecation_warn(name, date, msg: Incomplete | None = None) -> None: ... +def deprecation_warn_base(msg) -> None: ... +def deprecation_warning(date, msg: Incomplete | None = None): ... +def renamed_kwargs(old_name, new_name, value, kwargs): ... +def removed_arg(name, arg) -> None: ... +def removed_kwargs(name, kwargs) -> None: ... +def removed_function(name): ... +def deprecated(sdk_version: Incomplete | None = None): ... +def deprecated_function(func, name): ... +def deprecated_serialize(instance, name): ... +def deprecated_deserialize(instance, name): ... +def deprecated_class(cls, name): ... diff --git a/stubs/sagemaker/sagemaker/deserializers.pyi b/stubs/sagemaker/sagemaker/deserializers.pyi new file mode 100644 index 000000000000..58d61afe21d0 --- /dev/null +++ b/stubs/sagemaker/sagemaker/deserializers.pyi @@ -0,0 +1,30 @@ +from typing import List, Optional + +from sagemaker.base_deserializers import ( + BaseDeserializer, + BytesDeserializer as BytesDeserializer, + CSVDeserializer as CSVDeserializer, + DeferredError as DeferredError, + JSONDeserializer as JSONDeserializer, + JSONLinesDeserializer as JSONLinesDeserializer, + NumpyDeserializer as NumpyDeserializer, + PandasDeserializer as PandasDeserializer, + SimpleBaseDeserializer as SimpleBaseDeserializer, + StreamDeserializer as StreamDeserializer, + StringDeserializer as StringDeserializer, +) + +def retrieve_options( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> List[BaseDeserializer]: ... +def retrieve_default( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> BaseDeserializer: ... diff --git a/stubs/sagemaker/sagemaker/djl_inference/__init__.pyi b/stubs/sagemaker/sagemaker/djl_inference/__init__.pyi new file mode 100644 index 000000000000..4b57a35dc847 --- /dev/null +++ b/stubs/sagemaker/sagemaker/djl_inference/__init__.pyi @@ -0,0 +1,7 @@ +from sagemaker.djl_inference.model import ( + DeepSpeedModel as DeepSpeedModel, + DJLModel as DJLModel, + DJLPredictor as DJLPredictor, + FasterTransformerModel as FasterTransformerModel, + HuggingFaceAccelerateModel as HuggingFaceAccelerateModel, +) diff --git a/stubs/sagemaker/sagemaker/djl_inference/defaults.pyi b/stubs/sagemaker/sagemaker/djl_inference/defaults.pyi new file mode 100644 index 000000000000..d565a3cf7c18 --- /dev/null +++ b/stubs/sagemaker/sagemaker/djl_inference/defaults.pyi @@ -0,0 +1,8 @@ +from _typeshed import Incomplete + +STABLE_DIFFUSION_MODEL_TYPE: str +VALID_MODEL_CONFIG_FILES: Incomplete +DEEPSPEED_RECOMMENDED_ARCHITECTURES: Incomplete +FASTER_TRANSFORMER_RECOMMENDED_ARCHITECTURES: Incomplete +FASTER_TRANSFORMER_SUPPORTED_ARCHITECTURES: Incomplete +ALLOWED_INSTANCE_FAMILIES: Incomplete diff --git a/stubs/sagemaker/sagemaker/djl_inference/model.pyi b/stubs/sagemaker/sagemaker/djl_inference/model.pyi new file mode 100644 index 000000000000..433d3495c1a7 --- /dev/null +++ b/stubs/sagemaker/sagemaker/djl_inference/model.pyi @@ -0,0 +1,168 @@ +from _typeshed import Incomplete +from enum import Enum +from typing import Any, Dict, Optional, Union + +from sagemaker import Predictor +from sagemaker.deserializers import BaseDeserializer +from sagemaker.model import FrameworkModel +from sagemaker.serializers import BaseSerializer +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class DJLServingEngineEntryPointDefaults(Enum): + DEEPSPEED: Incomplete + HUGGINGFACE_ACCELERATE: Incomplete + STABLE_DIFFUSION: Incomplete + FASTER_TRANSFORMER: Incomplete + +class DJLPredictor(Predictor): + def __init__( + self, + endpoint_name: str, + sagemaker_session: Session = None, + serializer: BaseSerializer = ..., + deserializer: BaseDeserializer = ..., + ) -> None: ... + +class DJLModel(FrameworkModel): + def __new__(cls, model_id: str, *args, **kwargs): ... + model_id: Incomplete + djl_version: Incomplete + task: Incomplete + dtype: Incomplete + number_of_partitions: Incomplete + min_workers: Incomplete + max_workers: Incomplete + job_queue_size: Incomplete + parallel_loading: Incomplete + model_loading_timeout: Incomplete + prediction_timeout: Incomplete + sagemaker_session: Incomplete + save_mp_checkpoint_path: Incomplete + def __init__( + self, + model_id: str, + role: str, + djl_version: Optional[str] = None, + task: Optional[str] = None, + dtype: str = "fp32", + number_of_partitions: Optional[int] = None, + min_workers: Optional[int] = None, + max_workers: Optional[int] = None, + job_queue_size: Optional[int] = None, + parallel_loading: bool = False, + model_loading_timeout: Optional[int] = None, + prediction_timeout: Optional[int] = None, + entry_point: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + predictor_cls: callable = ..., + **kwargs, + ) -> None: ... + def package_for_edge(self, **_) -> None: ... + def compile(self, **_) -> None: ... + def transformer(self, **_) -> None: ... + def right_size(self, **_) -> None: ... + image_uri: Incomplete + def partition( + self, + instance_type: str, + s3_output_uri: str = None, + job_name: Optional[str] = None, + volume_kms_key: Optional[str] = None, + output_kms_key: Optional[str] = None, + use_spot_instances: bool = False, + max_wait: int = None, + enable_network_isolation: bool = False, + ): ... + def deploy( + self, + instance_type, + initial_instance_count: int = 1, + serializer: Incomplete | None = None, + deserializer: Incomplete | None = None, + endpoint_name: Incomplete | None = None, + tags: Incomplete | None = None, + kms_key: Incomplete | None = None, + wait: bool = True, + data_capture_config: Incomplete | None = None, + volume_size: Incomplete | None = None, + model_data_download_timeout: Incomplete | None = None, + container_startup_health_check_timeout: Incomplete | None = None, + ): ... + def prepare_container_def( + self, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... + def generate_serving_properties(self, serving_properties: Incomplete | None = None) -> Dict[str, str]: ... + def serving_image_uri(self, region_name): ... + +class DeepSpeedModel(DJLModel): + number_of_partitions: Incomplete + max_tokens: Incomplete + low_cpu_mem_usage: Incomplete + enable_cuda_graph: Incomplete + triangular_masking: Incomplete + return_tuple: Incomplete + save_mp_checkpoint_path: Incomplete + checkpoint: Incomplete + def __init__( + self, + model_id: str, + role: str, + tensor_parallel_degree: Optional[int] = None, + max_tokens: Optional[int] = None, + low_cpu_mem_usage: bool = False, + enable_cuda_graph: bool = False, + triangular_masking: bool = True, + return_tuple: bool = True, + **kwargs, + ) -> None: ... + def generate_serving_properties(self, serving_properties: Incomplete | None = None) -> Dict[str, Any]: ... + def partition( + self, + instance_type: str, + s3_output_uri: str = None, + job_name: Optional[str] = None, + volume_kms_key: Optional[str] = None, + output_kms_key: Optional[str] = None, + use_spot_instances: bool = False, + max_wait: int = None, + enable_network_isolation: bool = False, + ): ... + +class HuggingFaceAccelerateModel(DJLModel): + device_id: Incomplete + device_map: Incomplete + load_in_8bit: Incomplete + low_cpu_mem_usage: Incomplete + def __init__( + self, + model_id: str, + role: str, + number_of_partitions: Optional[int] = None, + device_id: Optional[int] = None, + device_map: Optional[Union[str, Dict[str, str]]] = None, + load_in_8bit: bool = False, + low_cpu_mem_usage: bool = False, + **kwargs, + ) -> None: ... + def generate_serving_properties(self, serving_properties: Incomplete | None = None) -> Dict[str, str]: ... + def partition( + self, + instance_type: str, + s3_output_uri: str = None, + job_name: Optional[str] = None, + volume_kms_key: Optional[str] = None, + output_kms_key: Optional[str] = None, + use_spot_instances: bool = False, + max_wait: int = None, + enable_network_isolation: bool = False, + ): ... + +class FasterTransformerModel(DJLModel): + number_of_partitions: Incomplete + def __init__(self, model_id: str, role: str, tensor_parallel_degree: Optional[int] = None, **kwargs) -> None: ... diff --git a/stubs/sagemaker/sagemaker/drift_check_baselines.pyi b/stubs/sagemaker/sagemaker/drift_check_baselines.pyi new file mode 100644 index 000000000000..f38f950ec08c --- /dev/null +++ b/stubs/sagemaker/sagemaker/drift_check_baselines.pyi @@ -0,0 +1,27 @@ +from _typeshed import Incomplete +from typing import Optional + +from sagemaker.model_metrics import FileSource, MetricsSource + +class DriftCheckBaselines: + model_statistics: Incomplete + model_constraints: Incomplete + model_data_statistics: Incomplete + model_data_constraints: Incomplete + bias_config_file: Incomplete + bias_pre_training_constraints: Incomplete + bias_post_training_constraints: Incomplete + explainability_constraints: Incomplete + explainability_config_file: Incomplete + def __init__( + self, + 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, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/environment_variables.pyi b/stubs/sagemaker/sagemaker/environment_variables.pyi new file mode 100644 index 000000000000..83c791d7d165 --- /dev/null +++ b/stubs/sagemaker/sagemaker/environment_variables.pyi @@ -0,0 +1,13 @@ +from _typeshed import Incomplete +from typing import Dict, Optional + +logger: Incomplete + +def retrieve_default( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, + include_aws_sdk_env_vars: bool = True, +) -> Dict[str, str]: ... diff --git a/stubs/sagemaker/sagemaker/estimator.pyi b/stubs/sagemaker/sagemaker/estimator.pyi new file mode 100644 index 000000000000..936972d060e5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/estimator.pyi @@ -0,0 +1,393 @@ +import abc +from _typeshed import Incomplete +from abc import abstractmethod +from typing import Dict, List, Optional, Union + +from sagemaker.debugger import ( + DEBUGGER_FLAG as DEBUGGER_FLAG, + DebuggerHookConfig, + ProfilerConfig, + RuleBase, + TensorBoardOutputConfig, + get_default_profiler_rule as get_default_profiler_rule, +) +from sagemaker.inputs import FileSystemInput, TrainingInput +from sagemaker.instance_group import InstanceGroup +from sagemaker.job import _Job +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class EstimatorBase(metaclass=abc.ABCMeta): + LAUNCH_PT_XLA_ENV_NAME: str + LAUNCH_PS_ENV_NAME: str + LAUNCH_MPI_ENV_NAME: str + LAUNCH_SM_DDP_ENV_NAME: str + LAUNCH_MWMS_ENV_NAME: str + INSTANCE_TYPE: str + MPI_NUM_PROCESSES_PER_HOST: str + MPI_CUSTOM_MPI_OPTIONS: str + SM_DDP_CUSTOM_MPI_OPTIONS: str + CONTAINER_CODE_CHANNEL_SOURCEDIR_PATH: str + JOB_CLASS_NAME: str + instance_count: Incomplete + instance_type: Incomplete + keep_alive_period_in_seconds: Incomplete + instance_groups: Incomplete + volume_size: Incomplete + max_run: Incomplete + input_mode: Incomplete + metric_definitions: Incomplete + model_uri: Incomplete + model_channel_name: Incomplete + code_uri: Incomplete + code_channel_name: str + source_dir: Incomplete + git_config: Incomplete + container_log_level: Incomplete + code_location: Incomplete + entry_point: Incomplete + dependencies: Incomplete + uploaded_code: Incomplete + tags: Incomplete + sagemaker_session: Incomplete + base_job_name: Incomplete + output_path: Incomplete + latest_training_job: Incomplete + jobs: Incomplete + deploy_instance_type: Incomplete + role: Incomplete + output_kms_key: Incomplete + volume_kms_key: Incomplete + subnets: Incomplete + security_group_ids: Incomplete + training_repository_access_mode: Incomplete + training_repository_credentials_provider_arn: Incomplete + container_entry_point: Incomplete + container_arguments: Incomplete + encrypt_inter_container_traffic: Incomplete + use_spot_instances: Incomplete + max_wait: Incomplete + checkpoint_s3_uri: Incomplete + checkpoint_local_path: Incomplete + rules: Incomplete + debugger_hook_config: Incomplete + tensorboard_output_config: Incomplete + debugger_rule_configs: Incomplete + collection_configs: Incomplete + enable_sagemaker_metrics: Incomplete + profiler_config: Incomplete + disable_profiler: Incomplete + environment: Incomplete + max_retry_attempts: Incomplete + profiler_rule_configs: Incomplete + profiler_rules: Incomplete + debugger_rules: Incomplete + disable_output_compression: Incomplete + def __init__( + self, + role: str = None, + instance_count: Optional[Union[int, PipelineVariable]] = None, + instance_type: Optional[Union[str, PipelineVariable]] = None, + keep_alive_period_in_seconds: Optional[Union[int, PipelineVariable]] = None, + volume_size: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + max_run: Union[int, PipelineVariable] = 86400, + 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] = None, + 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] = None, + profiler_config: Optional[ProfilerConfig] = None, + disable_profiler: bool = None, + environment: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + max_retry_attempts: Optional[Union[int, PipelineVariable]] = None, + source_dir: Optional[Union[str, PipelineVariable]] = None, + git_config: Optional[Dict[str, str]] = None, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + container_log_level: Union[int, PipelineVariable] = 20, + code_location: Optional[str] = None, + entry_point: Optional[Union[str, PipelineVariable]] = None, + dependencies: Optional[List[Union[str]]] = None, + instance_groups: Optional[List[InstanceGroup]] = None, + training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None, + training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None, + container_entry_point: Optional[List[str]] = None, + container_arguments: Optional[List[str]] = None, + disable_output_compression: bool = False, + **kwargs, + ) -> None: ... + @abstractmethod + def training_image_uri(self): ... + @abstractmethod + def hyperparameters(self): ... + def enable_network_isolation(self): ... + def prepare_workflow_for_training(self, job_name: Incomplete | None = None) -> None: ... + def latest_job_debugger_artifacts_path(self): ... + def latest_job_tensorboard_artifacts_path(self): ... + def latest_job_profiler_artifacts_path(self): ... + def fit( + self, + inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None, + wait: bool = True, + logs: str = "All", + job_name: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + ): ... + def compile_model( + self, + target_instance_family, + input_shape, + output_path, + framework: Incomplete | None = None, + framework_version: Incomplete | None = None, + compile_max_run=900, + tags: Incomplete | None = None, + target_platform_os: Incomplete | None = None, + target_platform_arch: Incomplete | None = None, + target_platform_accelerator: Incomplete | None = None, + compiler_options: Incomplete | None = None, + **kwargs, + ): ... + @classmethod + def attach(cls, training_job_name, sagemaker_session: Incomplete | None = None, model_channel_name: str = "model"): ... + def logs(self) -> None: ... + def deploy( + self, + initial_instance_count: Incomplete | None = None, + instance_type: Incomplete | None = None, + serializer: Incomplete | None = None, + deserializer: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + endpoint_name: Incomplete | None = None, + use_compiled_model: bool = False, + wait: bool = True, + model_name: Incomplete | None = None, + kms_key: Incomplete | None = None, + data_capture_config: Incomplete | None = None, + tags: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + async_inference_config: Incomplete | None = None, + volume_size: Incomplete | None = None, + model_data_download_timeout: Incomplete | None = None, + container_startup_health_check_timeout: Incomplete | None = None, + inference_recommendation_id: Incomplete | None = None, + explainer_config: Incomplete | None = None, + **kwargs, + ): ... + def register( + self, + content_types, + response_types, + inference_instances: Incomplete | None = None, + transform_instances: Incomplete | None = None, + image_uri: Incomplete | None = None, + model_package_name: Incomplete | None = None, + model_package_group_name: Incomplete | None = None, + model_metrics: Incomplete | None = None, + metadata_properties: Incomplete | None = None, + marketplace_cert: bool = False, + approval_status: Incomplete | None = None, + description: Incomplete | None = None, + compile_model_family: Incomplete | None = None, + model_name: Incomplete | None = None, + drift_check_baselines: Incomplete | None = None, + customer_metadata_properties: Incomplete | None = None, + domain: Incomplete | None = None, + sample_payload_url: Incomplete | None = None, + task: Incomplete | None = None, + framework: Incomplete | None = None, + framework_version: Incomplete | None = None, + nearest_model_name: Incomplete | None = None, + data_input_configuration: Incomplete | None = None, + **kwargs, + ): ... + @property + def model_data(self): ... + @abstractmethod + def create_model(self, **kwargs): ... + def transformer( + self, + instance_count, + instance_type, + strategy: Incomplete | None = None, + assemble_with: Incomplete | None = None, + output_path: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + accept: Incomplete | None = None, + env: Incomplete | None = None, + max_concurrent_transforms: Incomplete | None = None, + max_payload: Incomplete | None = None, + tags: Incomplete | None = None, + role: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + enable_network_isolation: Incomplete | None = None, + model_name: Incomplete | None = None, + ): ... + @property + def training_job_analytics(self): ... + def get_vpc_config(self, vpc_config_override="VPC_CONFIG_DEFAULT"): ... + delete_endpoint: Incomplete + def enable_default_profiling(self) -> None: ... + def disable_profiling(self) -> None: ... + def update_profiler( + self, + rules: Incomplete | None = None, + system_monitor_interval_millis: Incomplete | None = None, + s3_output_path: Incomplete | None = None, + framework_profile_params: Incomplete | None = None, + disable_framework_metrics: bool = False, + ) -> None: ... + +class _TrainingJob(_Job): + @classmethod + def start_new(cls, estimator, inputs, experiment_config): ... + @classmethod + def update( + cls, + estimator, + profiler_rule_configs: Incomplete | None = None, + profiler_config: Incomplete | None = None, + resource_config: Incomplete | None = None, + ): ... + def wait(self, logs: str = "All") -> None: ... + def describe(self): ... + def rule_job_summary(self): ... + def stop(self) -> None: ... + +class Estimator(EstimatorBase): + image_uri: Incomplete + def __init__( + self, + image_uri: Union[str, PipelineVariable], + role: str = None, + instance_count: Optional[Union[int, PipelineVariable]] = None, + instance_type: Optional[Union[str, PipelineVariable]] = None, + keep_alive_period_in_seconds: Optional[Union[int, PipelineVariable]] = None, + volume_size: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + max_run: Union[int, PipelineVariable] = 86400, + 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, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = 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] = None, + 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, + enable_network_isolation: Union[bool, PipelineVariable] = None, + rules: Optional[List[RuleBase]] = None, + debugger_hook_config: Optional[Union[DebuggerHookConfig, bool]] = None, + tensorboard_output_config: Optional[TensorBoardOutputConfig] = None, + enable_sagemaker_metrics: Optional[Union[bool, PipelineVariable]] = None, + 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[Union[str, PipelineVariable]] = None, + git_config: Optional[Dict[str, str]] = None, + container_log_level: Union[int, PipelineVariable] = 20, + code_location: Optional[str] = None, + entry_point: Optional[Union[str, PipelineVariable]] = None, + dependencies: Optional[List[str]] = None, + instance_groups: Optional[List[InstanceGroup]] = None, + training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None, + training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None, + container_entry_point: Optional[List[str]] = None, + container_arguments: Optional[List[str]] = None, + disable_output_compression: bool = False, + **kwargs, + ) -> None: ... + def training_image_uri(self): ... + def set_hyperparameters(self, **kwargs) -> None: ... + def hyperparameters(self): ... + def create_model( + self, + role: Incomplete | None = None, + image_uri: Incomplete | None = None, + predictor_cls: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + **kwargs, + ): ... + +class Framework(EstimatorBase, metaclass=abc.ABCMeta): + entry_point: Incomplete + git_config: Incomplete + source_dir: Incomplete + dependencies: Incomplete + uploaded_code: Incomplete + container_log_level: Incomplete + code_location: Incomplete + image_uri: Incomplete + checkpoint_s3_uri: Incomplete + checkpoint_local_path: Incomplete + enable_sagemaker_metrics: Incomplete + def __init__( + self, + entry_point: Union[str, PipelineVariable], + source_dir: Optional[Union[str, PipelineVariable]] = None, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + container_log_level: Union[int, PipelineVariable] = 20, + code_location: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + dependencies: Optional[List[str]] = None, + enable_network_isolation: Union[bool, PipelineVariable] = None, + git_config: Optional[Dict[str, str]] = None, + checkpoint_s3_uri: Optional[Union[str, PipelineVariable]] = None, + checkpoint_local_path: Optional[Union[str, PipelineVariable]] = None, + enable_sagemaker_metrics: Optional[Union[bool, PipelineVariable]] = None, + **kwargs, + ) -> None: ... + def set_hyperparameters(self, **kwargs) -> None: ... + def hyperparameters(self): ... + def training_image_uri(self, region: Incomplete | None = None): ... + @classmethod + def attach(cls, training_job_name, sagemaker_session: Incomplete | None = None, model_channel_name: str = "model"): ... + def transformer( + self, + instance_count, + instance_type, + strategy: Incomplete | None = None, + assemble_with: Incomplete | None = None, + output_path: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + accept: Incomplete | None = None, + env: Incomplete | None = None, + max_concurrent_transforms: Incomplete | None = None, + max_payload: Incomplete | None = None, + tags: Incomplete | None = None, + role: Incomplete | None = None, + model_server_workers: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + entry_point: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + enable_network_isolation: Incomplete | None = None, + model_name: Incomplete | None = None, + ): ... diff --git a/stubs/sagemaker/sagemaker/exceptions.pyi b/stubs/sagemaker/sagemaker/exceptions.pyi new file mode 100644 index 000000000000..647f5e1c8d8b --- /dev/null +++ b/stubs/sagemaker/sagemaker/exceptions.pyi @@ -0,0 +1,34 @@ +from _typeshed import Incomplete + +class UnexpectedStatusException(ValueError): + allowed_statuses: Incomplete + actual_status: Incomplete + def __init__(self, message, allowed_statuses, actual_status) -> None: ... + +class CapacityError(UnexpectedStatusException): ... + +class AsyncInferenceError(Exception): + fmt: str + kwargs: Incomplete + def __init__(self, **kwargs) -> None: ... + +class ObjectNotExistedError(AsyncInferenceError): + fmt: str + def __init__(self, message, output_path) -> None: ... + +class PollingTimeoutError(AsyncInferenceError): + fmt: str + def __init__(self, message, output_path, seconds) -> None: ... + +class UnexpectedClientError(AsyncInferenceError): + fmt: str + def __init__(self, message) -> None: ... + +class AutoMLStepInvalidModeError(Exception): + fmt: str + kwargs: Incomplete + def __init__(self, **kwargs) -> None: ... + +class AsyncInferenceModelError(AsyncInferenceError): + fmt: str + def __init__(self, message) -> None: ... diff --git a/stubs/sagemaker/sagemaker/experiments/__init__.pyi b/stubs/sagemaker/sagemaker/experiments/__init__.pyi new file mode 100644 index 000000000000..906c4d9aa1ad --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/__init__.pyi @@ -0,0 +1,8 @@ +from sagemaker.experiments.experiment import Experiment as Experiment +from sagemaker.experiments.run import ( + Run as Run, + SortByType as SortByType, + SortOrderType as SortOrderType, + list_runs as list_runs, + load_run as load_run, +) diff --git a/stubs/sagemaker/sagemaker/experiments/_api_types.pyi b/stubs/sagemaker/sagemaker/experiments/_api_types.pyi new file mode 100644 index 000000000000..5ed52742ef74 --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/_api_types.pyi @@ -0,0 +1,85 @@ +import enum +from _typeshed import Incomplete + +from sagemaker.apiutils import _base_types + +class TrialComponentMetricSummary(_base_types.ApiObject): + metric_name: Incomplete + source_arn: Incomplete + time_stamp: Incomplete + max: Incomplete + min: Incomplete + last: Incomplete + count: Incomplete + avg: Incomplete + std_dev: Incomplete + def __init__(self, metric_name: Incomplete | None = None, source_arn: Incomplete | None = None, **kwargs) -> None: ... + +class TrialComponentParameters(_base_types.ApiObject): + @classmethod + def from_boto(cls, boto_dict, **kwargs): ... + @classmethod + def to_boto(cls, parameters): ... + +class TrialComponentArtifact(_base_types.ApiObject): + value: Incomplete + media_type: Incomplete + def __init__(self, value: Incomplete | None = None, media_type: Incomplete | None = None, **kwargs) -> None: ... + +class _TrialComponentStatusType(enum.Enum): + InProgress: str + Completed: str + Failed: str + +class TrialComponentStatus(_base_types.ApiObject): + primary_status: Incomplete + message: Incomplete + def __init__(self, primary_status: Incomplete | None = None, message: Incomplete | None = None, **kwargs) -> None: ... + +class TrialComponentSummary(_base_types.ApiObject): + trial_component_name: Incomplete + trial_component_arn: Incomplete + display_name: Incomplete + source_arn: Incomplete + status: Incomplete + start_time: Incomplete + end_time: Incomplete + creation_time: Incomplete + created_by: Incomplete + last_modified_time: Incomplete + last_modified_by: Incomplete + +class TrialComponentSource(_base_types.ApiObject): + source_arn: Incomplete + def __init__(self, source_arn: Incomplete | None = None, **kwargs) -> None: ... + +class Parent(_base_types.ApiObject): + trial_name: Incomplete + experiment_name: Incomplete + run_name: Incomplete + +class TrialComponentSearchResult(_base_types.ApiObject): + trial_component_arn: Incomplete + trial_component_name: Incomplete + display_name: Incomplete + source: Incomplete + status: Incomplete + start_time: Incomplete + end_time: Incomplete + creation_time: Incomplete + created_by: Incomplete + last_modified_time: Incomplete + last_modified_by: Incomplete + parameters: Incomplete + input_artifacts: Incomplete + output_artifacts: Incomplete + metrics: Incomplete + source_detail: Incomplete + tags: Incomplete + parents: Incomplete + +class TrialSummary(_base_types.ApiObject): + trial_arn: Incomplete + trial_name: Incomplete + creation_time: Incomplete + last_modified_time: Incomplete diff --git a/stubs/sagemaker/sagemaker/experiments/_environment.pyi b/stubs/sagemaker/sagemaker/experiments/_environment.pyi new file mode 100644 index 000000000000..476f5e37aeed --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/_environment.pyi @@ -0,0 +1,28 @@ +import enum +from _typeshed import Incomplete + +from sagemaker import Session + +TRAINING_JOB_ARN_ENV: str +PROCESSING_JOB_CONFIG_PATH: str +TRANSFORM_JOB_ARN_ENV: str +MAX_RETRY_ATTEMPTS: int +logger: Incomplete + +class _EnvironmentType(enum.Enum): + SageMakerTrainingJob: int + SageMakerProcessingJob: int + SageMakerTransformJob: int + +class _RunEnvironment: + environment_type: Incomplete + source_arn: Incomplete + def __init__(self, environment_type: _EnvironmentType, source_arn: str) -> None: ... + @classmethod + def load( + cls, + training_job_arn_env: str = "TRAINING_JOB_ARN", + processing_job_config_path: str = "/opt/ml/config/processingjobconfig.json", + transform_job_arn_env: str = "TRANSFORM_JOB_ARN", + ): ... + def get_trial_component(self, sagemaker_session: Session): ... diff --git a/stubs/sagemaker/sagemaker/experiments/_helper.pyi b/stubs/sagemaker/sagemaker/experiments/_helper.pyi new file mode 100644 index 000000000000..426c723fad10 --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/_helper.pyi @@ -0,0 +1,47 @@ +from _typeshed import Incomplete + +logger: Incomplete + +class _ArtifactUploader: + sagemaker_session: Incomplete + trial_component_name: Incomplete + artifact_bucket: Incomplete + artifact_prefix: Incomplete + def __init__( + self, + trial_component_name, + sagemaker_session, + artifact_bucket: Incomplete | None = None, + artifact_prefix="trial-component-artifacts", + ) -> None: ... + def upload_artifact(self, file_path): ... + def upload_object_artifact(self, artifact_name, artifact_object, file_extension: Incomplete | None = None): ... + +class _LineageArtifactManager: + name: Incomplete + source_uri: Incomplete + etag: Incomplete + source_arn: Incomplete + dest_arn: Incomplete + artifact_arn: Incomplete + artifact_type: Incomplete + def __init__( + self, + name, + source_uri, + etag, + source_arn: Incomplete | None = None, + dest_arn: Incomplete | None = None, + artifact_type="Tracker", + ) -> None: ... + def create_artifact(self, sagemaker_session) -> None: ... + def add_association(self, sagemaker_session) -> None: ... + +class _LineageArtifactTracker: + trial_component_arn: Incomplete + sagemaker_session: Incomplete + artifacts: Incomplete + def __init__(self, trial_component_arn, sagemaker_session) -> None: ... + def add_input_artifact(self, name, source_uri, etag, artifact_type) -> None: ... + def add_output_artifact(self, name, source_uri, etag, artifact_type) -> None: ... + def save(self) -> None: ... diff --git a/stubs/sagemaker/sagemaker/experiments/_metrics.pyi b/stubs/sagemaker/sagemaker/experiments/_metrics.pyi new file mode 100644 index 000000000000..e8f624775da0 --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/_metrics.pyi @@ -0,0 +1,42 @@ +from _typeshed import Incomplete + +from sagemaker.session import Session + +METRICS_DIR: Incomplete +METRIC_TS_LOWER_BOUND_TO_NOW: int +METRIC_TS_UPPER_BOUND_FROM_NOW: int +BATCH_SIZE: int +logger: Incomplete + +class _RawMetricData: + MetricName: Incomplete + Value: Incomplete + Timestamp: Incomplete + Step: Incomplete + def __init__(self, metric_name, value, timestamp: Incomplete | None = None, step: Incomplete | None = None) -> None: ... + def to_record(self): ... + def to_raw_metric_data(self): ... + +class _MetricsManager: + sink: Incomplete + def __init__(self, trial_component_name: str, sagemaker_session: Session, sink: Incomplete | None = None) -> None: ... + def log_metric(self, metric_name, value, timestamp: Incomplete | None = None, step: Incomplete | None = None) -> None: ... + def __enter__(self): ... + def __exit__(self, exc_type, exc_value, exc_traceback) -> None: ... + def close(self) -> None: ... + +class _SyncMetricsSink: + def __init__(self, trial_component_name, metrics_client) -> None: ... + def log_metric(self, metric_data) -> None: ... + def close(self) -> None: ... + +class _MetricQueue: + def __init__(self, trial_component_name, metric_name, metrics_client) -> None: ... + def log_metric(self, metric_data) -> None: ... + def close(self) -> None: ... + def is_active(self): ... + +class _AsyncMetricsSink: + def __init__(self, trial_component_name, metrics_client) -> None: ... + def log_metric(self, metric_data) -> None: ... + def close(self): ... diff --git a/stubs/sagemaker/sagemaker/experiments/_run_context.pyi b/stubs/sagemaker/sagemaker/experiments/_run_context.pyi new file mode 100644 index 000000000000..09e718d46c9a --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/_run_context.pyi @@ -0,0 +1,9 @@ +from sagemaker.experiments import Run + +class _RunContext: + @classmethod + def add_run_object(cls, run: Run): ... + @classmethod + def drop_current_run(cls) -> Run: ... + @classmethod + def get_current_run(cls) -> Run: ... diff --git a/stubs/sagemaker/sagemaker/experiments/_utils.pyi b/stubs/sagemaker/sagemaker/experiments/_utils.pyi new file mode 100644 index 000000000000..e7c0e80e2917 --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/_utils.pyi @@ -0,0 +1,13 @@ +from typing import Optional + +from sagemaker import Session +from sagemaker.experiments._environment import _RunEnvironment + +def resolve_artifact_name(file_path): ... +def guess_media_type(file_path): ... +def verify_length_of_true_and_predicted(true_labels, predicted_attrs, predicted_attrs_name) -> None: ... +def validate_invoked_inside_run_context(func): ... +def is_already_exist_error(error): ... +def get_tc_and_exp_config_from_job_env(environment: _RunEnvironment, sagemaker_session: Session) -> dict: ... +def verify_load_input_names(run_name: Optional[str] = None, experiment_name: Optional[str] = None): ... +def is_run_trial_component(trial_component_name: str, sagemaker_session: Session) -> bool: ... diff --git a/stubs/sagemaker/sagemaker/experiments/experiment.pyi b/stubs/sagemaker/sagemaker/experiments/experiment.pyi new file mode 100644 index 000000000000..cc1e62c35be8 --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/experiment.pyi @@ -0,0 +1,29 @@ +from _typeshed import Incomplete + +from sagemaker.apiutils import _base_types + +class Experiment(_base_types.Record): + experiment_name: Incomplete + display_name: Incomplete + description: Incomplete + tags: Incomplete + def save(self): ... + def delete(self): ... + @classmethod + def load(cls, experiment_name, sagemaker_session: Incomplete | None = None): ... + @classmethod + def create( + cls, + experiment_name, + display_name: Incomplete | None = None, + description: Incomplete | None = None, + tags: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + ): ... + def list_trials( + self, + created_before: Incomplete | None = None, + created_after: Incomplete | None = None, + sort_by: Incomplete | None = None, + sort_order: Incomplete | None = None, + ): ... diff --git a/stubs/sagemaker/sagemaker/experiments/run.pyi b/stubs/sagemaker/sagemaker/experiments/run.pyi new file mode 100644 index 000000000000..5ef5f9f6cfe6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/run.pyi @@ -0,0 +1,81 @@ +import datetime +from _typeshed import Incomplete +from enum import Enum +from typing import Dict, List, Optional, Union + +from numpy import array +from sagemaker import Session + +logger: Incomplete +RUN_NAME_BASE: Incomplete +TRIAL_NAME_TEMPLATE: str +MAX_RUN_TC_ARTIFACTS_LEN: int +MAX_NAME_LEN_IN_BACKEND: int +EXPERIMENT_NAME: str +TRIAL_NAME: str +RUN_NAME: str +DELIMITER: str +RUN_TC_TAG_KEY: str +RUN_TC_TAG_VALUE: str +RUN_TC_TAG: Incomplete + +class SortByType(Enum): + CREATION_TIME: str + NAME: str + +class SortOrderType(Enum): + ASCENDING: str + DESCENDING: str + +class Run: + experiment_name: Incomplete + run_name: Incomplete + run_group_name: Incomplete + def __init__( + self, + experiment_name: str, + run_name: Optional[str] = None, + experiment_display_name: Optional[str] = None, + run_display_name: Optional[str] = None, + tags: Optional[List[Dict[str, str]]] = None, + sagemaker_session: Optional["Session"] = None, + ) -> None: ... + @property + def experiment_config(self) -> dict: ... + def log_parameter(self, name: str, value: Union[str, int, float]): ... + def log_parameters(self, parameters: Dict[str, Union[str, int, float]]): ... + def log_metric(self, name: str, value: float, timestamp: Optional[datetime.datetime] = None, step: Optional[int] = None): ... + def log_precision_recall( + self, + y_true: Union[list, array], + predicted_probabilities: Union[list, array], + positive_label: Optional[Union[str, int]] = None, + title: Optional[str] = None, + is_output: bool = True, + no_skill: Optional[int] = None, + ): ... + def log_roc_curve( + self, y_true: Union[list, array], y_score: Union[list, array], title: Optional[str] = None, is_output: bool = True + ): ... + def log_confusion_matrix( + self, y_true: Union[list, array], y_pred: Union[list, array], title: Optional[str] = None, is_output: bool = True + ): ... + def log_artifact(self, name: str, value: str, media_type: Optional[str] = None, is_output: bool = True): ... + def log_file(self, file_path: str, name: Optional[str] = None, media_type: Optional[str] = None, is_output: bool = True): ... + def close(self) -> None: ... + def __enter__(self): ... + def __exit__(self, exc_type, exc_value, exc_traceback) -> None: ... + +def load_run( + run_name: Optional[str] = None, experiment_name: Optional[str] = None, sagemaker_session: Optional["Session"] = None +) -> Run: ... +def list_runs( + experiment_name: str, + created_before: Optional[datetime.datetime] = None, + created_after: Optional[datetime.datetime] = None, + sagemaker_session: Optional["Session"] = None, + max_results: Optional[int] = None, + next_token: Optional[str] = None, + sort_by: SortByType = ..., + sort_order: SortOrderType = ..., +) -> list: ... diff --git a/stubs/sagemaker/sagemaker/experiments/trial.pyi b/stubs/sagemaker/sagemaker/experiments/trial.pyi new file mode 100644 index 000000000000..a5c9c8baa639 --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/trial.pyi @@ -0,0 +1,44 @@ +from _typeshed import Incomplete + +from sagemaker.apiutils import _base_types + +class _Trial(_base_types.Record): + trial_name: Incomplete + experiment_name: Incomplete + display_name: Incomplete + tags: Incomplete + def save(self): ... + def delete(self): ... + @classmethod + def load(cls, trial_name, sagemaker_session: Incomplete | None = None): ... + @classmethod + def create( + cls, + experiment_name, + trial_name, + display_name: Incomplete | None = None, + tags: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + ): ... + @classmethod + def list( + cls, + experiment_name: Incomplete | None = None, + trial_component_name: Incomplete | None = None, + created_before: Incomplete | None = None, + created_after: Incomplete | None = None, + sort_by: Incomplete | None = None, + sort_order: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + ): ... + def add_trial_component(self, trial_component) -> None: ... + def remove_trial_component(self, trial_component) -> None: ... + def list_trial_components( + self, + created_before: Incomplete | None = None, + created_after: Incomplete | None = None, + sort_by: Incomplete | None = None, + sort_order: Incomplete | None = None, + max_results: Incomplete | None = None, + next_token: Incomplete | None = None, + ): ... diff --git a/stubs/sagemaker/sagemaker/experiments/trial_component.pyi b/stubs/sagemaker/sagemaker/experiments/trial_component.pyi new file mode 100644 index 000000000000..1fe68e7afd38 --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/trial_component.pyi @@ -0,0 +1,60 @@ +from _typeshed import Incomplete + +from sagemaker.apiutils import _base_types + +class _TrialComponent(_base_types.Record): + trial_component_name: Incomplete + trial_component_arn: Incomplete + display_name: Incomplete + source: Incomplete + status: Incomplete + start_time: Incomplete + end_time: Incomplete + creation_time: Incomplete + created_by: Incomplete + last_modified_time: Incomplete + last_modified_by: Incomplete + parameters: Incomplete + input_artifacts: Incomplete + output_artifacts: Incomplete + metrics: Incomplete + parameters_to_remove: Incomplete + input_artifacts_to_remove: Incomplete + output_artifacts_to_remove: Incomplete + tags: Incomplete + def __init__(self, sagemaker_session: Incomplete | None = None, **kwargs) -> None: ... + def save(self): ... + def delete(self, force_disassociate: bool = False): ... + @classmethod + def load(cls, trial_component_name, sagemaker_session: Incomplete | None = None): ... + @classmethod + def create( + cls, + trial_component_name, + display_name: Incomplete | None = None, + tags: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + ): ... + @classmethod + def list( + cls, + source_arn: Incomplete | None = None, + created_before: Incomplete | None = None, + created_after: Incomplete | None = None, + sort_by: Incomplete | None = None, + sort_order: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + trial_name: Incomplete | None = None, + experiment_name: Incomplete | None = None, + max_results: Incomplete | None = None, + next_token: Incomplete | None = None, + ): ... + @classmethod + def search( + cls, + search_expression: Incomplete | None = None, + sort_by: Incomplete | None = None, + sort_order: Incomplete | None = None, + max_results: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + ): ... diff --git a/stubs/sagemaker/sagemaker/explainer/__init__.pyi b/stubs/sagemaker/sagemaker/explainer/__init__.pyi new file mode 100644 index 000000000000..eaef91fdb3a7 --- /dev/null +++ b/stubs/sagemaker/sagemaker/explainer/__init__.pyi @@ -0,0 +1,8 @@ +from sagemaker.explainer.clarify_explainer_config import ( + ClarifyExplainerConfig as ClarifyExplainerConfig, + ClarifyInferenceConfig as ClarifyInferenceConfig, + ClarifyShapBaselineConfig as ClarifyShapBaselineConfig, + ClarifyShapConfig as ClarifyShapConfig, + ClarifyTextConfig as ClarifyTextConfig, +) +from sagemaker.explainer.explainer_config import ExplainerConfig as ExplainerConfig diff --git a/stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi b/stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi new file mode 100644 index 000000000000..5e43de15c1d2 --- /dev/null +++ b/stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi @@ -0,0 +1,68 @@ +from _typeshed import Incomplete +from typing import List, Optional + +class ClarifyTextConfig: + language: Incomplete + granularity: Incomplete + def __init__(self, language: str, granularity: str) -> None: ... + +class ClarifyShapBaselineConfig: + mime_type: Incomplete + shap_baseline: Incomplete + shap_baseline_uri: Incomplete + def __init__( + self, mime_type: Optional[str] = "text/csv", shap_baseline: Optional[str] = None, shap_baseline_uri: Optional[str] = None + ) -> None: ... + +class ClarifyShapConfig: + number_of_samples: Incomplete + seed: Incomplete + shap_baseline_config: Incomplete + text_config: Incomplete + use_logit: Incomplete + def __init__( + self, + shap_baseline_config: ClarifyShapBaselineConfig, + number_of_samples: Optional[int] = None, + seed: Optional[int] = None, + use_logit: Optional[bool] = False, + text_config: Optional[ClarifyTextConfig] = None, + ) -> None: ... + +class ClarifyInferenceConfig: + feature_headers: Incomplete + feature_types: Incomplete + features_attribute: Incomplete + probability_index: Incomplete + probability_attribute: Incomplete + label_index: Incomplete + label_attribute: Incomplete + label_headers: Incomplete + max_payload_in_mb: Incomplete + max_record_count: Incomplete + content_template: Incomplete + def __init__( + self, + feature_headers: Optional[List[str]] = None, + feature_types: Optional[List[str]] = None, + features_attribute: Optional[str] = None, + probability_index: Optional[int] = None, + probability_attribute: Optional[str] = None, + label_index: Optional[int] = None, + label_attribute: Optional[str] = None, + label_headers: Optional[List[str]] = None, + max_payload_in_mb: Optional[int] = 6, + max_record_count: Optional[int] = None, + content_template: Optional[str] = None, + ) -> None: ... + +class ClarifyExplainerConfig: + enable_explanations: Incomplete + shap_config: Incomplete + inference_config: Incomplete + def __init__( + self, + shap_config: ClarifyShapConfig, + enable_explanations: Optional[str] = None, + inference_config: Optional[ClarifyInferenceConfig] = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/explainer/explainer_config.pyi b/stubs/sagemaker/sagemaker/explainer/explainer_config.pyi new file mode 100644 index 000000000000..16930ed5fbd3 --- /dev/null +++ b/stubs/sagemaker/sagemaker/explainer/explainer_config.pyi @@ -0,0 +1,8 @@ +from _typeshed import Incomplete +from typing import Optional + +from sagemaker.explainer.clarify_explainer_config import ClarifyExplainerConfig + +class ExplainerConfig: + clarify_explainer_config: Incomplete + def __init__(self, clarify_explainer_config: Optional[ClarifyExplainerConfig] = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/feature_store/__init__.pyi b/stubs/sagemaker/sagemaker/feature_store/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi b/stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi new file mode 100644 index 000000000000..0fb4080f3dba --- /dev/null +++ b/stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi @@ -0,0 +1,128 @@ +import datetime +from enum import Enum +from typing import List, Tuple + +import pandas as pd +from sagemaker.feature_store.feature_group import FeatureDefinition, FeatureGroup + +class TableType(Enum): + FEATURE_GROUP: str + DATA_FRAME: str + def __init__(self) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class JoinTypeEnum(Enum): + INNER_JOIN: str + LEFT_JOIN: str + RIGHT_JOIN: str + FULL_JOIN: str + CROSS_JOIN: str + def __init__(self) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class JoinComparatorEnum(Enum): + EQUALS: str + GREATER_THAN: str + GREATER_THAN_OR_EQUAL_TO: str + LESS_THAN: str + LESS_THAN_OR_EQUAL_TO: str + NOT_EQUAL_TO: str + def __init__(self) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class FeatureGroupToBeMerged: + features: List[str] + included_feature_names: List[str] + projected_feature_names: List[str] + catalog: str + database: str + table_name: str + record_identifier_feature_name: str + event_time_identifier_feature: FeatureDefinition + target_feature_name_in_base: str + table_type: TableType + feature_name_in_target: str + join_comparator: JoinComparatorEnum + join_type: JoinTypeEnum + def __init__( + self, + features, + included_feature_names, + projected_feature_names, + catalog, + database, + table_name, + record_identifier_feature_name, + event_time_identifier_feature, + target_feature_name_in_base, + table_type, + feature_name_in_target, + join_comparator, + join_type, + ) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +def construct_feature_group_to_be_merged( + target_feature_group: FeatureGroup, + included_feature_names: List[str], + target_feature_name_in_base: str = None, + feature_name_in_target: str = None, + join_comparator: JoinComparatorEnum = ..., + join_type: JoinTypeEnum = ..., +) -> FeatureGroupToBeMerged: ... + +class DatasetBuilder: + def with_feature_group( + self, + feature_group: FeatureGroup, + target_feature_name_in_base: str = None, + included_feature_names: List[str] = None, + feature_name_in_target: str = None, + join_comparator: JoinComparatorEnum = ..., + join_type: JoinTypeEnum = ..., + ): ... + def point_in_time_accurate_join(self): ... + def include_duplicated_records(self): ... + def include_deleted_records(self): ... + def with_number_of_recent_records_by_record_identifier(self, number_of_recent_records: int): ... + def with_number_of_records_from_query_results(self, number_of_records: int): ... + def as_of(self, timestamp: datetime.datetime): ... + def with_event_time_range(self, starting_timestamp: datetime.datetime = None, ending_timestamp: datetime.datetime = None): ... + def to_csv_file(self) -> Tuple[str, str]: ... + def to_dataframe(self) -> Tuple[pd.DataFrame, str]: ... + def __init__( + self, + sagemaker_session, + base, + output_path, + record_identifier_feature_name, + event_time_identifier_feature_name, + included_feature_names, + kms_key_id, + point_in_time_accurate_join, + include_duplicated_records, + include_deleted_records, + number_of_recent_records, + number_of_records, + write_time_ending_timestamp, + event_time_starting_timestamp, + event_time_ending_timestamp, + feature_groups_to_be_merged, + event_time_identifier_feature_type, + ) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi new file mode 100644 index 000000000000..d01e343caa13 --- /dev/null +++ b/stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi @@ -0,0 +1,28 @@ +from enum import Enum +from typing import Any, Dict + +from sagemaker.feature_store.inputs import Config + +class FeatureTypeEnum(Enum): + FRACTIONAL: str + INTEGRAL: str + STRING: str + +class FeatureDefinition(Config): + feature_name: str + feature_type: FeatureTypeEnum + def to_dict(self) -> Dict[str, Any]: ... + def __init__(self, feature_name, feature_type) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class FractionalFeatureDefinition(FeatureDefinition): + def __init__(self, feature_name: str) -> None: ... + +class IntegralFeatureDefinition(FeatureDefinition): + def __init__(self, feature_name: str) -> None: ... + +class StringFeatureDefinition(FeatureDefinition): + def __init__(self, feature_name: str) -> None: ... diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi new file mode 100644 index 000000000000..e7ea84cd71b0 --- /dev/null +++ b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi @@ -0,0 +1,117 @@ +from _typeshed import Incomplete +from multiprocessing.pool import AsyncResult as AsyncResult +from typing import Any, Dict, List, Sequence, Union + +from botocore.config import Config as Config +from pandas import DataFrame as DataFrame +from sagemaker.feature_store.feature_definition import FeatureDefinition, FeatureTypeEnum +from sagemaker.feature_store.inputs import DataCatalogConfig, DeletionModeEnum, FeatureParameter, FeatureValue, TableFormatEnum +from sagemaker.session import Session + +logger: Incomplete + +class AthenaQuery: + catalog: str + database: str + table_name: str + sagemaker_session: Session + def run(self, query_string: str, output_location: str, kms_key: str = None, workgroup: str = None) -> str: ... + def wait(self) -> None: ... + def get_query_execution(self) -> Dict[str, Any]: ... + def as_dataframe(self) -> DataFrame: ... + def __init__( + self, catalog, database, table_name, sagemaker_session, current_query_execution_id, result_bucket, result_file_prefix + ) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class IngestionManagerPandas: + feature_group_name: str + sagemaker_fs_runtime_client_config: Config + sagemaker_session: Session + max_workers: int + max_processes: int + profile_name: str + @property + def failed_rows(self) -> List[int]: ... + def wait(self, timeout: Incomplete | None = None) -> None: ... + def run(self, data_frame: DataFrame, wait: bool = True, timeout: Incomplete | None = None): ... + def __init__( + self, + feature_group_name, + sagemaker_fs_runtime_client_config, + sagemaker_session, + max_workers, + max_processes, + profile_name, + async_result, + processing_pool, + failed_indices, + ) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class IngestionError(Exception): + failed_rows: Incomplete + message: Incomplete + def __init__(self, failed_rows, message) -> None: ... + +class FeatureGroup: + name: str + sagemaker_session: Session + feature_definitions: Sequence[FeatureDefinition] + DTYPE_TO_FEATURE_DEFINITION_CLS_MAP: Dict[str, FeatureTypeEnum] + def create( + self, + s3_uri: Union[str, bool], + record_identifier_name: str, + event_time_feature_name: str, + role_arn: str = None, + online_store_kms_key_id: str = None, + enable_online_store: bool = False, + offline_store_kms_key_id: str = None, + disable_glue_table_creation: bool = False, + data_catalog_config: DataCatalogConfig = None, + description: str = None, + tags: List[Dict[str, str]] = None, + table_format: TableFormatEnum = None, + ) -> Dict[str, Any]: ... + def delete(self) -> None: ... + def describe(self, next_token: str = None) -> Dict[str, Any]: ... + def update(self, feature_additions: Sequence[FeatureDefinition]) -> Dict[str, Any]: ... + def update_feature_metadata( + self, + feature_name: str, + description: str = None, + parameter_additions: Sequence[FeatureParameter] = None, + parameter_removals: Sequence[str] = None, + ) -> Dict[str, Any]: ... + def describe_feature_metadata(self, feature_name: str) -> Dict[str, Any]: ... + def list_tags(self) -> Sequence[Dict[str, str]]: ... + def list_parameters_for_feature_metadata(self, feature_name: str) -> Sequence[Dict[str, str]]: ... + def load_feature_definitions(self, data_frame: DataFrame) -> Sequence[FeatureDefinition]: ... + def get_record( + self, record_identifier_value_as_string: str, feature_names: Sequence[str] = None + ) -> Sequence[Dict[str, str]]: ... + def put_record(self, record: Sequence[FeatureValue]): ... + def delete_record(self, record_identifier_value_as_string: str, event_time: str, deletion_mode: DeletionModeEnum = ...): ... + def ingest( + self, + data_frame: DataFrame, + max_workers: int = 1, + max_processes: int = 1, + wait: bool = True, + timeout: Union[int, float] = None, + profile_name: str = None, + ) -> IngestionManagerPandas: ... + def athena_query(self) -> AthenaQuery: ... + def as_hive_ddl(self, database: str = "sagemaker_featurestore", table_name: str = None) -> str: ... + def __init__(self, name, sagemaker_session, feature_definitions) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_processor/__init__.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_processor/__init__.pyi new file mode 100644 index 000000000000..7caf3f51ac80 --- /dev/null +++ b/stubs/sagemaker/sagemaker/feature_store/feature_processor/__init__.pyi @@ -0,0 +1,16 @@ +from sagemaker.feature_store.feature_processor._data_source import ( + CSVDataSource as CSVDataSource, + FeatureGroupDataSource as FeatureGroupDataSource, + ParquetDataSource as ParquetDataSource, +) +from sagemaker.feature_store.feature_processor._exceptions import IngestionError as IngestionError +from sagemaker.feature_store.feature_processor.feature_processor import feature_processor as feature_processor +from sagemaker.feature_store.feature_processor.feature_scheduler import ( + TransformationCode as TransformationCode, + delete_schedule as delete_schedule, + describe as describe, + execute as execute, + list_pipelines as list_pipelines, + schedule as schedule, + to_pipeline as to_pipeline, +) diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi new file mode 100644 index 000000000000..57bc87875a9e --- /dev/null +++ b/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi @@ -0,0 +1,48 @@ +import datetime +from typing import Any, Dict, Sequence, Union + +import pandas as pd +from sagemaker import Session +from sagemaker.feature_store.dataset_builder import DatasetBuilder +from sagemaker.feature_store.feature_group import FeatureGroup +from sagemaker.feature_store.inputs import Filter, Identifier, ResourceEnum, SearchOperatorEnum, SortOrderEnum + +class FeatureStore: + sagemaker_session: Session + def create_dataset( + self, + base: Union[FeatureGroup, pd.DataFrame], + output_path: str, + record_identifier_feature_name: str = None, + event_time_identifier_feature_name: str = None, + included_feature_names: Sequence[str] = None, + kms_key_id: str = None, + ) -> DatasetBuilder: ... + def list_feature_groups( + self, + name_contains: str = None, + feature_group_status_equals: str = None, + offline_store_status_equals: str = None, + creation_time_after: datetime.datetime = None, + creation_time_before: datetime.datetime = None, + sort_order: str = None, + sort_by: str = None, + max_results: int = None, + next_token: str = None, + ) -> Dict[str, Any]: ... + def batch_get_record(self, identifiers: Sequence[Identifier]) -> Dict[str, Any]: ... + def search( + self, + resource: ResourceEnum, + filters: Sequence[Filter] = None, + operator: SearchOperatorEnum = None, + sort_by: str = None, + sort_order: SortOrderEnum = None, + next_token: str = None, + max_results: int = None, + ) -> Dict[str, Any]: ... + def __init__(self, sagemaker_session) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/feature_store/inputs.pyi b/stubs/sagemaker/sagemaker/feature_store/inputs.pyi new file mode 100644 index 000000000000..786077ba1569 --- /dev/null +++ b/stubs/sagemaker/sagemaker/feature_store/inputs.pyi @@ -0,0 +1,135 @@ +import abc +from enum import Enum +from typing import Any, Dict, List + +class Config(abc.ABC, metaclass=abc.ABCMeta): + @abc.abstractmethod + def to_dict(self) -> Dict[str, Any]: ... + @classmethod + def construct_dict(cls, **kwargs) -> Dict[str, Any]: ... + +class OnlineStoreSecurityConfig(Config): + kms_key_id: str + def to_dict(self) -> Dict[str, Any]: ... + def __init__(self, kms_key_id) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class OnlineStoreConfig(Config): + enable_online_store: bool + online_store_security_config: OnlineStoreSecurityConfig + def to_dict(self) -> Dict[str, Any]: ... + def __init__(self, enable_online_store, online_store_security_config) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class S3StorageConfig(Config): + s3_uri: str + kms_key_id: str + def to_dict(self) -> Dict[str, Any]: ... + def __init__(self, s3_uri, kms_key_id) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class DataCatalogConfig(Config): + table_name: str + catalog: str + database: str + def to_dict(self) -> Dict[str, Any]: ... + def __init__(self, table_name, catalog, database) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class TableFormatEnum(Enum): + GLUE: str + ICEBERG: str + +class OfflineStoreConfig(Config): + s3_storage_config: S3StorageConfig + disable_glue_table_creation: bool + data_catalog_config: DataCatalogConfig + table_format: TableFormatEnum + def to_dict(self) -> Dict[str, Any]: ... + def __init__(self, s3_storage_config, disable_glue_table_creation, data_catalog_config, table_format) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class FeatureValue(Config): + feature_name: str + value_as_string: str + def to_dict(self) -> Dict[str, Any]: ... + def __init__(self, feature_name, value_as_string) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class FeatureParameter(Config): + key: str + value: str + def to_dict(self) -> Dict[str, Any]: ... + def __init__(self, key, value) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class ResourceEnum(Enum): + FEATURE_GROUP: str + FEATURE_METADATA: str + +class SearchOperatorEnum(Enum): + AND: str + OR: str + +class SortOrderEnum(Enum): + ASCENDING: str + DESCENDING: str + +class FilterOperatorEnum(Enum): + EQUALS: str + NOT_EQUALS: str + GREATER_THAN: str + GREATER_THAN_OR_EQUAL_TO: str + LESS_THAN: str + LESS_THAN_OR_EQUAL_TO: str + CONTAINS: str + EXISTS: str + NOT_EXISTS: str + IN: str + +class Filter(Config): + name: str + value: str + operator: FilterOperatorEnum + def to_dict(self) -> Dict[str, Any]: ... + def __init__(self, name, value, operator) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class Identifier(Config): + feature_group_name: str + record_identifiers_value_as_string: List[str] + feature_names: List[str] + def to_dict(self) -> Dict[str, Any]: ... + def __init__(self, feature_group_name, record_identifiers_value_as_string, feature_names) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class DeletionModeEnum(Enum): + SOFT_DELETE: str + HARD_DELETE: str diff --git a/stubs/sagemaker/sagemaker/fw_utils.pyi b/stubs/sagemaker/sagemaker/fw_utils.pyi new file mode 100644 index 000000000000..936b80a9bd5b --- /dev/null +++ b/stubs/sagemaker/sagemaker/fw_utils.pyi @@ -0,0 +1,74 @@ +from _typeshed import Incomplete +from typing import Dict, NamedTuple, Optional, Union + +from sagemaker.session_settings import SessionSettings +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class UploadedCode(NamedTuple): + s3_prefix: Incomplete + script_name: Incomplete + +PYTHON_2_DEPRECATION_WARNING: str +PARAMETER_SERVER_MULTI_GPU_WARNING: str +DEBUGGER_UNSUPPORTED_REGIONS: Incomplete +PROFILER_UNSUPPORTED_REGIONS: Incomplete +SINGLE_GPU_INSTANCE_TYPES: Incomplete +SM_DATAPARALLEL_SUPPORTED_INSTANCE_TYPES: Incomplete +SM_DATAPARALLEL_SUPPORTED_FRAMEWORK_VERSIONS: Incomplete +PYTORCHDDP_SUPPORTED_FRAMEWORK_VERSIONS: Incomplete +TORCH_DISTRIBUTED_GPU_SUPPORTED_FRAMEWORK_VERSIONS: Incomplete +TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES: Incomplete +TRAINIUM_SUPPORTED_TORCH_DISTRIBUTED_FRAMEWORK_VERSIONS: Incomplete +SMDISTRIBUTED_SUPPORTED_STRATEGIES: Incomplete +GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY: Incomplete +GRAVITON_ALLOWED_FRAMEWORKS: Incomplete + +def validate_source_dir(script, directory): ... +def validate_source_code_input_against_pipeline_variables( + entry_point: Optional[Union[str, PipelineVariable]] = None, + source_dir: Optional[Union[str, PipelineVariable]] = None, + git_config: Optional[Dict[str, str]] = None, + enable_network_isolation: Union[bool, PipelineVariable] = False, +): ... +def parse_mp_parameters(params): ... +def get_mp_parameters(distribution): ... +def validate_mp_config(config) -> None: ... +def tar_and_upload_dir( + session, + bucket, + s3_key_prefix, + script, + directory: Incomplete | None = None, + dependencies: Incomplete | None = None, + kms_key: Incomplete | None = None, + s3_resource: Incomplete | None = None, + settings: Optional[SessionSettings] = None, +) -> UploadedCode: ... +def framework_name_from_image(image_uri): ... +def framework_version_from_tag(image_tag): ... +def model_code_key_prefix(code_location_key_prefix, model_name, image): ... +def warn_if_parameter_server_with_multi_gpu(training_instance_type, distribution) -> None: ... +def profiler_config_deprecation_warning(profiler_config, image_uri, framework_name, framework_version) -> None: ... +def validate_smdistributed( + instance_type, framework_name, framework_version, py_version, distribution, image_uri: Incomplete | None = None +) -> None: ... +def validate_distribution(distribution, instance_groups, framework_name, framework_version, py_version, image_uri, kwargs): ... +def validate_distribution_for_instance_type(instance_type, distribution) -> None: ... +def validate_pytorch_distribution(distribution, framework_name, framework_version, py_version, image_uri) -> None: ... +def validate_torch_distributed_distribution( + instance_type, distribution, framework_version, py_version, image_uri, entry_point +) -> None: ... +def python_deprecation_warning(framework, latest_supported_version): ... +def validate_version_or_image_args(framework_version, py_version, image_uri) -> None: ... +def create_image_uri( + region, + framework, + instance_type, + framework_version, + py_version: Incomplete | None = None, + account: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + optimized_families: Incomplete | None = None, +): ... diff --git a/stubs/sagemaker/sagemaker/git_utils.pyi b/stubs/sagemaker/sagemaker/git_utils.pyi new file mode 100644 index 000000000000..384519e513b9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/git_utils.pyi @@ -0,0 +1,3 @@ +from _typeshed import Incomplete + +def git_clone_repo(git_config, entry_point, source_dir: Incomplete | None = None, dependencies: Incomplete | None = None): ... diff --git a/stubs/sagemaker/sagemaker/huggingface/__init__.pyi b/stubs/sagemaker/sagemaker/huggingface/__init__.pyi new file mode 100644 index 000000000000..4fb1a4d8ab29 --- /dev/null +++ b/stubs/sagemaker/sagemaker/huggingface/__init__.pyi @@ -0,0 +1,5 @@ +from sagemaker.huggingface.estimator import HuggingFace as HuggingFace +from sagemaker.huggingface.llm_utils import get_huggingface_llm_image_uri as get_huggingface_llm_image_uri +from sagemaker.huggingface.model import HuggingFaceModel as HuggingFaceModel, HuggingFacePredictor as HuggingFacePredictor +from sagemaker.huggingface.processing import HuggingFaceProcessor as HuggingFaceProcessor +from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig as TrainingCompilerConfig diff --git a/stubs/sagemaker/sagemaker/huggingface/estimator.pyi b/stubs/sagemaker/sagemaker/huggingface/estimator.pyi new file mode 100644 index 000000000000..8e44c797b66a --- /dev/null +++ b/stubs/sagemaker/sagemaker/huggingface/estimator.pyi @@ -0,0 +1,46 @@ +from _typeshed import Incomplete +from typing import Dict, Optional, Union + +from sagemaker.estimator import Framework +from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class HuggingFace(Framework): + LAUNCH_PYTORCH_DDP_ENV_NAME: str + LAUNCH_TORCH_DISTRIBUTED_ENV_NAME: str + INSTANCE_TYPE_ENV_NAME: str + framework_version: Incomplete + py_version: Incomplete + tensorflow_version: Incomplete + pytorch_version: Incomplete + base_framework_name: Incomplete + base_framework_version: Incomplete + distribution: Incomplete + compiler_config: Incomplete + def __init__( + self, + py_version: str, + entry_point: Union[str, PipelineVariable], + transformers_version: Optional[str] = None, + tensorflow_version: Optional[str] = None, + pytorch_version: Optional[str] = None, + source_dir: Optional[Union[str, PipelineVariable]] = None, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + distribution: Optional[Dict] = None, + compiler_config: Optional[TrainingCompilerConfig] = None, + **kwargs, + ) -> None: ... + def hyperparameters(self): ... + def create_model( + self, + model_server_workers: Incomplete | None = None, + role: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + entry_point: Incomplete | None = None, + source_dir: Incomplete | None = None, + dependencies: Incomplete | None = None, + **kwargs, + ): ... diff --git a/stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi b/stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi new file mode 100644 index 000000000000..01ea51b1a4ed --- /dev/null +++ b/stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi @@ -0,0 +1,7 @@ +from typing import Optional + +from sagemaker.session import Session + +def get_huggingface_llm_image_uri( + backend: str, session: Optional[Session] = None, region: Optional[str] = None, version: Optional[str] = None +) -> str: ... diff --git a/stubs/sagemaker/sagemaker/huggingface/model.pyi b/stubs/sagemaker/sagemaker/huggingface/model.pyi new file mode 100644 index 000000000000..58f76111630c --- /dev/null +++ b/stubs/sagemaker/sagemaker/huggingface/model.pyi @@ -0,0 +1,99 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker import ModelMetrics +from sagemaker.drift_check_baselines import DriftCheckBaselines +from sagemaker.metadata_properties import MetadataProperties +from sagemaker.model import FrameworkModel +from sagemaker.predictor import Predictor +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class HuggingFacePredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +def fetch_framework_and_framework_version(tensorflow_version, pytorch_version): ... + +class HuggingFaceModel(FrameworkModel): + framework_version: Incomplete + pytorch_version: Incomplete + tensorflow_version: Incomplete + py_version: Incomplete + sagemaker_session: Incomplete + model_server_workers: Incomplete + def __init__( + self, + role: Optional[str] = None, + model_data: Optional[Union[str, PipelineVariable]] = None, + entry_point: Optional[str] = None, + transformers_version: Optional[str] = None, + tensorflow_version: Optional[str] = None, + pytorch_version: Optional[str] = None, + py_version: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + predictor_cls: callable = ..., + model_server_workers: Optional[Union[int, PipelineVariable]] = None, + **kwargs, + ) -> None: ... + image_uri: Incomplete + def deploy( + self, + initial_instance_count: Incomplete | None = None, + instance_type: Incomplete | None = None, + serializer: Incomplete | None = None, + deserializer: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + endpoint_name: Incomplete | None = None, + tags: Incomplete | None = None, + kms_key: Incomplete | None = None, + wait: bool = True, + data_capture_config: Incomplete | None = None, + async_inference_config: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + volume_size: Incomplete | None = None, + model_data_download_timeout: Incomplete | None = None, + container_startup_health_check_timeout: Incomplete | None = None, + inference_recommendation_id: Incomplete | None = None, + explainer_config: Incomplete | None = None, + **kwargs, + ): ... + def register( + self, + 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, + domain: Optional[Union[str, PipelineVariable]] = None, + sample_payload_url: Optional[Union[str, PipelineVariable]] = None, + task: Optional[Union[str, PipelineVariable]] = None, + framework: Optional[Union[str, PipelineVariable]] = None, + framework_version: Optional[Union[str, PipelineVariable]] = None, + nearest_model_name: Optional[Union[str, PipelineVariable]] = None, + data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + ): ... + def prepare_container_def( + self, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + inference_tool: Incomplete | None = None, + ): ... + def serving_image_uri( + self, + region_name, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + inference_tool: Incomplete | None = None, + ): ... diff --git a/stubs/sagemaker/sagemaker/huggingface/processing.pyi b/stubs/sagemaker/sagemaker/huggingface/processing.pyi new file mode 100644 index 000000000000..4c63d3d8e4f6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/huggingface/processing.pyi @@ -0,0 +1,35 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker.huggingface.estimator import HuggingFace +from sagemaker.network import NetworkConfig +from sagemaker.processing import FrameworkProcessor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class HuggingFaceProcessor(FrameworkProcessor): + estimator_cls = HuggingFace + pytorch_version: Incomplete + tensorflow_version: Incomplete + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Union[int, PipelineVariable] = None, + instance_type: Union[str, PipelineVariable] = None, + transformers_version: Optional[str] = None, + tensorflow_version: Optional[str] = None, + pytorch_version: Optional[str] = None, + py_version: str = "py36", + image_uri: Optional[Union[str, PipelineVariable]] = None, + command: Optional[List[str]] = None, + volume_size_in_gb: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + output_kms_key: Optional[Union[str, PipelineVariable]] = None, + code_location: Optional[str] = None, + max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + base_job_name: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + network_config: Optional[NetworkConfig] = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/huggingface/training_compiler/__init__.pyi b/stubs/sagemaker/sagemaker/huggingface/training_compiler/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi b/stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi new file mode 100644 index 000000000000..90146f5a7e03 --- /dev/null +++ b/stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi @@ -0,0 +1,14 @@ +from _typeshed import Incomplete +from typing import Union + +from sagemaker.training_compiler.config import TrainingCompilerConfig as BaseConfig +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class TrainingCompilerConfig(BaseConfig): + SUPPORTED_INSTANCE_CLASS_PREFIXES: Incomplete + SUPPORTED_INSTANCE_TYPES_WITH_EFA: Incomplete + def __init__(self, enabled: Union[bool, PipelineVariable] = True, debug: Union[bool, PipelineVariable] = False) -> None: ... + @classmethod + def validate(cls, estimator) -> None: ... diff --git a/stubs/sagemaker/sagemaker/hyperparameters.pyi b/stubs/sagemaker/sagemaker/hyperparameters.pyi new file mode 100644 index 000000000000..93a00067c749 --- /dev/null +++ b/stubs/sagemaker/sagemaker/hyperparameters.pyi @@ -0,0 +1,22 @@ +from _typeshed import Incomplete +from typing import Dict, Optional + +from sagemaker.jumpstart.enums import HyperparameterValidationMode + +logger: Incomplete + +def retrieve_default( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + include_container_hyperparameters: bool = False, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> Dict[str, str]: ... +def validate( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + hyperparameters: Optional[dict] = None, + validation_mode: HyperparameterValidationMode = ..., +) -> None: ... diff --git a/stubs/sagemaker/sagemaker/image_uris.pyi b/stubs/sagemaker/sagemaker/image_uris.pyi new file mode 100644 index 000000000000..82d5c0840152 --- /dev/null +++ b/stubs/sagemaker/sagemaker/image_uris.pyi @@ -0,0 +1,46 @@ +from _typeshed import Incomplete + +logger: Incomplete +ECR_URI_TEMPLATE: str +HUGGING_FACE_FRAMEWORK: str +HUGGING_FACE_LLM_FRAMEWORK: str +XGBOOST_FRAMEWORK: str +SKLEARN_FRAMEWORK: str +TRAINIUM_ALLOWED_FRAMEWORKS: str +INFERENCE_GRAVITON: str +DATA_WRANGLER_FRAMEWORK: str + +def retrieve( + framework, + region, + version: Incomplete | None = None, + py_version: Incomplete | None = None, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + image_scope: Incomplete | None = None, + container_version: Incomplete | None = None, + distribution: Incomplete | None = None, + base_framework_version: Incomplete | None = None, + training_compiler_config: Incomplete | None = None, + model_id: Incomplete | None = None, + model_version: Incomplete | None = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, + sdk_version: Incomplete | None = None, + inference_tool: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, +) -> str: ... +def config_for_framework(framework): ... +def get_training_image_uri( + region, + framework, + framework_version: Incomplete | None = None, + py_version: Incomplete | None = None, + image_uri: Incomplete | None = None, + distribution: Incomplete | None = None, + compiler_config: Incomplete | None = None, + tensorflow_version: Incomplete | None = None, + pytorch_version: Incomplete | None = None, + instance_type: Incomplete | None = None, +) -> str: ... +def get_base_python_image_uri(region, py_version: str = "310") -> str: ... diff --git a/stubs/sagemaker/sagemaker/inference_recommender/__init__.pyi b/stubs/sagemaker/sagemaker/inference_recommender/__init__.pyi new file mode 100644 index 000000000000..12cbfca30457 --- /dev/null +++ b/stubs/sagemaker/sagemaker/inference_recommender/__init__.pyi @@ -0,0 +1,4 @@ +from sagemaker.inference_recommender.inference_recommender_mixin import ( + ModelLatencyThreshold as ModelLatencyThreshold, + Phase as Phase, +) diff --git a/stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi b/stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi new file mode 100644 index 000000000000..42955314cf60 --- /dev/null +++ b/stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi @@ -0,0 +1,36 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional + +from sagemaker.parameter import CategoricalParameter + +INFERENCE_RECOMMENDER_FRAMEWORK_MAPPING: Incomplete +LOGGER: Incomplete + +class Phase: + to_json: Incomplete + def __init__(self, duration_in_seconds: int, initial_number_of_users: int, spawn_rate: int) -> None: ... + +class ModelLatencyThreshold: + to_json: Incomplete + def __init__(self, percentile: str, value_in_milliseconds: int) -> None: ... + +class InferenceRecommenderMixin: + inference_recommender_job_results: Incomplete + inference_recommendations: Incomplete + def right_size( + self, + sample_payload_url: str = None, + supported_content_types: List[str] = None, + supported_instance_types: List[str] = None, + job_name: str = None, + framework: str = None, + job_duration_in_seconds: int = None, + hyperparameter_ranges: List[Dict[str, CategoricalParameter]] = None, + phases: List[Phase] = None, + traffic_type: str = None, + max_invocations: int = None, + model_latency_thresholds: List[ModelLatencyThreshold] = None, + max_tests: int = None, + max_parallel_tests: int = None, + log_level: Optional[str] = "Verbose", + ): ... diff --git a/stubs/sagemaker/sagemaker/inputs.pyi b/stubs/sagemaker/sagemaker/inputs.pyi new file mode 100644 index 000000000000..e8967bac6e9e --- /dev/null +++ b/stubs/sagemaker/sagemaker/inputs.pyi @@ -0,0 +1,83 @@ +from _typeshed import Incomplete +from typing import List, Optional, Union + +from sagemaker.workflow.entities import PipelineVariable + +FILE_SYSTEM_TYPES: Incomplete +FILE_SYSTEM_ACCESS_MODES: Incomplete + +class TrainingInput: + config: Incomplete + def __init__( + self, + s3_data: Union[str, PipelineVariable], + distribution: Optional[Union[str, PipelineVariable]] = None, + compression: Optional[Union[str, PipelineVariable]] = None, + content_type: Optional[Union[str, PipelineVariable]] = None, + record_wrapping: Optional[Union[str, PipelineVariable]] = None, + s3_data_type: Union[str, PipelineVariable] = "S3Prefix", + instance_groups: Optional[List[Union[str, PipelineVariable]]] = None, + input_mode: Optional[Union[str, PipelineVariable]] = None, + attribute_names: Optional[List[Union[str, PipelineVariable]]] = None, + target_attribute_name: Optional[Union[str, PipelineVariable]] = None, + shuffle_config: Optional["ShuffleConfig"] = None, + ) -> None: ... + +class ShuffleConfig: + seed: Incomplete + def __init__(self, seed) -> None: ... + +class CreateModelInput: + instance_type: str + accelerator_type: str + def __init__(self, instance_type, accelerator_type) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class TransformInput: + data: str + data_type: str + content_type: str + compression_type: str + split_type: str + input_filter: str + output_filter: str + join_source: str + model_client_config: dict + batch_data_capture_config: dict + def __init__( + self, + data, + data_type, + content_type, + compression_type, + split_type, + input_filter, + output_filter, + join_source, + model_client_config, + batch_data_capture_config, + ) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class FileSystemInput: + config: Incomplete + def __init__( + self, + file_system_id, + file_system_type, + directory_path, + file_system_access_mode: str = "ro", + content_type: Incomplete | None = None, + ) -> None: ... + +class BatchDataCaptureConfig: + destination_s3_uri: Incomplete + kms_key_id: Incomplete + generate_inference_id: Incomplete + def __init__(self, destination_s3_uri: str, kms_key_id: str = None, generate_inference_id: bool = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/instance_group.pyi b/stubs/sagemaker/sagemaker/instance_group.pyi new file mode 100644 index 000000000000..872a8b375563 --- /dev/null +++ b/stubs/sagemaker/sagemaker/instance_group.pyi @@ -0,0 +1,12 @@ +from _typeshed import Incomplete + +class InstanceGroup: + instance_group_name: Incomplete + instance_type: Incomplete + instance_count: Incomplete + def __init__( + self, + instance_group_name: Incomplete | None = None, + instance_type: Incomplete | None = None, + instance_count: Incomplete | None = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/instance_types.pyi b/stubs/sagemaker/sagemaker/instance_types.pyi new file mode 100644 index 000000000000..f6d25ce7eede --- /dev/null +++ b/stubs/sagemaker/sagemaker/instance_types.pyi @@ -0,0 +1,21 @@ +from _typeshed import Incomplete +from typing import List, Optional + +logger: Incomplete + +def retrieve_default( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + scope: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> str: ... +def retrieve( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + scope: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> List[str]: ... diff --git a/stubs/sagemaker/sagemaker/job.pyi b/stubs/sagemaker/sagemaker/job.pyi new file mode 100644 index 000000000000..15601fb8c668 --- /dev/null +++ b/stubs/sagemaker/sagemaker/job.pyi @@ -0,0 +1,18 @@ +import abc +from _typeshed import Incomplete +from abc import abstractmethod + +class _Job(metaclass=abc.ABCMeta): + sagemaker_session: Incomplete + job_name: Incomplete + def __init__(self, sagemaker_session, job_name) -> None: ... + @abstractmethod + def start_new(self, estimator, inputs): ... + @abstractmethod + def wait(self): ... + @abstractmethod + def describe(self): ... + @abstractmethod + def stop(self): ... + @property + def name(self): ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/__init__.pyi b/stubs/sagemaker/sagemaker/jumpstart/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/accessors.pyi b/stubs/sagemaker/sagemaker/jumpstart/accessors.pyi new file mode 100644 index 000000000000..cf1367aad09f --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/accessors.pyi @@ -0,0 +1,21 @@ +from typing import Any, Dict, List, Optional + +from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartModelSpecs + +class SageMakerSettings: + @staticmethod + def set_sagemaker_version(version: str) -> None: ... + @staticmethod + def get_sagemaker_version() -> str: ... + +class JumpStartModelsAccessor: + @staticmethod + def get_model_header(region: str, model_id: str, version: str) -> JumpStartModelHeader: ... + @staticmethod + def get_model_specs(region: str, model_id: str, version: str) -> JumpStartModelSpecs: ... + @staticmethod + def set_cache_kwargs(cache_kwargs: Dict[str, Any], region: str = None) -> None: ... + @staticmethod + def reset_cache(cache_kwargs: Dict[str, Any] = None, region: Optional[str] = None) -> None: ... + @staticmethod + def get_manifest(cache_kwargs: Optional[Dict[str, Any]] = ..., region: Optional[str] = ...) -> List[JumpStartModelHeader]: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/__init__.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/environment_variables.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/environment_variables.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/hyperparameters.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/hyperparameters.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/image_uris.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/image_uris.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/incremental_training.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/incremental_training.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/instance_types.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/instance_types.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/kwargs.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/kwargs.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/metric_definitions.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/metric_definitions.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/model_uris.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/model_uris.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/predictors.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/predictors.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/prepack.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/prepack.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/resource_names.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/resource_names.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/artifacts/script_uris.pyi b/stubs/sagemaker/sagemaker/jumpstart/artifacts/script_uris.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/cache.pyi b/stubs/sagemaker/sagemaker/jumpstart/cache.pyi new file mode 100644 index 000000000000..bac6d2a0e7a2 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/cache.pyi @@ -0,0 +1,30 @@ +import datetime +from _typeshed import Incomplete +from typing import List, Optional + +import botocore +from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartModelSpecs + +class JumpStartModelsCache: + s3_bucket_name: Incomplete + def __init__( + self, + region: str = "eu-west-1", + max_s3_cache_items: int = 20, + s3_cache_expiration_horizon: datetime.timedelta = ..., + max_semantic_version_cache_items: int = 20, + semantic_version_cache_expiration_horizon: datetime.timedelta = ..., + manifest_file_s3_key: str = "models_manifest.json", + s3_bucket_name: Optional[str] = None, + s3_client_config: Optional[botocore.config.Config] = None, + ) -> None: ... + def set_region(self, region: str) -> None: ... + def get_region(self) -> str: ... + def set_manifest_file_s3_key(self, key: str) -> None: ... + def get_manifest_file_s3_key(self) -> str: ... + def set_s3_bucket_name(self, s3_bucket_name: str) -> None: ... + def get_bucket(self) -> str: ... + def get_manifest(self) -> List[JumpStartModelHeader]: ... + def get_header(self, model_id: str, semantic_version_str: str) -> JumpStartModelHeader: ... + def get_specs(self, model_id: str, semantic_version_str: str) -> JumpStartModelSpecs: ... + def clear(self) -> None: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/constants.pyi b/stubs/sagemaker/sagemaker/jumpstart/constants.pyi new file mode 100644 index 000000000000..5064010a0302 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/constants.pyi @@ -0,0 +1,28 @@ +from _typeshed import Incomplete +from typing import Dict, Set, Type + +from sagemaker.base_deserializers import BaseDeserializer +from sagemaker.base_serializers import BaseSerializer +from sagemaker.jumpstart.enums import DeserializerType, MIMEType, SerializerType +from sagemaker.jumpstart.types import JumpStartLaunchedRegionInfo + +JUMPSTART_LAUNCHED_REGIONS: Set[JumpStartLaunchedRegionInfo] +JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT: Incomplete +JUMPSTART_REGION_NAME_SET: Incomplete +JUMPSTART_BUCKET_NAME_SET: Incomplete +JUMPSTART_DEFAULT_REGION_NAME: Incomplete +JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY: str +INFERENCE_ENTRY_POINT_SCRIPT_NAME: str +TRAINING_ENTRY_POINT_SCRIPT_NAME: str +SUPPORTED_JUMPSTART_SCOPES: Incomplete +ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE: str +ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE: str +ENV_VARIABLE_JUMPSTART_SCRIPT_ARTIFACT_BUCKET_OVERRIDE: str +ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE: str +ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE: str +JUMPSTART_RESOURCE_BASE_NAME: str +CONTENT_TYPE_TO_SERIALIZER_TYPE_MAP: Dict[MIMEType, SerializerType] +ACCEPT_TYPE_TO_DESERIALIZER_TYPE_MAP: Dict[MIMEType, DeserializerType] +SERIALIZER_TYPE_TO_CLASS_MAP: Dict[SerializerType, Type[BaseSerializer]] +DESERIALIZER_TYPE_TO_CLASS_MAP: Dict[DeserializerType, Type[BaseDeserializer]] +MODEL_ID_LIST_WEB_URL: str diff --git a/stubs/sagemaker/sagemaker/jumpstart/enums.pyi b/stubs/sagemaker/sagemaker/jumpstart/enums.pyi new file mode 100644 index 000000000000..111764984fc7 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/enums.pyi @@ -0,0 +1,55 @@ +from enum import Enum + +class ModelFramework(str, Enum): + PYTORCH: str + TENSORFLOW: str + MXNET: str + HUGGINGFACE: str + LIGHTGBM: str + CATBOOST: str + XGBOOST: str + SKLEARN: str + +class VariableScope(str, Enum): + CONTAINER: str + ALGORITHM: str + +class JumpStartScriptScope(str, Enum): + INFERENCE: str + TRAINING: str + +class HyperparameterValidationMode(str, Enum): + VALIDATE_PROVIDED: str + VALIDATE_ALGORITHM: str + VALIDATE_ALL: str + +class VariableTypes(str, Enum): + TEXT: str + INT: str + FLOAT: str + BOOL: str + +class JumpStartTag(str, Enum): + INFERENCE_MODEL_URI: str + INFERENCE_SCRIPT_URI: str + TRAINING_MODEL_URI: str + TRAINING_SCRIPT_URI: str + +class SerializerType(str, Enum): + TEXT: str + JSON: str + CSV: str + RAW_BYTES: str + +class DeserializerType(str, Enum): + JSON: str + +class MIMEType(str, Enum): + X_IMAGE: str + LIST_TEXT: str + X_TEXT: str + JSON: str + CSV: str + WAV: str + @staticmethod + def from_suffixed_type(mime_type_with_suffix: str) -> MIMEType: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi b/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi new file mode 100644 index 000000000000..11406710a414 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi @@ -0,0 +1,127 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker import session as session +from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig +from sagemaker.base_deserializers import BaseDeserializer +from sagemaker.base_serializers import BaseSerializer +from sagemaker.debugger.debugger import DebuggerHookConfig, RuleBase, TensorBoardOutputConfig +from sagemaker.debugger.profiler_config import ProfilerConfig +from sagemaker.estimator import Estimator +from sagemaker.explainer.explainer_config import ExplainerConfig +from sagemaker.inputs import FileSystemInput, TrainingInput +from sagemaker.instance_group import InstanceGroup +from sagemaker.model_monitor.data_capture_config import DataCaptureConfig +from sagemaker.predictor import PredictorBase +from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class JumpStartEstimator(Estimator): + model_id: Incomplete + model_version: Incomplete + instance_type: Incomplete + tolerate_deprecated_model: Incomplete + tolerate_vulnerable_model: Incomplete + instance_count: Incomplete + region: Incomplete + orig_predictor_cls: Incomplete + role: Incomplete + sagemaker_session: Incomplete + def __init__( + self, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: Optional[bool] = None, + tolerate_deprecated_model: Optional[bool] = None, + region: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str] = None, + instance_count: Optional[Union[int, PipelineVariable]] = None, + instance_type: Optional[Union[str, PipelineVariable]] = None, + keep_alive_period_in_seconds: Optional[Union[int, PipelineVariable]] = None, + volume_size: Optional[Union[int, PipelineVariable]] = None, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + max_run: Optional[Union[int, PipelineVariable]] = None, + input_mode: Optional[Union[str, PipelineVariable]] = None, + 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.Session] = None, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = 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: Optional[Union[str, PipelineVariable]] = None, + metric_definitions: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + encrypt_inter_container_traffic: Union[bool, PipelineVariable] = None, + use_spot_instances: Optional[Union[bool, PipelineVariable]] = None, + max_wait: Optional[Union[int, PipelineVariable]] = None, + checkpoint_s3_uri: Optional[Union[str, PipelineVariable]] = None, + checkpoint_local_path: Optional[Union[str, PipelineVariable]] = None, + enable_network_isolation: Union[bool, PipelineVariable] = None, + rules: Optional[List[RuleBase]] = None, + debugger_hook_config: Optional[Union[DebuggerHookConfig, bool]] = None, + tensorboard_output_config: Optional[TensorBoardOutputConfig] = None, + enable_sagemaker_metrics: Optional[Union[bool, PipelineVariable]] = None, + profiler_config: Optional[ProfilerConfig] = None, + disable_profiler: Optional[bool] = None, + environment: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + max_retry_attempts: Optional[Union[int, PipelineVariable]] = None, + source_dir: Optional[Union[str, PipelineVariable]] = None, + git_config: Optional[Dict[str, str]] = None, + container_log_level: Optional[Union[int, PipelineVariable]] = None, + code_location: Optional[str] = None, + entry_point: Optional[Union[str, PipelineVariable]] = None, + dependencies: Optional[List[str]] = None, + instance_groups: Optional[List[InstanceGroup]] = None, + training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None, + training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None, + ) -> None: ... + def fit( + self, + inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None, + wait: Optional[bool] = True, + logs: Optional[str] = None, + job_name: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + ) -> None: ... + def deploy( + self, + initial_instance_count: Optional[int] = None, + instance_type: Optional[str] = None, + serializer: Optional[BaseSerializer] = None, + deserializer: Optional[BaseDeserializer] = None, + accelerator_type: Optional[str] = None, + endpoint_name: Optional[str] = None, + tags: List[Dict[str, str]] = None, + kms_key: Optional[str] = None, + wait: Optional[bool] = True, + data_capture_config: Optional[DataCaptureConfig] = None, + async_inference_config: Optional[AsyncInferenceConfig] = None, + serverless_inference_config: Optional[ServerlessInferenceConfig] = None, + volume_size: Optional[int] = None, + model_data_download_timeout: Optional[int] = None, + container_startup_health_check_timeout: Optional[int] = None, + inference_recommendation_id: Optional[str] = None, + explainer_config: Optional[ExplainerConfig] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str] = 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, + sagemaker_session: Optional[session.Session] = None, + enable_network_isolation: Union[bool, PipelineVariable] = None, + 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: Optional[Union[int, PipelineVariable]] = None, + dependencies: Optional[List[str]] = None, + git_config: Optional[Dict[str, str]] = None, + use_compiled_model: bool = False, + ) -> PredictorBase: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi b/stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi new file mode 100644 index 000000000000..a415febf0ce1 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi @@ -0,0 +1,26 @@ +from _typeshed import Incomplete +from typing import List, Optional + +from sagemaker.jumpstart.constants import JumpStartScriptScope + +NO_AVAILABLE_INSTANCES_ERROR_MSG: str +INVALID_MODEL_ID_ERROR_MSG: Incomplete + +class JumpStartHyperparametersError(ValueError): + message: Incomplete + def __init__(self, message: Optional[str] = None) -> None: ... + +class VulnerableJumpStartModelError(ValueError): + message: Incomplete + def __init__( + self, + model_id: Optional[str] = None, + version: Optional[str] = None, + vulnerabilities: Optional[List[str]] = None, + scope: Optional[JumpStartScriptScope] = None, + message: Optional[str] = None, + ) -> None: ... + +class DeprecatedJumpStartModelError(ValueError): + message: Incomplete + def __init__(self, model_id: Optional[str] = None, version: Optional[str] = None, message: Optional[str] = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/factory/__init__.pyi b/stubs/sagemaker/sagemaker/jumpstart/factory/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi b/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi new file mode 100644 index 000000000000..f2c282bdb0d5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi @@ -0,0 +1,122 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig +from sagemaker.base_deserializers import BaseDeserializer +from sagemaker.base_serializers import BaseSerializer +from sagemaker.debugger.debugger import DebuggerHookConfig, RuleBase, TensorBoardOutputConfig +from sagemaker.debugger.profiler_config import ProfilerConfig +from sagemaker.explainer.explainer_config import ExplainerConfig +from sagemaker.inputs import FileSystemInput, TrainingInput +from sagemaker.instance_group import InstanceGroup +from sagemaker.jumpstart.types import JumpStartEstimatorDeployKwargs, JumpStartEstimatorFitKwargs, JumpStartEstimatorInitKwargs +from sagemaker.model_monitor.data_capture_config import DataCaptureConfig +from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +def get_init_kwargs( + model_id: str, + model_version: Optional[str] = None, + tolerate_vulnerable_model: Optional[bool] = None, + tolerate_deprecated_model: Optional[bool] = None, + region: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str] = None, + instance_count: Optional[Union[int, PipelineVariable]] = None, + instance_type: Optional[Union[str, PipelineVariable]] = None, + keep_alive_period_in_seconds: Optional[Union[int, PipelineVariable]] = None, + volume_size: Optional[Union[int, PipelineVariable]] = None, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + max_run: Optional[Union[int, PipelineVariable]] = None, + input_mode: Optional[Union[str, PipelineVariable]] = None, + 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, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = 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: Optional[Union[str, PipelineVariable]] = None, + metric_definitions: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + encrypt_inter_container_traffic: Union[bool, PipelineVariable] = None, + use_spot_instances: Optional[Union[bool, PipelineVariable]] = None, + max_wait: Optional[Union[int, PipelineVariable]] = None, + checkpoint_s3_uri: Optional[Union[str, PipelineVariable]] = None, + checkpoint_local_path: Optional[Union[str, PipelineVariable]] = None, + enable_network_isolation: Union[bool, PipelineVariable] = None, + rules: Optional[List[RuleBase]] = None, + debugger_hook_config: Optional[Union[DebuggerHookConfig, bool]] = None, + tensorboard_output_config: Optional[TensorBoardOutputConfig] = None, + enable_sagemaker_metrics: Optional[Union[bool, PipelineVariable]] = None, + profiler_config: Optional[ProfilerConfig] = None, + disable_profiler: Optional[bool] = None, + environment: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + max_retry_attempts: Optional[Union[int, PipelineVariable]] = None, + source_dir: Optional[Union[str, PipelineVariable]] = None, + git_config: Optional[Dict[str, str]] = None, + container_log_level: Optional[Union[int, PipelineVariable]] = None, + code_location: Optional[str] = None, + entry_point: Optional[Union[str, PipelineVariable]] = None, + dependencies: Optional[List[str]] = None, + instance_groups: Optional[List[InstanceGroup]] = None, + training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None, + training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None, +) -> JumpStartEstimatorInitKwargs: ... +def get_fit_kwargs( + model_id: str, + model_version: Optional[str] = None, + region: Optional[str] = None, + inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None, + wait: Optional[bool] = None, + logs: Optional[str] = None, + job_name: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + tolerate_vulnerable_model: Optional[bool] = None, + tolerate_deprecated_model: Optional[bool] = None, +) -> JumpStartEstimatorFitKwargs: ... +def get_deploy_kwargs( + model_id: str, + model_version: Optional[str] = None, + region: Optional[str] = None, + initial_instance_count: Optional[int] = None, + instance_type: Optional[str] = None, + serializer: Optional[BaseSerializer] = None, + deserializer: Optional[BaseDeserializer] = None, + accelerator_type: Optional[str] = None, + endpoint_name: Optional[str] = None, + tags: List[Dict[str, str]] = None, + kms_key: Optional[str] = None, + wait: Optional[bool] = None, + data_capture_config: Optional[DataCaptureConfig] = None, + async_inference_config: Optional[AsyncInferenceConfig] = None, + serverless_inference_config: Optional[ServerlessInferenceConfig] = None, + volume_size: Optional[int] = None, + model_data_download_timeout: Optional[int] = None, + container_startup_health_check_timeout: Optional[int] = None, + inference_recommendation_id: Optional[str] = None, + explainer_config: Optional[ExplainerConfig] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str] = 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, + enable_network_isolation: Union[bool, PipelineVariable] = None, + 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: Optional[Union[int, PipelineVariable]] = None, + dependencies: Optional[List[str]] = None, + git_config: Optional[Dict[str, str]] = None, + tolerate_deprecated_model: Optional[bool] = None, + tolerate_vulnerable_model: Optional[bool] = None, + use_compiled_model: Optional[bool] = None, + model_name: Optional[str] = None, +) -> JumpStartEstimatorDeployKwargs: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi b/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi new file mode 100644 index 000000000000..9bd15ddf2568 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi @@ -0,0 +1,74 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig +from sagemaker.base_deserializers import BaseDeserializer +from sagemaker.base_predictor import Predictor +from sagemaker.base_serializers import BaseSerializer +from sagemaker.explainer.explainer_config import ExplainerConfig +from sagemaker.jumpstart.types import JumpStartModelDeployKwargs, JumpStartModelInitKwargs +from sagemaker.model_monitor.data_capture_config import DataCaptureConfig +from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +def get_default_predictor( + predictor: Predictor, + model_id: str, + model_version: str, + region: str, + tolerate_vulnerable_model: bool, + tolerate_deprecated_model: bool, +) -> Predictor: ... +def get_deploy_kwargs( + model_id: str, + model_version: Optional[str] = None, + region: Optional[str] = None, + initial_instance_count: Optional[int] = None, + instance_type: Optional[str] = None, + serializer: Optional[BaseSerializer] = None, + deserializer: Optional[BaseDeserializer] = None, + accelerator_type: Optional[str] = None, + endpoint_name: Optional[str] = None, + tags: List[Dict[str, str]] = None, + kms_key: Optional[str] = None, + wait: Optional[bool] = None, + data_capture_config: Optional[DataCaptureConfig] = None, + async_inference_config: Optional[AsyncInferenceConfig] = None, + serverless_inference_config: Optional[ServerlessInferenceConfig] = None, + volume_size: Optional[int] = None, + model_data_download_timeout: Optional[int] = None, + container_startup_health_check_timeout: Optional[int] = None, + inference_recommendation_id: Optional[str] = None, + explainer_config: Optional[ExplainerConfig] = None, + tolerate_vulnerable_model: Optional[bool] = None, + tolerate_deprecated_model: Optional[bool] = None, +) -> JumpStartModelDeployKwargs: ... +def get_init_kwargs( + model_id: str, + model_from_estimator: bool = False, + model_version: Optional[str] = None, + tolerate_vulnerable_model: Optional[bool] = None, + tolerate_deprecated_model: Optional[bool] = None, + instance_type: Optional[str] = None, + region: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + 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] = None, + 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: Optional[Union[int, PipelineVariable]] = None, + dependencies: Optional[List[str]] = None, + git_config: Optional[Dict[str, str]] = None, +) -> JumpStartModelInitKwargs: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/filters.pyi b/stubs/sagemaker/sagemaker/jumpstart/filters.pyi new file mode 100644 index 000000000000..c092fff2c849 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/filters.pyi @@ -0,0 +1,87 @@ +from _typeshed import Incomplete +from enum import Enum +from typing import Any, Dict, List, Union + +from sagemaker.jumpstart.types import JumpStartDataHolderType + +class BooleanValues(str, Enum): + TRUE: str + FALSE: str + UNKNOWN: str + UNEVALUATED: str + +class FilterOperators(str, Enum): + EQUALS: str + NOT_EQUALS: str + IN: str + NOT_IN: str + +class SpecialSupportedFilterKeys(str, Enum): + TASK: str + FRAMEWORK: str + SUPPORTED_MODEL: str + +FILTER_OPERATOR_STRING_MAPPINGS: Incomplete +ACCEPTABLE_OPERATORS_IN_PARSE_ORDER: Incomplete +SPECIAL_SUPPORTED_FILTER_KEYS: Incomplete + +class Operand: + unresolved_value: Incomplete + def __init__(self, unresolved_value: Any, resolved_value: BooleanValues = ...) -> None: ... + def __iter__(self) -> Any: ... + def eval(self) -> None: ... + @property + def resolved_value(self) -> BooleanValues: ... + @resolved_value.setter + def resolved_value(self, new_resolved_value: Any) -> None: ... + @staticmethod + def validate_operand(operand: Any) -> Any: ... + +class Operator(Operand): + def __init__(self, resolved_value: BooleanValues = ..., unresolved_value: Any = None) -> None: ... + def eval(self) -> None: ... + def __iter__(self) -> Any: ... + +class And(Operator): + operands: Incomplete + def __init__(self, *operands: Union[Operand, str]) -> None: ... + resolved_value: Incomplete + def eval(self) -> None: ... + def __iter__(self) -> Any: ... + +class Constant(Operator): + def __init__(self, constant: BooleanValues) -> None: ... + def eval(self) -> None: ... + def __iter__(self) -> Any: ... + +class Identity(Operator): + operand: Incomplete + def __init__(self, operand: Union[Operand, str]) -> None: ... + def __iter__(self) -> Any: ... + resolved_value: Incomplete + def eval(self) -> None: ... + +class Or(Operator): + operands: Incomplete + def __init__(self, *operands: Union[Operand, str]) -> None: ... + resolved_value: Incomplete + def eval(self) -> None: ... + def __iter__(self) -> Any: ... + +class Not(Operator): + operand: Incomplete + def __init__(self, operand: Union[Operand, str]) -> None: ... + resolved_value: Incomplete + def eval(self) -> None: ... + def __iter__(self) -> Any: ... + +class ModelFilter(JumpStartDataHolderType): + key: Incomplete + value: Incomplete + operator: Incomplete + def __init__(self, key: str, value: str, operator: str) -> None: ... + +def parse_filter_string(filter_string: str) -> ModelFilter: ... +def evaluate_filter_expression( + model_filter: ModelFilter, cached_model_value: Union[str, bool, int, float, Dict[str, Any], List[Any]] +) -> BooleanValues: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/model.pyi b/stubs/sagemaker/sagemaker/jumpstart/model.pyi new file mode 100644 index 000000000000..d81ca06b29b1 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/model.pyi @@ -0,0 +1,70 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig +from sagemaker.base_deserializers import BaseDeserializer +from sagemaker.base_serializers import BaseSerializer +from sagemaker.explainer.explainer_config import ExplainerConfig +from sagemaker.model import Model +from sagemaker.model_monitor.data_capture_config import DataCaptureConfig +from sagemaker.predictor import PredictorBase +from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class JumpStartModel(Model): + orig_predictor_cls: Incomplete + model_id: Incomplete + model_version: Incomplete + instance_type: Incomplete + tolerate_vulnerable_model: Incomplete + tolerate_deprecated_model: Incomplete + region: Incomplete + def __init__( + self, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: Optional[bool] = None, + tolerate_deprecated_model: Optional[bool] = None, + region: Optional[str] = None, + instance_type: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + 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] = None, + 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: Optional[Union[int, PipelineVariable]] = None, + dependencies: Optional[List[str]] = None, + git_config: Optional[Dict[str, str]] = None, + ) -> None: ... + def deploy( + self, + initial_instance_count: Optional[int] = None, + instance_type: Optional[str] = None, + serializer: Optional[BaseSerializer] = None, + deserializer: Optional[BaseDeserializer] = None, + accelerator_type: Optional[str] = None, + endpoint_name: Optional[str] = None, + tags: List[Dict[str, str]] = None, + kms_key: Optional[str] = None, + wait: Optional[bool] = True, + data_capture_config: Optional[DataCaptureConfig] = None, + async_inference_config: Optional[AsyncInferenceConfig] = None, + serverless_inference_config: Optional[ServerlessInferenceConfig] = None, + volume_size: Optional[int] = None, + model_data_download_timeout: Optional[int] = None, + container_startup_health_check_timeout: Optional[int] = None, + inference_recommendation_id: Optional[str] = None, + explainer_config: Optional[ExplainerConfig] = None, + ) -> PredictorBase: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi b/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi new file mode 100644 index 000000000000..1972bb6b6efd --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi @@ -0,0 +1,16 @@ +from typing import List, Tuple, Union + +from sagemaker.jumpstart.filters import Operator + +def extract_framework_task_model(model_id: str) -> Tuple[str, str, str]: ... +def list_jumpstart_tasks(filter: Union[Operator, str] = ..., region: str = "eu-west-1") -> List[str]: ... +def list_jumpstart_frameworks(filter: Union[Operator, str] = ..., region: str = "eu-west-1") -> List[str]: ... +def list_jumpstart_scripts(filter: Union[Operator, str] = ..., region: str = "eu-west-1") -> List[str]: ... +def list_jumpstart_models( + filter: Union[Operator, str] = ..., + region: str = "eu-west-1", + list_incomplete_models: bool = False, + list_old_models: bool = False, + list_versions: bool = False, +) -> List[Union[Tuple[str], Tuple[str, str]]]: ... +def get_model_url(model_id: str, model_version: str, region: str = "eu-west-1") -> str: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/parameters.pyi b/stubs/sagemaker/sagemaker/jumpstart/parameters.pyi new file mode 100644 index 000000000000..7bc23e7d5312 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/parameters.pyi @@ -0,0 +1,6 @@ +from _typeshed import Incomplete + +JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS: int +JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS: int +JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON: Incomplete +JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON: Incomplete diff --git a/stubs/sagemaker/sagemaker/jumpstart/types.pyi b/stubs/sagemaker/sagemaker/jumpstart/types.pyi new file mode 100644 index 000000000000..901695ca2ed9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/types.pyi @@ -0,0 +1,449 @@ +from _typeshed import Incomplete +from enum import Enum +from typing import Any, Dict, List, Optional, Set, Union + +class JumpStartDataHolderType: + def __eq__(self, other: Any) -> bool: ... + def __hash__(self) -> int: ... + +class JumpStartS3FileType(str, Enum): + MANIFEST: str + SPECS: str + +class JumpStartLaunchedRegionInfo(JumpStartDataHolderType): + content_bucket: Incomplete + region_name: Incomplete + def __init__(self, content_bucket: str, region_name: str) -> None: ... + +class JumpStartModelHeader(JumpStartDataHolderType): + def __init__(self, header: Dict[str, str]) -> None: ... + def to_json(self) -> Dict[str, str]: ... + model_id: Incomplete + version: Incomplete + min_version: Incomplete + spec_key: Incomplete + def from_json(self, json_obj: Dict[str, str]) -> None: ... + +class JumpStartECRSpecs(JumpStartDataHolderType): + def __init__(self, spec: Dict[str, Any]) -> None: ... + framework: Incomplete + framework_version: Incomplete + py_version: Incomplete + huggingface_transformers_version: Incomplete + def from_json(self, json_obj: Dict[str, Any]) -> None: ... + def to_json(self) -> Dict[str, Any]: ... + +class JumpStartHyperparameter(JumpStartDataHolderType): + def __init__(self, spec: Dict[str, Any]) -> None: ... + name: Incomplete + type: Incomplete + default: Incomplete + scope: Incomplete + options: Incomplete + min: Incomplete + max: Incomplete + exclusive_min: Incomplete + exclusive_max: Incomplete + def from_json(self, json_obj: Dict[str, Any]) -> None: ... + def to_json(self) -> Dict[str, Any]: ... + +class JumpStartEnvironmentVariable(JumpStartDataHolderType): + def __init__(self, spec: Dict[str, Any]) -> None: ... + name: Incomplete + type: Incomplete + default: Incomplete + scope: Incomplete + required_for_model_class: Incomplete + def from_json(self, json_obj: Dict[str, Any]) -> None: ... + def to_json(self) -> Dict[str, Any]: ... + +class JumpStartPredictorSpecs(JumpStartDataHolderType): + def __init__(self, spec: Optional[Dict[str, Any]]) -> None: ... + default_content_type: Incomplete + supported_content_types: Incomplete + default_accept_type: Incomplete + supported_accept_types: Incomplete + def from_json(self, json_obj: Optional[Dict[str, Any]]) -> None: ... + def to_json(self) -> Dict[str, Any]: ... + +class JumpStartModelSpecs(JumpStartDataHolderType): + def __init__(self, spec: Dict[str, Any]) -> None: ... + model_id: Incomplete + url: Incomplete + version: Incomplete + min_sdk_version: Incomplete + incremental_training_supported: Incomplete + hosting_ecr_specs: Incomplete + hosting_artifact_key: Incomplete + hosting_script_key: Incomplete + training_supported: Incomplete + inference_environment_variables: Incomplete + inference_vulnerable: Incomplete + inference_dependencies: Incomplete + inference_vulnerabilities: Incomplete + training_vulnerable: Incomplete + training_dependencies: Incomplete + training_vulnerabilities: Incomplete + deprecated: Incomplete + default_inference_instance_type: Incomplete + default_training_instance_type: Incomplete + supported_inference_instance_types: Incomplete + supported_training_instance_types: Incomplete + metrics: Incomplete + training_prepacked_script_key: Incomplete + hosting_prepacked_artifact_key: Incomplete + model_kwargs: Incomplete + deploy_kwargs: Incomplete + predictor_specs: Incomplete + inference_volume_size: Incomplete + inference_enable_network_isolation: Incomplete + resource_name_base: Incomplete + training_ecr_specs: Incomplete + training_artifact_key: Incomplete + training_script_key: Incomplete + hyperparameters: Incomplete + estimator_kwargs: Incomplete + fit_kwargs: Incomplete + training_volume_size: Incomplete + training_enable_network_isolation: Incomplete + def from_json(self, json_obj: Dict[str, Any]) -> None: ... + def to_json(self) -> Dict[str, Any]: ... + def supports_prepacked_inference(self) -> bool: ... + def supports_incremental_training(self) -> bool: ... + +class JumpStartVersionedModelId(JumpStartDataHolderType): + model_id: Incomplete + version: Incomplete + def __init__(self, model_id: str, version: str) -> None: ... + +class JumpStartCachedS3ContentKey(JumpStartDataHolderType): + file_type: Incomplete + s3_key: Incomplete + def __init__(self, file_type: JumpStartS3FileType, s3_key: str) -> None: ... + +class JumpStartCachedS3ContentValue(JumpStartDataHolderType): + formatted_content: Incomplete + md5_hash: Incomplete + def __init__( + self, + formatted_content: Union[Dict[JumpStartVersionedModelId, JumpStartModelHeader], JumpStartModelSpecs], + md5_hash: Optional[str] = None, + ) -> None: ... + +class JumpStartKwargs(JumpStartDataHolderType): + SERIALIZATION_EXCLUSION_SET: Set[str] + def to_kwargs_dict(self): ... + +class JumpStartModelInitKwargs(JumpStartKwargs): + SERIALIZATION_EXCLUSION_SET: Incomplete + model_id: Incomplete + model_version: Incomplete + instance_type: Incomplete + region: Incomplete + image_uri: Incomplete + model_data: Incomplete + source_dir: Incomplete + entry_point: Incomplete + env: Incomplete + predictor_cls: Incomplete + role: Incomplete + name: Incomplete + vpc_config: Incomplete + sagemaker_session: Incomplete + enable_network_isolation: Incomplete + model_kms_key: Incomplete + image_config: Incomplete + code_location: Incomplete + container_log_level: Incomplete + dependencies: Incomplete + git_config: Incomplete + tolerate_deprecated_model: Incomplete + tolerate_vulnerable_model: Incomplete + def __init__( + self, + model_id: str, + model_version: Optional[str] = None, + region: Optional[str] = None, + instance_type: Optional[str] = None, + image_uri: Optional[Union[str, Any]] = None, + model_data: Optional[Union[str, Any]] = None, + role: Optional[str] = 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, + sagemaker_session: Optional[Any] = None, + enable_network_isolation: Union[bool, Any] = None, + model_kms_key: Optional[str] = None, + image_config: Optional[Dict[str, Union[str, Any]]] = None, + source_dir: Optional[str] = None, + code_location: Optional[str] = None, + entry_point: Optional[str] = None, + container_log_level: Optional[Union[int, Any]] = None, + dependencies: Optional[List[str]] = None, + git_config: Optional[Dict[str, str]] = None, + tolerate_vulnerable_model: Optional[bool] = None, + tolerate_deprecated_model: Optional[bool] = None, + ) -> None: ... + +class JumpStartModelDeployKwargs(JumpStartKwargs): + SERIALIZATION_EXCLUSION_SET: Incomplete + model_id: Incomplete + model_version: Incomplete + initial_instance_count: Incomplete + instance_type: Incomplete + region: Incomplete + serializer: Incomplete + deserializer: Incomplete + accelerator_type: Incomplete + endpoint_name: Incomplete + tags: Incomplete + kms_key: Incomplete + wait: Incomplete + data_capture_config: Incomplete + async_inference_config: Incomplete + serverless_inference_config: Incomplete + volume_size: Incomplete + model_data_download_timeout: Incomplete + container_startup_health_check_timeout: Incomplete + inference_recommendation_id: Incomplete + explainer_config: Incomplete + tolerate_vulnerable_model: Incomplete + tolerate_deprecated_model: Incomplete + def __init__( + self, + model_id: str, + model_version: Optional[str] = None, + region: Optional[str] = None, + initial_instance_count: Optional[int] = None, + instance_type: Optional[str] = None, + serializer: Optional[Any] = None, + deserializer: Optional[Any] = None, + accelerator_type: Optional[str] = None, + endpoint_name: Optional[str] = None, + tags: List[Dict[str, str]] = None, + kms_key: Optional[str] = None, + wait: Optional[bool] = None, + data_capture_config: Optional[Any] = None, + async_inference_config: Optional[Any] = None, + serverless_inference_config: Optional[Any] = None, + volume_size: Optional[int] = None, + model_data_download_timeout: Optional[int] = None, + container_startup_health_check_timeout: Optional[int] = None, + inference_recommendation_id: Optional[str] = None, + explainer_config: Optional[Any] = None, + tolerate_deprecated_model: Optional[bool] = None, + tolerate_vulnerable_model: Optional[bool] = None, + ) -> None: ... + +class JumpStartEstimatorInitKwargs(JumpStartKwargs): + SERIALIZATION_EXCLUSION_SET: Incomplete + model_id: Incomplete + model_version: Incomplete + instance_type: Incomplete + instance_count: Incomplete + region: Incomplete + image_uri: Incomplete + model_uri: Incomplete + source_dir: Incomplete + entry_point: Incomplete + hyperparameters: Incomplete + metric_definitions: Incomplete + role: Incomplete + keep_alive_period_in_seconds: Incomplete + volume_size: Incomplete + volume_kms_key: Incomplete + max_run: Incomplete + input_mode: Incomplete + output_path: Incomplete + output_kms_key: Incomplete + base_job_name: Incomplete + sagemaker_session: Incomplete + tags: Incomplete + subnets: Incomplete + security_group_ids: Incomplete + model_channel_name: Incomplete + encrypt_inter_container_traffic: Incomplete + use_spot_instances: Incomplete + max_wait: Incomplete + checkpoint_s3_uri: Incomplete + checkpoint_local_path: Incomplete + enable_network_isolation: Incomplete + rules: Incomplete + debugger_hook_config: Incomplete + tensorboard_output_config: Incomplete + enable_sagemaker_metrics: Incomplete + profiler_config: Incomplete + disable_profiler: Incomplete + environment: Incomplete + max_retry_attempts: Incomplete + git_config: Incomplete + container_log_level: Incomplete + code_location: Incomplete + dependencies: Incomplete + instance_groups: Incomplete + training_repository_access_mode: Incomplete + training_repository_credentials_provider_arn: Incomplete + tolerate_vulnerable_model: Incomplete + tolerate_deprecated_model: Incomplete + def __init__( + self, + model_id: str, + model_version: Optional[str] = None, + region: Optional[str] = None, + image_uri: Optional[Union[str, Any]] = None, + role: Optional[str] = None, + instance_count: Optional[Union[int, Any]] = None, + instance_type: Optional[Union[str, Any]] = None, + keep_alive_period_in_seconds: Optional[Union[int, Any]] = None, + volume_size: Optional[Union[int, Any]] = None, + volume_kms_key: Optional[Union[str, Any]] = None, + max_run: Optional[Union[int, Any]] = None, + input_mode: Optional[Union[str, Any]] = None, + output_path: Optional[Union[str, Any]] = None, + output_kms_key: Optional[Union[str, Any]] = None, + base_job_name: Optional[str] = None, + sagemaker_session: Optional[Any] = None, + hyperparameters: Optional[Dict[str, Union[str, Any]]] = None, + tags: Optional[List[Dict[str, Union[str, Any]]]] = None, + subnets: Optional[List[Union[str, Any]]] = None, + security_group_ids: Optional[List[Union[str, Any]]] = None, + model_uri: Optional[str] = None, + model_channel_name: Optional[Union[str, Any]] = None, + metric_definitions: Optional[List[Dict[str, Union[str, Any]]]] = None, + encrypt_inter_container_traffic: Union[bool, Any] = None, + use_spot_instances: Optional[Union[bool, Any]] = None, + max_wait: Optional[Union[int, Any]] = None, + checkpoint_s3_uri: Optional[Union[str, Any]] = None, + checkpoint_local_path: Optional[Union[str, Any]] = None, + enable_network_isolation: Union[bool, Any] = None, + rules: Optional[List[Any]] = None, + debugger_hook_config: Optional[Union[Any, bool]] = None, + tensorboard_output_config: Optional[Any] = None, + enable_sagemaker_metrics: Optional[Union[bool, Any]] = None, + profiler_config: Optional[Any] = None, + disable_profiler: Optional[bool] = None, + environment: Optional[Dict[str, Union[str, Any]]] = None, + max_retry_attempts: Optional[Union[int, Any]] = None, + source_dir: Optional[Union[str, Any]] = None, + git_config: Optional[Dict[str, str]] = None, + container_log_level: Optional[Union[int, Any]] = None, + code_location: Optional[str] = None, + entry_point: Optional[Union[str, Any]] = None, + dependencies: Optional[List[str]] = None, + instance_groups: Optional[List[Any]] = None, + training_repository_access_mode: Optional[Union[str, Any]] = None, + training_repository_credentials_provider_arn: Optional[Union[str, Any]] = None, + tolerate_vulnerable_model: Optional[bool] = None, + tolerate_deprecated_model: Optional[bool] = None, + ) -> None: ... + +class JumpStartEstimatorFitKwargs(JumpStartKwargs): + SERIALIZATION_EXCLUSION_SET: Incomplete + model_id: Incomplete + model_version: Incomplete + region: Incomplete + inputs: Incomplete + wait: Incomplete + logs: Incomplete + job_name: Incomplete + experiment_config: Incomplete + tolerate_deprecated_model: Incomplete + tolerate_vulnerable_model: Incomplete + def __init__( + self, + model_id: str, + model_version: Optional[str] = None, + region: Optional[str] = None, + inputs: Optional[Union[str, Dict, Any, Any]] = None, + wait: Optional[bool] = None, + logs: Optional[str] = None, + job_name: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + tolerate_deprecated_model: Optional[bool] = None, + tolerate_vulnerable_model: Optional[bool] = None, + ) -> None: ... + +class JumpStartEstimatorDeployKwargs(JumpStartKwargs): + SERIALIZATION_EXCLUSION_SET: Incomplete + model_id: Incomplete + model_version: Incomplete + instance_type: Incomplete + initial_instance_count: Incomplete + region: Incomplete + image_uri: Incomplete + source_dir: Incomplete + entry_point: Incomplete + env: Incomplete + predictor_cls: Incomplete + serializer: Incomplete + deserializer: Incomplete + accelerator_type: Incomplete + endpoint_name: Incomplete + tags: Incomplete + kms_key: Incomplete + wait: Incomplete + data_capture_config: Incomplete + async_inference_config: Incomplete + serverless_inference_config: Incomplete + volume_size: Incomplete + model_data_download_timeout: Incomplete + container_startup_health_check_timeout: Incomplete + inference_recommendation_id: Incomplete + explainer_config: Incomplete + role: Incomplete + model_name: Incomplete + vpc_config: Incomplete + sagemaker_session: Incomplete + enable_network_isolation: Incomplete + model_kms_key: Incomplete + image_config: Incomplete + code_location: Incomplete + container_log_level: Incomplete + dependencies: Incomplete + git_config: Incomplete + tolerate_deprecated_model: Incomplete + tolerate_vulnerable_model: Incomplete + use_compiled_model: Incomplete + def __init__( + self, + model_id: str, + model_version: Optional[str] = None, + region: Optional[str] = None, + initial_instance_count: Optional[int] = None, + instance_type: Optional[str] = None, + serializer: Optional[Any] = None, + deserializer: Optional[Any] = None, + accelerator_type: Optional[str] = None, + endpoint_name: Optional[str] = None, + tags: List[Dict[str, str]] = None, + kms_key: Optional[str] = None, + wait: Optional[bool] = None, + data_capture_config: Optional[Any] = None, + async_inference_config: Optional[Any] = None, + serverless_inference_config: Optional[Any] = None, + volume_size: Optional[int] = None, + model_data_download_timeout: Optional[int] = None, + container_startup_health_check_timeout: Optional[int] = None, + inference_recommendation_id: Optional[str] = None, + explainer_config: Optional[Any] = None, + image_uri: Optional[Union[str, Any]] = None, + role: Optional[str] = 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, + sagemaker_session: Optional[Any] = None, + enable_network_isolation: Union[bool, Any] = None, + model_kms_key: Optional[str] = None, + image_config: Optional[Dict[str, Union[str, Any]]] = None, + source_dir: Optional[str] = None, + code_location: Optional[str] = None, + entry_point: Optional[str] = None, + container_log_level: Optional[Union[int, Any]] = None, + dependencies: Optional[List[str]] = None, + git_config: Optional[Dict[str, str]] = None, + tolerate_deprecated_model: Optional[bool] = None, + tolerate_vulnerable_model: Optional[bool] = None, + use_compiled_model: bool = False, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/utils.pyi b/stubs/sagemaker/sagemaker/jumpstart/utils.pyi new file mode 100644 index 000000000000..7127090e6eb4 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/utils.pyi @@ -0,0 +1,53 @@ +from _typeshed import Incomplete +from typing import Any, Dict, List, Optional + +from sagemaker.jumpstart import enums +from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartModelSpecs, JumpStartVersionedModelId +from sagemaker.session import Session + +LOGGER: Incomplete + +def get_jumpstart_launched_regions_message() -> str: ... +def get_jumpstart_content_bucket(region: str = "eu-west-1") -> str: ... +def get_formatted_manifest(manifest: List[Dict]) -> Dict[JumpStartVersionedModelId, JumpStartModelHeader]: ... +def get_sagemaker_version() -> str: ... +def parse_sagemaker_version() -> str: ... +def is_jumpstart_model_input(model_id: Optional[str], version: Optional[str]) -> bool: ... +def is_jumpstart_model_uri(uri: Optional[str]) -> bool: ... +def tag_key_in_array(tag_key: str, tag_array: List[Dict[str, str]]) -> bool: ... +def get_tag_value(tag_key: str, tag_array: List[Dict[str, str]]) -> str: ... +def add_single_jumpstart_tag( + uri: str, tag_key: enums.JumpStartTag, curr_tags: Optional[List[Dict[str, str]]] +) -> Optional[List]: ... +def get_jumpstart_base_name_if_jumpstart_model(*uris: Optional[str]) -> Optional[str]: ... +def add_jumpstart_tags( + tags: Optional[List[Dict[str, str]]] = None, + inference_model_uri: Optional[str] = None, + inference_script_uri: Optional[str] = None, + training_model_uri: Optional[str] = None, + training_script_uri: Optional[str] = None, +) -> Optional[List[Dict[str, str]]]: ... +def update_inference_tags_with_jumpstart_training_tags( + inference_tags: Optional[List[Dict[str, str]]], training_tags: Optional[List[Dict[str, str]]] +) -> Optional[List[Dict[str, str]]]: ... +def verify_model_region_and_return_specs( + model_id: Optional[str], + version: Optional[str], + scope: Optional[str], + region: str, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> JumpStartModelSpecs: ... +def update_dict_if_key_not_present(dict_to_update: dict, key_to_add: Any, value_to_add: Any) -> dict: ... +def resolve_model_sagemaker_config_field( + field_name: str, field_val: Optional[Any], sagemaker_session: Session, default_value: Optional[str] = None +) -> Any: ... +def resolve_estimator_sagemaker_config_field( + field_name: str, field_val: Optional[Any], sagemaker_session: Session, default_value: Optional[str] = None +) -> Any: ... +def is_valid_model_id( + model_id: Optional[str], + region: Optional[str] = None, + model_version: Optional[str] = None, + script: enums.JumpStartScriptScope = ..., +) -> bool: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/validators.pyi b/stubs/sagemaker/sagemaker/jumpstart/validators.pyi new file mode 100644 index 000000000000..308276c81650 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/validators.pyi @@ -0,0 +1,11 @@ +from typing import Any, Dict, Optional + +from sagemaker.jumpstart.enums import HyperparameterValidationMode + +def validate_hyperparameters( + model_id: str, + model_version: str, + hyperparameters: Dict[str, Any], + validation_mode: HyperparameterValidationMode = ..., + region: Optional[str] = "eu-west-1", +) -> None: ... diff --git a/stubs/sagemaker/sagemaker/lambda_helper.pyi b/stubs/sagemaker/sagemaker/lambda_helper.pyi new file mode 100644 index 000000000000..e3cdbe4a068e --- /dev/null +++ b/stubs/sagemaker/sagemaker/lambda_helper.pyi @@ -0,0 +1,41 @@ +from _typeshed import Incomplete + +from sagemaker.session import Session + +class Lambda: + function_arn: Incomplete + function_name: Incomplete + zipped_code_dir: Incomplete + s3_bucket: Incomplete + script: Incomplete + handler: Incomplete + execution_role_arn: Incomplete + session: Incomplete + timeout: Incomplete + memory_size: Incomplete + runtime: Incomplete + vpc_config: Incomplete + environment: Incomplete + layers: Incomplete + def __init__( + self, + function_arn: str = None, + function_name: str = None, + execution_role_arn: str = None, + zipped_code_dir: str = None, + s3_bucket: str = None, + script: str = None, + handler: str = None, + session: Session = None, + timeout: int = 120, + memory_size: int = 128, + runtime: str = "python3.8", + vpc_config: dict = None, + environment: dict = None, + layers: list = None, + ) -> None: ... + def create(self): ... + def update(self): ... + def upsert(self): ... + def invoke(self): ... + def delete(self): ... diff --git a/stubs/sagemaker/sagemaker/lineage/__init__.pyi b/stubs/sagemaker/sagemaker/lineage/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/lineage/_api_types.pyi b/stubs/sagemaker/sagemaker/lineage/_api_types.pyi new file mode 100644 index 000000000000..098ac996e460 --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/_api_types.pyi @@ -0,0 +1,70 @@ +from _typeshed import Incomplete + +from sagemaker.apiutils import _base_types + +class ArtifactSource(_base_types.ApiObject): + source_uri: Incomplete + source_types: Incomplete + def __init__(self, source_uri: Incomplete | None = None, source_types: Incomplete | None = None, **kwargs) -> None: ... + +class ArtifactSourceType(_base_types.ApiObject): + source_id_type: Incomplete + value: Incomplete + def __init__(self, source_id_type: Incomplete | None = None, value: Incomplete | None = None, **kwargs) -> None: ... + +class ActionSource(_base_types.ApiObject): + source_uri: Incomplete + source_type: Incomplete + def __init__(self, source_uri: Incomplete | None = None, source_type: Incomplete | None = None, **kwargs) -> None: ... + +class ContextSource(_base_types.ApiObject): + source_uri: Incomplete + source_type: Incomplete + def __init__(self, source_uri: Incomplete | None = None, source_type: Incomplete | None = None, **kwargs) -> None: ... + +class ArtifactSummary(_base_types.ApiObject): + artifact_arn: Incomplete + artifact_name: Incomplete + source: Incomplete + artifact_type: Incomplete + creation_time: Incomplete + last_modified_time: Incomplete + +class ActionSummary(_base_types.ApiObject): + action_arn: Incomplete + action_name: Incomplete + source: Incomplete + action_type: Incomplete + status: Incomplete + creation_time: Incomplete + last_modified_time: Incomplete + +class ContextSummary(_base_types.ApiObject): + context_arn: Incomplete + context_name: Incomplete + source: Incomplete + context_type: Incomplete + creation_time: Incomplete + last_modified_time: Incomplete + +class UserContext(_base_types.ApiObject): + user_profile_arn: Incomplete + user_profile_name: Incomplete + domain_id: Incomplete + def __init__( + self, + user_profile_arn: Incomplete | None = None, + user_profile_name: Incomplete | None = None, + domain_id: Incomplete | None = None, + **kwargs, + ) -> None: ... + +class AssociationSummary(_base_types.ApiObject): + source_arn: Incomplete + source_name: Incomplete + destination_arn: Incomplete + destination_name: Incomplete + source_type: Incomplete + destination_type: Incomplete + association_type: Incomplete + creation_time: Incomplete diff --git a/stubs/sagemaker/sagemaker/lineage/_utils.pyi b/stubs/sagemaker/sagemaker/lineage/_utils.pyi new file mode 100644 index 000000000000..7a3fb4b1b7c5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/_utils.pyi @@ -0,0 +1 @@ +def get_resource_name_from_arn(arn): ... diff --git a/stubs/sagemaker/sagemaker/lineage/action.pyi b/stubs/sagemaker/sagemaker/lineage/action.pyi new file mode 100644 index 000000000000..c03020e8a82b --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/action.pyi @@ -0,0 +1,62 @@ +from _typeshed import Incomplete +from datetime import datetime +from typing import Iterator, List, Optional + +from sagemaker.apiutils import _base_types +from sagemaker.lineage._api_types import ActionSource, ActionSummary +from sagemaker.lineage.artifact import Artifact +from sagemaker.lineage.query import LineageQueryDirectionEnum +from sagemaker.session import Session + +class Action(_base_types.Record): + action_arn: str + action_name: str + action_type: str + description: str + status: str + source: ActionSource + properties: dict + properties_to_remove: list + tags: list + creation_time: datetime + created_by: str + last_modified_time: datetime + last_modified_by: str + def save(self) -> Action: ... + def delete(self, disassociate: bool = False): ... + @classmethod + def load(cls, action_name: str, sagemaker_session: Incomplete | None = None) -> Action: ... + def set_tag(self, tag: Incomplete | None = None): ... + def set_tags(self, tags: Incomplete | None = None): ... + @classmethod + def create( + cls, + action_name: str = None, + source_uri: str = None, + source_type: str = None, + action_type: str = None, + description: str = None, + status: str = None, + properties: dict = None, + tags: dict = None, + sagemaker_session: Session = None, + ) -> Action: ... + @classmethod + def list( + cls, + source_uri: Optional[str] = None, + action_type: Optional[str] = None, + created_after: Optional[datetime] = None, + created_before: Optional[datetime] = None, + sort_by: Optional[str] = None, + sort_order: Optional[str] = None, + sagemaker_session: Session = None, + max_results: Optional[int] = None, + next_token: Optional[str] = None, + ) -> Iterator[ActionSummary]: ... + def artifacts(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... + +class ModelPackageApprovalAction(Action): + def datasets(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... + def model_package(self): ... + def endpoints(self, direction: LineageQueryDirectionEnum = ...) -> List: ... diff --git a/stubs/sagemaker/sagemaker/lineage/artifact.pyi b/stubs/sagemaker/sagemaker/lineage/artifact.pyi new file mode 100644 index 000000000000..85303057fbf3 --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/artifact.pyi @@ -0,0 +1,73 @@ +from _typeshed import Incomplete +from datetime import datetime +from typing import Iterator, List, Optional + +from sagemaker.apiutils import _base_types +from sagemaker.lineage._api_types import ArtifactSource, ArtifactSummary +from sagemaker.lineage.association import Association +from sagemaker.lineage.context import Context +from sagemaker.lineage.query import LineageQueryDirectionEnum + +LOGGER: Incomplete + +class Artifact(_base_types.Record): + artifact_arn: str + artifact_name: str + artifact_type: str + source: ArtifactSource + properties: dict + tags: list + creation_time: datetime + created_by: str + last_modified_time: datetime + last_modified_by: str + def save(self) -> Artifact: ... + def delete(self, disassociate: bool = False): ... + @classmethod + def load(cls, artifact_arn: str, sagemaker_session: Incomplete | None = None) -> Artifact: ... + def downstream_trials(self, sagemaker_session: Incomplete | None = None) -> list: ... + def downstream_trials_v2(self) -> list: ... + def upstream_trials(self) -> List: ... + def set_tag(self, tag: Incomplete | None = None): ... + def set_tags(self, tags: Incomplete | None = None): ... + @classmethod + def create( + cls, + artifact_name: Optional[str] = None, + source_uri: Optional[str] = None, + source_types: Optional[list] = None, + artifact_type: Optional[str] = None, + properties: Optional[dict] = None, + tags: Optional[dict] = None, + sagemaker_session: Incomplete | None = None, + ) -> Artifact: ... + @classmethod + def list( + cls, + source_uri: Optional[str] = None, + artifact_type: Optional[str] = None, + created_before: Optional[datetime] = None, + created_after: Optional[datetime] = None, + sort_by: Optional[str] = None, + sort_order: Optional[str] = None, + max_results: Optional[int] = None, + next_token: Optional[str] = None, + sagemaker_session: Incomplete | None = None, + ) -> Iterator[ArtifactSummary]: ... + def s3_uri_artifacts(self, s3_uri: str) -> dict: ... + +class ModelArtifact(Artifact): + def endpoints(self) -> list: ... + def endpoint_contexts(self, direction: LineageQueryDirectionEnum = ...) -> List[Context]: ... + def dataset_artifacts(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... + def training_job_arns(self, direction: LineageQueryDirectionEnum = ...) -> List[str]: ... + def pipeline_execution_arn(self, direction: LineageQueryDirectionEnum = ...) -> str: ... + +class DatasetArtifact(Artifact): + def trained_models(self) -> List[Association]: ... + def endpoint_contexts(self, direction: LineageQueryDirectionEnum = ...) -> List[Context]: ... + def upstream_datasets(self) -> List[Artifact]: ... + def downstream_datasets(self) -> List[Artifact]: ... + +class ImageArtifact(Artifact): + def datasets(self, direction: LineageQueryDirectionEnum) -> List[Artifact]: ... diff --git a/stubs/sagemaker/sagemaker/lineage/association.pyi b/stubs/sagemaker/sagemaker/lineage/association.pyi new file mode 100644 index 000000000000..8cd92310e20d --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/association.pyi @@ -0,0 +1,35 @@ +from _typeshed import Incomplete +from datetime import datetime +from typing import Iterator, Optional + +from sagemaker.apiutils import _base_types +from sagemaker.lineage._api_types import AssociationSummary + +logger: Incomplete + +class Association(_base_types.Record): + source_arn: str + destination_arn: str + def delete(self) -> None: ... + def set_tag(self, tag: Incomplete | None = None): ... + def set_tags(self, tags: Incomplete | None = None): ... + @classmethod + def create( + cls, source_arn: str, destination_arn: str, association_type: str = None, sagemaker_session: Incomplete | None = None + ) -> Association: ... + @classmethod + def list( + cls, + source_arn: str = None, + destination_arn: str = None, + source_type: str = None, + destination_type: str = None, + association_type: str = None, + created_after: Optional[datetime] = None, + created_before: Optional[datetime] = None, + sort_by: Optional[str] = None, + sort_order: Optional[str] = None, + max_results: Optional[int] = None, + next_token: Optional[str] = None, + sagemaker_session: Incomplete | None = None, + ) -> Iterator[AssociationSummary]: ... diff --git a/stubs/sagemaker/sagemaker/lineage/context.pyi b/stubs/sagemaker/sagemaker/lineage/context.pyi new file mode 100644 index 000000000000..8dbd7ce06291 --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/context.pyi @@ -0,0 +1,67 @@ +from _typeshed import Incomplete +from datetime import datetime +from typing import Iterator, List, Optional + +from sagemaker.apiutils import _base_types +from sagemaker.lineage import association +from sagemaker.lineage._api_types import ContextSummary +from sagemaker.lineage.action import Action +from sagemaker.lineage.artifact import Artifact +from sagemaker.lineage.lineage_trial_component import LineageTrialComponent +from sagemaker.lineage.query import LineageQueryDirectionEnum + +class Context(_base_types.Record): + context_arn: str + context_name: str + context_type: str + properties: dict + tags: list + creation_time: datetime + created_by: str + last_modified_time: datetime + last_modified_by: str + def save(self) -> Context: ... + def delete(self, disassociate: bool = False): ... + def set_tag(self, tag: Incomplete | None = None): ... + def set_tags(self, tags: Incomplete | None = None): ... + @classmethod + def load(cls, context_name: str, sagemaker_session: Incomplete | None = None) -> Context: ... + @classmethod + def create( + cls, + context_name: str = None, + source_uri: str = None, + source_type: str = None, + context_type: str = None, + description: str = None, + properties: dict = None, + tags: dict = None, + sagemaker_session: Incomplete | None = None, + ) -> Context: ... + @classmethod + def list( + cls, + source_uri: Optional[str] = None, + context_type: Optional[str] = None, + created_after: Optional[datetime] = None, + created_before: Optional[datetime] = None, + sort_by: Optional[str] = None, + sort_order: Optional[str] = None, + max_results: Optional[int] = None, + next_token: Optional[str] = None, + sagemaker_session: Incomplete | None = None, + ) -> Iterator[ContextSummary]: ... + def actions(self, direction: LineageQueryDirectionEnum) -> List[Action]: ... + +class EndpointContext(Context): + def models(self) -> List[association.Association]: ... + def models_v2(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... + def dataset_artifacts(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... + def training_job_arns(self, direction: LineageQueryDirectionEnum = ...) -> List[str]: ... + def processing_jobs(self, direction: LineageQueryDirectionEnum = ...) -> List[LineageTrialComponent]: ... + def transform_jobs(self, direction: LineageQueryDirectionEnum = ...) -> List[LineageTrialComponent]: ... + def trial_components(self, direction: LineageQueryDirectionEnum = ...) -> List[LineageTrialComponent]: ... + def pipeline_execution_arn(self, direction: LineageQueryDirectionEnum = ...) -> str: ... + +class ModelPackageGroup(Context): + def pipeline_execution_arn(self) -> str: ... diff --git a/stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi b/stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi new file mode 100644 index 000000000000..f7cec0b6ecad --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi @@ -0,0 +1,34 @@ +from _typeshed import Incomplete +from typing import List + +from sagemaker.apiutils import _base_types +from sagemaker.lineage.artifact import Artifact +from sagemaker.lineage.query import LineageQueryDirectionEnum + +LOGGER: Incomplete + +class LineageTrialComponent(_base_types.Record): + trial_component_name: Incomplete + trial_component_arn: Incomplete + display_name: Incomplete + source: Incomplete + status: Incomplete + start_time: Incomplete + end_time: Incomplete + creation_time: Incomplete + created_by: Incomplete + last_modified_time: Incomplete + last_modified_by: Incomplete + parameters: Incomplete + input_artifacts: Incomplete + output_artifacts: Incomplete + metrics: Incomplete + parameters_to_remove: Incomplete + input_artifacts_to_remove: Incomplete + output_artifacts_to_remove: Incomplete + tags: Incomplete + @classmethod + def load(cls, trial_component_name: str, sagemaker_session: Incomplete | None = None) -> LineageTrialComponent: ... + def pipeline_execution_arn(self) -> str: ... + def dataset_artifacts(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... + def models(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... diff --git a/stubs/sagemaker/sagemaker/lineage/query.pyi b/stubs/sagemaker/sagemaker/lineage/query.pyi new file mode 100644 index 000000000000..b17fefaccba4 --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/query.pyi @@ -0,0 +1,91 @@ +from _typeshed import Incomplete +from datetime import datetime +from enum import Enum +from typing import Any, Dict, List, Optional, Union + +class LineageEntityEnum(Enum): + TRIAL: str + ACTION: str + ARTIFACT: str + CONTEXT: str + TRIAL_COMPONENT: str + +class LineageSourceEnum(Enum): + CHECKPOINT: str + DATASET: str + ENDPOINT: str + IMAGE: str + MODEL: str + MODEL_DATA: str + MODEL_DEPLOYMENT: str + MODEL_GROUP: str + MODEL_REPLACE: str + TENSORBOARD: str + TRAINING_JOB: str + APPROVAL: str + PROCESSING_JOB: str + TRANSFORM_JOB: str + +class LineageQueryDirectionEnum(Enum): + BOTH: str + ASCENDANTS: str + DESCENDANTS: str + +class Edge: + source_arn: Incomplete + destination_arn: Incomplete + association_type: Incomplete + def __init__(self, source_arn: str, destination_arn: str, association_type: str) -> None: ... + def __hash__(self): ... + def __eq__(self, other): ... + +class Vertex: + arn: Incomplete + lineage_entity: Incomplete + lineage_source: Incomplete + def __init__(self, arn: str, lineage_entity: str, lineage_source: str, sagemaker_session) -> None: ... + def __hash__(self): ... + def __eq__(self, other): ... + def to_lineage_object(self): ... + +class PyvisVisualizer: + graph_styles: Incomplete + def __init__(self, graph_styles, pyvis_options: Optional[Dict[str, Any]] = None) -> None: ... + def render(self, elements, path: str = "lineage_graph_pyvis.html"): ... + +class LineageQueryResult: + edges: Incomplete + vertices: Incomplete + startarn: Incomplete + def __init__(self, edges: List[Edge] = None, vertices: List[Vertex] = None, startarn: List[str] = None) -> None: ... + def visualize(self, path: Optional[str] = "lineage_graph_pyvis.html", pyvis_options: Optional[Dict[str, Any]] = None): ... + +class LineageFilter: + entities: Incomplete + sources: Incomplete + created_before: Incomplete + created_after: Incomplete + modified_before: Incomplete + modified_after: Incomplete + properties: Incomplete + def __init__( + self, + entities: Optional[List[Union[LineageEntityEnum, str]]] = None, + sources: Optional[List[Union[LineageSourceEnum, str]]] = None, + created_before: Optional[datetime] = None, + created_after: Optional[datetime] = None, + modified_before: Optional[datetime] = None, + modified_after: Optional[datetime] = None, + properties: Optional[Dict[str, str]] = None, + ) -> None: ... + +class LineageQuery: + def __init__(self, sagemaker_session) -> None: ... + def query( + self, + start_arns: List[str], + direction: LineageQueryDirectionEnum = ..., + include_edges: bool = True, + query_filter: LineageFilter = None, + max_depth: int = 10, + ) -> LineageQueryResult: ... diff --git a/stubs/sagemaker/sagemaker/lineage/visualizer.pyi b/stubs/sagemaker/sagemaker/lineage/visualizer.pyi new file mode 100644 index 000000000000..5b330b12238d --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/visualizer.pyi @@ -0,0 +1,18 @@ +from typing import Optional + +from pandas import DataFrame as DataFrame + +class LineageTableVisualizer: + def __init__(self, sagemaker_session) -> None: ... + def show( + self, + trial_component_name: Optional[str] = None, + training_job_name: Optional[str] = None, + processing_job_name: Optional[str] = None, + pipeline_execution_step: Optional[object] = None, + model_package_arn: Optional[str] = None, + endpoint_arn: Optional[str] = None, + artifact_arn: Optional[str] = None, + context_arn: Optional[str] = None, + actions_arn: Optional[str] = None, + ) -> DataFrame: ... diff --git a/stubs/sagemaker/sagemaker/local/__init__.pyi b/stubs/sagemaker/sagemaker/local/__init__.pyi new file mode 100644 index 000000000000..05c031959b89 --- /dev/null +++ b/stubs/sagemaker/sagemaker/local/__init__.pyi @@ -0,0 +1,6 @@ +from .local_session import ( + LocalSagemakerClient as LocalSagemakerClient, + LocalSagemakerRuntimeClient as LocalSagemakerRuntimeClient, + LocalSession as LocalSession, + file_input as file_input, +) diff --git a/stubs/sagemaker/sagemaker/local/data.pyi b/stubs/sagemaker/sagemaker/local/data.pyi new file mode 100644 index 000000000000..56914227490a --- /dev/null +++ b/stubs/sagemaker/sagemaker/local/data.pyi @@ -0,0 +1,51 @@ +import abc +from _typeshed import Incomplete +from abc import abstractmethod +from collections.abc import Generator + +def get_data_source_instance(data_source, sagemaker_session): ... +def get_splitter_instance(split_type): ... +def get_batch_strategy_instance(strategy, splitter): ... + +class DataSource(metaclass=abc.ABCMeta): + @abstractmethod + def get_file_list(self): ... + @abstractmethod + def get_root_dir(self): ... + +class LocalFileDataSource(DataSource): + root_path: Incomplete + def __init__(self, root_path) -> None: ... + def get_file_list(self): ... + def get_root_dir(self): ... + +class S3DataSource(DataSource): + files: Incomplete + def __init__(self, bucket, prefix, sagemaker_session) -> None: ... + def get_file_list(self): ... + def get_root_dir(self): ... + +class Splitter(metaclass=abc.ABCMeta): + @abstractmethod + def split(self, file): ... + +class NoneSplitter(Splitter): + def split(self, filename) -> Generator[Incomplete, None, None]: ... + +class LineSplitter(Splitter): + def split(self, file) -> Generator[Incomplete, None, None]: ... + +class RecordIOSplitter(Splitter): + def split(self, file) -> Generator[Incomplete, None, None]: ... + +class BatchStrategy(metaclass=abc.ABCMeta): + splitter: Incomplete + def __init__(self, splitter) -> None: ... + @abstractmethod + def pad(self, file, size): ... + +class MultiRecordStrategy(BatchStrategy): + def pad(self, file, size: int = 6) -> Generator[Incomplete, None, None]: ... + +class SingleRecordStrategy(BatchStrategy): + def pad(self, file, size: int = 6) -> Generator[Incomplete, None, None]: ... diff --git a/stubs/sagemaker/sagemaker/local/entities.pyi b/stubs/sagemaker/sagemaker/local/entities.pyi new file mode 100644 index 000000000000..6bf810d392b6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/local/entities.pyi @@ -0,0 +1,152 @@ +import enum +from _typeshed import Incomplete + +logger: Incomplete +HEALTH_CHECK_TIMEOUT_LIMIT: int + +class _LocalProcessingJob: + container: Incomplete + state: str + start_time: Incomplete + end_time: Incomplete + processing_job_name: str + processing_inputs: Incomplete + processing_output_config: Incomplete + environment: Incomplete + def __init__(self, container) -> None: ... + def start(self, processing_inputs, processing_output_config, environment, processing_job_name) -> None: ... + def describe(self): ... + +class _LocalTrainingJob: + container: Incomplete + model_artifacts: Incomplete + state: str + start_time: Incomplete + end_time: Incomplete + environment: Incomplete + training_job_name: str + def __init__(self, container) -> None: ... + def start(self, input_data_config, output_data_config, hyperparameters, environment, job_name) -> None: ... + def describe(self): ... + +class _LocalTransformJob: + local_session: Incomplete + name: Incomplete + model_name: Incomplete + primary_container: Incomplete + container: Incomplete + start_time: Incomplete + end_time: Incomplete + batch_strategy: Incomplete + transform_resources: Incomplete + input_data: Incomplete + output_data: Incomplete + environment: Incomplete + state: Incomplete + def __init__(self, transform_job_name, model_name, local_session: Incomplete | None = None) -> None: ... + def start(self, input_data, output_data, transform_resources, **kwargs) -> None: ... + def describe(self): ... + +class _LocalModel: + model_name: Incomplete + primary_container: Incomplete + creation_time: Incomplete + def __init__(self, model_name, primary_container) -> None: ... + def describe(self): ... + +class _LocalEndpointConfig: + name: Incomplete + production_variants: Incomplete + tags: Incomplete + creation_time: Incomplete + def __init__(self, config_name, production_variants, tags: Incomplete | None = None) -> None: ... + def describe(self): ... + +class _LocalEndpoint: + local_session: Incomplete + name: Incomplete + endpoint_config: Incomplete + production_variant: Incomplete + tags: Incomplete + primary_container: Incomplete + container: Incomplete + create_time: Incomplete + state: Incomplete + def __init__( + self, endpoint_name, endpoint_config_name, tags: Incomplete | None = None, local_session: Incomplete | None = None + ) -> None: ... + def serve(self) -> None: ... + def stop(self) -> None: ... + def describe(self): ... + +class _LocalPipeline: + local_session: Incomplete + pipeline: Incomplete + pipeline_description: Incomplete + creation_time: Incomplete + last_modified_time: Incomplete + def __init__( + self, pipeline, pipeline_description: Incomplete | None = None, local_session: Incomplete | None = None + ) -> None: ... + def describe(self): ... + def start(self, **kwargs): ... + +class _LocalPipelineExecution: + pipeline: Incomplete + pipeline_execution_name: Incomplete + pipeline_execution_description: Incomplete + pipeline_execution_display_name: Incomplete + status: Incomplete + failure_reason: Incomplete + creation_time: Incomplete + last_modified_time: Incomplete + step_execution: Incomplete + pipeline_dag: Incomplete + pipeline_parameters: Incomplete + def __init__( + self, + execution_id, + pipeline, + PipelineParameters: Incomplete | None = None, + PipelineExecutionDescription: Incomplete | None = None, + PipelineExecutionDisplayName: Incomplete | None = None, + ) -> None: ... + def describe(self): ... + def list_steps(self): ... + def update_execution_success(self) -> None: ... + def update_execution_failure(self, step_name, failure_message) -> None: ... + def update_step_properties(self, step_name, step_properties) -> None: ... + def update_step_failure(self, step_name, failure_message) -> None: ... + def mark_step_executing(self, step_name) -> None: ... + +class _LocalPipelineExecutionStep: + name: Incomplete + type: Incomplete + description: Incomplete + display_name: Incomplete + status: Incomplete + failure_reason: Incomplete + properties: Incomplete + start_time: Incomplete + end_time: Incomplete + def __init__( + self, + name, + step_type, + description, + display_name: Incomplete | None = None, + start_time: Incomplete | None = None, + end_time: Incomplete | None = None, + status: Incomplete | None = None, + properties: Incomplete | None = None, + failure_reason: Incomplete | None = None, + ) -> None: ... + def update_step_properties(self, properties) -> None: ... + def update_step_failure(self, failure_message) -> None: ... + def mark_step_executing(self) -> None: ... + def to_list_steps_response(self): ... + +class _LocalExecutionStatus(enum.Enum): + EXECUTING: str + SUCCEEDED: str + FAILED: str diff --git a/stubs/sagemaker/sagemaker/local/exceptions.pyi b/stubs/sagemaker/sagemaker/local/exceptions.pyi new file mode 100644 index 000000000000..8c319847a1d0 --- /dev/null +++ b/stubs/sagemaker/sagemaker/local/exceptions.pyi @@ -0,0 +1,6 @@ +from _typeshed import Incomplete + +class StepExecutionException(Exception): + message: Incomplete + step_name: Incomplete + def __init__(self, step_name, message) -> None: ... diff --git a/stubs/sagemaker/sagemaker/local/image.pyi b/stubs/sagemaker/sagemaker/local/image.pyi new file mode 100644 index 000000000000..c9e896f56a27 --- /dev/null +++ b/stubs/sagemaker/sagemaker/local/image.pyi @@ -0,0 +1,54 @@ +from _typeshed import Incomplete +from threading import Thread + +CONTAINER_PREFIX: str +DOCKER_COMPOSE_FILENAME: str +DOCKER_COMPOSE_HTTP_TIMEOUT_ENV: str +DOCKER_COMPOSE_HTTP_TIMEOUT: str +REGION_ENV_NAME: str +TRAINING_JOB_NAME_ENV_NAME: str +S3_ENDPOINT_URL_ENV_NAME: str +SELINUX_ENABLED: Incomplete +logger: Incomplete + +class _SageMakerContainer: + sagemaker_session: Incomplete + instance_type: Incomplete + instance_count: Incomplete + image: Incomplete + container_entrypoint: Incomplete + container_arguments: Incomplete + hosts: Incomplete + container_root: Incomplete + container: Incomplete + def __init__( + self, + instance_type, + instance_count, + image, + sagemaker_session: Incomplete | None = None, + container_entrypoint: Incomplete | None = None, + container_arguments: Incomplete | None = None, + ) -> None: ... + def process(self, processing_inputs, processing_output_config, environment, processing_job_name) -> None: ... + def train(self, input_data_config, output_data_config, hyperparameters, environment, job_name): ... + def serve(self, model_dir, environment) -> None: ... + def stop_serving(self) -> None: ... + def retrieve_artifacts(self, compose_data, output_data_config, job_name): ... + def write_processing_config_files( + self, host, environment, processing_inputs, processing_output_config, processing_job_name + ) -> None: ... + def write_config_files(self, host, hyperparameters, input_data_config) -> None: ... + +class _HostingContainer(Thread): + command: Incomplete + process: Incomplete + def __init__(self, command) -> None: ... + def run(self) -> None: ... + def down(self) -> None: ... + +class _Volume: + container_dir: Incomplete + host_dir: Incomplete + map: Incomplete + def __init__(self, host_dir, container_dir: Incomplete | None = None, channel: Incomplete | None = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/local/local_session.pyi b/stubs/sagemaker/sagemaker/local/local_session.pyi new file mode 100644 index 000000000000..24219e25563f --- /dev/null +++ b/stubs/sagemaker/sagemaker/local/local_session.pyi @@ -0,0 +1,85 @@ +from _typeshed import Incomplete + +from sagemaker.session import Session + +logger: Incomplete + +class LocalSagemakerClient: + sagemaker_session: Incomplete + def __init__(self, sagemaker_session: Incomplete | None = None) -> None: ... + def create_processing_job( + self, + ProcessingJobName, + AppSpecification, + ProcessingResources, + Environment: Incomplete | None = None, + ProcessingInputs: Incomplete | None = None, + ProcessingOutputConfig: Incomplete | None = None, + **kwargs, + ) -> None: ... + def describe_processing_job(self, ProcessingJobName): ... + def create_training_job( + self, + TrainingJobName, + AlgorithmSpecification, + OutputDataConfig, + ResourceConfig, + InputDataConfig: Incomplete | None = None, + Environment: Incomplete | None = None, + **kwargs, + ) -> None: ... + def describe_training_job(self, TrainingJobName): ... + def create_transform_job( + self, TransformJobName, ModelName, TransformInput, TransformOutput, TransformResources, **kwargs + ) -> None: ... + def describe_transform_job(self, TransformJobName): ... + def create_model(self, ModelName, PrimaryContainer, *args, **kwargs) -> None: ... + def describe_model(self, ModelName): ... + def describe_endpoint_config(self, EndpointConfigName): ... + def create_endpoint_config(self, EndpointConfigName, ProductionVariants, Tags: Incomplete | None = None) -> None: ... + def describe_endpoint(self, EndpointName): ... + def create_endpoint(self, EndpointName, EndpointConfigName, Tags: Incomplete | None = None) -> None: ... + def update_endpoint(self, EndpointName, EndpointConfigName) -> None: ... + def delete_endpoint(self, EndpointName) -> None: ... + def delete_endpoint_config(self, EndpointConfigName) -> None: ... + def delete_model(self, ModelName) -> None: ... + def create_pipeline(self, pipeline, pipeline_description, **kwargs): ... + def update_pipeline(self, pipeline, pipeline_description, **kwargs): ... + def describe_pipeline(self, PipelineName): ... + def delete_pipeline(self, PipelineName): ... + def start_pipeline_execution(self, PipelineName, **kwargs): ... + +class LocalSagemakerRuntimeClient: + http: Incomplete + serving_port: int + config: Incomplete + def __init__(self, config: Incomplete | None = None) -> None: ... + def invoke_endpoint( + self, + Body, + EndpointName, + ContentType: Incomplete | None = None, + Accept: Incomplete | None = None, + CustomAttributes: Incomplete | None = None, + TargetModel: Incomplete | None = None, + TargetVariant: Incomplete | None = None, + InferenceId: Incomplete | None = None, + ): ... + +class LocalSession(Session): + s3_endpoint_url: Incomplete + def __init__( + self, + boto_session: Incomplete | None = None, + default_bucket: Incomplete | None = None, + s3_endpoint_url: Incomplete | None = None, + disable_local_code: bool = False, + sagemaker_config: dict = None, + default_bucket_prefix: Incomplete | None = None, + ) -> None: ... + def logs_for_job(self, job_name, wait: bool = False, poll: int = 5, log_type: str = "All") -> None: ... + def logs_for_processing_job(self, job_name, wait: bool = False, poll: int = 10) -> None: ... + +class file_input: + config: Incomplete + def __init__(self, fileUri, content_type: Incomplete | None = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/local/pipeline.pyi b/stubs/sagemaker/sagemaker/local/pipeline.pyi new file mode 100644 index 000000000000..983f3c699680 --- /dev/null +++ b/stubs/sagemaker/sagemaker/local/pipeline.pyi @@ -0,0 +1,49 @@ +import abc +from _typeshed import Incomplete +from abc import ABC, abstractmethod +from typing import Dict + +from sagemaker.workflow.steps import Step + +PRIMITIVES: Incomplete +BINARY_CONDITION_TYPES: Incomplete + +class LocalPipelineExecutor: + sagemaker_session: Incomplete + execution: Incomplete + pipeline_dag: Incomplete + local_sagemaker_client: Incomplete + def __init__(self, execution, sagemaker_session) -> None: ... + def execute(self): ... + def evaluate_step_arguments(self, step): ... + def evaluate_pipeline_variable(self, pipeline_variable, step_name): ... + +class _StepExecutor(ABC, metaclass=abc.ABCMeta): + pipline_executor: Incomplete + step: Incomplete + def __init__(self, pipeline_executor: LocalPipelineExecutor, step: Step) -> None: ... + @abstractmethod + def execute(self) -> Dict: ... + +class _TrainingStepExecutor(_StepExecutor): + def execute(self): ... + +class _ProcessingStepExecutor(_StepExecutor): + def execute(self): ... + +class _ConditionStepExecutor(_StepExecutor): + def execute(self): ... + +class _TransformStepExecutor(_StepExecutor): + def execute(self): ... + +class _CreateModelStepExecutor(_StepExecutor): + def execute(self): ... + +class _FailStepExecutor(_StepExecutor): + def execute(self) -> None: ... + +class _StepExecutorFactory: + pipeline_executor: Incomplete + def __init__(self, pipeline_executor: LocalPipelineExecutor) -> None: ... + def get(self, step: Step) -> _StepExecutor: ... diff --git a/stubs/sagemaker/sagemaker/local/utils.pyi b/stubs/sagemaker/sagemaker/local/utils.pyi new file mode 100644 index 000000000000..9a589a5d2f99 --- /dev/null +++ b/stubs/sagemaker/sagemaker/local/utils.pyi @@ -0,0 +1,11 @@ +from _typeshed import Incomplete + +logger: Incomplete + +def copy_directory_structure(destination_directory, relative_path) -> None: ... +def move_to_destination(source, destination, job_name, sagemaker_session): ... +def recursive_copy(source, destination) -> None: ... +def kill_child_processes(pid) -> None: ... +def get_child_process_ids(pid): ... +def get_docker_host(): ... +def get_using_dot_notation(dictionary, keys): ... diff --git a/stubs/sagemaker/sagemaker/logs.pyi b/stubs/sagemaker/sagemaker/logs.pyi new file mode 100644 index 000000000000..9574c01061db --- /dev/null +++ b/stubs/sagemaker/sagemaker/logs.pyi @@ -0,0 +1,20 @@ +from _typeshed import Incomplete +from collections.abc import Generator +from typing import NamedTuple + +class ColorWrap: + colorize: Incomplete + def __init__(self, force: bool = False) -> None: ... + def __call__(self, index, s) -> None: ... + +def argmin(arr, f): ... +def some(arr): ... + +class Position(NamedTuple): + timestamp: Incomplete + skip: Incomplete + +def multi_stream_iter( + client, log_group, streams, positions: Incomplete | None = None +) -> Generator[Incomplete, None, Incomplete]: ... +def log_stream(client, log_group, stream_name, start_time: int = 0, skip: int = 0) -> Generator[Incomplete, None, None]: ... diff --git a/stubs/sagemaker/sagemaker/metadata_properties.pyi b/stubs/sagemaker/sagemaker/metadata_properties.pyi new file mode 100644 index 000000000000..39a6b1a215ad --- /dev/null +++ b/stubs/sagemaker/sagemaker/metadata_properties.pyi @@ -0,0 +1,17 @@ +from _typeshed import Incomplete +from typing import Optional, Union + +from sagemaker.workflow.entities import PipelineVariable + +class MetadataProperties: + commit_id: Incomplete + repository: Incomplete + generated_by: Incomplete + project_id: Incomplete + def __init__( + self, + 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, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/metric_definitions.pyi b/stubs/sagemaker/sagemaker/metric_definitions.pyi new file mode 100644 index 000000000000..fb5b40136bed --- /dev/null +++ b/stubs/sagemaker/sagemaker/metric_definitions.pyi @@ -0,0 +1,12 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional + +logger: Incomplete + +def retrieve_default( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> Optional[List[Dict[str, str]]]: ... diff --git a/stubs/sagemaker/sagemaker/model.pyi b/stubs/sagemaker/sagemaker/model.pyi new file mode 100644 index 000000000000..6fdfab7de786 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model.pyi @@ -0,0 +1,211 @@ +import abc +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker.drift_check_baselines import DriftCheckBaselines +from sagemaker.inference_recommender.inference_recommender_mixin import InferenceRecommenderMixin +from sagemaker.metadata_properties import MetadataProperties +from sagemaker.model_metrics import ModelMetrics +from sagemaker.predictor import PredictorBase +from sagemaker.serverless import ServerlessInferenceConfig +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +LOGGER: Incomplete +NEO_ALLOWED_FRAMEWORKS: Incomplete +NEO_IOC_TARGET_DEVICES: Incomplete +NEO_MULTIVERSION_UNSUPPORTED: Incomplete + +class ModelBase(abc.ABC, metaclass=abc.ABCMeta): + @abc.abstractmethod + def deploy(self, *args, **kwargs) -> PredictorBase: ... + @abc.abstractmethod + def delete_model(self, *args, **kwargs) -> None: ... + +SCRIPT_PARAM_NAME: str +DIR_PARAM_NAME: str +CONTAINER_LOG_LEVEL_PARAM_NAME: str +JOB_NAME_PARAM_NAME: str +MODEL_SERVER_WORKERS_PARAM_NAME: str +SAGEMAKER_REGION_PARAM_NAME: str +SAGEMAKER_OUTPUT_LOCATION: str + +class Model(ModelBase, InferenceRecommenderMixin): + model_data: Incomplete + image_uri: Incomplete + predictor_cls: Incomplete + name: Incomplete + sagemaker_session: Incomplete + role: Incomplete + vpc_config: Incomplete + endpoint_name: Incomplete + inference_recommender_job_results: Incomplete + inference_recommendations: Incomplete + env: Incomplete + model_kms_key: Incomplete + image_config: Incomplete + entry_point: Incomplete + source_dir: Incomplete + dependencies: Incomplete + git_config: Incomplete + container_log_level: Incomplete + uploaded_code: Incomplete + repacked_model_data: Incomplete + def __init__( + self, + 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] = None, + 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] = 20, + dependencies: Optional[List[str]] = None, + git_config: Optional[Dict[str, str]] = None, + ) -> None: ... + def register( + self, + 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, + task: Optional[Union[str, PipelineVariable]] = None, + sample_payload_url: Optional[Union[str, PipelineVariable]] = None, + framework: Optional[Union[str, PipelineVariable]] = None, + framework_version: Optional[Union[str, PipelineVariable]] = None, + nearest_model_name: Optional[Union[str, PipelineVariable]] = None, + data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + ): ... + def create( + self, + 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, + ): ... + def prepare_container_def( + self, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... + def enable_network_isolation(self): ... + def package_for_edge( + self, + output_path, + model_name, + model_version, + role: Incomplete | None = None, + job_name: Incomplete | None = None, + resource_key: Incomplete | None = None, + s3_kms_key: Incomplete | None = None, + tags: Incomplete | None = None, + ): ... + def compile( + self, + target_instance_family, + input_shape, + output_path, + role: Incomplete | None = None, + tags: Incomplete | None = None, + job_name: Incomplete | None = None, + compile_max_run=900, + framework: Incomplete | None = None, + framework_version: Incomplete | None = None, + target_platform_os: Incomplete | None = None, + target_platform_arch: Incomplete | None = None, + target_platform_accelerator: Incomplete | None = None, + compiler_options: Incomplete | None = None, + ): ... + def deploy( + self, + initial_instance_count: Incomplete | None = None, + instance_type: Incomplete | None = None, + serializer: Incomplete | None = None, + deserializer: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + endpoint_name: Incomplete | None = None, + tags: Incomplete | None = None, + kms_key: Incomplete | None = None, + wait: bool = True, + data_capture_config: Incomplete | None = None, + async_inference_config: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + volume_size: Incomplete | None = None, + model_data_download_timeout: Incomplete | None = None, + container_startup_health_check_timeout: Incomplete | None = None, + inference_recommendation_id: Incomplete | None = None, + explainer_config: Incomplete | None = None, + **kwargs, + ): ... + def transformer( + self, + instance_count, + instance_type, + strategy: Incomplete | None = None, + assemble_with: Incomplete | None = None, + output_path: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + accept: Incomplete | None = None, + env: Incomplete | None = None, + max_concurrent_transforms: Incomplete | None = None, + max_payload: Incomplete | None = None, + tags: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + ): ... + def delete_model(self) -> None: ... + +class FrameworkModel(Model): + def __init__( + self, + model_data: Union[str, PipelineVariable], + image_uri: Union[str, PipelineVariable], + role: Optional[str] = None, + entry_point: Optional[str] = None, + 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] = 20, + code_location: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + dependencies: Optional[List[str]] = None, + git_config: Optional[Dict[str, str]] = None, + **kwargs, + ) -> None: ... + +MODEL_PACKAGE_ARN_PATTERN: str + +class ModelPackage(Model): + algorithm_arn: Incomplete + model_data: Incomplete + model_package_arn: Incomplete + def __init__( + self, + role: Incomplete | None = None, + model_data: Incomplete | None = None, + algorithm_arn: Incomplete | None = None, + model_package_arn: Incomplete | None = None, + **kwargs, + ) -> None: ... + def enable_network_isolation(self): ... diff --git a/stubs/sagemaker/sagemaker/model_card/__init__.pyi b/stubs/sagemaker/sagemaker/model_card/__init__.pyi new file mode 100644 index 000000000000..4a9b62292b45 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_card/__init__.pyi @@ -0,0 +1,25 @@ +from sagemaker.model_card.evaluation_metric_parsers import EvaluationMetricTypeEnum as EvaluationMetricTypeEnum +from sagemaker.model_card.model_card import ( + AdditionalInformation as AdditionalInformation, + BusinessDetails as BusinessDetails, + Environment as Environment, + EvaluationJob as EvaluationJob, + Function as Function, + HyperParameter as HyperParameter, + IntendedUses as IntendedUses, + Metric as Metric, + MetricGroup as MetricGroup, + ModelCard as ModelCard, + ModelOverview as ModelOverview, + ObjectiveFunction as ObjectiveFunction, + TrainingDetails as TrainingDetails, + TrainingJobDetails as TrainingJobDetails, + TrainingMetric as TrainingMetric, +) +from sagemaker.model_card.schema_constraints import ( + FacetEnum as FacetEnum, + MetricTypeEnum as MetricTypeEnum, + ModelCardStatusEnum as ModelCardStatusEnum, + ObjectiveFunctionEnum as ObjectiveFunctionEnum, + RiskRatingEnum as RiskRatingEnum, +) diff --git a/stubs/sagemaker/sagemaker/model_card/evaluation_metric_parsers.pyi b/stubs/sagemaker/sagemaker/model_card/evaluation_metric_parsers.pyi new file mode 100644 index 000000000000..5b8668b5cb03 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_card/evaluation_metric_parsers.pyi @@ -0,0 +1,24 @@ +import abc +from _typeshed import Incomplete +from abc import ABC +from enum import Enum + +class EvaluationMetricTypeEnum(str, Enum): + MODEL_CARD_METRIC_SCHEMA: str + CLARIFY_BIAS: str + CLARIFY_EXPLAINABILITY: str + REGRESSION: str + BINARY_CLASSIFICATION: str + MULTICLASS_CLASSIFICATION: str + +class ParserBase(ABC, metaclass=abc.ABCMeta): + def run(self, json_data: dict): ... + +class DefaultParser(ParserBase): ... +class ClarifyBiasParser(ParserBase): ... +class ClarifyExplainabilityParser(ParserBase): ... +class ModelMonitorModelQualityParserBase(ParserBase, metaclass=abc.ABCMeta): ... +class RegressionParser(ModelMonitorModelQualityParserBase): ... +class ClassificationParser(ModelMonitorModelQualityParserBase): ... + +EVALUATION_METRIC_PARSERS: Incomplete diff --git a/stubs/sagemaker/sagemaker/model_card/helpers.pyi b/stubs/sagemaker/sagemaker/model_card/helpers.pyi new file mode 100644 index 000000000000..b127d8f01cbc --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_card/helpers.pyi @@ -0,0 +1,65 @@ +import abc +import collections +import json +from _typeshed import Incomplete +from abc import ABC, abstractmethod +from enum import Enum +from typing import Any, List, Optional, Union + +from boto3.session import Session as Session + +logger: Incomplete + +class _JSONEncoder(json.JSONEncoder): + def default(self, o): ... + +class _DefaultToRequestDict: ... +class _DefaultFromDict: ... + +class _DescriptorBase(ABC, metaclass=abc.ABCMeta): + private_name: Incomplete + def __set_name__(self, owner: type, name: str): ... + def __get__(self, obj, objtype: type = None): ... + def __set__(self, obj: object, value: Any): ... + @abstractmethod + def validate(self, value): ... + @abstractmethod + def require_decode(self, value: dict): ... + @abstractmethod + def decode(self, value: dict): ... + +class _OneOf(_DescriptorBase): + options: Incomplete + enumerator: Incomplete + enumerator_reverse: Incomplete + def __init__(self, enumerator: Enum) -> None: ... + def validate(self, value) -> None: ... + def require_decode(self, value: Union[Enum, str]): ... + def decode(self, value: str): ... + +class _IsList(_DescriptorBase): + item_type: Incomplete + max_size: Incomplete + def __init__(self, item_type: object, max_size: Optional[int] = None) -> None: ... + def validate(self, value: List): ... + def require_decode(self, value: List): ... + def decode(self, value: List): ... + +class _IsModelCardObject(_DescriptorBase): + custom_class: Incomplete + def __init__(self, custom_class: object) -> None: ... + def validate(self, value: Union[dict, object]): ... + def require_decode(self, value: Union[dict, object]): ... + def decode(self, value: dict): ... + +class _MaxSizeArray(collections.abc.MutableSequence): + list: Incomplete + def __init__(self, max_size: int, item_type: Any, array: List = None) -> None: ... + def __len__(self) -> int: ... + def __getitem__(self, index): ... + def __delitem__(self, index) -> None: ... + def __setitem__(self, index: int, value: Any): ... + def __eq__(self, other: Any): ... + def check(self, value: Any): ... + def insert(self, index, value) -> None: ... + def to_map(self, key_attribute: str): ... diff --git a/stubs/sagemaker/sagemaker/model_card/model_card.pyi b/stubs/sagemaker/sagemaker/model_card/model_card.pyi new file mode 100644 index 000000000000..6ca1dcc1e2c0 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_card/model_card.pyi @@ -0,0 +1,272 @@ +from _typeshed import Incomplete +from datetime import datetime +from typing import List, Optional, Union + +from boto3.session import Session as boto3_Session +from sagemaker.model_card.evaluation_metric_parsers import EvaluationMetricTypeEnum +from sagemaker.model_card.helpers import _DefaultFromDict, _DefaultToRequestDict +from sagemaker.model_card.schema_constraints import ( + FacetEnum, + MetricTypeEnum, + ModelCardStatusEnum, + ObjectiveFunctionEnum, + RiskRatingEnum, +) +from sagemaker.session import Session + +logger: Incomplete + +class Environment(_DefaultToRequestDict, _DefaultFromDict): + container_image: Incomplete + def __init__(self, container_image: List[str]) -> None: ... + +class ModelOverview(_DefaultToRequestDict, _DefaultFromDict): + model_artifact: Incomplete + inference_environment: Incomplete + model_id: Incomplete + model_name: Incomplete + model_description: Incomplete + model_version: Incomplete + problem_type: Incomplete + algorithm_type: Incomplete + model_creator: Incomplete + model_owner: Incomplete + def __init__( + self, + model_id: Optional[str] = None, + model_name: Optional[str] = None, + model_description: Optional[str] = None, + model_version: Optional[Union[int, float]] = None, + problem_type: Optional[str] = None, + algorithm_type: Optional[str] = None, + model_creator: Optional[str] = None, + model_owner: Optional[str] = None, + model_artifact: Optional[List[str]] = None, + inference_environment: Optional[Environment] = None, + ) -> None: ... + @classmethod + def from_model_name(cls, model_name: str, sagemaker_session: Session = None, **kwargs): ... + +class IntendedUses(_DefaultToRequestDict, _DefaultFromDict): + risk_rating: Incomplete + purpose_of_model: Incomplete + intended_uses: Incomplete + factors_affecting_model_efficiency: Incomplete + explanations_for_risk_rating: Incomplete + def __init__( + self, + purpose_of_model: Optional[str] = None, + intended_uses: Optional[str] = None, + factors_affecting_model_efficiency: Optional[str] = None, + risk_rating: Optional[Union[RiskRatingEnum, str]] = ..., + explanations_for_risk_rating: Optional[str] = None, + ) -> None: ... + +class BusinessDetails(_DefaultToRequestDict, _DefaultFromDict): + business_problem: Incomplete + business_stakeholders: Incomplete + line_of_business: Incomplete + def __init__( + self, + business_problem: Optional[str] = None, + business_stakeholders: Optional[str] = None, + line_of_business: Optional[str] = None, + ) -> None: ... + +class Function(_DefaultToRequestDict, _DefaultFromDict): + function: Incomplete + facet: Incomplete + condition: Incomplete + def __init__( + self, + function: Optional[Union[ObjectiveFunctionEnum, str]] = None, + facet: Optional[Union[FacetEnum, str]] = None, + condition: Optional[str] = None, + ) -> None: ... + +class ObjectiveFunction(_DefaultToRequestDict, _DefaultFromDict): + function: Incomplete + notes: Incomplete + def __init__(self, function: Function, notes: Optional[str] = None) -> None: ... + +class Metric(_DefaultToRequestDict, _DefaultFromDict): + type: Incomplete + name: Incomplete + notes: Incomplete + x_axis_name: Incomplete + y_axis_name: Incomplete + def __init__( + self, + name: str, + type: Union[MetricTypeEnum, str], + value: Union[int, float, str, bool, List], + notes: Optional[str] = None, + x_axis_name: Optional[Union[str, list]] = None, + y_axis_name: Optional[Union[str, list]] = None, + ) -> None: ... + @property + def value(self): ... + @value.setter + def value(self, val: Union[int, float, str, bool, List]): ... + +class TrainingMetric(_DefaultToRequestDict, _DefaultFromDict): + name: Incomplete + value: Incomplete + notes: Incomplete + def __init__(self, name: str, value: Union[int, float], notes: Optional[str] = None) -> None: ... + +class HyperParameter(_DefaultToRequestDict, _DefaultFromDict): + name: Incomplete + value: Incomplete + def __init__(self, name: str, value: str) -> None: ... + +class TrainingJobDetails(_DefaultToRequestDict, _DefaultFromDict): + training_datasets: Incomplete + training_metrics: Incomplete + user_provided_training_metrics: Incomplete + hyper_parameters: Incomplete + user_provided_hyper_parameters: Incomplete + training_environment: Incomplete + training_arn: Incomplete + def __init__( + self, + training_arn: Optional[str] = None, + training_datasets: Optional[List[str]] = None, + training_environment: Optional[Environment] = None, + training_metrics: Optional[List[TrainingMetric]] = None, + user_provided_training_metrics: Optional[List[TrainingMetric]] = None, + hyper_parameters: Optional[List[HyperParameter]] = None, + user_provided_hyper_parameters: Optional[List[HyperParameter]] = None, + ) -> None: ... + +class TrainingDetails(_DefaultToRequestDict, _DefaultFromDict): + objective_function: Incomplete + training_job_details: Incomplete + training_observations: Incomplete + def __init__( + self, + objective_function: Optional[ObjectiveFunction] = None, + training_observations: Optional[str] = None, + training_job_details: Optional[TrainingJobDetails] = None, + ) -> None: ... + @classmethod + def from_model_overview(cls, model_overview: ModelOverview, sagemaker_session: Session = None, **kwargs): ... + @classmethod + def from_training_job_name(cls, training_job_name: str, sagemaker_session: Session = None, **kwargs): ... + def add_metric(self, metric: TrainingMetric): ... + def add_parameter(self, parameter: HyperParameter): ... + +class MetricGroup(_DefaultToRequestDict, _DefaultFromDict): + metric_data: Incomplete + name: Incomplete + def __init__(self, name: str, metric_data: Optional[List[Metric]] = None) -> None: ... + def add_metric(self, metric: Metric): ... + +class EvaluationJob(_DefaultToRequestDict, _DefaultFromDict): + datasets: Incomplete + metric_groups: Incomplete + name: Incomplete + evaluation_observation: Incomplete + evaluation_job_arn: Incomplete + metadata: Incomplete + def __init__( + self, + name: str, + evaluation_observation: Optional[str] = None, + evaluation_job_arn: Optional[str] = None, + datasets: Optional[List[str]] = None, + metadata: Optional[dict] = None, + metric_groups: Optional[List[MetricGroup]] = None, + ) -> None: ... + def get_metric_group(self, group_name): ... + def add_metric_group(self, group_name: str): ... + def add_metric_group_from_json(self, json_path: str, metric_type: EvaluationMetricTypeEnum = ...): ... + def add_metric_group_from_s3(self, session: boto3_Session, s3_url: str, metric_type: EvaluationMetricTypeEnum = ...): ... + +class AdditionalInformation(_DefaultToRequestDict, _DefaultFromDict): + ethical_considerations: Incomplete + caveats_and_recommendations: Incomplete + custom_details: Incomplete + def __init__( + self, + ethical_considerations: Optional[str] = None, + caveats_and_recommendations: Optional[str] = None, + custom_details: Optional[dict] = None, + ) -> None: ... + +class ModelCard: + DECODER_ATTRIBUTE_MAP: Incomplete + CREATE_MODEL_CARD_REQUIRED: Incomplete + INITIAL_VERSION: int + status: Incomplete + model_overview: Incomplete + intended_uses: Incomplete + business_details: Incomplete + training_details: Incomplete + evaluation_details: Incomplete + additional_information: Incomplete + name: Incomplete + arn: Incomplete + version: Incomplete + created_time: Incomplete + created_by: Incomplete + last_modified_time: Incomplete + last_modified_by: Incomplete + sagemaker_session: Incomplete + def __init__( + self, + name: str, + status: Optional[Union[ModelCardStatusEnum, str]] = ..., + arn: Optional[str] = None, + version: Optional[int] = None, + created_time: Optional[datetime] = None, + created_by: Optional[dict] = None, + last_modified_time: Optional[datetime] = None, + last_modified_by: Optional[dict] = None, + model_overview: Optional[ModelOverview] = None, + intended_uses: Optional[IntendedUses] = None, + business_details: Optional[BusinessDetails] = None, + training_details: Optional[TrainingDetails] = None, + evaluation_details: Optional[List[EvaluationJob]] = None, + additional_information: Optional[AdditionalInformation] = None, + sagemaker_session: Optional[Session] = None, + ) -> None: ... + def create(self): ... + @classmethod + def load(cls, name: str, version: Optional[int] = None, sagemaker_session: Session = None): ... + def update(self, **kwargs): ... + def delete(self): ... + def export_pdf( + self, s3_output_path: str, export_job_name: Optional[str] = None, model_card_version: Optional[int] = None + ): ... + def list_export_jobs(self, **kwargs): ... + def get_version_history(self, **kwargs): ... + +class ModelCardExportJob: + EXPORT_JOB_POLLING_FREQUENCY: int + export_job_name: Incomplete + model_card_name: Incomplete + model_card_version: Incomplete + s3_output_path: Incomplete + s3_export_artifacts: Incomplete + sagemaker_session: Incomplete + export_job_arn: Incomplete + status: Incomplete + failure_reason: Incomplete + def __init__( + self, + export_job_name: str, + model_card_name: str, + model_card_version: int, + s3_output_path: str, + s3_export_artifacts: Optional[str] = None, + export_job_arn: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + status: Optional[str] = None, + failure_reason: Optional[str] = None, + ) -> None: ... + def create(self): ... + @classmethod + def load(cls, export_job_arn: str, sagemaker_session: Session = None): ... + @staticmethod + def list_export_jobs(model_card_name: str, sagemaker_session: Optional[Session] = None, **kwargs): ... diff --git a/stubs/sagemaker/sagemaker/model_card/schema_constraints.pyi b/stubs/sagemaker/sagemaker/model_card/schema_constraints.pyi new file mode 100644 index 000000000000..21ec7e0bd1ca --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_card/schema_constraints.pyi @@ -0,0 +1,44 @@ +from _typeshed import Incomplete +from enum import Enum + +class ModelCardStatusEnum(str, Enum): + DRAFT: str + PENDING_REVIEW: str + APPROVED: str + ARCHIVED: str + +class RiskRatingEnum(str, Enum): + HIGH: str + MEDIUM: str + LOW: str + UNKNOWN: str + +class ObjectiveFunctionEnum(str, Enum): + MAXIMIZE: str + MINIMIZE: str + +class FacetEnum(str, Enum): + LOSS: str + ACCURACY: str + RMSE: str + MAE: str + AUC: str + +class MetricTypeEnum(str, Enum): + NUMBER: str + LINEAR_GRAPH: str + STRING: str + BOOLEAN: str + MATRIX: str + BAR_CHART: str + +METRIC_VALUE_TYPE_MAP: Incomplete +PYTHON_TYPE_TO_METRIC_VALUE_TYPE: Incomplete +MODEL_ARTIFACT_MAX_SIZE: int +ENVIRONMENT_CONTAINER_IMAGES_MAX_SIZE: int +TRAINING_DATASETS_MAX_SIZE: int +TRAINING_METRICS_MAX_SIZE: int +USER_PROVIDED_TRAINING_METRICS_MAX_SIZE: int +HYPER_PARAMETERS_MAX_SIZE: int +USER_PROVIDED_HYPER_PARAMETERS_MAX_SIZE: int +EVALUATION_DATASETS_MAX_SIZE: int diff --git a/stubs/sagemaker/sagemaker/model_metrics.pyi b/stubs/sagemaker/sagemaker/model_metrics.pyi new file mode 100644 index 000000000000..073bd95bc731 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_metrics.pyi @@ -0,0 +1,47 @@ +from _typeshed import Incomplete +from typing import Optional, Union + +from sagemaker.workflow.entities import PipelineVariable + +class ModelMetrics: + model_statistics: Incomplete + model_constraints: Incomplete + model_data_statistics: Incomplete + model_data_constraints: Incomplete + bias: Incomplete + bias_pre_training: Incomplete + bias_post_training: Incomplete + explainability: Incomplete + def __init__( + self, + model_statistics: Optional["MetricsSource"] = None, + model_constraints: Optional["MetricsSource"] = None, + model_data_statistics: Optional["MetricsSource"] = None, + model_data_constraints: Optional["MetricsSource"] = None, + bias: Optional["MetricsSource"] = None, + explainability: Optional["MetricsSource"] = None, + bias_pre_training: Optional["MetricsSource"] = None, + bias_post_training: Optional["MetricsSource"] = None, + ) -> None: ... + +class MetricsSource: + content_type: Incomplete + s3_uri: Incomplete + content_digest: Incomplete + def __init__( + self, + content_type: Union[str, PipelineVariable], + s3_uri: Union[str, PipelineVariable], + content_digest: Optional[Union[str, PipelineVariable]] = None, + ) -> None: ... + +class FileSource: + content_type: Incomplete + s3_uri: Incomplete + content_digest: Incomplete + def __init__( + self, + s3_uri: Union[str, PipelineVariable], + content_digest: Optional[Union[str, PipelineVariable]] = None, + content_type: Optional[Union[str, PipelineVariable]] = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/model_monitor/__init__.pyi b/stubs/sagemaker/sagemaker/model_monitor/__init__.pyi new file mode 100644 index 000000000000..3ddced31ea70 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_monitor/__init__.pyi @@ -0,0 +1,32 @@ +from sagemaker.model_monitor.clarify_model_monitoring import ( + BiasAnalysisConfig as BiasAnalysisConfig, + ExplainabilityAnalysisConfig as ExplainabilityAnalysisConfig, + ModelBiasMonitor as ModelBiasMonitor, + ModelExplainabilityMonitor as ModelExplainabilityMonitor, +) +from sagemaker.model_monitor.cron_expression_generator import CronExpressionGenerator as CronExpressionGenerator +from sagemaker.model_monitor.data_capture_config import DataCaptureConfig as DataCaptureConfig +from sagemaker.model_monitor.data_quality_monitoring_config import ( + DataQualityDistributionConstraints as DataQualityDistributionConstraints, + DataQualityMonitoringConfig as DataQualityMonitoringConfig, +) +from sagemaker.model_monitor.dataset_format import ( + DatasetFormat as DatasetFormat, + MonitoringDatasetFormat as MonitoringDatasetFormat, +) +from sagemaker.model_monitor.model_monitoring import ( + BaseliningJob as BaseliningJob, + BatchTransformInput as BatchTransformInput, + DefaultModelMonitor as DefaultModelMonitor, + EndpointInput as EndpointInput, + ModelMonitor as ModelMonitor, + ModelQualityMonitor as ModelQualityMonitor, + MonitoringExecution as MonitoringExecution, + MonitoringOutput as MonitoringOutput, +) +from sagemaker.model_monitor.monitoring_files import ( + Constraints as Constraints, + ConstraintViolations as ConstraintViolations, + Statistics as Statistics, +) +from sagemaker.network import NetworkConfig as NetworkConfig diff --git a/stubs/sagemaker/sagemaker/model_monitor/clarify_model_monitoring.pyi b/stubs/sagemaker/sagemaker/model_monitor/clarify_model_monitoring.pyi new file mode 100644 index 000000000000..84ddf559964a --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_monitor/clarify_model_monitoring.pyi @@ -0,0 +1,186 @@ +from _typeshed import Incomplete + +from sagemaker.model_monitor import model_monitoring as mm + +class ClarifyModelMonitor(mm.ModelMonitor): + latest_baselining_job_config: Incomplete + def __init__( + self, + role: Incomplete | None = None, + instance_count: int = 1, + instance_type: str = "ml.m5.xlarge", + volume_size_in_gb: int = 30, + volume_kms_key: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + max_runtime_in_seconds: Incomplete | None = None, + base_job_name: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + env: Incomplete | None = None, + tags: Incomplete | None = None, + network_config: Incomplete | None = None, + ) -> None: ... + def run_baseline(self, **_) -> None: ... + def latest_monitoring_statistics(self, **_) -> None: ... + def list_executions(self): ... + +class ModelBiasMonitor(ClarifyModelMonitor): + JOB_DEFINITION_BASE_NAME: str + @classmethod + def monitoring_type(cls): ... + latest_baselining_job_config: Incomplete + latest_baselining_job_name: Incomplete + latest_baselining_job: Incomplete + def suggest_baseline( + self, + data_config, + bias_config, + model_config, + model_predicted_label_config: Incomplete | None = None, + wait: bool = False, + logs: bool = False, + job_name: Incomplete | None = None, + kms_key: Incomplete | None = None, + ): ... + job_definition_name: Incomplete + monitoring_schedule_name: Incomplete + def create_monitoring_schedule( + self, + endpoint_input: Incomplete | None = None, + ground_truth_input: Incomplete | None = None, + analysis_config: Incomplete | None = None, + output_s3_uri: Incomplete | None = None, + constraints: Incomplete | None = None, + monitor_schedule_name: Incomplete | None = None, + schedule_cron_expression: Incomplete | None = None, + enable_cloudwatch_metrics: bool = True, + batch_transform_input: Incomplete | None = None, + ) -> None: ... + role: Incomplete + instance_count: Incomplete + instance_type: Incomplete + volume_size_in_gb: Incomplete + volume_kms_key: Incomplete + output_kms_key: Incomplete + max_runtime_in_seconds: Incomplete + env: Incomplete + network_config: Incomplete + def update_monitoring_schedule( + self, + endpoint_input: Incomplete | None = None, + ground_truth_input: Incomplete | None = None, + analysis_config: Incomplete | None = None, + output_s3_uri: Incomplete | None = None, + constraints: Incomplete | None = None, + schedule_cron_expression: Incomplete | None = None, + enable_cloudwatch_metrics: Incomplete | None = None, + role: Incomplete | None = None, + instance_count: Incomplete | None = None, + instance_type: Incomplete | None = None, + volume_size_in_gb: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + max_runtime_in_seconds: Incomplete | None = None, + env: Incomplete | None = None, + network_config: Incomplete | None = None, + batch_transform_input: Incomplete | None = None, + ) -> None: ... + def delete_monitoring_schedule(self) -> None: ... + @classmethod + def attach(cls, monitor_schedule_name, sagemaker_session: Incomplete | None = None): ... + +class BiasAnalysisConfig: + analysis_config: Incomplete + def __init__(self, bias_config, headers: Incomplete | None = None, label: Incomplete | None = None) -> None: ... + +class ModelExplainabilityMonitor(ClarifyModelMonitor): + JOB_DEFINITION_BASE_NAME: str + @classmethod + def monitoring_type(cls): ... + latest_baselining_job_config: Incomplete + latest_baselining_job_name: Incomplete + latest_baselining_job: Incomplete + def suggest_baseline( + self, + data_config, + explainability_config, + model_config, + model_scores: Incomplete | None = None, + wait: bool = False, + logs: bool = False, + job_name: Incomplete | None = None, + kms_key: Incomplete | None = None, + ): ... + job_definition_name: Incomplete + monitoring_schedule_name: Incomplete + def create_monitoring_schedule( + self, + endpoint_input: Incomplete | None = None, + analysis_config: Incomplete | None = None, + output_s3_uri: Incomplete | None = None, + constraints: Incomplete | None = None, + monitor_schedule_name: Incomplete | None = None, + schedule_cron_expression: Incomplete | None = None, + enable_cloudwatch_metrics: bool = True, + batch_transform_input: Incomplete | None = None, + ) -> None: ... + role: Incomplete + instance_count: Incomplete + instance_type: Incomplete + volume_size_in_gb: Incomplete + volume_kms_key: Incomplete + output_kms_key: Incomplete + max_runtime_in_seconds: Incomplete + env: Incomplete + network_config: Incomplete + def update_monitoring_schedule( + self, + endpoint_input: Incomplete | None = None, + analysis_config: Incomplete | None = None, + output_s3_uri: Incomplete | None = None, + constraints: Incomplete | None = None, + schedule_cron_expression: Incomplete | None = None, + enable_cloudwatch_metrics: Incomplete | None = None, + role: Incomplete | None = None, + instance_count: Incomplete | None = None, + instance_type: Incomplete | None = None, + volume_size_in_gb: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + max_runtime_in_seconds: Incomplete | None = None, + env: Incomplete | None = None, + network_config: Incomplete | None = None, + batch_transform_input: Incomplete | None = None, + ) -> None: ... + def delete_monitoring_schedule(self) -> None: ... + @classmethod + def attach(cls, monitor_schedule_name, sagemaker_session: Incomplete | None = None): ... + +class ExplainabilityAnalysisConfig: + analysis_config: Incomplete + def __init__( + self, explainability_config, model_config, headers: Incomplete | None = None, label_headers: Incomplete | None = None + ) -> None: ... + +class ClarifyBaseliningConfig: + analysis_config: Incomplete + features_attribute: Incomplete + inference_attribute: Incomplete + probability_attribute: Incomplete + probability_threshold_attribute: Incomplete + def __init__( + self, + analysis_config, + features_attribute: Incomplete | None = None, + inference_attribute: Incomplete | None = None, + probability_attribute: Incomplete | None = None, + probability_threshold_attribute: Incomplete | None = None, + ) -> None: ... + +class ClarifyBaseliningJob(mm.BaseliningJob): + def __init__(self, processing_job) -> None: ... + def baseline_statistics(self, **_) -> None: ... + def suggested_constraints(self, file_name: Incomplete | None = None, kms_key: Incomplete | None = None): ... + +class ClarifyMonitoringExecution(mm.MonitoringExecution): + def __init__(self, sagemaker_session, job_name, inputs, output, output_kms_key: Incomplete | None = None) -> None: ... + def statistics(self, **_) -> None: ... diff --git a/stubs/sagemaker/sagemaker/model_monitor/cron_expression_generator.pyi b/stubs/sagemaker/sagemaker/model_monitor/cron_expression_generator.pyi new file mode 100644 index 000000000000..0db67cfad275 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_monitor/cron_expression_generator.pyi @@ -0,0 +1,7 @@ +class CronExpressionGenerator: + @staticmethod + def hourly(): ... + @staticmethod + def daily(hour: int = 0): ... + @staticmethod + def daily_every_x_hours(hour_interval, starting_hour: int = 0): ... diff --git a/stubs/sagemaker/sagemaker/model_monitor/data_capture_config.pyi b/stubs/sagemaker/sagemaker/model_monitor/data_capture_config.pyi new file mode 100644 index 000000000000..d6fa249b4b28 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_monitor/data_capture_config.pyi @@ -0,0 +1,22 @@ +from _typeshed import Incomplete + +class DataCaptureConfig: + API_MAPPING: Incomplete + enable_capture: Incomplete + sampling_percentage: Incomplete + destination_s3_uri: Incomplete + kms_key_id: Incomplete + capture_options: Incomplete + csv_content_types: Incomplete + json_content_types: Incomplete + def __init__( + self, + enable_capture, + sampling_percentage: int = 20, + destination_s3_uri: Incomplete | None = None, + kms_key_id: Incomplete | None = None, + capture_options: Incomplete | None = None, + csv_content_types: Incomplete | None = None, + json_content_types: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/model_monitor/data_quality_monitoring_config.pyi b/stubs/sagemaker/sagemaker/model_monitor/data_quality_monitoring_config.pyi new file mode 100644 index 000000000000..eebff2c9040a --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_monitor/data_quality_monitoring_config.pyi @@ -0,0 +1,18 @@ +from _typeshed import Incomplete + +CHI_SQUARED_METHOD: str +L_INFINITY_METHOD: str + +class DataQualityDistributionConstraints: + categorical_drift_method: Incomplete + def __init__(self, categorical_drift_method: str = None) -> None: ... + @staticmethod + def valid_distribution_constraints(distribution_constraints): ... + @staticmethod + def valid_categorical_drift_method(categorical_drift_method): ... + +class DataQualityMonitoringConfig: + distribution_constraints: Incomplete + def __init__(self, distribution_constraints: DataQualityDistributionConstraints = None) -> None: ... + @staticmethod + def valid_monitoring_config(monitoring_config): ... diff --git a/stubs/sagemaker/sagemaker/model_monitor/dataset_format.pyi b/stubs/sagemaker/sagemaker/model_monitor/dataset_format.pyi new file mode 100644 index 000000000000..4b9b057467a5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_monitor/dataset_format.pyi @@ -0,0 +1,15 @@ +class DatasetFormat: + @staticmethod + def csv(header: bool = True, output_columns_position: str = "START"): ... + @staticmethod + def json(lines: bool = True): ... + @staticmethod + def sagemaker_capture_json(): ... + +class MonitoringDatasetFormat: + @staticmethod + def csv(header: bool = True): ... + @staticmethod + def json(lines: bool = True): ... + @staticmethod + def parquet(): ... diff --git a/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi b/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi new file mode 100644 index 000000000000..7d995cc19868 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi @@ -0,0 +1,382 @@ +from _typeshed import Incomplete +from typing import Optional, Union + +from sagemaker.model_monitor.dataset_format import MonitoringDatasetFormat +from sagemaker.processing import ProcessingJob + +DEFAULT_REPOSITORY_NAME: str +STATISTICS_JSON_DEFAULT_FILE_NAME: str +CONSTRAINTS_JSON_DEFAULT_FILE_NAME: str +CONSTRAINT_VIOLATIONS_JSON_DEFAULT_FILE_NAME: str +framework_name: str + +class ModelMonitor: + image_uri: Incomplete + instance_count: Incomplete + instance_type: Incomplete + entrypoint: Incomplete + volume_size_in_gb: Incomplete + max_runtime_in_seconds: Incomplete + base_job_name: Incomplete + sagemaker_session: Incomplete + tags: Incomplete + baselining_jobs: Incomplete + latest_baselining_job: Incomplete + arguments: Incomplete + latest_baselining_job_name: Incomplete + monitoring_schedule_name: Incomplete + job_definition_name: Incomplete + role: Incomplete + volume_kms_key: Incomplete + output_kms_key: Incomplete + network_config: Incomplete + env: Incomplete + def __init__( + self, + role: Incomplete | None = None, + image_uri: Incomplete | None = None, + instance_count: int = 1, + instance_type: str = "ml.m5.xlarge", + entrypoint: Incomplete | None = None, + volume_size_in_gb: int = 30, + volume_kms_key: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + max_runtime_in_seconds: Incomplete | None = None, + base_job_name: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + env: Incomplete | None = None, + tags: Incomplete | None = None, + network_config: Incomplete | None = None, + ) -> None: ... + def run_baseline( + self, + baseline_inputs, + output, + arguments: Incomplete | None = None, + wait: bool = True, + logs: bool = True, + job_name: Incomplete | None = None, + ) -> None: ... + def create_monitoring_schedule( + self, + endpoint_input: Incomplete | None = None, + output: Incomplete | None = None, + statistics: Incomplete | None = None, + constraints: Incomplete | None = None, + monitor_schedule_name: Incomplete | None = None, + schedule_cron_expression: Incomplete | None = None, + batch_transform_input: Incomplete | None = None, + arguments: Incomplete | None = None, + ) -> None: ... + def update_monitoring_schedule( + self, + endpoint_input: Incomplete | None = None, + output: Incomplete | None = None, + statistics: Incomplete | None = None, + constraints: Incomplete | None = None, + schedule_cron_expression: Incomplete | None = None, + instance_count: Incomplete | None = None, + instance_type: Incomplete | None = None, + entrypoint: Incomplete | None = None, + volume_size_in_gb: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + arguments: Incomplete | None = None, + max_runtime_in_seconds: Incomplete | None = None, + env: Incomplete | None = None, + network_config: Incomplete | None = None, + role: Incomplete | None = None, + image_uri: Incomplete | None = None, + batch_transform_input: Incomplete | None = None, + ) -> None: ... + def start_monitoring_schedule(self) -> None: ... + def stop_monitoring_schedule(self) -> None: ... + def delete_monitoring_schedule(self) -> None: ... + def baseline_statistics(self, file_name="statistics.json"): ... + def suggested_constraints(self, file_name="constraints.json"): ... + def latest_monitoring_statistics(self, file_name="statistics.json"): ... + def latest_monitoring_constraint_violations(self, file_name="constraint_violations.json"): ... + def describe_latest_baselining_job(self): ... + def describe_schedule(self): ... + def list_executions(self): ... + def update_monitoring_alert( + self, monitoring_alert_name: str, data_points_to_alert: Optional[int], evaluation_period: Optional[int] + ): ... + def list_monitoring_alerts(self, next_token: Optional[str] = None, max_results: Optional[int] = 10): ... + def list_monitoring_alert_history( + self, + monitoring_alert_name: Optional[str] = None, + sort_by: Optional[str] = "CreationTime", + sort_order: Optional[str] = "Descending", + next_token: Optional[str] = None, + max_results: Optional[int] = 10, + creation_time_before: Optional[str] = None, + creation_time_after: Optional[str] = None, + status_equals: Optional[str] = None, + ): ... + @classmethod + def attach(cls, monitor_schedule_name, sagemaker_session: Incomplete | None = None): ... + @classmethod + def monitoring_type(cls) -> None: ... + +class DefaultModelMonitor(ModelMonitor): + JOB_DEFINITION_BASE_NAME: str + def __init__( + self, + role: Incomplete | None = None, + instance_count: int = 1, + instance_type: str = "ml.m5.xlarge", + volume_size_in_gb: int = 30, + volume_kms_key: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + max_runtime_in_seconds: Incomplete | None = None, + base_job_name: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + env: Incomplete | None = None, + tags: Incomplete | None = None, + network_config: Incomplete | None = None, + ) -> None: ... + @classmethod + def monitoring_type(cls): ... + latest_baselining_job_name: Incomplete + latest_baselining_job: Incomplete + def suggest_baseline( + self, + baseline_dataset, + dataset_format, + record_preprocessor_script: Incomplete | None = None, + post_analytics_processor_script: Incomplete | None = None, + output_s3_uri: Incomplete | None = None, + wait: bool = True, + logs: bool = True, + job_name: Incomplete | None = None, + monitoring_config_override: Incomplete | None = None, + ): ... + job_definition_name: Incomplete + monitoring_schedule_name: Incomplete + def create_monitoring_schedule( + self, + endpoint_input: Incomplete | None = None, + record_preprocessor_script: Incomplete | None = None, + post_analytics_processor_script: Incomplete | None = None, + output_s3_uri: Incomplete | None = None, + constraints: Incomplete | None = None, + statistics: Incomplete | None = None, + monitor_schedule_name: Incomplete | None = None, + schedule_cron_expression: Incomplete | None = None, + enable_cloudwatch_metrics: bool = True, + batch_transform_input: Incomplete | None = None, + ) -> None: ... + env: Incomplete + instance_type: Incomplete + instance_count: Incomplete + volume_size_in_gb: Incomplete + volume_kms_key: Incomplete + output_kms_key: Incomplete + max_runtime_in_seconds: Incomplete + network_config: Incomplete + role: Incomplete + def update_monitoring_schedule( + self, + endpoint_input: Incomplete | None = None, + record_preprocessor_script: Incomplete | None = None, + post_analytics_processor_script: Incomplete | None = None, + output_s3_uri: Incomplete | None = None, + statistics: Incomplete | None = None, + constraints: Incomplete | None = None, + schedule_cron_expression: Incomplete | None = None, + instance_count: Incomplete | None = None, + instance_type: Incomplete | None = None, + volume_size_in_gb: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + max_runtime_in_seconds: Incomplete | None = None, + env: Incomplete | None = None, + network_config: Incomplete | None = None, + enable_cloudwatch_metrics: Incomplete | None = None, + role: Incomplete | None = None, + batch_transform_input: Incomplete | None = None, + ) -> None: ... + def delete_monitoring_schedule(self) -> None: ... + def run_baseline(self) -> None: ... + @classmethod + def attach(cls, monitor_schedule_name, sagemaker_session: Incomplete | None = None): ... + def latest_monitoring_statistics(self): ... + def latest_monitoring_constraint_violations(self): ... + +class ModelQualityMonitor(ModelMonitor): + JOB_DEFINITION_BASE_NAME: str + def __init__( + self, + role: Incomplete | None = None, + instance_count: int = 1, + instance_type: str = "ml.m5.xlarge", + volume_size_in_gb: int = 30, + volume_kms_key: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + max_runtime_in_seconds: Incomplete | None = None, + base_job_name: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + env: Incomplete | None = None, + tags: Incomplete | None = None, + network_config: Incomplete | None = None, + ) -> None: ... + @classmethod + def monitoring_type(cls): ... + latest_baselining_job_name: Incomplete + latest_baselining_job: Incomplete + def suggest_baseline( + self, + baseline_dataset, + dataset_format, + problem_type, + inference_attribute: Incomplete | None = None, + probability_attribute: Incomplete | None = None, + ground_truth_attribute: Incomplete | None = None, + probability_threshold_attribute: Incomplete | None = None, + post_analytics_processor_script: Incomplete | None = None, + output_s3_uri: Incomplete | None = None, + wait: bool = False, + logs: bool = False, + job_name: Incomplete | None = None, + ): ... + job_definition_name: Incomplete + monitoring_schedule_name: Incomplete + def create_monitoring_schedule( + self, + endpoint_input: Incomplete | None = None, + ground_truth_input: Incomplete | None = None, + problem_type: Incomplete | None = None, + record_preprocessor_script: Incomplete | None = None, + post_analytics_processor_script: Incomplete | None = None, + output_s3_uri: Incomplete | None = None, + constraints: Incomplete | None = None, + monitor_schedule_name: Incomplete | None = None, + schedule_cron_expression: Incomplete | None = None, + enable_cloudwatch_metrics: bool = True, + batch_transform_input: Incomplete | None = None, + ) -> None: ... + role: Incomplete + instance_count: Incomplete + instance_type: Incomplete + volume_size_in_gb: Incomplete + volume_kms_key: Incomplete + output_kms_key: Incomplete + max_runtime_in_seconds: Incomplete + env: Incomplete + network_config: Incomplete + def update_monitoring_schedule( + self, + endpoint_input: Incomplete | None = None, + ground_truth_input: Incomplete | None = None, + problem_type: Incomplete | None = None, + record_preprocessor_script: Incomplete | None = None, + post_analytics_processor_script: Incomplete | None = None, + output_s3_uri: Incomplete | None = None, + constraints: Incomplete | None = None, + schedule_cron_expression: Incomplete | None = None, + enable_cloudwatch_metrics: Incomplete | None = None, + role: Incomplete | None = None, + instance_count: Incomplete | None = None, + instance_type: Incomplete | None = None, + volume_size_in_gb: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + max_runtime_in_seconds: Incomplete | None = None, + env: Incomplete | None = None, + network_config: Incomplete | None = None, + batch_transform_input: Incomplete | None = None, + ) -> None: ... + def delete_monitoring_schedule(self) -> None: ... + @classmethod + def attach(cls, monitor_schedule_name, sagemaker_session: Incomplete | None = None): ... + +class BaseliningJob(ProcessingJob): + inputs: Incomplete + outputs: Incomplete + def __init__(self, sagemaker_session, job_name, inputs, outputs, output_kms_key: Incomplete | None = None) -> None: ... + @classmethod + def from_processing_job(cls, processing_job): ... + def baseline_statistics(self, file_name="statistics.json", kms_key: Incomplete | None = None): ... + def suggested_constraints(self, file_name="constraints.json", kms_key: Incomplete | None = None): ... + +class MonitoringExecution(ProcessingJob): + output: Incomplete + def __init__(self, sagemaker_session, job_name, inputs, output, output_kms_key: Incomplete | None = None) -> None: ... + @classmethod + def from_processing_arn(cls, sagemaker_session, processing_job_arn): ... + def statistics(self, file_name="statistics.json", kms_key: Incomplete | None = None): ... + def constraint_violations(self, file_name="constraint_violations.json", kms_key: Incomplete | None = None): ... + +class EndpointInput: + endpoint_name: Incomplete + destination: Incomplete + s3_input_mode: Incomplete + s3_data_distribution_type: Incomplete + start_time_offset: Incomplete + end_time_offset: Incomplete + features_attribute: Incomplete + inference_attribute: Incomplete + probability_attribute: Incomplete + probability_threshold_attribute: Incomplete + def __init__( + self, + endpoint_name, + destination, + s3_input_mode: str = "File", + s3_data_distribution_type: str = "FullyReplicated", + start_time_offset: Incomplete | None = None, + end_time_offset: Incomplete | None = None, + features_attribute: Incomplete | None = None, + inference_attribute: Incomplete | None = None, + probability_attribute: Incomplete | None = None, + probability_threshold_attribute: Incomplete | None = None, + ) -> None: ... + +class MonitoringInput: + start_time_offset: str + end_time_offset: str + features_attribute: str + inference_attribute: str + probability_attribute: Union[str, int] + probability_threshold_attribute: float + def __init__( + self, + start_time_offset, + end_time_offset, + features_attribute, + inference_attribute, + probability_attribute, + probability_threshold_attribute, + ) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class BatchTransformInput(MonitoringInput): + data_captured_destination_s3_uri: Incomplete + destination: Incomplete + s3_input_mode: Incomplete + s3_data_distribution_type: Incomplete + dataset_format: Incomplete + def __init__( + self, + data_captured_destination_s3_uri: str, + destination: str, + dataset_format: MonitoringDatasetFormat, + s3_input_mode: str = "File", + s3_data_distribution_type: str = "FullyReplicated", + start_time_offset: str = None, + end_time_offset: str = None, + features_attribute: str = None, + inference_attribute: str = None, + probability_attribute: str = None, + probability_threshold_attribute: str = None, + ) -> None: ... + +class MonitoringOutput: + source: Incomplete + destination: Incomplete + s3_upload_mode: Incomplete + def __init__(self, source, destination: Incomplete | None = None, s3_upload_mode: str = "Continuous") -> None: ... diff --git a/stubs/sagemaker/sagemaker/model_monitor/monitoring_alert.pyi b/stubs/sagemaker/sagemaker/model_monitor/monitoring_alert.pyi new file mode 100644 index 000000000000..20765531585d --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_monitor/monitoring_alert.pyi @@ -0,0 +1,41 @@ +class ModelDashboardIndicatorAction: + enabled: bool + def __init__(self, enabled) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class MonitoringAlertActions: + model_dashboard_indicator: ModelDashboardIndicatorAction + def __init__(self, model_dashboard_indicator) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class MonitoringAlertSummary: + alert_name: str + creation_time: str + last_modified_time: str + alert_status: str + data_points_to_alert: int + evaluation_period: int + actions: MonitoringAlertActions + def __init__( + self, alert_name, creation_time, last_modified_time, alert_status, data_points_to_alert, evaluation_period, actions + ) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class MonitoringAlertHistorySummary: + alert_name: str + creation_time: str + alert_status: str + def __init__(self, alert_name, creation_time, alert_status) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/model_monitor/monitoring_files.pyi b/stubs/sagemaker/sagemaker/model_monitor/monitoring_files.pyi new file mode 100644 index 000000000000..c463094fb5b6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_monitor/monitoring_files.pyi @@ -0,0 +1,79 @@ +from _typeshed import Incomplete + +NO_SUCH_KEY_CODE: str + +class ModelMonitoringFile: + body_dict: Incomplete + file_s3_uri: Incomplete + kms_key: Incomplete + session: Incomplete + def __init__(self, body_dict, file_s3_uri, kms_key, sagemaker_session) -> None: ... + def save(self, new_save_location_s3_uri: Incomplete | None = None): ... + +class Statistics(ModelMonitoringFile): + def __init__( + self, body_dict, statistics_file_s3_uri, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None + ) -> None: ... + @classmethod + def from_s3_uri( + cls, statistics_file_s3_uri, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None + ): ... + @classmethod + def from_string( + cls, + statistics_file_string, + kms_key: Incomplete | None = None, + file_name: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + ): ... + @classmethod + def from_file_path( + cls, statistics_file_path, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None + ): ... + +class Constraints(ModelMonitoringFile): + def __init__( + self, body_dict, constraints_file_s3_uri, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None + ) -> None: ... + @classmethod + def from_s3_uri( + cls, constraints_file_s3_uri, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None + ): ... + @classmethod + def from_string( + cls, + constraints_file_string, + kms_key: Incomplete | None = None, + file_name: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + ): ... + @classmethod + def from_file_path( + cls, constraints_file_path, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None + ): ... + def set_monitoring(self, enable_monitoring, feature_name: Incomplete | None = None) -> None: ... + +class ConstraintViolations(ModelMonitoringFile): + def __init__( + self, + body_dict, + constraint_violations_file_s3_uri, + kms_key: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + ) -> None: ... + @classmethod + def from_s3_uri( + cls, constraint_violations_file_s3_uri, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None + ): ... + @classmethod + def from_string( + cls, + constraint_violations_file_string, + kms_key: Incomplete | None = None, + file_name: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + ): ... + @classmethod + def from_file_path( + cls, constraint_violations_file_path, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None + ): ... diff --git a/stubs/sagemaker/sagemaker/model_uris.pyi b/stubs/sagemaker/sagemaker/model_uris.pyi new file mode 100644 index 000000000000..b3bcf06d516f --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_uris.pyi @@ -0,0 +1,13 @@ +from _typeshed import Incomplete +from typing import Optional + +logger: Incomplete + +def retrieve( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + model_scope: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> str: ... diff --git a/stubs/sagemaker/sagemaker/multidatamodel.pyi b/stubs/sagemaker/sagemaker/multidatamodel.pyi new file mode 100644 index 000000000000..b59f1bf5b502 --- /dev/null +++ b/stubs/sagemaker/sagemaker/multidatamodel.pyi @@ -0,0 +1,50 @@ +from _typeshed import Incomplete +from collections.abc import Generator +from typing import Optional, Union + +from sagemaker.model import Model +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +MULTI_MODEL_CONTAINER_MODE: str + +class MultiDataModel(Model): + name: Incomplete + model_data_prefix: Incomplete + model: Incomplete + container_mode: Incomplete + sagemaker_session: Incomplete + endpoint_name: Incomplete + s3_client: Incomplete + def __init__( + self, + name: str, + model_data_prefix: str, + model: Optional[Model] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + **kwargs, + ) -> None: ... + def prepare_container_def( + self, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... + def deploy( + self, + initial_instance_count, + instance_type, + serializer: Incomplete | None = None, + deserializer: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + endpoint_name: Incomplete | None = None, + tags: Incomplete | None = None, + kms_key: Incomplete | None = None, + wait: bool = True, + data_capture_config: Incomplete | None = None, + **kwargs, + ): ... + def add_model(self, model_data_source, model_data_path: Incomplete | None = None): ... + def list_models(self) -> Generator[Incomplete, None, None]: ... diff --git a/stubs/sagemaker/sagemaker/mxnet/__init__.pyi b/stubs/sagemaker/sagemaker/mxnet/__init__.pyi new file mode 100644 index 000000000000..8597b1ad1123 --- /dev/null +++ b/stubs/sagemaker/sagemaker/mxnet/__init__.pyi @@ -0,0 +1,3 @@ +from sagemaker.mxnet.estimator import MXNet as MXNet +from sagemaker.mxnet.model import MXNetModel as MXNetModel, MXNetPredictor as MXNetPredictor +from sagemaker.mxnet.processing import MXNetProcessor as MXNetProcessor diff --git a/stubs/sagemaker/sagemaker/mxnet/defaults.pyi b/stubs/sagemaker/sagemaker/mxnet/defaults.pyi new file mode 100644 index 000000000000..0555246a1ab9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/mxnet/defaults.pyi @@ -0,0 +1 @@ +LATEST_PY2_VERSION: str diff --git a/stubs/sagemaker/sagemaker/mxnet/estimator.pyi b/stubs/sagemaker/sagemaker/mxnet/estimator.pyi new file mode 100644 index 000000000000..f50c25a42d32 --- /dev/null +++ b/stubs/sagemaker/sagemaker/mxnet/estimator.pyi @@ -0,0 +1,33 @@ +from _typeshed import Incomplete +from typing import Dict, Optional, Union + +from sagemaker.estimator import Framework +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class MXNet(Framework): + framework_version: Incomplete + py_version: Incomplete + def __init__( + self, + entry_point: Union[str, PipelineVariable], + framework_version: Optional[str] = None, + py_version: Optional[str] = None, + source_dir: Optional[Union[str, PipelineVariable]] = None, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + distribution: Optional[Dict[str, str]] = None, + **kwargs, + ) -> None: ... + def create_model( + self, + model_server_workers: Incomplete | None = None, + role: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + entry_point: Incomplete | None = None, + source_dir: Incomplete | None = None, + dependencies: Incomplete | None = None, + image_uri: Incomplete | None = None, + **kwargs, + ): ... diff --git a/stubs/sagemaker/sagemaker/mxnet/model.pyi b/stubs/sagemaker/sagemaker/mxnet/model.pyi new file mode 100644 index 000000000000..7fb974332ed4 --- /dev/null +++ b/stubs/sagemaker/sagemaker/mxnet/model.pyi @@ -0,0 +1,69 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker import ModelMetrics +from sagemaker.drift_check_baselines import DriftCheckBaselines +from sagemaker.metadata_properties import MetadataProperties +from sagemaker.model import FrameworkModel +from sagemaker.predictor import Predictor +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class MXNetPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class MXNetModel(FrameworkModel): + framework_version: Incomplete + py_version: Incomplete + model_server_workers: Incomplete + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + entry_point: Optional[str] = None, + framework_version: str = "1.4.0", + py_version: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + predictor_cls: callable = ..., + model_server_workers: Optional[Union[int, PipelineVariable]] = None, + **kwargs, + ) -> None: ... + image_uri: Incomplete + def register( + self, + 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, + domain: Optional[Union[str, PipelineVariable]] = None, + sample_payload_url: Optional[Union[str, PipelineVariable]] = None, + task: Optional[Union[str, PipelineVariable]] = None, + framework: Optional[Union[str, PipelineVariable]] = None, + framework_version: Optional[Union[str, PipelineVariable]] = None, + nearest_model_name: Optional[Union[str, PipelineVariable]] = None, + data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + ): ... + def prepare_container_def( + self, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... + def serving_image_uri( + self, + region_name, + instance_type, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... diff --git a/stubs/sagemaker/sagemaker/mxnet/processing.pyi b/stubs/sagemaker/sagemaker/mxnet/processing.pyi new file mode 100644 index 000000000000..208e5a2a4b2e --- /dev/null +++ b/stubs/sagemaker/sagemaker/mxnet/processing.pyi @@ -0,0 +1,30 @@ +from typing import Dict, List, Optional, Union + +from sagemaker.mxnet.estimator import MXNet +from sagemaker.network import NetworkConfig +from sagemaker.processing import FrameworkProcessor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class MXNetProcessor(FrameworkProcessor): + estimator_cls = MXNet + def __init__( + self, + framework_version: str, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Union[int, PipelineVariable] = None, + instance_type: Union[str, PipelineVariable] = None, + py_version: str = "py3", + image_uri: Optional[Union[str, PipelineVariable]] = None, + command: Optional[List[str]] = None, + volume_size_in_gb: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + output_kms_key: Optional[Union[str, PipelineVariable]] = None, + code_location: Optional[str] = None, + max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + base_job_name: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + network_config: Optional[NetworkConfig] = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/network.pyi b/stubs/sagemaker/sagemaker/network.pyi new file mode 100644 index 000000000000..0fdc3ec1ab8b --- /dev/null +++ b/stubs/sagemaker/sagemaker/network.pyi @@ -0,0 +1,17 @@ +from _typeshed import Incomplete +from typing import List, Optional, Union + +from sagemaker.workflow.entities import PipelineVariable + +class NetworkConfig: + enable_network_isolation: Incomplete + security_group_ids: Incomplete + subnets: Incomplete + encrypt_inter_container_traffic: Incomplete + def __init__( + self, + enable_network_isolation: Union[bool, PipelineVariable] = None, + security_group_ids: Optional[List[Union[str, PipelineVariable]]] = None, + subnets: Optional[List[Union[str, PipelineVariable]]] = None, + encrypt_inter_container_traffic: Optional[Union[bool, PipelineVariable]] = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/parameter.pyi b/stubs/sagemaker/sagemaker/parameter.pyi new file mode 100644 index 000000000000..21e9a7dcdba9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/parameter.pyi @@ -0,0 +1,40 @@ +from _typeshed import Incomplete +from typing import Union + +from sagemaker.workflow.entities import PipelineVariable + +class ParameterRange: + __all_types__: Incomplete + min_value: Incomplete + max_value: Incomplete + scaling_type: Incomplete + def __init__( + self, + min_value: Union[int, float, PipelineVariable], + max_value: Union[int, float, PipelineVariable], + scaling_type: Union[str, PipelineVariable] = "Auto", + ) -> None: ... + def is_valid(self, value): ... + @classmethod + def cast_to_type(cls, value): ... + def as_tuning_range(self, name): ... + +class ContinuousParameter(ParameterRange): + __name__: str + @classmethod + def cast_to_type(cls, value): ... + +class CategoricalParameter(ParameterRange): + __name__: str + values: Incomplete + def __init__(self, values) -> None: ... + def as_tuning_range(self, name): ... + def as_json_range(self, name): ... + def is_valid(self, value): ... + @classmethod + def cast_to_type(cls, value): ... + +class IntegerParameter(ParameterRange): + __name__: str + @classmethod + def cast_to_type(cls, value): ... diff --git a/stubs/sagemaker/sagemaker/pipeline.pyi b/stubs/sagemaker/sagemaker/pipeline.pyi new file mode 100644 index 000000000000..eaae3e1654ab --- /dev/null +++ b/stubs/sagemaker/sagemaker/pipeline.pyi @@ -0,0 +1,86 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker import Model, ModelMetrics +from sagemaker.drift_check_baselines import DriftCheckBaselines +from sagemaker.metadata_properties import MetadataProperties +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class PipelineModel: + models: Incomplete + predictor_cls: Incomplete + name: Incomplete + endpoint_name: Incomplete + sagemaker_session: Incomplete + role: Incomplete + vpc_config: Incomplete + enable_network_isolation: Incomplete + def __init__( + self, + models: List[Model], + role: str = 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, + enable_network_isolation: Union[bool, PipelineVariable] = None, + ) -> None: ... + def pipeline_container_def(self, instance_type: Incomplete | None = None): ... + def deploy( + self, + initial_instance_count, + instance_type, + serializer: Incomplete | None = None, + deserializer: Incomplete | None = None, + endpoint_name: Incomplete | None = None, + tags: Incomplete | None = None, + wait: bool = True, + update_endpoint: bool = False, + data_capture_config: Incomplete | None = None, + kms_key: Incomplete | None = None, + volume_size: Incomplete | None = None, + model_data_download_timeout: Incomplete | None = None, + container_startup_health_check_timeout: Incomplete | None = None, + ): ... + def create(self, instance_type: str): ... + def register( + self, + 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, + domain: Optional[Union[str, PipelineVariable]] = None, + sample_payload_url: Optional[Union[str, PipelineVariable]] = None, + task: Optional[Union[str, PipelineVariable]] = None, + framework: Optional[Union[str, PipelineVariable]] = None, + framework_version: Optional[Union[str, PipelineVariable]] = None, + nearest_model_name: Optional[Union[str, PipelineVariable]] = None, + data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + ): ... + def transformer( + self, + instance_count, + instance_type, + strategy: Incomplete | None = None, + assemble_with: Incomplete | None = None, + output_path: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + accept: Incomplete | None = None, + env: Incomplete | None = None, + max_concurrent_transforms: Incomplete | None = None, + max_payload: Incomplete | None = None, + tags: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + ): ... + def delete_model(self) -> None: ... diff --git a/stubs/sagemaker/sagemaker/predictor.pyi b/stubs/sagemaker/sagemaker/predictor.pyi new file mode 100644 index 000000000000..96b2d7f73e7f --- /dev/null +++ b/stubs/sagemaker/sagemaker/predictor.pyi @@ -0,0 +1,14 @@ +from typing import Optional + +from sagemaker.base_predictor import Predictor, PredictorBase as PredictorBase, RealTimePredictor as RealTimePredictor +from sagemaker.session import Session + +def retrieve_default( + endpoint_name: str, + sagemaker_session: Optional[Session] = None, + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> Predictor: ... diff --git a/stubs/sagemaker/sagemaker/predictor_async.pyi b/stubs/sagemaker/sagemaker/predictor_async.pyi new file mode 100644 index 000000000000..be0f26675f0a --- /dev/null +++ b/stubs/sagemaker/sagemaker/predictor_async.pyi @@ -0,0 +1,44 @@ +from _typeshed import Incomplete + +class AsyncPredictor: + predictor: Incomplete + endpoint_name: Incomplete + sagemaker_session: Incomplete + s3_client: Incomplete + serializer: Incomplete + deserializer: Incomplete + name: Incomplete + def __init__(self, predictor, name: Incomplete | None = None) -> None: ... + def predict( + self, + data: Incomplete | None = None, + input_path: Incomplete | None = None, + initial_args: Incomplete | None = None, + inference_id: Incomplete | None = None, + waiter_config=..., + ): ... + def predict_async( + self, + data: Incomplete | None = None, + input_path: Incomplete | None = None, + initial_args: Incomplete | None = None, + inference_id: Incomplete | None = None, + ): ... + def update_endpoint( + self, + initial_instance_count: Incomplete | None = None, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + model_name: Incomplete | None = None, + tags: Incomplete | None = None, + kms_key: Incomplete | None = None, + data_capture_config_dict: Incomplete | None = None, + wait: bool = True, + ) -> None: ... + def delete_endpoint(self, delete_endpoint_config: bool = True) -> None: ... + def delete_model(self) -> None: ... + def enable_data_capture(self) -> None: ... + def disable_data_capture(self) -> None: ... + def update_data_capture_config(self, data_capture_config) -> None: ... + def list_monitors(self): ... + def endpoint_context(self): ... diff --git a/stubs/sagemaker/sagemaker/processing.pyi b/stubs/sagemaker/sagemaker/processing.pyi new file mode 100644 index 000000000000..a373476d2e75 --- /dev/null +++ b/stubs/sagemaker/sagemaker/processing.pyi @@ -0,0 +1,230 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker.apiutils._base_types import ApiObject +from sagemaker.dataset_definition.inputs import DatasetDefinition, S3Input +from sagemaker.job import _Job +from sagemaker.network import NetworkConfig +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class Processor: + JOB_CLASS_NAME: str + image_uri: Incomplete + instance_count: Incomplete + instance_type: Incomplete + entrypoint: Incomplete + volume_size_in_gb: Incomplete + max_runtime_in_seconds: Incomplete + base_job_name: Incomplete + tags: Incomplete + jobs: Incomplete + latest_job: Incomplete + arguments: Incomplete + sagemaker_session: Incomplete + output_kms_key: Incomplete + volume_kms_key: Incomplete + network_config: Incomplete + role: Incomplete + env: Incomplete + def __init__( + self, + role: str = None, + image_uri: Union[str, PipelineVariable] = None, + instance_count: Union[int, PipelineVariable] = None, + instance_type: Union[str, PipelineVariable] = None, + entrypoint: Optional[List[Union[str, PipelineVariable]]] = None, + volume_size_in_gb: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + output_kms_key: Optional[Union[str, PipelineVariable]] = None, + max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + base_job_name: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + network_config: Optional[NetworkConfig] = None, + ) -> None: ... + def run( + self, + inputs: Optional[List["ProcessingInput"]] = None, + outputs: Optional[List["ProcessingOutput"]] = None, + arguments: Optional[List[Union[str, PipelineVariable]]] = None, + wait: bool = True, + logs: bool = True, + job_name: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + kms_key: Optional[str] = None, + ): ... + +class ScriptProcessor(Processor): + command: Incomplete + def __init__( + self, + role: Optional[Union[str, PipelineVariable]] = None, + image_uri: Union[str, PipelineVariable] = None, + command: List[str] = None, + instance_count: Union[int, PipelineVariable] = None, + instance_type: Union[str, PipelineVariable] = None, + volume_size_in_gb: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + output_kms_key: Optional[Union[str, PipelineVariable]] = None, + max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + base_job_name: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + network_config: Optional[NetworkConfig] = None, + ) -> None: ... + def get_run_args( + self, code, inputs: Incomplete | None = None, outputs: Incomplete | None = None, arguments: Incomplete | None = None + ): ... + latest_job: Incomplete + def run( + self, + code: str, + inputs: Optional[List["ProcessingInput"]] = None, + outputs: Optional[List["ProcessingOutput"]] = None, + arguments: Optional[List[Union[str, PipelineVariable]]] = None, + wait: bool = True, + logs: bool = True, + job_name: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + kms_key: Optional[str] = None, + ): ... + +class ProcessingJob(_Job): + inputs: Incomplete + outputs: Incomplete + output_kms_key: Incomplete + def __init__(self, sagemaker_session, job_name, inputs, outputs, output_kms_key: Incomplete | None = None) -> None: ... + @classmethod + def start_new(cls, processor, inputs, outputs, experiment_config): ... + @classmethod + def from_processing_name(cls, sagemaker_session, processing_job_name): ... + @classmethod + def from_processing_arn(cls, sagemaker_session, processing_job_arn): ... + def wait(self, logs: bool = True) -> None: ... + def describe(self): ... + def stop(self) -> None: ... + @staticmethod + def prepare_app_specification(container_arguments, container_entrypoint, image_uri): ... + @staticmethod + def prepare_output_config(kms_key_id, outputs): ... + @staticmethod + def prepare_processing_resources(instance_count, instance_type, volume_kms_key_id, volume_size_in_gb): ... + @staticmethod + def prepare_stopping_condition(max_runtime_in_seconds): ... + +class ProcessingInput: + source: Incomplete + destination: Incomplete + input_name: Incomplete + s3_data_type: Incomplete + s3_input_mode: Incomplete + s3_data_distribution_type: Incomplete + s3_compression_type: Incomplete + s3_input: Incomplete + dataset_definition: Incomplete + app_managed: Incomplete + def __init__( + self, + source: Optional[Union[str, PipelineVariable]] = None, + destination: Optional[Union[str, PipelineVariable]] = None, + input_name: Optional[Union[str, PipelineVariable]] = None, + s3_data_type: Union[str, PipelineVariable] = "S3Prefix", + s3_input_mode: Union[str, PipelineVariable] = "File", + s3_data_distribution_type: Union[str, PipelineVariable] = "FullyReplicated", + s3_compression_type: Union[str, PipelineVariable] = "None", + s3_input: Optional[S3Input] = None, + dataset_definition: Optional[DatasetDefinition] = None, + app_managed: Union[bool, PipelineVariable] = False, + ) -> None: ... + +class ProcessingOutput: + source: Incomplete + destination: Incomplete + output_name: Incomplete + s3_upload_mode: Incomplete + app_managed: Incomplete + feature_store_output: Incomplete + def __init__( + self, + source: Optional[Union[str, PipelineVariable]] = None, + destination: Optional[Union[str, PipelineVariable]] = None, + output_name: Optional[Union[str, PipelineVariable]] = None, + s3_upload_mode: Union[str, PipelineVariable] = "EndOfJob", + app_managed: Union[bool, PipelineVariable] = False, + feature_store_output: Optional["FeatureStoreOutput"] = None, + ) -> None: ... + +class RunArgs: + code: Incomplete + inputs: Incomplete + outputs: Incomplete + arguments: Incomplete + def __init__(self, code, inputs, outputs, arguments) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class FeatureStoreOutput(ApiObject): + feature_group_name: Incomplete + +class FrameworkProcessor(ScriptProcessor): + framework_entrypoint_command: Incomplete + estimator_cls: Incomplete + framework_version: Incomplete + py_version: Incomplete + code_location: Incomplete + image_uri: Incomplete + base_job_name: Incomplete + def __init__( + self, + estimator_cls: type, + framework_version: str, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Union[int, PipelineVariable] = None, + instance_type: Union[str, PipelineVariable] = None, + py_version: str = "py3", + image_uri: Optional[Union[str, PipelineVariable]] = None, + command: Optional[List[str]] = None, + volume_size_in_gb: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + output_kms_key: Optional[Union[str, PipelineVariable]] = None, + code_location: Optional[str] = None, + max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + base_job_name: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + network_config: Optional[NetworkConfig] = None, + ) -> None: ... + def get_run_args( + self, + code, + source_dir: Incomplete | None = None, + dependencies: Incomplete | None = None, + git_config: Incomplete | None = None, + inputs: Incomplete | None = None, + outputs: Incomplete | None = None, + arguments: Incomplete | None = None, + job_name: Incomplete | None = None, + ): ... + def run( + self, + code: str, + source_dir: Optional[str] = None, + dependencies: Optional[List[str]] = None, + git_config: Optional[Dict[str, str]] = None, + inputs: Optional[List[ProcessingInput]] = None, + outputs: Optional[List[ProcessingOutput]] = None, + arguments: Optional[List[Union[str, PipelineVariable]]] = None, + wait: bool = True, + logs: bool = True, + job_name: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + kms_key: Optional[str] = None, + ): ... diff --git a/stubs/sagemaker/sagemaker/pytorch/__init__.pyi b/stubs/sagemaker/sagemaker/pytorch/__init__.pyi new file mode 100644 index 000000000000..4de54989013d --- /dev/null +++ b/stubs/sagemaker/sagemaker/pytorch/__init__.pyi @@ -0,0 +1,4 @@ +from sagemaker.pytorch.estimator import PyTorch as PyTorch +from sagemaker.pytorch.model import PyTorchModel as PyTorchModel, PyTorchPredictor as PyTorchPredictor +from sagemaker.pytorch.processing import PyTorchProcessor as PyTorchProcessor +from sagemaker.pytorch.training_compiler.config import TrainingCompilerConfig as TrainingCompilerConfig diff --git a/stubs/sagemaker/sagemaker/pytorch/defaults.pyi b/stubs/sagemaker/sagemaker/pytorch/defaults.pyi new file mode 100644 index 000000000000..0555246a1ab9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/pytorch/defaults.pyi @@ -0,0 +1 @@ +LATEST_PY2_VERSION: str diff --git a/stubs/sagemaker/sagemaker/pytorch/estimator.pyi b/stubs/sagemaker/sagemaker/pytorch/estimator.pyi new file mode 100644 index 000000000000..c1c05e86503a --- /dev/null +++ b/stubs/sagemaker/sagemaker/pytorch/estimator.pyi @@ -0,0 +1,40 @@ +from _typeshed import Incomplete +from typing import Dict, Optional, Union + +from sagemaker.estimator import Framework +from sagemaker.pytorch.training_compiler.config import TrainingCompilerConfig +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class PyTorch(Framework): + LAUNCH_PYTORCH_DDP_ENV_NAME: str + LAUNCH_TORCH_DISTRIBUTED_ENV_NAME: str + INSTANCE_TYPE_ENV_NAME: str + framework_version: Incomplete + py_version: Incomplete + distribution: Incomplete + compiler_config: Incomplete + def __init__( + self, + entry_point: Union[str, PipelineVariable], + framework_version: Optional[str] = None, + py_version: Optional[str] = None, + source_dir: Optional[Union[str, PipelineVariable]] = None, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + distribution: Optional[Dict] = None, + compiler_config: Optional[TrainingCompilerConfig] = None, + **kwargs, + ) -> None: ... + def hyperparameters(self): ... + def create_model( + self, + model_server_workers: Incomplete | None = None, + role: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + entry_point: Incomplete | None = None, + source_dir: Incomplete | None = None, + dependencies: Incomplete | None = None, + **kwargs, + ): ... diff --git a/stubs/sagemaker/sagemaker/pytorch/model.pyi b/stubs/sagemaker/sagemaker/pytorch/model.pyi new file mode 100644 index 000000000000..97f981405a07 --- /dev/null +++ b/stubs/sagemaker/sagemaker/pytorch/model.pyi @@ -0,0 +1,69 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker import ModelMetrics +from sagemaker.drift_check_baselines import DriftCheckBaselines +from sagemaker.metadata_properties import MetadataProperties +from sagemaker.model import FrameworkModel +from sagemaker.predictor import Predictor +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class PyTorchPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class PyTorchModel(FrameworkModel): + framework_version: Incomplete + py_version: Incomplete + model_server_workers: Incomplete + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + entry_point: Optional[str] = None, + framework_version: str = "1.3", + py_version: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + predictor_cls: callable = ..., + model_server_workers: Optional[Union[int, PipelineVariable]] = None, + **kwargs, + ) -> None: ... + image_uri: Incomplete + def register( + self, + 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, + domain: Optional[Union[str, PipelineVariable]] = None, + sample_payload_url: Optional[Union[str, PipelineVariable]] = None, + task: Optional[Union[str, PipelineVariable]] = None, + framework: Optional[Union[str, PipelineVariable]] = None, + framework_version: Optional[Union[str, PipelineVariable]] = None, + nearest_model_name: Optional[Union[str, PipelineVariable]] = None, + data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + ): ... + def prepare_container_def( + self, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... + def serving_image_uri( + self, + region_name, + instance_type, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... diff --git a/stubs/sagemaker/sagemaker/pytorch/processing.pyi b/stubs/sagemaker/sagemaker/pytorch/processing.pyi new file mode 100644 index 000000000000..d9d3586ade4e --- /dev/null +++ b/stubs/sagemaker/sagemaker/pytorch/processing.pyi @@ -0,0 +1,30 @@ +from typing import Dict, List, Optional, Union + +from sagemaker.network import NetworkConfig +from sagemaker.processing import FrameworkProcessor +from sagemaker.pytorch.estimator import PyTorch +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class PyTorchProcessor(FrameworkProcessor): + estimator_cls = PyTorch + def __init__( + self, + framework_version: str, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Union[int, PipelineVariable] = None, + instance_type: Union[str, PipelineVariable] = None, + py_version: str = "py3", + image_uri: Optional[Union[str, PipelineVariable]] = None, + command: Optional[List[str]] = None, + volume_size_in_gb: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + output_kms_key: Optional[Union[str, PipelineVariable]] = None, + code_location: Optional[str] = None, + max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + base_job_name: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + network_config: Optional[NetworkConfig] = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/pytorch/training_compiler/__init__.pyi b/stubs/sagemaker/sagemaker/pytorch/training_compiler/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi b/stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi new file mode 100644 index 000000000000..90146f5a7e03 --- /dev/null +++ b/stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi @@ -0,0 +1,14 @@ +from _typeshed import Incomplete +from typing import Union + +from sagemaker.training_compiler.config import TrainingCompilerConfig as BaseConfig +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class TrainingCompilerConfig(BaseConfig): + SUPPORTED_INSTANCE_CLASS_PREFIXES: Incomplete + SUPPORTED_INSTANCE_TYPES_WITH_EFA: Incomplete + def __init__(self, enabled: Union[bool, PipelineVariable] = True, debug: Union[bool, PipelineVariable] = False) -> None: ... + @classmethod + def validate(cls, estimator) -> None: ... diff --git a/stubs/sagemaker/sagemaker/remote_function/__init__.pyi b/stubs/sagemaker/sagemaker/remote_function/__init__.pyi new file mode 100644 index 000000000000..9d00a0f4e7c2 --- /dev/null +++ b/stubs/sagemaker/sagemaker/remote_function/__init__.pyi @@ -0,0 +1 @@ +from sagemaker.remote_function.client import RemoteExecutor as RemoteExecutor, remote as remote diff --git a/stubs/sagemaker/sagemaker/remote_function/client.pyi b/stubs/sagemaker/sagemaker/remote_function/client.pyi new file mode 100644 index 000000000000..6469cf282286 --- /dev/null +++ b/stubs/sagemaker/sagemaker/remote_function/client.pyi @@ -0,0 +1,101 @@ +from _typeshed import Incomplete +from collections.abc import Generator +from typing import Any, Dict, List, Tuple + +from sagemaker.remote_function.job import _JobSettings +from sagemaker.remote_function.spark_config import SparkConfig +from sagemaker.session import Session + +logger: Incomplete + +def remote( + _func: Incomplete | None = None, + *, + dependencies: str = None, + pre_execution_commands: List[str] = None, + pre_execution_script: str = None, + environment_variables: Dict[str, str] = None, + image_uri: str = None, + include_local_workdir: bool = False, + instance_count: int = 1, + instance_type: str = None, + job_conda_env: str = None, + job_name_prefix: str = None, + keep_alive_period_in_seconds: int = 0, + max_retry_attempts: int = 1, + max_runtime_in_seconds: int = 86400, + role: str = None, + s3_kms_key: str = None, + s3_root_uri: str = None, + sagemaker_session: Session = None, + security_group_ids: List[str] = None, + subnets: List[str] = None, + tags: List[Tuple[str, str]] = None, + volume_kms_key: str = None, + volume_size: int = 30, + encrypt_inter_container_traffic: bool = None, + spark_config: SparkConfig = None, +): ... + +class _SubmitRequest: + future: Incomplete + job_settings: Incomplete + func: Incomplete + args: Incomplete + kwargs: Incomplete + run_info: Incomplete + def __init__( + self, future, job_settings: _JobSettings, func, func_args, func_kwargs, run_info: Incomplete | None = None + ) -> None: ... + +class RemoteExecutor: + max_parallel_jobs: Incomplete + job_settings: Incomplete + def __init__( + self, + *, + dependencies: str = None, + pre_execution_commands: List[str] = None, + pre_execution_script: str = None, + environment_variables: Dict[str, str] = None, + image_uri: str = None, + include_local_workdir: bool = False, + instance_count: int = 1, + instance_type: str = None, + job_conda_env: str = None, + job_name_prefix: str = None, + keep_alive_period_in_seconds: int = 0, + max_parallel_jobs: int = 1, + max_retry_attempts: int = 1, + max_runtime_in_seconds: int = 86400, + role: str = None, + s3_kms_key: str = None, + s3_root_uri: str = None, + sagemaker_session: Session = None, + security_group_ids: List[str] = None, + subnets: List[str] = None, + tags: List[Tuple[str, str]] = None, + volume_kms_key: str = None, + volume_size: int = 30, + encrypt_inter_container_traffic: bool = None, + spark_config: SparkConfig = None, + ) -> None: ... + def submit(self, func, *args, **kwargs): ... + def map(self, func, *iterables): ... + def shutdown(self) -> None: ... + def __enter__(self): ... + def __exit__(self, exc_type, exc_val, exc_tb): ... + +class Future: + def __init__(self) -> None: ... + @staticmethod + def from_describe_response(describe_training_job_response, sagemaker_session): ... + def result(self, timeout: float = None) -> Any: ... + def wait(self, timeout: int = None) -> None: ... + def cancel(self) -> bool: ... + def running(self) -> bool: ... + def cancelled(self) -> bool: ... + def done(self) -> bool: ... + +def get_future(job_name, sagemaker_session: Incomplete | None = None) -> Future: ... +def list_futures(job_name_prefix, sagemaker_session: Incomplete | None = None) -> Generator[Incomplete, None, None]: ... diff --git a/stubs/sagemaker/sagemaker/remote_function/errors.pyi b/stubs/sagemaker/sagemaker/remote_function/errors.pyi new file mode 100644 index 000000000000..be433348287a --- /dev/null +++ b/stubs/sagemaker/sagemaker/remote_function/errors.pyi @@ -0,0 +1,14 @@ +from _typeshed import Incomplete + +DEFAULT_FAILURE_CODE: int +FAILURE_REASON_PATH: str + +class RemoteFunctionError(Exception): + message: Incomplete + def __init__(self, message) -> None: ... + +class ServiceError(RemoteFunctionError): ... +class SerializationError(RemoteFunctionError): ... +class DeserializationError(RemoteFunctionError): ... + +def handle_error(error, sagemaker_session, s3_base_uri, s3_kms_key, hmac_key) -> int: ... diff --git a/stubs/sagemaker/sagemaker/remote_function/invoke_function.pyi b/stubs/sagemaker/sagemaker/remote_function/invoke_function.pyi new file mode 100644 index 000000000000..cc491d579998 --- /dev/null +++ b/stubs/sagemaker/sagemaker/remote_function/invoke_function.pyi @@ -0,0 +1,3 @@ +SUCCESS_EXIT_CODE: int + +def main() -> None: ... diff --git a/stubs/sagemaker/sagemaker/remote_function/job.pyi b/stubs/sagemaker/sagemaker/remote_function/job.pyi new file mode 100644 index 000000000000..b105cb8f7344 --- /dev/null +++ b/stubs/sagemaker/sagemaker/remote_function/job.pyi @@ -0,0 +1,102 @@ +from _typeshed import Incomplete +from typing import Dict, List, Tuple + +from sagemaker.remote_function.spark_config import SparkConfig +from sagemaker.session import Session + +BOOTSTRAP_SCRIPT_NAME: str +ENTRYPOINT_SCRIPT_NAME: str +PRE_EXECUTION_SCRIPT_NAME: str +RUNTIME_MANAGER_SCRIPT_NAME: str +SPARK_APP_SCRIPT_NAME: str +RUNTIME_SCRIPTS_CHANNEL_NAME: str +REMOTE_FUNCTION_WORKSPACE: str +JOB_REMOTE_FUNCTION_WORKSPACE: str +SPARK_CONF_CHANNEL_NAME: str +SPARK_CONF_FILE_NAME: str +SPARK_SUBMIT_JARS_WORKSPACE: str +SPARK_SUBMIT_PY_FILES_WORKSPACE: str +SPARK_SUBMIT_FILES_WORKSPACE: str +SPARK_CONF_WORKSPACE: str +DEFAULT_SPARK_VERSION: str +DEFAULT_SPARK_CONTAINER_VERSION: str +KEY_EXPERIMENT_NAME: str +KEY_RUN_NAME: str +JOBS_CONTAINER_ENTRYPOINT: Incomplete +SPARK_APP_SCRIPT_PATH: Incomplete +ENTRYPOINT_SCRIPT: Incomplete +SPARK_ENTRYPOINT_SCRIPT: Incomplete +logger: Incomplete + +class _JobSettings: + sagemaker_session: Incomplete + environment_variables: Incomplete + image_uri: Incomplete + dependencies: Incomplete + pre_execution_commands: Incomplete + pre_execution_script: Incomplete + include_local_workdir: Incomplete + instance_type: Incomplete + instance_count: Incomplete + volume_size: Incomplete + max_runtime_in_seconds: Incomplete + max_retry_attempts: Incomplete + keep_alive_period_in_seconds: Incomplete + spark_config: Incomplete + job_conda_env: Incomplete + job_name_prefix: Incomplete + encrypt_inter_container_traffic: Incomplete + enable_network_isolation: bool + role: Incomplete + s3_root_uri: Incomplete + s3_kms_key: Incomplete + volume_kms_key: Incomplete + vpc_config: Incomplete + tags: Incomplete + def __init__( + self, + *, + dependencies: str = None, + pre_execution_commands: List[str] = None, + pre_execution_script: str = None, + environment_variables: Dict[str, str] = None, + image_uri: str = None, + include_local_workdir: bool = None, + instance_count: int = 1, + instance_type: str = None, + job_conda_env: str = None, + job_name_prefix: str = None, + keep_alive_period_in_seconds: int = 0, + max_retry_attempts: int = 1, + max_runtime_in_seconds: int = 86400, + role: str = None, + s3_kms_key: str = None, + s3_root_uri: str = None, + sagemaker_session: Session = None, + security_group_ids: List[str] = None, + subnets: List[str] = None, + tags: List[Tuple[str, str]] = None, + volume_kms_key: str = None, + volume_size: int = 30, + encrypt_inter_container_traffic: bool = None, + spark_config: SparkConfig = None, + ) -> None: ... + +class _Job: + job_name: Incomplete + s3_uri: Incomplete + sagemaker_session: Incomplete + hmac_key: Incomplete + def __init__(self, job_name: str, s3_uri: str, sagemaker_session: Session, hmac_key: str) -> None: ... + @staticmethod + def from_describe_response(describe_training_job_response, sagemaker_session): ... + @staticmethod + def start(job_settings: _JobSettings, func, func_args, func_kwargs, run_info: Incomplete | None = None): ... + def describe(self): ... + def stop(self) -> None: ... + def wait(self, timeout: int = None): ... + +class _RunInfo: + experiment_name: str + run_name: str + def __init__(self, experiment_name, run_name) -> None: ... diff --git a/stubs/sagemaker/sagemaker/remote_function/logging_config.pyi b/stubs/sagemaker/sagemaker/remote_function/logging_config.pyi new file mode 100644 index 000000000000..d4be154fbf7c --- /dev/null +++ b/stubs/sagemaker/sagemaker/remote_function/logging_config.pyi @@ -0,0 +1,7 @@ +import logging +import time + +class _UTCFormatter(logging.Formatter): + converter = time.gmtime + +def get_logger(): ... diff --git a/stubs/sagemaker/sagemaker/remote_function/runtime_environment/__init__.pyi b/stubs/sagemaker/sagemaker/remote_function/runtime_environment/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.pyi b/stubs/sagemaker/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.pyi new file mode 100644 index 000000000000..7ea7c94f36ef --- /dev/null +++ b/stubs/sagemaker/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.pyi @@ -0,0 +1,12 @@ +from _typeshed import Incomplete + +SUCCESS_EXIT_CODE: int +DEFAULT_FAILURE_CODE: int +REMOTE_FUNCTION_WORKSPACE: str +BASE_CHANNEL_PATH: str +FAILURE_REASON_PATH: str +PRE_EXECUTION_SCRIPT_NAME: str +JOB_REMOTE_FUNCTION_WORKSPACE: str +logger: Incomplete + +def main() -> None: ... diff --git a/stubs/sagemaker/sagemaker/remote_function/runtime_environment/runtime_environment_manager.pyi b/stubs/sagemaker/sagemaker/remote_function/runtime_environment/runtime_environment_manager.pyi new file mode 100644 index 000000000000..bc461b922124 --- /dev/null +++ b/stubs/sagemaker/sagemaker/remote_function/runtime_environment/runtime_environment_manager.pyi @@ -0,0 +1,19 @@ +import logging +import time +from _typeshed import Incomplete + +class _UTCFormatter(logging.Formatter): + converter = time.gmtime + +def get_logger(): ... + +logger: Incomplete + +class RuntimeEnvironmentManager: + def snapshot(self, dependencies: str = None) -> str: ... + def bootstrap(self, local_dependencies_file: str, client_python_version: str, conda_env: str = None): ... + def run_pre_exec_script(self, pre_exec_script_path: str): ... + +class RuntimeEnvironmentError(Exception): + message: Incomplete + def __init__(self, message) -> None: ... diff --git a/stubs/sagemaker/sagemaker/remote_function/runtime_environment/spark_app.pyi b/stubs/sagemaker/sagemaker/remote_function/runtime_environment/spark_app.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi b/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi new file mode 100644 index 000000000000..47becee8d587 --- /dev/null +++ b/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi @@ -0,0 +1,13 @@ +from typing import Dict, List, Optional, Union + +class SparkConfig: + submit_jars: Optional[List[str]] + submit_py_files: Optional[List[str]] + submit_files: Optional[List[str]] + configuration: Optional[Union[List[Dict], Dict]] + spark_event_logs_uri: Optional[str] + def __init__(self, submit_jars, submit_py_files, submit_files, configuration, spark_event_logs_uri) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/rl/__init__.pyi b/stubs/sagemaker/sagemaker/rl/__init__.pyi new file mode 100644 index 000000000000..e297fbb2824a --- /dev/null +++ b/stubs/sagemaker/sagemaker/rl/__init__.pyi @@ -0,0 +1,6 @@ +from sagemaker.rl.estimator import ( + TOOLKIT_FRAMEWORK_VERSION_MAP as TOOLKIT_FRAMEWORK_VERSION_MAP, + RLEstimator as RLEstimator, + RLFramework as RLFramework, + RLToolkit as RLToolkit, +) diff --git a/stubs/sagemaker/sagemaker/rl/estimator.pyi b/stubs/sagemaker/sagemaker/rl/estimator.pyi new file mode 100644 index 000000000000..11bd475cf6b6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/rl/estimator.pyi @@ -0,0 +1,55 @@ +import enum +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker.estimator import Framework +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete +SAGEMAKER_ESTIMATOR: str +SAGEMAKER_ESTIMATOR_VALUE: str +PYTHON_VERSION: str +TOOLKIT_FRAMEWORK_VERSION_MAP: Incomplete + +class RLToolkit(enum.Enum): + COACH: str + RAY: str + +class RLFramework(enum.Enum): + TENSORFLOW: str + MXNET: str + PYTORCH: str + +class RLEstimator(Framework): + COACH_LATEST_VERSION_TF: str + COACH_LATEST_VERSION_MXNET: str + RAY_LATEST_VERSION: str + toolkit: Incomplete + toolkit_version: Incomplete + framework: Incomplete + framework_version: Incomplete + def __init__( + self, + entry_point: Union[str, PipelineVariable], + toolkit: Optional[RLToolkit] = None, + toolkit_version: Optional[str] = None, + framework: Optional[Framework] = None, + source_dir: Optional[Union[str, PipelineVariable]] = None, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + metric_definitions: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + **kwargs, + ) -> None: ... + def create_model( + self, + role: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + entry_point: Incomplete | None = None, + source_dir: Incomplete | None = None, + dependencies: Incomplete | None = None, + **kwargs, + ): ... + def training_image_uri(self): ... + def hyperparameters(self): ... + @classmethod + def default_metric_definitions(cls, toolkit): ... diff --git a/stubs/sagemaker/sagemaker/s3.pyi b/stubs/sagemaker/sagemaker/s3.pyi new file mode 100644 index 000000000000..ca4102758b49 --- /dev/null +++ b/stubs/sagemaker/sagemaker/s3.pyi @@ -0,0 +1,32 @@ +import io +from _typeshed import Incomplete +from typing import Union + +from sagemaker.s3_utils import determine_bucket_and_prefix as determine_bucket_and_prefix + +logger: Incomplete + +class S3Uploader: + @staticmethod + def upload(local_path, desired_s3_uri, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None): ... + @staticmethod + def upload_string_as_file_body( + body: str, + desired_s3_uri: Incomplete | None = None, + kms_key: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + ): ... + @staticmethod + def upload_bytes( + b: Union[bytes, io.BytesIO], s3_uri, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None + ): ... + +class S3Downloader: + @staticmethod + def download(s3_uri, local_path, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None): ... + @staticmethod + def read_file(s3_uri, sagemaker_session: Incomplete | None = None) -> str: ... + @staticmethod + def read_bytes(s3_uri, sagemaker_session: Incomplete | None = None) -> bytes: ... + @staticmethod + def list(s3_uri, sagemaker_session: Incomplete | None = None): ... diff --git a/stubs/sagemaker/sagemaker/s3_utils.pyi b/stubs/sagemaker/sagemaker/s3_utils.pyi new file mode 100644 index 000000000000..18aeb6b8183e --- /dev/null +++ b/stubs/sagemaker/sagemaker/s3_utils.pyi @@ -0,0 +1,10 @@ +from _typeshed import Incomplete +from typing import Optional + +logger: Incomplete + +def parse_s3_url(url): ... +def s3_path_join(*args, with_end_slash: bool = False): ... +def determine_bucket_and_prefix( + bucket: Optional[str] = None, key_prefix: Optional[str] = None, sagemaker_session: Incomplete | None = None +): ... diff --git a/stubs/sagemaker/sagemaker/script_uris.pyi b/stubs/sagemaker/sagemaker/script_uris.pyi new file mode 100644 index 000000000000..e7d49b650701 --- /dev/null +++ b/stubs/sagemaker/sagemaker/script_uris.pyi @@ -0,0 +1,13 @@ +from _typeshed import Incomplete +from typing import Optional + +logger: Incomplete + +def retrieve( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + script_scope: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> str: ... diff --git a/stubs/sagemaker/sagemaker/serializers.pyi b/stubs/sagemaker/sagemaker/serializers.pyi new file mode 100644 index 000000000000..2c3776d78e9e --- /dev/null +++ b/stubs/sagemaker/sagemaker/serializers.pyi @@ -0,0 +1,29 @@ +from typing import List, Optional + +from sagemaker.base_serializers import ( + BaseSerializer, + CSVSerializer as CSVSerializer, + DataSerializer as DataSerializer, + IdentitySerializer as IdentitySerializer, + JSONLinesSerializer as JSONLinesSerializer, + JSONSerializer as JSONSerializer, + LibSVMSerializer as LibSVMSerializer, + NumpySerializer as NumpySerializer, + SimpleBaseSerializer as SimpleBaseSerializer, + SparseMatrixSerializer as SparseMatrixSerializer, +) + +def retrieve_options( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> List[BaseSerializer]: ... +def retrieve_default( + region: Optional[str] = None, + model_id: Optional[str] = None, + model_version: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, +) -> BaseSerializer: ... diff --git a/stubs/sagemaker/sagemaker/serverless/__init__.pyi b/stubs/sagemaker/sagemaker/serverless/__init__.pyi new file mode 100644 index 000000000000..575835c85dcd --- /dev/null +++ b/stubs/sagemaker/sagemaker/serverless/__init__.pyi @@ -0,0 +1,3 @@ +from sagemaker.serverless.model import LambdaModel as LambdaModel +from sagemaker.serverless.predictor import LambdaPredictor as LambdaPredictor +from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig as ServerlessInferenceConfig diff --git a/stubs/sagemaker/sagemaker/serverless/model.pyi b/stubs/sagemaker/sagemaker/serverless/model.pyi new file mode 100644 index 000000000000..1fb86950ce0c --- /dev/null +++ b/stubs/sagemaker/sagemaker/serverless/model.pyi @@ -0,0 +1 @@ +class LambdaModel: ... diff --git a/stubs/sagemaker/sagemaker/serverless/predictor.pyi b/stubs/sagemaker/sagemaker/serverless/predictor.pyi new file mode 100644 index 000000000000..d3e6f24393c5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/serverless/predictor.pyi @@ -0,0 +1 @@ +class LambdaPredictor: ... diff --git a/stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi b/stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi new file mode 100644 index 000000000000..211edd548eeb --- /dev/null +++ b/stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi @@ -0,0 +1,10 @@ +from _typeshed import Incomplete +from typing import Optional + +class ServerlessInferenceConfig: + memory_size_in_mb: Incomplete + max_concurrency: Incomplete + provisioned_concurrency: Incomplete + def __init__( + self, memory_size_in_mb: int = 2048, max_concurrency: int = 5, provisioned_concurrency: Optional[int] = None + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/session.pyi b/stubs/sagemaker/sagemaker/session.pyi new file mode 100644 index 000000000000..64eb768a84c7 --- /dev/null +++ b/stubs/sagemaker/sagemaker/session.pyi @@ -0,0 +1,585 @@ +from _typeshed import Incomplete +from typing import Any, Dict, List, Optional, Sequence + +from sagemaker.inputs import BatchDataCaptureConfig + +LOGGER: Incomplete +NOTEBOOK_METADATA_FILE: str + +class LogState: + STARTING: int + WAIT_IN_PROGRESS: int + TAILING: int + JOB_COMPLETE: int + COMPLETE: int + +class Session: + default_bucket_prefix: Incomplete + s3_resource: Incomplete + s3_client: Incomplete + resource_groups_client: Incomplete + resource_group_tagging_client: Incomplete + config: Incomplete + lambda_client: Incomplete + settings: Incomplete + def __init__( + self, + boto_session: Incomplete | None = None, + sagemaker_client: Incomplete | None = None, + sagemaker_runtime_client: Incomplete | None = None, + sagemaker_featurestore_runtime_client: Incomplete | None = None, + default_bucket: Incomplete | None = None, + settings=..., + sagemaker_metrics_client: Incomplete | None = None, + sagemaker_config: dict = None, + default_bucket_prefix: str = None, + ) -> None: ... + @property + def boto_region_name(self): ... + def upload_data( + self, path, bucket: Incomplete | None = None, key_prefix: str = "data", extra_args: Incomplete | None = None + ): ... + def upload_string_as_file_body(self, body, bucket, key, kms_key: Incomplete | None = None): ... + def download_data(self, path, bucket, key_prefix: str = "", extra_args: Incomplete | None = None): ... + def read_s3_file(self, bucket, key_prefix): ... + def list_s3_files(self, bucket, key_prefix): ... + def default_bucket(self): ... + def train( + self, + input_mode, + input_config, + role: Incomplete | None = None, + job_name: Incomplete | None = None, + output_config: Incomplete | None = None, + resource_config: Incomplete | None = None, + vpc_config: Incomplete | None = None, + hyperparameters: Incomplete | None = None, + stop_condition: Incomplete | None = None, + tags: Incomplete | None = None, + metric_definitions: Incomplete | None = None, + enable_network_isolation: Incomplete | None = None, + image_uri: Incomplete | None = None, + training_image_config: Incomplete | None = None, + container_entry_point: Incomplete | None = None, + container_arguments: Incomplete | None = None, + algorithm_arn: Incomplete | None = None, + encrypt_inter_container_traffic: Incomplete | None = None, + use_spot_instances: bool = False, + checkpoint_s3_uri: Incomplete | None = None, + checkpoint_local_path: Incomplete | None = None, + experiment_config: Incomplete | None = None, + debugger_rule_configs: Incomplete | None = None, + debugger_hook_config: Incomplete | None = None, + tensorboard_output_config: Incomplete | None = None, + enable_sagemaker_metrics: Incomplete | None = None, + profiler_rule_configs: Incomplete | None = None, + profiler_config: Incomplete | None = None, + environment: Optional[Dict[str, str]] = None, + retry_strategy: Incomplete | None = None, + ): ... + def update_training_job( + self, + job_name, + profiler_rule_configs: Incomplete | None = None, + profiler_config: Incomplete | None = None, + resource_config: Incomplete | None = None, + ) -> None: ... + def process( + self, + inputs, + output_config, + job_name, + resources, + stopping_condition, + app_specification, + environment: Optional[Dict[str, str]] = None, + network_config: Incomplete | None = None, + role_arn: Incomplete | None = None, + tags: Incomplete | None = None, + experiment_config: Incomplete | None = None, + ): ... + def create_monitoring_schedule( + self, + monitoring_schedule_name, + schedule_expression, + statistics_s3_uri, + constraints_s3_uri, + monitoring_inputs, + monitoring_output_config, + instance_count, + instance_type, + volume_size_in_gb, + volume_kms_key: Incomplete | None = None, + image_uri: Incomplete | None = None, + entrypoint: Incomplete | None = None, + arguments: Incomplete | None = None, + record_preprocessor_source_uri: Incomplete | None = None, + post_analytics_processor_source_uri: Incomplete | None = None, + max_runtime_in_seconds: Incomplete | None = None, + environment: Incomplete | None = None, + network_config: Incomplete | None = None, + role_arn: Incomplete | None = None, + tags: Incomplete | None = None, + ) -> None: ... + def update_monitoring_schedule( + self, + monitoring_schedule_name, + schedule_expression: Incomplete | None = None, + statistics_s3_uri: Incomplete | None = None, + constraints_s3_uri: Incomplete | None = None, + monitoring_inputs: Incomplete | None = None, + monitoring_output_config: Incomplete | None = None, + instance_count: Incomplete | None = None, + instance_type: Incomplete | None = None, + volume_size_in_gb: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + image_uri: Incomplete | None = None, + entrypoint: Incomplete | None = None, + arguments: Incomplete | None = None, + record_preprocessor_source_uri: Incomplete | None = None, + post_analytics_processor_source_uri: Incomplete | None = None, + max_runtime_in_seconds: Incomplete | None = None, + environment: Incomplete | None = None, + network_config: Incomplete | None = None, + role_arn: Incomplete | None = None, + ) -> None: ... + def start_monitoring_schedule(self, monitoring_schedule_name) -> None: ... + def stop_monitoring_schedule(self, monitoring_schedule_name) -> None: ... + def delete_monitoring_schedule(self, monitoring_schedule_name) -> None: ... + def describe_monitoring_schedule(self, monitoring_schedule_name): ... + def list_monitoring_executions( + self, monitoring_schedule_name, sort_by: str = "ScheduledTime", sort_order: str = "Descending", max_results: int = 100 + ): ... + def list_monitoring_schedules( + self, + endpoint_name: Incomplete | None = None, + sort_by: str = "CreationTime", + sort_order: str = "Descending", + max_results: int = 100, + ): ... + def update_monitoring_alert( + self, monitoring_schedule_name: str, monitoring_alert_name: str, data_points_to_alert: int, evaluation_period: int + ): ... + def list_monitoring_alerts( + self, monitoring_schedule_name: str, next_token: Optional[str] = None, max_results: Optional[int] = 10 + ) -> Dict: ... + def list_monitoring_alert_history( + self, + monitoring_schedule_name: Optional[str] = None, + monitoring_alert_name: Optional[str] = None, + sort_by: Optional[str] = "CreationTime", + sort_order: Optional[str] = "Descending", + next_token: Optional[str] = None, + max_results: Optional[int] = 10, + creation_time_before: Optional[str] = None, + creation_time_after: Optional[str] = None, + status_equals: Optional[str] = None, + ) -> Dict: ... + def was_processing_job_successful(self, job_name): ... + def describe_processing_job(self, job_name): ... + def stop_processing_job(self, job_name) -> None: ... + def stop_training_job(self, job_name) -> None: ... + def describe_training_job(self, job_name): ... + def auto_ml( + self, + input_config, + output_config, + auto_ml_job_config, + role: Incomplete | None = None, + job_name: Incomplete | None = None, + problem_type: Incomplete | None = None, + job_objective: Incomplete | None = None, + generate_candidate_definitions_only: bool = False, + tags: Incomplete | None = None, + model_deploy_config: Incomplete | None = None, + ) -> None: ... + def describe_auto_ml_job(self, job_name): ... + def list_candidates( + self, + job_name, + status_equals: Incomplete | None = None, + candidate_name: Incomplete | None = None, + candidate_arn: Incomplete | None = None, + sort_order: Incomplete | None = None, + sort_by: Incomplete | None = None, + max_results: Incomplete | None = None, + ): ... + def wait_for_auto_ml_job(self, job, poll: int = 5): ... + def logs_for_auto_ml_job(self, job_name, wait: bool = False, poll: int = 10): ... + def compile_model( + self, + input_model_config, + output_model_config, + role: Incomplete | None = None, + job_name: Incomplete | None = None, + stop_condition: Incomplete | None = None, + tags: Incomplete | None = None, + ) -> None: ... + def package_model_for_edge( + self, + output_model_config, + role: Incomplete | None = None, + job_name: Incomplete | None = None, + compilation_job_name: Incomplete | None = None, + model_name: Incomplete | None = None, + model_version: Incomplete | None = None, + resource_key: Incomplete | None = None, + tags: Incomplete | None = None, + ) -> None: ... + def tune( + self, + job_name, + strategy, + objective_type, + objective_metric_name, + max_jobs, + max_parallel_jobs, + parameter_ranges, + static_hyperparameters, + input_mode, + metric_definitions, + role, + input_config, + output_config, + resource_config, + stop_condition, + tags, + warm_start_config, + max_runtime_in_seconds: Incomplete | None = None, + strategy_config: Incomplete | None = None, + completion_criteria_config: Incomplete | None = None, + enable_network_isolation: bool = False, + image_uri: Incomplete | None = None, + algorithm_arn: Incomplete | None = None, + early_stopping_type: str = "Off", + encrypt_inter_container_traffic: bool = False, + vpc_config: Incomplete | None = None, + use_spot_instances: bool = False, + checkpoint_s3_uri: Incomplete | None = None, + checkpoint_local_path: Incomplete | None = None, + random_seed: Incomplete | None = None, + environment: Incomplete | None = None, + hpo_resource_config: Incomplete | None = None, + autotune: bool = False, + auto_parameters: Incomplete | None = None, + ) -> None: ... + def create_tuning_job( + self, + job_name, + tuning_config, + training_config: Incomplete | None = None, + training_config_list: Incomplete | None = None, + warm_start_config: Incomplete | None = None, + tags: Incomplete | None = None, + autotune: bool = False, + ) -> None: ... + def describe_tuning_job(self, job_name): ... + def stop_tuning_job(self, name) -> None: ... + def transform( + self, + job_name, + model_name, + strategy, + max_concurrent_transforms, + max_payload, + input_config, + output_config, + resource_config, + experiment_config, + env: Optional[Dict[str, str]] = None, + tags: Incomplete | None = None, + data_processing: Incomplete | None = None, + model_client_config: Incomplete | None = None, + batch_data_capture_config: BatchDataCaptureConfig = None, + ): ... + def create_model( + self, + name, + role: Incomplete | None = None, + container_defs: Incomplete | None = None, + vpc_config: Incomplete | None = None, + enable_network_isolation: Incomplete | None = None, + primary_container: Incomplete | None = None, + tags: Incomplete | None = None, + ): ... + def create_model_from_job( + self, + training_job_name, + name: Incomplete | None = None, + role: Incomplete | None = None, + image_uri: Incomplete | None = None, + model_data_url: Incomplete | None = None, + env: Incomplete | None = None, + enable_network_isolation: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + tags: Incomplete | None = None, + ): ... + def create_model_package_from_algorithm(self, name, description, algorithm_arn, model_data) -> None: ... + def create_model_package_from_containers( + self, + containers: Incomplete | None = None, + content_types: Incomplete | None = None, + response_types: Incomplete | None = None, + inference_instances: Incomplete | None = None, + transform_instances: Incomplete | None = None, + model_package_name: Incomplete | None = None, + model_package_group_name: Incomplete | None = None, + model_metrics: Incomplete | None = None, + metadata_properties: Incomplete | None = None, + marketplace_cert: bool = False, + approval_status: str = "PendingManualApproval", + description: Incomplete | None = None, + drift_check_baselines: Incomplete | None = None, + customer_metadata_properties: Incomplete | None = None, + validation_specification: Incomplete | None = None, + domain: Incomplete | None = None, + sample_payload_url: Incomplete | None = None, + task: Incomplete | None = None, + ): ... + def wait_for_model_package(self, model_package_name, poll: int = 5): ... + def describe_model(self, name): ... + def create_endpoint_config( + self, + name, + model_name, + initial_instance_count, + instance_type, + accelerator_type: Incomplete | None = None, + tags: Incomplete | None = None, + kms_key: Incomplete | None = None, + data_capture_config_dict: Incomplete | None = None, + volume_size: Incomplete | None = None, + model_data_download_timeout: Incomplete | None = None, + container_startup_health_check_timeout: Incomplete | None = None, + explainer_config_dict: Incomplete | None = None, + ): ... + def create_endpoint_config_from_existing( + self, + existing_config_name, + new_config_name, + new_tags: Incomplete | None = None, + new_kms_key: Incomplete | None = None, + new_data_capture_config_dict: Incomplete | None = None, + new_production_variants: Incomplete | None = None, + new_explainer_config_dict: Incomplete | None = None, + ) -> None: ... + def create_endpoint(self, endpoint_name, config_name, tags: Incomplete | None = None, wait: bool = True): ... + def update_endpoint(self, endpoint_name, endpoint_config_name, wait: bool = True): ... + def delete_endpoint(self, endpoint_name) -> None: ... + def delete_endpoint_config(self, endpoint_config_name) -> None: ... + def delete_model(self, model_name) -> None: ... + def list_group_resources(self, group, filters, next_token: str = ""): ... + def delete_resource_group(self, group): ... + def get_resource_group_query(self, group): ... + def get_tagging_resources(self, tag_filters, resource_type_filters): ... + def create_group(self, name, resource_query, tags): ... + def list_tags(self, resource_arn, max_results: int = 50): ... + def wait_for_job(self, job, poll: int = 5): ... + def wait_for_processing_job(self, job, poll: int = 5): ... + def wait_for_compilation_job(self, job, poll: int = 5): ... + def wait_for_edge_packaging_job(self, job, poll: int = 5): ... + def wait_for_tuning_job(self, job, poll: int = 5): ... + def describe_transform_job(self, job_name): ... + def wait_for_transform_job(self, job, poll: int = 5): ... + def stop_transform_job(self, name) -> None: ... + def wait_for_endpoint(self, endpoint, poll: int = 30): ... + def endpoint_from_job( + self, + job_name, + initial_instance_count, + instance_type, + image_uri: Incomplete | None = None, + name: Incomplete | None = None, + role: Incomplete | None = None, + wait: bool = True, + model_environment_vars: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + accelerator_type: Incomplete | None = None, + data_capture_config: Incomplete | None = None, + ): ... + def endpoint_from_model_data( + self, + model_s3_location, + image_uri, + initial_instance_count, + instance_type, + name: Incomplete | None = None, + role: Incomplete | None = None, + wait: bool = True, + model_environment_vars: Incomplete | None = None, + model_vpc_config: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + data_capture_config: Incomplete | None = None, + tags: Incomplete | None = None, + ): ... + def endpoint_from_production_variants( + self, + name, + production_variants, + tags: Incomplete | None = None, + kms_key: Incomplete | None = None, + wait: bool = True, + data_capture_config_dict: Incomplete | None = None, + async_inference_config_dict: Incomplete | None = None, + explainer_config_dict: Incomplete | None = None, + ): ... + def expand_role(self, role): ... + def get_caller_identity_arn(self): ... + def logs_for_job( + self, job_name, wait: bool = False, poll: int = 10, log_type: str = "All", timeout: Incomplete | None = None + ) -> None: ... + def logs_for_processing_job(self, job_name, wait: bool = False, poll: int = 10): ... + def logs_for_transform_job(self, job_name, wait: bool = False, poll: int = 10): ... + def delete_feature_group(self, feature_group_name: str): ... + def create_feature_group( + self, + feature_group_name: str, + record_identifier_name: str, + event_time_feature_name: str, + feature_definitions: Sequence[Dict[str, str]], + role_arn: str = None, + online_store_config: Dict[str, str] = None, + offline_store_config: Dict[str, str] = None, + description: str = None, + tags: List[Dict[str, str]] = None, + ) -> Dict[str, Any]: ... + def describe_feature_group(self, feature_group_name: str, next_token: str = None) -> Dict[str, Any]: ... + def update_feature_group(self, feature_group_name: str, feature_additions: Sequence[Dict[str, str]]) -> Dict[str, Any]: ... + def list_feature_groups( + self, + name_contains, + feature_group_status_equals, + offline_store_status_equals, + creation_time_after, + creation_time_before, + sort_order, + sort_by, + max_results, + next_token, + ) -> Dict[str, Any]: ... + def update_feature_metadata( + self, + feature_group_name: str, + feature_name: str, + description: str = None, + parameter_additions: Sequence[Dict[str, str]] = None, + parameter_removals: Sequence[str] = None, + ) -> Dict[str, Any]: ... + def describe_feature_metadata(self, feature_group_name: str, feature_name: str) -> Dict[str, Any]: ... + def search( + self, + resource: str, + search_expression: Dict[str, any] = None, + sort_by: str = None, + sort_order: str = None, + next_token: str = None, + max_results: int = None, + ) -> Dict[str, Any]: ... + def put_record(self, feature_group_name: str, record: Sequence[Dict[str, str]]): ... + def delete_record( + self, feature_group_name: str, record_identifier_value_as_string: str, event_time: str, deletion_mode: str = None + ): ... + def get_record( + self, record_identifier_value_as_string: str, feature_group_name: str, feature_names: Sequence[str] + ) -> Dict[str, Sequence[Dict[str, str]]]: ... + def batch_get_record(self, identifiers: Sequence[Dict[str, Any]]) -> Dict[str, Any]: ... + def start_query_execution( + self, catalog: str, database: str, query_string: str, output_location: str, kms_key: str = None, workgroup: str = None + ) -> Dict[str, str]: ... + def get_query_execution(self, query_execution_id: str) -> Dict[str, Any]: ... + def wait_for_athena_query(self, query_execution_id: str, poll: int = 5): ... + def download_athena_query_result(self, bucket: str, prefix: str, query_execution_id: str, filename: str): ... + def account_id(self) -> str: ... + def create_inference_recommendations_job( + self, + role: str, + sample_payload_url: str, + supported_content_types: List[str], + job_name: str = None, + job_type: str = "Default", + model_name: str = None, + model_package_version_arn: str = None, + job_duration_in_seconds: int = None, + nearest_model_name: str = None, + supported_instance_types: List[str] = None, + framework: str = None, + framework_version: str = None, + endpoint_configurations: List[Dict[str, any]] = None, + traffic_pattern: Dict[str, any] = None, + stopping_conditions: Dict[str, any] = None, + resource_limit: Dict[str, any] = None, + ): ... + def wait_for_inference_recommendations_job( + self, job_name: str, poll: int = 120, log_level: str = "Verbose" + ) -> Dict[str, Any]: ... + +def get_model_package_args( + content_types, + response_types, + inference_instances: Incomplete | None = None, + transform_instances: Incomplete | None = None, + model_package_name: Incomplete | None = None, + model_package_group_name: Incomplete | None = None, + model_data: Incomplete | None = None, + image_uri: Incomplete | None = None, + model_metrics: Incomplete | None = None, + metadata_properties: Incomplete | None = None, + marketplace_cert: bool = False, + approval_status: Incomplete | None = None, + description: Incomplete | None = None, + tags: Incomplete | None = None, + container_def_list: Incomplete | None = None, + drift_check_baselines: Incomplete | None = None, + customer_metadata_properties: Incomplete | None = None, + validation_specification: Incomplete | None = None, + domain: Incomplete | None = None, + sample_payload_url: Incomplete | None = None, + task: Incomplete | None = None, +): ... +def get_create_model_package_request( + model_package_name: Incomplete | None = None, + model_package_group_name: Incomplete | None = None, + containers: Incomplete | None = None, + content_types: Incomplete | None = None, + response_types: Incomplete | None = None, + inference_instances: Incomplete | None = None, + transform_instances: Incomplete | None = None, + model_metrics: Incomplete | None = None, + metadata_properties: Incomplete | None = None, + marketplace_cert: bool = False, + approval_status: str = "PendingManualApproval", + description: Incomplete | None = None, + tags: Incomplete | None = None, + drift_check_baselines: Incomplete | None = None, + customer_metadata_properties: Incomplete | None = None, + validation_specification: Incomplete | None = None, + domain: Incomplete | None = None, + sample_payload_url: Incomplete | None = None, + task: Incomplete | None = None, +): ... +def update_args(args: Dict[str, Any], **kwargs): ... +def container_def( + image_uri, + model_data_url: Incomplete | None = None, + env: Incomplete | None = None, + container_mode: Incomplete | None = None, + image_config: Incomplete | None = None, +): ... +def pipeline_container_def(models, instance_type: Incomplete | None = None): ... +def production_variant( + model_name, + instance_type: Incomplete | None = None, + initial_instance_count: Incomplete | None = None, + variant_name: str = "AllTraffic", + initial_weight: int = 1, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + volume_size: Incomplete | None = None, + model_data_download_timeout: Incomplete | None = None, + container_startup_health_check_timeout: Incomplete | None = None, +): ... +def get_execution_role(sagemaker_session: Incomplete | None = None): ... +def generate_default_sagemaker_bucket_name(boto_session): ... +def get_log_events_for_inference_recommender(cw_client, log_group_name, log_stream_name): ... + +s3_input: Incomplete diff --git a/stubs/sagemaker/sagemaker/session_settings.pyi b/stubs/sagemaker/sagemaker/session_settings.pyi new file mode 100644 index 000000000000..c37b26853aab --- /dev/null +++ b/stubs/sagemaker/sagemaker/session_settings.pyi @@ -0,0 +1,8 @@ +from _typeshed import Incomplete + +class SessionSettings: + def __init__(self, encrypt_repacked_artifacts: bool = True, local_download_dir: Incomplete | None = None) -> None: ... + @property + def encrypt_repacked_artifacts(self) -> bool: ... + @property + def local_download_dir(self) -> str: ... diff --git a/stubs/sagemaker/sagemaker/sklearn/__init__.pyi b/stubs/sagemaker/sagemaker/sklearn/__init__.pyi new file mode 100644 index 000000000000..0eeb8bb193db --- /dev/null +++ b/stubs/sagemaker/sagemaker/sklearn/__init__.pyi @@ -0,0 +1,3 @@ +from sagemaker.sklearn.estimator import SKLearn as SKLearn +from sagemaker.sklearn.model import SKLearnModel as SKLearnModel, SKLearnPredictor as SKLearnPredictor +from sagemaker.sklearn.processing import SKLearnProcessor as SKLearnProcessor diff --git a/stubs/sagemaker/sagemaker/sklearn/defaults.pyi b/stubs/sagemaker/sagemaker/sklearn/defaults.pyi new file mode 100644 index 000000000000..ff170251ca2a --- /dev/null +++ b/stubs/sagemaker/sagemaker/sklearn/defaults.pyi @@ -0,0 +1 @@ +SKLEARN_NAME: str diff --git a/stubs/sagemaker/sagemaker/sklearn/estimator.pyi b/stubs/sagemaker/sagemaker/sklearn/estimator.pyi new file mode 100644 index 000000000000..3f37e15e0e74 --- /dev/null +++ b/stubs/sagemaker/sagemaker/sklearn/estimator.pyi @@ -0,0 +1,33 @@ +from _typeshed import Incomplete +from typing import Dict, Optional, Union + +from sagemaker.estimator import Framework +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class SKLearn(Framework): + framework_version: Incomplete + py_version: Incomplete + image_uri: Incomplete + def __init__( + self, + entry_point: Union[str, PipelineVariable], + framework_version: Optional[str] = None, + py_version: str = "py3", + source_dir: Optional[Union[str, PipelineVariable]] = None, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri_region: Optional[str] = None, + **kwargs, + ) -> None: ... + def create_model( + self, + model_server_workers: Incomplete | None = None, + role: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + entry_point: Incomplete | None = None, + source_dir: Incomplete | None = None, + dependencies: Incomplete | None = None, + **kwargs, + ): ... diff --git a/stubs/sagemaker/sagemaker/sklearn/model.pyi b/stubs/sagemaker/sagemaker/sklearn/model.pyi new file mode 100644 index 000000000000..8726615f7a8e --- /dev/null +++ b/stubs/sagemaker/sagemaker/sklearn/model.pyi @@ -0,0 +1,63 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker import ModelMetrics +from sagemaker.drift_check_baselines import DriftCheckBaselines +from sagemaker.metadata_properties import MetadataProperties +from sagemaker.model import FrameworkModel +from sagemaker.predictor import Predictor +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class SKLearnPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class SKLearnModel(FrameworkModel): + framework_version: Incomplete + py_version: Incomplete + model_server_workers: Incomplete + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + entry_point: Optional[str] = None, + framework_version: Optional[str] = None, + py_version: str = "py3", + image_uri: Optional[Union[str, PipelineVariable]] = None, + predictor_cls: callable = ..., + model_server_workers: Optional[Union[int, PipelineVariable]] = None, + **kwargs, + ) -> None: ... + image_uri: Incomplete + def register( + self, + 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, + domain: Optional[Union[str, PipelineVariable]] = None, + sample_payload_url: Optional[Union[str, PipelineVariable]] = None, + task: Optional[Union[str, PipelineVariable]] = None, + framework: Optional[Union[str, PipelineVariable]] = None, + framework_version: Optional[Union[str, PipelineVariable]] = None, + nearest_model_name: Optional[Union[str, PipelineVariable]] = None, + data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + ): ... + def prepare_container_def( + self, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... + def serving_image_uri(self, region_name, instance_type, serverless_inference_config: Incomplete | None = None): ... diff --git a/stubs/sagemaker/sagemaker/sklearn/processing.pyi b/stubs/sagemaker/sagemaker/sklearn/processing.pyi new file mode 100644 index 000000000000..d8e3eeccfd72 --- /dev/null +++ b/stubs/sagemaker/sagemaker/sklearn/processing.pyi @@ -0,0 +1,25 @@ +from typing import Dict, List, Optional, Union + +from sagemaker import Session +from sagemaker.network import NetworkConfig +from sagemaker.processing import ScriptProcessor +from sagemaker.workflow.entities import PipelineVariable + +class SKLearnProcessor(ScriptProcessor): + def __init__( + self, + framework_version: str, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Union[int, PipelineVariable] = None, + instance_type: Union[str, PipelineVariable] = None, + command: Optional[List[str]] = None, + volume_size_in_gb: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + output_kms_key: Optional[Union[str, PipelineVariable]] = None, + max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + base_job_name: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + network_config: Optional[NetworkConfig] = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/spark/__init__.pyi b/stubs/sagemaker/sagemaker/spark/__init__.pyi new file mode 100644 index 000000000000..3d2c4fac6e8c --- /dev/null +++ b/stubs/sagemaker/sagemaker/spark/__init__.pyi @@ -0,0 +1 @@ +from sagemaker.spark.processing import PySparkProcessor as PySparkProcessor, SparkJarProcessor as SparkJarProcessor diff --git a/stubs/sagemaker/sagemaker/spark/defaults.pyi b/stubs/sagemaker/sagemaker/spark/defaults.pyi new file mode 100644 index 000000000000..642f4dea2c92 --- /dev/null +++ b/stubs/sagemaker/sagemaker/spark/defaults.pyi @@ -0,0 +1 @@ +SPARK_NAME: str diff --git a/stubs/sagemaker/sagemaker/spark/processing.pyi b/stubs/sagemaker/sagemaker/spark/processing.pyi new file mode 100644 index 000000000000..6e32888a55d1 --- /dev/null +++ b/stubs/sagemaker/sagemaker/spark/processing.pyi @@ -0,0 +1,182 @@ +from _typeshed import Incomplete +from enum import Enum +from typing import Dict, List, Optional, Union + +from sagemaker.network import NetworkConfig +from sagemaker.processing import ProcessingInput, ProcessingOutput, ScriptProcessor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class _SparkProcessorBase(ScriptProcessor): + configuration_location: Incomplete + dependency_location: Incomplete + history_server: Incomplete + image_uri: Incomplete + def __init__( + self, + role: Incomplete | None = None, + instance_type: Incomplete | None = None, + instance_count: Incomplete | None = None, + framework_version: Incomplete | None = None, + py_version: Incomplete | None = None, + container_version: Incomplete | None = None, + image_uri: Incomplete | None = None, + volume_size_in_gb: int = 30, + volume_kms_key: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + configuration_location: Optional[str] = None, + dependency_location: Optional[str] = None, + max_runtime_in_seconds: Incomplete | None = None, + base_job_name: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + env: Incomplete | None = None, + tags: Incomplete | None = None, + network_config: Incomplete | None = None, + ) -> None: ... + def get_run_args( + self, code, inputs: Incomplete | None = None, outputs: Incomplete | None = None, arguments: Incomplete | None = None + ): ... + def run( + self, + submit_app, + inputs: Incomplete | None = None, + outputs: Incomplete | None = None, + arguments: Incomplete | None = None, + wait: bool = True, + logs: bool = True, + job_name: Incomplete | None = None, + experiment_config: Incomplete | None = None, + kms_key: Incomplete | None = None, + ): ... + def start_history_server(self, spark_event_logs_s3_uri: Incomplete | None = None) -> None: ... + def terminate_history_server(self) -> None: ... + +class PySparkProcessor(_SparkProcessorBase): + def __init__( + self, + role: str = None, + instance_type: Union[str, PipelineVariable] = None, + instance_count: Union[int, PipelineVariable] = None, + framework_version: Optional[str] = None, + py_version: Optional[str] = None, + container_version: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + volume_size_in_gb: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + output_kms_key: Optional[Union[str, PipelineVariable]] = None, + configuration_location: Optional[str] = None, + dependency_location: Optional[str] = None, + max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + base_job_name: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + network_config: Optional[NetworkConfig] = None, + ) -> None: ... + def get_run_args( + self, + submit_app, + submit_py_files: Incomplete | None = None, + submit_jars: Incomplete | None = None, + submit_files: Incomplete | None = None, + inputs: Incomplete | None = None, + outputs: Incomplete | None = None, + arguments: Incomplete | None = None, + job_name: Incomplete | None = None, + configuration: Incomplete | None = None, + spark_event_logs_s3_uri: Incomplete | None = None, + ): ... + def run( + self, + submit_app: str, + submit_py_files: Optional[List[Union[str, PipelineVariable]]] = None, + submit_jars: Optional[List[Union[str, PipelineVariable]]] = None, + submit_files: Optional[List[Union[str, PipelineVariable]]] = None, + inputs: Optional[List[ProcessingInput]] = None, + outputs: Optional[List[ProcessingOutput]] = None, + arguments: Optional[List[Union[str, PipelineVariable]]] = None, + wait: bool = True, + logs: bool = True, + job_name: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + configuration: Optional[Union[List[Dict], Dict]] = None, + spark_event_logs_s3_uri: Optional[Union[str, PipelineVariable]] = None, + kms_key: Optional[str] = None, + ): ... + +class SparkJarProcessor(_SparkProcessorBase): + def __init__( + self, + role: str = None, + instance_type: Union[str, PipelineVariable] = None, + instance_count: Union[int, PipelineVariable] = None, + framework_version: Optional[str] = None, + py_version: Optional[str] = None, + container_version: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + volume_size_in_gb: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + output_kms_key: Optional[Union[str, PipelineVariable]] = None, + configuration_location: Optional[str] = None, + dependency_location: Optional[str] = None, + max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + base_job_name: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + network_config: Optional[NetworkConfig] = None, + ) -> None: ... + def get_run_args( + self, + submit_app, + submit_class: Incomplete | None = None, + submit_jars: Incomplete | None = None, + submit_files: Incomplete | None = None, + inputs: Incomplete | None = None, + outputs: Incomplete | None = None, + arguments: Incomplete | None = None, + job_name: Incomplete | None = None, + configuration: Incomplete | None = None, + spark_event_logs_s3_uri: Incomplete | None = None, + ): ... + def run( + self, + submit_app: str, + submit_class: Union[str, PipelineVariable], + submit_jars: Optional[List[Union[str, PipelineVariable]]] = None, + submit_files: Optional[List[Union[str, PipelineVariable]]] = None, + inputs: Optional[List[ProcessingInput]] = None, + outputs: Optional[List[ProcessingOutput]] = None, + arguments: Optional[List[Union[str, PipelineVariable]]] = None, + wait: bool = True, + logs: bool = True, + job_name: Optional[str] = None, + experiment_config: Optional[Dict[str, str]] = None, + configuration: Optional[Union[List[Dict], Dict]] = None, + spark_event_logs_s3_uri: Optional[Union[str, PipelineVariable]] = None, + kms_key: Optional[str] = None, + ): ... + +class _HistoryServer: + arg_event_logs_s3_uri: str + arg_remote_domain_name: str + cli_args: Incomplete + image_uri: Incomplete + network_config: Incomplete + run_history_server_command: Incomplete + def __init__(self, cli_args, image_uri, network_config) -> None: ... + def run(self) -> None: ... + def down(self) -> None: ... + +class FileType(Enum): + JAR: int + PYTHON: int + FILE: int + +class SparkConfigUtils: + @staticmethod + def validate_configuration(configuration: Dict): ... + @staticmethod + def validate_s3_uri(spark_output_s3_path) -> None: ... diff --git a/stubs/sagemaker/sagemaker/sparkml/__init__.pyi b/stubs/sagemaker/sagemaker/sparkml/__init__.pyi new file mode 100644 index 000000000000..b17a95e1452f --- /dev/null +++ b/stubs/sagemaker/sagemaker/sparkml/__init__.pyi @@ -0,0 +1 @@ +from sagemaker.sparkml.model import SparkMLModel as SparkMLModel, SparkMLPredictor as SparkMLPredictor diff --git a/stubs/sagemaker/sagemaker/sparkml/model.pyi b/stubs/sagemaker/sagemaker/sparkml/model.pyi new file mode 100644 index 000000000000..3e46af3941db --- /dev/null +++ b/stubs/sagemaker/sagemaker/sparkml/model.pyi @@ -0,0 +1,20 @@ +from _typeshed import Incomplete +from typing import Optional, Union + +from sagemaker import Model, Predictor, Session +from sagemaker.workflow.entities import PipelineVariable + +framework_name: str + +class SparkMLPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., **kwargs) -> None: ... + +class SparkMLModel(Model): + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: Optional[str] = None, + spark_version: str = "3.3", + sagemaker_session: Optional[Session] = None, + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/tensorflow/__init__.pyi b/stubs/sagemaker/sagemaker/tensorflow/__init__.pyi new file mode 100644 index 000000000000..52c8ef6a448a --- /dev/null +++ b/stubs/sagemaker/sagemaker/tensorflow/__init__.pyi @@ -0,0 +1,4 @@ +from sagemaker.tensorflow.estimator import TensorFlow as TensorFlow +from sagemaker.tensorflow.model import TensorFlowModel as TensorFlowModel, TensorFlowPredictor as TensorFlowPredictor +from sagemaker.tensorflow.processing import TensorFlowProcessor as TensorFlowProcessor +from sagemaker.tensorflow.training_compiler.config import TrainingCompilerConfig as TrainingCompilerConfig diff --git a/stubs/sagemaker/sagemaker/tensorflow/defaults.pyi b/stubs/sagemaker/sagemaker/tensorflow/defaults.pyi new file mode 100644 index 000000000000..0555246a1ab9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/tensorflow/defaults.pyi @@ -0,0 +1 @@ +LATEST_PY2_VERSION: str diff --git a/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi b/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi new file mode 100644 index 000000000000..c9cf390cf2ae --- /dev/null +++ b/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi @@ -0,0 +1,56 @@ +from _typeshed import Incomplete +from typing import Dict, Optional, Union + +from sagemaker.estimator import Framework +from sagemaker.tensorflow.training_compiler.config import TrainingCompilerConfig +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class TensorFlow(Framework): + framework_version: Incomplete + py_version: Incomplete + instance_type: Incomplete + model_dir: Incomplete + distribution: Incomplete + compiler_config: Incomplete + def __init__( + self, + py_version: Optional[str] = None, + framework_version: Optional[str] = None, + model_dir: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + distribution: Optional[Dict[str, str]] = None, + compiler_config: Optional[TrainingCompilerConfig] = None, + **kwargs, + ) -> None: ... + def create_model( + self, + role: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + entry_point: Incomplete | None = None, + source_dir: Incomplete | None = None, + dependencies: Incomplete | None = None, + **kwargs, + ): ... + def hyperparameters(self): ... + def transformer( + self, + instance_count, + instance_type, + strategy: Incomplete | None = None, + assemble_with: Incomplete | None = None, + output_path: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + accept: Incomplete | None = None, + env: Incomplete | None = None, + max_concurrent_transforms: Incomplete | None = None, + max_payload: Incomplete | None = None, + tags: Incomplete | None = None, + role: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + entry_point: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + enable_network_isolation: Incomplete | None = None, + model_name: Incomplete | None = None, + ): ... diff --git a/stubs/sagemaker/sagemaker/tensorflow/model.pyi b/stubs/sagemaker/sagemaker/tensorflow/model.pyi new file mode 100644 index 000000000000..a4c94427b71c --- /dev/null +++ b/stubs/sagemaker/sagemaker/tensorflow/model.pyi @@ -0,0 +1,103 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +import sagemaker +from sagemaker import ModelMetrics +from sagemaker.drift_check_baselines import DriftCheckBaselines +from sagemaker.metadata_properties import MetadataProperties +from sagemaker.predictor import Predictor +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class TensorFlowPredictor(Predictor): + def __init__( + self, + endpoint_name, + sagemaker_session: Incomplete | None = None, + serializer=..., + deserializer=..., + model_name: Incomplete | None = None, + model_version: Incomplete | None = None, + **kwargs, + ) -> None: ... + def classify(self, data): ... + def regress(self, data): ... + def predict(self, data, initial_args: Incomplete | None = None): ... + +class TensorFlowModel(sagemaker.model.FrameworkModel): + LOG_LEVEL_PARAM_NAME: str + LOG_LEVEL_MAP: Incomplete + LATEST_EIA_VERSION: Incomplete + framework_version: Incomplete + inference_framework_version: Incomplete + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: str = None, + entry_point: Optional[str] = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + framework_version: Optional[str] = None, + container_log_level: Optional[int] = None, + predictor_cls: callable = ..., + **kwargs, + ) -> None: ... + image_uri: Incomplete + def register( + self, + 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, + domain: Optional[Union[str, PipelineVariable]] = None, + sample_payload_url: Optional[Union[str, PipelineVariable]] = None, + task: Optional[Union[str, PipelineVariable]] = None, + framework: Optional[Union[str, PipelineVariable]] = None, + framework_version: Optional[Union[str, PipelineVariable]] = None, + nearest_model_name: Optional[Union[str, PipelineVariable]] = None, + data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + ): ... + def deploy( + self, + initial_instance_count: Incomplete | None = None, + instance_type: Incomplete | None = None, + serializer: Incomplete | None = None, + deserializer: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + endpoint_name: Incomplete | None = None, + tags: Incomplete | None = None, + kms_key: Incomplete | None = None, + wait: bool = True, + data_capture_config: Incomplete | None = None, + async_inference_config: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + volume_size: Incomplete | None = None, + model_data_download_timeout: Incomplete | None = None, + container_startup_health_check_timeout: Incomplete | None = None, + inference_recommendation_id: Incomplete | None = None, + explainer_config: Incomplete | None = None, + **kwargs, + ): ... + def prepare_container_def( + self, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... + def serving_image_uri( + self, + region_name, + instance_type, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... diff --git a/stubs/sagemaker/sagemaker/tensorflow/processing.pyi b/stubs/sagemaker/sagemaker/tensorflow/processing.pyi new file mode 100644 index 000000000000..3bacec0a4740 --- /dev/null +++ b/stubs/sagemaker/sagemaker/tensorflow/processing.pyi @@ -0,0 +1,30 @@ +from typing import Dict, List, Optional, Union + +from sagemaker.network import NetworkConfig +from sagemaker.processing import FrameworkProcessor +from sagemaker.session import Session +from sagemaker.tensorflow.estimator import TensorFlow +from sagemaker.workflow.entities import PipelineVariable + +class TensorFlowProcessor(FrameworkProcessor): + estimator_cls = TensorFlow + def __init__( + self, + framework_version: str, + role: Optional[Union[str, PipelineVariable]] = None, + instance_count: Union[int, PipelineVariable] = None, + instance_type: Union[str, PipelineVariable] = None, + py_version: str = "py3", + image_uri: Optional[Union[str, PipelineVariable]] = None, + command: Optional[List[str]] = None, + volume_size_in_gb: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + output_kms_key: Optional[Union[str, PipelineVariable]] = None, + code_location: Optional[str] = None, + max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + base_job_name: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + network_config: Optional[NetworkConfig] = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/tensorflow/serving.pyi b/stubs/sagemaker/sagemaker/tensorflow/serving.pyi new file mode 100644 index 000000000000..3f9d8f1eba99 --- /dev/null +++ b/stubs/sagemaker/sagemaker/tensorflow/serving.pyi @@ -0,0 +1,4 @@ +from _typeshed import Incomplete + +Model: Incomplete +Predictor: Incomplete diff --git a/stubs/sagemaker/sagemaker/tensorflow/training_compiler/__init__.pyi b/stubs/sagemaker/sagemaker/tensorflow/training_compiler/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/tensorflow/training_compiler/config.pyi b/stubs/sagemaker/sagemaker/tensorflow/training_compiler/config.pyi new file mode 100644 index 000000000000..163046e33288 --- /dev/null +++ b/stubs/sagemaker/sagemaker/tensorflow/training_compiler/config.pyi @@ -0,0 +1,13 @@ +from _typeshed import Incomplete + +from sagemaker.training_compiler.config import TrainingCompilerConfig as BaseConfig + +logger: Incomplete + +class TrainingCompilerConfig(BaseConfig): + SUPPORTED_INSTANCE_CLASS_PREFIXES: Incomplete + MIN_SUPPORTED_VERSION: str + MAX_SUPPORTED_VERSION: str + def __init__(self, enabled: bool = True, debug: bool = False) -> None: ... + @classmethod + def validate(cls, estimator) -> None: ... diff --git a/stubs/sagemaker/sagemaker/training_compiler/__init__.pyi b/stubs/sagemaker/sagemaker/training_compiler/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/training_compiler/config.pyi b/stubs/sagemaker/sagemaker/training_compiler/config.pyi new file mode 100644 index 000000000000..6cb0bec952c7 --- /dev/null +++ b/stubs/sagemaker/sagemaker/training_compiler/config.pyi @@ -0,0 +1,16 @@ +from _typeshed import Incomplete + +logger: Incomplete + +class TrainingCompilerConfig: + DEBUG_PATH: str + SUPPORTED_INSTANCE_CLASS_PREFIXES: Incomplete + HP_ENABLE_COMPILER: str + HP_ENABLE_DEBUG: str + enabled: Incomplete + debug: Incomplete + def __init__(self, enabled: bool = True, debug: bool = False) -> None: ... + def __nonzero__(self): ... + def disclaimers_and_warnings(self) -> None: ... + @classmethod + def validate(cls, estimator) -> None: ... diff --git a/stubs/sagemaker/sagemaker/transformer.pyi b/stubs/sagemaker/sagemaker/transformer.pyi new file mode 100644 index 000000000000..bd9ea6f0ae4d --- /dev/null +++ b/stubs/sagemaker/sagemaker/transformer.pyi @@ -0,0 +1,108 @@ +import abc +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker.inputs import BatchDataCaptureConfig +from sagemaker.job import _Job +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable + +class Transformer: + JOB_CLASS_NAME: str + model_name: Incomplete + strategy: Incomplete + output_path: Incomplete + accept: Incomplete + assemble_with: Incomplete + instance_count: Incomplete + instance_type: Incomplete + max_concurrent_transforms: Incomplete + max_payload: Incomplete + tags: Incomplete + base_transform_job_name: Incomplete + latest_transform_job: Incomplete + sagemaker_session: Incomplete + volume_kms_key: Incomplete + output_kms_key: Incomplete + env: Incomplete + def __init__( + self, + model_name: Union[str, PipelineVariable], + instance_count: Union[int, PipelineVariable], + instance_type: Union[str, PipelineVariable], + strategy: Optional[Union[str, PipelineVariable]] = None, + assemble_with: Optional[Union[str, PipelineVariable]] = None, + output_path: Optional[Union[str, PipelineVariable]] = None, + output_kms_key: Optional[Union[str, PipelineVariable]] = None, + accept: Optional[Union[str, PipelineVariable]] = None, + max_concurrent_transforms: Optional[Union[int, PipelineVariable]] = None, + max_payload: Optional[Union[int, PipelineVariable]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + base_transform_job_name: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + ) -> None: ... + def transform( + self, + data: Union[str, PipelineVariable], + data_type: Union[str, PipelineVariable] = "S3Prefix", + content_type: Optional[Union[str, PipelineVariable]] = None, + compression_type: Optional[Union[str, PipelineVariable]] = None, + split_type: Optional[Union[str, PipelineVariable]] = None, + job_name: Optional[str] = None, + input_filter: Optional[Union[str, PipelineVariable]] = None, + output_filter: Optional[Union[str, PipelineVariable]] = None, + join_source: Optional[Union[str, PipelineVariable]] = None, + experiment_config: Optional[Dict[str, str]] = None, + model_client_config: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + batch_data_capture_config: BatchDataCaptureConfig = None, + wait: bool = True, + logs: bool = True, + ): ... + def transform_with_monitoring( + self, + monitoring_config, + monitoring_resource_config, + data: str, + data_type: str = "S3Prefix", + content_type: str = None, + compression_type: str = None, + split_type: str = None, + input_filter: str = None, + output_filter: str = None, + join_source: str = None, + model_client_config: Dict[str, str] = None, + batch_data_capture_config: BatchDataCaptureConfig = None, + monitor_before_transform: bool = False, + supplied_baseline_statistics: str = None, + supplied_baseline_constraints: str = None, + wait: bool = True, + pipeline_name: str = None, + role: str = None, + ): ... + def delete_model(self) -> None: ... + def wait(self, logs: bool = True) -> None: ... + def stop_transform_job(self, wait: bool = True) -> None: ... + @classmethod + def attach(cls, transform_job_name, sagemaker_session: Incomplete | None = None): ... + +class _TransformJob(_Job, metaclass=abc.ABCMeta): + @classmethod + def start_new( + cls, + transformer, + data, + data_type, + content_type, + compression_type, + split_type, + input_filter, + output_filter, + join_source, + experiment_config, + model_client_config, + batch_data_capture_config, + ): ... + def wait(self, logs: bool = True) -> None: ... + def stop(self) -> None: ... diff --git a/stubs/sagemaker/sagemaker/tuner.pyi b/stubs/sagemaker/sagemaker/tuner.pyi new file mode 100644 index 000000000000..bb9f6317d850 --- /dev/null +++ b/stubs/sagemaker/sagemaker/tuner.pyi @@ -0,0 +1,223 @@ +import abc +from _typeshed import Incomplete +from enum import Enum +from typing import Dict, List, Optional, Set, Union + +from sagemaker.amazon.amazon_estimator import FileSystemRecordSet, RecordSet +from sagemaker.estimator import EstimatorBase +from sagemaker.inputs import FileSystemInput, TrainingInput +from sagemaker.job import _Job +from sagemaker.parameter import ParameterRange +from sagemaker.workflow.entities import PipelineVariable + +AMAZON_ESTIMATOR_MODULE: str +AMAZON_ESTIMATOR_CLS_NAMES: Incomplete +HYPERPARAMETER_TUNING_JOB_NAME: str +PARENT_HYPERPARAMETER_TUNING_JOBS: str +WARM_START_TYPE: str +HYPERBAND_STRATEGY_CONFIG: str +HYPERBAND_MIN_RESOURCE: str +HYPERBAND_MAX_RESOURCE: str +GRID_SEARCH: str +MAX_NUMBER_OF_TRAINING_JOBS_NOT_IMPROVING: str +BEST_OBJECTIVE_NOT_IMPROVING: str +CONVERGENCE_DETECTED: str +COMPLETE_ON_CONVERGENCE_DETECTED: str +TARGET_OBJECTIVE_METRIC_VALUE: str +MAX_RUNTIME_IN_SECONDS: str +logger: Incomplete + +class WarmStartTypes(Enum): + IDENTICAL_DATA_AND_ALGORITHM: str + TRANSFER_LEARNING: str + +class WarmStartConfig: + type: Incomplete + parents: Incomplete + def __init__(self, warm_start_type: WarmStartTypes, parents: Set[Union[str, PipelineVariable]]) -> None: ... + @classmethod + def from_job_desc(cls, warm_start_config): ... + def to_input_req(self): ... + +class HyperbandStrategyConfig: + min_resource: Incomplete + max_resource: Incomplete + def __init__(self, max_resource: int, min_resource: int) -> None: ... + @classmethod + def from_job_desc(cls, hyperband_strategy_config): ... + def to_input_req(self): ... + +class StrategyConfig: + hyperband_strategy_config: Incomplete + def __init__(self, hyperband_strategy_config: HyperbandStrategyConfig) -> None: ... + @classmethod + def from_job_desc(cls, strategy_config): ... + def to_input_req(self): ... + +class InstanceConfig: + instance_count: Incomplete + instance_type: Incomplete + volume_size: Incomplete + def __init__( + self, + instance_count: Union[int, PipelineVariable] = None, + instance_type: Union[str, PipelineVariable] = None, + volume_size: Union[int, PipelineVariable] = 30, + ) -> None: ... + @classmethod + def from_job_desc(cls, instance_config): ... + def to_input_req(self): ... + +class TuningJobCompletionCriteriaConfig: + max_number_of_training_jobs_not_improving: Incomplete + complete_on_convergence: Incomplete + target_objective_metric_value: Incomplete + def __init__( + self, + max_number_of_training_jobs_not_improving: int = None, + complete_on_convergence: bool = None, + target_objective_metric_value: float = None, + ) -> None: ... + @classmethod + def from_job_desc(cls, completion_criteria_config): ... + def to_input_req(self): ... + +class HyperparameterTuner: + TUNING_JOB_NAME_MAX_LENGTH: int + SAGEMAKER_ESTIMATOR_MODULE: str + SAGEMAKER_ESTIMATOR_CLASS_NAME: str + DEFAULT_ESTIMATOR_MODULE: str + DEFAULT_ESTIMATOR_CLS_NAME: str + estimator: Incomplete + objective_metric_name: Incomplete + metric_definitions: Incomplete + estimator_dict: Incomplete + objective_metric_name_dict: Incomplete + metric_definitions_dict: Incomplete + static_hyperparameters: Incomplete + auto_parameters: Incomplete + auto_parameters_dict: Incomplete + hyperparameters_to_keep_static: Incomplete + hyperparameters_to_keep_static_dict: Incomplete + static_hyperparameters_dict: Incomplete + strategy: Incomplete + strategy_config: Incomplete + completion_criteria_config: Incomplete + objective_type: Incomplete + max_jobs: Incomplete + max_parallel_jobs: Incomplete + max_runtime_in_seconds: Incomplete + tags: Incomplete + base_tuning_job_name: Incomplete + latest_tuning_job: Incomplete + warm_start_config: Incomplete + early_stopping_type: Incomplete + random_seed: Incomplete + instance_configs_dict: Incomplete + instance_configs: Incomplete + autotune: Incomplete + def __init__( + self, + estimator: EstimatorBase, + objective_metric_name: Union[str, PipelineVariable], + hyperparameter_ranges: Dict[str, ParameterRange], + metric_definitions: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + strategy: Union[str, PipelineVariable] = "Bayesian", + objective_type: Union[str, PipelineVariable] = "Maximize", + max_jobs: Union[int, PipelineVariable] = None, + max_parallel_jobs: Union[int, PipelineVariable] = 1, + max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + base_tuning_job_name: Optional[str] = None, + warm_start_config: Optional[WarmStartConfig] = None, + strategy_config: Optional[StrategyConfig] = None, + completion_criteria_config: Optional[TuningJobCompletionCriteriaConfig] = None, + early_stopping_type: Union[str, PipelineVariable] = "Off", + estimator_name: Optional[str] = None, + random_seed: Optional[int] = None, + autotune: bool = False, + hyperparameters_to_keep_static: Optional[List[str]] = None, + ) -> None: ... + def override_resource_config(self, instance_configs: Union[List[InstanceConfig], Dict[str, List[InstanceConfig]]]): ... + def fit( + self, + inputs: Optional[Union[str, Dict, List, TrainingInput, FileSystemInput, RecordSet, FileSystemRecordSet]] = None, + job_name: Optional[str] = None, + include_cls_metadata: Union[bool, Dict[str, bool]] = False, + estimator_kwargs: Optional[Dict[str, dict]] = None, + wait: bool = True, + **kwargs, + ): ... + @classmethod + def attach( + cls, + tuning_job_name, + sagemaker_session: Incomplete | None = None, + job_details: Incomplete | None = None, + estimator_cls: Incomplete | None = None, + ): ... + def deploy( + self, + initial_instance_count, + instance_type, + serializer: Incomplete | None = None, + deserializer: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + endpoint_name: Incomplete | None = None, + wait: bool = True, + model_name: Incomplete | None = None, + kms_key: Incomplete | None = None, + data_capture_config: Incomplete | None = None, + **kwargs, + ): ... + def stop_tuning_job(self) -> None: ... + def describe(self): ... + def wait(self) -> None: ... + def best_estimator(self, best_training_job: Incomplete | None = None): ... + def best_training_job(self): ... + def hyperparameter_ranges(self): ... + def hyperparameter_ranges_dict(self): ... + @property + def sagemaker_session(self): ... + def analytics(self): ... + def transfer_learning_tuner(self, additional_parents: Incomplete | None = None, estimator: Incomplete | None = None): ... + def identical_dataset_and_algorithm_tuner(self, additional_parents: Incomplete | None = None): ... + @classmethod + def create( + cls, + estimator_dict, + objective_metric_name_dict, + hyperparameter_ranges_dict, + metric_definitions_dict: Incomplete | None = None, + base_tuning_job_name: Incomplete | None = None, + strategy: str = "Bayesian", + strategy_config: Incomplete | None = None, + completion_criteria_config: Incomplete | None = None, + objective_type: str = "Maximize", + max_jobs: Incomplete | None = None, + max_parallel_jobs: int = 1, + max_runtime_in_seconds: Incomplete | None = None, + tags: Incomplete | None = None, + warm_start_config: Incomplete | None = None, + early_stopping_type: str = "Off", + random_seed: Incomplete | None = None, + autotune: bool = False, + hyperparameters_to_keep_static_dict: Incomplete | None = None, + ): ... + delete_endpoint: Incomplete + +class _TuningJob(_Job, metaclass=abc.ABCMeta): + @classmethod + def start_new(cls, tuner, inputs): ... + def stop(self) -> None: ... + def wait(self) -> None: ... + +def create_identical_dataset_and_algorithm_tuner( + parent, additional_parents: Incomplete | None = None, sagemaker_session: Incomplete | None = None +): ... +def create_transfer_learning_tuner( + parent, + additional_parents: Incomplete | None = None, + estimator: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, +): ... diff --git a/stubs/sagemaker/sagemaker/user_agent.pyi b/stubs/sagemaker/sagemaker/user_agent.pyi new file mode 100644 index 000000000000..c35abbb919b5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/user_agent.pyi @@ -0,0 +1,10 @@ +from _typeshed import Incomplete + +SDK_VERSION: Incomplete +OS_NAME: Incomplete +OS_VERSION: Incomplete +OS_NAME_VERSION: Incomplete +PYTHON_VERSION: Incomplete + +def determine_prefix(user_agent: str = ""): ... +def prepend_user_agent(client) -> None: ... diff --git a/stubs/sagemaker/sagemaker/utilities/__init__.pyi b/stubs/sagemaker/sagemaker/utilities/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/utilities/cache.pyi b/stubs/sagemaker/sagemaker/utilities/cache.pyi new file mode 100644 index 000000000000..0229648f7e3d --- /dev/null +++ b/stubs/sagemaker/sagemaker/utilities/cache.pyi @@ -0,0 +1,24 @@ +import datetime +from _typeshed import Incomplete +from typing import Callable, Optional, TypeVar + +KeyType = TypeVar("KeyType") +ValType = TypeVar("ValType") + +class LRUCache: + class Element: + value: Incomplete + creation_time: Incomplete + def __init__(self, value: ValType, creation_time: datetime.datetime) -> None: ... + + def __init__( + self, + max_cache_items: int, + expiration_horizon: datetime.timedelta, + retrieval_function: Callable[[KeyType, ValType], ValType], + ) -> None: ... + def __len__(self) -> int: ... + def __contains__(self, key: KeyType) -> bool: ... + def clear(self) -> None: ... + def get(self, key: KeyType, data_source_fallback: Optional[bool] = True) -> ValType: ... + def put(self, key: KeyType, value: Optional[ValType] = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/utilities/search_expression.pyi b/stubs/sagemaker/sagemaker/utilities/search_expression.pyi new file mode 100644 index 000000000000..6cad2c5a0123 --- /dev/null +++ b/stubs/sagemaker/sagemaker/utilities/search_expression.pyi @@ -0,0 +1,47 @@ +from _typeshed import Incomplete +from enum import Enum + +from sagemaker.apiutils._base_types import ApiObject + +class Operator(Enum): + EQUALS: str + NOT_EQUALS: str + GREATER_THAN: str + GREATER_THAN_OR_EQUAL: str + LESS_THAN: str + LESS_THAN_OR_EQUAL: str + CONTAINS: str + EXISTS: str + NOT_EXISTS: str + +class BooleanOperator(Enum): + AND: str + OR: str + +class SearchObject(ApiObject): + def to_boto(self): ... + +class Filter(SearchObject): + name: Incomplete + operator: Incomplete + value: Incomplete + def __init__(self, name, operator: Incomplete | None = None, value: Incomplete | None = None, **kwargs) -> None: ... + +class NestedFilter(SearchObject): + nested_property_name: Incomplete + filters: Incomplete + def __init__(self, property_name, filters, **kwargs) -> None: ... + +class SearchExpression(SearchObject): + filters: Incomplete + nested_filters: Incomplete + operator: Incomplete + sub_expressions: Incomplete + def __init__( + self, + filters: Incomplete | None = None, + nested_filters: Incomplete | None = None, + sub_expressions: Incomplete | None = None, + boolean_operator=..., + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/utils.pyi b/stubs/sagemaker/sagemaker/utils.pyi new file mode 100644 index 000000000000..1cd58a102e0f --- /dev/null +++ b/stubs/sagemaker/sagemaker/utils.pyi @@ -0,0 +1,113 @@ +import abc +from _typeshed import Incomplete +from collections.abc import Generator +from typing import Any, List, Optional + +ECR_URI_PATTERN: str +MAX_BUCKET_PATHS_COUNT: int +S3_PREFIX: str +HTTP_PREFIX: str +HTTPS_PREFIX: str +DEFAULT_SLEEP_TIME_SECONDS: int +WAITING_DOT_NUMBER: int +logger: Incomplete + +def name_from_image(image, max_length: int = 63): ... +def name_from_base(base, max_length: int = 63, short: bool = False): ... +def unique_name_from_base(base, max_length: int = 63): ... +def base_name_from_image(image, default_base_name: Incomplete | None = None): ... +def base_from_name(name): ... +def sagemaker_timestamp(): ... +def sagemaker_short_timestamp(): ... +def build_dict(key, value): ... +def get_config_value(key_path, config): ... +def get_nested_value(dictionary: dict, nested_keys: List[str]): ... +def set_nested_value(dictionary: dict, nested_keys: List[str], value_to_set: object): ... +def get_short_version(framework_version): ... +def secondary_training_status_changed(current_job_description, prev_job_description): ... +def secondary_training_status_message(job_description, prev_description): ... +def download_folder(bucket_name, prefix, target, sagemaker_session) -> None: ... +def create_tar_file(source_files, target: Incomplete | None = None): ... +def repack_model( + inference_script, + source_directory, + dependencies, + model_uri, + repacked_model_uri, + sagemaker_session, + kms_key: Incomplete | None = None, +) -> None: ... +def download_file_from_url(url, dst, sagemaker_session) -> None: ... +def download_file(bucket_name, path, target, sagemaker_session) -> None: ... +def sts_regional_endpoint(region): ... +def retries(max_retry_count, exception_message_prefix, seconds_to_sleep=10) -> Generator[Incomplete, None, None]: ... +def retry_with_backoff(callable_func, num_attempts: int = 8, botocore_client_error_code: Incomplete | None = None): ... + +class DeferredError: + exc: Incomplete + def __init__(self, exception) -> None: ... + def __getattr__(self, name) -> None: ... + +class DataConfig(abc.ABC, metaclass=abc.ABCMeta): + @abc.abstractmethod + def fetch_data_config(self): ... + +class S3DataConfig(DataConfig): + bucket_name: Incomplete + prefix: Incomplete + sagemaker_session: Incomplete + def __init__(self, sagemaker_session, bucket_name, prefix) -> None: ... + def fetch_data_config(self): ... + def get_data_bucket(self, region_requested: Incomplete | None = None): ... + +get_ecr_image_uri_prefix: Incomplete + +def update_container_with_inference_params( + framework: Incomplete | None = None, + framework_version: Incomplete | None = None, + nearest_model_name: Incomplete | None = None, + data_input_configuration: Incomplete | None = None, + container_def: Incomplete | None = None, + container_list: Incomplete | None = None, +): ... +def construct_container_object(obj, data_input_configuration, framework, framework_version, nearest_model_name): ... +def pop_out_unused_kwarg(arg_name: str, kwargs: dict, override_val: Optional[str] = None): ... +def to_string(obj: object): ... +def get_module(module_name): ... +def check_and_get_run_experiment_config(experiment_config: Optional[dict] = None) -> dict: ... +def resolve_value_from_config( + direct_input: Incomplete | None = None, + config_path: str = None, + default_value: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + sagemaker_config: dict = None, +): ... +def get_sagemaker_config_value(sagemaker_session, key, sagemaker_config: dict = None): ... +def resolve_class_attribute_from_config( + clazz: Optional[type], + instance: Optional[object], + attribute: str, + config_path: str, + default_value: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, +): ... +def resolve_nested_dict_value_from_config( + dictionary: dict, + nested_keys: List[str], + config_path: str, + default_value: object = None, + sagemaker_session: Incomplete | None = None, +): ... +def update_list_of_dicts_with_values_from_config( + input_list, + config_key_path, + required_key_paths: List[str] = None, + union_key_paths: List[List[str]] = None, + sagemaker_session: Incomplete | None = None, +): ... +def update_nested_dictionary_with_values_from_config( + source_dict, config_key_path, sagemaker_session: Incomplete | None = None +) -> dict: ... +def stringify_object(obj: Any) -> str: ... +def volume_size_supported(instance_type: str) -> bool: ... +def instance_supports_kms(instance_type: str) -> bool: ... diff --git a/stubs/sagemaker/sagemaker/vpc_utils.pyi b/stubs/sagemaker/sagemaker/vpc_utils.pyi new file mode 100644 index 000000000000..4cc80b644223 --- /dev/null +++ b/stubs/sagemaker/sagemaker/vpc_utils.pyi @@ -0,0 +1,8 @@ +SUBNETS_KEY: str +SECURITY_GROUP_IDS_KEY: str +VPC_CONFIG_KEY: str +VPC_CONFIG_DEFAULT: str + +def to_dict(subnets, security_group_ids): ... +def from_dict(vpc_config, do_sanitize: bool = False): ... +def sanitize(vpc_config): ... diff --git a/stubs/sagemaker/sagemaker/workflow/__init__.pyi b/stubs/sagemaker/sagemaker/workflow/__init__.pyi new file mode 100644 index 000000000000..3d798d52b642 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/__init__.pyi @@ -0,0 +1,2 @@ +def is_pipeline_variable(var: object) -> bool: ... +def is_pipeline_parameter_string(var: object) -> bool: ... diff --git a/stubs/sagemaker/sagemaker/workflow/_repack_model.pyi b/stubs/sagemaker/sagemaker/workflow/_repack_model.pyi new file mode 100644 index 000000000000..6ddc85724700 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/_repack_model.pyi @@ -0,0 +1,5 @@ +from _typeshed import Incomplete + +def repack( + inference_script, model_archive, dependencies: Incomplete | None = None, source_dir: Incomplete | None = None +) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/_utils.pyi b/stubs/sagemaker/sagemaker/workflow/_utils.pyi new file mode 100644 index 000000000000..7f65070b08ba --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/_utils.pyi @@ -0,0 +1,97 @@ +from _typeshed import Incomplete +from typing import List, Optional, Union + +from sagemaker.estimator import EstimatorBase +from sagemaker.workflow.entities import RequestType as RequestType +from sagemaker.workflow.retry import RetryPolicy +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.workflow.steps import ConfigurableRetryStep, Step, TrainingStep + +logger: Incomplete +FRAMEWORK_VERSION: str +INSTANCE_TYPE: str +REPACK_SCRIPT: str +REPACK_SCRIPT_LAUNCHER: str +LAUNCH_REPACK_SCRIPT_CMD: str + +class _RepackModelStep(TrainingStep): + sagemaker_session: Incomplete + role: Incomplete + def __init__( + self, + name: str, + sagemaker_session, + role, + model_data: str, + entry_point: str, + display_name: str = None, + description: str = None, + source_dir: str = None, + dependencies: List = None, + depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + retry_policies: List[RetryPolicy] = None, + subnets: Incomplete | None = None, + security_group_ids: Incomplete | None = None, + **kwargs, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self): ... + +class _RegisterModelStep(ConfigurableRetryStep): + step_args: Incomplete + estimator: Incomplete + model_data: Incomplete + content_types: Incomplete + response_types: Incomplete + inference_instances: Incomplete + transform_instances: Incomplete + model_package_group_name: Incomplete + tags: Incomplete + model_metrics: Incomplete + drift_check_baselines: Incomplete + customer_metadata_properties: Incomplete + domain: Incomplete + sample_payload_url: Incomplete + task: Incomplete + metadata_properties: Incomplete + approval_status: Incomplete + image_uri: Incomplete + compile_model_family: Incomplete + description: Incomplete + kwargs: Incomplete + container_def_list: Incomplete + def __init__( + self, + name: str, + step_args: Optional[dict] = None, + content_types: Optional[list] = None, + response_types: Optional[list] = None, + inference_instances: Optional[list] = None, + transform_instances: Optional[list] = None, + estimator: EstimatorBase = None, + model_data: Incomplete | None = None, + model_package_group_name: Incomplete | None = None, + model_metrics: Incomplete | None = None, + metadata_properties: Incomplete | None = None, + approval_status: str = "PendingManualApproval", + image_uri: Incomplete | None = None, + compile_model_family: Incomplete | None = None, + display_name: str = None, + description: Incomplete | None = None, + depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + retry_policies: Optional[List[RetryPolicy]] = None, + tags: Incomplete | None = None, + container_def_list: Incomplete | None = None, + drift_check_baselines: Incomplete | None = None, + customer_metadata_properties: Incomplete | None = None, + domain: Incomplete | None = None, + sample_payload_url: Incomplete | None = None, + task: Incomplete | None = None, + **kwargs, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self): ... diff --git a/stubs/sagemaker/sagemaker/workflow/airflow.pyi b/stubs/sagemaker/sagemaker/workflow/airflow.pyi new file mode 100644 index 000000000000..fee8e98a15d9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/airflow.pyi @@ -0,0 +1,102 @@ +from _typeshed import Incomplete + +def prepare_framework(estimator, s3_operations) -> None: ... +def prepare_amazon_algorithm_estimator(estimator, inputs, mini_batch_size: Incomplete | None = None) -> None: ... +def training_base_config( + estimator, inputs: Incomplete | None = None, job_name: Incomplete | None = None, mini_batch_size: Incomplete | None = None +): ... +def training_config( + estimator, inputs: Incomplete | None = None, job_name: Incomplete | None = None, mini_batch_size: Incomplete | None = None +): ... +def tuning_config( + tuner, + inputs, + job_name: Incomplete | None = None, + include_cls_metadata: bool = False, + mini_batch_size: Incomplete | None = None, +): ... +def update_submit_s3_uri(estimator, job_name) -> None: ... +def update_estimator_from_task(estimator, task_id, task_type) -> None: ... +def prepare_framework_container_def(model, instance_type, s3_operations): ... +def model_config( + model, instance_type: Incomplete | None = None, role: Incomplete | None = None, image_uri: Incomplete | None = None +): ... +def model_config_from_estimator( + estimator, + task_id, + task_type, + instance_type: Incomplete | None = None, + role: Incomplete | None = None, + image_uri: Incomplete | None = None, + name: Incomplete | None = None, + model_server_workers: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", +): ... +def transform_config( + transformer, + data, + data_type: str = "S3Prefix", + content_type: Incomplete | None = None, + compression_type: Incomplete | None = None, + split_type: Incomplete | None = None, + job_name: Incomplete | None = None, + input_filter: Incomplete | None = None, + output_filter: Incomplete | None = None, + join_source: Incomplete | None = None, +): ... +def transform_config_from_estimator( + estimator, + task_id, + task_type, + instance_count, + instance_type, + data, + data_type: str = "S3Prefix", + content_type: Incomplete | None = None, + compression_type: Incomplete | None = None, + split_type: Incomplete | None = None, + job_name: Incomplete | None = None, + model_name: Incomplete | None = None, + strategy: Incomplete | None = None, + assemble_with: Incomplete | None = None, + output_path: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + accept: Incomplete | None = None, + env: Incomplete | None = None, + max_concurrent_transforms: Incomplete | None = None, + max_payload: Incomplete | None = None, + tags: Incomplete | None = None, + role: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + model_server_workers: Incomplete | None = None, + image_uri: Incomplete | None = None, + vpc_config_override: Incomplete | None = None, + input_filter: Incomplete | None = None, + output_filter: Incomplete | None = None, + join_source: Incomplete | None = None, +): ... +def deploy_config( + model, initial_instance_count, instance_type, endpoint_name: Incomplete | None = None, tags: Incomplete | None = None +): ... +def deploy_config_from_estimator( + estimator, + task_id, + task_type, + initial_instance_count, + instance_type, + model_name: Incomplete | None = None, + endpoint_name: Incomplete | None = None, + tags: Incomplete | None = None, + **kwargs, +): ... +def processing_config( + processor, + inputs: Incomplete | None = None, + outputs: Incomplete | None = None, + job_name: Incomplete | None = None, + experiment_config: Incomplete | None = None, + container_arguments: Incomplete | None = None, + container_entrypoint: Incomplete | None = None, + kms_key_id: Incomplete | None = None, +): ... +def input_output_list_converter(object_list): ... diff --git a/stubs/sagemaker/sagemaker/workflow/automl_step.pyi b/stubs/sagemaker/sagemaker/workflow/automl_step.pyi new file mode 100644 index 000000000000..15d77943ea7b --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/automl_step.pyi @@ -0,0 +1,28 @@ +from _typeshed import Incomplete +from typing import List, Optional, Union + +from sagemaker.workflow.entities import RequestType as RequestType +from sagemaker.workflow.pipeline_context import _JobStepArguments +from sagemaker.workflow.retry import RetryPolicy +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.workflow.steps import CacheConfig, ConfigurableRetryStep, Step + +class AutoMLStep(ConfigurableRetryStep): + step_args: Incomplete + cache_config: Incomplete + def __init__( + self, + name: str, + step_args: _JobStepArguments, + display_name: str = None, + description: str = None, + cache_config: CacheConfig = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + retry_policies: List[RetryPolicy] = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self): ... + def to_request(self) -> RequestType: ... + def get_best_auto_ml_model(self, role, sagemaker_session: Incomplete | None = None): ... diff --git a/stubs/sagemaker/sagemaker/workflow/callback_step.pyi b/stubs/sagemaker/sagemaker/workflow/callback_step.pyi new file mode 100644 index 000000000000..ad94b479878c --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/callback_step.pyi @@ -0,0 +1,46 @@ +from _typeshed import Incomplete +from enum import Enum +from typing import Dict, List, Optional, Union + +from sagemaker.workflow.entities import DefaultEnumMeta, RequestType as RequestType +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.workflow.steps import CacheConfig, Step + +class CallbackOutputTypeEnum(Enum, metaclass=DefaultEnumMeta): + String: str + Integer: str + Boolean: str + Float: str + +class CallbackOutput: + output_name: str + output_type: CallbackOutputTypeEnum + def to_request(self) -> RequestType: ... + def expr(self, step_name) -> Dict[str, str]: ... + def __init__(self, output_name, output_type) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class CallbackStep(Step): + sqs_queue_url: Incomplete + outputs: Incomplete + cache_config: Incomplete + inputs: Incomplete + def __init__( + self, + name: str, + sqs_queue_url: str, + inputs: dict, + outputs: List[CallbackOutput], + display_name: str = None, + description: str = None, + cache_config: CacheConfig = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self): ... + def to_request(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/check_job_config.pyi b/stubs/sagemaker/sagemaker/workflow/check_job_config.pyi new file mode 100644 index 000000000000..1daf8b3a1424 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/check_job_config.pyi @@ -0,0 +1,30 @@ +from _typeshed import Incomplete + +class CheckJobConfig: + role: Incomplete + instance_count: Incomplete + instance_type: Incomplete + volume_size_in_gb: Incomplete + volume_kms_key: Incomplete + output_kms_key: Incomplete + max_runtime_in_seconds: Incomplete + base_job_name: Incomplete + sagemaker_session: Incomplete + env: Incomplete + tags: Incomplete + network_config: Incomplete + def __init__( + self, + role, + instance_count: int = 1, + instance_type: str = "ml.m5.xlarge", + volume_size_in_gb: int = 30, + volume_kms_key: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + max_runtime_in_seconds: Incomplete | None = None, + base_job_name: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + env: Incomplete | None = None, + tags: Incomplete | None = None, + network_config: Incomplete | None = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi b/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi new file mode 100644 index 000000000000..ae7fd5e1a7fd --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi @@ -0,0 +1,90 @@ +from _typeshed import Incomplete +from abc import ABC +from typing import List, Optional, Union + +from sagemaker.clarify import BiasConfig, DataConfig, ModelConfig, ModelPredictedLabelConfig, SHAPConfig +from sagemaker.workflow.check_job_config import CheckJobConfig +from sagemaker.workflow.entities import PipelineVariable, RequestType as RequestType +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.workflow.steps import CacheConfig, Step + +class ClarifyCheckConfig(ABC): + data_config: DataConfig + kms_key: str + monitoring_analysis_config_uri: str + def __init__(self, data_config, kms_key, monitoring_analysis_config_uri) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class DataBiasCheckConfig(ClarifyCheckConfig): + data_bias_config: BiasConfig + methods: Union[str, List[str]] + def __init__(self, data_config, kms_key, monitoring_analysis_config_uri, data_bias_config, methods) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class ModelBiasCheckConfig(ClarifyCheckConfig): + data_bias_config: BiasConfig + model_config: ModelConfig + model_predicted_label_config: ModelPredictedLabelConfig + methods: Union[str, List[str]] + def __init__( + self, + data_config, + kms_key, + monitoring_analysis_config_uri, + data_bias_config, + model_config, + model_predicted_label_config, + methods, + ) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class ModelExplainabilityCheckConfig(ClarifyCheckConfig): + model_config: ModelConfig + explainability_config: SHAPConfig + model_scores: Union[str, int, ModelPredictedLabelConfig] + def __init__( + self, data_config, kms_key, monitoring_analysis_config_uri, model_config, explainability_config, model_scores + ) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class ClarifyCheckStep(Step): + skip_check: Incomplete + fail_on_violation: Incomplete + register_new_baseline: Incomplete + clarify_check_config: Incomplete + check_job_config: Incomplete + model_package_group_name: Incomplete + supplied_baseline_constraints: Incomplete + cache_config: Incomplete + def __init__( + self, + name: str, + clarify_check_config: ClarifyCheckConfig, + check_job_config: CheckJobConfig, + skip_check: Union[bool, PipelineVariable] = False, + fail_on_violation: Union[bool, PipelineVariable] = True, + register_new_baseline: Union[bool, PipelineVariable] = False, + model_package_group_name: Union[str, PipelineVariable] = None, + supplied_baseline_constraints: Union[str, PipelineVariable] = None, + display_name: str = None, + description: str = None, + cache_config: CacheConfig = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self): ... + def to_request(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/condition_step.pyi b/stubs/sagemaker/sagemaker/workflow/condition_step.pyi new file mode 100644 index 000000000000..b82ddd1929e3 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/condition_step.pyi @@ -0,0 +1,33 @@ +from _typeshed import Incomplete +from typing import List, Optional, Union + +from sagemaker.workflow.conditions import Condition +from sagemaker.workflow.entities import RequestType as RequestType +from sagemaker.workflow.functions import JsonGet as NewJsonGet +from sagemaker.workflow.properties import PropertyFile +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.workflow.steps import Step + +class ConditionStep(Step): + conditions: Incomplete + if_steps: Incomplete + else_steps: Incomplete + def __init__( + self, + name: str, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + display_name: str = None, + description: str = None, + conditions: List[Condition] = None, + if_steps: List[Union[Step, StepCollection]] = None, + else_steps: List[Union[Step, StepCollection]] = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def step_only_arguments(self): ... + @property + def properties(self): ... + +class JsonGet(NewJsonGet): + def __init__(self, step: Step, property_file: Union[PropertyFile, str], json_path: str) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/conditions.pyi b/stubs/sagemaker/sagemaker/workflow/conditions.pyi new file mode 100644 index 000000000000..e24ccaa5048e --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/conditions.pyi @@ -0,0 +1,92 @@ +import abc +from _typeshed import Incomplete +from enum import Enum +from typing import Dict, List, Union + +from sagemaker.workflow.entities import ( + DefaultEnumMeta, + Entity, + Expression, + PrimitiveType as PrimitiveType, + RequestType as RequestType, +) +from sagemaker.workflow.execution_variables import ExecutionVariable +from sagemaker.workflow.parameters import Parameter +from sagemaker.workflow.properties import Properties + +ConditionValueType = Union[ExecutionVariable, Parameter, Properties] + +class ConditionTypeEnum(Enum, metaclass=DefaultEnumMeta): + EQ: str + GT: str + GTE: str + IN: str + LT: str + LTE: str + NOT: str + OR: str + +class Condition(Entity, metaclass=abc.ABCMeta): + condition_type: ConditionTypeEnum + def __init__(self, condition_type) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class ConditionComparison(Condition): + left: Union[ConditionValueType, PrimitiveType] + right: Union[ConditionValueType, PrimitiveType] + def to_request(self) -> RequestType: ... + def __init__(self, condition_type, left, right) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class ConditionEquals(ConditionComparison): + def __init__( + self, left: Union[ConditionValueType, PrimitiveType], right: Union[ConditionValueType, PrimitiveType] + ) -> None: ... + +class ConditionGreaterThan(ConditionComparison): + def __init__( + self, left: Union[ConditionValueType, PrimitiveType], right: Union[ConditionValueType, PrimitiveType] + ) -> None: ... + +class ConditionGreaterThanOrEqualTo(ConditionComparison): + def __init__( + self, left: Union[ConditionValueType, PrimitiveType], right: Union[ConditionValueType, PrimitiveType] + ) -> None: ... + +class ConditionLessThan(ConditionComparison): + def __init__( + self, left: Union[ConditionValueType, PrimitiveType], right: Union[ConditionValueType, PrimitiveType] + ) -> None: ... + +class ConditionLessThanOrEqualTo(ConditionComparison): + def __init__( + self, left: Union[ConditionValueType, PrimitiveType], right: Union[ConditionValueType, PrimitiveType] + ) -> None: ... + +class ConditionIn(Condition): + value: Incomplete + in_values: Incomplete + def __init__( + self, value: Union[ConditionValueType, PrimitiveType], in_values: List[Union[ConditionValueType, PrimitiveType]] + ) -> None: ... + def to_request(self) -> RequestType: ... + +class ConditionNot(Condition): + expression: Incomplete + def __init__(self, expression: Condition) -> None: ... + def to_request(self) -> RequestType: ... + +class ConditionOr(Condition): + conditions: Incomplete + def __init__(self, conditions: List[Condition] = None) -> None: ... + def to_request(self) -> RequestType: ... + +def primitive_or_expr( + value: Union[ExecutionVariable, Expression, PrimitiveType, Parameter, Properties] +) -> Union[Dict[str, str], PrimitiveType]: ... diff --git a/stubs/sagemaker/sagemaker/workflow/emr_step.pyi b/stubs/sagemaker/sagemaker/workflow/emr_step.pyi new file mode 100644 index 000000000000..480e21608c87 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/emr_step.pyi @@ -0,0 +1,46 @@ +from _typeshed import Incomplete +from typing import Any, Dict, List, Optional, Union + +from sagemaker.workflow.entities import RequestType as RequestType +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.workflow.steps import CacheConfig, Step + +class EMRStepConfig: + jar: Incomplete + args: Incomplete + main_class: Incomplete + properties: Incomplete + def __init__(self, jar, args: List[str] = None, main_class: str = None, properties: List[dict] = None) -> None: ... + def to_request(self) -> RequestType: ... + +INSTANCES: str +INSTANCEGROUPS: str +INSTANCEFLEETS: str +ERR_STR_WITH_NAME_AUTO_TERMINATION_OR_STEPS: str +ERR_STR_WITHOUT_INSTANCE: Incomplete +ERR_STR_WITH_KEEPJOBFLOW_OR_TERMINATIONPROTECTED: Incomplete +ERR_STR_BOTH_OR_NONE_INSTANCEGROUPS_OR_INSTANCEFLEETS: Incomplete +ERR_STR_WITH_BOTH_CLUSTER_ID_AND_CLUSTER_CFG: str +ERR_STR_WITH_EXEC_ROLE_ARN_AND_WITHOUT_CLUSTER_ID: str +ERR_STR_WITHOUT_CLUSTER_ID_AND_CLUSTER_CFG: str + +class EMRStep(Step): + args: Incomplete + cache_config: Incomplete + def __init__( + self, + name: str, + display_name: str, + description: str, + cluster_id: str, + step_config: EMRStepConfig, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + cache_config: CacheConfig = None, + cluster_config: Dict[str, Any] = None, + execution_role_arn: str = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self) -> RequestType: ... + def to_request(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/entities.pyi b/stubs/sagemaker/sagemaker/workflow/entities.pyi new file mode 100644 index 000000000000..f8c6383b9b32 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/entities.pyi @@ -0,0 +1,30 @@ +import abc +from _typeshed import Incomplete +from enum import EnumMeta +from typing import Any, Dict, List, Union + +PrimitiveType = Union[str, int, bool, float, None] +RequestType = Union[Dict[str, Any], List[Dict[str, Any]]] + +class Entity(abc.ABC, metaclass=abc.ABCMeta): + @abc.abstractmethod + def to_request(self) -> RequestType: ... + +class DefaultEnumMeta(EnumMeta): + default: Incomplete + def __call__(cls, *args, value=..., **kwargs): ... + factory = __call__ + +class Expression(abc.ABC, metaclass=abc.ABCMeta): + @property + @abc.abstractmethod + def expr(self) -> RequestType: ... + +class PipelineVariable(Expression, metaclass=abc.ABCMeta): + def __add__(self, other: Union[Expression, PrimitiveType]): ... + def __int__(self) -> int: ... + def __float__(self) -> float: ... + def to_string(self): ... + @property + @abc.abstractmethod + def expr(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/execution_variables.pyi b/stubs/sagemaker/sagemaker/workflow/execution_variables.pyi new file mode 100644 index 000000000000..7208797c3bc3 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/execution_variables.pyi @@ -0,0 +1,21 @@ +from _typeshed import Incomplete + +from sagemaker.workflow.entities import PipelineVariable, RequestType as RequestType + +class ExecutionVariable(PipelineVariable): + name: Incomplete + def __init__(self, name: str) -> None: ... + def __eq__(self, other): ... + def to_string(self) -> PipelineVariable: ... + @property + def expr(self) -> RequestType: ... + +class ExecutionVariables: + START_DATETIME: Incomplete + CURRENT_DATETIME: Incomplete + PIPELINE_NAME: Incomplete + PIPELINE_ARN: Incomplete + PIPELINE_EXECUTION_ID: Incomplete + PIPELINE_EXECUTION_ARN: Incomplete + TRAINING_JOB_NAME: Incomplete + PROCESSING_JOB_NAME: Incomplete diff --git a/stubs/sagemaker/sagemaker/workflow/fail_step.pyi b/stubs/sagemaker/sagemaker/workflow/fail_step.pyi new file mode 100644 index 000000000000..2868b3cb71f6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/fail_step.pyi @@ -0,0 +1,21 @@ +from _typeshed import Incomplete +from typing import List, Optional, Union + +from sagemaker.workflow.entities import PipelineVariable, RequestType as RequestType +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.workflow.steps import Step + +class FailStep(Step): + error_message: Incomplete + def __init__( + self, + name: str, + error_message: Union[str, PipelineVariable] = None, + display_name: str = None, + description: str = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/functions.pyi b/stubs/sagemaker/sagemaker/workflow/functions.pyi new file mode 100644 index 000000000000..5313530ff8a5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/functions.pyi @@ -0,0 +1,28 @@ +from typing import List, Union + +from sagemaker.workflow.entities import PipelineVariable +from sagemaker.workflow.properties import PropertyFile + +class Join(PipelineVariable): + on: str + values: List + def to_string(self) -> PipelineVariable: ... + @property + def expr(self): ... + def __init__(self, on, values) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class JsonGet(PipelineVariable): + step_name: str + property_file: Union[PropertyFile, str] + json_path: str + @property + def expr(self): ... + def __init__(self, step_name, property_file, json_path) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi b/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi new file mode 100644 index 000000000000..4c39ca89b56d --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi @@ -0,0 +1,47 @@ +from _typeshed import Incomplete +from enum import Enum +from typing import Dict, List, Optional, Union + +from sagemaker.lambda_helper import Lambda +from sagemaker.workflow.entities import DefaultEnumMeta, RequestType as RequestType +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.workflow.steps import CacheConfig, Step + +class LambdaOutputTypeEnum(Enum, metaclass=DefaultEnumMeta): + String: str + Integer: str + Boolean: str + Float: str + +class LambdaOutput: + output_name: str + output_type: LambdaOutputTypeEnum + def to_request(self) -> RequestType: ... + def expr(self, step_name) -> Dict[str, str]: ... + def __init__(self, output_name, output_type) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class LambdaStep(Step): + lambda_func: Incomplete + outputs: Incomplete + cache_config: Incomplete + inputs: Incomplete + def __init__( + self, + name: str, + lambda_func: Lambda, + display_name: str = None, + description: str = None, + inputs: dict = None, + outputs: List[LambdaOutput] = None, + cache_config: CacheConfig = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self): ... + def to_request(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/model_step.pyi b/stubs/sagemaker/sagemaker/workflow/model_step.pyi new file mode 100644 index 000000000000..63302f38720e --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/model_step.pyi @@ -0,0 +1,25 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker.workflow.pipeline_context import _ModelStepArguments +from sagemaker.workflow.retry import RetryPolicy +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.workflow.steps import Step + +class ModelStep(StepCollection): + name: Incomplete + step_args: Incomplete + depends_on: Incomplete + retry_policies: Incomplete + display_name: Incomplete + description: Incomplete + steps: Incomplete + def __init__( + self, + name: str, + step_args: _ModelStepArguments, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + retry_policies: Optional[Union[List[RetryPolicy], Dict[str, List[RetryPolicy]]]] = None, + display_name: Optional[str] = None, + description: Optional[str] = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi b/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi new file mode 100644 index 000000000000..0bf48558b0cb --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi @@ -0,0 +1,26 @@ +from _typeshed import Incomplete +from typing import Optional, Union + +from sagemaker.workflow.check_job_config import CheckJobConfig +from sagemaker.workflow.clarify_check_step import ClarifyCheckConfig +from sagemaker.workflow.entities import PipelineVariable +from sagemaker.workflow.pipeline_context import _JobStepArguments +from sagemaker.workflow.quality_check_step import QualityCheckConfig +from sagemaker.workflow.step_collections import StepCollection + +class MonitorBatchTransformStep(StepCollection): + name: Incomplete + steps: Incomplete + def __init__( + self, + name: str, + transform_step_args: _JobStepArguments, + monitor_configuration: Union[QualityCheckConfig, ClarifyCheckConfig], + check_job_configuration: CheckJobConfig, + monitor_before_transform: bool = False, + fail_on_violation: Union[bool, PipelineVariable] = True, + supplied_baseline_statistics: Union[str, PipelineVariable] = None, + supplied_baseline_constraints: Union[str, PipelineVariable] = None, + display_name: Optional[str] = None, + description: Optional[str] = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/parallelism_config.pyi b/stubs/sagemaker/sagemaker/workflow/parallelism_config.pyi new file mode 100644 index 000000000000..7143ece59ecd --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/parallelism_config.pyi @@ -0,0 +1,8 @@ +from _typeshed import Incomplete + +from sagemaker.workflow.entities import RequestType as RequestType + +class ParallelismConfiguration: + max_parallel_execution_steps: Incomplete + def __init__(self, max_parallel_execution_steps: int) -> None: ... + def to_request(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/parameters.pyi b/stubs/sagemaker/sagemaker/workflow/parameters.pyi new file mode 100644 index 000000000000..c1ded205badc --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/parameters.pyi @@ -0,0 +1,47 @@ +from _typeshed import Incomplete +from enum import Enum +from typing import Dict, List, Type + +from sagemaker.workflow.entities import ( + DefaultEnumMeta, + Entity, + PipelineVariable, + PrimitiveType as PrimitiveType, + RequestType as RequestType, +) + +class ParameterTypeEnum(Enum, metaclass=DefaultEnumMeta): + STRING: str + INTEGER: str + BOOLEAN: str + FLOAT: str + @property + def python_type(self) -> Type: ... + +class Parameter(PipelineVariable, Entity): + name: str + parameter_type: ParameterTypeEnum + default_value: PrimitiveType + def to_request(self) -> RequestType: ... + @property + def expr(self) -> Dict[str, str]: ... + def __init__(self, name, parameter_type, default_value) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +ParameterBoolean: Incomplete + +class ParameterString(Parameter): + enum_values: Incomplete + def __init__(self, name: str, default_value: str = None, enum_values: List[str] = None) -> None: ... + def __hash__(self): ... + def to_string(self) -> PipelineVariable: ... + def to_request(self) -> RequestType: ... + +class ParameterInteger(Parameter): + def __init__(self, name: str, default_value: int = None) -> None: ... + +class ParameterFloat(Parameter): + def __init__(self, name: str, default_value: float = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/pipeline.pyi b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi new file mode 100644 index 000000000000..54d014fc080a --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi @@ -0,0 +1,91 @@ +from _typeshed import Incomplete +from typing import Any, Dict, List, Optional, Sequence, Set, Union + +from sagemaker.session import Session +from sagemaker.workflow.entities import Entity, RequestType as RequestType +from sagemaker.workflow.parallelism_config import ParallelismConfiguration +from sagemaker.workflow.parameters import Parameter +from sagemaker.workflow.pipeline_experiment_config import PipelineExperimentConfig +from sagemaker.workflow.selective_execution_config import SelectiveExecutionConfig +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.workflow.steps import Step + +logger: Incomplete + +class Pipeline(Entity): + name: Incomplete + parameters: Incomplete + pipeline_experiment_config: Incomplete + steps: Incomplete + sagemaker_session: Incomplete + def __init__( + self, + name: str = "", + parameters: Optional[Sequence[Parameter]] = None, + pipeline_experiment_config: Optional[PipelineExperimentConfig] = ..., + steps: Optional[Sequence[Union[Step, StepCollection]]] = None, + sagemaker_session: Optional[Session] = None, + ) -> None: ... + def to_request(self) -> RequestType: ... + def create( + self, + role_arn: str = None, + description: str = None, + tags: List[Dict[str, str]] = None, + parallelism_config: ParallelismConfiguration = None, + ) -> Dict[str, Any]: ... + def describe(self) -> Dict[str, Any]: ... + def update( + self, role_arn: str = None, description: str = None, parallelism_config: ParallelismConfiguration = None + ) -> Dict[str, Any]: ... + def upsert( + self, + role_arn: str = None, + description: str = None, + tags: List[Dict[str, str]] = None, + parallelism_config: ParallelismConfiguration = None, + ) -> Dict[str, Any]: ... + def delete(self) -> Dict[str, Any]: ... + def start( + self, + parameters: Dict[str, Union[str, bool, int, float]] = None, + execution_display_name: str = None, + execution_description: str = None, + parallelism_config: ParallelismConfiguration = None, + selective_execution_config: SelectiveExecutionConfig = None, + ): ... + def definition(self) -> str: ... + def list_executions( + self, sort_by: str = None, sort_order: str = None, max_results: int = None, next_token: str = None + ) -> Dict[str, Any]: ... + +def format_start_parameters(parameters: Dict[str, Any]) -> List[Dict[str, Any]]: ... +def interpolate( + request_obj: RequestType, callback_output_to_step_map: Dict[str, str], lambda_output_to_step_map: Dict[str, str] +) -> RequestType: ... +def update_args(args: Dict[str, Any], **kwargs): ... + +class _PipelineExecution: + arn: str + sagemaker_session: Session + def stop(self): ... + def describe(self): ... + def list_steps(self): ... + def wait(self, delay: int = 30, max_attempts: int = 60) -> None: ... + def __init__(self, arn, sagemaker_session) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class PipelineGraph: + step_map: Incomplete + adjacency_list: Incomplete + def __init__(self, steps: Sequence[Union[Step, StepCollection]]) -> None: ... + @classmethod + def from_pipeline(cls, pipeline: Pipeline): ... + def is_cyclic(self) -> bool: ... + def get_steps_in_sub_dag(self, current_step: Union[Step, StepCollection], sub_dag_steps: Set[str] = None) -> Set[str]: ... + stack: Incomplete + def __iter__(self): ... + def __next__(self) -> Step: ... diff --git a/stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi b/stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi new file mode 100644 index 000000000000..ed5fdb97b84e --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi @@ -0,0 +1,60 @@ +from _typeshed import Incomplete +from typing import Callable, Optional + +from sagemaker.local import LocalSession +from sagemaker.session import Session + +class _StepArguments: + caller_name: Incomplete + func: Incomplete + func_args: Incomplete + func_kwargs: Incomplete + def __init__(self, caller_name: str = None, func: Callable = None, *func_args, **func_kwargs) -> None: ... + +class _JobStepArguments(_StepArguments): + args: Incomplete + def __init__(self, caller_name: str, args: dict) -> None: ... + +class _ModelStepArguments(_StepArguments): + model: Incomplete + create_model_package_request: Incomplete + create_model_request: Incomplete + need_runtime_repack: Incomplete + runtime_repack_output_prefix: Incomplete + def __init__(self, model) -> None: ... + +class _PipelineConfig: + pipeline_name: Incomplete + step_name: Incomplete + code_hash: Incomplete + config_hash: Incomplete + def __init__(self, pipeline_name, step_name, code_hash, config_hash) -> None: ... + +class PipelineSession(Session): + def __init__( + self, + boto_session: Incomplete | None = None, + sagemaker_client: Incomplete | None = None, + default_bucket: Incomplete | None = None, + settings=..., + sagemaker_config: dict = None, + default_bucket_prefix: str = None, + ) -> None: ... + @property + def context(self): ... + @context.setter + def context(self, value: Optional[_StepArguments] = ...): ... + def init_model_step_arguments(self, model) -> None: ... + +class LocalPipelineSession(LocalSession, PipelineSession): + def __init__( + self, + boto_session: Incomplete | None = None, + default_bucket: Incomplete | None = None, + s3_endpoint_url: Incomplete | None = None, + disable_local_code: bool = False, + default_bucket_prefix: Incomplete | None = None, + ) -> None: ... + +def runnable_by_pipeline(run_func): ... +def retrieve_caller_name(job_instance): ... diff --git a/stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi b/stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi new file mode 100644 index 000000000000..8ea15168b974 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi @@ -0,0 +1,26 @@ +from _typeshed import Incomplete +from typing import Union + +from sagemaker.workflow.entities import Entity, Expression, RequestType as RequestType +from sagemaker.workflow.execution_variables import ExecutionVariable +from sagemaker.workflow.parameters import Parameter + +class PipelineExperimentConfig(Entity): + experiment_name: Incomplete + trial_name: Incomplete + def __init__( + self, + experiment_name: Union[str, Parameter, ExecutionVariable, Expression], + trial_name: Union[str, Parameter, ExecutionVariable, Expression], + ) -> None: ... + def to_request(self) -> RequestType: ... + +class PipelineExperimentConfigProperty(Expression): + name: Incomplete + def __init__(self, name: str) -> None: ... + @property + def expr(self) -> RequestType: ... + +class PipelineExperimentConfigProperties: + EXPERIMENT_NAME: Incomplete + TRIAL_NAME: Incomplete diff --git a/stubs/sagemaker/sagemaker/workflow/properties.pyi b/stubs/sagemaker/sagemaker/workflow/properties.pyi new file mode 100644 index 000000000000..47b3f72c9ae8 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/properties.pyi @@ -0,0 +1,46 @@ +from _typeshed import Incomplete +from abc import ABCMeta +from typing import Dict, List, Union + +from sagemaker.workflow.entities import Expression, PipelineVariable + +class PropertiesMeta(ABCMeta): + def __new__(mcs, *args, **kwargs): ... + +class Properties(PipelineVariable, metaclass=PropertiesMeta): + step_name: Incomplete + path: Incomplete + def __init__( + self, + step_name: str, + path: str = None, + shape_name: str = None, + shape_names: List[str] = None, + service_name: str = "sagemaker", + ) -> None: ... + @property + def expr(self): ... + +class PropertiesList(Properties): + shape_name: Incomplete + service_name: Incomplete + def __init__(self, step_name: str, path: str, shape_name: str = None, service_name: str = "sagemaker") -> None: ... + def __getitem__(self, item: Union[int, str]): ... + +class PropertiesMap(Properties): + shape_name: Incomplete + service_name: Incomplete + def __init__(self, step_name: str, path: str, shape_name: str = None, service_name: str = "sagemaker") -> None: ... + def __getitem__(self, item: Union[int, str]): ... + +class PropertyFile(Expression): + name: str + output_name: str + path: str + @property + def expr(self) -> Dict[str, str]: ... + def __init__(self, name, output_name, path) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi b/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi new file mode 100644 index 000000000000..aeaade118d81 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi @@ -0,0 +1,84 @@ +from _typeshed import Incomplete +from abc import ABC +from typing import List, Optional, Union + +from sagemaker.workflow.check_job_config import CheckJobConfig +from sagemaker.workflow.entities import PipelineVariable, RequestType as RequestType +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.workflow.steps import CacheConfig, Step + +class QualityCheckConfig(ABC): + baseline_dataset: Union[str, PipelineVariable] + dataset_format: dict + output_s3_uri: Union[str, PipelineVariable] + post_analytics_processor_script: str + def __init__(self, baseline_dataset, dataset_format, output_s3_uri, post_analytics_processor_script) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class DataQualityCheckConfig(QualityCheckConfig): + record_preprocessor_script: str + def __init__( + self, baseline_dataset, dataset_format, output_s3_uri, post_analytics_processor_script, record_preprocessor_script + ) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class ModelQualityCheckConfig(QualityCheckConfig): + problem_type: Union[str, PipelineVariable] + inference_attribute: Union[str, PipelineVariable] + probability_attribute: Union[str, PipelineVariable] + ground_truth_attribute: Union[str, PipelineVariable] + probability_threshold_attribute: Union[str, PipelineVariable] + def __init__( + self, + baseline_dataset, + dataset_format, + output_s3_uri, + post_analytics_processor_script, + problem_type, + inference_attribute, + probability_attribute, + ground_truth_attribute, + probability_threshold_attribute, + ) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class QualityCheckStep(Step): + skip_check: Incomplete + fail_on_violation: Incomplete + register_new_baseline: Incomplete + check_job_config: Incomplete + quality_check_config: Incomplete + model_package_group_name: Incomplete + supplied_baseline_statistics: Incomplete + supplied_baseline_constraints: Incomplete + cache_config: Incomplete + def __init__( + self, + name: str, + quality_check_config: QualityCheckConfig, + check_job_config: CheckJobConfig, + skip_check: Union[bool, PipelineVariable] = False, + fail_on_violation: Union[bool, PipelineVariable] = True, + register_new_baseline: Union[bool, PipelineVariable] = False, + model_package_group_name: Union[str, PipelineVariable] = None, + supplied_baseline_statistics: Union[str, PipelineVariable] = None, + supplied_baseline_constraints: Union[str, PipelineVariable] = None, + display_name: str = None, + description: str = None, + cache_config: CacheConfig = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self): ... + def to_request(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/retry.pyi b/stubs/sagemaker/sagemaker/workflow/retry.pyi new file mode 100644 index 000000000000..18c7335be39f --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/retry.pyi @@ -0,0 +1,62 @@ +from _typeshed import Incomplete +from enum import Enum +from typing import List + +from sagemaker.workflow.entities import DefaultEnumMeta, Entity, RequestType as RequestType + +DEFAULT_BACKOFF_RATE: float +DEFAULT_INTERVAL_SECONDS: int +MAX_ATTEMPTS_CAP: int +MAX_EXPIRE_AFTER_MIN: int + +class StepExceptionTypeEnum(Enum, metaclass=DefaultEnumMeta): + SERVICE_FAULT: str + THROTTLING: str + +class SageMakerJobExceptionTypeEnum(Enum, metaclass=DefaultEnumMeta): + INTERNAL_ERROR: str + CAPACITY_ERROR: str + RESOURCE_LIMIT: str + +class RetryPolicy(Entity): + backoff_rate: float + interval_seconds: int + max_attempts: int + expire_after_mins: int + def validate_backoff_rate(self, _, value) -> None: ... + def validate_interval_seconds(self, _, value) -> None: ... + def validate_max_attempts(self, _, value) -> None: ... + def validate_expire_after_mins(self, _, value) -> None: ... + def to_request(self) -> RequestType: ... + def __init__(self, backoff_rate, interval_seconds, max_attempts, expire_after_mins) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class StepRetryPolicy(RetryPolicy): + exception_types: Incomplete + def __init__( + self, + exception_types: List[StepExceptionTypeEnum], + backoff_rate: float = 2.0, + interval_seconds: int = 1, + max_attempts: int = None, + expire_after_mins: int = None, + ) -> None: ... + def to_request(self) -> RequestType: ... + def __hash__(self): ... + +class SageMakerJobStepRetryPolicy(RetryPolicy): + exception_type_list: Incomplete + def __init__( + self, + exception_types: List[SageMakerJobExceptionTypeEnum] = None, + failure_reason_types: List[SageMakerJobExceptionTypeEnum] = None, + backoff_rate: float = 2.0, + interval_seconds: int = 1, + max_attempts: int = None, + expire_after_mins: int = None, + ) -> None: ... + def to_request(self) -> RequestType: ... + def __hash__(self): ... diff --git a/stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi b/stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi new file mode 100644 index 000000000000..f28ce9bdc5e2 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi @@ -0,0 +1,10 @@ +from _typeshed import Incomplete +from typing import List + +from sagemaker.workflow.entities import RequestType as RequestType + +class SelectiveExecutionConfig: + source_pipeline_execution_arn: Incomplete + selected_steps: Incomplete + def __init__(self, selected_steps: List[str], source_pipeline_execution_arn: str = None) -> None: ... + def to_request(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/step_collections.pyi b/stubs/sagemaker/sagemaker/workflow/step_collections.pyi new file mode 100644 index 000000000000..29964ced9093 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/step_collections.pyi @@ -0,0 +1,92 @@ +from _typeshed import Incomplete +from typing import List, Optional, Union + +from sagemaker import PipelineModel +from sagemaker.estimator import EstimatorBase +from sagemaker.model import Model +from sagemaker.workflow.entities import RequestType as RequestType +from sagemaker.workflow.retry import RetryPolicy +from sagemaker.workflow.steps import Step + +class StepCollection: + name: str + steps: List[Step] + def request_dicts(self) -> List[RequestType]: ... + @property + def properties(self): ... + def __init__(self, name, steps) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class RegisterModel(StepCollection): + name: Incomplete + model_list: Incomplete + container_def_list: Incomplete + steps: Incomplete + def __init__( + self, + name: str, + content_types, + response_types, + inference_instances: Incomplete | None = None, + transform_instances: Incomplete | None = None, + estimator: EstimatorBase = None, + model_data: Incomplete | None = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + repack_model_step_retry_policies: List[RetryPolicy] = None, + register_model_step_retry_policies: List[RetryPolicy] = None, + model_package_group_name: Incomplete | None = None, + model_metrics: Incomplete | None = None, + approval_status: Incomplete | None = None, + image_uri: Incomplete | None = None, + compile_model_family: Incomplete | None = None, + display_name: Incomplete | None = None, + description: Incomplete | None = None, + tags: Incomplete | None = None, + model: Union[Model, PipelineModel] = None, + drift_check_baselines: Incomplete | None = None, + customer_metadata_properties: Incomplete | None = None, + domain: Incomplete | None = None, + sample_payload_url: Incomplete | None = None, + task: Incomplete | None = None, + framework: Incomplete | None = None, + framework_version: Incomplete | None = None, + nearest_model_name: Incomplete | None = None, + data_input_configuration: Incomplete | None = None, + **kwargs, + ) -> None: ... + +class EstimatorTransformer(StepCollection): + name: Incomplete + steps: Incomplete + def __init__( + self, + name: str, + estimator: EstimatorBase, + model_data, + model_inputs, + instance_count, + instance_type, + transform_inputs, + description: str = None, + display_name: str = None, + image_uri: Incomplete | None = None, + predictor_cls: Incomplete | None = None, + env: Incomplete | None = None, + strategy: Incomplete | None = None, + assemble_with: Incomplete | None = None, + output_path: Incomplete | None = None, + output_kms_key: Incomplete | None = None, + accept: Incomplete | None = None, + max_concurrent_transforms: Incomplete | None = None, + max_payload: Incomplete | None = None, + tags: Incomplete | None = None, + volume_kms_key: Incomplete | None = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + repack_model_step_retry_policies: List[RetryPolicy] = None, + model_step_retry_policies: List[RetryPolicy] = None, + transform_step_retry_policies: List[RetryPolicy] = None, + **kwargs, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/steps.pyi b/stubs/sagemaker/sagemaker/workflow/steps.pyi new file mode 100644 index 000000000000..c4622be62499 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/steps.pyi @@ -0,0 +1,210 @@ +import abc +from _typeshed import Incomplete +from enum import Enum +from typing import Dict, List, Optional, Union + +from sagemaker.estimator import EstimatorBase +from sagemaker.inputs import CreateModelInput, FileSystemInput, TrainingInput, TransformInput +from sagemaker.model import Model +from sagemaker.pipeline import PipelineModel +from sagemaker.processing import ProcessingInput, ProcessingOutput, Processor +from sagemaker.transformer import Transformer +from sagemaker.tuner import HyperparameterTuner +from sagemaker.workflow.entities import DefaultEnumMeta, Entity, RequestType as RequestType +from sagemaker.workflow.functions import Join +from sagemaker.workflow.pipeline_context import _JobStepArguments +from sagemaker.workflow.properties import PropertyFile +from sagemaker.workflow.retry import RetryPolicy +from sagemaker.workflow.step_collections import StepCollection + +class StepTypeEnum(Enum, metaclass=DefaultEnumMeta): + CONDITION: str + CREATE_MODEL: str + PROCESSING: str + REGISTER_MODEL: str + TRAINING: str + TRANSFORM: str + CALLBACK: str + TUNING: str + LAMBDA: str + QUALITY_CHECK: str + CLARIFY_CHECK: str + EMR: str + FAIL: str + AUTOML: str + +class Step(Entity, metaclass=abc.ABCMeta): + name: str + display_name: Optional[str] + description: Optional[str] + step_type: StepTypeEnum + depends_on: Optional[List[Union[str, "Step", "StepCollection"]]] + @property + @abc.abstractmethod + def arguments(self) -> RequestType: ... + @property + def step_only_arguments(self) -> RequestType: ... + @property + @abc.abstractmethod + def properties(self): ... + def to_request(self) -> RequestType: ... + def add_depends_on(self, step_names: List[Union[str, "Step", "StepCollection"]]): ... + @property + def ref(self) -> Dict[str, str]: ... + def __init__(self, name, display_name, description, step_type, depends_on) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class CacheConfig: + enable_caching: bool + expire_after: Incomplete + @property + def config(self): ... + def __init__(self, enable_caching, expire_after) -> None: ... + def __lt__(self, other): ... + def __le__(self, other): ... + def __gt__(self, other): ... + def __ge__(self, other): ... + +class ConfigurableRetryStep(Step, metaclass=abc.ABCMeta): + retry_policies: Incomplete + def __init__( + self, + name: str, + step_type: StepTypeEnum, + display_name: str = None, + description: str = None, + depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + retry_policies: List[RetryPolicy] = None, + ) -> None: ... + def add_retry_policy(self, retry_policy: RetryPolicy): ... + def to_request(self) -> RequestType: ... + +class TrainingStep(ConfigurableRetryStep): + step_args: Incomplete + estimator: Incomplete + inputs: Incomplete + cache_config: Incomplete + job_name: Incomplete + def __init__( + self, + name: str, + step_args: _JobStepArguments = None, + estimator: EstimatorBase = None, + display_name: str = None, + description: str = None, + inputs: Union[TrainingInput, dict, str, FileSystemInput] = None, + cache_config: CacheConfig = None, + depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + retry_policies: List[RetryPolicy] = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self): ... + def to_request(self) -> RequestType: ... + +class CreateModelStep(ConfigurableRetryStep): + step_args: Incomplete + model: Incomplete + inputs: Incomplete + def __init__( + self, + name: str, + step_args: Optional[dict] = None, + model: Optional[Union[Model, PipelineModel]] = None, + inputs: Optional[CreateModelInput] = None, + depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + retry_policies: Optional[List[RetryPolicy]] = None, + display_name: Optional[str] = None, + description: Optional[str] = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self): ... + +class TransformStep(ConfigurableRetryStep): + step_args: Incomplete + transformer: Incomplete + inputs: Incomplete + cache_config: Incomplete + def __init__( + self, + name: str, + step_args: _JobStepArguments = None, + transformer: Transformer = None, + inputs: TransformInput = None, + display_name: str = None, + description: str = None, + cache_config: CacheConfig = None, + depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + retry_policies: List[RetryPolicy] = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self): ... + def to_request(self) -> RequestType: ... + +class ProcessingStep(ConfigurableRetryStep): + step_args: Incomplete + processor: Incomplete + inputs: Incomplete + outputs: Incomplete + job_arguments: Incomplete + code: Incomplete + property_files: Incomplete + job_name: Incomplete + kms_key: Incomplete + cache_config: Incomplete + def __init__( + self, + name: str, + step_args: _JobStepArguments = None, + processor: Processor = None, + display_name: str = None, + description: str = None, + inputs: List[ProcessingInput] = None, + outputs: List[ProcessingOutput] = None, + job_arguments: List[str] = None, + code: str = None, + property_files: List[PropertyFile] = None, + cache_config: CacheConfig = None, + depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + retry_policies: List[RetryPolicy] = None, + kms_key: Incomplete | None = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self): ... + def to_request(self) -> RequestType: ... + +class TuningStep(ConfigurableRetryStep): + step_args: Incomplete + tuner: Incomplete + inputs: Incomplete + job_arguments: Incomplete + cache_config: Incomplete + def __init__( + self, + name: str, + step_args: _JobStepArguments = None, + tuner: HyperparameterTuner = None, + display_name: str = None, + description: str = None, + inputs: Incomplete | None = None, + job_arguments: List[str] = None, + cache_config: CacheConfig = None, + depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + retry_policies: List[RetryPolicy] = None, + ) -> None: ... + @property + def arguments(self) -> RequestType: ... + @property + def properties(self): ... + def to_request(self) -> RequestType: ... + def get_top_model_s3_uri(self, top_k: int, s3_bucket: str, prefix: str = "") -> Join: ... diff --git a/stubs/sagemaker/sagemaker/workflow/utilities.pyi b/stubs/sagemaker/sagemaker/workflow/utilities.pyi new file mode 100644 index 000000000000..89acb50bc08f --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/utilities.pyi @@ -0,0 +1,23 @@ +from _typeshed import Incomplete +from typing import List, Sequence, Set, Union + +from sagemaker.workflow.entities import Entity, RequestType as RequestType +from sagemaker.workflow.pipeline_context import _StepArguments +from sagemaker.workflow.step_collections import StepCollection + +logger: Incomplete +BUF_SIZE: int + +def list_to_request(entities: Sequence[Union[Entity, "StepCollection"]]) -> List[RequestType]: ... +def build_steps(steps: Sequence[Entity], pipeline_name: str): ... +def get_code_hash(step: Entity) -> str: ... +def get_processing_dependencies(dependency_args: List[List[str]]) -> List[str]: ... +def get_processing_code_hash(code: str, source_dir: str, dependencies: List[str]) -> str: ... +def get_training_code_hash(entry_point: str, source_dir: str, dependencies: List[str]) -> str: ... +def get_config_hash(step: Entity): ... +def hash_object(obj) -> str: ... +def hash_file(path: str) -> str: ... +def hash_files_or_dirs(paths: List[str]) -> str: ... +def validate_step_args_input(step_args: _StepArguments, expected_caller: Set[str], error_message: str): ... +def override_pipeline_parameter_var(func): ... +def execute_job_functions(step_args: _StepArguments): ... diff --git a/stubs/sagemaker/sagemaker/wrangler/__init__.pyi b/stubs/sagemaker/sagemaker/wrangler/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/stubs/sagemaker/sagemaker/wrangler/ingestion.pyi b/stubs/sagemaker/sagemaker/wrangler/ingestion.pyi new file mode 100644 index 000000000000..796076df0ac0 --- /dev/null +++ b/stubs/sagemaker/sagemaker/wrangler/ingestion.pyi @@ -0,0 +1,18 @@ +from typing import Dict + +from sagemaker.dataset_definition.inputs import AthenaDatasetDefinition, RedshiftDatasetDefinition + +def generate_data_ingestion_flow_from_s3_input( + input_name: str, + s3_uri: str, + s3_content_type: str = "csv", + s3_has_header: bool = False, + operator_version: str = "0.1", + schema: Dict = None, +): ... +def generate_data_ingestion_flow_from_athena_dataset_definition( + input_name: str, athena_dataset_definition: AthenaDatasetDefinition, operator_version: str = "0.1", schema: Dict = None +): ... +def generate_data_ingestion_flow_from_redshift_dataset_definition( + input_name: str, redshift_dataset_definition: RedshiftDatasetDefinition, operator_version: str = "0.1", schema: Dict = None +): ... diff --git a/stubs/sagemaker/sagemaker/wrangler/processing.pyi b/stubs/sagemaker/sagemaker/wrangler/processing.pyi new file mode 100644 index 000000000000..de9051a81302 --- /dev/null +++ b/stubs/sagemaker/sagemaker/wrangler/processing.pyi @@ -0,0 +1,26 @@ +from _typeshed import Incomplete +from typing import Dict, List + +from sagemaker.network import NetworkConfig +from sagemaker.processing import Processor +from sagemaker.session import Session + +class DataWranglerProcessor(Processor): + data_wrangler_flow_source: Incomplete + sagemaker_session: Incomplete + def __init__( + self, + role: str = None, + data_wrangler_flow_source: str = None, + instance_count: int = None, + instance_type: str = None, + volume_size_in_gb: int = 30, + volume_kms_key: str = None, + output_kms_key: str = None, + max_runtime_in_seconds: int = None, + base_job_name: str = None, + sagemaker_session: Session = None, + env: Dict[str, str] = None, + tags: List[dict] = None, + network_config: NetworkConfig = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/xgboost/__init__.pyi b/stubs/sagemaker/sagemaker/xgboost/__init__.pyi new file mode 100644 index 000000000000..2cf2154cc4f6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/xgboost/__init__.pyi @@ -0,0 +1,4 @@ +from sagemaker.xgboost.defaults import XGBOOST_NAME as XGBOOST_NAME +from sagemaker.xgboost.estimator import XGBoost as XGBoost +from sagemaker.xgboost.model import XGBoostModel as XGBoostModel, XGBoostPredictor as XGBoostPredictor +from sagemaker.xgboost.processing import XGBoostProcessor as XGBoostProcessor diff --git a/stubs/sagemaker/sagemaker/xgboost/defaults.pyi b/stubs/sagemaker/sagemaker/xgboost/defaults.pyi new file mode 100644 index 000000000000..1fbd9555e5b5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/xgboost/defaults.pyi @@ -0,0 +1,4 @@ +from _typeshed import Incomplete + +XGBOOST_NAME: str +XGBOOST_UNSUPPORTED_VERSIONS: Incomplete diff --git a/stubs/sagemaker/sagemaker/xgboost/estimator.pyi b/stubs/sagemaker/sagemaker/xgboost/estimator.pyi new file mode 100644 index 000000000000..21c50a492253 --- /dev/null +++ b/stubs/sagemaker/sagemaker/xgboost/estimator.pyi @@ -0,0 +1,35 @@ +from _typeshed import Incomplete +from typing import Dict, Optional, Union + +from sagemaker.estimator import Framework +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class XGBoost(Framework): + py_version: Incomplete + framework_version: Incomplete + image_uri: Incomplete + def __init__( + self, + entry_point: Union[str, PipelineVariable], + framework_version: str, + source_dir: Optional[Union[str, PipelineVariable]] = None, + hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + py_version: str = "py3", + image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri_region: Optional[str] = None, + **kwargs, + ) -> None: ... + def create_model( + self, + model_server_workers: Incomplete | None = None, + role: Incomplete | None = None, + vpc_config_override="VPC_CONFIG_DEFAULT", + entry_point: Incomplete | None = None, + source_dir: Incomplete | None = None, + dependencies: Incomplete | None = None, + **kwargs, + ): ... + @classmethod + def attach(cls, training_job_name, sagemaker_session: Incomplete | None = None, model_channel_name: str = "model"): ... diff --git a/stubs/sagemaker/sagemaker/xgboost/model.pyi b/stubs/sagemaker/sagemaker/xgboost/model.pyi new file mode 100644 index 000000000000..b4faae2961db --- /dev/null +++ b/stubs/sagemaker/sagemaker/xgboost/model.pyi @@ -0,0 +1,63 @@ +from _typeshed import Incomplete +from typing import Dict, List, Optional, Union + +from sagemaker import ModelMetrics +from sagemaker.drift_check_baselines import DriftCheckBaselines +from sagemaker.metadata_properties import MetadataProperties +from sagemaker.model import FrameworkModel +from sagemaker.predictor import Predictor +from sagemaker.workflow.entities import PipelineVariable + +logger: Incomplete + +class XGBoostPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class XGBoostModel(FrameworkModel): + py_version: Incomplete + framework_version: Incomplete + model_server_workers: Incomplete + def __init__( + self, + model_data: Union[str, PipelineVariable], + role: str = None, + entry_point: str = None, + framework_version: str = None, + image_uri: Optional[Union[str, PipelineVariable]] = None, + py_version: str = "py3", + predictor_cls: callable = ..., + model_server_workers: Optional[Union[int, PipelineVariable]] = None, + **kwargs, + ) -> None: ... + image_uri: Incomplete + def register( + self, + 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, + domain: Optional[Union[str, PipelineVariable]] = None, + sample_payload_url: Optional[Union[str, PipelineVariable]] = None, + task: Optional[Union[str, PipelineVariable]] = None, + framework: Optional[Union[str, PipelineVariable]] = None, + framework_version: Optional[Union[str, PipelineVariable]] = None, + nearest_model_name: Optional[Union[str, PipelineVariable]] = None, + data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + ): ... + def prepare_container_def( + self, + instance_type: Incomplete | None = None, + accelerator_type: Incomplete | None = None, + serverless_inference_config: Incomplete | None = None, + ): ... + def serving_image_uri(self, region_name, instance_type, serverless_inference_config: Incomplete | None = None): ... diff --git a/stubs/sagemaker/sagemaker/xgboost/processing.pyi b/stubs/sagemaker/sagemaker/xgboost/processing.pyi new file mode 100644 index 000000000000..2985764e42c5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/xgboost/processing.pyi @@ -0,0 +1,30 @@ +from typing import Dict, List, Optional, Union + +from sagemaker.network import NetworkConfig +from sagemaker.processing import FrameworkProcessor +from sagemaker.session import Session +from sagemaker.workflow.entities import PipelineVariable +from sagemaker.xgboost.estimator import XGBoost + +class XGBoostProcessor(FrameworkProcessor): + estimator_cls = XGBoost + def __init__( + self, + framework_version: str, + role: str = None, + instance_count: Union[int, PipelineVariable] = None, + instance_type: Union[str, PipelineVariable] = None, + py_version: str = "py3", + image_uri: Optional[Union[str, PipelineVariable]] = None, + command: Optional[List[str]] = None, + volume_size_in_gb: Union[int, PipelineVariable] = 30, + volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + output_kms_key: Optional[Union[str, PipelineVariable]] = None, + code_location: Optional[str] = None, + max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + base_job_name: Optional[str] = None, + sagemaker_session: Optional[Session] = None, + env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + network_config: Optional[NetworkConfig] = None, + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/xgboost/utils.pyi b/stubs/sagemaker/sagemaker/xgboost/utils.pyi new file mode 100644 index 000000000000..a026e3c36777 --- /dev/null +++ b/stubs/sagemaker/sagemaker/xgboost/utils.pyi @@ -0,0 +1,2 @@ +def validate_py_version(py_version) -> None: ... +def validate_framework_version(framework_version) -> None: ... From 45c2b1538488c2ced747b6ffc64bbc9355611a80 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 23 Jun 2023 15:26:36 +0000 Subject: [PATCH 02/10] [pre-commit.ci] auto fixes from pre-commit.com hooks --- stubs/sagemaker/METADATA.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stubs/sagemaker/METADATA.toml b/stubs/sagemaker/METADATA.toml index 4b4e7f8b507a..4398389b5581 100644 --- a/stubs/sagemaker/METADATA.toml +++ b/stubs/sagemaker/METADATA.toml @@ -1 +1 @@ -version = "2.168.*" \ No newline at end of file +version = "2.168.*" From 0ade8cd5585b7614956f8834d41e5a1937c2849f Mon Sep 17 00:00:00 2001 From: DemyCode Date: Fri, 23 Jun 2023 18:34:37 +0200 Subject: [PATCH 03/10] fix: fixing Union to PEP 604 --- stubs/sagemaker/sagemaker/algorithm.pyi | 36 ++--- .../sagemaker/amazon/amazon_estimator.pyi | 14 +- .../amazon/factorization_machines.pyi | 8 +- .../sagemaker/sagemaker/amazon/ipinsights.pyi | 8 +- stubs/sagemaker/sagemaker/amazon/kmeans.pyi | 10 +- stubs/sagemaker/sagemaker/amazon/knn.pyi | 8 +- stubs/sagemaker/sagemaker/amazon/lda.pyi | 6 +- .../sagemaker/amazon/linear_learner.pyi | 8 +- stubs/sagemaker/sagemaker/amazon/ntm.pyi | 8 +- .../sagemaker/sagemaker/amazon/object2vec.pyi | 8 +- stubs/sagemaker/sagemaker/amazon/pca.pyi | 8 +- .../sagemaker/amazon/randomcutforest.pyi | 8 +- .../sagemaker/sagemaker/chainer/estimator.pyi | 16 +- stubs/sagemaker/sagemaker/chainer/model.pyi | 38 ++--- stubs/sagemaker/sagemaker/clarify.pyi | 52 +++---- .../sagemaker/sagemaker/debugger/debugger.pyi | 28 ++-- .../sagemaker/debugger/profiler_config.pyi | 6 +- .../sagemaker/djl_inference/model.pyi | 4 +- stubs/sagemaker/sagemaker/estimator.pyi | 144 +++++++++--------- stubs/sagemaker/sagemaker/experiments/run.pyi | 14 +- .../sagemaker/feature_store/feature_group.pyi | 4 +- .../sagemaker/feature_store/feature_store.pyi | 2 +- stubs/sagemaker/sagemaker/fw_utils.pyi | 6 +- .../sagemaker/huggingface/estimator.pyi | 8 +- .../sagemaker/sagemaker/huggingface/model.pyi | 38 ++--- .../sagemaker/huggingface/processing.pyi | 20 +-- .../huggingface/training_compiler/config.pyi | 2 +- stubs/sagemaker/sagemaker/inputs.pyi | 20 +-- .../sagemaker/jumpstart/estimator.pyi | 76 ++++----- .../sagemaker/jumpstart/factory/estimator.pyi | 76 ++++----- .../sagemaker/jumpstart/factory/model.pyi | 14 +- .../sagemaker/sagemaker/jumpstart/filters.pyi | 10 +- stubs/sagemaker/sagemaker/jumpstart/model.pyi | 14 +- .../sagemaker/jumpstart/notebook_utils.pyi | 10 +- stubs/sagemaker/sagemaker/jumpstart/types.pyi | 92 +++++------ stubs/sagemaker/sagemaker/lineage/query.pyi | 4 +- .../sagemaker/metadata_properties.pyi | 8 +- stubs/sagemaker/sagemaker/model.pyi | 58 +++---- .../sagemaker/model_card/helpers.pyi | 6 +- .../sagemaker/model_card/model_card.pyi | 22 +-- stubs/sagemaker/sagemaker/model_metrics.pyi | 12 +- .../model_monitor/model_monitoring.pyi | 2 +- stubs/sagemaker/sagemaker/multidatamodel.pyi | 2 +- stubs/sagemaker/sagemaker/mxnet/estimator.pyi | 8 +- stubs/sagemaker/sagemaker/mxnet/model.pyi | 38 ++--- .../sagemaker/sagemaker/mxnet/processing.pyi | 20 +-- stubs/sagemaker/sagemaker/network.pyi | 8 +- stubs/sagemaker/sagemaker/parameter.pyi | 6 +- stubs/sagemaker/sagemaker/pipeline.pyi | 36 ++--- stubs/sagemaker/sagemaker/processing.pyi | 92 +++++------ .../sagemaker/sagemaker/pytorch/estimator.pyi | 8 +- stubs/sagemaker/sagemaker/pytorch/model.pyi | 38 ++--- .../sagemaker/pytorch/processing.pyi | 20 +-- .../pytorch/training_compiler/config.pyi | 2 +- .../remote_function/spark_config.pyi | 2 +- stubs/sagemaker/sagemaker/rl/estimator.pyi | 10 +- stubs/sagemaker/sagemaker/s3.pyi | 2 +- .../sagemaker/sagemaker/sklearn/estimator.pyi | 8 +- stubs/sagemaker/sagemaker/sklearn/model.pyi | 38 ++--- .../sagemaker/sklearn/processing.pyi | 18 +-- .../sagemaker/sagemaker/spark/processing.pyi | 60 ++++---- stubs/sagemaker/sagemaker/sparkml/model.pyi | 2 +- .../sagemaker/tensorflow/estimator.pyi | 4 +- .../sagemaker/sagemaker/tensorflow/model.pyi | 36 ++--- .../sagemaker/tensorflow/processing.pyi | 20 +-- stubs/sagemaker/sagemaker/transformer.pyi | 44 +++--- stubs/sagemaker/sagemaker/tuner.pyi | 32 ++-- stubs/sagemaker/sagemaker/workflow/_utils.pyi | 4 +- .../sagemaker/workflow/automl_step.pyi | 2 +- .../sagemaker/workflow/callback_step.pyi | 2 +- .../sagemaker/workflow/clarify_check_step.pyi | 18 +-- .../sagemaker/workflow/condition_step.pyi | 8 +- .../sagemaker/workflow/conditions.pyi | 22 +-- .../sagemaker/sagemaker/workflow/emr_step.pyi | 2 +- .../sagemaker/sagemaker/workflow/entities.pyi | 6 +- .../sagemaker/workflow/fail_step.pyi | 4 +- .../sagemaker/workflow/functions.pyi | 2 +- .../sagemaker/workflow/lambda_step.pyi | 2 +- .../sagemaker/workflow/model_step.pyi | 4 +- .../workflow/monitor_batch_transform_step.pyi | 8 +- .../sagemaker/sagemaker/workflow/pipeline.pyi | 8 +- .../workflow/pipeline_experiment_config.pyi | 4 +- .../sagemaker/workflow/properties.pyi | 4 +- .../sagemaker/workflow/quality_check_step.pyi | 28 ++-- .../sagemaker/workflow/step_collections.pyi | 6 +- stubs/sagemaker/sagemaker/workflow/steps.pyi | 20 +-- .../sagemaker/workflow/utilities.pyi | 2 +- .../sagemaker/sagemaker/xgboost/estimator.pyi | 8 +- stubs/sagemaker/sagemaker/xgboost/model.pyi | 38 ++--- .../sagemaker/xgboost/processing.pyi | 18 +-- 90 files changed, 855 insertions(+), 855 deletions(-) diff --git a/stubs/sagemaker/sagemaker/algorithm.pyi b/stubs/sagemaker/sagemaker/algorithm.pyi index 602bbf46bd23..9b4a6ae90e8f 100644 --- a/stubs/sagemaker/sagemaker/algorithm.pyi +++ b/stubs/sagemaker/sagemaker/algorithm.pyi @@ -14,26 +14,26 @@ class AlgorithmEstimator(EstimatorBase): self, algorithm_arn: str, role: str = None, - 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] = 86400, - input_mode: Union[str, PipelineVariable] = "File", - output_path: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, + volume_size: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + max_run: int | PipelineVariable = 86400, + input_mode: str | PipelineVariable = "File", + output_path: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = 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, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, + subnets: Optional[List[str | PipelineVariable]] = None, + security_group_ids: Optional[List[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, + model_channel_name: str | PipelineVariable = "model", + metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, + encrypt_inter_container_traffic: bool | PipelineVariable = False, + use_spot_instances: bool | PipelineVariable = False, + max_wait: Optional[int | PipelineVariable] = None, **kwargs, ) -> None: ... def validate_train_spec(self) -> None: ... @@ -68,7 +68,7 @@ class AlgorithmEstimator(EstimatorBase): ): ... def fit( self, - inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None, + inputs: Optional[str | Dict | TrainingInput | FileSystemInput] = None, wait: bool = True, logs: bool = True, job_name: Optional[str] = None, diff --git a/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi b/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi index 77fa3bcf2c54..6fe1300e7da9 100644 --- a/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi +++ b/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi @@ -16,11 +16,11 @@ class AmazonAlgorithmEstimatorBase(EstimatorBase, metaclass=abc.ABCMeta): DEFAULT_MINI_BATCH_SIZE: Optional[int] def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, data_location: Optional[str] = None, - enable_network_isolation: Union[bool, PipelineVariable] = False, + enable_network_isolation: bool | PipelineVariable = False, **kwargs, ) -> None: ... def training_image_uri(self): ... @@ -52,11 +52,11 @@ class RecordSet: channel: Incomplete def __init__( self, - s3_data: Union[str, PipelineVariable], + s3_data: str | PipelineVariable, num_records: int, feature_dim: int, - s3_data_type: Union[str, PipelineVariable] = "ManifestFile", - channel: Union[str, PipelineVariable] = "train", + s3_data_type: str | PipelineVariable = "ManifestFile", + channel: str | PipelineVariable = "train", ) -> None: ... def data_channel(self): ... def records_s3_input(self): ... diff --git a/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi b/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi index 3d70a2034c8b..5c04f9333064 100644 --- a/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi +++ b/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi @@ -37,9 +37,9 @@ class FactorizationMachines(AmazonAlgorithmEstimatorBase): factors_init_value: hp def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, num_factors: Optional[int] = None, predictor_type: Optional[str] = None, epochs: Optional[int] = None, @@ -74,7 +74,7 @@ class FactorizationMachinesPredictor(Predictor): class FactorizationMachinesModel(Model): def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, sagemaker_session: Optional[Session] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi b/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi index f893b5c38b13..d7fec5412f8c 100644 --- a/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi +++ b/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi @@ -23,9 +23,9 @@ class IPInsights(AmazonAlgorithmEstimatorBase): weight_decay: hp def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, num_entity_vectors: Optional[int] = None, vector_dim: Optional[int] = None, batch_metrics_publish_interval: Optional[int] = None, @@ -45,7 +45,7 @@ class IPInsightsPredictor(Predictor): class IPInsightsModel(Model): def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, sagemaker_session: Optional[Session] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/amazon/kmeans.pyi b/stubs/sagemaker/sagemaker/amazon/kmeans.pyi index 121cfa8088f8..fbc0be928537 100644 --- a/stubs/sagemaker/sagemaker/amazon/kmeans.pyi +++ b/stubs/sagemaker/sagemaker/amazon/kmeans.pyi @@ -23,9 +23,9 @@ class KMeans(AmazonAlgorithmEstimatorBase): eval_metrics: hp def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, k: Optional[int] = None, init_method: Optional[str] = None, max_iterations: Optional[int] = None, @@ -35,7 +35,7 @@ class KMeans(AmazonAlgorithmEstimatorBase): half_life_time_size: Optional[int] = None, epochs: Optional[int] = None, center_factor: Optional[int] = None, - eval_metrics: Optional[List[Union[str, PipelineVariable]]] = None, + eval_metrics: Optional[List[str | PipelineVariable]] = None, **kwargs, ) -> None: ... def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... @@ -47,7 +47,7 @@ class KMeansPredictor(Predictor): class KMeansModel(Model): def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, sagemaker_session: Optional[Session] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/amazon/knn.pyi b/stubs/sagemaker/sagemaker/amazon/knn.pyi index 5db619fd0e0f..093b6186eb6a 100644 --- a/stubs/sagemaker/sagemaker/amazon/knn.pyi +++ b/stubs/sagemaker/sagemaker/amazon/knn.pyi @@ -22,9 +22,9 @@ class KNN(AmazonAlgorithmEstimatorBase): faiss_index_pq_m: hp def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, k: Optional[int] = None, sample_size: Optional[int] = None, predictor_type: Optional[str] = None, @@ -44,7 +44,7 @@ class KNNPredictor(Predictor): class KNNModel(Model): def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, sagemaker_session: Optional[Session] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/amazon/lda.pyi b/stubs/sagemaker/sagemaker/amazon/lda.pyi index 49efbaf7a1db..cbe389ebe9e0 100644 --- a/stubs/sagemaker/sagemaker/amazon/lda.pyi +++ b/stubs/sagemaker/sagemaker/amazon/lda.pyi @@ -20,8 +20,8 @@ class LDA(AmazonAlgorithmEstimatorBase): tol: hp def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, num_topics: Optional[int] = None, alpha0: Optional[float] = None, max_restarts: Optional[int] = None, @@ -37,7 +37,7 @@ class LDAPredictor(Predictor): class LDAModel(Model): def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, sagemaker_session: Optional[Session] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi b/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi index 867d6450a795..35eaca2bf260 100644 --- a/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi +++ b/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi @@ -58,9 +58,9 @@ class LinearLearner(AmazonAlgorithmEstimatorBase): balance_multiclass_weights: hp def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, predictor_type: Optional[str] = None, binary_classifier_model_selection_criteria: Optional[str] = None, target_recall: Optional[float] = None, @@ -113,7 +113,7 @@ class LinearLearnerPredictor(Predictor): class LinearLearnerModel(Model): def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, sagemaker_session: Optional[Session] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/amazon/ntm.pyi b/stubs/sagemaker/sagemaker/amazon/ntm.pyi index 4ca8a2ab3a5e..06ac9a112442 100644 --- a/stubs/sagemaker/sagemaker/amazon/ntm.pyi +++ b/stubs/sagemaker/sagemaker/amazon/ntm.pyi @@ -25,9 +25,9 @@ class NTM(AmazonAlgorithmEstimatorBase): learning_rate: hp def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, num_topics: Optional[int] = None, encoder_layers: Optional[List] = None, epochs: Optional[int] = None, @@ -50,7 +50,7 @@ class NTMPredictor(Predictor): class NTMModel(Model): def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, sagemaker_session: Optional[Session] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/amazon/object2vec.pyi b/stubs/sagemaker/sagemaker/amazon/object2vec.pyi index f537d79b6621..feecabb546c8 100644 --- a/stubs/sagemaker/sagemaker/amazon/object2vec.pyi +++ b/stubs/sagemaker/sagemaker/amazon/object2vec.pyi @@ -45,9 +45,9 @@ class Object2Vec(AmazonAlgorithmEstimatorBase): enc1_freeze_pretrained_embedding: hp def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, epochs: Optional[int] = None, enc0_max_seq_len: Optional[int] = None, enc0_vocab_size: Optional[int] = None, @@ -88,7 +88,7 @@ class Object2Vec(AmazonAlgorithmEstimatorBase): class Object2VecModel(Model): def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, sagemaker_session: Optional[Session] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/amazon/pca.pyi b/stubs/sagemaker/sagemaker/amazon/pca.pyi index af1a2fa8ce63..9981d23d8947 100644 --- a/stubs/sagemaker/sagemaker/amazon/pca.pyi +++ b/stubs/sagemaker/sagemaker/amazon/pca.pyi @@ -18,9 +18,9 @@ class PCA(AmazonAlgorithmEstimatorBase): extra_components: hp def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, num_components: Optional[int] = None, algorithm_mode: Optional[str] = None, subtract_mean: Optional[bool] = None, @@ -35,7 +35,7 @@ class PCAPredictor(Predictor): class PCAModel(Model): def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, sagemaker_session: Optional[Session] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi b/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi index f646ff1eb145..26525d786425 100644 --- a/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi +++ b/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi @@ -18,9 +18,9 @@ class RandomCutForest(AmazonAlgorithmEstimatorBase): feature_dim: hp def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, num_samples_per_tree: Optional[int] = None, num_trees: Optional[int] = None, eval_metrics: Optional[List] = None, @@ -34,7 +34,7 @@ class RandomCutForestPredictor(Predictor): class RandomCutForestModel(Model): def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, sagemaker_session: Optional[Session] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/chainer/estimator.pyi b/stubs/sagemaker/sagemaker/chainer/estimator.pyi index 9e1cd059a122..5fd5d03bc746 100644 --- a/stubs/sagemaker/sagemaker/chainer/estimator.pyi +++ b/stubs/sagemaker/sagemaker/chainer/estimator.pyi @@ -15,16 +15,16 @@ class Chainer(Framework): additional_mpi_options: Incomplete def __init__( self, - entry_point: Union[str, PipelineVariable], - use_mpi: Optional[Union[bool, PipelineVariable]] = None, - num_processes: Optional[Union[int, PipelineVariable]] = None, - process_slots_per_host: Optional[Union[int, PipelineVariable]] = None, - additional_mpi_options: Optional[Union[str, PipelineVariable]] = None, - source_dir: Optional[Union[str, PipelineVariable]] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + entry_point: str | PipelineVariable, + use_mpi: Optional[bool | PipelineVariable] = None, + num_processes: Optional[int | PipelineVariable] = None, + process_slots_per_host: Optional[int | PipelineVariable] = None, + additional_mpi_options: Optional[str | PipelineVariable] = None, + source_dir: Optional[str | PipelineVariable] = None, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, framework_version: Optional[str] = None, py_version: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, **kwargs, ) -> None: ... def hyperparameters(self): ... diff --git a/stubs/sagemaker/sagemaker/chainer/model.pyi b/stubs/sagemaker/sagemaker/chainer/model.pyi index 9a03e7b2ddc1..2aa2a804ffde 100644 --- a/stubs/sagemaker/sagemaker/chainer/model.pyi +++ b/stubs/sagemaker/sagemaker/chainer/model.pyi @@ -19,40 +19,40 @@ class ChainerModel(FrameworkModel): model_server_workers: Incomplete def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, entry_point: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, framework_version: Optional[str] = None, py_version: Optional[str] = None, predictor_cls: callable = ..., - model_server_workers: Optional[Union[int, PipelineVariable]] = None, + model_server_workers: Optional[int | PipelineVariable] = None, **kwargs, ) -> None: ... image_uri: Incomplete def register( self, - 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, + content_types: List[str | PipelineVariable], + response_types: List[str | PipelineVariable], + inference_instances: Optional[List[str | PipelineVariable]] = None, + transform_instances: Optional[List[str | PipelineVariable]] = None, + model_package_name: Optional[str | PipelineVariable] = None, + model_package_group_name: Optional[str | PipelineVariable] = None, + image_uri: Optional[str | PipelineVariable] = None, model_metrics: Optional[ModelMetrics] = None, metadata_properties: Optional[MetadataProperties] = None, marketplace_cert: bool = False, - approval_status: Optional[Union[str, PipelineVariable]] = None, + approval_status: Optional[str | PipelineVariable] = None, description: Optional[str] = None, drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - domain: Optional[Union[str, PipelineVariable]] = None, - sample_payload_url: Optional[Union[str, PipelineVariable]] = None, - task: Optional[Union[str, PipelineVariable]] = None, - framework: Optional[Union[str, PipelineVariable]] = None, - framework_version: Optional[Union[str, PipelineVariable]] = None, - nearest_model_name: Optional[Union[str, PipelineVariable]] = None, - data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, + domain: Optional[str | PipelineVariable] = None, + sample_payload_url: Optional[str | PipelineVariable] = None, + task: Optional[str | PipelineVariable] = None, + framework: Optional[str | PipelineVariable] = None, + framework_version: Optional[str | PipelineVariable] = None, + nearest_model_name: Optional[str | PipelineVariable] = None, + data_input_configuration: Optional[str | PipelineVariable] = None, ): ... def prepare_container_def( self, diff --git a/stubs/sagemaker/sagemaker/clarify.pyi b/stubs/sagemaker/sagemaker/clarify.pyi index 739411a48042..8641c1734f6b 100644 --- a/stubs/sagemaker/sagemaker/clarify.pyi +++ b/stubs/sagemaker/sagemaker/clarify.pyi @@ -45,13 +45,13 @@ class DataConfig: features: Optional[str] = None, dataset_type: str = "text/csv", s3_compression_type: str = "None", - joinsource: Optional[Union[str, int]] = None, + joinsource: Optional[str | int] = None, facet_dataset_uri: Optional[str] = None, facet_headers: Optional[List[str]] = None, predicted_label_dataset_uri: Optional[str] = None, predicted_label_headers: Optional[List[str]] = None, - predicted_label: Optional[Union[str, int]] = None, - excluded_columns: Optional[Union[List[int], List[str]]] = None, + predicted_label: Optional[str | int] = None, + excluded_columns: Optional[List[int, List[str]]] = None, ) -> None: ... def get_config(self): ... @@ -59,9 +59,9 @@ class BiasConfig: analysis_config: Incomplete def __init__( self, - label_values_or_threshold: Union[int, float, str], - facet_name: Union[str, int, List[str], List[int]], - facet_values_or_threshold: Optional[Union[int, float, str]] = None, + label_values_or_threshold: int | float | str, + facet_name: str | int | List[str, List[int]], + facet_values_or_threshold: Optional[int | float | str] = None, group_name: Optional[str] = None, ) -> None: ... def get_config(self): ... @@ -93,8 +93,8 @@ class ModelPredictedLabelConfig: predictor_config: Incomplete def __init__( self, - label: Optional[Union[str, int]] = None, - probability: Optional[Union[str, int]] = None, + label: Optional[str | int] = None, + probability: Optional[str | int] = None, probability_threshold: Optional[float] = None, label_headers: Optional[List[str]] = None, ) -> None: ... @@ -132,7 +132,7 @@ class SHAPConfig(ExplainabilityConfig): shap_config: Incomplete def __init__( self, - baseline: Optional[Union[str, List, Dict]] = None, + baseline: Optional[str | List | Dict] = None, num_samples: Optional[int] = None, agg_method: Optional[str] = None, use_logit: bool = False, @@ -169,7 +169,7 @@ class SageMakerClarifyProcessor(Processor): self, data_config: DataConfig, data_bias_config: BiasConfig, - methods: Union[str, List[str]] = "all", + methods: str | List[str] = "all", wait: bool = True, logs: bool = True, job_name: Optional[str] = None, @@ -182,7 +182,7 @@ class SageMakerClarifyProcessor(Processor): data_bias_config: BiasConfig, model_config: Optional[ModelConfig] = None, model_predicted_label_config: Optional[ModelPredictedLabelConfig] = None, - methods: Union[str, List[str]] = "all", + methods: str | List[str] = "all", wait: bool = True, logs: bool = True, job_name: Optional[str] = None, @@ -195,8 +195,8 @@ class SageMakerClarifyProcessor(Processor): bias_config: BiasConfig, model_config: Optional[ModelConfig] = None, model_predicted_label_config: Optional[ModelPredictedLabelConfig] = None, - pre_training_methods: Union[str, List[str]] = "all", - post_training_methods: Union[str, List[str]] = "all", + pre_training_methods: str | List[str] = "all", + post_training_methods: str | List[str] = "all", wait: bool = True, logs: bool = True, job_name: Optional[str] = None, @@ -207,8 +207,8 @@ class SageMakerClarifyProcessor(Processor): self, data_config: DataConfig, model_config: ModelConfig, - explainability_config: Union[ExplainabilityConfig, List], - model_scores: Optional[Union[int, str, ModelPredictedLabelConfig]] = None, + explainability_config: ExplainabilityConfig | List, + model_scores: Optional[int | str | ModelPredictedLabelConfig] = None, wait: bool = True, logs: bool = True, job_name: Optional[str] = None, @@ -219,10 +219,10 @@ class SageMakerClarifyProcessor(Processor): self, data_config: DataConfig, model_config: ModelConfig, - explainability_config: Union[ExplainabilityConfig, List[ExplainabilityConfig]], + explainability_config: ExplainabilityConfig | List[ExplainabilityConfig], bias_config: BiasConfig, - pre_training_methods: Union[str, List[str]] = "all", - post_training_methods: Union[str, List[str]] = "all", + pre_training_methods: str | List[str] = "all", + post_training_methods: str | List[str] = "all", model_predicted_label_config: ModelPredictedLabelConfig = None, wait: bool = True, logs: bool = True, @@ -238,10 +238,10 @@ class _AnalysisConfigGenerator: data_config: DataConfig, model_config: ModelConfig, model_predicted_label_config: ModelPredictedLabelConfig, - explainability_config: Union[ExplainabilityConfig, List[ExplainabilityConfig]], + explainability_config: ExplainabilityConfig | List[ExplainabilityConfig], bias_config: BiasConfig, - pre_training_methods: Union[str, List[str]] = "all", - post_training_methods: Union[str, List[str]] = "all", + pre_training_methods: str | List[str] = "all", + post_training_methods: str | List[str] = "all", ): ... @classmethod def explainability( @@ -249,17 +249,17 @@ class _AnalysisConfigGenerator: data_config: DataConfig, model_config: ModelConfig, model_predicted_label_config: ModelPredictedLabelConfig, - explainability_config: Union[ExplainabilityConfig, List[ExplainabilityConfig]], + explainability_config: ExplainabilityConfig | List[ExplainabilityConfig], ): ... @classmethod - def bias_pre_training(cls, data_config: DataConfig, bias_config: BiasConfig, methods: Union[str, List[str]]): ... + def bias_pre_training(cls, data_config: DataConfig, bias_config: BiasConfig, methods: str | List[str]): ... @classmethod def bias_post_training( cls, data_config: DataConfig, bias_config: BiasConfig, model_predicted_label_config: ModelPredictedLabelConfig, - methods: Union[str, List[str]], + methods: str | List[str], model_config: ModelConfig, ): ... @classmethod @@ -269,8 +269,8 @@ class _AnalysisConfigGenerator: bias_config: BiasConfig, model_config: ModelConfig, model_predicted_label_config: ModelPredictedLabelConfig, - pre_training_methods: Union[str, List[str]] = "all", - post_training_methods: Union[str, List[str]] = "all", + pre_training_methods: str | List[str] = "all", + post_training_methods: str | List[str] = "all", ): ... class ProcessingOutputHandler: diff --git a/stubs/sagemaker/sagemaker/debugger/debugger.pyi b/stubs/sagemaker/sagemaker/debugger/debugger.pyi index ae44e407fff1..4bad04bf37f8 100644 --- a/stubs/sagemaker/sagemaker/debugger/debugger.pyi +++ b/stubs/sagemaker/sagemaker/debugger/debugger.pyi @@ -57,15 +57,15 @@ class Rule(RuleBase): def custom( cls, name: str, - image_uri: Union[str, PipelineVariable], - instance_type: Union[str, PipelineVariable], - volume_size_in_gb: Union[int, PipelineVariable], + image_uri: str | PipelineVariable, + instance_type: str | PipelineVariable, + volume_size_in_gb: int | PipelineVariable, source: Optional[str] = None, - rule_to_invoke: Optional[Union[str, PipelineVariable]] = None, - container_local_output_path: Optional[Union[str, PipelineVariable]] = None, - s3_output_path: Optional[Union[str, PipelineVariable]] = None, - other_trials_s3_input_paths: Optional[List[Union[str, PipelineVariable]]] = None, - rule_parameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + rule_to_invoke: Optional[str | PipelineVariable] = None, + container_local_output_path: Optional[str | PipelineVariable] = None, + s3_output_path: Optional[str | PipelineVariable] = None, + other_trials_s3_input_paths: Optional[List[str | PipelineVariable]] = None, + rule_parameters: Optional[Dict[str, str | PipelineVariable]] = None, collections_to_save: Optional[List["CollectionConfig"]] = None, actions: Incomplete | None = None, ): ... @@ -103,9 +103,9 @@ class DebuggerHookConfig: collection_configs: Incomplete def __init__( self, - s3_output_path: Optional[Union[str, PipelineVariable]] = None, - container_local_output_path: Optional[Union[str, PipelineVariable]] = None, - hook_parameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + s3_output_path: Optional[str | PipelineVariable] = None, + container_local_output_path: Optional[str | PipelineVariable] = None, + hook_parameters: Optional[Dict[str, str | PipelineVariable]] = None, collection_configs: Optional[List["CollectionConfig"]] = None, ) -> None: ... @@ -114,15 +114,15 @@ class TensorBoardOutputConfig: container_local_output_path: Incomplete def __init__( self, - s3_output_path: Union[str, PipelineVariable], - container_local_output_path: Optional[Union[str, PipelineVariable]] = None, + s3_output_path: str | PipelineVariable, + container_local_output_path: Optional[str | PipelineVariable] = None, ) -> None: ... class CollectionConfig: name: Incomplete parameters: Incomplete def __init__( - self, name: Union[str, PipelineVariable], parameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None + self, name: str | PipelineVariable, parameters: Optional[Dict[str, str | PipelineVariable]] = None ) -> None: ... def __eq__(self, other): ... def __ne__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi b/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi index f5e7cb8f2dec..11e37d5c275c 100644 --- a/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi +++ b/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi @@ -13,8 +13,8 @@ class ProfilerConfig: disable_profiler: Incomplete def __init__( self, - s3_output_path: Optional[Union[str, PipelineVariable]] = None, - system_monitor_interval_millis: Optional[Union[int, PipelineVariable]] = None, + s3_output_path: Optional[str | PipelineVariable] = None, + system_monitor_interval_millis: Optional[int | PipelineVariable] = None, framework_profile_params: Optional[FrameworkProfile] = None, - disable_profiler: Optional[Union[str, PipelineVariable]] = False, + disable_profiler: Optional[str | PipelineVariable] = False, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/djl_inference/model.pyi b/stubs/sagemaker/sagemaker/djl_inference/model.pyi index 433d3495c1a7..abfe9fd0e60e 100644 --- a/stubs/sagemaker/sagemaker/djl_inference/model.pyi +++ b/stubs/sagemaker/sagemaker/djl_inference/model.pyi @@ -56,7 +56,7 @@ class DJLModel(FrameworkModel): model_loading_timeout: Optional[int] = None, prediction_timeout: Optional[int] = None, entry_point: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, predictor_cls: callable = ..., **kwargs, ) -> None: ... @@ -145,7 +145,7 @@ class HuggingFaceAccelerateModel(DJLModel): role: str, number_of_partitions: Optional[int] = None, device_id: Optional[int] = None, - device_map: Optional[Union[str, Dict[str, str]]] = None, + device_map: Optional[str | Dict[str | str]] = None, load_in_8bit: bool = False, low_cpu_mem_usage: bool = False, **kwargs, diff --git a/stubs/sagemaker/sagemaker/estimator.pyi b/stubs/sagemaker/sagemaker/estimator.pyi index 936972d060e5..0452c4dc26c4 100644 --- a/stubs/sagemaker/sagemaker/estimator.pyi +++ b/stubs/sagemaker/sagemaker/estimator.pyi @@ -88,47 +88,47 @@ class EstimatorBase(metaclass=abc.ABCMeta): def __init__( self, role: str = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, - keep_alive_period_in_seconds: Optional[Union[int, PipelineVariable]] = None, - volume_size: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - max_run: Union[int, PipelineVariable] = 86400, - input_mode: Union[str, PipelineVariable] = "File", - output_path: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, + keep_alive_period_in_seconds: Optional[int | PipelineVariable] = None, + volume_size: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + max_run: int | PipelineVariable = 86400, + input_mode: str | PipelineVariable = "File", + output_path: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[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, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, + subnets: Optional[List[str | PipelineVariable]] = None, + security_group_ids: Optional[List[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] = None, - 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, + model_channel_name: str | PipelineVariable = "model", + metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, + encrypt_inter_container_traffic: bool | PipelineVariable = None, + use_spot_instances: bool | PipelineVariable = False, + max_wait: Optional[int | PipelineVariable] = None, + checkpoint_s3_uri: Optional[str | PipelineVariable] = None, + checkpoint_local_path: Optional[str | PipelineVariable] = None, rules: Optional[List[RuleBase]] = None, - debugger_hook_config: Optional[Union[bool, DebuggerHookConfig]] = None, + debugger_hook_config: Optional[bool | DebuggerHookConfig] = None, tensorboard_output_config: Optional[TensorBoardOutputConfig] = None, - enable_sagemaker_metrics: Optional[Union[bool, PipelineVariable]] = None, - enable_network_isolation: Union[bool, PipelineVariable] = None, + enable_sagemaker_metrics: Optional[bool | PipelineVariable] = None, + enable_network_isolation: bool | PipelineVariable = None, profiler_config: Optional[ProfilerConfig] = None, disable_profiler: bool = None, - environment: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - max_retry_attempts: Optional[Union[int, PipelineVariable]] = None, - source_dir: Optional[Union[str, PipelineVariable]] = None, + environment: Optional[Dict[str, str | PipelineVariable]] = None, + max_retry_attempts: Optional[int | PipelineVariable] = None, + source_dir: Optional[str | PipelineVariable] = None, git_config: Optional[Dict[str, str]] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - container_log_level: Union[int, PipelineVariable] = 20, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + container_log_level: int | PipelineVariable = 20, code_location: Optional[str] = None, - entry_point: Optional[Union[str, PipelineVariable]] = None, - dependencies: Optional[List[Union[str]]] = None, + entry_point: Optional[str | PipelineVariable] = None, + dependencies: Optional[List[str]] = None, instance_groups: Optional[List[InstanceGroup]] = None, - training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None, - training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None, + training_repository_access_mode: Optional[str | PipelineVariable] = None, + training_repository_credentials_provider_arn: Optional[str | PipelineVariable] = None, container_entry_point: Optional[List[str]] = None, container_arguments: Optional[List[str]] = None, disable_output_compression: bool = False, @@ -145,7 +145,7 @@ class EstimatorBase(metaclass=abc.ABCMeta): def latest_job_profiler_artifacts_path(self): ... def fit( self, - inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None, + inputs: Optional[str | Dict | TrainingInput | FileSystemInput] = None, wait: bool = True, logs: str = "All", job_name: Optional[str] = None, @@ -277,49 +277,49 @@ class Estimator(EstimatorBase): image_uri: Incomplete def __init__( self, - image_uri: Union[str, PipelineVariable], + image_uri: str | PipelineVariable, role: str = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, - keep_alive_period_in_seconds: Optional[Union[int, PipelineVariable]] = None, - volume_size: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - max_run: Union[int, PipelineVariable] = 86400, - input_mode: Union[str, PipelineVariable] = "File", - output_path: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, + keep_alive_period_in_seconds: Optional[int | PipelineVariable] = None, + volume_size: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + max_run: int | PipelineVariable = 86400, + input_mode: str | PipelineVariable = "File", + output_path: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = 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, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, + subnets: Optional[List[str | PipelineVariable]] = None, + security_group_ids: Optional[List[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] = None, - 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, - enable_network_isolation: Union[bool, PipelineVariable] = None, + model_channel_name: str | PipelineVariable = "model", + metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, + encrypt_inter_container_traffic: bool | PipelineVariable = None, + use_spot_instances: bool | PipelineVariable = False, + max_wait: Optional[int | PipelineVariable] = None, + checkpoint_s3_uri: Optional[str | PipelineVariable] = None, + checkpoint_local_path: Optional[str | PipelineVariable] = None, + enable_network_isolation: bool | PipelineVariable = None, rules: Optional[List[RuleBase]] = None, - debugger_hook_config: Optional[Union[DebuggerHookConfig, bool]] = None, + debugger_hook_config: Optional[DebuggerHookConfig | bool] = None, tensorboard_output_config: Optional[TensorBoardOutputConfig] = None, - enable_sagemaker_metrics: Optional[Union[bool, PipelineVariable]] = None, + enable_sagemaker_metrics: Optional[bool | PipelineVariable] = None, 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[Union[str, PipelineVariable]] = None, + environment: Optional[Dict[str, str | PipelineVariable]] = None, + max_retry_attempts: Optional[int | PipelineVariable] = None, + source_dir: Optional[str | PipelineVariable] = None, git_config: Optional[Dict[str, str]] = None, - container_log_level: Union[int, PipelineVariable] = 20, + container_log_level: int | PipelineVariable = 20, code_location: Optional[str] = None, - entry_point: Optional[Union[str, PipelineVariable]] = None, + entry_point: Optional[str | PipelineVariable] = None, dependencies: Optional[List[str]] = None, instance_groups: Optional[List[InstanceGroup]] = None, - training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None, - training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None, + training_repository_access_mode: Optional[str | PipelineVariable] = None, + training_repository_credentials_provider_arn: Optional[str | PipelineVariable] = None, container_entry_point: Optional[List[str]] = None, container_arguments: Optional[List[str]] = None, disable_output_compression: bool = False, @@ -351,18 +351,18 @@ class Framework(EstimatorBase, metaclass=abc.ABCMeta): enable_sagemaker_metrics: Incomplete def __init__( self, - entry_point: Union[str, PipelineVariable], - source_dir: Optional[Union[str, PipelineVariable]] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - container_log_level: Union[int, PipelineVariable] = 20, + entry_point: str | PipelineVariable, + source_dir: Optional[str | PipelineVariable] = None, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + container_log_level: int | PipelineVariable = 20, code_location: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, dependencies: Optional[List[str]] = None, - enable_network_isolation: Union[bool, PipelineVariable] = None, + enable_network_isolation: bool | PipelineVariable = None, git_config: Optional[Dict[str, str]] = None, - checkpoint_s3_uri: Optional[Union[str, PipelineVariable]] = None, - checkpoint_local_path: Optional[Union[str, PipelineVariable]] = None, - enable_sagemaker_metrics: Optional[Union[bool, PipelineVariable]] = None, + checkpoint_s3_uri: Optional[str | PipelineVariable] = None, + checkpoint_local_path: Optional[str | PipelineVariable] = None, + enable_sagemaker_metrics: Optional[bool | PipelineVariable] = None, **kwargs, ) -> None: ... def set_hyperparameters(self, **kwargs) -> None: ... diff --git a/stubs/sagemaker/sagemaker/experiments/run.pyi b/stubs/sagemaker/sagemaker/experiments/run.pyi index 5ef5f9f6cfe6..0c11710a76df 100644 --- a/stubs/sagemaker/sagemaker/experiments/run.pyi +++ b/stubs/sagemaker/sagemaker/experiments/run.pyi @@ -42,23 +42,23 @@ class Run: ) -> None: ... @property def experiment_config(self) -> dict: ... - def log_parameter(self, name: str, value: Union[str, int, float]): ... - def log_parameters(self, parameters: Dict[str, Union[str, int, float]]): ... + def log_parameter(self, name: str, value: str | int | float): ... + def log_parameters(self, parameters: Dict[str, str | int | float]): ... def log_metric(self, name: str, value: float, timestamp: Optional[datetime.datetime] = None, step: Optional[int] = None): ... def log_precision_recall( self, - y_true: Union[list, array], - predicted_probabilities: Union[list, array], - positive_label: Optional[Union[str, int]] = None, + y_true: list | array, + predicted_probabilities: list | array, + positive_label: Optional[str | int] = None, title: Optional[str] = None, is_output: bool = True, no_skill: Optional[int] = None, ): ... def log_roc_curve( - self, y_true: Union[list, array], y_score: Union[list, array], title: Optional[str] = None, is_output: bool = True + self, y_true: list | array, y_score: list | array, title: Optional[str] = None, is_output: bool = True ): ... def log_confusion_matrix( - self, y_true: Union[list, array], y_pred: Union[list, array], title: Optional[str] = None, is_output: bool = True + self, y_true: list | array, y_pred: list | array, title: Optional[str] = None, is_output: bool = True ): ... def log_artifact(self, name: str, value: str, media_type: Optional[str] = None, is_output: bool = True): ... def log_file(self, file_path: str, name: Optional[str] = None, media_type: Optional[str] = None, is_output: bool = True): ... diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi index e7ea84cd71b0..24b82ae0d3a7 100644 --- a/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi @@ -67,7 +67,7 @@ class FeatureGroup: DTYPE_TO_FEATURE_DEFINITION_CLS_MAP: Dict[str, FeatureTypeEnum] def create( self, - s3_uri: Union[str, bool], + s3_uri: str | bool, record_identifier_name: str, event_time_feature_name: str, role_arn: str = None, @@ -105,7 +105,7 @@ class FeatureGroup: max_workers: int = 1, max_processes: int = 1, wait: bool = True, - timeout: Union[int, float] = None, + timeout: int | float = None, profile_name: str = None, ) -> IngestionManagerPandas: ... def athena_query(self) -> AthenaQuery: ... diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi index 57bc87875a9e..f4f51c3d7a81 100644 --- a/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi @@ -11,7 +11,7 @@ class FeatureStore: sagemaker_session: Session def create_dataset( self, - base: Union[FeatureGroup, pd.DataFrame], + base: FeatureGroup | pd.DataFrame, output_path: str, record_identifier_feature_name: str = None, event_time_identifier_feature_name: str = None, diff --git a/stubs/sagemaker/sagemaker/fw_utils.pyi b/stubs/sagemaker/sagemaker/fw_utils.pyi index 936b80a9bd5b..6b9d7fc951c3 100644 --- a/stubs/sagemaker/sagemaker/fw_utils.pyi +++ b/stubs/sagemaker/sagemaker/fw_utils.pyi @@ -27,10 +27,10 @@ GRAVITON_ALLOWED_FRAMEWORKS: Incomplete def validate_source_dir(script, directory): ... def validate_source_code_input_against_pipeline_variables( - entry_point: Optional[Union[str, PipelineVariable]] = None, - source_dir: Optional[Union[str, PipelineVariable]] = None, + entry_point: Optional[str | PipelineVariable] = None, + source_dir: Optional[str | PipelineVariable] = None, git_config: Optional[Dict[str, str]] = None, - enable_network_isolation: Union[bool, PipelineVariable] = False, + enable_network_isolation: bool | PipelineVariable = False, ): ... def parse_mp_parameters(params): ... def get_mp_parameters(distribution): ... diff --git a/stubs/sagemaker/sagemaker/huggingface/estimator.pyi b/stubs/sagemaker/sagemaker/huggingface/estimator.pyi index 8e44c797b66a..e8d44f2211ae 100644 --- a/stubs/sagemaker/sagemaker/huggingface/estimator.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/estimator.pyi @@ -22,13 +22,13 @@ class HuggingFace(Framework): def __init__( self, py_version: str, - entry_point: Union[str, PipelineVariable], + entry_point: str | PipelineVariable, transformers_version: Optional[str] = None, tensorflow_version: Optional[str] = None, pytorch_version: Optional[str] = None, - source_dir: Optional[Union[str, PipelineVariable]] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + source_dir: Optional[str | PipelineVariable] = None, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, distribution: Optional[Dict] = None, compiler_config: Optional[TrainingCompilerConfig] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/huggingface/model.pyi b/stubs/sagemaker/sagemaker/huggingface/model.pyi index 58f76111630c..ef886b6f32b7 100644 --- a/stubs/sagemaker/sagemaker/huggingface/model.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/model.pyi @@ -25,15 +25,15 @@ class HuggingFaceModel(FrameworkModel): def __init__( self, role: Optional[str] = None, - model_data: Optional[Union[str, PipelineVariable]] = None, + model_data: Optional[str | PipelineVariable] = None, entry_point: Optional[str] = None, transformers_version: Optional[str] = None, tensorflow_version: Optional[str] = None, pytorch_version: Optional[str] = None, py_version: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, predictor_cls: callable = ..., - model_server_workers: Optional[Union[int, PipelineVariable]] = None, + model_server_workers: Optional[int | PipelineVariable] = None, **kwargs, ) -> None: ... image_uri: Incomplete @@ -60,27 +60,27 @@ class HuggingFaceModel(FrameworkModel): ): ... def register( self, - 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, + content_types: List[str | PipelineVariable], + response_types: List[str | PipelineVariable], + inference_instances: Optional[List[str | PipelineVariable]] = None, + transform_instances: Optional[List[str | PipelineVariable]] = None, + model_package_name: Optional[str | PipelineVariable] = None, + model_package_group_name: Optional[str | PipelineVariable] = None, + image_uri: Optional[str | PipelineVariable] = None, model_metrics: Optional[ModelMetrics] = None, metadata_properties: Optional[MetadataProperties] = None, marketplace_cert: bool = False, - approval_status: Optional[Union[str, PipelineVariable]] = None, + approval_status: Optional[str | PipelineVariable] = None, description: Optional[str] = None, drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - domain: Optional[Union[str, PipelineVariable]] = None, - sample_payload_url: Optional[Union[str, PipelineVariable]] = None, - task: Optional[Union[str, PipelineVariable]] = None, - framework: Optional[Union[str, PipelineVariable]] = None, - framework_version: Optional[Union[str, PipelineVariable]] = None, - nearest_model_name: Optional[Union[str, PipelineVariable]] = None, - data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, + domain: Optional[str | PipelineVariable] = None, + sample_payload_url: Optional[str | PipelineVariable] = None, + task: Optional[str | PipelineVariable] = None, + framework: Optional[str | PipelineVariable] = None, + framework_version: Optional[str | PipelineVariable] = None, + nearest_model_name: Optional[str | PipelineVariable] = None, + data_input_configuration: Optional[str | PipelineVariable] = None, ): ... def prepare_container_def( self, diff --git a/stubs/sagemaker/sagemaker/huggingface/processing.pyi b/stubs/sagemaker/sagemaker/huggingface/processing.pyi index 4c63d3d8e4f6..708c3f58fa7d 100644 --- a/stubs/sagemaker/sagemaker/huggingface/processing.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/processing.pyi @@ -13,23 +13,23 @@ class HuggingFaceProcessor(FrameworkProcessor): tensorflow_version: Incomplete def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Union[int, PipelineVariable] = None, - instance_type: Union[str, PipelineVariable] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: int | PipelineVariable = None, + instance_type: str | PipelineVariable = None, transformers_version: Optional[str] = None, tensorflow_version: Optional[str] = None, pytorch_version: Optional[str] = None, py_version: str = "py36", - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, command: Optional[List[str]] = None, - volume_size_in_gb: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + volume_size_in_gb: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, code_location: Optional[str] = None, - max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + max_runtime_in_seconds: Optional[int | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, network_config: Optional[NetworkConfig] = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi b/stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi index 90146f5a7e03..272711ce08e7 100644 --- a/stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi @@ -9,6 +9,6 @@ logger: Incomplete class TrainingCompilerConfig(BaseConfig): SUPPORTED_INSTANCE_CLASS_PREFIXES: Incomplete SUPPORTED_INSTANCE_TYPES_WITH_EFA: Incomplete - def __init__(self, enabled: Union[bool, PipelineVariable] = True, debug: Union[bool, PipelineVariable] = False) -> None: ... + def __init__(self, enabled: bool | PipelineVariable = True, debug: bool | PipelineVariable = False) -> None: ... @classmethod def validate(cls, estimator) -> None: ... diff --git a/stubs/sagemaker/sagemaker/inputs.pyi b/stubs/sagemaker/sagemaker/inputs.pyi index e8967bac6e9e..41e845506f5f 100644 --- a/stubs/sagemaker/sagemaker/inputs.pyi +++ b/stubs/sagemaker/sagemaker/inputs.pyi @@ -10,16 +10,16 @@ class TrainingInput: config: Incomplete def __init__( self, - s3_data: Union[str, PipelineVariable], - distribution: Optional[Union[str, PipelineVariable]] = None, - compression: Optional[Union[str, PipelineVariable]] = None, - content_type: Optional[Union[str, PipelineVariable]] = None, - record_wrapping: Optional[Union[str, PipelineVariable]] = None, - s3_data_type: Union[str, PipelineVariable] = "S3Prefix", - instance_groups: Optional[List[Union[str, PipelineVariable]]] = None, - input_mode: Optional[Union[str, PipelineVariable]] = None, - attribute_names: Optional[List[Union[str, PipelineVariable]]] = None, - target_attribute_name: Optional[Union[str, PipelineVariable]] = None, + s3_data: str | PipelineVariable, + distribution: Optional[str | PipelineVariable] = None, + compression: Optional[str | PipelineVariable] = None, + content_type: Optional[str | PipelineVariable] = None, + record_wrapping: Optional[str | PipelineVariable] = None, + s3_data_type: str | PipelineVariable = "S3Prefix", + instance_groups: Optional[List[str | PipelineVariable]] = None, + input_mode: Optional[str | PipelineVariable] = None, + attribute_names: Optional[List[str | PipelineVariable]] = None, + target_attribute_name: Optional[str | PipelineVariable] = None, shuffle_config: Optional["ShuffleConfig"] = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi b/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi index 11406710a414..eeffbe8a0dc3 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi @@ -36,53 +36,53 @@ class JumpStartEstimator(Estimator): tolerate_vulnerable_model: Optional[bool] = None, tolerate_deprecated_model: Optional[bool] = None, region: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, role: Optional[str] = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, - keep_alive_period_in_seconds: Optional[Union[int, PipelineVariable]] = None, - volume_size: Optional[Union[int, PipelineVariable]] = None, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - max_run: Optional[Union[int, PipelineVariable]] = None, - input_mode: Optional[Union[str, PipelineVariable]] = None, - output_path: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, + keep_alive_period_in_seconds: Optional[int | PipelineVariable] = None, + volume_size: Optional[int | PipelineVariable] = None, + volume_kms_key: Optional[str | PipelineVariable] = None, + max_run: Optional[int | PipelineVariable] = None, + input_mode: Optional[str | PipelineVariable] = None, + output_path: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[session.Session] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = 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, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, + subnets: Optional[List[str | PipelineVariable]] = None, + security_group_ids: Optional[List[str | PipelineVariable]] = None, model_uri: Optional[str] = None, - model_channel_name: Optional[Union[str, PipelineVariable]] = None, - metric_definitions: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, - encrypt_inter_container_traffic: Union[bool, PipelineVariable] = None, - use_spot_instances: Optional[Union[bool, PipelineVariable]] = None, - max_wait: Optional[Union[int, PipelineVariable]] = None, - checkpoint_s3_uri: Optional[Union[str, PipelineVariable]] = None, - checkpoint_local_path: Optional[Union[str, PipelineVariable]] = None, - enable_network_isolation: Union[bool, PipelineVariable] = None, + model_channel_name: Optional[str | PipelineVariable] = None, + metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, + encrypt_inter_container_traffic: bool | PipelineVariable = None, + use_spot_instances: Optional[bool | PipelineVariable] = None, + max_wait: Optional[int | PipelineVariable] = None, + checkpoint_s3_uri: Optional[str | PipelineVariable] = None, + checkpoint_local_path: Optional[str | PipelineVariable] = None, + enable_network_isolation: bool | PipelineVariable = None, rules: Optional[List[RuleBase]] = None, - debugger_hook_config: Optional[Union[DebuggerHookConfig, bool]] = None, + debugger_hook_config: Optional[DebuggerHookConfig | bool] = None, tensorboard_output_config: Optional[TensorBoardOutputConfig] = None, - enable_sagemaker_metrics: Optional[Union[bool, PipelineVariable]] = None, + enable_sagemaker_metrics: Optional[bool | PipelineVariable] = None, profiler_config: Optional[ProfilerConfig] = None, disable_profiler: Optional[bool] = None, - environment: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - max_retry_attempts: Optional[Union[int, PipelineVariable]] = None, - source_dir: Optional[Union[str, PipelineVariable]] = None, + environment: Optional[Dict[str, str | PipelineVariable]] = None, + max_retry_attempts: Optional[int | PipelineVariable] = None, + source_dir: Optional[str | PipelineVariable] = None, git_config: Optional[Dict[str, str]] = None, - container_log_level: Optional[Union[int, PipelineVariable]] = None, + container_log_level: Optional[int | PipelineVariable] = None, code_location: Optional[str] = None, - entry_point: Optional[Union[str, PipelineVariable]] = None, + entry_point: Optional[str | PipelineVariable] = None, dependencies: Optional[List[str]] = None, instance_groups: Optional[List[InstanceGroup]] = None, - training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None, - training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None, + training_repository_access_mode: Optional[str | PipelineVariable] = None, + training_repository_credentials_provider_arn: Optional[str | PipelineVariable] = None, ) -> None: ... def fit( self, - inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None, + inputs: Optional[str | Dict | TrainingInput | FileSystemInput] = None, wait: Optional[bool] = True, logs: Optional[str] = None, job_name: Optional[str] = None, @@ -107,20 +107,20 @@ class JumpStartEstimator(Estimator): container_startup_health_check_timeout: Optional[int] = None, inference_recommendation_id: Optional[str] = None, explainer_config: Optional[ExplainerConfig] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, role: Optional[str] = None, predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, model_name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = None, + vpc_config: Optional[Dict[str, List[str | PipelineVariable]]] = None, sagemaker_session: Optional[session.Session] = None, - enable_network_isolation: Union[bool, PipelineVariable] = None, + enable_network_isolation: bool | PipelineVariable = None, model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + image_config: Optional[Dict[str, str | PipelineVariable]] = None, source_dir: Optional[str] = None, code_location: Optional[str] = None, entry_point: Optional[str] = None, - container_log_level: Optional[Union[int, PipelineVariable]] = None, + container_log_level: Optional[int | PipelineVariable] = None, dependencies: Optional[List[str]] = None, git_config: Optional[Dict[str, str]] = None, use_compiled_model: bool = False, diff --git a/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi b/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi index f2c282bdb0d5..6773faa713bf 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi @@ -23,55 +23,55 @@ def get_init_kwargs( tolerate_vulnerable_model: Optional[bool] = None, tolerate_deprecated_model: Optional[bool] = None, region: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, role: Optional[str] = None, - instance_count: Optional[Union[int, PipelineVariable]] = None, - instance_type: Optional[Union[str, PipelineVariable]] = None, - keep_alive_period_in_seconds: Optional[Union[int, PipelineVariable]] = None, - volume_size: Optional[Union[int, PipelineVariable]] = None, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - max_run: Optional[Union[int, PipelineVariable]] = None, - input_mode: Optional[Union[str, PipelineVariable]] = None, - output_path: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + instance_count: Optional[int | PipelineVariable] = None, + instance_type: Optional[str | PipelineVariable] = None, + keep_alive_period_in_seconds: Optional[int | PipelineVariable] = None, + volume_size: Optional[int | PipelineVariable] = None, + volume_kms_key: Optional[str | PipelineVariable] = None, + max_run: Optional[int | PipelineVariable] = None, + input_mode: Optional[str | PipelineVariable] = None, + output_path: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = 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, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, + subnets: Optional[List[str | PipelineVariable]] = None, + security_group_ids: Optional[List[str | PipelineVariable]] = None, model_uri: Optional[str] = None, - model_channel_name: Optional[Union[str, PipelineVariable]] = None, - metric_definitions: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, - encrypt_inter_container_traffic: Union[bool, PipelineVariable] = None, - use_spot_instances: Optional[Union[bool, PipelineVariable]] = None, - max_wait: Optional[Union[int, PipelineVariable]] = None, - checkpoint_s3_uri: Optional[Union[str, PipelineVariable]] = None, - checkpoint_local_path: Optional[Union[str, PipelineVariable]] = None, - enable_network_isolation: Union[bool, PipelineVariable] = None, + model_channel_name: Optional[str | PipelineVariable] = None, + metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, + encrypt_inter_container_traffic: bool | PipelineVariable = None, + use_spot_instances: Optional[bool | PipelineVariable] = None, + max_wait: Optional[int | PipelineVariable] = None, + checkpoint_s3_uri: Optional[str | PipelineVariable] = None, + checkpoint_local_path: Optional[str | PipelineVariable] = None, + enable_network_isolation: bool | PipelineVariable = None, rules: Optional[List[RuleBase]] = None, - debugger_hook_config: Optional[Union[DebuggerHookConfig, bool]] = None, + debugger_hook_config: Optional[DebuggerHookConfig | bool] = None, tensorboard_output_config: Optional[TensorBoardOutputConfig] = None, - enable_sagemaker_metrics: Optional[Union[bool, PipelineVariable]] = None, + enable_sagemaker_metrics: Optional[bool | PipelineVariable] = None, profiler_config: Optional[ProfilerConfig] = None, disable_profiler: Optional[bool] = None, - environment: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - max_retry_attempts: Optional[Union[int, PipelineVariable]] = None, - source_dir: Optional[Union[str, PipelineVariable]] = None, + environment: Optional[Dict[str, str | PipelineVariable]] = None, + max_retry_attempts: Optional[int | PipelineVariable] = None, + source_dir: Optional[str | PipelineVariable] = None, git_config: Optional[Dict[str, str]] = None, - container_log_level: Optional[Union[int, PipelineVariable]] = None, + container_log_level: Optional[int | PipelineVariable] = None, code_location: Optional[str] = None, - entry_point: Optional[Union[str, PipelineVariable]] = None, + entry_point: Optional[str | PipelineVariable] = None, dependencies: Optional[List[str]] = None, instance_groups: Optional[List[InstanceGroup]] = None, - training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None, - training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None, + training_repository_access_mode: Optional[str | PipelineVariable] = None, + training_repository_credentials_provider_arn: Optional[str | PipelineVariable] = None, ) -> JumpStartEstimatorInitKwargs: ... def get_fit_kwargs( model_id: str, model_version: Optional[str] = None, region: Optional[str] = None, - inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None, + inputs: Optional[str | Dict | TrainingInput | FileSystemInput] = None, wait: Optional[bool] = None, logs: Optional[str] = None, job_name: Optional[str] = None, @@ -100,19 +100,19 @@ def get_deploy_kwargs( container_startup_health_check_timeout: Optional[int] = None, inference_recommendation_id: Optional[str] = None, explainer_config: Optional[ExplainerConfig] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, role: Optional[str] = None, predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, + vpc_config: Optional[Dict[str, List[str | PipelineVariable]]] = None, sagemaker_session: Optional[Session] = None, - enable_network_isolation: Union[bool, PipelineVariable] = None, + enable_network_isolation: bool | PipelineVariable = None, model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + image_config: Optional[Dict[str, str | PipelineVariable]] = None, source_dir: Optional[str] = None, code_location: Optional[str] = None, entry_point: Optional[str] = None, - container_log_level: Optional[Union[int, PipelineVariable]] = None, + container_log_level: Optional[int | PipelineVariable] = None, dependencies: Optional[List[str]] = None, git_config: Optional[Dict[str, str]] = None, tolerate_deprecated_model: Optional[bool] = None, diff --git a/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi b/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi index 9bd15ddf2568..190d5d0535e1 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi @@ -54,21 +54,21 @@ def get_init_kwargs( tolerate_deprecated_model: Optional[bool] = None, instance_type: Optional[str] = None, region: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, - model_data: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, + model_data: Optional[str | PipelineVariable] = None, role: Optional[str] = None, predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = None, + vpc_config: Optional[Dict[str, List[str | PipelineVariable]]] = None, sagemaker_session: Optional[Session] = None, - enable_network_isolation: Union[bool, PipelineVariable] = None, + enable_network_isolation: bool | PipelineVariable = None, model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + image_config: Optional[Dict[str, str | PipelineVariable]] = None, source_dir: Optional[str] = None, code_location: Optional[str] = None, entry_point: Optional[str] = None, - container_log_level: Optional[Union[int, PipelineVariable]] = None, + container_log_level: Optional[int | PipelineVariable] = None, dependencies: Optional[List[str]] = None, git_config: Optional[Dict[str, str]] = None, ) -> JumpStartModelInitKwargs: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/filters.pyi b/stubs/sagemaker/sagemaker/jumpstart/filters.pyi index c092fff2c849..c2d7ecc84e82 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/filters.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/filters.pyi @@ -44,7 +44,7 @@ class Operator(Operand): class And(Operator): operands: Incomplete - def __init__(self, *operands: Union[Operand, str]) -> None: ... + def __init__(self, *operands: Operand | str) -> None: ... resolved_value: Incomplete def eval(self) -> None: ... def __iter__(self) -> Any: ... @@ -56,21 +56,21 @@ class Constant(Operator): class Identity(Operator): operand: Incomplete - def __init__(self, operand: Union[Operand, str]) -> None: ... + def __init__(self, operand: Operand | str) -> None: ... def __iter__(self) -> Any: ... resolved_value: Incomplete def eval(self) -> None: ... class Or(Operator): operands: Incomplete - def __init__(self, *operands: Union[Operand, str]) -> None: ... + def __init__(self, *operands: Operand | str) -> None: ... resolved_value: Incomplete def eval(self) -> None: ... def __iter__(self) -> Any: ... class Not(Operator): operand: Incomplete - def __init__(self, operand: Union[Operand, str]) -> None: ... + def __init__(self, operand: Operand | str) -> None: ... resolved_value: Incomplete def eval(self) -> None: ... def __iter__(self) -> Any: ... @@ -83,5 +83,5 @@ class ModelFilter(JumpStartDataHolderType): def parse_filter_string(filter_string: str) -> ModelFilter: ... def evaluate_filter_expression( - model_filter: ModelFilter, cached_model_value: Union[str, bool, int, float, Dict[str, Any], List[Any]] + model_filter: ModelFilter, cached_model_value: str | bool | int | float | Dict[str | Any, List[Any]] ) -> BooleanValues: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/model.pyi b/stubs/sagemaker/sagemaker/jumpstart/model.pyi index d81ca06b29b1..2d1d6c4fd675 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/model.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/model.pyi @@ -30,21 +30,21 @@ class JumpStartModel(Model): tolerate_deprecated_model: Optional[bool] = None, region: Optional[str] = None, instance_type: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, - model_data: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, + model_data: Optional[str | PipelineVariable] = None, role: Optional[str] = None, predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = None, + vpc_config: Optional[Dict[str, List[str | PipelineVariable]]] = None, sagemaker_session: Optional[Session] = None, - enable_network_isolation: Union[bool, PipelineVariable] = None, + enable_network_isolation: bool | PipelineVariable = None, model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + image_config: Optional[Dict[str, str | PipelineVariable]] = None, source_dir: Optional[str] = None, code_location: Optional[str] = None, entry_point: Optional[str] = None, - container_log_level: Optional[Union[int, PipelineVariable]] = None, + container_log_level: Optional[int | PipelineVariable] = None, dependencies: Optional[List[str]] = None, git_config: Optional[Dict[str, str]] = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi b/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi index 1972bb6b6efd..14fea565372c 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi @@ -3,14 +3,14 @@ from typing import List, Tuple, Union from sagemaker.jumpstart.filters import Operator def extract_framework_task_model(model_id: str) -> Tuple[str, str, str]: ... -def list_jumpstart_tasks(filter: Union[Operator, str] = ..., region: str = "eu-west-1") -> List[str]: ... -def list_jumpstart_frameworks(filter: Union[Operator, str] = ..., region: str = "eu-west-1") -> List[str]: ... -def list_jumpstart_scripts(filter: Union[Operator, str] = ..., region: str = "eu-west-1") -> List[str]: ... +def list_jumpstart_tasks(filter: Operator | str = ..., region: str = "eu-west-1") -> List[str]: ... +def list_jumpstart_frameworks(filter: Operator | str = ..., region: str = "eu-west-1") -> List[str]: ... +def list_jumpstart_scripts(filter: Operator | str = ..., region: str = "eu-west-1") -> List[str]: ... def list_jumpstart_models( - filter: Union[Operator, str] = ..., + filter: Operator | str = ..., region: str = "eu-west-1", list_incomplete_models: bool = False, list_old_models: bool = False, list_versions: bool = False, -) -> List[Union[Tuple[str], Tuple[str, str]]]: ... +) -> List[Tuple[str, Tuple[str, str]]]: ... def get_model_url(model_id: str, model_version: str, region: str = "eu-west-1") -> str: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/types.pyi b/stubs/sagemaker/sagemaker/jumpstart/types.pyi index 901695ca2ed9..d1ecff0c2321 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/types.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/types.pyi @@ -126,7 +126,7 @@ class JumpStartCachedS3ContentValue(JumpStartDataHolderType): md5_hash: Incomplete def __init__( self, - formatted_content: Union[Dict[JumpStartVersionedModelId, JumpStartModelHeader], JumpStartModelSpecs], + formatted_content: Dict[JumpStartVersionedModelId | JumpStartModelHeader, JumpStartModelSpecs], md5_hash: Optional[str] = None, ) -> None: ... @@ -165,21 +165,21 @@ class JumpStartModelInitKwargs(JumpStartKwargs): model_version: Optional[str] = None, region: Optional[str] = None, instance_type: Optional[str] = None, - image_uri: Optional[Union[str, Any]] = None, - model_data: Optional[Union[str, Any]] = None, + image_uri: Optional[str | Any] = None, + model_data: Optional[str | Any] = None, role: Optional[str] = None, predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, Union[str, Any]]] = None, + env: Optional[Dict[str, str | Any]] = None, name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[Union[str, Any]]]] = None, + vpc_config: Optional[Dict[str, List[str | Any]]] = None, sagemaker_session: Optional[Any] = None, - enable_network_isolation: Union[bool, Any] = None, + enable_network_isolation: bool | Any = None, model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, Union[str, Any]]] = None, + image_config: Optional[Dict[str, str | Any]] = None, source_dir: Optional[str] = None, code_location: Optional[str] = None, entry_point: Optional[str] = None, - container_log_level: Optional[Union[int, Any]] = None, + container_log_level: Optional[int | Any] = None, dependencies: Optional[List[str]] = None, git_config: Optional[Dict[str, str]] = None, tolerate_vulnerable_model: Optional[bool] = None, @@ -291,49 +291,49 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs): model_id: str, model_version: Optional[str] = None, region: Optional[str] = None, - image_uri: Optional[Union[str, Any]] = None, + image_uri: Optional[str | Any] = None, role: Optional[str] = None, - instance_count: Optional[Union[int, Any]] = None, - instance_type: Optional[Union[str, Any]] = None, - keep_alive_period_in_seconds: Optional[Union[int, Any]] = None, - volume_size: Optional[Union[int, Any]] = None, - volume_kms_key: Optional[Union[str, Any]] = None, - max_run: Optional[Union[int, Any]] = None, - input_mode: Optional[Union[str, Any]] = None, - output_path: Optional[Union[str, Any]] = None, - output_kms_key: Optional[Union[str, Any]] = None, + instance_count: Optional[int | Any] = None, + instance_type: Optional[str | Any] = None, + keep_alive_period_in_seconds: Optional[int | Any] = None, + volume_size: Optional[int | Any] = None, + volume_kms_key: Optional[str | Any] = None, + max_run: Optional[int | Any] = None, + input_mode: Optional[str | Any] = None, + output_path: Optional[str | Any] = None, + output_kms_key: Optional[str | Any] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Any] = None, - hyperparameters: Optional[Dict[str, Union[str, Any]]] = None, - tags: Optional[List[Dict[str, Union[str, Any]]]] = None, - subnets: Optional[List[Union[str, Any]]] = None, - security_group_ids: Optional[List[Union[str, Any]]] = None, + hyperparameters: Optional[Dict[str, str | Any]] = None, + tags: Optional[List[Dict[str, str | Any]]] = None, + subnets: Optional[List[str | Any]] = None, + security_group_ids: Optional[List[str | Any]] = None, model_uri: Optional[str] = None, - model_channel_name: Optional[Union[str, Any]] = None, - metric_definitions: Optional[List[Dict[str, Union[str, Any]]]] = None, - encrypt_inter_container_traffic: Union[bool, Any] = None, - use_spot_instances: Optional[Union[bool, Any]] = None, - max_wait: Optional[Union[int, Any]] = None, - checkpoint_s3_uri: Optional[Union[str, Any]] = None, - checkpoint_local_path: Optional[Union[str, Any]] = None, - enable_network_isolation: Union[bool, Any] = None, + model_channel_name: Optional[str | Any] = None, + metric_definitions: Optional[List[Dict[str, str | Any]]] = None, + encrypt_inter_container_traffic: bool | Any = None, + use_spot_instances: Optional[bool | Any] = None, + max_wait: Optional[int | Any] = None, + checkpoint_s3_uri: Optional[str | Any] = None, + checkpoint_local_path: Optional[str | Any] = None, + enable_network_isolation: bool | Any = None, rules: Optional[List[Any]] = None, - debugger_hook_config: Optional[Union[Any, bool]] = None, + debugger_hook_config: Optional[Any | bool] = None, tensorboard_output_config: Optional[Any] = None, - enable_sagemaker_metrics: Optional[Union[bool, Any]] = None, + enable_sagemaker_metrics: Optional[bool | Any] = None, profiler_config: Optional[Any] = None, disable_profiler: Optional[bool] = None, - environment: Optional[Dict[str, Union[str, Any]]] = None, - max_retry_attempts: Optional[Union[int, Any]] = None, - source_dir: Optional[Union[str, Any]] = None, + environment: Optional[Dict[str, str | Any]] = None, + max_retry_attempts: Optional[int | Any] = None, + source_dir: Optional[str | Any] = None, git_config: Optional[Dict[str, str]] = None, - container_log_level: Optional[Union[int, Any]] = None, + container_log_level: Optional[int | Any] = None, code_location: Optional[str] = None, - entry_point: Optional[Union[str, Any]] = None, + entry_point: Optional[str | Any] = None, dependencies: Optional[List[str]] = None, instance_groups: Optional[List[Any]] = None, - training_repository_access_mode: Optional[Union[str, Any]] = None, - training_repository_credentials_provider_arn: Optional[Union[str, Any]] = None, + training_repository_access_mode: Optional[str | Any] = None, + training_repository_credentials_provider_arn: Optional[str | Any] = None, tolerate_vulnerable_model: Optional[bool] = None, tolerate_deprecated_model: Optional[bool] = None, ) -> None: ... @@ -355,7 +355,7 @@ class JumpStartEstimatorFitKwargs(JumpStartKwargs): model_id: str, model_version: Optional[str] = None, region: Optional[str] = None, - inputs: Optional[Union[str, Dict, Any, Any]] = None, + inputs: Optional[str | Dict | Any | Any] = None, wait: Optional[bool] = None, logs: Optional[str] = None, job_name: Optional[str] = None, @@ -427,20 +427,20 @@ class JumpStartEstimatorDeployKwargs(JumpStartKwargs): container_startup_health_check_timeout: Optional[int] = None, inference_recommendation_id: Optional[str] = None, explainer_config: Optional[Any] = None, - image_uri: Optional[Union[str, Any]] = None, + image_uri: Optional[str | Any] = None, role: Optional[str] = None, predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, Union[str, Any]]] = None, + env: Optional[Dict[str, str | Any]] = None, model_name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[Union[str, Any]]]] = None, + vpc_config: Optional[Dict[str, List[str | Any]]] = None, sagemaker_session: Optional[Any] = None, - enable_network_isolation: Union[bool, Any] = None, + enable_network_isolation: bool | Any = None, model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, Union[str, Any]]] = None, + image_config: Optional[Dict[str, str | Any]] = None, source_dir: Optional[str] = None, code_location: Optional[str] = None, entry_point: Optional[str] = None, - container_log_level: Optional[Union[int, Any]] = None, + container_log_level: Optional[int | Any] = None, dependencies: Optional[List[str]] = None, git_config: Optional[Dict[str, str]] = None, tolerate_deprecated_model: Optional[bool] = None, diff --git a/stubs/sagemaker/sagemaker/lineage/query.pyi b/stubs/sagemaker/sagemaker/lineage/query.pyi index b17fefaccba4..22330d7d94dd 100644 --- a/stubs/sagemaker/sagemaker/lineage/query.pyi +++ b/stubs/sagemaker/sagemaker/lineage/query.pyi @@ -70,8 +70,8 @@ class LineageFilter: properties: Incomplete def __init__( self, - entities: Optional[List[Union[LineageEntityEnum, str]]] = None, - sources: Optional[List[Union[LineageSourceEnum, str]]] = None, + entities: Optional[List[LineageEntityEnum | str]] = None, + sources: Optional[List[LineageSourceEnum | str]] = None, created_before: Optional[datetime] = None, created_after: Optional[datetime] = None, modified_before: Optional[datetime] = None, diff --git a/stubs/sagemaker/sagemaker/metadata_properties.pyi b/stubs/sagemaker/sagemaker/metadata_properties.pyi index 39a6b1a215ad..baf0929c699b 100644 --- a/stubs/sagemaker/sagemaker/metadata_properties.pyi +++ b/stubs/sagemaker/sagemaker/metadata_properties.pyi @@ -10,8 +10,8 @@ class MetadataProperties: project_id: Incomplete def __init__( self, - 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, + commit_id: Optional[str | PipelineVariable] = None, + repository: Optional[str | PipelineVariable] = None, + generated_by: Optional[str | PipelineVariable] = None, + project_id: Optional[str | PipelineVariable] = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/model.pyi b/stubs/sagemaker/sagemaker/model.pyi index 6fdfab7de786..56c60f5a4210 100644 --- a/stubs/sagemaker/sagemaker/model.pyi +++ b/stubs/sagemaker/sagemaker/model.pyi @@ -53,55 +53,55 @@ class Model(ModelBase, InferenceRecommenderMixin): repacked_model_data: Incomplete def __init__( self, - image_uri: Union[str, PipelineVariable], - model_data: Optional[Union[str, PipelineVariable]] = None, + image_uri: str | PipelineVariable, + model_data: Optional[str | PipelineVariable] = None, role: Optional[str] = None, predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = None, + vpc_config: Optional[Dict[str, List[str | PipelineVariable]]] = None, sagemaker_session: Optional[Session] = None, - enable_network_isolation: Union[bool, PipelineVariable] = None, + enable_network_isolation: bool | PipelineVariable = None, model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + image_config: Optional[Dict[str, str | PipelineVariable]] = None, source_dir: Optional[str] = None, code_location: Optional[str] = None, entry_point: Optional[str] = None, - container_log_level: Union[int, PipelineVariable] = 20, + container_log_level: int | PipelineVariable = 20, dependencies: Optional[List[str]] = None, git_config: Optional[Dict[str, str]] = None, ) -> None: ... def register( self, - 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, + content_types: List[str | PipelineVariable], + response_types: List[str | PipelineVariable], + inference_instances: Optional[List[str | PipelineVariable]] = None, + transform_instances: Optional[List[str | PipelineVariable]] = None, + model_package_name: Optional[str | PipelineVariable] = None, + model_package_group_name: Optional[str | PipelineVariable] = None, + image_uri: Optional[str | PipelineVariable] = None, model_metrics: Optional[ModelMetrics] = None, metadata_properties: Optional[MetadataProperties] = None, marketplace_cert: bool = False, - approval_status: Optional[Union[str, PipelineVariable]] = None, + approval_status: Optional[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, - task: Optional[Union[str, PipelineVariable]] = None, - sample_payload_url: Optional[Union[str, PipelineVariable]] = None, - framework: Optional[Union[str, PipelineVariable]] = None, - framework_version: Optional[Union[str, PipelineVariable]] = None, - nearest_model_name: Optional[Union[str, PipelineVariable]] = None, - data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, + validation_specification: Optional[str | PipelineVariable] = None, + domain: Optional[str | PipelineVariable] = None, + task: Optional[str | PipelineVariable] = None, + sample_payload_url: Optional[str | PipelineVariable] = None, + framework: Optional[str | PipelineVariable] = None, + framework_version: Optional[str | PipelineVariable] = None, + nearest_model_name: Optional[str | PipelineVariable] = None, + data_input_configuration: Optional[str | PipelineVariable] = None, ): ... def create( self, 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, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, ): ... def prepare_container_def( self, @@ -178,15 +178,15 @@ class Model(ModelBase, InferenceRecommenderMixin): class FrameworkModel(Model): def __init__( self, - model_data: Union[str, PipelineVariable], - image_uri: Union[str, PipelineVariable], + model_data: str | PipelineVariable, + image_uri: str | PipelineVariable, role: Optional[str] = None, entry_point: Optional[str] = None, source_dir: Optional[str] = None, predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, name: Optional[str] = None, - container_log_level: Union[int, PipelineVariable] = 20, + container_log_level: int | PipelineVariable = 20, code_location: Optional[str] = None, sagemaker_session: Optional[Session] = None, dependencies: Optional[List[str]] = None, diff --git a/stubs/sagemaker/sagemaker/model_card/helpers.pyi b/stubs/sagemaker/sagemaker/model_card/helpers.pyi index b127d8f01cbc..990111b69921 100644 --- a/stubs/sagemaker/sagemaker/model_card/helpers.pyi +++ b/stubs/sagemaker/sagemaker/model_card/helpers.pyi @@ -34,7 +34,7 @@ class _OneOf(_DescriptorBase): enumerator_reverse: Incomplete def __init__(self, enumerator: Enum) -> None: ... def validate(self, value) -> None: ... - def require_decode(self, value: Union[Enum, str]): ... + def require_decode(self, value: Enum | str): ... def decode(self, value: str): ... class _IsList(_DescriptorBase): @@ -48,8 +48,8 @@ class _IsList(_DescriptorBase): class _IsModelCardObject(_DescriptorBase): custom_class: Incomplete def __init__(self, custom_class: object) -> None: ... - def validate(self, value: Union[dict, object]): ... - def require_decode(self, value: Union[dict, object]): ... + def validate(self, value: dict | object): ... + def require_decode(self, value: dict | object): ... def decode(self, value: dict): ... class _MaxSizeArray(collections.abc.MutableSequence): diff --git a/stubs/sagemaker/sagemaker/model_card/model_card.pyi b/stubs/sagemaker/sagemaker/model_card/model_card.pyi index 6ca1dcc1e2c0..454d6660ecc3 100644 --- a/stubs/sagemaker/sagemaker/model_card/model_card.pyi +++ b/stubs/sagemaker/sagemaker/model_card/model_card.pyi @@ -36,7 +36,7 @@ class ModelOverview(_DefaultToRequestDict, _DefaultFromDict): model_id: Optional[str] = None, model_name: Optional[str] = None, model_description: Optional[str] = None, - model_version: Optional[Union[int, float]] = None, + model_version: Optional[int | float] = None, problem_type: Optional[str] = None, algorithm_type: Optional[str] = None, model_creator: Optional[str] = None, @@ -58,7 +58,7 @@ class IntendedUses(_DefaultToRequestDict, _DefaultFromDict): purpose_of_model: Optional[str] = None, intended_uses: Optional[str] = None, factors_affecting_model_efficiency: Optional[str] = None, - risk_rating: Optional[Union[RiskRatingEnum, str]] = ..., + risk_rating: Optional[RiskRatingEnum | str] = ..., explanations_for_risk_rating: Optional[str] = None, ) -> None: ... @@ -79,8 +79,8 @@ class Function(_DefaultToRequestDict, _DefaultFromDict): condition: Incomplete def __init__( self, - function: Optional[Union[ObjectiveFunctionEnum, str]] = None, - facet: Optional[Union[FacetEnum, str]] = None, + function: Optional[ObjectiveFunctionEnum | str] = None, + facet: Optional[FacetEnum | str] = None, condition: Optional[str] = None, ) -> None: ... @@ -98,22 +98,22 @@ class Metric(_DefaultToRequestDict, _DefaultFromDict): def __init__( self, name: str, - type: Union[MetricTypeEnum, str], - value: Union[int, float, str, bool, List], + type: MetricTypeEnum | str, + value: int | float | str | bool | List, notes: Optional[str] = None, - x_axis_name: Optional[Union[str, list]] = None, - y_axis_name: Optional[Union[str, list]] = None, + x_axis_name: Optional[str | list] = None, + y_axis_name: Optional[str | list] = None, ) -> None: ... @property def value(self): ... @value.setter - def value(self, val: Union[int, float, str, bool, List]): ... + def value(self, val: int | float | str | bool | List): ... class TrainingMetric(_DefaultToRequestDict, _DefaultFromDict): name: Incomplete value: Incomplete notes: Incomplete - def __init__(self, name: str, value: Union[int, float], notes: Optional[str] = None) -> None: ... + def __init__(self, name: str, value: int | float, notes: Optional[str] = None) -> None: ... class HyperParameter(_DefaultToRequestDict, _DefaultFromDict): name: Incomplete @@ -216,7 +216,7 @@ class ModelCard: def __init__( self, name: str, - status: Optional[Union[ModelCardStatusEnum, str]] = ..., + status: Optional[ModelCardStatusEnum | str] = ..., arn: Optional[str] = None, version: Optional[int] = None, created_time: Optional[datetime] = None, diff --git a/stubs/sagemaker/sagemaker/model_metrics.pyi b/stubs/sagemaker/sagemaker/model_metrics.pyi index 073bd95bc731..f022d6f9a4f3 100644 --- a/stubs/sagemaker/sagemaker/model_metrics.pyi +++ b/stubs/sagemaker/sagemaker/model_metrics.pyi @@ -30,9 +30,9 @@ class MetricsSource: content_digest: Incomplete def __init__( self, - content_type: Union[str, PipelineVariable], - s3_uri: Union[str, PipelineVariable], - content_digest: Optional[Union[str, PipelineVariable]] = None, + content_type: str | PipelineVariable, + s3_uri: str | PipelineVariable, + content_digest: Optional[str | PipelineVariable] = None, ) -> None: ... class FileSource: @@ -41,7 +41,7 @@ class FileSource: content_digest: Incomplete def __init__( self, - s3_uri: Union[str, PipelineVariable], - content_digest: Optional[Union[str, PipelineVariable]] = None, - content_type: Optional[Union[str, PipelineVariable]] = None, + s3_uri: str | PipelineVariable, + content_digest: Optional[str | PipelineVariable] = None, + content_type: Optional[str | PipelineVariable] = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi b/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi index 7d995cc19868..3c680aa58307 100644 --- a/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi +++ b/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi @@ -338,7 +338,7 @@ class MonitoringInput: end_time_offset: str features_attribute: str inference_attribute: str - probability_attribute: Union[str, int] + probability_attribute: str | int probability_threshold_attribute: float def __init__( self, diff --git a/stubs/sagemaker/sagemaker/multidatamodel.pyi b/stubs/sagemaker/sagemaker/multidatamodel.pyi index b59f1bf5b502..851305ce895e 100644 --- a/stubs/sagemaker/sagemaker/multidatamodel.pyi +++ b/stubs/sagemaker/sagemaker/multidatamodel.pyi @@ -21,7 +21,7 @@ class MultiDataModel(Model): name: str, model_data_prefix: str, model: Optional[Model] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, role: Optional[str] = None, sagemaker_session: Optional[Session] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/mxnet/estimator.pyi b/stubs/sagemaker/sagemaker/mxnet/estimator.pyi index f50c25a42d32..1d43b7192c04 100644 --- a/stubs/sagemaker/sagemaker/mxnet/estimator.pyi +++ b/stubs/sagemaker/sagemaker/mxnet/estimator.pyi @@ -11,12 +11,12 @@ class MXNet(Framework): py_version: Incomplete def __init__( self, - entry_point: Union[str, PipelineVariable], + entry_point: str | PipelineVariable, framework_version: Optional[str] = None, py_version: Optional[str] = None, - source_dir: Optional[Union[str, PipelineVariable]] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + source_dir: Optional[str | PipelineVariable] = None, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, distribution: Optional[Dict[str, str]] = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/mxnet/model.pyi b/stubs/sagemaker/sagemaker/mxnet/model.pyi index 7fb974332ed4..5855c68f96b4 100644 --- a/stubs/sagemaker/sagemaker/mxnet/model.pyi +++ b/stubs/sagemaker/sagemaker/mxnet/model.pyi @@ -19,40 +19,40 @@ class MXNetModel(FrameworkModel): model_server_workers: Incomplete def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, entry_point: Optional[str] = None, framework_version: str = "1.4.0", py_version: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, predictor_cls: callable = ..., - model_server_workers: Optional[Union[int, PipelineVariable]] = None, + model_server_workers: Optional[int | PipelineVariable] = None, **kwargs, ) -> None: ... image_uri: Incomplete def register( self, - 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, + content_types: List[str | PipelineVariable], + response_types: List[str | PipelineVariable], + inference_instances: Optional[List[str | PipelineVariable]] = None, + transform_instances: Optional[List[str | PipelineVariable]] = None, + model_package_name: Optional[str | PipelineVariable] = None, + model_package_group_name: Optional[str | PipelineVariable] = None, + image_uri: Optional[str | PipelineVariable] = None, model_metrics: Optional[ModelMetrics] = None, metadata_properties: Optional[MetadataProperties] = None, marketplace_cert: bool = False, - approval_status: Optional[Union[str, PipelineVariable]] = None, + approval_status: Optional[str | PipelineVariable] = None, description: Optional[str] = None, drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - domain: Optional[Union[str, PipelineVariable]] = None, - sample_payload_url: Optional[Union[str, PipelineVariable]] = None, - task: Optional[Union[str, PipelineVariable]] = None, - framework: Optional[Union[str, PipelineVariable]] = None, - framework_version: Optional[Union[str, PipelineVariable]] = None, - nearest_model_name: Optional[Union[str, PipelineVariable]] = None, - data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, + domain: Optional[str | PipelineVariable] = None, + sample_payload_url: Optional[str | PipelineVariable] = None, + task: Optional[str | PipelineVariable] = None, + framework: Optional[str | PipelineVariable] = None, + framework_version: Optional[str | PipelineVariable] = None, + nearest_model_name: Optional[str | PipelineVariable] = None, + data_input_configuration: Optional[str | PipelineVariable] = None, ): ... def prepare_container_def( self, diff --git a/stubs/sagemaker/sagemaker/mxnet/processing.pyi b/stubs/sagemaker/sagemaker/mxnet/processing.pyi index 208e5a2a4b2e..bf71be3b63cb 100644 --- a/stubs/sagemaker/sagemaker/mxnet/processing.pyi +++ b/stubs/sagemaker/sagemaker/mxnet/processing.pyi @@ -11,20 +11,20 @@ class MXNetProcessor(FrameworkProcessor): def __init__( self, framework_version: str, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Union[int, PipelineVariable] = None, - instance_type: Union[str, PipelineVariable] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: int | PipelineVariable = None, + instance_type: str | PipelineVariable = None, py_version: str = "py3", - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, command: Optional[List[str]] = None, - volume_size_in_gb: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + volume_size_in_gb: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, code_location: Optional[str] = None, - max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + max_runtime_in_seconds: Optional[int | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, network_config: Optional[NetworkConfig] = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/network.pyi b/stubs/sagemaker/sagemaker/network.pyi index 0fdc3ec1ab8b..bd1dd7b114d6 100644 --- a/stubs/sagemaker/sagemaker/network.pyi +++ b/stubs/sagemaker/sagemaker/network.pyi @@ -10,8 +10,8 @@ class NetworkConfig: encrypt_inter_container_traffic: Incomplete def __init__( self, - enable_network_isolation: Union[bool, PipelineVariable] = None, - security_group_ids: Optional[List[Union[str, PipelineVariable]]] = None, - subnets: Optional[List[Union[str, PipelineVariable]]] = None, - encrypt_inter_container_traffic: Optional[Union[bool, PipelineVariable]] = None, + enable_network_isolation: bool | PipelineVariable = None, + security_group_ids: Optional[List[str | PipelineVariable]] = None, + subnets: Optional[List[str | PipelineVariable]] = None, + encrypt_inter_container_traffic: Optional[bool | PipelineVariable] = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/parameter.pyi b/stubs/sagemaker/sagemaker/parameter.pyi index 21e9a7dcdba9..9e32d87cf1da 100644 --- a/stubs/sagemaker/sagemaker/parameter.pyi +++ b/stubs/sagemaker/sagemaker/parameter.pyi @@ -10,9 +10,9 @@ class ParameterRange: scaling_type: Incomplete def __init__( self, - min_value: Union[int, float, PipelineVariable], - max_value: Union[int, float, PipelineVariable], - scaling_type: Union[str, PipelineVariable] = "Auto", + min_value: int | float | PipelineVariable, + max_value: int | float | PipelineVariable, + scaling_type: str | PipelineVariable = "Auto", ) -> None: ... def is_valid(self, value): ... @classmethod diff --git a/stubs/sagemaker/sagemaker/pipeline.pyi b/stubs/sagemaker/sagemaker/pipeline.pyi index eaae3e1654ab..5ff19a13a317 100644 --- a/stubs/sagemaker/sagemaker/pipeline.pyi +++ b/stubs/sagemaker/sagemaker/pipeline.pyi @@ -22,9 +22,9 @@ class PipelineModel: role: str = None, predictor_cls: Optional[callable] = None, name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = None, + vpc_config: Optional[Dict[str, List[str | PipelineVariable]]] = None, sagemaker_session: Optional[Session] = None, - enable_network_isolation: Union[bool, PipelineVariable] = None, + enable_network_isolation: bool | PipelineVariable = None, ) -> None: ... def pipeline_container_def(self, instance_type: Incomplete | None = None): ... def deploy( @@ -46,27 +46,27 @@ class PipelineModel: def create(self, instance_type: str): ... def register( self, - 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, + content_types: List[str | PipelineVariable], + response_types: List[str | PipelineVariable], + inference_instances: Optional[List[str | PipelineVariable]] = None, + transform_instances: Optional[List[str | PipelineVariable]] = None, + model_package_name: Optional[str | PipelineVariable] = None, + model_package_group_name: Optional[str | PipelineVariable] = None, + image_uri: Optional[str | PipelineVariable] = None, model_metrics: Optional[ModelMetrics] = None, metadata_properties: Optional[MetadataProperties] = None, marketplace_cert: bool = False, - approval_status: Optional[Union[str, PipelineVariable]] = None, + approval_status: Optional[str | PipelineVariable] = None, description: Optional[str] = None, drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - domain: Optional[Union[str, PipelineVariable]] = None, - sample_payload_url: Optional[Union[str, PipelineVariable]] = None, - task: Optional[Union[str, PipelineVariable]] = None, - framework: Optional[Union[str, PipelineVariable]] = None, - framework_version: Optional[Union[str, PipelineVariable]] = None, - nearest_model_name: Optional[Union[str, PipelineVariable]] = None, - data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, + domain: Optional[str | PipelineVariable] = None, + sample_payload_url: Optional[str | PipelineVariable] = None, + task: Optional[str | PipelineVariable] = None, + framework: Optional[str | PipelineVariable] = None, + framework_version: Optional[str | PipelineVariable] = None, + nearest_model_name: Optional[str | PipelineVariable] = None, + data_input_configuration: Optional[str | PipelineVariable] = None, ): ... def transformer( self, diff --git a/stubs/sagemaker/sagemaker/processing.pyi b/stubs/sagemaker/sagemaker/processing.pyi index a373476d2e75..af93e89b5c3e 100644 --- a/stubs/sagemaker/sagemaker/processing.pyi +++ b/stubs/sagemaker/sagemaker/processing.pyi @@ -32,25 +32,25 @@ class Processor: def __init__( self, role: str = None, - image_uri: Union[str, PipelineVariable] = None, - instance_count: Union[int, PipelineVariable] = None, - instance_type: Union[str, PipelineVariable] = None, - entrypoint: Optional[List[Union[str, PipelineVariable]]] = None, - volume_size_in_gb: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, - max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + image_uri: str | PipelineVariable = None, + instance_count: int | PipelineVariable = None, + instance_type: str | PipelineVariable = None, + entrypoint: Optional[List[str | PipelineVariable]] = None, + volume_size_in_gb: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, + max_runtime_in_seconds: Optional[int | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, network_config: Optional[NetworkConfig] = None, ) -> None: ... def run( self, inputs: Optional[List["ProcessingInput"]] = None, outputs: Optional[List["ProcessingOutput"]] = None, - arguments: Optional[List[Union[str, PipelineVariable]]] = None, + arguments: Optional[List[str | PipelineVariable]] = None, wait: bool = True, logs: bool = True, job_name: Optional[str] = None, @@ -62,19 +62,19 @@ class ScriptProcessor(Processor): command: Incomplete def __init__( self, - role: Optional[Union[str, PipelineVariable]] = None, - image_uri: Union[str, PipelineVariable] = None, + role: Optional[str | PipelineVariable] = None, + image_uri: str | PipelineVariable = None, command: List[str] = None, - instance_count: Union[int, PipelineVariable] = None, - instance_type: Union[str, PipelineVariable] = None, - volume_size_in_gb: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, - max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + instance_count: int | PipelineVariable = None, + instance_type: str | PipelineVariable = None, + volume_size_in_gb: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, + max_runtime_in_seconds: Optional[int | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, network_config: Optional[NetworkConfig] = None, ) -> None: ... def get_run_args( @@ -86,7 +86,7 @@ class ScriptProcessor(Processor): code: str, inputs: Optional[List["ProcessingInput"]] = None, outputs: Optional[List["ProcessingOutput"]] = None, - arguments: Optional[List[Union[str, PipelineVariable]]] = None, + arguments: Optional[List[str | PipelineVariable]] = None, wait: bool = True, logs: bool = True, job_name: Optional[str] = None, @@ -130,16 +130,16 @@ class ProcessingInput: app_managed: Incomplete def __init__( self, - source: Optional[Union[str, PipelineVariable]] = None, - destination: Optional[Union[str, PipelineVariable]] = None, - input_name: Optional[Union[str, PipelineVariable]] = None, - s3_data_type: Union[str, PipelineVariable] = "S3Prefix", - s3_input_mode: Union[str, PipelineVariable] = "File", - s3_data_distribution_type: Union[str, PipelineVariable] = "FullyReplicated", - s3_compression_type: Union[str, PipelineVariable] = "None", + source: Optional[str | PipelineVariable] = None, + destination: Optional[str | PipelineVariable] = None, + input_name: Optional[str | PipelineVariable] = None, + s3_data_type: str | PipelineVariable = "S3Prefix", + s3_input_mode: str | PipelineVariable = "File", + s3_data_distribution_type: str | PipelineVariable = "FullyReplicated", + s3_compression_type: str | PipelineVariable = "None", s3_input: Optional[S3Input] = None, dataset_definition: Optional[DatasetDefinition] = None, - app_managed: Union[bool, PipelineVariable] = False, + app_managed: bool | PipelineVariable = False, ) -> None: ... class ProcessingOutput: @@ -151,11 +151,11 @@ class ProcessingOutput: feature_store_output: Incomplete def __init__( self, - source: Optional[Union[str, PipelineVariable]] = None, - destination: Optional[Union[str, PipelineVariable]] = None, - output_name: Optional[Union[str, PipelineVariable]] = None, - s3_upload_mode: Union[str, PipelineVariable] = "EndOfJob", - app_managed: Union[bool, PipelineVariable] = False, + source: Optional[str | PipelineVariable] = None, + destination: Optional[str | PipelineVariable] = None, + output_name: Optional[str | PipelineVariable] = None, + s3_upload_mode: str | PipelineVariable = "EndOfJob", + app_managed: bool | PipelineVariable = False, feature_store_output: Optional["FeatureStoreOutput"] = None, ) -> None: ... @@ -185,21 +185,21 @@ class FrameworkProcessor(ScriptProcessor): self, estimator_cls: type, framework_version: str, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Union[int, PipelineVariable] = None, - instance_type: Union[str, PipelineVariable] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: int | PipelineVariable = None, + instance_type: str | PipelineVariable = None, py_version: str = "py3", - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, command: Optional[List[str]] = None, - volume_size_in_gb: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + volume_size_in_gb: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, code_location: Optional[str] = None, - max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + max_runtime_in_seconds: Optional[int | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, network_config: Optional[NetworkConfig] = None, ) -> None: ... def get_run_args( @@ -221,7 +221,7 @@ class FrameworkProcessor(ScriptProcessor): git_config: Optional[Dict[str, str]] = None, inputs: Optional[List[ProcessingInput]] = None, outputs: Optional[List[ProcessingOutput]] = None, - arguments: Optional[List[Union[str, PipelineVariable]]] = None, + arguments: Optional[List[str | PipelineVariable]] = None, wait: bool = True, logs: bool = True, job_name: Optional[str] = None, diff --git a/stubs/sagemaker/sagemaker/pytorch/estimator.pyi b/stubs/sagemaker/sagemaker/pytorch/estimator.pyi index c1c05e86503a..f43715d384c4 100644 --- a/stubs/sagemaker/sagemaker/pytorch/estimator.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/estimator.pyi @@ -17,12 +17,12 @@ class PyTorch(Framework): compiler_config: Incomplete def __init__( self, - entry_point: Union[str, PipelineVariable], + entry_point: str | PipelineVariable, framework_version: Optional[str] = None, py_version: Optional[str] = None, - source_dir: Optional[Union[str, PipelineVariable]] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + source_dir: Optional[str | PipelineVariable] = None, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, distribution: Optional[Dict] = None, compiler_config: Optional[TrainingCompilerConfig] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/pytorch/model.pyi b/stubs/sagemaker/sagemaker/pytorch/model.pyi index 97f981405a07..4a97b9d33616 100644 --- a/stubs/sagemaker/sagemaker/pytorch/model.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/model.pyi @@ -19,40 +19,40 @@ class PyTorchModel(FrameworkModel): model_server_workers: Incomplete def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, entry_point: Optional[str] = None, framework_version: str = "1.3", py_version: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, predictor_cls: callable = ..., - model_server_workers: Optional[Union[int, PipelineVariable]] = None, + model_server_workers: Optional[int | PipelineVariable] = None, **kwargs, ) -> None: ... image_uri: Incomplete def register( self, - 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, + content_types: List[str | PipelineVariable], + response_types: List[str | PipelineVariable], + inference_instances: Optional[List[str | PipelineVariable]] = None, + transform_instances: Optional[List[str | PipelineVariable]] = None, + model_package_name: Optional[str | PipelineVariable] = None, + model_package_group_name: Optional[str | PipelineVariable] = None, + image_uri: Optional[str | PipelineVariable] = None, model_metrics: Optional[ModelMetrics] = None, metadata_properties: Optional[MetadataProperties] = None, marketplace_cert: bool = False, - approval_status: Optional[Union[str, PipelineVariable]] = None, + approval_status: Optional[str | PipelineVariable] = None, description: Optional[str] = None, drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - domain: Optional[Union[str, PipelineVariable]] = None, - sample_payload_url: Optional[Union[str, PipelineVariable]] = None, - task: Optional[Union[str, PipelineVariable]] = None, - framework: Optional[Union[str, PipelineVariable]] = None, - framework_version: Optional[Union[str, PipelineVariable]] = None, - nearest_model_name: Optional[Union[str, PipelineVariable]] = None, - data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, + domain: Optional[str | PipelineVariable] = None, + sample_payload_url: Optional[str | PipelineVariable] = None, + task: Optional[str | PipelineVariable] = None, + framework: Optional[str | PipelineVariable] = None, + framework_version: Optional[str | PipelineVariable] = None, + nearest_model_name: Optional[str | PipelineVariable] = None, + data_input_configuration: Optional[str | PipelineVariable] = None, ): ... def prepare_container_def( self, diff --git a/stubs/sagemaker/sagemaker/pytorch/processing.pyi b/stubs/sagemaker/sagemaker/pytorch/processing.pyi index d9d3586ade4e..94ce85a7e28e 100644 --- a/stubs/sagemaker/sagemaker/pytorch/processing.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/processing.pyi @@ -11,20 +11,20 @@ class PyTorchProcessor(FrameworkProcessor): def __init__( self, framework_version: str, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Union[int, PipelineVariable] = None, - instance_type: Union[str, PipelineVariable] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: int | PipelineVariable = None, + instance_type: str | PipelineVariable = None, py_version: str = "py3", - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, command: Optional[List[str]] = None, - volume_size_in_gb: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + volume_size_in_gb: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, code_location: Optional[str] = None, - max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + max_runtime_in_seconds: Optional[int | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, network_config: Optional[NetworkConfig] = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi b/stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi index 90146f5a7e03..272711ce08e7 100644 --- a/stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi @@ -9,6 +9,6 @@ logger: Incomplete class TrainingCompilerConfig(BaseConfig): SUPPORTED_INSTANCE_CLASS_PREFIXES: Incomplete SUPPORTED_INSTANCE_TYPES_WITH_EFA: Incomplete - def __init__(self, enabled: Union[bool, PipelineVariable] = True, debug: Union[bool, PipelineVariable] = False) -> None: ... + def __init__(self, enabled: bool | PipelineVariable = True, debug: bool | PipelineVariable = False) -> None: ... @classmethod def validate(cls, estimator) -> None: ... diff --git a/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi b/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi index 47becee8d587..7ca69b64f673 100644 --- a/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi +++ b/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi @@ -4,7 +4,7 @@ class SparkConfig: submit_jars: Optional[List[str]] submit_py_files: Optional[List[str]] submit_files: Optional[List[str]] - configuration: Optional[Union[List[Dict], Dict]] + configuration: Optional[List[Dict, Dict]] spark_event_logs_uri: Optional[str] def __init__(self, submit_jars, submit_py_files, submit_files, configuration, spark_event_logs_uri) -> None: ... def __lt__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/rl/estimator.pyi b/stubs/sagemaker/sagemaker/rl/estimator.pyi index 11bd475cf6b6..876350132d31 100644 --- a/stubs/sagemaker/sagemaker/rl/estimator.pyi +++ b/stubs/sagemaker/sagemaker/rl/estimator.pyi @@ -30,14 +30,14 @@ class RLEstimator(Framework): framework_version: Incomplete def __init__( self, - entry_point: Union[str, PipelineVariable], + entry_point: str | PipelineVariable, toolkit: Optional[RLToolkit] = None, toolkit_version: Optional[str] = None, framework: Optional[Framework] = None, - source_dir: Optional[Union[str, PipelineVariable]] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, - metric_definitions: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + source_dir: Optional[str | PipelineVariable] = None, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, + metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, **kwargs, ) -> None: ... def create_model( diff --git a/stubs/sagemaker/sagemaker/s3.pyi b/stubs/sagemaker/sagemaker/s3.pyi index ca4102758b49..6e72c7613a81 100644 --- a/stubs/sagemaker/sagemaker/s3.pyi +++ b/stubs/sagemaker/sagemaker/s3.pyi @@ -18,7 +18,7 @@ class S3Uploader: ): ... @staticmethod def upload_bytes( - b: Union[bytes, io.BytesIO], s3_uri, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None + b: bytes | io.BytesIO, s3_uri, kms_key: Incomplete | None = None, sagemaker_session: Incomplete | None = None ): ... class S3Downloader: diff --git a/stubs/sagemaker/sagemaker/sklearn/estimator.pyi b/stubs/sagemaker/sagemaker/sklearn/estimator.pyi index 3f37e15e0e74..50fac3ab5ca5 100644 --- a/stubs/sagemaker/sagemaker/sklearn/estimator.pyi +++ b/stubs/sagemaker/sagemaker/sklearn/estimator.pyi @@ -12,12 +12,12 @@ class SKLearn(Framework): image_uri: Incomplete def __init__( self, - entry_point: Union[str, PipelineVariable], + entry_point: str | PipelineVariable, framework_version: Optional[str] = None, py_version: str = "py3", - source_dir: Optional[Union[str, PipelineVariable]] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + source_dir: Optional[str | PipelineVariable] = None, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, image_uri_region: Optional[str] = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/sklearn/model.pyi b/stubs/sagemaker/sagemaker/sklearn/model.pyi index 8726615f7a8e..a39f0f9a6b86 100644 --- a/stubs/sagemaker/sagemaker/sklearn/model.pyi +++ b/stubs/sagemaker/sagemaker/sklearn/model.pyi @@ -19,40 +19,40 @@ class SKLearnModel(FrameworkModel): model_server_workers: Incomplete def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, entry_point: Optional[str] = None, framework_version: Optional[str] = None, py_version: str = "py3", - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, predictor_cls: callable = ..., - model_server_workers: Optional[Union[int, PipelineVariable]] = None, + model_server_workers: Optional[int | PipelineVariable] = None, **kwargs, ) -> None: ... image_uri: Incomplete def register( self, - 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, + content_types: List[str | PipelineVariable], + response_types: List[str | PipelineVariable], + inference_instances: Optional[List[str | PipelineVariable]] = None, + transform_instances: Optional[List[str | PipelineVariable]] = None, + model_package_name: Optional[str | PipelineVariable] = None, + model_package_group_name: Optional[str | PipelineVariable] = None, + image_uri: Optional[str | PipelineVariable] = None, model_metrics: Optional[ModelMetrics] = None, metadata_properties: Optional[MetadataProperties] = None, marketplace_cert: bool = False, - approval_status: Optional[Union[str, PipelineVariable]] = None, + approval_status: Optional[str | PipelineVariable] = None, description: Optional[str] = None, drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - domain: Optional[Union[str, PipelineVariable]] = None, - sample_payload_url: Optional[Union[str, PipelineVariable]] = None, - task: Optional[Union[str, PipelineVariable]] = None, - framework: Optional[Union[str, PipelineVariable]] = None, - framework_version: Optional[Union[str, PipelineVariable]] = None, - nearest_model_name: Optional[Union[str, PipelineVariable]] = None, - data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, + domain: Optional[str | PipelineVariable] = None, + sample_payload_url: Optional[str | PipelineVariable] = None, + task: Optional[str | PipelineVariable] = None, + framework: Optional[str | PipelineVariable] = None, + framework_version: Optional[str | PipelineVariable] = None, + nearest_model_name: Optional[str | PipelineVariable] = None, + data_input_configuration: Optional[str | PipelineVariable] = None, ): ... def prepare_container_def( self, diff --git a/stubs/sagemaker/sagemaker/sklearn/processing.pyi b/stubs/sagemaker/sagemaker/sklearn/processing.pyi index d8e3eeccfd72..fbfc41df84c7 100644 --- a/stubs/sagemaker/sagemaker/sklearn/processing.pyi +++ b/stubs/sagemaker/sagemaker/sklearn/processing.pyi @@ -9,17 +9,17 @@ class SKLearnProcessor(ScriptProcessor): def __init__( self, framework_version: str, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Union[int, PipelineVariable] = None, - instance_type: Union[str, PipelineVariable] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: int | PipelineVariable = None, + instance_type: str | PipelineVariable = None, command: Optional[List[str]] = None, - volume_size_in_gb: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, - max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + volume_size_in_gb: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, + max_runtime_in_seconds: Optional[int | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, network_config: Optional[NetworkConfig] = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/spark/processing.pyi b/stubs/sagemaker/sagemaker/spark/processing.pyi index 6e32888a55d1..c563b2335f11 100644 --- a/stubs/sagemaker/sagemaker/spark/processing.pyi +++ b/stubs/sagemaker/sagemaker/spark/processing.pyi @@ -57,22 +57,22 @@ class PySparkProcessor(_SparkProcessorBase): def __init__( self, role: str = None, - instance_type: Union[str, PipelineVariable] = None, - instance_count: Union[int, PipelineVariable] = None, + instance_type: str | PipelineVariable = None, + instance_count: int | PipelineVariable = None, framework_version: Optional[str] = None, py_version: Optional[str] = None, container_version: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, - volume_size_in_gb: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, + volume_size_in_gb: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, configuration_location: Optional[str] = None, dependency_location: Optional[str] = None, - max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + max_runtime_in_seconds: Optional[int | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, network_config: Optional[NetworkConfig] = None, ) -> None: ... def get_run_args( @@ -91,18 +91,18 @@ class PySparkProcessor(_SparkProcessorBase): def run( self, submit_app: str, - submit_py_files: Optional[List[Union[str, PipelineVariable]]] = None, - submit_jars: Optional[List[Union[str, PipelineVariable]]] = None, - submit_files: Optional[List[Union[str, PipelineVariable]]] = None, + submit_py_files: Optional[List[str | PipelineVariable]] = None, + submit_jars: Optional[List[str | PipelineVariable]] = None, + submit_files: Optional[List[str | PipelineVariable]] = None, inputs: Optional[List[ProcessingInput]] = None, outputs: Optional[List[ProcessingOutput]] = None, - arguments: Optional[List[Union[str, PipelineVariable]]] = None, + arguments: Optional[List[str | PipelineVariable]] = None, wait: bool = True, logs: bool = True, job_name: Optional[str] = None, experiment_config: Optional[Dict[str, str]] = None, - configuration: Optional[Union[List[Dict], Dict]] = None, - spark_event_logs_s3_uri: Optional[Union[str, PipelineVariable]] = None, + configuration: Optional[List[Dict, Dict]] = None, + spark_event_logs_s3_uri: Optional[str | PipelineVariable] = None, kms_key: Optional[str] = None, ): ... @@ -110,22 +110,22 @@ class SparkJarProcessor(_SparkProcessorBase): def __init__( self, role: str = None, - instance_type: Union[str, PipelineVariable] = None, - instance_count: Union[int, PipelineVariable] = None, + instance_type: str | PipelineVariable = None, + instance_count: int | PipelineVariable = None, framework_version: Optional[str] = None, py_version: Optional[str] = None, container_version: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, - volume_size_in_gb: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, + volume_size_in_gb: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, configuration_location: Optional[str] = None, dependency_location: Optional[str] = None, - max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + max_runtime_in_seconds: Optional[int | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, network_config: Optional[NetworkConfig] = None, ) -> None: ... def get_run_args( @@ -144,18 +144,18 @@ class SparkJarProcessor(_SparkProcessorBase): def run( self, submit_app: str, - submit_class: Union[str, PipelineVariable], - submit_jars: Optional[List[Union[str, PipelineVariable]]] = None, - submit_files: Optional[List[Union[str, PipelineVariable]]] = None, + submit_class: str | PipelineVariable, + submit_jars: Optional[List[str | PipelineVariable]] = None, + submit_files: Optional[List[str | PipelineVariable]] = None, inputs: Optional[List[ProcessingInput]] = None, outputs: Optional[List[ProcessingOutput]] = None, - arguments: Optional[List[Union[str, PipelineVariable]]] = None, + arguments: Optional[List[str | PipelineVariable]] = None, wait: bool = True, logs: bool = True, job_name: Optional[str] = None, experiment_config: Optional[Dict[str, str]] = None, - configuration: Optional[Union[List[Dict], Dict]] = None, - spark_event_logs_s3_uri: Optional[Union[str, PipelineVariable]] = None, + configuration: Optional[List[Dict, Dict]] = None, + spark_event_logs_s3_uri: Optional[str | PipelineVariable] = None, kms_key: Optional[str] = None, ): ... diff --git a/stubs/sagemaker/sagemaker/sparkml/model.pyi b/stubs/sagemaker/sagemaker/sparkml/model.pyi index 3e46af3941db..67d4ec2c16ad 100644 --- a/stubs/sagemaker/sagemaker/sparkml/model.pyi +++ b/stubs/sagemaker/sagemaker/sparkml/model.pyi @@ -12,7 +12,7 @@ class SparkMLPredictor(Predictor): class SparkMLModel(Model): def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: Optional[str] = None, spark_version: str = "3.3", sagemaker_session: Optional[Session] = None, diff --git a/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi b/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi index c9cf390cf2ae..ed02d9cd8d1e 100644 --- a/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi +++ b/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi @@ -18,8 +18,8 @@ class TensorFlow(Framework): self, py_version: Optional[str] = None, framework_version: Optional[str] = None, - model_dir: Optional[Union[str, PipelineVariable]] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + model_dir: Optional[str | PipelineVariable] = None, + image_uri: Optional[str | PipelineVariable] = None, distribution: Optional[Dict[str, str]] = None, compiler_config: Optional[TrainingCompilerConfig] = None, **kwargs, diff --git a/stubs/sagemaker/sagemaker/tensorflow/model.pyi b/stubs/sagemaker/sagemaker/tensorflow/model.pyi index a4c94427b71c..1c5d8577183e 100644 --- a/stubs/sagemaker/sagemaker/tensorflow/model.pyi +++ b/stubs/sagemaker/sagemaker/tensorflow/model.pyi @@ -33,10 +33,10 @@ class TensorFlowModel(sagemaker.model.FrameworkModel): inference_framework_version: Incomplete def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: str = None, entry_point: Optional[str] = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, framework_version: Optional[str] = None, container_log_level: Optional[int] = None, predictor_cls: callable = ..., @@ -45,27 +45,27 @@ class TensorFlowModel(sagemaker.model.FrameworkModel): image_uri: Incomplete def register( self, - 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, + content_types: List[str | PipelineVariable], + response_types: List[str | PipelineVariable], + inference_instances: Optional[List[str | PipelineVariable]] = None, + transform_instances: Optional[List[str | PipelineVariable]] = None, + model_package_name: Optional[str | PipelineVariable] = None, + model_package_group_name: Optional[str | PipelineVariable] = None, + image_uri: Optional[str | PipelineVariable] = None, model_metrics: Optional[ModelMetrics] = None, metadata_properties: Optional[MetadataProperties] = None, marketplace_cert: bool = False, - approval_status: Optional[Union[str, PipelineVariable]] = None, + approval_status: Optional[str | PipelineVariable] = None, description: Optional[str] = None, drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - domain: Optional[Union[str, PipelineVariable]] = None, - sample_payload_url: Optional[Union[str, PipelineVariable]] = None, - task: Optional[Union[str, PipelineVariable]] = None, - framework: Optional[Union[str, PipelineVariable]] = None, - framework_version: Optional[Union[str, PipelineVariable]] = None, - nearest_model_name: Optional[Union[str, PipelineVariable]] = None, - data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, + domain: Optional[str | PipelineVariable] = None, + sample_payload_url: Optional[str | PipelineVariable] = None, + task: Optional[str | PipelineVariable] = None, + framework: Optional[str | PipelineVariable] = None, + framework_version: Optional[str | PipelineVariable] = None, + nearest_model_name: Optional[str | PipelineVariable] = None, + data_input_configuration: Optional[str | PipelineVariable] = None, ): ... def deploy( self, diff --git a/stubs/sagemaker/sagemaker/tensorflow/processing.pyi b/stubs/sagemaker/sagemaker/tensorflow/processing.pyi index 3bacec0a4740..562bb08c86c8 100644 --- a/stubs/sagemaker/sagemaker/tensorflow/processing.pyi +++ b/stubs/sagemaker/sagemaker/tensorflow/processing.pyi @@ -11,20 +11,20 @@ class TensorFlowProcessor(FrameworkProcessor): def __init__( self, framework_version: str, - role: Optional[Union[str, PipelineVariable]] = None, - instance_count: Union[int, PipelineVariable] = None, - instance_type: Union[str, PipelineVariable] = None, + role: Optional[str | PipelineVariable] = None, + instance_count: int | PipelineVariable = None, + instance_type: str | PipelineVariable = None, py_version: str = "py3", - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, command: Optional[List[str]] = None, - volume_size_in_gb: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + volume_size_in_gb: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, code_location: Optional[str] = None, - max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + max_runtime_in_seconds: Optional[int | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, network_config: Optional[NetworkConfig] = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/transformer.pyi b/stubs/sagemaker/sagemaker/transformer.pyi index bd9ea6f0ae4d..4b9f2a7fb9a1 100644 --- a/stubs/sagemaker/sagemaker/transformer.pyi +++ b/stubs/sagemaker/sagemaker/transformer.pyi @@ -27,35 +27,35 @@ class Transformer: env: Incomplete def __init__( self, - model_name: Union[str, PipelineVariable], - instance_count: Union[int, PipelineVariable], - instance_type: Union[str, PipelineVariable], - strategy: Optional[Union[str, PipelineVariable]] = None, - assemble_with: Optional[Union[str, PipelineVariable]] = None, - output_path: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, - accept: Optional[Union[str, PipelineVariable]] = None, - max_concurrent_transforms: Optional[Union[int, PipelineVariable]] = None, - max_payload: Optional[Union[int, PipelineVariable]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + model_name: str | PipelineVariable, + instance_count: int | PipelineVariable, + instance_type: str | PipelineVariable, + strategy: Optional[str | PipelineVariable] = None, + assemble_with: Optional[str | PipelineVariable] = None, + output_path: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, + accept: Optional[str | PipelineVariable] = None, + max_concurrent_transforms: Optional[int | PipelineVariable] = None, + max_payload: Optional[int | PipelineVariable] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, base_transform_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, + volume_kms_key: Optional[str | PipelineVariable] = None, ) -> None: ... def transform( self, - data: Union[str, PipelineVariable], - data_type: Union[str, PipelineVariable] = "S3Prefix", - content_type: Optional[Union[str, PipelineVariable]] = None, - compression_type: Optional[Union[str, PipelineVariable]] = None, - split_type: Optional[Union[str, PipelineVariable]] = None, + data: str | PipelineVariable, + data_type: str | PipelineVariable = "S3Prefix", + content_type: Optional[str | PipelineVariable] = None, + compression_type: Optional[str | PipelineVariable] = None, + split_type: Optional[str | PipelineVariable] = None, job_name: Optional[str] = None, - input_filter: Optional[Union[str, PipelineVariable]] = None, - output_filter: Optional[Union[str, PipelineVariable]] = None, - join_source: Optional[Union[str, PipelineVariable]] = None, + input_filter: Optional[str | PipelineVariable] = None, + output_filter: Optional[str | PipelineVariable] = None, + join_source: Optional[str | PipelineVariable] = None, experiment_config: Optional[Dict[str, str]] = None, - model_client_config: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + model_client_config: Optional[Dict[str, str | PipelineVariable]] = None, batch_data_capture_config: BatchDataCaptureConfig = None, wait: bool = True, logs: bool = True, diff --git a/stubs/sagemaker/sagemaker/tuner.pyi b/stubs/sagemaker/sagemaker/tuner.pyi index bb9f6317d850..465e90501d20 100644 --- a/stubs/sagemaker/sagemaker/tuner.pyi +++ b/stubs/sagemaker/sagemaker/tuner.pyi @@ -34,7 +34,7 @@ class WarmStartTypes(Enum): class WarmStartConfig: type: Incomplete parents: Incomplete - def __init__(self, warm_start_type: WarmStartTypes, parents: Set[Union[str, PipelineVariable]]) -> None: ... + def __init__(self, warm_start_type: WarmStartTypes, parents: Set[str | PipelineVariable]) -> None: ... @classmethod def from_job_desc(cls, warm_start_config): ... def to_input_req(self): ... @@ -60,9 +60,9 @@ class InstanceConfig: volume_size: Incomplete def __init__( self, - instance_count: Union[int, PipelineVariable] = None, - instance_type: Union[str, PipelineVariable] = None, - volume_size: Union[int, PipelineVariable] = 30, + instance_count: int | PipelineVariable = None, + instance_type: str | PipelineVariable = None, + volume_size: int | PipelineVariable = 30, ) -> None: ... @classmethod def from_job_desc(cls, instance_config): ... @@ -119,31 +119,31 @@ class HyperparameterTuner: def __init__( self, estimator: EstimatorBase, - objective_metric_name: Union[str, PipelineVariable], + objective_metric_name: str | PipelineVariable, hyperparameter_ranges: Dict[str, ParameterRange], - metric_definitions: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, - strategy: Union[str, PipelineVariable] = "Bayesian", - objective_type: Union[str, PipelineVariable] = "Maximize", - max_jobs: Union[int, PipelineVariable] = None, - max_parallel_jobs: Union[int, PipelineVariable] = 1, - max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, + strategy: str | PipelineVariable = "Bayesian", + objective_type: str | PipelineVariable = "Maximize", + max_jobs: int | PipelineVariable = None, + max_parallel_jobs: int | PipelineVariable = 1, + max_runtime_in_seconds: Optional[int | PipelineVariable] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, base_tuning_job_name: Optional[str] = None, warm_start_config: Optional[WarmStartConfig] = None, strategy_config: Optional[StrategyConfig] = None, completion_criteria_config: Optional[TuningJobCompletionCriteriaConfig] = None, - early_stopping_type: Union[str, PipelineVariable] = "Off", + early_stopping_type: str | PipelineVariable = "Off", estimator_name: Optional[str] = None, random_seed: Optional[int] = None, autotune: bool = False, hyperparameters_to_keep_static: Optional[List[str]] = None, ) -> None: ... - def override_resource_config(self, instance_configs: Union[List[InstanceConfig], Dict[str, List[InstanceConfig]]]): ... + def override_resource_config(self, instance_configs: List[InstanceConfig, Dict[str, List[InstanceConfig]]]): ... def fit( self, - inputs: Optional[Union[str, Dict, List, TrainingInput, FileSystemInput, RecordSet, FileSystemRecordSet]] = None, + inputs: Optional[str | Dict | List | TrainingInput | FileSystemInput | RecordSet | FileSystemRecordSet] = None, job_name: Optional[str] = None, - include_cls_metadata: Union[bool, Dict[str, bool]] = False, + include_cls_metadata: bool | Dict[str | bool] = False, estimator_kwargs: Optional[Dict[str, dict]] = None, wait: bool = True, **kwargs, diff --git a/stubs/sagemaker/sagemaker/workflow/_utils.pyi b/stubs/sagemaker/sagemaker/workflow/_utils.pyi index 7f65070b08ba..c48cbe9c67df 100644 --- a/stubs/sagemaker/sagemaker/workflow/_utils.pyi +++ b/stubs/sagemaker/sagemaker/workflow/_utils.pyi @@ -28,7 +28,7 @@ class _RepackModelStep(TrainingStep): description: str = None, source_dir: str = None, dependencies: List = None, - depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + depends_on: Optional[List[str | Step | "StepCollection"]] = None, retry_policies: List[RetryPolicy] = None, subnets: Incomplete | None = None, security_group_ids: Incomplete | None = None, @@ -80,7 +80,7 @@ class _RegisterModelStep(ConfigurableRetryStep): compile_model_family: Incomplete | None = None, display_name: str = None, description: Incomplete | None = None, - depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + depends_on: Optional[List[str | Step | "StepCollection"]] = None, retry_policies: Optional[List[RetryPolicy]] = None, tags: Incomplete | None = None, container_def_list: Incomplete | None = None, diff --git a/stubs/sagemaker/sagemaker/workflow/automl_step.pyi b/stubs/sagemaker/sagemaker/workflow/automl_step.pyi index 15d77943ea7b..1a67ba36d1f9 100644 --- a/stubs/sagemaker/sagemaker/workflow/automl_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/automl_step.pyi @@ -17,7 +17,7 @@ class AutoMLStep(ConfigurableRetryStep): display_name: str = None, description: str = None, cache_config: CacheConfig = None, - depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + depends_on: Optional[List[str | Step | StepCollection]] = None, retry_policies: List[RetryPolicy] = None, ) -> None: ... @property diff --git a/stubs/sagemaker/sagemaker/workflow/callback_step.pyi b/stubs/sagemaker/sagemaker/workflow/callback_step.pyi index ad94b479878c..671e82c5e3c3 100644 --- a/stubs/sagemaker/sagemaker/workflow/callback_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/callback_step.pyi @@ -37,7 +37,7 @@ class CallbackStep(Step): display_name: str = None, description: str = None, cache_config: CacheConfig = None, - depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + depends_on: Optional[List[str | Step | StepCollection]] = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi b/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi index ae7fd5e1a7fd..4f34112593c8 100644 --- a/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi @@ -20,7 +20,7 @@ class ClarifyCheckConfig(ABC): class DataBiasCheckConfig(ClarifyCheckConfig): data_bias_config: BiasConfig - methods: Union[str, List[str]] + methods: str | List[str] def __init__(self, data_config, kms_key, monitoring_analysis_config_uri, data_bias_config, methods) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -31,7 +31,7 @@ class ModelBiasCheckConfig(ClarifyCheckConfig): data_bias_config: BiasConfig model_config: ModelConfig model_predicted_label_config: ModelPredictedLabelConfig - methods: Union[str, List[str]] + methods: str | List[str] def __init__( self, data_config, @@ -50,7 +50,7 @@ class ModelBiasCheckConfig(ClarifyCheckConfig): class ModelExplainabilityCheckConfig(ClarifyCheckConfig): model_config: ModelConfig explainability_config: SHAPConfig - model_scores: Union[str, int, ModelPredictedLabelConfig] + model_scores: str | int | ModelPredictedLabelConfig def __init__( self, data_config, kms_key, monitoring_analysis_config_uri, model_config, explainability_config, model_scores ) -> None: ... @@ -73,15 +73,15 @@ class ClarifyCheckStep(Step): name: str, clarify_check_config: ClarifyCheckConfig, check_job_config: CheckJobConfig, - skip_check: Union[bool, PipelineVariable] = False, - fail_on_violation: Union[bool, PipelineVariable] = True, - register_new_baseline: Union[bool, PipelineVariable] = False, - model_package_group_name: Union[str, PipelineVariable] = None, - supplied_baseline_constraints: Union[str, PipelineVariable] = None, + skip_check: bool | PipelineVariable = False, + fail_on_violation: bool | PipelineVariable = True, + register_new_baseline: bool | PipelineVariable = False, + model_package_group_name: str | PipelineVariable = None, + supplied_baseline_constraints: str | PipelineVariable = None, display_name: str = None, description: str = None, cache_config: CacheConfig = None, - depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + depends_on: Optional[List[str | Step | StepCollection]] = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/condition_step.pyi b/stubs/sagemaker/sagemaker/workflow/condition_step.pyi index b82ddd1929e3..bf1b47b32ddd 100644 --- a/stubs/sagemaker/sagemaker/workflow/condition_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/condition_step.pyi @@ -15,12 +15,12 @@ class ConditionStep(Step): def __init__( self, name: str, - depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + depends_on: Optional[List[str | Step | StepCollection]] = None, display_name: str = None, description: str = None, conditions: List[Condition] = None, - if_steps: List[Union[Step, StepCollection]] = None, - else_steps: List[Union[Step, StepCollection]] = None, + if_steps: List[Step | StepCollection] = None, + else_steps: List[Step | StepCollection] = None, ) -> None: ... @property def arguments(self) -> RequestType: ... @@ -30,4 +30,4 @@ class ConditionStep(Step): def properties(self): ... class JsonGet(NewJsonGet): - def __init__(self, step: Step, property_file: Union[PropertyFile, str], json_path: str) -> None: ... + def __init__(self, step: Step, property_file: PropertyFile | str, json_path: str) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/conditions.pyi b/stubs/sagemaker/sagemaker/workflow/conditions.pyi index e24ccaa5048e..29187b98b603 100644 --- a/stubs/sagemaker/sagemaker/workflow/conditions.pyi +++ b/stubs/sagemaker/sagemaker/workflow/conditions.pyi @@ -14,7 +14,7 @@ from sagemaker.workflow.execution_variables import ExecutionVariable from sagemaker.workflow.parameters import Parameter from sagemaker.workflow.properties import Properties -ConditionValueType = Union[ExecutionVariable, Parameter, Properties] +ConditionValueType = ExecutionVariable | Parameter | Properties class ConditionTypeEnum(Enum, metaclass=DefaultEnumMeta): EQ: str @@ -35,8 +35,8 @@ class Condition(Entity, metaclass=abc.ABCMeta): def __ge__(self, other): ... class ConditionComparison(Condition): - left: Union[ConditionValueType, PrimitiveType] - right: Union[ConditionValueType, PrimitiveType] + left: ConditionValueType | PrimitiveType + right: ConditionValueType | PrimitiveType def to_request(self) -> RequestType: ... def __init__(self, condition_type, left, right) -> None: ... def __lt__(self, other): ... @@ -46,34 +46,34 @@ class ConditionComparison(Condition): class ConditionEquals(ConditionComparison): def __init__( - self, left: Union[ConditionValueType, PrimitiveType], right: Union[ConditionValueType, PrimitiveType] + self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType ) -> None: ... class ConditionGreaterThan(ConditionComparison): def __init__( - self, left: Union[ConditionValueType, PrimitiveType], right: Union[ConditionValueType, PrimitiveType] + self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType ) -> None: ... class ConditionGreaterThanOrEqualTo(ConditionComparison): def __init__( - self, left: Union[ConditionValueType, PrimitiveType], right: Union[ConditionValueType, PrimitiveType] + self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType ) -> None: ... class ConditionLessThan(ConditionComparison): def __init__( - self, left: Union[ConditionValueType, PrimitiveType], right: Union[ConditionValueType, PrimitiveType] + self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType ) -> None: ... class ConditionLessThanOrEqualTo(ConditionComparison): def __init__( - self, left: Union[ConditionValueType, PrimitiveType], right: Union[ConditionValueType, PrimitiveType] + self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType ) -> None: ... class ConditionIn(Condition): value: Incomplete in_values: Incomplete def __init__( - self, value: Union[ConditionValueType, PrimitiveType], in_values: List[Union[ConditionValueType, PrimitiveType]] + self, value: ConditionValueType | PrimitiveType, in_values: List[ConditionValueType | PrimitiveType] ) -> None: ... def to_request(self) -> RequestType: ... @@ -88,5 +88,5 @@ class ConditionOr(Condition): def to_request(self) -> RequestType: ... def primitive_or_expr( - value: Union[ExecutionVariable, Expression, PrimitiveType, Parameter, Properties] -) -> Union[Dict[str, str], PrimitiveType]: ... + value: ExecutionVariable | Expression | PrimitiveType | Parameter | Properties +) -> Dict[str | str, PrimitiveType]: ... diff --git a/stubs/sagemaker/sagemaker/workflow/emr_step.pyi b/stubs/sagemaker/sagemaker/workflow/emr_step.pyi index 480e21608c87..4581eed59b2f 100644 --- a/stubs/sagemaker/sagemaker/workflow/emr_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/emr_step.pyi @@ -34,7 +34,7 @@ class EMRStep(Step): description: str, cluster_id: str, step_config: EMRStepConfig, - depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + depends_on: Optional[List[str | Step | StepCollection]] = None, cache_config: CacheConfig = None, cluster_config: Dict[str, Any] = None, execution_role_arn: str = None, diff --git a/stubs/sagemaker/sagemaker/workflow/entities.pyi b/stubs/sagemaker/sagemaker/workflow/entities.pyi index f8c6383b9b32..efc729637e22 100644 --- a/stubs/sagemaker/sagemaker/workflow/entities.pyi +++ b/stubs/sagemaker/sagemaker/workflow/entities.pyi @@ -3,8 +3,8 @@ from _typeshed import Incomplete from enum import EnumMeta from typing import Any, Dict, List, Union -PrimitiveType = Union[str, int, bool, float, None] -RequestType = Union[Dict[str, Any], List[Dict[str, Any]]] +PrimitiveType = str | int | bool | float | None +RequestType = Dict[str | Any, List[Dict[str, Any]]] class Entity(abc.ABC, metaclass=abc.ABCMeta): @abc.abstractmethod @@ -21,7 +21,7 @@ class Expression(abc.ABC, metaclass=abc.ABCMeta): def expr(self) -> RequestType: ... class PipelineVariable(Expression, metaclass=abc.ABCMeta): - def __add__(self, other: Union[Expression, PrimitiveType]): ... + def __add__(self, other: Expression | PrimitiveType): ... def __int__(self) -> int: ... def __float__(self) -> float: ... def to_string(self): ... diff --git a/stubs/sagemaker/sagemaker/workflow/fail_step.pyi b/stubs/sagemaker/sagemaker/workflow/fail_step.pyi index 2868b3cb71f6..67a872311071 100644 --- a/stubs/sagemaker/sagemaker/workflow/fail_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/fail_step.pyi @@ -10,10 +10,10 @@ class FailStep(Step): def __init__( self, name: str, - error_message: Union[str, PipelineVariable] = None, + error_message: str | PipelineVariable = None, display_name: str = None, description: str = None, - depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + depends_on: Optional[List[str | Step | StepCollection]] = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/functions.pyi b/stubs/sagemaker/sagemaker/workflow/functions.pyi index 5313530ff8a5..50d522031e2e 100644 --- a/stubs/sagemaker/sagemaker/workflow/functions.pyi +++ b/stubs/sagemaker/sagemaker/workflow/functions.pyi @@ -17,7 +17,7 @@ class Join(PipelineVariable): class JsonGet(PipelineVariable): step_name: str - property_file: Union[PropertyFile, str] + property_file: PropertyFile | str json_path: str @property def expr(self): ... diff --git a/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi b/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi index 4c39ca89b56d..e5c91cedff98 100644 --- a/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi @@ -38,7 +38,7 @@ class LambdaStep(Step): inputs: dict = None, outputs: List[LambdaOutput] = None, cache_config: CacheConfig = None, - depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + depends_on: Optional[List[str | Step | StepCollection]] = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/model_step.pyi b/stubs/sagemaker/sagemaker/workflow/model_step.pyi index 63302f38720e..9e86b7241e3d 100644 --- a/stubs/sagemaker/sagemaker/workflow/model_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/model_step.pyi @@ -18,8 +18,8 @@ class ModelStep(StepCollection): self, name: str, step_args: _ModelStepArguments, - depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, - retry_policies: Optional[Union[List[RetryPolicy], Dict[str, List[RetryPolicy]]]] = None, + depends_on: Optional[List[str | Step | StepCollection]] = None, + retry_policies: Optional[List[RetryPolicy, Dict[str, List[RetryPolicy]]]] = None, display_name: Optional[str] = None, description: Optional[str] = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi b/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi index 0bf48558b0cb..50eee0e8d351 100644 --- a/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi @@ -15,12 +15,12 @@ class MonitorBatchTransformStep(StepCollection): self, name: str, transform_step_args: _JobStepArguments, - monitor_configuration: Union[QualityCheckConfig, ClarifyCheckConfig], + monitor_configuration: QualityCheckConfig | ClarifyCheckConfig, check_job_configuration: CheckJobConfig, monitor_before_transform: bool = False, - fail_on_violation: Union[bool, PipelineVariable] = True, - supplied_baseline_statistics: Union[str, PipelineVariable] = None, - supplied_baseline_constraints: Union[str, PipelineVariable] = None, + fail_on_violation: bool | PipelineVariable = True, + supplied_baseline_statistics: str | PipelineVariable = None, + supplied_baseline_constraints: str | PipelineVariable = None, display_name: Optional[str] = None, description: Optional[str] = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/pipeline.pyi b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi index 54d014fc080a..5622736ad516 100644 --- a/stubs/sagemaker/sagemaker/workflow/pipeline.pyi +++ b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi @@ -23,7 +23,7 @@ class Pipeline(Entity): name: str = "", parameters: Optional[Sequence[Parameter]] = None, pipeline_experiment_config: Optional[PipelineExperimentConfig] = ..., - steps: Optional[Sequence[Union[Step, StepCollection]]] = None, + steps: Optional[Sequence[Step | StepCollection]] = None, sagemaker_session: Optional[Session] = None, ) -> None: ... def to_request(self) -> RequestType: ... @@ -48,7 +48,7 @@ class Pipeline(Entity): def delete(self) -> Dict[str, Any]: ... def start( self, - parameters: Dict[str, Union[str, bool, int, float]] = None, + parameters: Dict[str, str | bool | int | float] = None, execution_display_name: str = None, execution_description: str = None, parallelism_config: ParallelismConfiguration = None, @@ -81,11 +81,11 @@ class _PipelineExecution: class PipelineGraph: step_map: Incomplete adjacency_list: Incomplete - def __init__(self, steps: Sequence[Union[Step, StepCollection]]) -> None: ... + def __init__(self, steps: Sequence[Step | StepCollection]) -> None: ... @classmethod def from_pipeline(cls, pipeline: Pipeline): ... def is_cyclic(self) -> bool: ... - def get_steps_in_sub_dag(self, current_step: Union[Step, StepCollection], sub_dag_steps: Set[str] = None) -> Set[str]: ... + def get_steps_in_sub_dag(self, current_step: Step | StepCollection, sub_dag_steps: Set[str] = None) -> Set[str]: ... stack: Incomplete def __iter__(self): ... def __next__(self) -> Step: ... diff --git a/stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi b/stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi index 8ea15168b974..51cd4602567b 100644 --- a/stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi +++ b/stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi @@ -10,8 +10,8 @@ class PipelineExperimentConfig(Entity): trial_name: Incomplete def __init__( self, - experiment_name: Union[str, Parameter, ExecutionVariable, Expression], - trial_name: Union[str, Parameter, ExecutionVariable, Expression], + experiment_name: str | Parameter | ExecutionVariable | Expression, + trial_name: str | Parameter | ExecutionVariable | Expression, ) -> None: ... def to_request(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/properties.pyi b/stubs/sagemaker/sagemaker/workflow/properties.pyi index 47b3f72c9ae8..dca7119e4d6f 100644 --- a/stubs/sagemaker/sagemaker/workflow/properties.pyi +++ b/stubs/sagemaker/sagemaker/workflow/properties.pyi @@ -25,13 +25,13 @@ class PropertiesList(Properties): shape_name: Incomplete service_name: Incomplete def __init__(self, step_name: str, path: str, shape_name: str = None, service_name: str = "sagemaker") -> None: ... - def __getitem__(self, item: Union[int, str]): ... + def __getitem__(self, item: int | str): ... class PropertiesMap(Properties): shape_name: Incomplete service_name: Incomplete def __init__(self, step_name: str, path: str, shape_name: str = None, service_name: str = "sagemaker") -> None: ... - def __getitem__(self, item: Union[int, str]): ... + def __getitem__(self, item: int | str): ... class PropertyFile(Expression): name: str diff --git a/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi b/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi index aeaade118d81..b13daa368b8c 100644 --- a/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi @@ -8,9 +8,9 @@ from sagemaker.workflow.step_collections import StepCollection from sagemaker.workflow.steps import CacheConfig, Step class QualityCheckConfig(ABC): - baseline_dataset: Union[str, PipelineVariable] + baseline_dataset: str | PipelineVariable dataset_format: dict - output_s3_uri: Union[str, PipelineVariable] + output_s3_uri: str | PipelineVariable post_analytics_processor_script: str def __init__(self, baseline_dataset, dataset_format, output_s3_uri, post_analytics_processor_script) -> None: ... def __lt__(self, other): ... @@ -29,11 +29,11 @@ class DataQualityCheckConfig(QualityCheckConfig): def __ge__(self, other): ... class ModelQualityCheckConfig(QualityCheckConfig): - problem_type: Union[str, PipelineVariable] - inference_attribute: Union[str, PipelineVariable] - probability_attribute: Union[str, PipelineVariable] - ground_truth_attribute: Union[str, PipelineVariable] - probability_threshold_attribute: Union[str, PipelineVariable] + problem_type: str | PipelineVariable + inference_attribute: str | PipelineVariable + probability_attribute: str | PipelineVariable + ground_truth_attribute: str | PipelineVariable + probability_threshold_attribute: str | PipelineVariable def __init__( self, baseline_dataset, @@ -66,16 +66,16 @@ class QualityCheckStep(Step): name: str, quality_check_config: QualityCheckConfig, check_job_config: CheckJobConfig, - skip_check: Union[bool, PipelineVariable] = False, - fail_on_violation: Union[bool, PipelineVariable] = True, - register_new_baseline: Union[bool, PipelineVariable] = False, - model_package_group_name: Union[str, PipelineVariable] = None, - supplied_baseline_statistics: Union[str, PipelineVariable] = None, - supplied_baseline_constraints: Union[str, PipelineVariable] = None, + skip_check: bool | PipelineVariable = False, + fail_on_violation: bool | PipelineVariable = True, + register_new_baseline: bool | PipelineVariable = False, + model_package_group_name: str | PipelineVariable = None, + supplied_baseline_statistics: str | PipelineVariable = None, + supplied_baseline_constraints: str | PipelineVariable = None, display_name: str = None, description: str = None, cache_config: CacheConfig = None, - depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + depends_on: Optional[List[str | Step | StepCollection]] = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/step_collections.pyi b/stubs/sagemaker/sagemaker/workflow/step_collections.pyi index 29964ced9093..09a887a67e9e 100644 --- a/stubs/sagemaker/sagemaker/workflow/step_collections.pyi +++ b/stubs/sagemaker/sagemaker/workflow/step_collections.pyi @@ -34,7 +34,7 @@ class RegisterModel(StepCollection): transform_instances: Incomplete | None = None, estimator: EstimatorBase = None, model_data: Incomplete | None = None, - depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + depends_on: Optional[List[str | Step | StepCollection]] = None, repack_model_step_retry_policies: List[RetryPolicy] = None, register_model_step_retry_policies: List[RetryPolicy] = None, model_package_group_name: Incomplete | None = None, @@ -45,7 +45,7 @@ class RegisterModel(StepCollection): display_name: Incomplete | None = None, description: Incomplete | None = None, tags: Incomplete | None = None, - model: Union[Model, PipelineModel] = None, + model: Model | PipelineModel = None, drift_check_baselines: Incomplete | None = None, customer_metadata_properties: Incomplete | None = None, domain: Incomplete | None = None, @@ -84,7 +84,7 @@ class EstimatorTransformer(StepCollection): max_payload: Incomplete | None = None, tags: Incomplete | None = None, volume_kms_key: Incomplete | None = None, - depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + depends_on: Optional[List[str | Step | StepCollection]] = None, repack_model_step_retry_policies: List[RetryPolicy] = None, model_step_retry_policies: List[RetryPolicy] = None, transform_step_retry_policies: List[RetryPolicy] = None, diff --git a/stubs/sagemaker/sagemaker/workflow/steps.pyi b/stubs/sagemaker/sagemaker/workflow/steps.pyi index c4622be62499..1b7b9bfe6d60 100644 --- a/stubs/sagemaker/sagemaker/workflow/steps.pyi +++ b/stubs/sagemaker/sagemaker/workflow/steps.pyi @@ -38,7 +38,7 @@ class Step(Entity, metaclass=abc.ABCMeta): display_name: Optional[str] description: Optional[str] step_type: StepTypeEnum - depends_on: Optional[List[Union[str, "Step", "StepCollection"]]] + depends_on: Optional[List[str | "Step" | "StepCollection"]] @property @abc.abstractmethod def arguments(self) -> RequestType: ... @@ -48,7 +48,7 @@ class Step(Entity, metaclass=abc.ABCMeta): @abc.abstractmethod def properties(self): ... def to_request(self) -> RequestType: ... - def add_depends_on(self, step_names: List[Union[str, "Step", "StepCollection"]]): ... + def add_depends_on(self, step_names: List[str | "Step" | "StepCollection"]): ... @property def ref(self) -> Dict[str, str]: ... def __init__(self, name, display_name, description, step_type, depends_on) -> None: ... @@ -76,7 +76,7 @@ class ConfigurableRetryStep(Step, metaclass=abc.ABCMeta): step_type: StepTypeEnum, display_name: str = None, description: str = None, - depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + depends_on: Optional[List[str | Step | "StepCollection"]] = None, retry_policies: List[RetryPolicy] = None, ) -> None: ... def add_retry_policy(self, retry_policy: RetryPolicy): ... @@ -95,9 +95,9 @@ class TrainingStep(ConfigurableRetryStep): estimator: EstimatorBase = None, display_name: str = None, description: str = None, - inputs: Union[TrainingInput, dict, str, FileSystemInput] = None, + inputs: TrainingInput | dict | str | FileSystemInput = None, cache_config: CacheConfig = None, - depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + depends_on: Optional[List[str | Step | "StepCollection"]] = None, retry_policies: List[RetryPolicy] = None, ) -> None: ... @property @@ -114,9 +114,9 @@ class CreateModelStep(ConfigurableRetryStep): self, name: str, step_args: Optional[dict] = None, - model: Optional[Union[Model, PipelineModel]] = None, + model: Optional[Model | PipelineModel] = None, inputs: Optional[CreateModelInput] = None, - depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + depends_on: Optional[List[str | Step | "StepCollection"]] = None, retry_policies: Optional[List[RetryPolicy]] = None, display_name: Optional[str] = None, description: Optional[str] = None, @@ -140,7 +140,7 @@ class TransformStep(ConfigurableRetryStep): display_name: str = None, description: str = None, cache_config: CacheConfig = None, - depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + depends_on: Optional[List[str | Step | "StepCollection"]] = None, retry_policies: List[RetryPolicy] = None, ) -> None: ... @property @@ -173,7 +173,7 @@ class ProcessingStep(ConfigurableRetryStep): code: str = None, property_files: List[PropertyFile] = None, cache_config: CacheConfig = None, - depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + depends_on: Optional[List[str | Step | "StepCollection"]] = None, retry_policies: List[RetryPolicy] = None, kms_key: Incomplete | None = None, ) -> None: ... @@ -199,7 +199,7 @@ class TuningStep(ConfigurableRetryStep): inputs: Incomplete | None = None, job_arguments: List[str] = None, cache_config: CacheConfig = None, - depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None, + depends_on: Optional[List[str | Step | "StepCollection"]] = None, retry_policies: List[RetryPolicy] = None, ) -> None: ... @property diff --git a/stubs/sagemaker/sagemaker/workflow/utilities.pyi b/stubs/sagemaker/sagemaker/workflow/utilities.pyi index 89acb50bc08f..f57c8f393728 100644 --- a/stubs/sagemaker/sagemaker/workflow/utilities.pyi +++ b/stubs/sagemaker/sagemaker/workflow/utilities.pyi @@ -8,7 +8,7 @@ from sagemaker.workflow.step_collections import StepCollection logger: Incomplete BUF_SIZE: int -def list_to_request(entities: Sequence[Union[Entity, "StepCollection"]]) -> List[RequestType]: ... +def list_to_request(entities: Sequence[Entity | "StepCollection"]) -> List[RequestType]: ... def build_steps(steps: Sequence[Entity], pipeline_name: str): ... def get_code_hash(step: Entity) -> str: ... def get_processing_dependencies(dependency_args: List[List[str]]) -> List[str]: ... diff --git a/stubs/sagemaker/sagemaker/xgboost/estimator.pyi b/stubs/sagemaker/sagemaker/xgboost/estimator.pyi index 21c50a492253..a57dbedc9235 100644 --- a/stubs/sagemaker/sagemaker/xgboost/estimator.pyi +++ b/stubs/sagemaker/sagemaker/xgboost/estimator.pyi @@ -12,12 +12,12 @@ class XGBoost(Framework): image_uri: Incomplete def __init__( self, - entry_point: Union[str, PipelineVariable], + entry_point: str | PipelineVariable, framework_version: str, - source_dir: Optional[Union[str, PipelineVariable]] = None, - hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None, + source_dir: Optional[str | PipelineVariable] = None, + hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, py_version: str = "py3", - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, image_uri_region: Optional[str] = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/xgboost/model.pyi b/stubs/sagemaker/sagemaker/xgboost/model.pyi index b4faae2961db..548b3028c656 100644 --- a/stubs/sagemaker/sagemaker/xgboost/model.pyi +++ b/stubs/sagemaker/sagemaker/xgboost/model.pyi @@ -19,40 +19,40 @@ class XGBoostModel(FrameworkModel): model_server_workers: Incomplete def __init__( self, - model_data: Union[str, PipelineVariable], + model_data: str | PipelineVariable, role: str = None, entry_point: str = None, framework_version: str = None, - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, py_version: str = "py3", predictor_cls: callable = ..., - model_server_workers: Optional[Union[int, PipelineVariable]] = None, + model_server_workers: Optional[int | PipelineVariable] = None, **kwargs, ) -> None: ... image_uri: Incomplete def register( self, - 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, + content_types: List[str | PipelineVariable], + response_types: List[str | PipelineVariable], + inference_instances: Optional[List[str | PipelineVariable]] = None, + transform_instances: Optional[List[str | PipelineVariable]] = None, + model_package_name: Optional[str | PipelineVariable] = None, + model_package_group_name: Optional[str | PipelineVariable] = None, + image_uri: Optional[str | PipelineVariable] = None, model_metrics: Optional[ModelMetrics] = None, metadata_properties: Optional[MetadataProperties] = None, marketplace_cert: bool = False, - approval_status: Optional[Union[str, PipelineVariable]] = None, + approval_status: Optional[str | PipelineVariable] = None, description: Optional[str] = None, drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - domain: Optional[Union[str, PipelineVariable]] = None, - sample_payload_url: Optional[Union[str, PipelineVariable]] = None, - task: Optional[Union[str, PipelineVariable]] = None, - framework: Optional[Union[str, PipelineVariable]] = None, - framework_version: Optional[Union[str, PipelineVariable]] = None, - nearest_model_name: Optional[Union[str, PipelineVariable]] = None, - data_input_configuration: Optional[Union[str, PipelineVariable]] = None, + customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, + domain: Optional[str | PipelineVariable] = None, + sample_payload_url: Optional[str | PipelineVariable] = None, + task: Optional[str | PipelineVariable] = None, + framework: Optional[str | PipelineVariable] = None, + framework_version: Optional[str | PipelineVariable] = None, + nearest_model_name: Optional[str | PipelineVariable] = None, + data_input_configuration: Optional[str | PipelineVariable] = None, ): ... def prepare_container_def( self, diff --git a/stubs/sagemaker/sagemaker/xgboost/processing.pyi b/stubs/sagemaker/sagemaker/xgboost/processing.pyi index 2985764e42c5..efaf7365e677 100644 --- a/stubs/sagemaker/sagemaker/xgboost/processing.pyi +++ b/stubs/sagemaker/sagemaker/xgboost/processing.pyi @@ -12,19 +12,19 @@ class XGBoostProcessor(FrameworkProcessor): self, framework_version: str, role: str = None, - instance_count: Union[int, PipelineVariable] = None, - instance_type: Union[str, PipelineVariable] = None, + instance_count: int | PipelineVariable = None, + instance_type: str | PipelineVariable = None, py_version: str = "py3", - image_uri: Optional[Union[str, PipelineVariable]] = None, + image_uri: Optional[str | PipelineVariable] = None, command: Optional[List[str]] = None, - volume_size_in_gb: Union[int, PipelineVariable] = 30, - volume_kms_key: Optional[Union[str, PipelineVariable]] = None, - output_kms_key: Optional[Union[str, PipelineVariable]] = None, + volume_size_in_gb: int | PipelineVariable = 30, + volume_kms_key: Optional[str | PipelineVariable] = None, + output_kms_key: Optional[str | PipelineVariable] = None, code_location: Optional[str] = None, - max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None, + max_runtime_in_seconds: Optional[int | PipelineVariable] = None, base_job_name: Optional[str] = None, sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, - tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None, + env: Optional[Dict[str, str | PipelineVariable]] = None, + tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, network_config: Optional[NetworkConfig] = None, ) -> None: ... From 2cfc053e1c93dc6420bdf75ded99510f650031c8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 23 Jun 2023 16:37:59 +0000 Subject: [PATCH 04/10] [pre-commit.ci] auto fixes from pre-commit.com hooks --- stubs/sagemaker/sagemaker/algorithm.pyi | 2 +- .../sagemaker/amazon/amazon_estimator.pyi | 2 +- .../amazon/factorization_machines.pyi | 2 +- .../sagemaker/sagemaker/amazon/ipinsights.pyi | 2 +- stubs/sagemaker/sagemaker/amazon/kmeans.pyi | 2 +- stubs/sagemaker/sagemaker/amazon/knn.pyi | 2 +- stubs/sagemaker/sagemaker/amazon/lda.pyi | 2 +- .../sagemaker/amazon/linear_learner.pyi | 2 +- stubs/sagemaker/sagemaker/amazon/ntm.pyi | 2 +- .../sagemaker/sagemaker/amazon/object2vec.pyi | 2 +- stubs/sagemaker/sagemaker/amazon/pca.pyi | 2 +- .../sagemaker/amazon/randomcutforest.pyi | 2 +- .../sagemaker/sagemaker/chainer/estimator.pyi | 2 +- stubs/sagemaker/sagemaker/chainer/model.pyi | 2 +- stubs/sagemaker/sagemaker/clarify.pyi | 2 +- .../sagemaker/sagemaker/debugger/debugger.pyi | 10 +++----- .../sagemaker/debugger/profiler_config.pyi | 2 +- .../sagemaker/djl_inference/model.pyi | 2 +- stubs/sagemaker/sagemaker/estimator.pyi | 2 +- stubs/sagemaker/sagemaker/experiments/run.pyi | 6 ++--- .../sagemaker/feature_store/feature_group.pyi | 2 +- .../sagemaker/feature_store/feature_store.pyi | 2 +- stubs/sagemaker/sagemaker/fw_utils.pyi | 2 +- .../sagemaker/huggingface/estimator.pyi | 2 +- .../sagemaker/sagemaker/huggingface/model.pyi | 2 +- .../sagemaker/huggingface/processing.pyi | 2 +- .../huggingface/training_compiler/config.pyi | 1 - stubs/sagemaker/sagemaker/inputs.pyi | 2 +- .../sagemaker/jumpstart/estimator.pyi | 2 +- .../sagemaker/jumpstart/factory/estimator.pyi | 2 +- .../sagemaker/jumpstart/factory/model.pyi | 2 +- .../sagemaker/sagemaker/jumpstart/filters.pyi | 2 +- stubs/sagemaker/sagemaker/jumpstart/model.pyi | 2 +- .../sagemaker/jumpstart/notebook_utils.pyi | 2 +- stubs/sagemaker/sagemaker/jumpstart/types.pyi | 2 +- stubs/sagemaker/sagemaker/lineage/query.pyi | 2 +- .../sagemaker/metadata_properties.pyi | 2 +- stubs/sagemaker/sagemaker/model.pyi | 2 +- .../sagemaker/model_card/helpers.pyi | 2 +- .../sagemaker/model_card/model_card.pyi | 2 +- stubs/sagemaker/sagemaker/model_metrics.pyi | 2 +- .../model_monitor/model_monitoring.pyi | 2 +- stubs/sagemaker/sagemaker/multidatamodel.pyi | 2 +- stubs/sagemaker/sagemaker/mxnet/estimator.pyi | 2 +- stubs/sagemaker/sagemaker/mxnet/model.pyi | 2 +- .../sagemaker/sagemaker/mxnet/processing.pyi | 2 +- stubs/sagemaker/sagemaker/network.pyi | 2 +- stubs/sagemaker/sagemaker/parameter.pyi | 1 - stubs/sagemaker/sagemaker/pipeline.pyi | 2 +- stubs/sagemaker/sagemaker/processing.pyi | 2 +- .../sagemaker/sagemaker/pytorch/estimator.pyi | 2 +- stubs/sagemaker/sagemaker/pytorch/model.pyi | 2 +- .../sagemaker/pytorch/processing.pyi | 2 +- .../pytorch/training_compiler/config.pyi | 1 - .../remote_function/spark_config.pyi | 2 +- stubs/sagemaker/sagemaker/rl/estimator.pyi | 2 +- stubs/sagemaker/sagemaker/s3.pyi | 1 - .../sagemaker/sagemaker/sklearn/estimator.pyi | 2 +- stubs/sagemaker/sagemaker/sklearn/model.pyi | 2 +- .../sagemaker/sklearn/processing.pyi | 2 +- .../sagemaker/sagemaker/spark/processing.pyi | 2 +- stubs/sagemaker/sagemaker/sparkml/model.pyi | 2 +- .../sagemaker/tensorflow/estimator.pyi | 2 +- .../sagemaker/sagemaker/tensorflow/model.pyi | 2 +- .../sagemaker/tensorflow/processing.pyi | 2 +- stubs/sagemaker/sagemaker/transformer.pyi | 2 +- stubs/sagemaker/sagemaker/tuner.pyi | 2 +- stubs/sagemaker/sagemaker/workflow/_utils.pyi | 3 +-- .../sagemaker/workflow/automl_step.pyi | 2 +- .../sagemaker/workflow/callback_step.pyi | 2 +- .../sagemaker/workflow/clarify_check_step.pyi | 2 +- .../sagemaker/workflow/condition_step.pyi | 2 +- .../sagemaker/workflow/conditions.pyi | 24 ++++++------------- .../sagemaker/sagemaker/workflow/emr_step.pyi | 2 +- .../sagemaker/sagemaker/workflow/entities.pyi | 2 +- .../sagemaker/workflow/fail_step.pyi | 2 +- .../sagemaker/workflow/functions.pyi | 2 +- .../sagemaker/workflow/lambda_step.pyi | 2 +- .../sagemaker/workflow/model_step.pyi | 2 +- .../workflow/monitor_batch_transform_step.pyi | 2 +- .../sagemaker/sagemaker/workflow/pipeline.pyi | 2 +- .../workflow/pipeline_experiment_config.pyi | 1 - .../sagemaker/workflow/properties.pyi | 2 +- .../sagemaker/workflow/quality_check_step.pyi | 2 +- .../sagemaker/workflow/step_collections.pyi | 2 +- stubs/sagemaker/sagemaker/workflow/steps.pyi | 3 +-- .../sagemaker/workflow/utilities.pyi | 3 +-- .../sagemaker/sagemaker/xgboost/estimator.pyi | 2 +- stubs/sagemaker/sagemaker/xgboost/model.pyi | 2 +- .../sagemaker/xgboost/processing.pyi | 2 +- 90 files changed, 94 insertions(+), 118 deletions(-) diff --git a/stubs/sagemaker/sagemaker/algorithm.pyi b/stubs/sagemaker/sagemaker/algorithm.pyi index 9b4a6ae90e8f..1bf184fea132 100644 --- a/stubs/sagemaker/sagemaker/algorithm.pyi +++ b/stubs/sagemaker/sagemaker/algorithm.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.estimator import EstimatorBase from sagemaker.inputs import FileSystemInput, TrainingInput diff --git a/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi b/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi index 6fe1300e7da9..9ab83b9e93da 100644 --- a/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi +++ b/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi @@ -1,6 +1,6 @@ import abc from _typeshed import Incomplete -from typing import Dict, Optional, Union +from typing import Dict, Optional from sagemaker.amazon.hyperparameter import Hyperparameter as hp from sagemaker.estimator import EstimatorBase diff --git a/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi b/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi index 5c04f9333064..4137eefeac57 100644 --- a/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi +++ b/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional, Union +from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp diff --git a/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi b/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi index d7fec5412f8c..6f208e9a13b1 100644 --- a/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi +++ b/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional, Union +from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp diff --git a/stubs/sagemaker/sagemaker/amazon/kmeans.pyi b/stubs/sagemaker/sagemaker/amazon/kmeans.pyi index fbc0be928537..f99bc3da1309 100644 --- a/stubs/sagemaker/sagemaker/amazon/kmeans.pyi +++ b/stubs/sagemaker/sagemaker/amazon/kmeans.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import List, Optional, Union +from typing import List, Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp diff --git a/stubs/sagemaker/sagemaker/amazon/knn.pyi b/stubs/sagemaker/sagemaker/amazon/knn.pyi index 093b6186eb6a..808e9883e233 100644 --- a/stubs/sagemaker/sagemaker/amazon/knn.pyi +++ b/stubs/sagemaker/sagemaker/amazon/knn.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional, Union +from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp diff --git a/stubs/sagemaker/sagemaker/amazon/lda.pyi b/stubs/sagemaker/sagemaker/amazon/lda.pyi index cbe389ebe9e0..44920b303c97 100644 --- a/stubs/sagemaker/sagemaker/amazon/lda.pyi +++ b/stubs/sagemaker/sagemaker/amazon/lda.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional, Union +from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp diff --git a/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi b/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi index 35eaca2bf260..9443a79d2d79 100644 --- a/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi +++ b/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional, Union +from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp diff --git a/stubs/sagemaker/sagemaker/amazon/ntm.pyi b/stubs/sagemaker/sagemaker/amazon/ntm.pyi index 06ac9a112442..c7ab3b7b4685 100644 --- a/stubs/sagemaker/sagemaker/amazon/ntm.pyi +++ b/stubs/sagemaker/sagemaker/amazon/ntm.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import List, Optional, Union +from typing import List, Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp diff --git a/stubs/sagemaker/sagemaker/amazon/object2vec.pyi b/stubs/sagemaker/sagemaker/amazon/object2vec.pyi index feecabb546c8..3153b860cafb 100644 --- a/stubs/sagemaker/sagemaker/amazon/object2vec.pyi +++ b/stubs/sagemaker/sagemaker/amazon/object2vec.pyi @@ -1,4 +1,4 @@ -from typing import Optional, Union +from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp diff --git a/stubs/sagemaker/sagemaker/amazon/pca.pyi b/stubs/sagemaker/sagemaker/amazon/pca.pyi index 9981d23d8947..a1bd6fe56697 100644 --- a/stubs/sagemaker/sagemaker/amazon/pca.pyi +++ b/stubs/sagemaker/sagemaker/amazon/pca.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional, Union +from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp diff --git a/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi b/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi index 26525d786425..cc8391aa9183 100644 --- a/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi +++ b/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import List, Optional, Union +from typing import List, Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp diff --git a/stubs/sagemaker/sagemaker/chainer/estimator.pyi b/stubs/sagemaker/sagemaker/chainer/estimator.pyi index 5fd5d03bc746..2295931d7429 100644 --- a/stubs/sagemaker/sagemaker/chainer/estimator.pyi +++ b/stubs/sagemaker/sagemaker/chainer/estimator.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, Optional, Union +from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/chainer/model.pyi b/stubs/sagemaker/sagemaker/chainer/model.pyi index 2aa2a804ffde..64f15000b7b2 100644 --- a/stubs/sagemaker/sagemaker/chainer/model.pyi +++ b/stubs/sagemaker/sagemaker/chainer/model.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker import ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/clarify.pyi b/stubs/sagemaker/sagemaker/clarify.pyi index 8641c1734f6b..0fd3204e8d4a 100644 --- a/stubs/sagemaker/sagemaker/clarify.pyi +++ b/stubs/sagemaker/sagemaker/clarify.pyi @@ -2,7 +2,7 @@ import abc from _typeshed import Incomplete from abc import ABC, abstractmethod from enum import Enum -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional from sagemaker.network import NetworkConfig from sagemaker.processing import Processor diff --git a/stubs/sagemaker/sagemaker/debugger/debugger.pyi b/stubs/sagemaker/sagemaker/debugger/debugger.pyi index 4bad04bf37f8..5653efe75e27 100644 --- a/stubs/sagemaker/sagemaker/debugger/debugger.pyi +++ b/stubs/sagemaker/sagemaker/debugger/debugger.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from abc import ABC -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.workflow.entities import PipelineVariable @@ -113,17 +113,13 @@ class TensorBoardOutputConfig: s3_output_path: Incomplete container_local_output_path: Incomplete def __init__( - self, - s3_output_path: str | PipelineVariable, - container_local_output_path: Optional[str | PipelineVariable] = None, + self, s3_output_path: str | PipelineVariable, container_local_output_path: Optional[str | PipelineVariable] = None ) -> None: ... class CollectionConfig: name: Incomplete parameters: Incomplete - def __init__( - self, name: str | PipelineVariable, parameters: Optional[Dict[str, str | PipelineVariable]] = None - ) -> None: ... + def __init__(self, name: str | PipelineVariable, parameters: Optional[Dict[str, str | PipelineVariable]] = None) -> None: ... def __eq__(self, other): ... def __ne__(self, other): ... def __hash__(self): ... diff --git a/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi b/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi index 11e37d5c275c..b862e13c7519 100644 --- a/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi +++ b/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional, Union +from typing import Optional from sagemaker.debugger.framework_profile import FrameworkProfile from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/djl_inference/model.pyi b/stubs/sagemaker/sagemaker/djl_inference/model.pyi index abfe9fd0e60e..e7d1a8efbc57 100644 --- a/stubs/sagemaker/sagemaker/djl_inference/model.pyi +++ b/stubs/sagemaker/sagemaker/djl_inference/model.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from enum import Enum -from typing import Any, Dict, Optional, Union +from typing import Any, Dict, Optional from sagemaker import Predictor from sagemaker.deserializers import BaseDeserializer diff --git a/stubs/sagemaker/sagemaker/estimator.pyi b/stubs/sagemaker/sagemaker/estimator.pyi index 0452c4dc26c4..39b620aeb575 100644 --- a/stubs/sagemaker/sagemaker/estimator.pyi +++ b/stubs/sagemaker/sagemaker/estimator.pyi @@ -1,7 +1,7 @@ import abc from _typeshed import Incomplete from abc import abstractmethod -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.debugger import ( DEBUGGER_FLAG as DEBUGGER_FLAG, diff --git a/stubs/sagemaker/sagemaker/experiments/run.pyi b/stubs/sagemaker/sagemaker/experiments/run.pyi index 0c11710a76df..a8c6a4727cc0 100644 --- a/stubs/sagemaker/sagemaker/experiments/run.pyi +++ b/stubs/sagemaker/sagemaker/experiments/run.pyi @@ -1,7 +1,7 @@ import datetime from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from numpy import array from sagemaker import Session @@ -54,9 +54,7 @@ class Run: is_output: bool = True, no_skill: Optional[int] = None, ): ... - def log_roc_curve( - self, y_true: list | array, y_score: list | array, title: Optional[str] = None, is_output: bool = True - ): ... + def log_roc_curve(self, y_true: list | array, y_score: list | array, title: Optional[str] = None, is_output: bool = True): ... def log_confusion_matrix( self, y_true: list | array, y_pred: list | array, title: Optional[str] = None, is_output: bool = True ): ... diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi index 24b82ae0d3a7..d431a357bf5d 100644 --- a/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from multiprocessing.pool import AsyncResult as AsyncResult -from typing import Any, Dict, List, Sequence, Union +from typing import Any, Dict, List, Sequence from botocore.config import Config as Config from pandas import DataFrame as DataFrame diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi index f4f51c3d7a81..91409de5e356 100644 --- a/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi @@ -1,5 +1,5 @@ import datetime -from typing import Any, Dict, Sequence, Union +from typing import Any, Dict, Sequence import pandas as pd from sagemaker import Session diff --git a/stubs/sagemaker/sagemaker/fw_utils.pyi b/stubs/sagemaker/sagemaker/fw_utils.pyi index 6b9d7fc951c3..af6fa158e4ba 100644 --- a/stubs/sagemaker/sagemaker/fw_utils.pyi +++ b/stubs/sagemaker/sagemaker/fw_utils.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, NamedTuple, Optional, Union +from typing import Dict, NamedTuple, Optional from sagemaker.session_settings import SessionSettings from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/huggingface/estimator.pyi b/stubs/sagemaker/sagemaker/huggingface/estimator.pyi index e8d44f2211ae..a3202da34b39 100644 --- a/stubs/sagemaker/sagemaker/huggingface/estimator.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/estimator.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, Optional, Union +from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig diff --git a/stubs/sagemaker/sagemaker/huggingface/model.pyi b/stubs/sagemaker/sagemaker/huggingface/model.pyi index ef886b6f32b7..d1aa7f3fc516 100644 --- a/stubs/sagemaker/sagemaker/huggingface/model.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/model.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker import ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/huggingface/processing.pyi b/stubs/sagemaker/sagemaker/huggingface/processing.pyi index 708c3f58fa7d..383279314a10 100644 --- a/stubs/sagemaker/sagemaker/huggingface/processing.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/processing.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.huggingface.estimator import HuggingFace from sagemaker.network import NetworkConfig diff --git a/stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi b/stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi index 272711ce08e7..7d5f4a15d459 100644 --- a/stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Union from sagemaker.training_compiler.config import TrainingCompilerConfig as BaseConfig from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/inputs.pyi b/stubs/sagemaker/sagemaker/inputs.pyi index 41e845506f5f..0eacb811ab70 100644 --- a/stubs/sagemaker/sagemaker/inputs.pyi +++ b/stubs/sagemaker/sagemaker/inputs.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import List, Optional, Union +from typing import List, Optional from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi b/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi index eeffbe8a0dc3..8e87c1e384f9 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker import session as session from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig diff --git a/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi b/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi index 6773faa713bf..19fee6bb3395 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig from sagemaker.base_deserializers import BaseDeserializer diff --git a/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi b/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi index 190d5d0535e1..deffcd4f5b52 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig from sagemaker.base_deserializers import BaseDeserializer diff --git a/stubs/sagemaker/sagemaker/jumpstart/filters.pyi b/stubs/sagemaker/sagemaker/jumpstart/filters.pyi index c2d7ecc84e82..9dc9e8c8397b 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/filters.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/filters.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from enum import Enum -from typing import Any, Dict, List, Union +from typing import Any, Dict, List from sagemaker.jumpstart.types import JumpStartDataHolderType diff --git a/stubs/sagemaker/sagemaker/jumpstart/model.pyi b/stubs/sagemaker/sagemaker/jumpstart/model.pyi index 2d1d6c4fd675..0149adcefa18 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/model.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/model.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig from sagemaker.base_deserializers import BaseDeserializer diff --git a/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi b/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi index 14fea565372c..4e026fdb42ac 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi @@ -1,4 +1,4 @@ -from typing import List, Tuple, Union +from typing import List, Tuple from sagemaker.jumpstart.filters import Operator diff --git a/stubs/sagemaker/sagemaker/jumpstart/types.pyi b/stubs/sagemaker/sagemaker/jumpstart/types.pyi index d1ecff0c2321..397cc6f6eb19 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/types.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/types.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from enum import Enum -from typing import Any, Dict, List, Optional, Set, Union +from typing import Any, Dict, List, Optional, Set class JumpStartDataHolderType: def __eq__(self, other: Any) -> bool: ... diff --git a/stubs/sagemaker/sagemaker/lineage/query.pyi b/stubs/sagemaker/sagemaker/lineage/query.pyi index 22330d7d94dd..024814f7f2a8 100644 --- a/stubs/sagemaker/sagemaker/lineage/query.pyi +++ b/stubs/sagemaker/sagemaker/lineage/query.pyi @@ -1,7 +1,7 @@ from _typeshed import Incomplete from datetime import datetime from enum import Enum -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional class LineageEntityEnum(Enum): TRIAL: str diff --git a/stubs/sagemaker/sagemaker/metadata_properties.pyi b/stubs/sagemaker/sagemaker/metadata_properties.pyi index baf0929c699b..405816373717 100644 --- a/stubs/sagemaker/sagemaker/metadata_properties.pyi +++ b/stubs/sagemaker/sagemaker/metadata_properties.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional, Union +from typing import Optional from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/model.pyi b/stubs/sagemaker/sagemaker/model.pyi index 56c60f5a4210..9b978abc36e1 100644 --- a/stubs/sagemaker/sagemaker/model.pyi +++ b/stubs/sagemaker/sagemaker/model.pyi @@ -1,6 +1,6 @@ import abc from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.drift_check_baselines import DriftCheckBaselines from sagemaker.inference_recommender.inference_recommender_mixin import InferenceRecommenderMixin diff --git a/stubs/sagemaker/sagemaker/model_card/helpers.pyi b/stubs/sagemaker/sagemaker/model_card/helpers.pyi index 990111b69921..ffbfbf756384 100644 --- a/stubs/sagemaker/sagemaker/model_card/helpers.pyi +++ b/stubs/sagemaker/sagemaker/model_card/helpers.pyi @@ -4,7 +4,7 @@ import json from _typeshed import Incomplete from abc import ABC, abstractmethod from enum import Enum -from typing import Any, List, Optional, Union +from typing import Any, List, Optional from boto3.session import Session as Session diff --git a/stubs/sagemaker/sagemaker/model_card/model_card.pyi b/stubs/sagemaker/sagemaker/model_card/model_card.pyi index 454d6660ecc3..fd5c9203c676 100644 --- a/stubs/sagemaker/sagemaker/model_card/model_card.pyi +++ b/stubs/sagemaker/sagemaker/model_card/model_card.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from datetime import datetime -from typing import List, Optional, Union +from typing import List, Optional from boto3.session import Session as boto3_Session from sagemaker.model_card.evaluation_metric_parsers import EvaluationMetricTypeEnum diff --git a/stubs/sagemaker/sagemaker/model_metrics.pyi b/stubs/sagemaker/sagemaker/model_metrics.pyi index f022d6f9a4f3..080278693967 100644 --- a/stubs/sagemaker/sagemaker/model_metrics.pyi +++ b/stubs/sagemaker/sagemaker/model_metrics.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional, Union +from typing import Optional from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi b/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi index 3c680aa58307..3354c733b444 100644 --- a/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi +++ b/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional, Union +from typing import Optional from sagemaker.model_monitor.dataset_format import MonitoringDatasetFormat from sagemaker.processing import ProcessingJob diff --git a/stubs/sagemaker/sagemaker/multidatamodel.pyi b/stubs/sagemaker/sagemaker/multidatamodel.pyi index 851305ce895e..8f44b3357ab8 100644 --- a/stubs/sagemaker/sagemaker/multidatamodel.pyi +++ b/stubs/sagemaker/sagemaker/multidatamodel.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from collections.abc import Generator -from typing import Optional, Union +from typing import Optional from sagemaker.model import Model from sagemaker.session import Session diff --git a/stubs/sagemaker/sagemaker/mxnet/estimator.pyi b/stubs/sagemaker/sagemaker/mxnet/estimator.pyi index 1d43b7192c04..eec5b37d8a19 100644 --- a/stubs/sagemaker/sagemaker/mxnet/estimator.pyi +++ b/stubs/sagemaker/sagemaker/mxnet/estimator.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, Optional, Union +from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/mxnet/model.pyi b/stubs/sagemaker/sagemaker/mxnet/model.pyi index 5855c68f96b4..510167b2907a 100644 --- a/stubs/sagemaker/sagemaker/mxnet/model.pyi +++ b/stubs/sagemaker/sagemaker/mxnet/model.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker import ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/mxnet/processing.pyi b/stubs/sagemaker/sagemaker/mxnet/processing.pyi index bf71be3b63cb..ca8359d6de5f 100644 --- a/stubs/sagemaker/sagemaker/mxnet/processing.pyi +++ b/stubs/sagemaker/sagemaker/mxnet/processing.pyi @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.mxnet.estimator import MXNet from sagemaker.network import NetworkConfig diff --git a/stubs/sagemaker/sagemaker/network.pyi b/stubs/sagemaker/sagemaker/network.pyi index bd1dd7b114d6..327771cd1066 100644 --- a/stubs/sagemaker/sagemaker/network.pyi +++ b/stubs/sagemaker/sagemaker/network.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import List, Optional, Union +from typing import List, Optional from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/parameter.pyi b/stubs/sagemaker/sagemaker/parameter.pyi index 9e32d87cf1da..6b5fe6129022 100644 --- a/stubs/sagemaker/sagemaker/parameter.pyi +++ b/stubs/sagemaker/sagemaker/parameter.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Union from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/pipeline.pyi b/stubs/sagemaker/sagemaker/pipeline.pyi index 5ff19a13a317..f3f872e309e5 100644 --- a/stubs/sagemaker/sagemaker/pipeline.pyi +++ b/stubs/sagemaker/sagemaker/pipeline.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker import Model, ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/processing.pyi b/stubs/sagemaker/sagemaker/processing.pyi index af93e89b5c3e..da9c8cb939e3 100644 --- a/stubs/sagemaker/sagemaker/processing.pyi +++ b/stubs/sagemaker/sagemaker/processing.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.apiutils._base_types import ApiObject from sagemaker.dataset_definition.inputs import DatasetDefinition, S3Input diff --git a/stubs/sagemaker/sagemaker/pytorch/estimator.pyi b/stubs/sagemaker/sagemaker/pytorch/estimator.pyi index f43715d384c4..3087bbe8be5c 100644 --- a/stubs/sagemaker/sagemaker/pytorch/estimator.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/estimator.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, Optional, Union +from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.pytorch.training_compiler.config import TrainingCompilerConfig diff --git a/stubs/sagemaker/sagemaker/pytorch/model.pyi b/stubs/sagemaker/sagemaker/pytorch/model.pyi index 4a97b9d33616..34009ceb0fe3 100644 --- a/stubs/sagemaker/sagemaker/pytorch/model.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/model.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker import ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/pytorch/processing.pyi b/stubs/sagemaker/sagemaker/pytorch/processing.pyi index 94ce85a7e28e..8193fdbb673e 100644 --- a/stubs/sagemaker/sagemaker/pytorch/processing.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/processing.pyi @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.network import NetworkConfig from sagemaker.processing import FrameworkProcessor diff --git a/stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi b/stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi index 272711ce08e7..7d5f4a15d459 100644 --- a/stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Union from sagemaker.training_compiler.config import TrainingCompilerConfig as BaseConfig from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi b/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi index 7ca69b64f673..21d36b1c5ac4 100644 --- a/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi +++ b/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional class SparkConfig: submit_jars: Optional[List[str]] diff --git a/stubs/sagemaker/sagemaker/rl/estimator.pyi b/stubs/sagemaker/sagemaker/rl/estimator.pyi index 876350132d31..17d02e4a0bf3 100644 --- a/stubs/sagemaker/sagemaker/rl/estimator.pyi +++ b/stubs/sagemaker/sagemaker/rl/estimator.pyi @@ -1,6 +1,6 @@ import enum from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.estimator import Framework from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/s3.pyi b/stubs/sagemaker/sagemaker/s3.pyi index 6e72c7613a81..45ecce1fd8fe 100644 --- a/stubs/sagemaker/sagemaker/s3.pyi +++ b/stubs/sagemaker/sagemaker/s3.pyi @@ -1,6 +1,5 @@ import io from _typeshed import Incomplete -from typing import Union from sagemaker.s3_utils import determine_bucket_and_prefix as determine_bucket_and_prefix diff --git a/stubs/sagemaker/sagemaker/sklearn/estimator.pyi b/stubs/sagemaker/sagemaker/sklearn/estimator.pyi index 50fac3ab5ca5..58ba3d43bd25 100644 --- a/stubs/sagemaker/sagemaker/sklearn/estimator.pyi +++ b/stubs/sagemaker/sagemaker/sklearn/estimator.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, Optional, Union +from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/sklearn/model.pyi b/stubs/sagemaker/sagemaker/sklearn/model.pyi index a39f0f9a6b86..61fd7877063b 100644 --- a/stubs/sagemaker/sagemaker/sklearn/model.pyi +++ b/stubs/sagemaker/sagemaker/sklearn/model.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker import ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/sklearn/processing.pyi b/stubs/sagemaker/sagemaker/sklearn/processing.pyi index fbfc41df84c7..ec6e9c62b22d 100644 --- a/stubs/sagemaker/sagemaker/sklearn/processing.pyi +++ b/stubs/sagemaker/sagemaker/sklearn/processing.pyi @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker import Session from sagemaker.network import NetworkConfig diff --git a/stubs/sagemaker/sagemaker/spark/processing.pyi b/stubs/sagemaker/sagemaker/spark/processing.pyi index c563b2335f11..0020f95cd2fb 100644 --- a/stubs/sagemaker/sagemaker/spark/processing.pyi +++ b/stubs/sagemaker/sagemaker/spark/processing.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.network import NetworkConfig from sagemaker.processing import ProcessingInput, ProcessingOutput, ScriptProcessor diff --git a/stubs/sagemaker/sagemaker/sparkml/model.pyi b/stubs/sagemaker/sagemaker/sparkml/model.pyi index 67d4ec2c16ad..79985c7057e6 100644 --- a/stubs/sagemaker/sagemaker/sparkml/model.pyi +++ b/stubs/sagemaker/sagemaker/sparkml/model.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional, Union +from typing import Optional from sagemaker import Model, Predictor, Session from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi b/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi index ed02d9cd8d1e..322cb7a5b5cd 100644 --- a/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi +++ b/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, Optional, Union +from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.tensorflow.training_compiler.config import TrainingCompilerConfig diff --git a/stubs/sagemaker/sagemaker/tensorflow/model.pyi b/stubs/sagemaker/sagemaker/tensorflow/model.pyi index 1c5d8577183e..df35c8a673eb 100644 --- a/stubs/sagemaker/sagemaker/tensorflow/model.pyi +++ b/stubs/sagemaker/sagemaker/tensorflow/model.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional import sagemaker from sagemaker import ModelMetrics diff --git a/stubs/sagemaker/sagemaker/tensorflow/processing.pyi b/stubs/sagemaker/sagemaker/tensorflow/processing.pyi index 562bb08c86c8..d0fa1560e444 100644 --- a/stubs/sagemaker/sagemaker/tensorflow/processing.pyi +++ b/stubs/sagemaker/sagemaker/tensorflow/processing.pyi @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.network import NetworkConfig from sagemaker.processing import FrameworkProcessor diff --git a/stubs/sagemaker/sagemaker/transformer.pyi b/stubs/sagemaker/sagemaker/transformer.pyi index 4b9f2a7fb9a1..21796c04ff87 100644 --- a/stubs/sagemaker/sagemaker/transformer.pyi +++ b/stubs/sagemaker/sagemaker/transformer.pyi @@ -1,6 +1,6 @@ import abc from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.inputs import BatchDataCaptureConfig from sagemaker.job import _Job diff --git a/stubs/sagemaker/sagemaker/tuner.pyi b/stubs/sagemaker/sagemaker/tuner.pyi index 465e90501d20..533a4d896887 100644 --- a/stubs/sagemaker/sagemaker/tuner.pyi +++ b/stubs/sagemaker/sagemaker/tuner.pyi @@ -1,7 +1,7 @@ import abc from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Optional, Set, Union +from typing import Dict, List, Optional, Set from sagemaker.amazon.amazon_estimator import FileSystemRecordSet, RecordSet from sagemaker.estimator import EstimatorBase diff --git a/stubs/sagemaker/sagemaker/workflow/_utils.pyi b/stubs/sagemaker/sagemaker/workflow/_utils.pyi index c48cbe9c67df..3af080afd34c 100644 --- a/stubs/sagemaker/sagemaker/workflow/_utils.pyi +++ b/stubs/sagemaker/sagemaker/workflow/_utils.pyi @@ -1,10 +1,9 @@ from _typeshed import Incomplete -from typing import List, Optional, Union +from typing import List, Optional from sagemaker.estimator import EstimatorBase from sagemaker.workflow.entities import RequestType as RequestType from sagemaker.workflow.retry import RetryPolicy -from sagemaker.workflow.step_collections import StepCollection from sagemaker.workflow.steps import ConfigurableRetryStep, Step, TrainingStep logger: Incomplete diff --git a/stubs/sagemaker/sagemaker/workflow/automl_step.pyi b/stubs/sagemaker/sagemaker/workflow/automl_step.pyi index 1a67ba36d1f9..bdfa9e538b4b 100644 --- a/stubs/sagemaker/sagemaker/workflow/automl_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/automl_step.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import List, Optional, Union +from typing import List, Optional from sagemaker.workflow.entities import RequestType as RequestType from sagemaker.workflow.pipeline_context import _JobStepArguments diff --git a/stubs/sagemaker/sagemaker/workflow/callback_step.pyi b/stubs/sagemaker/sagemaker/workflow/callback_step.pyi index 671e82c5e3c3..8cdac4bd9204 100644 --- a/stubs/sagemaker/sagemaker/workflow/callback_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/callback_step.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.workflow.entities import DefaultEnumMeta, RequestType as RequestType from sagemaker.workflow.step_collections import StepCollection diff --git a/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi b/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi index 4f34112593c8..27ae9ed9f67f 100644 --- a/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from abc import ABC -from typing import List, Optional, Union +from typing import List, Optional from sagemaker.clarify import BiasConfig, DataConfig, ModelConfig, ModelPredictedLabelConfig, SHAPConfig from sagemaker.workflow.check_job_config import CheckJobConfig diff --git a/stubs/sagemaker/sagemaker/workflow/condition_step.pyi b/stubs/sagemaker/sagemaker/workflow/condition_step.pyi index bf1b47b32ddd..05b6a09d72cd 100644 --- a/stubs/sagemaker/sagemaker/workflow/condition_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/condition_step.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import List, Optional, Union +from typing import List, Optional from sagemaker.workflow.conditions import Condition from sagemaker.workflow.entities import RequestType as RequestType diff --git a/stubs/sagemaker/sagemaker/workflow/conditions.pyi b/stubs/sagemaker/sagemaker/workflow/conditions.pyi index 29187b98b603..a43ccf05a6e9 100644 --- a/stubs/sagemaker/sagemaker/workflow/conditions.pyi +++ b/stubs/sagemaker/sagemaker/workflow/conditions.pyi @@ -1,7 +1,7 @@ import abc from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Union +from typing import Dict, List from sagemaker.workflow.entities import ( DefaultEnumMeta, @@ -45,29 +45,19 @@ class ConditionComparison(Condition): def __ge__(self, other): ... class ConditionEquals(ConditionComparison): - def __init__( - self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType - ) -> None: ... + def __init__(self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType) -> None: ... class ConditionGreaterThan(ConditionComparison): - def __init__( - self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType - ) -> None: ... + def __init__(self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType) -> None: ... class ConditionGreaterThanOrEqualTo(ConditionComparison): - def __init__( - self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType - ) -> None: ... + def __init__(self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType) -> None: ... class ConditionLessThan(ConditionComparison): - def __init__( - self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType - ) -> None: ... + def __init__(self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType) -> None: ... class ConditionLessThanOrEqualTo(ConditionComparison): - def __init__( - self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType - ) -> None: ... + def __init__(self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType) -> None: ... class ConditionIn(Condition): value: Incomplete @@ -88,5 +78,5 @@ class ConditionOr(Condition): def to_request(self) -> RequestType: ... def primitive_or_expr( - value: ExecutionVariable | Expression | PrimitiveType | Parameter | Properties + value: ExecutionVariable | Expression | PrimitiveType | Parameter | Properties, ) -> Dict[str | str, PrimitiveType]: ... diff --git a/stubs/sagemaker/sagemaker/workflow/emr_step.pyi b/stubs/sagemaker/sagemaker/workflow/emr_step.pyi index 4581eed59b2f..46b355128d8d 100644 --- a/stubs/sagemaker/sagemaker/workflow/emr_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/emr_step.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional from sagemaker.workflow.entities import RequestType as RequestType from sagemaker.workflow.step_collections import StepCollection diff --git a/stubs/sagemaker/sagemaker/workflow/entities.pyi b/stubs/sagemaker/sagemaker/workflow/entities.pyi index efc729637e22..873279ea031a 100644 --- a/stubs/sagemaker/sagemaker/workflow/entities.pyi +++ b/stubs/sagemaker/sagemaker/workflow/entities.pyi @@ -1,7 +1,7 @@ import abc from _typeshed import Incomplete from enum import EnumMeta -from typing import Any, Dict, List, Union +from typing import Any, Dict, List PrimitiveType = str | int | bool | float | None RequestType = Dict[str | Any, List[Dict[str, Any]]] diff --git a/stubs/sagemaker/sagemaker/workflow/fail_step.pyi b/stubs/sagemaker/sagemaker/workflow/fail_step.pyi index 67a872311071..6b7a3bfeb2a6 100644 --- a/stubs/sagemaker/sagemaker/workflow/fail_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/fail_step.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import List, Optional, Union +from typing import List, Optional from sagemaker.workflow.entities import PipelineVariable, RequestType as RequestType from sagemaker.workflow.step_collections import StepCollection diff --git a/stubs/sagemaker/sagemaker/workflow/functions.pyi b/stubs/sagemaker/sagemaker/workflow/functions.pyi index 50d522031e2e..93f2bb85c9b5 100644 --- a/stubs/sagemaker/sagemaker/workflow/functions.pyi +++ b/stubs/sagemaker/sagemaker/workflow/functions.pyi @@ -1,4 +1,4 @@ -from typing import List, Union +from typing import List from sagemaker.workflow.entities import PipelineVariable from sagemaker.workflow.properties import PropertyFile diff --git a/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi b/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi index e5c91cedff98..c5495579961a 100644 --- a/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.lambda_helper import Lambda from sagemaker.workflow.entities import DefaultEnumMeta, RequestType as RequestType diff --git a/stubs/sagemaker/sagemaker/workflow/model_step.pyi b/stubs/sagemaker/sagemaker/workflow/model_step.pyi index 9e86b7241e3d..43c1c82f118f 100644 --- a/stubs/sagemaker/sagemaker/workflow/model_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/model_step.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.workflow.pipeline_context import _ModelStepArguments from sagemaker.workflow.retry import RetryPolicy diff --git a/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi b/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi index 50eee0e8d351..09355569bd83 100644 --- a/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional, Union +from typing import Optional from sagemaker.workflow.check_job_config import CheckJobConfig from sagemaker.workflow.clarify_check_step import ClarifyCheckConfig diff --git a/stubs/sagemaker/sagemaker/workflow/pipeline.pyi b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi index 5622736ad516..7b3abf2b45a1 100644 --- a/stubs/sagemaker/sagemaker/workflow/pipeline.pyi +++ b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Any, Dict, List, Optional, Sequence, Set, Union +from typing import Any, Dict, List, Optional, Sequence, Set from sagemaker.session import Session from sagemaker.workflow.entities import Entity, RequestType as RequestType diff --git a/stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi b/stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi index 51cd4602567b..ff69e202de8c 100644 --- a/stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi +++ b/stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Union from sagemaker.workflow.entities import Entity, Expression, RequestType as RequestType from sagemaker.workflow.execution_variables import ExecutionVariable diff --git a/stubs/sagemaker/sagemaker/workflow/properties.pyi b/stubs/sagemaker/sagemaker/workflow/properties.pyi index dca7119e4d6f..3dc9f41c2740 100644 --- a/stubs/sagemaker/sagemaker/workflow/properties.pyi +++ b/stubs/sagemaker/sagemaker/workflow/properties.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from abc import ABCMeta -from typing import Dict, List, Union +from typing import Dict, List from sagemaker.workflow.entities import Expression, PipelineVariable diff --git a/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi b/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi index b13daa368b8c..de6b5321cf23 100644 --- a/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from abc import ABC -from typing import List, Optional, Union +from typing import List, Optional from sagemaker.workflow.check_job_config import CheckJobConfig from sagemaker.workflow.entities import PipelineVariable, RequestType as RequestType diff --git a/stubs/sagemaker/sagemaker/workflow/step_collections.pyi b/stubs/sagemaker/sagemaker/workflow/step_collections.pyi index 09a887a67e9e..1063d4816c2d 100644 --- a/stubs/sagemaker/sagemaker/workflow/step_collections.pyi +++ b/stubs/sagemaker/sagemaker/workflow/step_collections.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import List, Optional, Union +from typing import List, Optional from sagemaker import PipelineModel from sagemaker.estimator import EstimatorBase diff --git a/stubs/sagemaker/sagemaker/workflow/steps.pyi b/stubs/sagemaker/sagemaker/workflow/steps.pyi index 1b7b9bfe6d60..ac3f5bbfc615 100644 --- a/stubs/sagemaker/sagemaker/workflow/steps.pyi +++ b/stubs/sagemaker/sagemaker/workflow/steps.pyi @@ -1,7 +1,7 @@ import abc from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.estimator import EstimatorBase from sagemaker.inputs import CreateModelInput, FileSystemInput, TrainingInput, TransformInput @@ -15,7 +15,6 @@ from sagemaker.workflow.functions import Join from sagemaker.workflow.pipeline_context import _JobStepArguments from sagemaker.workflow.properties import PropertyFile from sagemaker.workflow.retry import RetryPolicy -from sagemaker.workflow.step_collections import StepCollection class StepTypeEnum(Enum, metaclass=DefaultEnumMeta): CONDITION: str diff --git a/stubs/sagemaker/sagemaker/workflow/utilities.pyi b/stubs/sagemaker/sagemaker/workflow/utilities.pyi index f57c8f393728..67660d63c110 100644 --- a/stubs/sagemaker/sagemaker/workflow/utilities.pyi +++ b/stubs/sagemaker/sagemaker/workflow/utilities.pyi @@ -1,9 +1,8 @@ from _typeshed import Incomplete -from typing import List, Sequence, Set, Union +from typing import List, Sequence, Set from sagemaker.workflow.entities import Entity, RequestType as RequestType from sagemaker.workflow.pipeline_context import _StepArguments -from sagemaker.workflow.step_collections import StepCollection logger: Incomplete BUF_SIZE: int diff --git a/stubs/sagemaker/sagemaker/xgboost/estimator.pyi b/stubs/sagemaker/sagemaker/xgboost/estimator.pyi index a57dbedc9235..993065991227 100644 --- a/stubs/sagemaker/sagemaker/xgboost/estimator.pyi +++ b/stubs/sagemaker/sagemaker/xgboost/estimator.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, Optional, Union +from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/xgboost/model.pyi b/stubs/sagemaker/sagemaker/xgboost/model.pyi index 548b3028c656..f90f3f040839 100644 --- a/stubs/sagemaker/sagemaker/xgboost/model.pyi +++ b/stubs/sagemaker/sagemaker/xgboost/model.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker import ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/xgboost/processing.pyi b/stubs/sagemaker/sagemaker/xgboost/processing.pyi index efaf7365e677..e73715950366 100644 --- a/stubs/sagemaker/sagemaker/xgboost/processing.pyi +++ b/stubs/sagemaker/sagemaker/xgboost/processing.pyi @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional from sagemaker.network import NetworkConfig from sagemaker.processing import FrameworkProcessor From 67ee2b16013dda611bd890998bf9b81fcd76387c Mon Sep 17 00:00:00 2001 From: AlexWaygood Date: Sun, 25 Jun 2023 15:07:21 +0100 Subject: [PATCH 05/10] Fix many things using ruff --- stubs/sagemaker/sagemaker/accept_types.pyi | 14 +- stubs/sagemaker/sagemaker/algorithm.pyi | 34 +- .../sagemaker/amazon/amazon_estimator.pyi | 20 +- .../amazon/factorization_machines.pyi | 58 ++-- .../sagemaker/sagemaker/amazon/ipinsights.pyi | 28 +- stubs/sagemaker/sagemaker/amazon/kmeans.pyi | 30 +- stubs/sagemaker/sagemaker/amazon/knn.pyi | 28 +- stubs/sagemaker/sagemaker/amazon/lda.pyi | 18 +- .../sagemaker/amazon/linear_learner.pyi | 94 +++--- stubs/sagemaker/sagemaker/amazon/ntm.pyi | 34 +- .../sagemaker/sagemaker/amazon/object2vec.pyi | 76 ++--- stubs/sagemaker/sagemaker/amazon/pca.pyi | 18 +- .../sagemaker/amazon/randomcutforest.pyi | 16 +- stubs/sagemaker/sagemaker/automl/automl.pyi | 50 +-- stubs/sagemaker/sagemaker/base_predictor.pyi | 2 +- .../sagemaker/sagemaker/chainer/estimator.pyi | 18 +- stubs/sagemaker/sagemaker/chainer/model.pyi | 52 +-- stubs/sagemaker/sagemaker/clarify.pyi | 182 +++++------ stubs/sagemaker/sagemaker/collection.pyi | 10 +- stubs/sagemaker/sagemaker/config/config.pyi | 4 +- stubs/sagemaker/sagemaker/content_types.pyi | 14 +- .../sagemaker/sagemaker/debugger/debugger.pyi | 26 +- .../sagemaker/debugger/profiler_config.pyi | 8 +- stubs/sagemaker/sagemaker/deserializers.pyi | 14 +- .../sagemaker/djl_inference/model.pyi | 70 ++-- .../sagemaker/drift_check_baselines.pyi | 18 +- .../sagemaker/environment_variables.pyi | 8 +- stubs/sagemaker/sagemaker/estimator.pyi | 176 +++++----- .../sagemaker/experiments/_utils.pyi | 2 +- stubs/sagemaker/sagemaker/experiments/run.pyi | 36 +-- .../explainer/clarify_explainer_config.pyi | 36 +-- .../sagemaker/explainer/explainer_config.pyi | 2 +- .../feature_store/dataset_builder.pyi | 24 +- .../feature_store/feature_definition.pyi | 2 +- .../sagemaker/feature_store/feature_group.pyi | 55 ++-- .../sagemaker/feature_store/feature_store.pyi | 47 +-- .../sagemaker/feature_store/inputs.pyi | 26 +- stubs/sagemaker/sagemaker/fw_utils.pyi | 8 +- .../sagemaker/huggingface/estimator.pyi | 16 +- .../sagemaker/huggingface/llm_utils.pyi | 2 +- .../sagemaker/sagemaker/huggingface/model.pyi | 58 ++-- .../sagemaker/huggingface/processing.pyi | 34 +- stubs/sagemaker/sagemaker/hyperparameters.pyi | 16 +- .../inference_recommender_mixin.pyi | 28 +- stubs/sagemaker/sagemaker/inputs.pyi | 18 +- stubs/sagemaker/sagemaker/instance_types.pyi | 18 +- .../sagemaker/jumpstart/accessors.pyi | 6 +- stubs/sagemaker/sagemaker/jumpstart/cache.pyi | 6 +- .../sagemaker/jumpstart/constants.pyi | 10 +- .../sagemaker/jumpstart/estimator.pyi | 172 +++++----- .../sagemaker/jumpstart/exceptions.pyi | 14 +- .../sagemaker/jumpstart/factory/estimator.pyi | 188 +++++------ .../sagemaker/jumpstart/factory/model.pyi | 86 ++--- .../sagemaker/sagemaker/jumpstart/filters.pyi | 2 +- stubs/sagemaker/sagemaker/jumpstart/model.pyi | 80 ++--- .../sagemaker/jumpstart/notebook_utils.pyi | 10 +- stubs/sagemaker/sagemaker/jumpstart/types.pyi | 306 +++++++++--------- stubs/sagemaker/sagemaker/jumpstart/utils.pyi | 48 +-- .../sagemaker/jumpstart/validators.pyi | 4 +- stubs/sagemaker/sagemaker/lambda_helper.pyi | 22 +- stubs/sagemaker/sagemaker/lineage/action.pyi | 45 +-- .../sagemaker/sagemaker/lineage/artifact.pyi | 49 +-- .../sagemaker/lineage/association.pyi | 27 +- stubs/sagemaker/sagemaker/lineage/context.pyi | 49 +-- .../lineage/lineage_trial_component.pyi | 4 +- stubs/sagemaker/sagemaker/lineage/query.pyi | 24 +- .../sagemaker/lineage/visualizer.pyi | 18 +- .../sagemaker/local/local_session.pyi | 2 +- stubs/sagemaker/sagemaker/local/pipeline.pyi | 2 +- .../sagemaker/metadata_properties.pyi | 8 +- .../sagemaker/metric_definitions.pyi | 8 +- stubs/sagemaker/sagemaker/model.pyi | 100 +++--- .../sagemaker/model_card/helpers.pyi | 14 +- .../sagemaker/model_card/model_card.pyi | 148 ++++----- stubs/sagemaker/sagemaker/model_metrics.pyi | 6 +- .../data_quality_monitoring_config.pyi | 4 +- .../model_monitor/model_monitoring.pyi | 32 +- stubs/sagemaker/sagemaker/model_uris.pyi | 8 +- stubs/sagemaker/sagemaker/multidatamodel.pyi | 8 +- stubs/sagemaker/sagemaker/mxnet/estimator.pyi | 12 +- stubs/sagemaker/sagemaker/mxnet/model.pyi | 50 +-- .../sagemaker/sagemaker/mxnet/processing.pyi | 28 +- stubs/sagemaker/sagemaker/network.pyi | 8 +- stubs/sagemaker/sagemaker/pipeline.pyi | 54 ++-- stubs/sagemaker/sagemaker/predictor.pyi | 8 +- stubs/sagemaker/sagemaker/processing.pyi | 138 ++++---- .../sagemaker/sagemaker/pytorch/estimator.pyi | 14 +- stubs/sagemaker/sagemaker/pytorch/model.pyi | 50 +-- .../sagemaker/pytorch/processing.pyi | 28 +- .../sagemaker/remote_function/client.pyi | 76 ++--- .../sagemaker/remote_function/job.pyi | 40 +-- .../runtime_environment_manager.pyi | 4 +- .../remote_function/spark_config.pyi | 10 +- stubs/sagemaker/sagemaker/rl/estimator.pyi | 14 +- stubs/sagemaker/sagemaker/s3_utils.pyi | 2 +- stubs/sagemaker/sagemaker/script_uris.pyi | 8 +- stubs/sagemaker/sagemaker/serializers.pyi | 14 +- .../serverless_inference_config.pyi | 2 +- stubs/sagemaker/sagemaker/session.pyi | 125 +++---- .../sagemaker/sagemaker/sklearn/estimator.pyi | 10 +- stubs/sagemaker/sagemaker/sklearn/model.pyi | 50 +-- .../sagemaker/sklearn/processing.pyi | 24 +- .../sagemaker/sagemaker/spark/processing.pyi | 116 +++---- stubs/sagemaker/sagemaker/sparkml/model.pyi | 4 +- .../sagemaker/tensorflow/estimator.pyi | 12 +- .../sagemaker/sagemaker/tensorflow/model.pyi | 50 +-- .../sagemaker/tensorflow/processing.pyi | 28 +- stubs/sagemaker/sagemaker/transformer.pyi | 68 ++-- stubs/sagemaker/sagemaker/tuner.pyi | 46 +-- stubs/sagemaker/sagemaker/utilities/cache.pyi | 7 +- stubs/sagemaker/sagemaker/utils.pyi | 24 +- stubs/sagemaker/sagemaker/workflow/_utils.pyi | 30 +- .../sagemaker/workflow/automl_step.pyi | 10 +- .../sagemaker/workflow/callback_step.pyi | 12 +- .../sagemaker/workflow/clarify_check_step.pyi | 16 +- .../sagemaker/workflow/condition_step.pyi | 12 +- .../sagemaker/workflow/conditions.pyi | 6 +- .../sagemaker/sagemaker/workflow/emr_step.pyi | 10 +- .../sagemaker/sagemaker/workflow/entities.pyi | 2 +- .../sagemaker/workflow/fail_step.pyi | 8 +- .../sagemaker/workflow/functions.pyi | 2 +- .../sagemaker/workflow/lambda_step.pyi | 14 +- .../sagemaker/workflow/model_step.pyi | 8 +- .../workflow/monitor_batch_transform_step.pyi | 8 +- .../sagemaker/workflow/parameters.pyi | 10 +- .../sagemaker/sagemaker/workflow/pipeline.pyi | 61 ++-- .../sagemaker/workflow/pipeline_context.pyi | 11 +- .../sagemaker/workflow/properties.pyi | 12 +- .../sagemaker/workflow/quality_check_step.pyi | 14 +- stubs/sagemaker/sagemaker/workflow/retry.pyi | 14 +- .../workflow/selective_execution_config.pyi | 2 +- .../sagemaker/workflow/step_collections.pyi | 26 +- stubs/sagemaker/sagemaker/workflow/steps.pyi | 104 +++--- .../sagemaker/workflow/utilities.pyi | 15 +- .../sagemaker/wrangler/ingestion.pyi | 6 +- .../sagemaker/wrangler/processing.pyi | 24 +- .../sagemaker/sagemaker/xgboost/estimator.pyi | 8 +- stubs/sagemaker/sagemaker/xgboost/model.pyi | 50 +-- .../sagemaker/xgboost/processing.pyi | 28 +- 139 files changed, 2358 insertions(+), 2347 deletions(-) diff --git a/stubs/sagemaker/sagemaker/accept_types.pyi b/stubs/sagemaker/sagemaker/accept_types.pyi index 9f6a7a2b80fe..08282c899c31 100644 --- a/stubs/sagemaker/sagemaker/accept_types.pyi +++ b/stubs/sagemaker/sagemaker/accept_types.pyi @@ -1,16 +1,16 @@ from typing import List, Optional def retrieve_options( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, -) -> List[str]: ... +) -> list[str]: ... def retrieve_default( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, ) -> str: ... diff --git a/stubs/sagemaker/sagemaker/algorithm.pyi b/stubs/sagemaker/sagemaker/algorithm.pyi index 1bf184fea132..e93a25e20479 100644 --- a/stubs/sagemaker/sagemaker/algorithm.pyi +++ b/stubs/sagemaker/sagemaker/algorithm.pyi @@ -13,27 +13,27 @@ class AlgorithmEstimator(EstimatorBase): def __init__( self, algorithm_arn: str, - role: str = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, + role: str | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, volume_size: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, + volume_kms_key: str | PipelineVariable | None = None, max_run: int | PipelineVariable = 86400, input_mode: str | PipelineVariable = "File", - output_path: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - subnets: Optional[List[str | PipelineVariable]] = None, - security_group_ids: Optional[List[str | PipelineVariable]] = None, - model_uri: Optional[str] = None, + output_path: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + subnets: list[str | PipelineVariable] | None = None, + security_group_ids: list[str | PipelineVariable] | None = None, + model_uri: str | None = None, model_channel_name: str | PipelineVariable = "model", - metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, + metric_definitions: list[dict[str, str | PipelineVariable]] | None = None, encrypt_inter_container_traffic: bool | PipelineVariable = False, use_spot_instances: bool | PipelineVariable = False, - max_wait: Optional[int | PipelineVariable] = None, + max_wait: int | PipelineVariable | None = None, **kwargs, ) -> None: ... def validate_train_spec(self) -> None: ... @@ -68,8 +68,8 @@ class AlgorithmEstimator(EstimatorBase): ): ... def fit( self, - inputs: Optional[str | Dict | TrainingInput | FileSystemInput] = None, + inputs: str | dict | TrainingInput | FileSystemInput | None = None, wait: bool = True, logs: bool = True, - job_name: Optional[str] = None, + job_name: str | None = None, ): ... diff --git a/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi b/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi index 9ab83b9e93da..8f491abc4744 100644 --- a/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi +++ b/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi @@ -11,15 +11,15 @@ logger: Incomplete class AmazonAlgorithmEstimatorBase(EstimatorBase, metaclass=abc.ABCMeta): feature_dim: hp mini_batch_size: hp - repo_name: Optional[str] - repo_version: Optional[str] - DEFAULT_MINI_BATCH_SIZE: Optional[int] + repo_name: str | None + repo_version: str | None + DEFAULT_MINI_BATCH_SIZE: int | None def __init__( self, - role: Optional[str | PipelineVariable] = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, - data_location: Optional[str] = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + data_location: str | None = None, enable_network_isolation: bool | PipelineVariable = False, **kwargs, ) -> None: ... @@ -36,11 +36,11 @@ class AmazonAlgorithmEstimatorBase(EstimatorBase, metaclass=abc.ABCMeta): def fit( self, records: RecordSet, - mini_batch_size: Optional[int] = None, + mini_batch_size: int | None = None, wait: bool = True, logs: bool = True, - job_name: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, + job_name: str | None = None, + experiment_config: dict[str, str] | None = None, ): ... def record_set(self, train, labels: Incomplete | None = None, channel: str = "train", encrypt: bool = False): ... diff --git a/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi b/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi index 4137eefeac57..68668c6f3207 100644 --- a/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi +++ b/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi @@ -37,33 +37,33 @@ class FactorizationMachines(AmazonAlgorithmEstimatorBase): factors_init_value: hp def __init__( self, - role: Optional[str | PipelineVariable] = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[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, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + num_factors: int | None = None, + predictor_type: str | None = None, + epochs: int | None = None, + clip_gradient: float | None = None, + eps: float | None = None, + rescale_grad: float | None = None, + bias_lr: float | None = None, + linear_lr: float | None = None, + factors_lr: float | None = None, + bias_wd: float | None = None, + linear_wd: float | None = None, + factors_wd: float | None = None, + bias_init_method: str | None = None, + bias_init_scale: float | None = None, + bias_init_sigma: float | None = None, + bias_init_value: float | None = None, + linear_init_method: str | None = None, + linear_init_scale: float | None = None, + linear_init_sigma: float | None = None, + linear_init_value: float | None = None, + factors_init_method: str | None = None, + factors_init_scale: float | None = None, + factors_init_sigma: float | None = None, + factors_init_value: float | None = None, **kwargs, ) -> None: ... def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... @@ -75,7 +75,7 @@ class FactorizationMachinesModel(Model): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - sagemaker_session: Optional[Session] = None, + role: str | None = None, + sagemaker_session: Session | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi b/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi index 6f208e9a13b1..a60ceea84067 100644 --- a/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi +++ b/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi @@ -23,18 +23,18 @@ class IPInsights(AmazonAlgorithmEstimatorBase): weight_decay: hp def __init__( self, - role: Optional[str | PipelineVariable] = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[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, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + num_entity_vectors: int | None = None, + vector_dim: int | None = None, + batch_metrics_publish_interval: int | None = None, + epochs: int | None = None, + learning_rate: float | None = None, + num_ip_encoder_layers: int | None = None, + random_negative_sampling_rate: int | None = None, + shuffled_negative_sampling_rate: int | None = None, + weight_decay: float | None = None, **kwargs, ) -> None: ... def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... @@ -46,7 +46,7 @@ class IPInsightsModel(Model): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - sagemaker_session: Optional[Session] = None, + role: str | None = None, + sagemaker_session: Session | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/kmeans.pyi b/stubs/sagemaker/sagemaker/amazon/kmeans.pyi index f99bc3da1309..c2bd3bed0dac 100644 --- a/stubs/sagemaker/sagemaker/amazon/kmeans.pyi +++ b/stubs/sagemaker/sagemaker/amazon/kmeans.pyi @@ -23,19 +23,19 @@ class KMeans(AmazonAlgorithmEstimatorBase): eval_metrics: hp def __init__( self, - role: Optional[str | PipelineVariable] = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[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[str | PipelineVariable]] = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + k: int | None = None, + init_method: str | None = None, + max_iterations: int | None = None, + tol: float | None = None, + num_trials: int | None = None, + local_init_method: str | None = None, + half_life_time_size: int | None = None, + epochs: int | None = None, + center_factor: int | None = None, + eval_metrics: list[str | PipelineVariable] | None = None, **kwargs, ) -> None: ... def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... @@ -48,7 +48,7 @@ class KMeansModel(Model): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - sagemaker_session: Optional[Session] = None, + role: str | None = None, + sagemaker_session: Session | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/knn.pyi b/stubs/sagemaker/sagemaker/amazon/knn.pyi index 808e9883e233..06636307df50 100644 --- a/stubs/sagemaker/sagemaker/amazon/knn.pyi +++ b/stubs/sagemaker/sagemaker/amazon/knn.pyi @@ -22,18 +22,18 @@ class KNN(AmazonAlgorithmEstimatorBase): faiss_index_pq_m: hp def __init__( self, - role: Optional[str | PipelineVariable] = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, - k: Optional[int] = None, - sample_size: Optional[int] = None, - predictor_type: Optional[str] = None, - dimension_reduction_type: Optional[str] = None, - dimension_reduction_target: Optional[int] = None, - index_type: Optional[str] = None, - index_metric: Optional[str] = None, - faiss_index_ivf_nlists: Optional[str] = None, - faiss_index_pq_m: Optional[int] = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + k: int | None = None, + sample_size: int | None = None, + predictor_type: str | None = None, + dimension_reduction_type: str | None = None, + dimension_reduction_target: int | None = None, + index_type: str | None = None, + index_metric: str | None = None, + faiss_index_ivf_nlists: str | None = None, + faiss_index_pq_m: int | None = None, **kwargs, ) -> None: ... def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... @@ -45,7 +45,7 @@ class KNNModel(Model): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - sagemaker_session: Optional[Session] = None, + role: str | None = None, + sagemaker_session: Session | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/lda.pyi b/stubs/sagemaker/sagemaker/amazon/lda.pyi index 44920b303c97..e0b1ec0c7247 100644 --- a/stubs/sagemaker/sagemaker/amazon/lda.pyi +++ b/stubs/sagemaker/sagemaker/amazon/lda.pyi @@ -20,13 +20,13 @@ class LDA(AmazonAlgorithmEstimatorBase): tol: hp def __init__( self, - role: Optional[str | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, - num_topics: Optional[int] = None, - alpha0: Optional[float] = None, - max_restarts: Optional[int] = None, - max_iterations: Optional[int] = None, - tol: Optional[float] = None, + role: str | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + num_topics: int | None = None, + alpha0: float | None = None, + max_restarts: int | None = None, + max_iterations: int | None = None, + tol: float | None = None, **kwargs, ) -> None: ... def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... @@ -38,7 +38,7 @@ class LDAModel(Model): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - sagemaker_session: Optional[Session] = None, + role: str | None = None, + sagemaker_session: Session | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi b/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi index 9443a79d2d79..2011a49f3374 100644 --- a/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi +++ b/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi @@ -58,51 +58,51 @@ class LinearLearner(AmazonAlgorithmEstimatorBase): balance_multiclass_weights: hp def __init__( self, - role: Optional[str | PipelineVariable] = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, - predictor_type: Optional[str] = None, - binary_classifier_model_selection_criteria: Optional[str] = None, - target_recall: Optional[float] = None, - target_precision: Optional[float] = None, - positive_example_weight_mult: Optional[float] = None, - epochs: Optional[int] = None, - use_bias: Optional[bool] = None, - num_models: Optional[int] = None, - num_calibration_samples: Optional[int] = None, - init_method: Optional[str] = None, - init_scale: Optional[float] = None, - init_sigma: Optional[float] = None, - init_bias: Optional[float] = None, - optimizer: Optional[str] = None, - loss: Optional[str] = None, - wd: Optional[float] = None, - l1: Optional[float] = None, - momentum: Optional[float] = None, - learning_rate: Optional[float] = None, - beta_1: Optional[float] = None, - beta_2: Optional[float] = None, - bias_lr_mult: Optional[float] = None, - bias_wd_mult: Optional[float] = None, - use_lr_scheduler: Optional[bool] = None, - lr_scheduler_step: Optional[int] = None, - lr_scheduler_factor: Optional[float] = None, - lr_scheduler_minimum_lr: Optional[float] = None, - normalize_data: Optional[bool] = None, - normalize_label: Optional[bool] = None, - unbias_data: Optional[bool] = None, - unbias_label: Optional[bool] = None, - num_point_for_scaler: Optional[int] = None, - margin: Optional[float] = None, - quantile: Optional[float] = None, - loss_insensitivity: Optional[float] = None, - huber_delta: Optional[float] = None, - early_stopping_patience: Optional[int] = None, - early_stopping_tolerance: Optional[float] = None, - num_classes: Optional[int] = None, - accuracy_top_k: Optional[int] = None, - f_beta: Optional[float] = None, - balance_multiclass_weights: Optional[bool] = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + predictor_type: str | None = None, + binary_classifier_model_selection_criteria: str | None = None, + target_recall: float | None = None, + target_precision: float | None = None, + positive_example_weight_mult: float | None = None, + epochs: int | None = None, + use_bias: bool | None = None, + num_models: int | None = None, + num_calibration_samples: int | None = None, + init_method: str | None = None, + init_scale: float | None = None, + init_sigma: float | None = None, + init_bias: float | None = None, + optimizer: str | None = None, + loss: str | None = None, + wd: float | None = None, + l1: float | None = None, + momentum: float | None = None, + learning_rate: float | None = None, + beta_1: float | None = None, + beta_2: float | None = None, + bias_lr_mult: float | None = None, + bias_wd_mult: float | None = None, + use_lr_scheduler: bool | None = None, + lr_scheduler_step: int | None = None, + lr_scheduler_factor: float | None = None, + lr_scheduler_minimum_lr: float | None = None, + normalize_data: bool | None = None, + normalize_label: bool | None = None, + unbias_data: bool | None = None, + unbias_label: bool | None = None, + num_point_for_scaler: int | None = None, + margin: float | None = None, + quantile: float | None = None, + loss_insensitivity: float | None = None, + huber_delta: float | None = None, + early_stopping_patience: int | None = None, + early_stopping_tolerance: float | None = None, + num_classes: int | None = None, + accuracy_top_k: int | None = None, + f_beta: float | None = None, + balance_multiclass_weights: bool | None = None, **kwargs, ) -> None: ... def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... @@ -114,7 +114,7 @@ class LinearLearnerModel(Model): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - sagemaker_session: Optional[Session] = None, + role: str | None = None, + sagemaker_session: Session | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/ntm.pyi b/stubs/sagemaker/sagemaker/amazon/ntm.pyi index c7ab3b7b4685..5d8b67dfa1d1 100644 --- a/stubs/sagemaker/sagemaker/amazon/ntm.pyi +++ b/stubs/sagemaker/sagemaker/amazon/ntm.pyi @@ -25,21 +25,21 @@ class NTM(AmazonAlgorithmEstimatorBase): learning_rate: hp def __init__( self, - role: Optional[str | PipelineVariable] = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, - num_topics: Optional[int] = None, - encoder_layers: Optional[List] = None, - epochs: Optional[int] = None, - encoder_layers_activation: Optional[str] = None, - optimizer: Optional[str] = None, - tolerance: Optional[float] = None, - num_patience_epochs: Optional[int] = None, - batch_norm: Optional[bool] = None, - rescale_gradient: Optional[float] = None, - clip_gradient: Optional[float] = None, - weight_decay: Optional[float] = None, - learning_rate: Optional[float] = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + num_topics: int | None = None, + encoder_layers: list | None = None, + epochs: int | None = None, + encoder_layers_activation: str | None = None, + optimizer: str | None = None, + tolerance: float | None = None, + num_patience_epochs: int | None = None, + batch_norm: bool | None = None, + rescale_gradient: float | None = None, + clip_gradient: float | None = None, + weight_decay: float | None = None, + learning_rate: float | None = None, **kwargs, ) -> None: ... def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... @@ -51,7 +51,7 @@ class NTMModel(Model): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - sagemaker_session: Optional[Session] = None, + role: str | None = None, + sagemaker_session: Session | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/object2vec.pyi b/stubs/sagemaker/sagemaker/amazon/object2vec.pyi index 3153b860cafb..cdad95098444 100644 --- a/stubs/sagemaker/sagemaker/amazon/object2vec.pyi +++ b/stubs/sagemaker/sagemaker/amazon/object2vec.pyi @@ -45,42 +45,42 @@ class Object2Vec(AmazonAlgorithmEstimatorBase): enc1_freeze_pretrained_embedding: hp def __init__( self, - role: Optional[str | PipelineVariable] = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, - epochs: Optional[int] = None, - enc0_max_seq_len: Optional[int] = None, - enc0_vocab_size: Optional[int] = None, - enc_dim: Optional[int] = None, - mini_batch_size: Optional[int] = None, - early_stopping_patience: Optional[int] = None, - early_stopping_tolerance: Optional[float] = None, - dropout: Optional[float] = None, - weight_decay: Optional[float] = None, - bucket_width: Optional[int] = None, - num_classes: Optional[int] = None, - mlp_layers: Optional[int] = None, - mlp_dim: Optional[int] = None, - mlp_activation: Optional[str] = None, - output_layer: Optional[str] = None, - optimizer: Optional[str] = None, - learning_rate: Optional[float] = None, - negative_sampling_rate: Optional[int] = None, - comparator_list: Optional[str] = None, - tied_token_embedding_weight: Optional[bool] = None, - token_embedding_storage_type: Optional[str] = None, - enc0_network: Optional[str] = None, - enc1_network: Optional[str] = None, - enc0_cnn_filter_width: Optional[int] = None, - enc1_cnn_filter_width: Optional[int] = None, - enc1_max_seq_len: Optional[int] = None, - enc0_token_embedding_dim: Optional[int] = None, - enc1_token_embedding_dim: Optional[int] = None, - enc1_vocab_size: Optional[int] = None, - enc0_layers: Optional[int] = None, - enc1_layers: Optional[int] = None, - enc0_freeze_pretrained_embedding: Optional[bool] = None, - enc1_freeze_pretrained_embedding: Optional[bool] = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + epochs: int | None = None, + enc0_max_seq_len: int | None = None, + enc0_vocab_size: int | None = None, + enc_dim: int | None = None, + mini_batch_size: int | None = None, + early_stopping_patience: int | None = None, + early_stopping_tolerance: float | None = None, + dropout: float | None = None, + weight_decay: float | None = None, + bucket_width: int | None = None, + num_classes: int | None = None, + mlp_layers: int | None = None, + mlp_dim: int | None = None, + mlp_activation: str | None = None, + output_layer: str | None = None, + optimizer: str | None = None, + learning_rate: float | None = None, + negative_sampling_rate: int | None = None, + comparator_list: str | None = None, + tied_token_embedding_weight: bool | None = None, + token_embedding_storage_type: str | None = None, + enc0_network: str | None = None, + enc1_network: str | None = None, + enc0_cnn_filter_width: int | None = None, + enc1_cnn_filter_width: int | None = None, + enc1_max_seq_len: int | None = None, + enc0_token_embedding_dim: int | None = None, + enc1_token_embedding_dim: int | None = None, + enc1_vocab_size: int | None = None, + enc0_layers: int | None = None, + enc1_layers: int | None = None, + enc0_freeze_pretrained_embedding: bool | None = None, + enc1_freeze_pretrained_embedding: bool | None = None, **kwargs, ) -> None: ... def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... @@ -89,7 +89,7 @@ class Object2VecModel(Model): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - sagemaker_session: Optional[Session] = None, + role: str | None = None, + sagemaker_session: Session | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/pca.pyi b/stubs/sagemaker/sagemaker/amazon/pca.pyi index a1bd6fe56697..958f083fe6e9 100644 --- a/stubs/sagemaker/sagemaker/amazon/pca.pyi +++ b/stubs/sagemaker/sagemaker/amazon/pca.pyi @@ -18,13 +18,13 @@ class PCA(AmazonAlgorithmEstimatorBase): extra_components: hp def __init__( self, - role: Optional[str | PipelineVariable] = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, - num_components: Optional[int] = None, - algorithm_mode: Optional[str] = None, - subtract_mean: Optional[bool] = None, - extra_components: Optional[int] = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + num_components: int | None = None, + algorithm_mode: str | None = None, + subtract_mean: bool | None = None, + extra_components: int | None = None, **kwargs, ) -> None: ... def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... @@ -36,7 +36,7 @@ class PCAModel(Model): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - sagemaker_session: Optional[Session] = None, + role: str | None = None, + sagemaker_session: Session | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi b/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi index cc8391aa9183..6d79c22be3c5 100644 --- a/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi +++ b/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi @@ -18,12 +18,12 @@ class RandomCutForest(AmazonAlgorithmEstimatorBase): feature_dim: hp def __init__( self, - role: Optional[str | PipelineVariable] = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, - num_samples_per_tree: Optional[int] = None, - num_trees: Optional[int] = None, - eval_metrics: Optional[List] = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + num_samples_per_tree: int | None = None, + num_trees: int | None = None, + eval_metrics: list | None = None, **kwargs, ) -> None: ... def create_model(self, vpc_config_override="VPC_CONFIG_DEFAULT", **kwargs): ... @@ -35,7 +35,7 @@ class RandomCutForestModel(Model): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - sagemaker_session: Optional[Session] = None, + role: str | None = None, + sagemaker_session: Session | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/automl/automl.pyi b/stubs/sagemaker/sagemaker/automl/automl.pyi index aff51516981b..c222638df7d1 100644 --- a/stubs/sagemaker/sagemaker/automl/automl.pyi +++ b/stubs/sagemaker/sagemaker/automl/automl.pyi @@ -56,31 +56,31 @@ class AutoML: sample_weight_attribute_name: Incomplete def __init__( self, - role: Optional[str] = None, - target_attribute_name: str = None, - output_kms_key: Optional[str] = None, - output_path: Optional[str] = None, - base_job_name: Optional[str] = None, - compression_type: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - volume_kms_key: Optional[str] = None, - encrypt_inter_container_traffic: Optional[bool] = None, - vpc_config: Optional[Dict[str, List]] = None, - problem_type: Optional[str] = None, - max_candidates: Optional[int] = None, - max_runtime_per_training_job_in_seconds: Optional[int] = None, - total_job_runtime_in_seconds: Optional[int] = None, - job_objective: Optional[Dict[str, str]] = None, - generate_candidate_definitions_only: Optional[bool] = False, - tags: Optional[List[Dict[str, str]]] = None, - content_type: Optional[str] = None, - s3_data_type: Optional[str] = None, - feature_specification_s3_uri: Optional[str] = None, - validation_fraction: Optional[float] = None, - mode: Optional[str] = None, - auto_generate_endpoint_name: Optional[bool] = None, - endpoint_name: Optional[str] = None, - sample_weight_attribute_name: str = None, + role: str | None = None, + target_attribute_name: str | None = None, + output_kms_key: str | None = None, + output_path: str | None = None, + base_job_name: str | None = None, + compression_type: str | None = None, + sagemaker_session: Session | None = None, + volume_kms_key: str | None = None, + encrypt_inter_container_traffic: bool | None = None, + vpc_config: dict[str, list] | None = None, + problem_type: str | None = None, + max_candidates: int | None = None, + max_runtime_per_training_job_in_seconds: int | None = None, + total_job_runtime_in_seconds: int | None = None, + job_objective: dict[str, str] | None = None, + generate_candidate_definitions_only: bool | None = False, + tags: list[dict[str, str]] | None = None, + content_type: str | None = None, + s3_data_type: str | None = None, + feature_specification_s3_uri: str | None = None, + validation_fraction: float | None = None, + mode: str | None = None, + auto_generate_endpoint_name: bool | None = None, + endpoint_name: str | None = None, + sample_weight_attribute_name: str | None = None, ) -> None: ... latest_auto_ml_job: Incomplete def fit( diff --git a/stubs/sagemaker/sagemaker/base_predictor.pyi b/stubs/sagemaker/sagemaker/base_predictor.pyi index 7d6dd88bb122..1f87d79c2cb4 100644 --- a/stubs/sagemaker/sagemaker/base_predictor.pyi +++ b/stubs/sagemaker/sagemaker/base_predictor.pyi @@ -14,7 +14,7 @@ class PredictorBase(abc.ABC, metaclass=abc.ABCMeta): def content_type(self) -> str: ... @property @abc.abstractmethod - def accept(self) -> Tuple[str]: ... + def accept(self) -> tuple[str]: ... class Predictor(PredictorBase): endpoint_name: Incomplete diff --git a/stubs/sagemaker/sagemaker/chainer/estimator.pyi b/stubs/sagemaker/sagemaker/chainer/estimator.pyi index 2295931d7429..9fefc5424cc0 100644 --- a/stubs/sagemaker/sagemaker/chainer/estimator.pyi +++ b/stubs/sagemaker/sagemaker/chainer/estimator.pyi @@ -16,15 +16,15 @@ class Chainer(Framework): def __init__( self, entry_point: str | PipelineVariable, - use_mpi: Optional[bool | PipelineVariable] = None, - num_processes: Optional[int | PipelineVariable] = None, - process_slots_per_host: Optional[int | PipelineVariable] = None, - additional_mpi_options: Optional[str | PipelineVariable] = None, - source_dir: Optional[str | PipelineVariable] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, - framework_version: Optional[str] = None, - py_version: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, + use_mpi: bool | PipelineVariable | None = None, + num_processes: int | PipelineVariable | None = None, + process_slots_per_host: int | PipelineVariable | None = None, + additional_mpi_options: str | PipelineVariable | None = None, + source_dir: str | PipelineVariable | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, + framework_version: str | None = None, + py_version: str | None = None, + image_uri: str | PipelineVariable | None = None, **kwargs, ) -> None: ... def hyperparameters(self): ... diff --git a/stubs/sagemaker/sagemaker/chainer/model.pyi b/stubs/sagemaker/sagemaker/chainer/model.pyi index 64f15000b7b2..a2ea2c387fac 100644 --- a/stubs/sagemaker/sagemaker/chainer/model.pyi +++ b/stubs/sagemaker/sagemaker/chainer/model.pyi @@ -20,39 +20,39 @@ class ChainerModel(FrameworkModel): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - entry_point: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, - framework_version: Optional[str] = None, - py_version: Optional[str] = None, + role: str | None = None, + entry_point: str | None = None, + image_uri: str | PipelineVariable | None = None, + framework_version: str | None = None, + py_version: str | None = None, predictor_cls: callable = ..., - model_server_workers: Optional[int | PipelineVariable] = None, + model_server_workers: int | PipelineVariable | None = None, **kwargs, ) -> None: ... image_uri: Incomplete def register( self, - content_types: List[str | PipelineVariable], - response_types: List[str | PipelineVariable], - inference_instances: Optional[List[str | PipelineVariable]] = None, - transform_instances: Optional[List[str | PipelineVariable]] = None, - model_package_name: Optional[str | PipelineVariable] = None, - model_package_group_name: Optional[str | PipelineVariable] = None, - image_uri: Optional[str | PipelineVariable] = None, - model_metrics: Optional[ModelMetrics] = None, - metadata_properties: Optional[MetadataProperties] = None, + content_types: list[str | PipelineVariable], + response_types: list[str | PipelineVariable], + inference_instances: list[str | PipelineVariable] | None = None, + transform_instances: list[str | PipelineVariable] | None = None, + model_package_name: str | PipelineVariable | None = None, + model_package_group_name: str | PipelineVariable | None = None, + image_uri: str | PipelineVariable | None = None, + model_metrics: ModelMetrics | None = None, + metadata_properties: MetadataProperties | None = None, marketplace_cert: bool = False, - approval_status: Optional[str | PipelineVariable] = None, - description: Optional[str] = None, - drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, - domain: Optional[str | PipelineVariable] = None, - sample_payload_url: Optional[str | PipelineVariable] = None, - task: Optional[str | PipelineVariable] = None, - framework: Optional[str | PipelineVariable] = None, - framework_version: Optional[str | PipelineVariable] = None, - nearest_model_name: Optional[str | PipelineVariable] = None, - data_input_configuration: Optional[str | PipelineVariable] = None, + approval_status: str | PipelineVariable | None = None, + description: str | None = None, + drift_check_baselines: DriftCheckBaselines | None = None, + customer_metadata_properties: dict[str, str | PipelineVariable] | None = None, + domain: str | PipelineVariable | None = None, + sample_payload_url: str | PipelineVariable | None = None, + task: str | PipelineVariable | None = None, + framework: str | PipelineVariable | None = None, + framework_version: str | PipelineVariable | None = None, + nearest_model_name: str | PipelineVariable | None = None, + data_input_configuration: str | PipelineVariable | None = None, ): ... def prepare_container_def( self, diff --git a/stubs/sagemaker/sagemaker/clarify.pyi b/stubs/sagemaker/sagemaker/clarify.pyi index 0fd3204e8d4a..c027928f133a 100644 --- a/stubs/sagemaker/sagemaker/clarify.pyi +++ b/stubs/sagemaker/sagemaker/clarify.pyi @@ -39,19 +39,19 @@ class DataConfig: self, s3_data_input_path: str, s3_output_path: str, - s3_analysis_config_output_path: Optional[str] = None, - label: Optional[str] = None, - headers: Optional[List[str]] = None, - features: Optional[str] = None, + s3_analysis_config_output_path: str | None = None, + label: str | None = None, + headers: list[str] | None = None, + features: str | None = None, dataset_type: str = "text/csv", s3_compression_type: str = "None", - joinsource: Optional[str | int] = None, - facet_dataset_uri: Optional[str] = None, - facet_headers: Optional[List[str]] = None, - predicted_label_dataset_uri: Optional[str] = None, - predicted_label_headers: Optional[List[str]] = None, - predicted_label: Optional[str | int] = None, - excluded_columns: Optional[List[int, List[str]]] = None, + joinsource: str | int | None = None, + facet_dataset_uri: str | None = None, + facet_headers: list[str] | None = None, + predicted_label_dataset_uri: str | None = None, + predicted_label_headers: list[str] | None = None, + predicted_label: str | int | None = None, + excluded_columns: list[int, list[str]] | None = None, ) -> None: ... def get_config(self): ... @@ -60,9 +60,9 @@ class BiasConfig: def __init__( self, label_values_or_threshold: int | float | str, - facet_name: str | int | List[str, List[int]], - facet_values_or_threshold: Optional[int | float | str] = None, - group_name: Optional[str] = None, + facet_name: str | int | list[str, list[int]], + facet_values_or_threshold: int | float | str | None = None, + group_name: str | None = None, ) -> None: ... def get_config(self): ... @@ -70,18 +70,18 @@ class ModelConfig: predictor_config: Incomplete def __init__( self, - model_name: Optional[str] = None, - instance_count: Optional[int] = None, - instance_type: Optional[str] = None, - accept_type: Optional[str] = None, - content_type: Optional[str] = None, - content_template: Optional[str] = None, - record_template: Optional[str] = None, - custom_attributes: Optional[str] = None, - accelerator_type: Optional[str] = None, - endpoint_name_prefix: Optional[str] = None, - target_model: Optional[str] = None, - endpoint_name: Optional[str] = None, + model_name: str | None = None, + instance_count: int | None = None, + instance_type: str | None = None, + accept_type: str | None = None, + content_type: str | None = None, + content_template: str | None = None, + record_template: str | None = None, + custom_attributes: str | None = None, + accelerator_type: str | None = None, + endpoint_name_prefix: str | None = None, + target_model: str | None = None, + endpoint_name: str | None = None, ) -> None: ... def get_predictor_config(self): ... @@ -93,10 +93,10 @@ class ModelPredictedLabelConfig: predictor_config: Incomplete def __init__( self, - label: Optional[str | int] = None, - probability: Optional[str | int] = None, - probability_threshold: Optional[float] = None, - label_headers: Optional[List[str]] = None, + label: str | int | None = None, + probability: str | int | None = None, + probability_threshold: float | None = None, + label_headers: list[str] | None = None, ) -> None: ... def get_predictor_config(self): ... @@ -106,7 +106,7 @@ class ExplainabilityConfig(ABC, metaclass=abc.ABCMeta): class PDPConfig(ExplainabilityConfig): pdp_config: Incomplete - def __init__(self, features: Optional[List] = None, grid_resolution: int = 15, top_k_features: int = 10) -> None: ... + def __init__(self, features: list | None = None, grid_resolution: int = 15, top_k_features: int = 10) -> None: ... def get_explainability_config(self): ... class TextConfig: @@ -119,12 +119,12 @@ class ImageConfig: def __init__( self, model_type: str, - num_segments: Optional[int] = None, - feature_extraction_method: Optional[str] = None, - segment_compactness: Optional[float] = None, - max_objects: Optional[int] = None, - iou_threshold: Optional[float] = None, - context: Optional[float] = None, + num_segments: int | None = None, + feature_extraction_method: str | None = None, + segment_compactness: float | None = None, + max_objects: int | None = None, + iou_threshold: float | None = None, + context: float | None = None, ) -> None: ... def get_image_config(self): ... @@ -132,15 +132,15 @@ class SHAPConfig(ExplainabilityConfig): shap_config: Incomplete def __init__( self, - baseline: Optional[str | List | Dict] = None, - num_samples: Optional[int] = None, - agg_method: Optional[str] = None, + baseline: str | list | dict | None = None, + num_samples: int | None = None, + agg_method: str | None = None, use_logit: bool = False, save_local_shap_values: bool = True, - seed: Optional[int] = None, - num_clusters: Optional[int] = None, - text_config: Optional[TextConfig] = None, - image_config: Optional[ImageConfig] = None, + seed: int | None = None, + num_clusters: int | None = None, + text_config: TextConfig | None = None, + image_config: ImageConfig | None = None, ) -> None: ... def get_explainability_config(self): ... @@ -149,19 +149,19 @@ class SageMakerClarifyProcessor(Processor): skip_early_validation: Incomplete def __init__( self, - role: Optional[str] = None, - instance_count: int = None, - instance_type: str = None, + role: str | None = None, + instance_count: int | None = None, + instance_type: str | None = None, volume_size_in_gb: int = 30, - volume_kms_key: Optional[str] = None, - output_kms_key: Optional[str] = None, - max_runtime_in_seconds: Optional[int] = None, - sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, str]] = None, - tags: Optional[List[Dict[str, str]]] = None, - network_config: Optional[NetworkConfig] = None, - job_name_prefix: Optional[str] = None, - version: Optional[str] = None, + volume_kms_key: str | None = None, + output_kms_key: str | None = None, + max_runtime_in_seconds: int | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str] | None = None, + tags: list[dict[str, str]] | None = None, + network_config: NetworkConfig | None = None, + job_name_prefix: str | None = None, + version: str | None = None, skip_early_validation: bool = False, ) -> None: ... def run(self, **_) -> None: ... @@ -169,61 +169,61 @@ class SageMakerClarifyProcessor(Processor): self, data_config: DataConfig, data_bias_config: BiasConfig, - methods: str | List[str] = "all", + methods: str | list[str] = "all", wait: bool = True, logs: bool = True, - job_name: Optional[str] = None, - kms_key: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, + job_name: str | None = None, + kms_key: str | None = None, + experiment_config: dict[str, str] | None = None, ): ... def run_post_training_bias( self, data_config: DataConfig, data_bias_config: BiasConfig, - model_config: Optional[ModelConfig] = None, - model_predicted_label_config: Optional[ModelPredictedLabelConfig] = None, - methods: str | List[str] = "all", + model_config: ModelConfig | None = None, + model_predicted_label_config: ModelPredictedLabelConfig | None = None, + methods: str | list[str] = "all", wait: bool = True, logs: bool = True, - job_name: Optional[str] = None, - kms_key: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, + job_name: str | None = None, + kms_key: str | None = None, + experiment_config: dict[str, str] | None = None, ): ... def run_bias( self, data_config: DataConfig, bias_config: BiasConfig, - model_config: Optional[ModelConfig] = None, - model_predicted_label_config: Optional[ModelPredictedLabelConfig] = None, - pre_training_methods: str | List[str] = "all", - post_training_methods: str | List[str] = "all", + model_config: ModelConfig | None = None, + model_predicted_label_config: ModelPredictedLabelConfig | None = None, + pre_training_methods: str | list[str] = "all", + post_training_methods: str | list[str] = "all", wait: bool = True, logs: bool = True, - job_name: Optional[str] = None, - kms_key: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, + job_name: str | None = None, + kms_key: str | None = None, + experiment_config: dict[str, str] | None = None, ): ... def run_explainability( self, data_config: DataConfig, model_config: ModelConfig, - explainability_config: ExplainabilityConfig | List, - model_scores: Optional[int | str | ModelPredictedLabelConfig] = None, + explainability_config: ExplainabilityConfig | list, + model_scores: int | str | ModelPredictedLabelConfig | None = None, wait: bool = True, logs: bool = True, - job_name: Optional[str] = None, - kms_key: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, + job_name: str | None = None, + kms_key: str | None = None, + experiment_config: dict[str, str] | None = None, ): ... def run_bias_and_explainability( self, data_config: DataConfig, model_config: ModelConfig, - explainability_config: ExplainabilityConfig | List[ExplainabilityConfig], + explainability_config: ExplainabilityConfig | list[ExplainabilityConfig], bias_config: BiasConfig, - pre_training_methods: str | List[str] = "all", - post_training_methods: str | List[str] = "all", - model_predicted_label_config: ModelPredictedLabelConfig = None, + pre_training_methods: str | list[str] = "all", + post_training_methods: str | list[str] = "all", + model_predicted_label_config: ModelPredictedLabelConfig | None = None, wait: bool = True, logs: bool = True, job_name: Incomplete | None = None, @@ -238,10 +238,10 @@ class _AnalysisConfigGenerator: data_config: DataConfig, model_config: ModelConfig, model_predicted_label_config: ModelPredictedLabelConfig, - explainability_config: ExplainabilityConfig | List[ExplainabilityConfig], + explainability_config: ExplainabilityConfig | list[ExplainabilityConfig], bias_config: BiasConfig, - pre_training_methods: str | List[str] = "all", - post_training_methods: str | List[str] = "all", + pre_training_methods: str | list[str] = "all", + post_training_methods: str | list[str] = "all", ): ... @classmethod def explainability( @@ -249,17 +249,17 @@ class _AnalysisConfigGenerator: data_config: DataConfig, model_config: ModelConfig, model_predicted_label_config: ModelPredictedLabelConfig, - explainability_config: ExplainabilityConfig | List[ExplainabilityConfig], + explainability_config: ExplainabilityConfig | list[ExplainabilityConfig], ): ... @classmethod - def bias_pre_training(cls, data_config: DataConfig, bias_config: BiasConfig, methods: str | List[str]): ... + def bias_pre_training(cls, data_config: DataConfig, bias_config: BiasConfig, methods: str | list[str]): ... @classmethod def bias_post_training( cls, data_config: DataConfig, bias_config: BiasConfig, model_predicted_label_config: ModelPredictedLabelConfig, - methods: str | List[str], + methods: str | list[str], model_config: ModelConfig, ): ... @classmethod @@ -269,8 +269,8 @@ class _AnalysisConfigGenerator: bias_config: BiasConfig, model_config: ModelConfig, model_predicted_label_config: ModelPredictedLabelConfig, - pre_training_methods: str | List[str] = "all", - post_training_methods: str | List[str] = "all", + pre_training_methods: str | list[str] = "all", + post_training_methods: str | list[str] = "all", ): ... class ProcessingOutputHandler: @@ -278,4 +278,4 @@ class ProcessingOutputHandler: CONTINUOUS: str ENDOFJOB: str @classmethod - def get_s3_upload_mode(cls, analysis_config: Dict[str, Any]) -> str: ... + def get_s3_upload_mode(cls, analysis_config: dict[str, Any]) -> str: ... diff --git a/stubs/sagemaker/sagemaker/collection.pyi b/stubs/sagemaker/sagemaker/collection.pyi index 447c831f8f74..b1704d849589 100644 --- a/stubs/sagemaker/sagemaker/collection.pyi +++ b/stubs/sagemaker/sagemaker/collection.pyi @@ -4,9 +4,9 @@ from typing import List class Collection: sagemaker_session: Incomplete def __init__(self, sagemaker_session) -> None: ... - def create(self, collection_name: str, parent_collection_name: str = None): ... - def delete(self, collections: List[str]): ... - def add_model_groups(self, collection_name: str, model_groups: List[str]): ... - def remove_model_groups(self, collection_name: str, model_groups: List[str]): ... + def create(self, collection_name: str, parent_collection_name: str | None = None): ... + def delete(self, collections: list[str]): ... + def add_model_groups(self, collection_name: str, model_groups: list[str]): ... + def remove_model_groups(self, collection_name: str, model_groups: list[str]): ... def move_model_group(self, source_collection_name: str, model_group: str, destination_collection_name: str): ... - def list_collection(self, collection_name: str = None): ... + def list_collection(self, collection_name: str | None = None): ... diff --git a/stubs/sagemaker/sagemaker/config/config.pyi b/stubs/sagemaker/sagemaker/config/config.pyi index 44b25ad4d987..4d73876d6281 100644 --- a/stubs/sagemaker/sagemaker/config/config.pyi +++ b/stubs/sagemaker/sagemaker/config/config.pyi @@ -6,5 +6,5 @@ ENV_VARIABLE_ADMIN_CONFIG_OVERRIDE: str ENV_VARIABLE_USER_CONFIG_OVERRIDE: str S3_PREFIX: str -def load_sagemaker_config(additional_config_paths: List[str] = None, s3_resource: Incomplete | None = None) -> dict: ... -def validate_sagemaker_config(sagemaker_config: dict = None): ... +def load_sagemaker_config(additional_config_paths: list[str] | None = None, s3_resource: Incomplete | None = None) -> dict: ... +def validate_sagemaker_config(sagemaker_config: dict | None = None): ... diff --git a/stubs/sagemaker/sagemaker/content_types.pyi b/stubs/sagemaker/sagemaker/content_types.pyi index b8d3fac6a6f3..2852ed5e7419 100644 --- a/stubs/sagemaker/sagemaker/content_types.pyi +++ b/stubs/sagemaker/sagemaker/content_types.pyi @@ -1,16 +1,16 @@ from typing import List, Optional def retrieve_options( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, -) -> List[str]: ... +) -> list[str]: ... def retrieve_default( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, ) -> str: ... diff --git a/stubs/sagemaker/sagemaker/debugger/debugger.pyi b/stubs/sagemaker/sagemaker/debugger/debugger.pyi index 5653efe75e27..15fe3c9785b0 100644 --- a/stubs/sagemaker/sagemaker/debugger/debugger.pyi +++ b/stubs/sagemaker/sagemaker/debugger/debugger.pyi @@ -60,13 +60,13 @@ class Rule(RuleBase): image_uri: str | PipelineVariable, instance_type: str | PipelineVariable, volume_size_in_gb: int | PipelineVariable, - source: Optional[str] = None, - rule_to_invoke: Optional[str | PipelineVariable] = None, - container_local_output_path: Optional[str | PipelineVariable] = None, - s3_output_path: Optional[str | PipelineVariable] = None, - other_trials_s3_input_paths: Optional[List[str | PipelineVariable]] = None, - rule_parameters: Optional[Dict[str, str | PipelineVariable]] = None, - collections_to_save: Optional[List["CollectionConfig"]] = None, + source: str | None = None, + rule_to_invoke: str | PipelineVariable | None = None, + container_local_output_path: str | PipelineVariable | None = None, + s3_output_path: str | PipelineVariable | None = None, + other_trials_s3_input_paths: list[str | PipelineVariable] | None = None, + rule_parameters: dict[str, str | PipelineVariable] | None = None, + collections_to_save: list["CollectionConfig"] | None = None, actions: Incomplete | None = None, ): ... def prepare_actions(self, training_job_name) -> None: ... @@ -103,23 +103,23 @@ class DebuggerHookConfig: collection_configs: Incomplete def __init__( self, - s3_output_path: Optional[str | PipelineVariable] = None, - container_local_output_path: Optional[str | PipelineVariable] = None, - hook_parameters: Optional[Dict[str, str | PipelineVariable]] = None, - collection_configs: Optional[List["CollectionConfig"]] = None, + s3_output_path: str | PipelineVariable | None = None, + container_local_output_path: str | PipelineVariable | None = None, + hook_parameters: dict[str, str | PipelineVariable] | None = None, + collection_configs: list["CollectionConfig"] | None = None, ) -> None: ... class TensorBoardOutputConfig: s3_output_path: Incomplete container_local_output_path: Incomplete def __init__( - self, s3_output_path: str | PipelineVariable, container_local_output_path: Optional[str | PipelineVariable] = None + self, s3_output_path: str | PipelineVariable, container_local_output_path: str | PipelineVariable | None = None ) -> None: ... class CollectionConfig: name: Incomplete parameters: Incomplete - def __init__(self, name: str | PipelineVariable, parameters: Optional[Dict[str, str | PipelineVariable]] = None) -> None: ... + def __init__(self, name: str | PipelineVariable, parameters: dict[str, str | PipelineVariable] | None = None) -> None: ... def __eq__(self, other): ... def __ne__(self, other): ... def __hash__(self): ... diff --git a/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi b/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi index b862e13c7519..9c56853c61c5 100644 --- a/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi +++ b/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi @@ -13,8 +13,8 @@ class ProfilerConfig: disable_profiler: Incomplete def __init__( self, - s3_output_path: Optional[str | PipelineVariable] = None, - system_monitor_interval_millis: Optional[int | PipelineVariable] = None, - framework_profile_params: Optional[FrameworkProfile] = None, - disable_profiler: Optional[str | PipelineVariable] = False, + s3_output_path: str | PipelineVariable | None = None, + system_monitor_interval_millis: int | PipelineVariable | None = None, + framework_profile_params: FrameworkProfile | None = None, + disable_profiler: str | PipelineVariable | None = False, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/deserializers.pyi b/stubs/sagemaker/sagemaker/deserializers.pyi index 58d61afe21d0..cf345a84a5c3 100644 --- a/stubs/sagemaker/sagemaker/deserializers.pyi +++ b/stubs/sagemaker/sagemaker/deserializers.pyi @@ -15,16 +15,16 @@ from sagemaker.base_deserializers import ( ) def retrieve_options( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, -) -> List[BaseDeserializer]: ... +) -> list[BaseDeserializer]: ... def retrieve_default( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, ) -> BaseDeserializer: ... diff --git a/stubs/sagemaker/sagemaker/djl_inference/model.pyi b/stubs/sagemaker/sagemaker/djl_inference/model.pyi index e7d1a8efbc57..a809237fa7fe 100644 --- a/stubs/sagemaker/sagemaker/djl_inference/model.pyi +++ b/stubs/sagemaker/sagemaker/djl_inference/model.pyi @@ -21,7 +21,7 @@ class DJLPredictor(Predictor): def __init__( self, endpoint_name: str, - sagemaker_session: Session = None, + sagemaker_session: Session | None = None, serializer: BaseSerializer = ..., deserializer: BaseDeserializer = ..., ) -> None: ... @@ -45,18 +45,18 @@ class DJLModel(FrameworkModel): self, model_id: str, role: str, - djl_version: Optional[str] = None, - task: Optional[str] = None, + djl_version: str | None = None, + task: str | None = None, dtype: str = "fp32", - number_of_partitions: Optional[int] = None, - min_workers: Optional[int] = None, - max_workers: Optional[int] = None, - job_queue_size: Optional[int] = None, + number_of_partitions: int | None = None, + min_workers: int | None = None, + max_workers: int | None = None, + job_queue_size: int | None = None, parallel_loading: bool = False, - model_loading_timeout: Optional[int] = None, - prediction_timeout: Optional[int] = None, - entry_point: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, + model_loading_timeout: int | None = None, + prediction_timeout: int | None = None, + entry_point: str | None = None, + image_uri: str | PipelineVariable | None = None, predictor_cls: callable = ..., **kwargs, ) -> None: ... @@ -68,12 +68,12 @@ class DJLModel(FrameworkModel): def partition( self, instance_type: str, - s3_output_uri: str = None, - job_name: Optional[str] = None, - volume_kms_key: Optional[str] = None, - output_kms_key: Optional[str] = None, + s3_output_uri: str | None = None, + job_name: str | None = None, + volume_kms_key: str | None = None, + output_kms_key: str | None = None, use_spot_instances: bool = False, - max_wait: int = None, + max_wait: int | None = None, enable_network_isolation: bool = False, ): ... def deploy( @@ -97,7 +97,7 @@ class DJLModel(FrameworkModel): accelerator_type: Incomplete | None = None, serverless_inference_config: Incomplete | None = None, ): ... - def generate_serving_properties(self, serving_properties: Incomplete | None = None) -> Dict[str, str]: ... + def generate_serving_properties(self, serving_properties: Incomplete | None = None) -> dict[str, str]: ... def serving_image_uri(self, region_name): ... class DeepSpeedModel(DJLModel): @@ -113,24 +113,24 @@ class DeepSpeedModel(DJLModel): self, model_id: str, role: str, - tensor_parallel_degree: Optional[int] = None, - max_tokens: Optional[int] = None, + tensor_parallel_degree: int | None = None, + max_tokens: int | None = None, low_cpu_mem_usage: bool = False, enable_cuda_graph: bool = False, triangular_masking: bool = True, return_tuple: bool = True, **kwargs, ) -> None: ... - def generate_serving_properties(self, serving_properties: Incomplete | None = None) -> Dict[str, Any]: ... + def generate_serving_properties(self, serving_properties: Incomplete | None = None) -> dict[str, Any]: ... def partition( self, instance_type: str, - s3_output_uri: str = None, - job_name: Optional[str] = None, - volume_kms_key: Optional[str] = None, - output_kms_key: Optional[str] = None, + s3_output_uri: str | None = None, + job_name: str | None = None, + volume_kms_key: str | None = None, + output_kms_key: str | None = None, use_spot_instances: bool = False, - max_wait: int = None, + max_wait: int | None = None, enable_network_isolation: bool = False, ): ... @@ -143,26 +143,26 @@ class HuggingFaceAccelerateModel(DJLModel): self, model_id: str, role: str, - number_of_partitions: Optional[int] = None, - device_id: Optional[int] = None, - device_map: Optional[str | Dict[str | str]] = None, + number_of_partitions: int | None = None, + device_id: int | None = None, + device_map: str | dict[str | str] | None = None, load_in_8bit: bool = False, low_cpu_mem_usage: bool = False, **kwargs, ) -> None: ... - def generate_serving_properties(self, serving_properties: Incomplete | None = None) -> Dict[str, str]: ... + def generate_serving_properties(self, serving_properties: Incomplete | None = None) -> dict[str, str]: ... def partition( self, instance_type: str, - s3_output_uri: str = None, - job_name: Optional[str] = None, - volume_kms_key: Optional[str] = None, - output_kms_key: Optional[str] = None, + s3_output_uri: str | None = None, + job_name: str | None = None, + volume_kms_key: str | None = None, + output_kms_key: str | None = None, use_spot_instances: bool = False, - max_wait: int = None, + max_wait: int | None = None, enable_network_isolation: bool = False, ): ... class FasterTransformerModel(DJLModel): number_of_partitions: Incomplete - def __init__(self, model_id: str, role: str, tensor_parallel_degree: Optional[int] = None, **kwargs) -> None: ... + def __init__(self, model_id: str, role: str, tensor_parallel_degree: int | None = None, **kwargs) -> None: ... diff --git a/stubs/sagemaker/sagemaker/drift_check_baselines.pyi b/stubs/sagemaker/sagemaker/drift_check_baselines.pyi index f38f950ec08c..784fd440a9c9 100644 --- a/stubs/sagemaker/sagemaker/drift_check_baselines.pyi +++ b/stubs/sagemaker/sagemaker/drift_check_baselines.pyi @@ -15,13 +15,13 @@ class DriftCheckBaselines: explainability_config_file: Incomplete def __init__( self, - 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, + model_statistics: MetricsSource | None = None, + model_constraints: MetricsSource | None = None, + model_data_statistics: MetricsSource | None = None, + model_data_constraints: MetricsSource | None = None, + bias_config_file: FileSource | None = None, + bias_pre_training_constraints: MetricsSource | None = None, + bias_post_training_constraints: MetricsSource | None = None, + explainability_constraints: MetricsSource | None = None, + explainability_config_file: FileSource | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/environment_variables.pyi b/stubs/sagemaker/sagemaker/environment_variables.pyi index 83c791d7d165..208e30e9e19e 100644 --- a/stubs/sagemaker/sagemaker/environment_variables.pyi +++ b/stubs/sagemaker/sagemaker/environment_variables.pyi @@ -4,10 +4,10 @@ from typing import Dict, Optional logger: Incomplete def retrieve_default( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, include_aws_sdk_env_vars: bool = True, -) -> Dict[str, str]: ... +) -> dict[str, str]: ... diff --git a/stubs/sagemaker/sagemaker/estimator.pyi b/stubs/sagemaker/sagemaker/estimator.pyi index 39b620aeb575..60ba0099b77b 100644 --- a/stubs/sagemaker/sagemaker/estimator.pyi +++ b/stubs/sagemaker/sagemaker/estimator.pyi @@ -87,50 +87,50 @@ class EstimatorBase(metaclass=abc.ABCMeta): disable_output_compression: Incomplete def __init__( self, - role: str = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, - keep_alive_period_in_seconds: Optional[int | PipelineVariable] = None, + role: str | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + keep_alive_period_in_seconds: int | PipelineVariable | None = None, volume_size: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, + volume_kms_key: str | PipelineVariable | None = None, max_run: int | PipelineVariable = 86400, input_mode: str | PipelineVariable = "File", - output_path: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - subnets: Optional[List[str | PipelineVariable]] = None, - security_group_ids: Optional[List[str | PipelineVariable]] = None, - model_uri: Optional[str] = None, + output_path: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + subnets: list[str | PipelineVariable] | None = None, + security_group_ids: list[str | PipelineVariable] | None = None, + model_uri: str | None = None, model_channel_name: str | PipelineVariable = "model", - metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, - encrypt_inter_container_traffic: bool | PipelineVariable = None, + metric_definitions: list[dict[str, str | PipelineVariable]] | None = None, + encrypt_inter_container_traffic: bool | PipelineVariable | None = None, use_spot_instances: bool | PipelineVariable = False, - max_wait: Optional[int | PipelineVariable] = None, - checkpoint_s3_uri: Optional[str | PipelineVariable] = None, - checkpoint_local_path: Optional[str | PipelineVariable] = None, - rules: Optional[List[RuleBase]] = None, - debugger_hook_config: Optional[bool | DebuggerHookConfig] = None, - tensorboard_output_config: Optional[TensorBoardOutputConfig] = None, - enable_sagemaker_metrics: Optional[bool | PipelineVariable] = None, - enable_network_isolation: bool | PipelineVariable = None, - profiler_config: Optional[ProfilerConfig] = None, - disable_profiler: bool = None, - environment: Optional[Dict[str, str | PipelineVariable]] = None, - max_retry_attempts: Optional[int | PipelineVariable] = None, - source_dir: Optional[str | PipelineVariable] = None, - git_config: Optional[Dict[str, str]] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + max_wait: int | PipelineVariable | None = None, + checkpoint_s3_uri: str | PipelineVariable | None = None, + checkpoint_local_path: str | PipelineVariable | None = None, + rules: list[RuleBase] | None = None, + debugger_hook_config: bool | DebuggerHookConfig | None = None, + tensorboard_output_config: TensorBoardOutputConfig | None = None, + enable_sagemaker_metrics: bool | PipelineVariable | None = None, + enable_network_isolation: bool | PipelineVariable | None = None, + profiler_config: ProfilerConfig | None = None, + disable_profiler: bool | None = None, + environment: dict[str, str | PipelineVariable] | None = None, + max_retry_attempts: int | PipelineVariable | None = None, + source_dir: str | PipelineVariable | None = None, + git_config: dict[str, str] | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, container_log_level: int | PipelineVariable = 20, - code_location: Optional[str] = None, - entry_point: Optional[str | PipelineVariable] = None, - dependencies: Optional[List[str]] = None, - instance_groups: Optional[List[InstanceGroup]] = None, - training_repository_access_mode: Optional[str | PipelineVariable] = None, - training_repository_credentials_provider_arn: Optional[str | PipelineVariable] = None, - container_entry_point: Optional[List[str]] = None, - container_arguments: Optional[List[str]] = None, + code_location: str | None = None, + entry_point: str | PipelineVariable | None = None, + dependencies: list[str] | None = None, + instance_groups: list[InstanceGroup] | None = None, + training_repository_access_mode: str | PipelineVariable | None = None, + training_repository_credentials_provider_arn: str | PipelineVariable | None = None, + container_entry_point: list[str] | None = None, + container_arguments: list[str] | None = None, disable_output_compression: bool = False, **kwargs, ) -> None: ... @@ -145,11 +145,11 @@ class EstimatorBase(metaclass=abc.ABCMeta): def latest_job_profiler_artifacts_path(self): ... def fit( self, - inputs: Optional[str | Dict | TrainingInput | FileSystemInput] = None, + inputs: str | dict | TrainingInput | FileSystemInput | None = None, wait: bool = True, logs: str = "All", - job_name: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, + job_name: str | None = None, + experiment_config: dict[str, str] | None = None, ): ... def compile_model( self, @@ -278,50 +278,50 @@ class Estimator(EstimatorBase): def __init__( self, image_uri: str | PipelineVariable, - role: str = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, - keep_alive_period_in_seconds: Optional[int | PipelineVariable] = None, + role: str | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + keep_alive_period_in_seconds: int | PipelineVariable | None = None, volume_size: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, + volume_kms_key: str | PipelineVariable | None = None, max_run: int | PipelineVariable = 86400, input_mode: str | PipelineVariable = "File", - output_path: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - subnets: Optional[List[str | PipelineVariable]] = None, - security_group_ids: Optional[List[str | PipelineVariable]] = None, - model_uri: Optional[str] = None, + output_path: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + subnets: list[str | PipelineVariable] | None = None, + security_group_ids: list[str | PipelineVariable] | None = None, + model_uri: str | None = None, model_channel_name: str | PipelineVariable = "model", - metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, - encrypt_inter_container_traffic: bool | PipelineVariable = None, + metric_definitions: list[dict[str, str | PipelineVariable]] | None = None, + encrypt_inter_container_traffic: bool | PipelineVariable | None = None, use_spot_instances: bool | PipelineVariable = False, - max_wait: Optional[int | PipelineVariable] = None, - checkpoint_s3_uri: Optional[str | PipelineVariable] = None, - checkpoint_local_path: Optional[str | PipelineVariable] = None, - enable_network_isolation: bool | PipelineVariable = None, - rules: Optional[List[RuleBase]] = None, - debugger_hook_config: Optional[DebuggerHookConfig | bool] = None, - tensorboard_output_config: Optional[TensorBoardOutputConfig] = None, - enable_sagemaker_metrics: Optional[bool | PipelineVariable] = None, - profiler_config: Optional[ProfilerConfig] = None, + max_wait: int | PipelineVariable | None = None, + checkpoint_s3_uri: str | PipelineVariable | None = None, + checkpoint_local_path: str | PipelineVariable | None = None, + enable_network_isolation: bool | PipelineVariable | None = None, + rules: list[RuleBase] | None = None, + debugger_hook_config: DebuggerHookConfig | bool | None = None, + tensorboard_output_config: TensorBoardOutputConfig | None = None, + enable_sagemaker_metrics: bool | PipelineVariable | None = None, + profiler_config: ProfilerConfig | None = None, disable_profiler: bool = False, - environment: Optional[Dict[str, str | PipelineVariable]] = None, - max_retry_attempts: Optional[int | PipelineVariable] = None, - source_dir: Optional[str | PipelineVariable] = None, - git_config: Optional[Dict[str, str]] = None, + environment: dict[str, str | PipelineVariable] | None = None, + max_retry_attempts: int | PipelineVariable | None = None, + source_dir: str | PipelineVariable | None = None, + git_config: dict[str, str] | None = None, container_log_level: int | PipelineVariable = 20, - code_location: Optional[str] = None, - entry_point: Optional[str | PipelineVariable] = None, - dependencies: Optional[List[str]] = None, - instance_groups: Optional[List[InstanceGroup]] = None, - training_repository_access_mode: Optional[str | PipelineVariable] = None, - training_repository_credentials_provider_arn: Optional[str | PipelineVariable] = None, - container_entry_point: Optional[List[str]] = None, - container_arguments: Optional[List[str]] = None, + code_location: str | None = None, + entry_point: str | PipelineVariable | None = None, + dependencies: list[str] | None = None, + instance_groups: list[InstanceGroup] | None = None, + training_repository_access_mode: str | PipelineVariable | None = None, + training_repository_credentials_provider_arn: str | PipelineVariable | None = None, + container_entry_point: list[str] | None = None, + container_arguments: list[str] | None = None, disable_output_compression: bool = False, **kwargs, ) -> None: ... @@ -352,17 +352,17 @@ class Framework(EstimatorBase, metaclass=abc.ABCMeta): def __init__( self, entry_point: str | PipelineVariable, - source_dir: Optional[str | PipelineVariable] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + source_dir: str | PipelineVariable | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, container_log_level: int | PipelineVariable = 20, - code_location: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, - dependencies: Optional[List[str]] = None, - enable_network_isolation: bool | PipelineVariable = None, - git_config: Optional[Dict[str, str]] = None, - checkpoint_s3_uri: Optional[str | PipelineVariable] = None, - checkpoint_local_path: Optional[str | PipelineVariable] = None, - enable_sagemaker_metrics: Optional[bool | PipelineVariable] = None, + code_location: str | None = None, + image_uri: str | PipelineVariable | None = None, + dependencies: list[str] | None = None, + enable_network_isolation: bool | PipelineVariable | None = None, + git_config: dict[str, str] | None = None, + checkpoint_s3_uri: str | PipelineVariable | None = None, + checkpoint_local_path: str | PipelineVariable | None = None, + enable_sagemaker_metrics: bool | PipelineVariable | None = None, **kwargs, ) -> None: ... def set_hyperparameters(self, **kwargs) -> None: ... diff --git a/stubs/sagemaker/sagemaker/experiments/_utils.pyi b/stubs/sagemaker/sagemaker/experiments/_utils.pyi index e7c0e80e2917..138a27c84289 100644 --- a/stubs/sagemaker/sagemaker/experiments/_utils.pyi +++ b/stubs/sagemaker/sagemaker/experiments/_utils.pyi @@ -9,5 +9,5 @@ def verify_length_of_true_and_predicted(true_labels, predicted_attrs, predicted_ def validate_invoked_inside_run_context(func): ... def is_already_exist_error(error): ... def get_tc_and_exp_config_from_job_env(environment: _RunEnvironment, sagemaker_session: Session) -> dict: ... -def verify_load_input_names(run_name: Optional[str] = None, experiment_name: Optional[str] = None): ... +def verify_load_input_names(run_name: str | None = None, experiment_name: str | None = None): ... def is_run_trial_component(trial_component_name: str, sagemaker_session: Session) -> bool: ... diff --git a/stubs/sagemaker/sagemaker/experiments/run.pyi b/stubs/sagemaker/sagemaker/experiments/run.pyi index a8c6a4727cc0..166d2118c78d 100644 --- a/stubs/sagemaker/sagemaker/experiments/run.pyi +++ b/stubs/sagemaker/sagemaker/experiments/run.pyi @@ -34,46 +34,46 @@ class Run: def __init__( self, experiment_name: str, - run_name: Optional[str] = None, - experiment_display_name: Optional[str] = None, - run_display_name: Optional[str] = None, - tags: Optional[List[Dict[str, str]]] = None, + run_name: str | None = None, + experiment_display_name: str | None = None, + run_display_name: str | None = None, + tags: list[dict[str, str]] | None = None, sagemaker_session: Optional["Session"] = None, ) -> None: ... @property def experiment_config(self) -> dict: ... def log_parameter(self, name: str, value: str | int | float): ... - def log_parameters(self, parameters: Dict[str, str | int | float]): ... - def log_metric(self, name: str, value: float, timestamp: Optional[datetime.datetime] = None, step: Optional[int] = None): ... + def log_parameters(self, parameters: dict[str, str | int | float]): ... + def log_metric(self, name: str, value: float, timestamp: datetime.datetime | None = None, step: int | None = None): ... def log_precision_recall( self, y_true: list | array, predicted_probabilities: list | array, - positive_label: Optional[str | int] = None, - title: Optional[str] = None, + positive_label: str | int | None = None, + title: str | None = None, is_output: bool = True, - no_skill: Optional[int] = None, + no_skill: int | None = None, ): ... - def log_roc_curve(self, y_true: list | array, y_score: list | array, title: Optional[str] = None, is_output: bool = True): ... + def log_roc_curve(self, y_true: list | array, y_score: list | array, title: str | None = None, is_output: bool = True): ... def log_confusion_matrix( - self, y_true: list | array, y_pred: list | array, title: Optional[str] = None, is_output: bool = True + self, y_true: list | array, y_pred: list | array, title: str | None = None, is_output: bool = True ): ... - def log_artifact(self, name: str, value: str, media_type: Optional[str] = None, is_output: bool = True): ... - def log_file(self, file_path: str, name: Optional[str] = None, media_type: Optional[str] = None, is_output: bool = True): ... + def log_artifact(self, name: str, value: str, media_type: str | None = None, is_output: bool = True): ... + def log_file(self, file_path: str, name: str | None = None, media_type: str | None = None, is_output: bool = True): ... def close(self) -> None: ... def __enter__(self): ... def __exit__(self, exc_type, exc_value, exc_traceback) -> None: ... def load_run( - run_name: Optional[str] = None, experiment_name: Optional[str] = None, sagemaker_session: Optional["Session"] = None + run_name: str | None = None, experiment_name: str | None = None, sagemaker_session: Optional["Session"] = None ) -> Run: ... def list_runs( experiment_name: str, - created_before: Optional[datetime.datetime] = None, - created_after: Optional[datetime.datetime] = None, + created_before: datetime.datetime | None = None, + created_after: datetime.datetime | None = None, sagemaker_session: Optional["Session"] = None, - max_results: Optional[int] = None, - next_token: Optional[str] = None, + max_results: int | None = None, + next_token: str | None = None, sort_by: SortByType = ..., sort_order: SortOrderType = ..., ) -> list: ... diff --git a/stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi b/stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi index 5e43de15c1d2..fb034b672b25 100644 --- a/stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi +++ b/stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi @@ -11,7 +11,7 @@ class ClarifyShapBaselineConfig: shap_baseline: Incomplete shap_baseline_uri: Incomplete def __init__( - self, mime_type: Optional[str] = "text/csv", shap_baseline: Optional[str] = None, shap_baseline_uri: Optional[str] = None + self, mime_type: str | None = "text/csv", shap_baseline: str | None = None, shap_baseline_uri: str | None = None ) -> None: ... class ClarifyShapConfig: @@ -23,10 +23,10 @@ class ClarifyShapConfig: def __init__( self, shap_baseline_config: ClarifyShapBaselineConfig, - number_of_samples: Optional[int] = None, - seed: Optional[int] = None, - use_logit: Optional[bool] = False, - text_config: Optional[ClarifyTextConfig] = None, + number_of_samples: int | None = None, + seed: int | None = None, + use_logit: bool | None = False, + text_config: ClarifyTextConfig | None = None, ) -> None: ... class ClarifyInferenceConfig: @@ -43,17 +43,17 @@ class ClarifyInferenceConfig: content_template: Incomplete def __init__( self, - feature_headers: Optional[List[str]] = None, - feature_types: Optional[List[str]] = None, - features_attribute: Optional[str] = None, - probability_index: Optional[int] = None, - probability_attribute: Optional[str] = None, - label_index: Optional[int] = None, - label_attribute: Optional[str] = None, - label_headers: Optional[List[str]] = None, - max_payload_in_mb: Optional[int] = 6, - max_record_count: Optional[int] = None, - content_template: Optional[str] = None, + feature_headers: list[str] | None = None, + feature_types: list[str] | None = None, + features_attribute: str | None = None, + probability_index: int | None = None, + probability_attribute: str | None = None, + label_index: int | None = None, + label_attribute: str | None = None, + label_headers: list[str] | None = None, + max_payload_in_mb: int | None = 6, + max_record_count: int | None = None, + content_template: str | None = None, ) -> None: ... class ClarifyExplainerConfig: @@ -63,6 +63,6 @@ class ClarifyExplainerConfig: def __init__( self, shap_config: ClarifyShapConfig, - enable_explanations: Optional[str] = None, - inference_config: Optional[ClarifyInferenceConfig] = None, + enable_explanations: str | None = None, + inference_config: ClarifyInferenceConfig | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/explainer/explainer_config.pyi b/stubs/sagemaker/sagemaker/explainer/explainer_config.pyi index 16930ed5fbd3..fa0b5557d107 100644 --- a/stubs/sagemaker/sagemaker/explainer/explainer_config.pyi +++ b/stubs/sagemaker/sagemaker/explainer/explainer_config.pyi @@ -5,4 +5,4 @@ from sagemaker.explainer.clarify_explainer_config import ClarifyExplainerConfig class ExplainerConfig: clarify_explainer_config: Incomplete - def __init__(self, clarify_explainer_config: Optional[ClarifyExplainerConfig] = None) -> None: ... + def __init__(self, clarify_explainer_config: ClarifyExplainerConfig | None = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi b/stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi index 0fb4080f3dba..386f2babf3b8 100644 --- a/stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi @@ -40,9 +40,9 @@ class JoinComparatorEnum(Enum): def __ge__(self, other): ... class FeatureGroupToBeMerged: - features: List[str] - included_feature_names: List[str] - projected_feature_names: List[str] + features: list[str] + included_feature_names: list[str] + projected_feature_names: list[str] catalog: str database: str table_name: str @@ -76,9 +76,9 @@ class FeatureGroupToBeMerged: def construct_feature_group_to_be_merged( target_feature_group: FeatureGroup, - included_feature_names: List[str], - target_feature_name_in_base: str = None, - feature_name_in_target: str = None, + included_feature_names: list[str], + target_feature_name_in_base: str | None = None, + feature_name_in_target: str | None = None, join_comparator: JoinComparatorEnum = ..., join_type: JoinTypeEnum = ..., ) -> FeatureGroupToBeMerged: ... @@ -87,9 +87,9 @@ class DatasetBuilder: def with_feature_group( self, feature_group: FeatureGroup, - target_feature_name_in_base: str = None, - included_feature_names: List[str] = None, - feature_name_in_target: str = None, + target_feature_name_in_base: str | None = None, + included_feature_names: list[str] | None = None, + feature_name_in_target: str | None = None, join_comparator: JoinComparatorEnum = ..., join_type: JoinTypeEnum = ..., ): ... @@ -99,9 +99,9 @@ class DatasetBuilder: def with_number_of_recent_records_by_record_identifier(self, number_of_recent_records: int): ... def with_number_of_records_from_query_results(self, number_of_records: int): ... def as_of(self, timestamp: datetime.datetime): ... - def with_event_time_range(self, starting_timestamp: datetime.datetime = None, ending_timestamp: datetime.datetime = None): ... - def to_csv_file(self) -> Tuple[str, str]: ... - def to_dataframe(self) -> Tuple[pd.DataFrame, str]: ... + def with_event_time_range(self, starting_timestamp: datetime.datetime | None = None, ending_timestamp: datetime.datetime | None = None): ... + def to_csv_file(self) -> tuple[str, str]: ... + def to_dataframe(self) -> tuple[pd.DataFrame, str]: ... def __init__( self, sagemaker_session, diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi index d01e343caa13..1e5c8ec2ce06 100644 --- a/stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi @@ -11,7 +11,7 @@ class FeatureTypeEnum(Enum): class FeatureDefinition(Config): feature_name: str feature_type: FeatureTypeEnum - def to_dict(self) -> Dict[str, Any]: ... + def to_dict(self) -> dict[str, Any]: ... def __init__(self, feature_name, feature_type) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi index d431a357bf5d..792918b8e623 100644 --- a/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi @@ -1,6 +1,7 @@ from _typeshed import Incomplete from multiprocessing.pool import AsyncResult as AsyncResult -from typing import Any, Dict, List, Sequence +from typing import Any, Dict, List +from collections.abc import Sequence from botocore.config import Config as Config from pandas import DataFrame as DataFrame @@ -15,9 +16,9 @@ class AthenaQuery: database: str table_name: str sagemaker_session: Session - def run(self, query_string: str, output_location: str, kms_key: str = None, workgroup: str = None) -> str: ... + def run(self, query_string: str, output_location: str, kms_key: str | None = None, workgroup: str | None = None) -> str: ... def wait(self) -> None: ... - def get_query_execution(self) -> Dict[str, Any]: ... + def get_query_execution(self) -> dict[str, Any]: ... def as_dataframe(self) -> DataFrame: ... def __init__( self, catalog, database, table_name, sagemaker_session, current_query_execution_id, result_bucket, result_file_prefix @@ -35,7 +36,7 @@ class IngestionManagerPandas: max_processes: int profile_name: str @property - def failed_rows(self) -> List[int]: ... + def failed_rows(self) -> list[int]: ... def wait(self, timeout: Incomplete | None = None) -> None: ... def run(self, data_frame: DataFrame, wait: bool = True, timeout: Incomplete | None = None): ... def __init__( @@ -64,39 +65,39 @@ class FeatureGroup: name: str sagemaker_session: Session feature_definitions: Sequence[FeatureDefinition] - DTYPE_TO_FEATURE_DEFINITION_CLS_MAP: Dict[str, FeatureTypeEnum] + DTYPE_TO_FEATURE_DEFINITION_CLS_MAP: dict[str, FeatureTypeEnum] def create( self, s3_uri: str | bool, record_identifier_name: str, event_time_feature_name: str, - role_arn: str = None, - online_store_kms_key_id: str = None, + role_arn: str | None = None, + online_store_kms_key_id: str | None = None, enable_online_store: bool = False, - offline_store_kms_key_id: str = None, + offline_store_kms_key_id: str | None = None, disable_glue_table_creation: bool = False, - data_catalog_config: DataCatalogConfig = None, - description: str = None, - tags: List[Dict[str, str]] = None, - table_format: TableFormatEnum = None, - ) -> Dict[str, Any]: ... + data_catalog_config: DataCatalogConfig | None = None, + description: str | None = None, + tags: list[dict[str, str]] | None = None, + table_format: TableFormatEnum | None = None, + ) -> dict[str, Any]: ... def delete(self) -> None: ... - def describe(self, next_token: str = None) -> Dict[str, Any]: ... - def update(self, feature_additions: Sequence[FeatureDefinition]) -> Dict[str, Any]: ... + def describe(self, next_token: str | None = None) -> dict[str, Any]: ... + def update(self, feature_additions: Sequence[FeatureDefinition]) -> dict[str, Any]: ... def update_feature_metadata( self, feature_name: str, - description: str = None, - parameter_additions: Sequence[FeatureParameter] = None, - parameter_removals: Sequence[str] = None, - ) -> Dict[str, Any]: ... - def describe_feature_metadata(self, feature_name: str) -> Dict[str, Any]: ... - def list_tags(self) -> Sequence[Dict[str, str]]: ... - def list_parameters_for_feature_metadata(self, feature_name: str) -> Sequence[Dict[str, str]]: ... + description: str | None = None, + parameter_additions: Sequence[FeatureParameter] | None = None, + parameter_removals: Sequence[str] | None = None, + ) -> dict[str, Any]: ... + def describe_feature_metadata(self, feature_name: str) -> dict[str, Any]: ... + def list_tags(self) -> Sequence[dict[str, str]]: ... + def list_parameters_for_feature_metadata(self, feature_name: str) -> Sequence[dict[str, str]]: ... def load_feature_definitions(self, data_frame: DataFrame) -> Sequence[FeatureDefinition]: ... def get_record( - self, record_identifier_value_as_string: str, feature_names: Sequence[str] = None - ) -> Sequence[Dict[str, str]]: ... + self, record_identifier_value_as_string: str, feature_names: Sequence[str] | None = None + ) -> Sequence[dict[str, str]]: ... def put_record(self, record: Sequence[FeatureValue]): ... def delete_record(self, record_identifier_value_as_string: str, event_time: str, deletion_mode: DeletionModeEnum = ...): ... def ingest( @@ -105,11 +106,11 @@ class FeatureGroup: max_workers: int = 1, max_processes: int = 1, wait: bool = True, - timeout: int | float = None, - profile_name: str = None, + timeout: int | float | None = None, + profile_name: str | None = None, ) -> IngestionManagerPandas: ... def athena_query(self) -> AthenaQuery: ... - def as_hive_ddl(self, database: str = "sagemaker_featurestore", table_name: str = None) -> str: ... + def as_hive_ddl(self, database: str = "sagemaker_featurestore", table_name: str | None = None) -> str: ... def __init__(self, name, sagemaker_session, feature_definitions) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi index 91409de5e356..0b9acd9cd8d1 100644 --- a/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi @@ -1,5 +1,6 @@ import datetime -from typing import Any, Dict, Sequence +from typing import Any, Dict +from collections.abc import Sequence import pandas as pd from sagemaker import Session @@ -13,34 +14,34 @@ class FeatureStore: self, base: FeatureGroup | pd.DataFrame, output_path: str, - record_identifier_feature_name: str = None, - event_time_identifier_feature_name: str = None, - included_feature_names: Sequence[str] = None, - kms_key_id: str = None, + record_identifier_feature_name: str | None = None, + event_time_identifier_feature_name: str | None = None, + included_feature_names: Sequence[str] | None = None, + kms_key_id: str | None = None, ) -> DatasetBuilder: ... def list_feature_groups( self, - name_contains: str = None, - feature_group_status_equals: str = None, - offline_store_status_equals: str = None, - creation_time_after: datetime.datetime = None, - creation_time_before: datetime.datetime = None, - sort_order: str = None, - sort_by: str = None, - max_results: int = None, - next_token: str = None, - ) -> Dict[str, Any]: ... - def batch_get_record(self, identifiers: Sequence[Identifier]) -> Dict[str, Any]: ... + name_contains: str | None = None, + feature_group_status_equals: str | None = None, + offline_store_status_equals: str | None = None, + creation_time_after: datetime.datetime | None = None, + creation_time_before: datetime.datetime | None = None, + sort_order: str | None = None, + sort_by: str | None = None, + max_results: int | None = None, + next_token: str | None = None, + ) -> dict[str, Any]: ... + def batch_get_record(self, identifiers: Sequence[Identifier]) -> dict[str, Any]: ... def search( self, resource: ResourceEnum, - filters: Sequence[Filter] = None, - operator: SearchOperatorEnum = None, - sort_by: str = None, - sort_order: SortOrderEnum = None, - next_token: str = None, - max_results: int = None, - ) -> Dict[str, Any]: ... + filters: Sequence[Filter] | None = None, + operator: SearchOperatorEnum | None = None, + sort_by: str | None = None, + sort_order: SortOrderEnum | None = None, + next_token: str | None = None, + max_results: int | None = None, + ) -> dict[str, Any]: ... def __init__(self, sagemaker_session) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/feature_store/inputs.pyi b/stubs/sagemaker/sagemaker/feature_store/inputs.pyi index 786077ba1569..8e43957f0750 100644 --- a/stubs/sagemaker/sagemaker/feature_store/inputs.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/inputs.pyi @@ -4,13 +4,13 @@ from typing import Any, Dict, List class Config(abc.ABC, metaclass=abc.ABCMeta): @abc.abstractmethod - def to_dict(self) -> Dict[str, Any]: ... + def to_dict(self) -> dict[str, Any]: ... @classmethod - def construct_dict(cls, **kwargs) -> Dict[str, Any]: ... + def construct_dict(cls, **kwargs) -> dict[str, Any]: ... class OnlineStoreSecurityConfig(Config): kms_key_id: str - def to_dict(self) -> Dict[str, Any]: ... + def to_dict(self) -> dict[str, Any]: ... def __init__(self, kms_key_id) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -20,7 +20,7 @@ class OnlineStoreSecurityConfig(Config): class OnlineStoreConfig(Config): enable_online_store: bool online_store_security_config: OnlineStoreSecurityConfig - def to_dict(self) -> Dict[str, Any]: ... + def to_dict(self) -> dict[str, Any]: ... def __init__(self, enable_online_store, online_store_security_config) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -30,7 +30,7 @@ class OnlineStoreConfig(Config): class S3StorageConfig(Config): s3_uri: str kms_key_id: str - def to_dict(self) -> Dict[str, Any]: ... + def to_dict(self) -> dict[str, Any]: ... def __init__(self, s3_uri, kms_key_id) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -41,7 +41,7 @@ class DataCatalogConfig(Config): table_name: str catalog: str database: str - def to_dict(self) -> Dict[str, Any]: ... + def to_dict(self) -> dict[str, Any]: ... def __init__(self, table_name, catalog, database) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -57,7 +57,7 @@ class OfflineStoreConfig(Config): disable_glue_table_creation: bool data_catalog_config: DataCatalogConfig table_format: TableFormatEnum - def to_dict(self) -> Dict[str, Any]: ... + def to_dict(self) -> dict[str, Any]: ... def __init__(self, s3_storage_config, disable_glue_table_creation, data_catalog_config, table_format) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -67,7 +67,7 @@ class OfflineStoreConfig(Config): class FeatureValue(Config): feature_name: str value_as_string: str - def to_dict(self) -> Dict[str, Any]: ... + def to_dict(self) -> dict[str, Any]: ... def __init__(self, feature_name, value_as_string) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -77,7 +77,7 @@ class FeatureValue(Config): class FeatureParameter(Config): key: str value: str - def to_dict(self) -> Dict[str, Any]: ... + def to_dict(self) -> dict[str, Any]: ... def __init__(self, key, value) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -112,7 +112,7 @@ class Filter(Config): name: str value: str operator: FilterOperatorEnum - def to_dict(self) -> Dict[str, Any]: ... + def to_dict(self) -> dict[str, Any]: ... def __init__(self, name, value, operator) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -121,9 +121,9 @@ class Filter(Config): class Identifier(Config): feature_group_name: str - record_identifiers_value_as_string: List[str] - feature_names: List[str] - def to_dict(self) -> Dict[str, Any]: ... + record_identifiers_value_as_string: list[str] + feature_names: list[str] + def to_dict(self) -> dict[str, Any]: ... def __init__(self, feature_group_name, record_identifiers_value_as_string, feature_names) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/fw_utils.pyi b/stubs/sagemaker/sagemaker/fw_utils.pyi index af6fa158e4ba..f787c8a4aa5d 100644 --- a/stubs/sagemaker/sagemaker/fw_utils.pyi +++ b/stubs/sagemaker/sagemaker/fw_utils.pyi @@ -27,9 +27,9 @@ GRAVITON_ALLOWED_FRAMEWORKS: Incomplete def validate_source_dir(script, directory): ... def validate_source_code_input_against_pipeline_variables( - entry_point: Optional[str | PipelineVariable] = None, - source_dir: Optional[str | PipelineVariable] = None, - git_config: Optional[Dict[str, str]] = None, + entry_point: str | PipelineVariable | None = None, + source_dir: str | PipelineVariable | None = None, + git_config: dict[str, str] | None = None, enable_network_isolation: bool | PipelineVariable = False, ): ... def parse_mp_parameters(params): ... @@ -44,7 +44,7 @@ def tar_and_upload_dir( dependencies: Incomplete | None = None, kms_key: Incomplete | None = None, s3_resource: Incomplete | None = None, - settings: Optional[SessionSettings] = None, + settings: SessionSettings | None = None, ) -> UploadedCode: ... def framework_name_from_image(image_uri): ... def framework_version_from_tag(image_tag): ... diff --git a/stubs/sagemaker/sagemaker/huggingface/estimator.pyi b/stubs/sagemaker/sagemaker/huggingface/estimator.pyi index a3202da34b39..d76c839250f7 100644 --- a/stubs/sagemaker/sagemaker/huggingface/estimator.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/estimator.pyi @@ -23,14 +23,14 @@ class HuggingFace(Framework): self, py_version: str, entry_point: str | PipelineVariable, - transformers_version: Optional[str] = None, - tensorflow_version: Optional[str] = None, - pytorch_version: Optional[str] = None, - source_dir: Optional[str | PipelineVariable] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, - image_uri: Optional[str | PipelineVariable] = None, - distribution: Optional[Dict] = None, - compiler_config: Optional[TrainingCompilerConfig] = None, + transformers_version: str | None = None, + tensorflow_version: str | None = None, + pytorch_version: str | None = None, + source_dir: str | PipelineVariable | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, + image_uri: str | PipelineVariable | None = None, + distribution: dict | None = None, + compiler_config: TrainingCompilerConfig | None = None, **kwargs, ) -> None: ... def hyperparameters(self): ... diff --git a/stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi b/stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi index 01ea51b1a4ed..31eb491b3c67 100644 --- a/stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi @@ -3,5 +3,5 @@ from typing import Optional from sagemaker.session import Session def get_huggingface_llm_image_uri( - backend: str, session: Optional[Session] = None, region: Optional[str] = None, version: Optional[str] = None + backend: str, session: Session | None = None, region: str | None = None, version: str | None = None ) -> str: ... diff --git a/stubs/sagemaker/sagemaker/huggingface/model.pyi b/stubs/sagemaker/sagemaker/huggingface/model.pyi index d1aa7f3fc516..6cdf18caf685 100644 --- a/stubs/sagemaker/sagemaker/huggingface/model.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/model.pyi @@ -24,16 +24,16 @@ class HuggingFaceModel(FrameworkModel): model_server_workers: Incomplete def __init__( self, - role: Optional[str] = None, - model_data: Optional[str | PipelineVariable] = None, - entry_point: Optional[str] = None, - transformers_version: Optional[str] = None, - tensorflow_version: Optional[str] = None, - pytorch_version: Optional[str] = None, - py_version: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, + role: str | None = None, + model_data: str | PipelineVariable | None = None, + entry_point: str | None = None, + transformers_version: str | None = None, + tensorflow_version: str | None = None, + pytorch_version: str | None = None, + py_version: str | None = None, + image_uri: str | PipelineVariable | None = None, predictor_cls: callable = ..., - model_server_workers: Optional[int | PipelineVariable] = None, + model_server_workers: int | PipelineVariable | None = None, **kwargs, ) -> None: ... image_uri: Incomplete @@ -60,27 +60,27 @@ class HuggingFaceModel(FrameworkModel): ): ... def register( self, - content_types: List[str | PipelineVariable], - response_types: List[str | PipelineVariable], - inference_instances: Optional[List[str | PipelineVariable]] = None, - transform_instances: Optional[List[str | PipelineVariable]] = None, - model_package_name: Optional[str | PipelineVariable] = None, - model_package_group_name: Optional[str | PipelineVariable] = None, - image_uri: Optional[str | PipelineVariable] = None, - model_metrics: Optional[ModelMetrics] = None, - metadata_properties: Optional[MetadataProperties] = None, + content_types: list[str | PipelineVariable], + response_types: list[str | PipelineVariable], + inference_instances: list[str | PipelineVariable] | None = None, + transform_instances: list[str | PipelineVariable] | None = None, + model_package_name: str | PipelineVariable | None = None, + model_package_group_name: str | PipelineVariable | None = None, + image_uri: str | PipelineVariable | None = None, + model_metrics: ModelMetrics | None = None, + metadata_properties: MetadataProperties | None = None, marketplace_cert: bool = False, - approval_status: Optional[str | PipelineVariable] = None, - description: Optional[str] = None, - drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, - domain: Optional[str | PipelineVariable] = None, - sample_payload_url: Optional[str | PipelineVariable] = None, - task: Optional[str | PipelineVariable] = None, - framework: Optional[str | PipelineVariable] = None, - framework_version: Optional[str | PipelineVariable] = None, - nearest_model_name: Optional[str | PipelineVariable] = None, - data_input_configuration: Optional[str | PipelineVariable] = None, + approval_status: str | PipelineVariable | None = None, + description: str | None = None, + drift_check_baselines: DriftCheckBaselines | None = None, + customer_metadata_properties: dict[str, str | PipelineVariable] | None = None, + domain: str | PipelineVariable | None = None, + sample_payload_url: str | PipelineVariable | None = None, + task: str | PipelineVariable | None = None, + framework: str | PipelineVariable | None = None, + framework_version: str | PipelineVariable | None = None, + nearest_model_name: str | PipelineVariable | None = None, + data_input_configuration: str | PipelineVariable | None = None, ): ... def prepare_container_def( self, diff --git a/stubs/sagemaker/sagemaker/huggingface/processing.pyi b/stubs/sagemaker/sagemaker/huggingface/processing.pyi index 383279314a10..162b9d200351 100644 --- a/stubs/sagemaker/sagemaker/huggingface/processing.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/processing.pyi @@ -13,23 +13,23 @@ class HuggingFaceProcessor(FrameworkProcessor): tensorflow_version: Incomplete def __init__( self, - role: Optional[str | PipelineVariable] = None, - instance_count: int | PipelineVariable = None, - instance_type: str | PipelineVariable = None, - transformers_version: Optional[str] = None, - tensorflow_version: Optional[str] = None, - pytorch_version: Optional[str] = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + transformers_version: str | None = None, + tensorflow_version: str | None = None, + pytorch_version: str | None = None, py_version: str = "py36", - image_uri: Optional[str | PipelineVariable] = None, - command: Optional[List[str]] = None, + image_uri: str | PipelineVariable | None = None, + command: list[str] | None = None, volume_size_in_gb: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - code_location: Optional[str] = None, - max_runtime_in_seconds: Optional[int | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - network_config: Optional[NetworkConfig] = None, + volume_kms_key: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + code_location: str | None = None, + max_runtime_in_seconds: int | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + network_config: NetworkConfig | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/hyperparameters.pyi b/stubs/sagemaker/sagemaker/hyperparameters.pyi index 93a00067c749..a808d5b6974d 100644 --- a/stubs/sagemaker/sagemaker/hyperparameters.pyi +++ b/stubs/sagemaker/sagemaker/hyperparameters.pyi @@ -6,17 +6,17 @@ from sagemaker.jumpstart.enums import HyperparameterValidationMode logger: Incomplete def retrieve_default( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, include_container_hyperparameters: bool = False, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, -) -> Dict[str, str]: ... +) -> dict[str, str]: ... def validate( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, - hyperparameters: Optional[dict] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, + hyperparameters: dict | None = None, validation_mode: HyperparameterValidationMode = ..., ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi b/stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi index 42955314cf60..f942f08a9749 100644 --- a/stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi +++ b/stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi @@ -19,18 +19,18 @@ class InferenceRecommenderMixin: inference_recommendations: Incomplete def right_size( self, - sample_payload_url: str = None, - supported_content_types: List[str] = None, - supported_instance_types: List[str] = None, - job_name: str = None, - framework: str = None, - job_duration_in_seconds: int = None, - hyperparameter_ranges: List[Dict[str, CategoricalParameter]] = None, - phases: List[Phase] = None, - traffic_type: str = None, - max_invocations: int = None, - model_latency_thresholds: List[ModelLatencyThreshold] = None, - max_tests: int = None, - max_parallel_tests: int = None, - log_level: Optional[str] = "Verbose", + sample_payload_url: str | None = None, + supported_content_types: list[str] | None = None, + supported_instance_types: list[str] | None = None, + job_name: str | None = None, + framework: str | None = None, + job_duration_in_seconds: int | None = None, + hyperparameter_ranges: list[dict[str, CategoricalParameter]] | None = None, + phases: list[Phase] | None = None, + traffic_type: str | None = None, + max_invocations: int | None = None, + model_latency_thresholds: list[ModelLatencyThreshold] | None = None, + max_tests: int | None = None, + max_parallel_tests: int | None = None, + log_level: str | None = "Verbose", ): ... diff --git a/stubs/sagemaker/sagemaker/inputs.pyi b/stubs/sagemaker/sagemaker/inputs.pyi index 0eacb811ab70..c61e7ae20dac 100644 --- a/stubs/sagemaker/sagemaker/inputs.pyi +++ b/stubs/sagemaker/sagemaker/inputs.pyi @@ -11,15 +11,15 @@ class TrainingInput: def __init__( self, s3_data: str | PipelineVariable, - distribution: Optional[str | PipelineVariable] = None, - compression: Optional[str | PipelineVariable] = None, - content_type: Optional[str | PipelineVariable] = None, - record_wrapping: Optional[str | PipelineVariable] = None, + distribution: str | PipelineVariable | None = None, + compression: str | PipelineVariable | None = None, + content_type: str | PipelineVariable | None = None, + record_wrapping: str | PipelineVariable | None = None, s3_data_type: str | PipelineVariable = "S3Prefix", - instance_groups: Optional[List[str | PipelineVariable]] = None, - input_mode: Optional[str | PipelineVariable] = None, - attribute_names: Optional[List[str | PipelineVariable]] = None, - target_attribute_name: Optional[str | PipelineVariable] = None, + instance_groups: list[str | PipelineVariable] | None = None, + input_mode: str | PipelineVariable | None = None, + attribute_names: list[str | PipelineVariable] | None = None, + target_attribute_name: str | PipelineVariable | None = None, shuffle_config: Optional["ShuffleConfig"] = None, ) -> None: ... @@ -80,4 +80,4 @@ class BatchDataCaptureConfig: destination_s3_uri: Incomplete kms_key_id: Incomplete generate_inference_id: Incomplete - def __init__(self, destination_s3_uri: str, kms_key_id: str = None, generate_inference_id: bool = None) -> None: ... + def __init__(self, destination_s3_uri: str, kms_key_id: str | None = None, generate_inference_id: bool | None = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/instance_types.pyi b/stubs/sagemaker/sagemaker/instance_types.pyi index f6d25ce7eede..5d86e5786160 100644 --- a/stubs/sagemaker/sagemaker/instance_types.pyi +++ b/stubs/sagemaker/sagemaker/instance_types.pyi @@ -4,18 +4,18 @@ from typing import List, Optional logger: Incomplete def retrieve_default( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, - scope: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, + scope: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, ) -> str: ... def retrieve( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, - scope: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, + scope: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, -) -> List[str]: ... +) -> list[str]: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/accessors.pyi b/stubs/sagemaker/sagemaker/jumpstart/accessors.pyi index cf1367aad09f..097c23013cb9 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/accessors.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/accessors.pyi @@ -14,8 +14,8 @@ class JumpStartModelsAccessor: @staticmethod def get_model_specs(region: str, model_id: str, version: str) -> JumpStartModelSpecs: ... @staticmethod - def set_cache_kwargs(cache_kwargs: Dict[str, Any], region: str = None) -> None: ... + def set_cache_kwargs(cache_kwargs: dict[str, Any], region: str | None = None) -> None: ... @staticmethod - def reset_cache(cache_kwargs: Dict[str, Any] = None, region: Optional[str] = None) -> None: ... + def reset_cache(cache_kwargs: dict[str, Any] | None = None, region: str | None = None) -> None: ... @staticmethod - def get_manifest(cache_kwargs: Optional[Dict[str, Any]] = ..., region: Optional[str] = ...) -> List[JumpStartModelHeader]: ... + def get_manifest(cache_kwargs: dict[str, Any] | None = ..., region: str | None = ...) -> list[JumpStartModelHeader]: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/cache.pyi b/stubs/sagemaker/sagemaker/jumpstart/cache.pyi index bac6d2a0e7a2..ebbbdac1fb56 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/cache.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/cache.pyi @@ -15,8 +15,8 @@ class JumpStartModelsCache: max_semantic_version_cache_items: int = 20, semantic_version_cache_expiration_horizon: datetime.timedelta = ..., manifest_file_s3_key: str = "models_manifest.json", - s3_bucket_name: Optional[str] = None, - s3_client_config: Optional[botocore.config.Config] = None, + s3_bucket_name: str | None = None, + s3_client_config: botocore.config.Config | None = None, ) -> None: ... def set_region(self, region: str) -> None: ... def get_region(self) -> str: ... @@ -24,7 +24,7 @@ class JumpStartModelsCache: def get_manifest_file_s3_key(self) -> str: ... def set_s3_bucket_name(self, s3_bucket_name: str) -> None: ... def get_bucket(self) -> str: ... - def get_manifest(self) -> List[JumpStartModelHeader]: ... + def get_manifest(self) -> list[JumpStartModelHeader]: ... def get_header(self, model_id: str, semantic_version_str: str) -> JumpStartModelHeader: ... def get_specs(self, model_id: str, semantic_version_str: str) -> JumpStartModelSpecs: ... def clear(self) -> None: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/constants.pyi b/stubs/sagemaker/sagemaker/jumpstart/constants.pyi index 5064010a0302..c095b070d3de 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/constants.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/constants.pyi @@ -6,7 +6,7 @@ from sagemaker.base_serializers import BaseSerializer from sagemaker.jumpstart.enums import DeserializerType, MIMEType, SerializerType from sagemaker.jumpstart.types import JumpStartLaunchedRegionInfo -JUMPSTART_LAUNCHED_REGIONS: Set[JumpStartLaunchedRegionInfo] +JUMPSTART_LAUNCHED_REGIONS: set[JumpStartLaunchedRegionInfo] JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT: Incomplete JUMPSTART_REGION_NAME_SET: Incomplete JUMPSTART_BUCKET_NAME_SET: Incomplete @@ -21,8 +21,8 @@ ENV_VARIABLE_JUMPSTART_SCRIPT_ARTIFACT_BUCKET_OVERRIDE: str ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE: str ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE: str JUMPSTART_RESOURCE_BASE_NAME: str -CONTENT_TYPE_TO_SERIALIZER_TYPE_MAP: Dict[MIMEType, SerializerType] -ACCEPT_TYPE_TO_DESERIALIZER_TYPE_MAP: Dict[MIMEType, DeserializerType] -SERIALIZER_TYPE_TO_CLASS_MAP: Dict[SerializerType, Type[BaseSerializer]] -DESERIALIZER_TYPE_TO_CLASS_MAP: Dict[DeserializerType, Type[BaseDeserializer]] +CONTENT_TYPE_TO_SERIALIZER_TYPE_MAP: dict[MIMEType, SerializerType] +ACCEPT_TYPE_TO_DESERIALIZER_TYPE_MAP: dict[MIMEType, DeserializerType] +SERIALIZER_TYPE_TO_CLASS_MAP: dict[SerializerType, type[BaseSerializer]] +DESERIALIZER_TYPE_TO_CLASS_MAP: dict[DeserializerType, type[BaseDeserializer]] MODEL_ID_LIST_WEB_URL: str diff --git a/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi b/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi index 8e87c1e384f9..041a0ad70610 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi @@ -31,97 +31,97 @@ class JumpStartEstimator(Estimator): sagemaker_session: Incomplete def __init__( self, - model_id: Optional[str] = None, - model_version: Optional[str] = None, - tolerate_vulnerable_model: Optional[bool] = None, - tolerate_deprecated_model: Optional[bool] = None, - region: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, - role: Optional[str] = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, - keep_alive_period_in_seconds: Optional[int | PipelineVariable] = None, - volume_size: Optional[int | PipelineVariable] = None, - volume_kms_key: Optional[str | PipelineVariable] = None, - max_run: Optional[int | PipelineVariable] = None, - input_mode: Optional[str | PipelineVariable] = None, - output_path: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[session.Session] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - subnets: Optional[List[str | PipelineVariable]] = None, - security_group_ids: Optional[List[str | PipelineVariable]] = None, - model_uri: Optional[str] = None, - model_channel_name: Optional[str | PipelineVariable] = None, - metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, - encrypt_inter_container_traffic: bool | PipelineVariable = None, - use_spot_instances: Optional[bool | PipelineVariable] = None, - max_wait: Optional[int | PipelineVariable] = None, - checkpoint_s3_uri: Optional[str | PipelineVariable] = None, - checkpoint_local_path: Optional[str | PipelineVariable] = None, - enable_network_isolation: bool | PipelineVariable = None, - rules: Optional[List[RuleBase]] = None, - debugger_hook_config: Optional[DebuggerHookConfig | bool] = None, - tensorboard_output_config: Optional[TensorBoardOutputConfig] = None, - enable_sagemaker_metrics: Optional[bool | PipelineVariable] = None, - profiler_config: Optional[ProfilerConfig] = None, - disable_profiler: Optional[bool] = None, - environment: Optional[Dict[str, str | PipelineVariable]] = None, - max_retry_attempts: Optional[int | PipelineVariable] = None, - source_dir: Optional[str | PipelineVariable] = None, - git_config: Optional[Dict[str, str]] = None, - container_log_level: Optional[int | PipelineVariable] = None, - code_location: Optional[str] = None, - entry_point: Optional[str | PipelineVariable] = None, - dependencies: Optional[List[str]] = None, - instance_groups: Optional[List[InstanceGroup]] = None, - training_repository_access_mode: Optional[str | PipelineVariable] = None, - training_repository_credentials_provider_arn: Optional[str | PipelineVariable] = None, + model_id: str | None = None, + model_version: str | None = None, + tolerate_vulnerable_model: bool | None = None, + tolerate_deprecated_model: bool | None = None, + region: str | None = None, + image_uri: str | PipelineVariable | None = None, + role: str | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + keep_alive_period_in_seconds: int | PipelineVariable | None = None, + volume_size: int | PipelineVariable | None = None, + volume_kms_key: str | PipelineVariable | None = None, + max_run: int | PipelineVariable | None = None, + input_mode: str | PipelineVariable | None = None, + output_path: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: session.Session | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + subnets: list[str | PipelineVariable] | None = None, + security_group_ids: list[str | PipelineVariable] | None = None, + model_uri: str | None = None, + model_channel_name: str | PipelineVariable | None = None, + metric_definitions: list[dict[str, str | PipelineVariable]] | None = None, + encrypt_inter_container_traffic: bool | PipelineVariable | None = None, + use_spot_instances: bool | PipelineVariable | None = None, + max_wait: int | PipelineVariable | None = None, + checkpoint_s3_uri: str | PipelineVariable | None = None, + checkpoint_local_path: str | PipelineVariable | None = None, + enable_network_isolation: bool | PipelineVariable | None = None, + rules: list[RuleBase] | None = None, + debugger_hook_config: DebuggerHookConfig | bool | None = None, + tensorboard_output_config: TensorBoardOutputConfig | None = None, + enable_sagemaker_metrics: bool | PipelineVariable | None = None, + profiler_config: ProfilerConfig | None = None, + disable_profiler: bool | None = None, + environment: dict[str, str | PipelineVariable] | None = None, + max_retry_attempts: int | PipelineVariable | None = None, + source_dir: str | PipelineVariable | None = None, + git_config: dict[str, str] | None = None, + container_log_level: int | PipelineVariable | None = None, + code_location: str | None = None, + entry_point: str | PipelineVariable | None = None, + dependencies: list[str] | None = None, + instance_groups: list[InstanceGroup] | None = None, + training_repository_access_mode: str | PipelineVariable | None = None, + training_repository_credentials_provider_arn: str | PipelineVariable | None = None, ) -> None: ... def fit( self, - inputs: Optional[str | Dict | TrainingInput | FileSystemInput] = None, - wait: Optional[bool] = True, - logs: Optional[str] = None, - job_name: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, + inputs: str | dict | TrainingInput | FileSystemInput | None = None, + wait: bool | None = True, + logs: str | None = None, + job_name: str | None = None, + experiment_config: dict[str, str] | None = None, ) -> None: ... def deploy( self, - initial_instance_count: Optional[int] = None, - instance_type: Optional[str] = None, - serializer: Optional[BaseSerializer] = None, - deserializer: Optional[BaseDeserializer] = None, - accelerator_type: Optional[str] = None, - endpoint_name: Optional[str] = None, - tags: List[Dict[str, str]] = None, - kms_key: Optional[str] = None, - wait: Optional[bool] = True, - data_capture_config: Optional[DataCaptureConfig] = None, - async_inference_config: Optional[AsyncInferenceConfig] = None, - serverless_inference_config: Optional[ServerlessInferenceConfig] = None, - volume_size: Optional[int] = None, - model_data_download_timeout: Optional[int] = None, - container_startup_health_check_timeout: Optional[int] = None, - inference_recommendation_id: Optional[str] = None, - explainer_config: Optional[ExplainerConfig] = None, - image_uri: Optional[str | PipelineVariable] = None, - role: Optional[str] = None, - predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - model_name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[str | PipelineVariable]]] = None, - sagemaker_session: Optional[session.Session] = None, - enable_network_isolation: bool | PipelineVariable = None, - model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, str | PipelineVariable]] = None, - source_dir: Optional[str] = None, - code_location: Optional[str] = None, - entry_point: Optional[str] = None, - container_log_level: Optional[int | PipelineVariable] = None, - dependencies: Optional[List[str]] = None, - git_config: Optional[Dict[str, str]] = None, + initial_instance_count: int | None = None, + instance_type: str | None = None, + serializer: BaseSerializer | None = None, + deserializer: BaseDeserializer | None = None, + accelerator_type: str | None = None, + endpoint_name: str | None = None, + tags: list[dict[str, str]] | None = None, + kms_key: str | None = None, + wait: bool | None = True, + data_capture_config: DataCaptureConfig | None = None, + async_inference_config: AsyncInferenceConfig | None = None, + serverless_inference_config: ServerlessInferenceConfig | None = None, + volume_size: int | None = None, + model_data_download_timeout: int | None = None, + container_startup_health_check_timeout: int | None = None, + inference_recommendation_id: str | None = None, + explainer_config: ExplainerConfig | None = None, + image_uri: str | PipelineVariable | None = None, + role: str | None = None, + predictor_cls: callable | None = None, + env: dict[str, str | PipelineVariable] | None = None, + model_name: str | None = None, + vpc_config: dict[str, list[str | PipelineVariable]] | None = None, + sagemaker_session: session.Session | None = None, + enable_network_isolation: bool | PipelineVariable | None = None, + model_kms_key: str | None = None, + image_config: dict[str, str | PipelineVariable] | None = None, + source_dir: str | None = None, + code_location: str | None = None, + entry_point: str | None = None, + container_log_level: int | PipelineVariable | None = None, + dependencies: list[str] | None = None, + git_config: dict[str, str] | None = None, use_compiled_model: bool = False, ) -> PredictorBase: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi b/stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi index a415febf0ce1..60be283cc732 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi @@ -8,19 +8,19 @@ INVALID_MODEL_ID_ERROR_MSG: Incomplete class JumpStartHyperparametersError(ValueError): message: Incomplete - def __init__(self, message: Optional[str] = None) -> None: ... + def __init__(self, message: str | None = None) -> None: ... class VulnerableJumpStartModelError(ValueError): message: Incomplete def __init__( self, - model_id: Optional[str] = None, - version: Optional[str] = None, - vulnerabilities: Optional[List[str]] = None, - scope: Optional[JumpStartScriptScope] = None, - message: Optional[str] = None, + model_id: str | None = None, + version: str | None = None, + vulnerabilities: list[str] | None = None, + scope: JumpStartScriptScope | None = None, + message: str | None = None, ) -> None: ... class DeprecatedJumpStartModelError(ValueError): message: Incomplete - def __init__(self, model_id: Optional[str] = None, version: Optional[str] = None, message: Optional[str] = None) -> None: ... + def __init__(self, model_id: str | None = None, version: str | None = None, message: str | None = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi b/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi index 19fee6bb3395..630e15be1a1e 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi @@ -19,104 +19,104 @@ logger: Incomplete def get_init_kwargs( model_id: str, - model_version: Optional[str] = None, - tolerate_vulnerable_model: Optional[bool] = None, - tolerate_deprecated_model: Optional[bool] = None, - region: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, - role: Optional[str] = None, - instance_count: Optional[int | PipelineVariable] = None, - instance_type: Optional[str | PipelineVariable] = None, - keep_alive_period_in_seconds: Optional[int | PipelineVariable] = None, - volume_size: Optional[int | PipelineVariable] = None, - volume_kms_key: Optional[str | PipelineVariable] = None, - max_run: Optional[int | PipelineVariable] = None, - input_mode: Optional[str | PipelineVariable] = None, - output_path: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - subnets: Optional[List[str | PipelineVariable]] = None, - security_group_ids: Optional[List[str | PipelineVariable]] = None, - model_uri: Optional[str] = None, - model_channel_name: Optional[str | PipelineVariable] = None, - metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, - encrypt_inter_container_traffic: bool | PipelineVariable = None, - use_spot_instances: Optional[bool | PipelineVariable] = None, - max_wait: Optional[int | PipelineVariable] = None, - checkpoint_s3_uri: Optional[str | PipelineVariable] = None, - checkpoint_local_path: Optional[str | PipelineVariable] = None, - enable_network_isolation: bool | PipelineVariable = None, - rules: Optional[List[RuleBase]] = None, - debugger_hook_config: Optional[DebuggerHookConfig | bool] = None, - tensorboard_output_config: Optional[TensorBoardOutputConfig] = None, - enable_sagemaker_metrics: Optional[bool | PipelineVariable] = None, - profiler_config: Optional[ProfilerConfig] = None, - disable_profiler: Optional[bool] = None, - environment: Optional[Dict[str, str | PipelineVariable]] = None, - max_retry_attempts: Optional[int | PipelineVariable] = None, - source_dir: Optional[str | PipelineVariable] = None, - git_config: Optional[Dict[str, str]] = None, - container_log_level: Optional[int | PipelineVariable] = None, - code_location: Optional[str] = None, - entry_point: Optional[str | PipelineVariable] = None, - dependencies: Optional[List[str]] = None, - instance_groups: Optional[List[InstanceGroup]] = None, - training_repository_access_mode: Optional[str | PipelineVariable] = None, - training_repository_credentials_provider_arn: Optional[str | PipelineVariable] = None, + model_version: str | None = None, + tolerate_vulnerable_model: bool | None = None, + tolerate_deprecated_model: bool | None = None, + region: str | None = None, + image_uri: str | PipelineVariable | None = None, + role: str | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + keep_alive_period_in_seconds: int | PipelineVariable | None = None, + volume_size: int | PipelineVariable | None = None, + volume_kms_key: str | PipelineVariable | None = None, + max_run: int | PipelineVariable | None = None, + input_mode: str | PipelineVariable | None = None, + output_path: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + subnets: list[str | PipelineVariable] | None = None, + security_group_ids: list[str | PipelineVariable] | None = None, + model_uri: str | None = None, + model_channel_name: str | PipelineVariable | None = None, + metric_definitions: list[dict[str, str | PipelineVariable]] | None = None, + encrypt_inter_container_traffic: bool | PipelineVariable | None = None, + use_spot_instances: bool | PipelineVariable | None = None, + max_wait: int | PipelineVariable | None = None, + checkpoint_s3_uri: str | PipelineVariable | None = None, + checkpoint_local_path: str | PipelineVariable | None = None, + enable_network_isolation: bool | PipelineVariable | None = None, + rules: list[RuleBase] | None = None, + debugger_hook_config: DebuggerHookConfig | bool | None = None, + tensorboard_output_config: TensorBoardOutputConfig | None = None, + enable_sagemaker_metrics: bool | PipelineVariable | None = None, + profiler_config: ProfilerConfig | None = None, + disable_profiler: bool | None = None, + environment: dict[str, str | PipelineVariable] | None = None, + max_retry_attempts: int | PipelineVariable | None = None, + source_dir: str | PipelineVariable | None = None, + git_config: dict[str, str] | None = None, + container_log_level: int | PipelineVariable | None = None, + code_location: str | None = None, + entry_point: str | PipelineVariable | None = None, + dependencies: list[str] | None = None, + instance_groups: list[InstanceGroup] | None = None, + training_repository_access_mode: str | PipelineVariable | None = None, + training_repository_credentials_provider_arn: str | PipelineVariable | None = None, ) -> JumpStartEstimatorInitKwargs: ... def get_fit_kwargs( model_id: str, - model_version: Optional[str] = None, - region: Optional[str] = None, - inputs: Optional[str | Dict | TrainingInput | FileSystemInput] = None, - wait: Optional[bool] = None, - logs: Optional[str] = None, - job_name: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, - tolerate_vulnerable_model: Optional[bool] = None, - tolerate_deprecated_model: Optional[bool] = None, + model_version: str | None = None, + region: str | None = None, + inputs: str | dict | TrainingInput | FileSystemInput | None = None, + wait: bool | None = None, + logs: str | None = None, + job_name: str | None = None, + experiment_config: dict[str, str] | None = None, + tolerate_vulnerable_model: bool | None = None, + tolerate_deprecated_model: bool | None = None, ) -> JumpStartEstimatorFitKwargs: ... def get_deploy_kwargs( model_id: str, - model_version: Optional[str] = None, - region: Optional[str] = None, - initial_instance_count: Optional[int] = None, - instance_type: Optional[str] = None, - serializer: Optional[BaseSerializer] = None, - deserializer: Optional[BaseDeserializer] = None, - accelerator_type: Optional[str] = None, - endpoint_name: Optional[str] = None, - tags: List[Dict[str, str]] = None, - kms_key: Optional[str] = None, - wait: Optional[bool] = None, - data_capture_config: Optional[DataCaptureConfig] = None, - async_inference_config: Optional[AsyncInferenceConfig] = None, - serverless_inference_config: Optional[ServerlessInferenceConfig] = None, - volume_size: Optional[int] = None, - model_data_download_timeout: Optional[int] = None, - container_startup_health_check_timeout: Optional[int] = None, - inference_recommendation_id: Optional[str] = None, - explainer_config: Optional[ExplainerConfig] = None, - image_uri: Optional[str | PipelineVariable] = None, - role: Optional[str] = None, - predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - vpc_config: Optional[Dict[str, List[str | PipelineVariable]]] = None, - sagemaker_session: Optional[Session] = None, - enable_network_isolation: bool | PipelineVariable = None, - model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, str | PipelineVariable]] = None, - source_dir: Optional[str] = None, - code_location: Optional[str] = None, - entry_point: Optional[str] = None, - container_log_level: Optional[int | PipelineVariable] = None, - dependencies: Optional[List[str]] = None, - git_config: Optional[Dict[str, str]] = None, - tolerate_deprecated_model: Optional[bool] = None, - tolerate_vulnerable_model: Optional[bool] = None, - use_compiled_model: Optional[bool] = None, - model_name: Optional[str] = None, + model_version: str | None = None, + region: str | None = None, + initial_instance_count: int | None = None, + instance_type: str | None = None, + serializer: BaseSerializer | None = None, + deserializer: BaseDeserializer | None = None, + accelerator_type: str | None = None, + endpoint_name: str | None = None, + tags: list[dict[str, str]] | None = None, + kms_key: str | None = None, + wait: bool | None = None, + data_capture_config: DataCaptureConfig | None = None, + async_inference_config: AsyncInferenceConfig | None = None, + serverless_inference_config: ServerlessInferenceConfig | None = None, + volume_size: int | None = None, + model_data_download_timeout: int | None = None, + container_startup_health_check_timeout: int | None = None, + inference_recommendation_id: str | None = None, + explainer_config: ExplainerConfig | None = None, + image_uri: str | PipelineVariable | None = None, + role: str | None = None, + predictor_cls: callable | None = None, + env: dict[str, str | PipelineVariable] | None = None, + vpc_config: dict[str, list[str | PipelineVariable]] | None = None, + sagemaker_session: Session | None = None, + enable_network_isolation: bool | PipelineVariable | None = None, + model_kms_key: str | None = None, + image_config: dict[str, str | PipelineVariable] | None = None, + source_dir: str | None = None, + code_location: str | None = None, + entry_point: str | None = None, + container_log_level: int | PipelineVariable | None = None, + dependencies: list[str] | None = None, + git_config: dict[str, str] | None = None, + tolerate_deprecated_model: bool | None = None, + tolerate_vulnerable_model: bool | None = None, + use_compiled_model: bool | None = None, + model_name: str | None = None, ) -> JumpStartEstimatorDeployKwargs: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi b/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi index deffcd4f5b52..b8e531fdb354 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi @@ -24,51 +24,51 @@ def get_default_predictor( ) -> Predictor: ... def get_deploy_kwargs( model_id: str, - model_version: Optional[str] = None, - region: Optional[str] = None, - initial_instance_count: Optional[int] = None, - instance_type: Optional[str] = None, - serializer: Optional[BaseSerializer] = None, - deserializer: Optional[BaseDeserializer] = None, - accelerator_type: Optional[str] = None, - endpoint_name: Optional[str] = None, - tags: List[Dict[str, str]] = None, - kms_key: Optional[str] = None, - wait: Optional[bool] = None, - data_capture_config: Optional[DataCaptureConfig] = None, - async_inference_config: Optional[AsyncInferenceConfig] = None, - serverless_inference_config: Optional[ServerlessInferenceConfig] = None, - volume_size: Optional[int] = None, - model_data_download_timeout: Optional[int] = None, - container_startup_health_check_timeout: Optional[int] = None, - inference_recommendation_id: Optional[str] = None, - explainer_config: Optional[ExplainerConfig] = None, - tolerate_vulnerable_model: Optional[bool] = None, - tolerate_deprecated_model: Optional[bool] = None, + model_version: str | None = None, + region: str | None = None, + initial_instance_count: int | None = None, + instance_type: str | None = None, + serializer: BaseSerializer | None = None, + deserializer: BaseDeserializer | None = None, + accelerator_type: str | None = None, + endpoint_name: str | None = None, + tags: list[dict[str, str]] | None = None, + kms_key: str | None = None, + wait: bool | None = None, + data_capture_config: DataCaptureConfig | None = None, + async_inference_config: AsyncInferenceConfig | None = None, + serverless_inference_config: ServerlessInferenceConfig | None = None, + volume_size: int | None = None, + model_data_download_timeout: int | None = None, + container_startup_health_check_timeout: int | None = None, + inference_recommendation_id: str | None = None, + explainer_config: ExplainerConfig | None = None, + tolerate_vulnerable_model: bool | None = None, + tolerate_deprecated_model: bool | None = None, ) -> JumpStartModelDeployKwargs: ... def get_init_kwargs( model_id: str, model_from_estimator: bool = False, - model_version: Optional[str] = None, - tolerate_vulnerable_model: Optional[bool] = None, - tolerate_deprecated_model: Optional[bool] = None, - instance_type: Optional[str] = None, - region: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, - model_data: Optional[str | PipelineVariable] = None, - role: Optional[str] = None, - predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[str | PipelineVariable]]] = None, - sagemaker_session: Optional[Session] = None, - enable_network_isolation: bool | PipelineVariable = None, - model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, str | PipelineVariable]] = None, - source_dir: Optional[str] = None, - code_location: Optional[str] = None, - entry_point: Optional[str] = None, - container_log_level: Optional[int | PipelineVariable] = None, - dependencies: Optional[List[str]] = None, - git_config: Optional[Dict[str, str]] = None, + model_version: str | None = None, + tolerate_vulnerable_model: bool | None = None, + tolerate_deprecated_model: bool | None = None, + instance_type: str | None = None, + region: str | None = None, + image_uri: str | PipelineVariable | None = None, + model_data: str | PipelineVariable | None = None, + role: str | None = None, + predictor_cls: callable | None = None, + env: dict[str, str | PipelineVariable] | None = None, + name: str | None = None, + vpc_config: dict[str, list[str | PipelineVariable]] | None = None, + sagemaker_session: Session | None = None, + enable_network_isolation: bool | PipelineVariable | None = None, + model_kms_key: str | None = None, + image_config: dict[str, str | PipelineVariable] | None = None, + source_dir: str | None = None, + code_location: str | None = None, + entry_point: str | None = None, + container_log_level: int | PipelineVariable | None = None, + dependencies: list[str] | None = None, + git_config: dict[str, str] | None = None, ) -> JumpStartModelInitKwargs: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/filters.pyi b/stubs/sagemaker/sagemaker/jumpstart/filters.pyi index 9dc9e8c8397b..e344a892d92d 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/filters.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/filters.pyi @@ -83,5 +83,5 @@ class ModelFilter(JumpStartDataHolderType): def parse_filter_string(filter_string: str) -> ModelFilter: ... def evaluate_filter_expression( - model_filter: ModelFilter, cached_model_value: str | bool | int | float | Dict[str | Any, List[Any]] + model_filter: ModelFilter, cached_model_value: str | bool | int | float | dict[str | Any, list[Any]] ) -> BooleanValues: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/model.pyi b/stubs/sagemaker/sagemaker/jumpstart/model.pyi index 0149adcefa18..1e7ba70eb65f 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/model.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/model.pyi @@ -24,47 +24,47 @@ class JumpStartModel(Model): region: Incomplete def __init__( self, - model_id: Optional[str] = None, - model_version: Optional[str] = None, - tolerate_vulnerable_model: Optional[bool] = None, - tolerate_deprecated_model: Optional[bool] = None, - region: Optional[str] = None, - instance_type: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, - model_data: Optional[str | PipelineVariable] = None, - role: Optional[str] = None, - predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[str | PipelineVariable]]] = None, - sagemaker_session: Optional[Session] = None, - enable_network_isolation: bool | PipelineVariable = None, - model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, str | PipelineVariable]] = None, - source_dir: Optional[str] = None, - code_location: Optional[str] = None, - entry_point: Optional[str] = None, - container_log_level: Optional[int | PipelineVariable] = None, - dependencies: Optional[List[str]] = None, - git_config: Optional[Dict[str, str]] = None, + model_id: str | None = None, + model_version: str | None = None, + tolerate_vulnerable_model: bool | None = None, + tolerate_deprecated_model: bool | None = None, + region: str | None = None, + instance_type: str | None = None, + image_uri: str | PipelineVariable | None = None, + model_data: str | PipelineVariable | None = None, + role: str | None = None, + predictor_cls: callable | None = None, + env: dict[str, str | PipelineVariable] | None = None, + name: str | None = None, + vpc_config: dict[str, list[str | PipelineVariable]] | None = None, + sagemaker_session: Session | None = None, + enable_network_isolation: bool | PipelineVariable | None = None, + model_kms_key: str | None = None, + image_config: dict[str, str | PipelineVariable] | None = None, + source_dir: str | None = None, + code_location: str | None = None, + entry_point: str | None = None, + container_log_level: int | PipelineVariable | None = None, + dependencies: list[str] | None = None, + git_config: dict[str, str] | None = None, ) -> None: ... def deploy( self, - initial_instance_count: Optional[int] = None, - instance_type: Optional[str] = None, - serializer: Optional[BaseSerializer] = None, - deserializer: Optional[BaseDeserializer] = None, - accelerator_type: Optional[str] = None, - endpoint_name: Optional[str] = None, - tags: List[Dict[str, str]] = None, - kms_key: Optional[str] = None, - wait: Optional[bool] = True, - data_capture_config: Optional[DataCaptureConfig] = None, - async_inference_config: Optional[AsyncInferenceConfig] = None, - serverless_inference_config: Optional[ServerlessInferenceConfig] = None, - volume_size: Optional[int] = None, - model_data_download_timeout: Optional[int] = None, - container_startup_health_check_timeout: Optional[int] = None, - inference_recommendation_id: Optional[str] = None, - explainer_config: Optional[ExplainerConfig] = None, + initial_instance_count: int | None = None, + instance_type: str | None = None, + serializer: BaseSerializer | None = None, + deserializer: BaseDeserializer | None = None, + accelerator_type: str | None = None, + endpoint_name: str | None = None, + tags: list[dict[str, str]] | None = None, + kms_key: str | None = None, + wait: bool | None = True, + data_capture_config: DataCaptureConfig | None = None, + async_inference_config: AsyncInferenceConfig | None = None, + serverless_inference_config: ServerlessInferenceConfig | None = None, + volume_size: int | None = None, + model_data_download_timeout: int | None = None, + container_startup_health_check_timeout: int | None = None, + inference_recommendation_id: str | None = None, + explainer_config: ExplainerConfig | None = None, ) -> PredictorBase: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi b/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi index 4e026fdb42ac..5290926abc8d 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi @@ -2,15 +2,15 @@ from typing import List, Tuple from sagemaker.jumpstart.filters import Operator -def extract_framework_task_model(model_id: str) -> Tuple[str, str, str]: ... -def list_jumpstart_tasks(filter: Operator | str = ..., region: str = "eu-west-1") -> List[str]: ... -def list_jumpstart_frameworks(filter: Operator | str = ..., region: str = "eu-west-1") -> List[str]: ... -def list_jumpstart_scripts(filter: Operator | str = ..., region: str = "eu-west-1") -> List[str]: ... +def extract_framework_task_model(model_id: str) -> tuple[str, str, str]: ... +def list_jumpstart_tasks(filter: Operator | str = ..., region: str = "eu-west-1") -> list[str]: ... +def list_jumpstart_frameworks(filter: Operator | str = ..., region: str = "eu-west-1") -> list[str]: ... +def list_jumpstart_scripts(filter: Operator | str = ..., region: str = "eu-west-1") -> list[str]: ... def list_jumpstart_models( filter: Operator | str = ..., region: str = "eu-west-1", list_incomplete_models: bool = False, list_old_models: bool = False, list_versions: bool = False, -) -> List[Tuple[str, Tuple[str, str]]]: ... +) -> list[tuple[str, tuple[str, str]]]: ... def get_model_url(model_id: str, model_version: str, region: str = "eu-west-1") -> str: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/types.pyi b/stubs/sagemaker/sagemaker/jumpstart/types.pyi index 397cc6f6eb19..97f93a5fd1cc 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/types.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/types.pyi @@ -16,25 +16,25 @@ class JumpStartLaunchedRegionInfo(JumpStartDataHolderType): def __init__(self, content_bucket: str, region_name: str) -> None: ... class JumpStartModelHeader(JumpStartDataHolderType): - def __init__(self, header: Dict[str, str]) -> None: ... - def to_json(self) -> Dict[str, str]: ... + def __init__(self, header: dict[str, str]) -> None: ... + def to_json(self) -> dict[str, str]: ... model_id: Incomplete version: Incomplete min_version: Incomplete spec_key: Incomplete - def from_json(self, json_obj: Dict[str, str]) -> None: ... + def from_json(self, json_obj: dict[str, str]) -> None: ... class JumpStartECRSpecs(JumpStartDataHolderType): - def __init__(self, spec: Dict[str, Any]) -> None: ... + def __init__(self, spec: dict[str, Any]) -> None: ... framework: Incomplete framework_version: Incomplete py_version: Incomplete huggingface_transformers_version: Incomplete - def from_json(self, json_obj: Dict[str, Any]) -> None: ... - def to_json(self) -> Dict[str, Any]: ... + def from_json(self, json_obj: dict[str, Any]) -> None: ... + def to_json(self) -> dict[str, Any]: ... class JumpStartHyperparameter(JumpStartDataHolderType): - def __init__(self, spec: Dict[str, Any]) -> None: ... + def __init__(self, spec: dict[str, Any]) -> None: ... name: Incomplete type: Incomplete default: Incomplete @@ -44,30 +44,30 @@ class JumpStartHyperparameter(JumpStartDataHolderType): max: Incomplete exclusive_min: Incomplete exclusive_max: Incomplete - def from_json(self, json_obj: Dict[str, Any]) -> None: ... - def to_json(self) -> Dict[str, Any]: ... + def from_json(self, json_obj: dict[str, Any]) -> None: ... + def to_json(self) -> dict[str, Any]: ... class JumpStartEnvironmentVariable(JumpStartDataHolderType): - def __init__(self, spec: Dict[str, Any]) -> None: ... + def __init__(self, spec: dict[str, Any]) -> None: ... name: Incomplete type: Incomplete default: Incomplete scope: Incomplete required_for_model_class: Incomplete - def from_json(self, json_obj: Dict[str, Any]) -> None: ... - def to_json(self) -> Dict[str, Any]: ... + def from_json(self, json_obj: dict[str, Any]) -> None: ... + def to_json(self) -> dict[str, Any]: ... class JumpStartPredictorSpecs(JumpStartDataHolderType): - def __init__(self, spec: Optional[Dict[str, Any]]) -> None: ... + def __init__(self, spec: dict[str, Any] | None) -> None: ... default_content_type: Incomplete supported_content_types: Incomplete default_accept_type: Incomplete supported_accept_types: Incomplete - def from_json(self, json_obj: Optional[Dict[str, Any]]) -> None: ... - def to_json(self) -> Dict[str, Any]: ... + def from_json(self, json_obj: dict[str, Any] | None) -> None: ... + def to_json(self) -> dict[str, Any]: ... class JumpStartModelSpecs(JumpStartDataHolderType): - def __init__(self, spec: Dict[str, Any]) -> None: ... + def __init__(self, spec: dict[str, Any]) -> None: ... model_id: Incomplete url: Incomplete version: Incomplete @@ -106,8 +106,8 @@ class JumpStartModelSpecs(JumpStartDataHolderType): fit_kwargs: Incomplete training_volume_size: Incomplete training_enable_network_isolation: Incomplete - def from_json(self, json_obj: Dict[str, Any]) -> None: ... - def to_json(self) -> Dict[str, Any]: ... + def from_json(self, json_obj: dict[str, Any]) -> None: ... + def to_json(self) -> dict[str, Any]: ... def supports_prepacked_inference(self) -> bool: ... def supports_incremental_training(self) -> bool: ... @@ -126,12 +126,12 @@ class JumpStartCachedS3ContentValue(JumpStartDataHolderType): md5_hash: Incomplete def __init__( self, - formatted_content: Dict[JumpStartVersionedModelId | JumpStartModelHeader, JumpStartModelSpecs], - md5_hash: Optional[str] = None, + formatted_content: dict[JumpStartVersionedModelId | JumpStartModelHeader, JumpStartModelSpecs], + md5_hash: str | None = None, ) -> None: ... class JumpStartKwargs(JumpStartDataHolderType): - SERIALIZATION_EXCLUSION_SET: Set[str] + SERIALIZATION_EXCLUSION_SET: set[str] def to_kwargs_dict(self): ... class JumpStartModelInitKwargs(JumpStartKwargs): @@ -162,28 +162,28 @@ class JumpStartModelInitKwargs(JumpStartKwargs): def __init__( self, model_id: str, - model_version: Optional[str] = None, - region: Optional[str] = None, - instance_type: Optional[str] = None, - image_uri: Optional[str | Any] = None, - model_data: Optional[str | Any] = None, - role: Optional[str] = None, - predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, str | Any]] = None, - name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[str | Any]]] = None, - sagemaker_session: Optional[Any] = None, + model_version: str | None = None, + region: str | None = None, + instance_type: str | None = None, + image_uri: str | Any | None = None, + model_data: str | Any | None = None, + role: str | None = None, + predictor_cls: callable | None = None, + env: dict[str, str | Any] | None = None, + name: str | None = None, + vpc_config: dict[str, list[str | Any]] | None = None, + sagemaker_session: Any | None = None, enable_network_isolation: bool | Any = None, - model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, str | Any]] = None, - source_dir: Optional[str] = None, - code_location: Optional[str] = None, - entry_point: Optional[str] = None, - container_log_level: Optional[int | Any] = None, - dependencies: Optional[List[str]] = None, - git_config: Optional[Dict[str, str]] = None, - tolerate_vulnerable_model: Optional[bool] = None, - tolerate_deprecated_model: Optional[bool] = None, + model_kms_key: str | None = None, + image_config: dict[str, str | Any] | None = None, + source_dir: str | None = None, + code_location: str | None = None, + entry_point: str | None = None, + container_log_level: int | Any | None = None, + dependencies: list[str] | None = None, + git_config: dict[str, str] | None = None, + tolerate_vulnerable_model: bool | None = None, + tolerate_deprecated_model: bool | None = None, ) -> None: ... class JumpStartModelDeployKwargs(JumpStartKwargs): @@ -213,27 +213,27 @@ class JumpStartModelDeployKwargs(JumpStartKwargs): def __init__( self, model_id: str, - model_version: Optional[str] = None, - region: Optional[str] = None, - initial_instance_count: Optional[int] = None, - instance_type: Optional[str] = None, - serializer: Optional[Any] = None, - deserializer: Optional[Any] = None, - accelerator_type: Optional[str] = None, - endpoint_name: Optional[str] = None, - tags: List[Dict[str, str]] = None, - kms_key: Optional[str] = None, - wait: Optional[bool] = None, - data_capture_config: Optional[Any] = None, - async_inference_config: Optional[Any] = None, - serverless_inference_config: Optional[Any] = None, - volume_size: Optional[int] = None, - model_data_download_timeout: Optional[int] = None, - container_startup_health_check_timeout: Optional[int] = None, - inference_recommendation_id: Optional[str] = None, - explainer_config: Optional[Any] = None, - tolerate_deprecated_model: Optional[bool] = None, - tolerate_vulnerable_model: Optional[bool] = None, + model_version: str | None = None, + region: str | None = None, + initial_instance_count: int | None = None, + instance_type: str | None = None, + serializer: Any | None = None, + deserializer: Any | None = None, + accelerator_type: str | None = None, + endpoint_name: str | None = None, + tags: list[dict[str, str]] | None = None, + kms_key: str | None = None, + wait: bool | None = None, + data_capture_config: Any | None = None, + async_inference_config: Any | None = None, + serverless_inference_config: Any | None = None, + volume_size: int | None = None, + model_data_download_timeout: int | None = None, + container_startup_health_check_timeout: int | None = None, + inference_recommendation_id: str | None = None, + explainer_config: Any | None = None, + tolerate_deprecated_model: bool | None = None, + tolerate_vulnerable_model: bool | None = None, ) -> None: ... class JumpStartEstimatorInitKwargs(JumpStartKwargs): @@ -289,53 +289,53 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs): def __init__( self, model_id: str, - model_version: Optional[str] = None, - region: Optional[str] = None, - image_uri: Optional[str | Any] = None, - role: Optional[str] = None, - instance_count: Optional[int | Any] = None, - instance_type: Optional[str | Any] = None, - keep_alive_period_in_seconds: Optional[int | Any] = None, - volume_size: Optional[int | Any] = None, - volume_kms_key: Optional[str | Any] = None, - max_run: Optional[int | Any] = None, - input_mode: Optional[str | Any] = None, - output_path: Optional[str | Any] = None, - output_kms_key: Optional[str | Any] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Any] = None, - hyperparameters: Optional[Dict[str, str | Any]] = None, - tags: Optional[List[Dict[str, str | Any]]] = None, - subnets: Optional[List[str | Any]] = None, - security_group_ids: Optional[List[str | Any]] = None, - model_uri: Optional[str] = None, - model_channel_name: Optional[str | Any] = None, - metric_definitions: Optional[List[Dict[str, str | Any]]] = None, + model_version: str | None = None, + region: str | None = None, + image_uri: str | Any | None = None, + role: str | None = None, + instance_count: int | Any | None = None, + instance_type: str | Any | None = None, + keep_alive_period_in_seconds: int | Any | None = None, + volume_size: int | Any | None = None, + volume_kms_key: str | Any | None = None, + max_run: int | Any | None = None, + input_mode: str | Any | None = None, + output_path: str | Any | None = None, + output_kms_key: str | Any | None = None, + base_job_name: str | None = None, + sagemaker_session: Any | None = None, + hyperparameters: dict[str, str | Any] | None = None, + tags: list[dict[str, str | Any]] | None = None, + subnets: list[str | Any] | None = None, + security_group_ids: list[str | Any] | None = None, + model_uri: str | None = None, + model_channel_name: str | Any | None = None, + metric_definitions: list[dict[str, str | Any]] | None = None, encrypt_inter_container_traffic: bool | Any = None, - use_spot_instances: Optional[bool | Any] = None, - max_wait: Optional[int | Any] = None, - checkpoint_s3_uri: Optional[str | Any] = None, - checkpoint_local_path: Optional[str | Any] = None, + use_spot_instances: bool | Any | None = None, + max_wait: int | Any | None = None, + checkpoint_s3_uri: str | Any | None = None, + checkpoint_local_path: str | Any | None = None, enable_network_isolation: bool | Any = None, - rules: Optional[List[Any]] = None, - debugger_hook_config: Optional[Any | bool] = None, - tensorboard_output_config: Optional[Any] = None, - enable_sagemaker_metrics: Optional[bool | Any] = None, - profiler_config: Optional[Any] = None, - disable_profiler: Optional[bool] = None, - environment: Optional[Dict[str, str | Any]] = None, - max_retry_attempts: Optional[int | Any] = None, - source_dir: Optional[str | Any] = None, - git_config: Optional[Dict[str, str]] = None, - container_log_level: Optional[int | Any] = None, - code_location: Optional[str] = None, - entry_point: Optional[str | Any] = None, - dependencies: Optional[List[str]] = None, - instance_groups: Optional[List[Any]] = None, - training_repository_access_mode: Optional[str | Any] = None, - training_repository_credentials_provider_arn: Optional[str | Any] = None, - tolerate_vulnerable_model: Optional[bool] = None, - tolerate_deprecated_model: Optional[bool] = None, + rules: list[Any] | None = None, + debugger_hook_config: Any | bool | None = None, + tensorboard_output_config: Any | None = None, + enable_sagemaker_metrics: bool | Any | None = None, + profiler_config: Any | None = None, + disable_profiler: bool | None = None, + environment: dict[str, str | Any] | None = None, + max_retry_attempts: int | Any | None = None, + source_dir: str | Any | None = None, + git_config: dict[str, str] | None = None, + container_log_level: int | Any | None = None, + code_location: str | None = None, + entry_point: str | Any | None = None, + dependencies: list[str] | None = None, + instance_groups: list[Any] | None = None, + training_repository_access_mode: str | Any | None = None, + training_repository_credentials_provider_arn: str | Any | None = None, + tolerate_vulnerable_model: bool | None = None, + tolerate_deprecated_model: bool | None = None, ) -> None: ... class JumpStartEstimatorFitKwargs(JumpStartKwargs): @@ -353,15 +353,15 @@ class JumpStartEstimatorFitKwargs(JumpStartKwargs): def __init__( self, model_id: str, - model_version: Optional[str] = None, - region: Optional[str] = None, - inputs: Optional[str | Dict | Any | Any] = None, - wait: Optional[bool] = None, - logs: Optional[str] = None, - job_name: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, - tolerate_deprecated_model: Optional[bool] = None, - tolerate_vulnerable_model: Optional[bool] = None, + model_version: str | None = None, + region: str | None = None, + inputs: str | dict | Any | Any | None = None, + wait: bool | None = None, + logs: str | None = None, + job_name: str | None = None, + experiment_config: dict[str, str] | None = None, + tolerate_deprecated_model: bool | None = None, + tolerate_vulnerable_model: bool | None = None, ) -> None: ... class JumpStartEstimatorDeployKwargs(JumpStartKwargs): @@ -408,42 +408,42 @@ class JumpStartEstimatorDeployKwargs(JumpStartKwargs): def __init__( self, model_id: str, - model_version: Optional[str] = None, - region: Optional[str] = None, - initial_instance_count: Optional[int] = None, - instance_type: Optional[str] = None, - serializer: Optional[Any] = None, - deserializer: Optional[Any] = None, - accelerator_type: Optional[str] = None, - endpoint_name: Optional[str] = None, - tags: List[Dict[str, str]] = None, - kms_key: Optional[str] = None, - wait: Optional[bool] = None, - data_capture_config: Optional[Any] = None, - async_inference_config: Optional[Any] = None, - serverless_inference_config: Optional[Any] = None, - volume_size: Optional[int] = None, - model_data_download_timeout: Optional[int] = None, - container_startup_health_check_timeout: Optional[int] = None, - inference_recommendation_id: Optional[str] = None, - explainer_config: Optional[Any] = None, - image_uri: Optional[str | Any] = None, - role: Optional[str] = None, - predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, str | Any]] = None, - model_name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[str | Any]]] = None, - sagemaker_session: Optional[Any] = None, + model_version: str | None = None, + region: str | None = None, + initial_instance_count: int | None = None, + instance_type: str | None = None, + serializer: Any | None = None, + deserializer: Any | None = None, + accelerator_type: str | None = None, + endpoint_name: str | None = None, + tags: list[dict[str, str]] | None = None, + kms_key: str | None = None, + wait: bool | None = None, + data_capture_config: Any | None = None, + async_inference_config: Any | None = None, + serverless_inference_config: Any | None = None, + volume_size: int | None = None, + model_data_download_timeout: int | None = None, + container_startup_health_check_timeout: int | None = None, + inference_recommendation_id: str | None = None, + explainer_config: Any | None = None, + image_uri: str | Any | None = None, + role: str | None = None, + predictor_cls: callable | None = None, + env: dict[str, str | Any] | None = None, + model_name: str | None = None, + vpc_config: dict[str, list[str | Any]] | None = None, + sagemaker_session: Any | None = None, enable_network_isolation: bool | Any = None, - model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, str | Any]] = None, - source_dir: Optional[str] = None, - code_location: Optional[str] = None, - entry_point: Optional[str] = None, - container_log_level: Optional[int | Any] = None, - dependencies: Optional[List[str]] = None, - git_config: Optional[Dict[str, str]] = None, - tolerate_deprecated_model: Optional[bool] = None, - tolerate_vulnerable_model: Optional[bool] = None, + model_kms_key: str | None = None, + image_config: dict[str, str | Any] | None = None, + source_dir: str | None = None, + code_location: str | None = None, + entry_point: str | None = None, + container_log_level: int | Any | None = None, + dependencies: list[str] | None = None, + git_config: dict[str, str] | None = None, + tolerate_deprecated_model: bool | None = None, + tolerate_vulnerable_model: bool | None = None, use_compiled_model: bool = False, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/utils.pyi b/stubs/sagemaker/sagemaker/jumpstart/utils.pyi index 7127090e6eb4..d14dae7de992 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/utils.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/utils.pyi @@ -9,45 +9,45 @@ LOGGER: Incomplete def get_jumpstart_launched_regions_message() -> str: ... def get_jumpstart_content_bucket(region: str = "eu-west-1") -> str: ... -def get_formatted_manifest(manifest: List[Dict]) -> Dict[JumpStartVersionedModelId, JumpStartModelHeader]: ... +def get_formatted_manifest(manifest: list[dict]) -> dict[JumpStartVersionedModelId, JumpStartModelHeader]: ... def get_sagemaker_version() -> str: ... def parse_sagemaker_version() -> str: ... -def is_jumpstart_model_input(model_id: Optional[str], version: Optional[str]) -> bool: ... -def is_jumpstart_model_uri(uri: Optional[str]) -> bool: ... -def tag_key_in_array(tag_key: str, tag_array: List[Dict[str, str]]) -> bool: ... -def get_tag_value(tag_key: str, tag_array: List[Dict[str, str]]) -> str: ... +def is_jumpstart_model_input(model_id: str | None, version: str | None) -> bool: ... +def is_jumpstart_model_uri(uri: str | None) -> bool: ... +def tag_key_in_array(tag_key: str, tag_array: list[dict[str, str]]) -> bool: ... +def get_tag_value(tag_key: str, tag_array: list[dict[str, str]]) -> str: ... def add_single_jumpstart_tag( - uri: str, tag_key: enums.JumpStartTag, curr_tags: Optional[List[Dict[str, str]]] -) -> Optional[List]: ... -def get_jumpstart_base_name_if_jumpstart_model(*uris: Optional[str]) -> Optional[str]: ... + uri: str, tag_key: enums.JumpStartTag, curr_tags: list[dict[str, str]] | None +) -> list | None: ... +def get_jumpstart_base_name_if_jumpstart_model(*uris: str | None) -> str | None: ... def add_jumpstart_tags( - tags: Optional[List[Dict[str, str]]] = None, - inference_model_uri: Optional[str] = None, - inference_script_uri: Optional[str] = None, - training_model_uri: Optional[str] = None, - training_script_uri: Optional[str] = None, -) -> Optional[List[Dict[str, str]]]: ... + tags: list[dict[str, str]] | None = None, + inference_model_uri: str | None = None, + inference_script_uri: str | None = None, + training_model_uri: str | None = None, + training_script_uri: str | None = None, +) -> list[dict[str, str]] | None: ... def update_inference_tags_with_jumpstart_training_tags( - inference_tags: Optional[List[Dict[str, str]]], training_tags: Optional[List[Dict[str, str]]] -) -> Optional[List[Dict[str, str]]]: ... + inference_tags: list[dict[str, str]] | None, training_tags: list[dict[str, str]] | None +) -> list[dict[str, str]] | None: ... def verify_model_region_and_return_specs( - model_id: Optional[str], - version: Optional[str], - scope: Optional[str], + model_id: str | None, + version: str | None, + scope: str | None, region: str, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, ) -> JumpStartModelSpecs: ... def update_dict_if_key_not_present(dict_to_update: dict, key_to_add: Any, value_to_add: Any) -> dict: ... def resolve_model_sagemaker_config_field( - field_name: str, field_val: Optional[Any], sagemaker_session: Session, default_value: Optional[str] = None + field_name: str, field_val: Any | None, sagemaker_session: Session, default_value: str | None = None ) -> Any: ... def resolve_estimator_sagemaker_config_field( - field_name: str, field_val: Optional[Any], sagemaker_session: Session, default_value: Optional[str] = None + field_name: str, field_val: Any | None, sagemaker_session: Session, default_value: str | None = None ) -> Any: ... def is_valid_model_id( - model_id: Optional[str], - region: Optional[str] = None, - model_version: Optional[str] = None, + model_id: str | None, + region: str | None = None, + model_version: str | None = None, script: enums.JumpStartScriptScope = ..., ) -> bool: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/validators.pyi b/stubs/sagemaker/sagemaker/jumpstart/validators.pyi index 308276c81650..36c366b6b0c3 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/validators.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/validators.pyi @@ -5,7 +5,7 @@ from sagemaker.jumpstart.enums import HyperparameterValidationMode def validate_hyperparameters( model_id: str, model_version: str, - hyperparameters: Dict[str, Any], + hyperparameters: dict[str, Any], validation_mode: HyperparameterValidationMode = ..., - region: Optional[str] = "eu-west-1", + region: str | None = "eu-west-1", ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/lambda_helper.pyi b/stubs/sagemaker/sagemaker/lambda_helper.pyi index e3cdbe4a068e..7dee30dd812a 100644 --- a/stubs/sagemaker/sagemaker/lambda_helper.pyi +++ b/stubs/sagemaker/sagemaker/lambda_helper.pyi @@ -19,20 +19,20 @@ class Lambda: layers: Incomplete def __init__( self, - function_arn: str = None, - function_name: str = None, - execution_role_arn: str = None, - zipped_code_dir: str = None, - s3_bucket: str = None, - script: str = None, - handler: str = None, - session: Session = None, + function_arn: str | None = None, + function_name: str | None = None, + execution_role_arn: str | None = None, + zipped_code_dir: str | None = None, + s3_bucket: str | None = None, + script: str | None = None, + handler: str | None = None, + session: Session | None = None, timeout: int = 120, memory_size: int = 128, runtime: str = "python3.8", - vpc_config: dict = None, - environment: dict = None, - layers: list = None, + vpc_config: dict | None = None, + environment: dict | None = None, + layers: list | None = None, ) -> None: ... def create(self): ... def update(self): ... diff --git a/stubs/sagemaker/sagemaker/lineage/action.pyi b/stubs/sagemaker/sagemaker/lineage/action.pyi index c03020e8a82b..537b8254a365 100644 --- a/stubs/sagemaker/sagemaker/lineage/action.pyi +++ b/stubs/sagemaker/sagemaker/lineage/action.pyi @@ -1,6 +1,7 @@ from _typeshed import Incomplete from datetime import datetime -from typing import Iterator, List, Optional +from typing import Optional +from collections.abc import Iterator from sagemaker.apiutils import _base_types from sagemaker.lineage._api_types import ActionSource, ActionSummary @@ -31,32 +32,32 @@ class Action(_base_types.Record): @classmethod def create( cls, - action_name: str = None, - source_uri: str = None, - source_type: str = None, - action_type: str = None, - description: str = None, - status: str = None, - properties: dict = None, - tags: dict = None, - sagemaker_session: Session = None, + action_name: str | None = None, + source_uri: str | None = None, + source_type: str | None = None, + action_type: str | None = None, + description: str | None = None, + status: str | None = None, + properties: dict | None = None, + tags: dict | None = None, + sagemaker_session: Session | None = None, ) -> Action: ... @classmethod def list( cls, - source_uri: Optional[str] = None, - action_type: Optional[str] = None, - created_after: Optional[datetime] = None, - created_before: Optional[datetime] = None, - sort_by: Optional[str] = None, - sort_order: Optional[str] = None, - sagemaker_session: Session = None, - max_results: Optional[int] = None, - next_token: Optional[str] = None, + source_uri: str | None = None, + action_type: str | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + sort_by: str | None = None, + sort_order: str | None = None, + sagemaker_session: Session | None = None, + max_results: int | None = None, + next_token: str | None = None, ) -> Iterator[ActionSummary]: ... - def artifacts(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... + def artifacts(self, direction: LineageQueryDirectionEnum = ...) -> list[Artifact]: ... class ModelPackageApprovalAction(Action): - def datasets(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... + def datasets(self, direction: LineageQueryDirectionEnum = ...) -> list[Artifact]: ... def model_package(self): ... - def endpoints(self, direction: LineageQueryDirectionEnum = ...) -> List: ... + def endpoints(self, direction: LineageQueryDirectionEnum = ...) -> list: ... diff --git a/stubs/sagemaker/sagemaker/lineage/artifact.pyi b/stubs/sagemaker/sagemaker/lineage/artifact.pyi index 85303057fbf3..a1bb108e6b72 100644 --- a/stubs/sagemaker/sagemaker/lineage/artifact.pyi +++ b/stubs/sagemaker/sagemaker/lineage/artifact.pyi @@ -1,6 +1,7 @@ from _typeshed import Incomplete from datetime import datetime -from typing import Iterator, List, Optional +from typing import Optional +from collections.abc import Iterator from sagemaker.apiutils import _base_types from sagemaker.lineage._api_types import ArtifactSource, ArtifactSummary @@ -27,47 +28,47 @@ class Artifact(_base_types.Record): def load(cls, artifact_arn: str, sagemaker_session: Incomplete | None = None) -> Artifact: ... def downstream_trials(self, sagemaker_session: Incomplete | None = None) -> list: ... def downstream_trials_v2(self) -> list: ... - def upstream_trials(self) -> List: ... + def upstream_trials(self) -> list: ... def set_tag(self, tag: Incomplete | None = None): ... def set_tags(self, tags: Incomplete | None = None): ... @classmethod def create( cls, - artifact_name: Optional[str] = None, - source_uri: Optional[str] = None, - source_types: Optional[list] = None, - artifact_type: Optional[str] = None, - properties: Optional[dict] = None, - tags: Optional[dict] = None, + artifact_name: str | None = None, + source_uri: str | None = None, + source_types: list | None = None, + artifact_type: str | None = None, + properties: dict | None = None, + tags: dict | None = None, sagemaker_session: Incomplete | None = None, ) -> Artifact: ... @classmethod def list( cls, - source_uri: Optional[str] = None, - artifact_type: Optional[str] = None, - created_before: Optional[datetime] = None, - created_after: Optional[datetime] = None, - sort_by: Optional[str] = None, - sort_order: Optional[str] = None, - max_results: Optional[int] = None, - next_token: Optional[str] = None, + source_uri: str | None = None, + artifact_type: str | None = None, + created_before: datetime | None = None, + created_after: datetime | None = None, + sort_by: str | None = None, + sort_order: str | None = None, + max_results: int | None = None, + next_token: str | None = None, sagemaker_session: Incomplete | None = None, ) -> Iterator[ArtifactSummary]: ... def s3_uri_artifacts(self, s3_uri: str) -> dict: ... class ModelArtifact(Artifact): def endpoints(self) -> list: ... - def endpoint_contexts(self, direction: LineageQueryDirectionEnum = ...) -> List[Context]: ... - def dataset_artifacts(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... - def training_job_arns(self, direction: LineageQueryDirectionEnum = ...) -> List[str]: ... + def endpoint_contexts(self, direction: LineageQueryDirectionEnum = ...) -> list[Context]: ... + def dataset_artifacts(self, direction: LineageQueryDirectionEnum = ...) -> list[Artifact]: ... + def training_job_arns(self, direction: LineageQueryDirectionEnum = ...) -> list[str]: ... def pipeline_execution_arn(self, direction: LineageQueryDirectionEnum = ...) -> str: ... class DatasetArtifact(Artifact): - def trained_models(self) -> List[Association]: ... - def endpoint_contexts(self, direction: LineageQueryDirectionEnum = ...) -> List[Context]: ... - def upstream_datasets(self) -> List[Artifact]: ... - def downstream_datasets(self) -> List[Artifact]: ... + def trained_models(self) -> list[Association]: ... + def endpoint_contexts(self, direction: LineageQueryDirectionEnum = ...) -> list[Context]: ... + def upstream_datasets(self) -> list[Artifact]: ... + def downstream_datasets(self) -> list[Artifact]: ... class ImageArtifact(Artifact): - def datasets(self, direction: LineageQueryDirectionEnum) -> List[Artifact]: ... + def datasets(self, direction: LineageQueryDirectionEnum) -> list[Artifact]: ... diff --git a/stubs/sagemaker/sagemaker/lineage/association.pyi b/stubs/sagemaker/sagemaker/lineage/association.pyi index 8cd92310e20d..ef758e2fc502 100644 --- a/stubs/sagemaker/sagemaker/lineage/association.pyi +++ b/stubs/sagemaker/sagemaker/lineage/association.pyi @@ -1,6 +1,7 @@ from _typeshed import Incomplete from datetime import datetime -from typing import Iterator, Optional +from typing import Optional +from collections.abc import Iterator from sagemaker.apiutils import _base_types from sagemaker.lineage._api_types import AssociationSummary @@ -15,21 +16,21 @@ class Association(_base_types.Record): def set_tags(self, tags: Incomplete | None = None): ... @classmethod def create( - cls, source_arn: str, destination_arn: str, association_type: str = None, sagemaker_session: Incomplete | None = None + cls, source_arn: str, destination_arn: str, association_type: str | None = None, sagemaker_session: Incomplete | None = None ) -> Association: ... @classmethod def list( cls, - source_arn: str = None, - destination_arn: str = None, - source_type: str = None, - destination_type: str = None, - association_type: str = None, - created_after: Optional[datetime] = None, - created_before: Optional[datetime] = None, - sort_by: Optional[str] = None, - sort_order: Optional[str] = None, - max_results: Optional[int] = None, - next_token: Optional[str] = None, + source_arn: str | None = None, + destination_arn: str | None = None, + source_type: str | None = None, + destination_type: str | None = None, + association_type: str | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + sort_by: str | None = None, + sort_order: str | None = None, + max_results: int | None = None, + next_token: str | None = None, sagemaker_session: Incomplete | None = None, ) -> Iterator[AssociationSummary]: ... diff --git a/stubs/sagemaker/sagemaker/lineage/context.pyi b/stubs/sagemaker/sagemaker/lineage/context.pyi index 8dbd7ce06291..180c3e015837 100644 --- a/stubs/sagemaker/sagemaker/lineage/context.pyi +++ b/stubs/sagemaker/sagemaker/lineage/context.pyi @@ -1,6 +1,7 @@ from _typeshed import Incomplete from datetime import datetime -from typing import Iterator, List, Optional +from typing import Optional +from collections.abc import Iterator from sagemaker.apiutils import _base_types from sagemaker.lineage import association @@ -29,38 +30,38 @@ class Context(_base_types.Record): @classmethod def create( cls, - context_name: str = None, - source_uri: str = None, - source_type: str = None, - context_type: str = None, - description: str = None, - properties: dict = None, - tags: dict = None, + context_name: str | None = None, + source_uri: str | None = None, + source_type: str | None = None, + context_type: str | None = None, + description: str | None = None, + properties: dict | None = None, + tags: dict | None = None, sagemaker_session: Incomplete | None = None, ) -> Context: ... @classmethod def list( cls, - source_uri: Optional[str] = None, - context_type: Optional[str] = None, - created_after: Optional[datetime] = None, - created_before: Optional[datetime] = None, - sort_by: Optional[str] = None, - sort_order: Optional[str] = None, - max_results: Optional[int] = None, - next_token: Optional[str] = None, + source_uri: str | None = None, + context_type: str | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + sort_by: str | None = None, + sort_order: str | None = None, + max_results: int | None = None, + next_token: str | None = None, sagemaker_session: Incomplete | None = None, ) -> Iterator[ContextSummary]: ... - def actions(self, direction: LineageQueryDirectionEnum) -> List[Action]: ... + def actions(self, direction: LineageQueryDirectionEnum) -> list[Action]: ... class EndpointContext(Context): - def models(self) -> List[association.Association]: ... - def models_v2(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... - def dataset_artifacts(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... - def training_job_arns(self, direction: LineageQueryDirectionEnum = ...) -> List[str]: ... - def processing_jobs(self, direction: LineageQueryDirectionEnum = ...) -> List[LineageTrialComponent]: ... - def transform_jobs(self, direction: LineageQueryDirectionEnum = ...) -> List[LineageTrialComponent]: ... - def trial_components(self, direction: LineageQueryDirectionEnum = ...) -> List[LineageTrialComponent]: ... + def models(self) -> list[association.Association]: ... + def models_v2(self, direction: LineageQueryDirectionEnum = ...) -> list[Artifact]: ... + def dataset_artifacts(self, direction: LineageQueryDirectionEnum = ...) -> list[Artifact]: ... + def training_job_arns(self, direction: LineageQueryDirectionEnum = ...) -> list[str]: ... + def processing_jobs(self, direction: LineageQueryDirectionEnum = ...) -> list[LineageTrialComponent]: ... + def transform_jobs(self, direction: LineageQueryDirectionEnum = ...) -> list[LineageTrialComponent]: ... + def trial_components(self, direction: LineageQueryDirectionEnum = ...) -> list[LineageTrialComponent]: ... def pipeline_execution_arn(self, direction: LineageQueryDirectionEnum = ...) -> str: ... class ModelPackageGroup(Context): diff --git a/stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi b/stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi index f7cec0b6ecad..d78438f5beb6 100644 --- a/stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi +++ b/stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi @@ -30,5 +30,5 @@ class LineageTrialComponent(_base_types.Record): @classmethod def load(cls, trial_component_name: str, sagemaker_session: Incomplete | None = None) -> LineageTrialComponent: ... def pipeline_execution_arn(self) -> str: ... - def dataset_artifacts(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... - def models(self, direction: LineageQueryDirectionEnum = ...) -> List[Artifact]: ... + def dataset_artifacts(self, direction: LineageQueryDirectionEnum = ...) -> list[Artifact]: ... + def models(self, direction: LineageQueryDirectionEnum = ...) -> list[Artifact]: ... diff --git a/stubs/sagemaker/sagemaker/lineage/query.pyi b/stubs/sagemaker/sagemaker/lineage/query.pyi index 024814f7f2a8..2e22c56382df 100644 --- a/stubs/sagemaker/sagemaker/lineage/query.pyi +++ b/stubs/sagemaker/sagemaker/lineage/query.pyi @@ -50,15 +50,15 @@ class Vertex: class PyvisVisualizer: graph_styles: Incomplete - def __init__(self, graph_styles, pyvis_options: Optional[Dict[str, Any]] = None) -> None: ... + def __init__(self, graph_styles, pyvis_options: dict[str, Any] | None = None) -> None: ... def render(self, elements, path: str = "lineage_graph_pyvis.html"): ... class LineageQueryResult: edges: Incomplete vertices: Incomplete startarn: Incomplete - def __init__(self, edges: List[Edge] = None, vertices: List[Vertex] = None, startarn: List[str] = None) -> None: ... - def visualize(self, path: Optional[str] = "lineage_graph_pyvis.html", pyvis_options: Optional[Dict[str, Any]] = None): ... + def __init__(self, edges: list[Edge] | None = None, vertices: list[Vertex] | None = None, startarn: list[str] | None = None) -> None: ... + def visualize(self, path: str | None = "lineage_graph_pyvis.html", pyvis_options: dict[str, Any] | None = None): ... class LineageFilter: entities: Incomplete @@ -70,22 +70,22 @@ class LineageFilter: properties: Incomplete def __init__( self, - entities: Optional[List[LineageEntityEnum | str]] = None, - sources: Optional[List[LineageSourceEnum | str]] = None, - created_before: Optional[datetime] = None, - created_after: Optional[datetime] = None, - modified_before: Optional[datetime] = None, - modified_after: Optional[datetime] = None, - properties: Optional[Dict[str, str]] = None, + entities: list[LineageEntityEnum | str] | None = None, + sources: list[LineageSourceEnum | str] | None = None, + created_before: datetime | None = None, + created_after: datetime | None = None, + modified_before: datetime | None = None, + modified_after: datetime | None = None, + properties: dict[str, str] | None = None, ) -> None: ... class LineageQuery: def __init__(self, sagemaker_session) -> None: ... def query( self, - start_arns: List[str], + start_arns: list[str], direction: LineageQueryDirectionEnum = ..., include_edges: bool = True, - query_filter: LineageFilter = None, + query_filter: LineageFilter | None = None, max_depth: int = 10, ) -> LineageQueryResult: ... diff --git a/stubs/sagemaker/sagemaker/lineage/visualizer.pyi b/stubs/sagemaker/sagemaker/lineage/visualizer.pyi index 5b330b12238d..f24bb67a39ec 100644 --- a/stubs/sagemaker/sagemaker/lineage/visualizer.pyi +++ b/stubs/sagemaker/sagemaker/lineage/visualizer.pyi @@ -6,13 +6,13 @@ class LineageTableVisualizer: def __init__(self, sagemaker_session) -> None: ... def show( self, - trial_component_name: Optional[str] = None, - training_job_name: Optional[str] = None, - processing_job_name: Optional[str] = None, - pipeline_execution_step: Optional[object] = None, - model_package_arn: Optional[str] = None, - endpoint_arn: Optional[str] = None, - artifact_arn: Optional[str] = None, - context_arn: Optional[str] = None, - actions_arn: Optional[str] = None, + trial_component_name: str | None = None, + training_job_name: str | None = None, + processing_job_name: str | None = None, + pipeline_execution_step: object | None = None, + model_package_arn: str | None = None, + endpoint_arn: str | None = None, + artifact_arn: str | None = None, + context_arn: str | None = None, + actions_arn: str | None = None, ) -> DataFrame: ... diff --git a/stubs/sagemaker/sagemaker/local/local_session.pyi b/stubs/sagemaker/sagemaker/local/local_session.pyi index 24219e25563f..3175cff45bb6 100644 --- a/stubs/sagemaker/sagemaker/local/local_session.pyi +++ b/stubs/sagemaker/sagemaker/local/local_session.pyi @@ -74,7 +74,7 @@ class LocalSession(Session): default_bucket: Incomplete | None = None, s3_endpoint_url: Incomplete | None = None, disable_local_code: bool = False, - sagemaker_config: dict = None, + sagemaker_config: dict | None = None, default_bucket_prefix: Incomplete | None = None, ) -> None: ... def logs_for_job(self, job_name, wait: bool = False, poll: int = 5, log_type: str = "All") -> None: ... diff --git a/stubs/sagemaker/sagemaker/local/pipeline.pyi b/stubs/sagemaker/sagemaker/local/pipeline.pyi index 983f3c699680..b1ee84ca1901 100644 --- a/stubs/sagemaker/sagemaker/local/pipeline.pyi +++ b/stubs/sagemaker/sagemaker/local/pipeline.pyi @@ -23,7 +23,7 @@ class _StepExecutor(ABC, metaclass=abc.ABCMeta): step: Incomplete def __init__(self, pipeline_executor: LocalPipelineExecutor, step: Step) -> None: ... @abstractmethod - def execute(self) -> Dict: ... + def execute(self) -> dict: ... class _TrainingStepExecutor(_StepExecutor): def execute(self): ... diff --git a/stubs/sagemaker/sagemaker/metadata_properties.pyi b/stubs/sagemaker/sagemaker/metadata_properties.pyi index 405816373717..c63450183f4c 100644 --- a/stubs/sagemaker/sagemaker/metadata_properties.pyi +++ b/stubs/sagemaker/sagemaker/metadata_properties.pyi @@ -10,8 +10,8 @@ class MetadataProperties: project_id: Incomplete def __init__( self, - commit_id: Optional[str | PipelineVariable] = None, - repository: Optional[str | PipelineVariable] = None, - generated_by: Optional[str | PipelineVariable] = None, - project_id: Optional[str | PipelineVariable] = None, + commit_id: str | PipelineVariable | None = None, + repository: str | PipelineVariable | None = None, + generated_by: str | PipelineVariable | None = None, + project_id: str | PipelineVariable | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/metric_definitions.pyi b/stubs/sagemaker/sagemaker/metric_definitions.pyi index fb5b40136bed..a0d688803871 100644 --- a/stubs/sagemaker/sagemaker/metric_definitions.pyi +++ b/stubs/sagemaker/sagemaker/metric_definitions.pyi @@ -4,9 +4,9 @@ from typing import Dict, List, Optional logger: Incomplete def retrieve_default( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, -) -> Optional[List[Dict[str, str]]]: ... +) -> list[dict[str, str]] | None: ... diff --git a/stubs/sagemaker/sagemaker/model.pyi b/stubs/sagemaker/sagemaker/model.pyi index 9b978abc36e1..125b98a886f4 100644 --- a/stubs/sagemaker/sagemaker/model.pyi +++ b/stubs/sagemaker/sagemaker/model.pyi @@ -54,54 +54,54 @@ class Model(ModelBase, InferenceRecommenderMixin): def __init__( self, image_uri: str | PipelineVariable, - model_data: Optional[str | PipelineVariable] = None, - role: Optional[str] = None, - predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[str | PipelineVariable]]] = None, - sagemaker_session: Optional[Session] = None, - enable_network_isolation: bool | PipelineVariable = None, - model_kms_key: Optional[str] = None, - image_config: Optional[Dict[str, str | PipelineVariable]] = None, - source_dir: Optional[str] = None, - code_location: Optional[str] = None, - entry_point: Optional[str] = None, + model_data: str | PipelineVariable | None = None, + role: str | None = None, + predictor_cls: callable | None = None, + env: dict[str, str | PipelineVariable] | None = None, + name: str | None = None, + vpc_config: dict[str, list[str | PipelineVariable]] | None = None, + sagemaker_session: Session | None = None, + enable_network_isolation: bool | PipelineVariable | None = None, + model_kms_key: str | None = None, + image_config: dict[str, str | PipelineVariable] | None = None, + source_dir: str | None = None, + code_location: str | None = None, + entry_point: str | None = None, container_log_level: int | PipelineVariable = 20, - dependencies: Optional[List[str]] = None, - git_config: Optional[Dict[str, str]] = None, + dependencies: list[str] | None = None, + git_config: dict[str, str] | None = None, ) -> None: ... def register( self, - content_types: List[str | PipelineVariable], - response_types: List[str | PipelineVariable], - inference_instances: Optional[List[str | PipelineVariable]] = None, - transform_instances: Optional[List[str | PipelineVariable]] = None, - model_package_name: Optional[str | PipelineVariable] = None, - model_package_group_name: Optional[str | PipelineVariable] = None, - image_uri: Optional[str | PipelineVariable] = None, - model_metrics: Optional[ModelMetrics] = None, - metadata_properties: Optional[MetadataProperties] = None, + content_types: list[str | PipelineVariable], + response_types: list[str | PipelineVariable], + inference_instances: list[str | PipelineVariable] | None = None, + transform_instances: list[str | PipelineVariable] | None = None, + model_package_name: str | PipelineVariable | None = None, + model_package_group_name: str | PipelineVariable | None = None, + image_uri: str | PipelineVariable | None = None, + model_metrics: ModelMetrics | None = None, + metadata_properties: MetadataProperties | None = None, marketplace_cert: bool = False, - approval_status: Optional[str | PipelineVariable] = None, - description: Optional[str] = None, - drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, - validation_specification: Optional[str | PipelineVariable] = None, - domain: Optional[str | PipelineVariable] = None, - task: Optional[str | PipelineVariable] = None, - sample_payload_url: Optional[str | PipelineVariable] = None, - framework: Optional[str | PipelineVariable] = None, - framework_version: Optional[str | PipelineVariable] = None, - nearest_model_name: Optional[str | PipelineVariable] = None, - data_input_configuration: Optional[str | PipelineVariable] = None, + approval_status: str | PipelineVariable | None = None, + description: str | None = None, + drift_check_baselines: DriftCheckBaselines | None = None, + customer_metadata_properties: dict[str, str | PipelineVariable] | None = None, + validation_specification: str | PipelineVariable | None = None, + domain: str | PipelineVariable | None = None, + task: str | PipelineVariable | None = None, + sample_payload_url: str | PipelineVariable | None = None, + framework: str | PipelineVariable | None = None, + framework_version: str | PipelineVariable | None = None, + nearest_model_name: str | PipelineVariable | None = None, + data_input_configuration: str | PipelineVariable | None = None, ): ... def create( self, - instance_type: Optional[str] = None, - accelerator_type: Optional[str] = None, - serverless_inference_config: Optional[ServerlessInferenceConfig] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, + instance_type: str | None = None, + accelerator_type: str | None = None, + serverless_inference_config: ServerlessInferenceConfig | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, ): ... def prepare_container_def( self, @@ -180,17 +180,17 @@ class FrameworkModel(Model): self, model_data: str | PipelineVariable, image_uri: str | PipelineVariable, - role: Optional[str] = None, - entry_point: Optional[str] = None, - source_dir: Optional[str] = None, - predictor_cls: Optional[callable] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - name: Optional[str] = None, + role: str | None = None, + entry_point: str | None = None, + source_dir: str | None = None, + predictor_cls: callable | None = None, + env: dict[str, str | PipelineVariable] | None = None, + name: str | None = None, container_log_level: int | PipelineVariable = 20, - code_location: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - dependencies: Optional[List[str]] = None, - git_config: Optional[Dict[str, str]] = None, + code_location: str | None = None, + sagemaker_session: Session | None = None, + dependencies: list[str] | None = None, + git_config: dict[str, str] | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/model_card/helpers.pyi b/stubs/sagemaker/sagemaker/model_card/helpers.pyi index ffbfbf756384..815343a726f4 100644 --- a/stubs/sagemaker/sagemaker/model_card/helpers.pyi +++ b/stubs/sagemaker/sagemaker/model_card/helpers.pyi @@ -4,7 +4,7 @@ import json from _typeshed import Incomplete from abc import ABC, abstractmethod from enum import Enum -from typing import Any, List, Optional +from typing import Any, Optional from boto3.session import Session as Session @@ -19,7 +19,7 @@ class _DefaultFromDict: ... class _DescriptorBase(ABC, metaclass=abc.ABCMeta): private_name: Incomplete def __set_name__(self, owner: type, name: str): ... - def __get__(self, obj, objtype: type = None): ... + def __get__(self, obj, objtype: type | None = None): ... def __set__(self, obj: object, value: Any): ... @abstractmethod def validate(self, value): ... @@ -40,10 +40,10 @@ class _OneOf(_DescriptorBase): class _IsList(_DescriptorBase): item_type: Incomplete max_size: Incomplete - def __init__(self, item_type: object, max_size: Optional[int] = None) -> None: ... - def validate(self, value: List): ... - def require_decode(self, value: List): ... - def decode(self, value: List): ... + def __init__(self, item_type: object, max_size: int | None = None) -> None: ... + def validate(self, value: list): ... + def require_decode(self, value: list): ... + def decode(self, value: list): ... class _IsModelCardObject(_DescriptorBase): custom_class: Incomplete @@ -54,7 +54,7 @@ class _IsModelCardObject(_DescriptorBase): class _MaxSizeArray(collections.abc.MutableSequence): list: Incomplete - def __init__(self, max_size: int, item_type: Any, array: List = None) -> None: ... + def __init__(self, max_size: int, item_type: Any, array: list | None = None) -> None: ... def __len__(self) -> int: ... def __getitem__(self, index): ... def __delitem__(self, index) -> None: ... diff --git a/stubs/sagemaker/sagemaker/model_card/model_card.pyi b/stubs/sagemaker/sagemaker/model_card/model_card.pyi index fd5c9203c676..0e23f9645505 100644 --- a/stubs/sagemaker/sagemaker/model_card/model_card.pyi +++ b/stubs/sagemaker/sagemaker/model_card/model_card.pyi @@ -18,7 +18,7 @@ logger: Incomplete class Environment(_DefaultToRequestDict, _DefaultFromDict): container_image: Incomplete - def __init__(self, container_image: List[str]) -> None: ... + def __init__(self, container_image: list[str]) -> None: ... class ModelOverview(_DefaultToRequestDict, _DefaultFromDict): model_artifact: Incomplete @@ -33,19 +33,19 @@ class ModelOverview(_DefaultToRequestDict, _DefaultFromDict): model_owner: Incomplete def __init__( self, - model_id: Optional[str] = None, - model_name: Optional[str] = None, - model_description: Optional[str] = None, - model_version: Optional[int | float] = None, - problem_type: Optional[str] = None, - algorithm_type: Optional[str] = None, - model_creator: Optional[str] = None, - model_owner: Optional[str] = None, - model_artifact: Optional[List[str]] = None, - inference_environment: Optional[Environment] = None, + model_id: str | None = None, + model_name: str | None = None, + model_description: str | None = None, + model_version: int | float | None = None, + problem_type: str | None = None, + algorithm_type: str | None = None, + model_creator: str | None = None, + model_owner: str | None = None, + model_artifact: list[str] | None = None, + inference_environment: Environment | None = None, ) -> None: ... @classmethod - def from_model_name(cls, model_name: str, sagemaker_session: Session = None, **kwargs): ... + def from_model_name(cls, model_name: str, sagemaker_session: Session | None = None, **kwargs): ... class IntendedUses(_DefaultToRequestDict, _DefaultFromDict): risk_rating: Incomplete @@ -55,11 +55,11 @@ class IntendedUses(_DefaultToRequestDict, _DefaultFromDict): explanations_for_risk_rating: Incomplete def __init__( self, - purpose_of_model: Optional[str] = None, - intended_uses: Optional[str] = None, - factors_affecting_model_efficiency: Optional[str] = None, - risk_rating: Optional[RiskRatingEnum | str] = ..., - explanations_for_risk_rating: Optional[str] = None, + purpose_of_model: str | None = None, + intended_uses: str | None = None, + factors_affecting_model_efficiency: str | None = None, + risk_rating: RiskRatingEnum | str | None = ..., + explanations_for_risk_rating: str | None = None, ) -> None: ... class BusinessDetails(_DefaultToRequestDict, _DefaultFromDict): @@ -68,9 +68,9 @@ class BusinessDetails(_DefaultToRequestDict, _DefaultFromDict): line_of_business: Incomplete def __init__( self, - business_problem: Optional[str] = None, - business_stakeholders: Optional[str] = None, - line_of_business: Optional[str] = None, + business_problem: str | None = None, + business_stakeholders: str | None = None, + line_of_business: str | None = None, ) -> None: ... class Function(_DefaultToRequestDict, _DefaultFromDict): @@ -79,15 +79,15 @@ class Function(_DefaultToRequestDict, _DefaultFromDict): condition: Incomplete def __init__( self, - function: Optional[ObjectiveFunctionEnum | str] = None, - facet: Optional[FacetEnum | str] = None, - condition: Optional[str] = None, + function: ObjectiveFunctionEnum | str | None = None, + facet: FacetEnum | str | None = None, + condition: str | None = None, ) -> None: ... class ObjectiveFunction(_DefaultToRequestDict, _DefaultFromDict): function: Incomplete notes: Incomplete - def __init__(self, function: Function, notes: Optional[str] = None) -> None: ... + def __init__(self, function: Function, notes: str | None = None) -> None: ... class Metric(_DefaultToRequestDict, _DefaultFromDict): type: Incomplete @@ -99,21 +99,21 @@ class Metric(_DefaultToRequestDict, _DefaultFromDict): self, name: str, type: MetricTypeEnum | str, - value: int | float | str | bool | List, - notes: Optional[str] = None, - x_axis_name: Optional[str | list] = None, - y_axis_name: Optional[str | list] = None, + value: int | float | str | bool | list, + notes: str | None = None, + x_axis_name: str | list | None = None, + y_axis_name: str | list | None = None, ) -> None: ... @property def value(self): ... @value.setter - def value(self, val: int | float | str | bool | List): ... + def value(self, val: int | float | str | bool | list): ... class TrainingMetric(_DefaultToRequestDict, _DefaultFromDict): name: Incomplete value: Incomplete notes: Incomplete - def __init__(self, name: str, value: int | float, notes: Optional[str] = None) -> None: ... + def __init__(self, name: str, value: int | float, notes: str | None = None) -> None: ... class HyperParameter(_DefaultToRequestDict, _DefaultFromDict): name: Incomplete @@ -130,13 +130,13 @@ class TrainingJobDetails(_DefaultToRequestDict, _DefaultFromDict): training_arn: Incomplete def __init__( self, - training_arn: Optional[str] = None, - training_datasets: Optional[List[str]] = None, - training_environment: Optional[Environment] = None, - training_metrics: Optional[List[TrainingMetric]] = None, - user_provided_training_metrics: Optional[List[TrainingMetric]] = None, - hyper_parameters: Optional[List[HyperParameter]] = None, - user_provided_hyper_parameters: Optional[List[HyperParameter]] = None, + training_arn: str | None = None, + training_datasets: list[str] | None = None, + training_environment: Environment | None = None, + training_metrics: list[TrainingMetric] | None = None, + user_provided_training_metrics: list[TrainingMetric] | None = None, + hyper_parameters: list[HyperParameter] | None = None, + user_provided_hyper_parameters: list[HyperParameter] | None = None, ) -> None: ... class TrainingDetails(_DefaultToRequestDict, _DefaultFromDict): @@ -145,21 +145,21 @@ class TrainingDetails(_DefaultToRequestDict, _DefaultFromDict): training_observations: Incomplete def __init__( self, - objective_function: Optional[ObjectiveFunction] = None, - training_observations: Optional[str] = None, - training_job_details: Optional[TrainingJobDetails] = None, + objective_function: ObjectiveFunction | None = None, + training_observations: str | None = None, + training_job_details: TrainingJobDetails | None = None, ) -> None: ... @classmethod - def from_model_overview(cls, model_overview: ModelOverview, sagemaker_session: Session = None, **kwargs): ... + def from_model_overview(cls, model_overview: ModelOverview, sagemaker_session: Session | None = None, **kwargs): ... @classmethod - def from_training_job_name(cls, training_job_name: str, sagemaker_session: Session = None, **kwargs): ... + def from_training_job_name(cls, training_job_name: str, sagemaker_session: Session | None = None, **kwargs): ... def add_metric(self, metric: TrainingMetric): ... def add_parameter(self, parameter: HyperParameter): ... class MetricGroup(_DefaultToRequestDict, _DefaultFromDict): metric_data: Incomplete name: Incomplete - def __init__(self, name: str, metric_data: Optional[List[Metric]] = None) -> None: ... + def __init__(self, name: str, metric_data: list[Metric] | None = None) -> None: ... def add_metric(self, metric: Metric): ... class EvaluationJob(_DefaultToRequestDict, _DefaultFromDict): @@ -172,11 +172,11 @@ class EvaluationJob(_DefaultToRequestDict, _DefaultFromDict): def __init__( self, name: str, - evaluation_observation: Optional[str] = None, - evaluation_job_arn: Optional[str] = None, - datasets: Optional[List[str]] = None, - metadata: Optional[dict] = None, - metric_groups: Optional[List[MetricGroup]] = None, + evaluation_observation: str | None = None, + evaluation_job_arn: str | None = None, + datasets: list[str] | None = None, + metadata: dict | None = None, + metric_groups: list[MetricGroup] | None = None, ) -> None: ... def get_metric_group(self, group_name): ... def add_metric_group(self, group_name: str): ... @@ -189,9 +189,9 @@ class AdditionalInformation(_DefaultToRequestDict, _DefaultFromDict): custom_details: Incomplete def __init__( self, - ethical_considerations: Optional[str] = None, - caveats_and_recommendations: Optional[str] = None, - custom_details: Optional[dict] = None, + ethical_considerations: str | None = None, + caveats_and_recommendations: str | None = None, + custom_details: dict | None = None, ) -> None: ... class ModelCard: @@ -216,28 +216,28 @@ class ModelCard: def __init__( self, name: str, - status: Optional[ModelCardStatusEnum | str] = ..., - arn: Optional[str] = None, - version: Optional[int] = None, - created_time: Optional[datetime] = None, - created_by: Optional[dict] = None, - last_modified_time: Optional[datetime] = None, - last_modified_by: Optional[dict] = None, - model_overview: Optional[ModelOverview] = None, - intended_uses: Optional[IntendedUses] = None, - business_details: Optional[BusinessDetails] = None, - training_details: Optional[TrainingDetails] = None, - evaluation_details: Optional[List[EvaluationJob]] = None, - additional_information: Optional[AdditionalInformation] = None, - sagemaker_session: Optional[Session] = None, + status: ModelCardStatusEnum | str | None = ..., + arn: str | None = None, + version: int | None = None, + created_time: datetime | None = None, + created_by: dict | None = None, + last_modified_time: datetime | None = None, + last_modified_by: dict | None = None, + model_overview: ModelOverview | None = None, + intended_uses: IntendedUses | None = None, + business_details: BusinessDetails | None = None, + training_details: TrainingDetails | None = None, + evaluation_details: list[EvaluationJob] | None = None, + additional_information: AdditionalInformation | None = None, + sagemaker_session: Session | None = None, ) -> None: ... def create(self): ... @classmethod - def load(cls, name: str, version: Optional[int] = None, sagemaker_session: Session = None): ... + def load(cls, name: str, version: int | None = None, sagemaker_session: Session | None = None): ... def update(self, **kwargs): ... def delete(self): ... def export_pdf( - self, s3_output_path: str, export_job_name: Optional[str] = None, model_card_version: Optional[int] = None + self, s3_output_path: str, export_job_name: str | None = None, model_card_version: int | None = None ): ... def list_export_jobs(self, **kwargs): ... def get_version_history(self, **kwargs): ... @@ -259,14 +259,14 @@ class ModelCardExportJob: model_card_name: str, model_card_version: int, s3_output_path: str, - s3_export_artifacts: Optional[str] = None, - export_job_arn: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - status: Optional[str] = None, - failure_reason: Optional[str] = None, + s3_export_artifacts: str | None = None, + export_job_arn: str | None = None, + sagemaker_session: Session | None = None, + status: str | None = None, + failure_reason: str | None = None, ) -> None: ... def create(self): ... @classmethod - def load(cls, export_job_arn: str, sagemaker_session: Session = None): ... + def load(cls, export_job_arn: str, sagemaker_session: Session | None = None): ... @staticmethod - def list_export_jobs(model_card_name: str, sagemaker_session: Optional[Session] = None, **kwargs): ... + def list_export_jobs(model_card_name: str, sagemaker_session: Session | None = None, **kwargs): ... diff --git a/stubs/sagemaker/sagemaker/model_metrics.pyi b/stubs/sagemaker/sagemaker/model_metrics.pyi index 080278693967..0b0f54e13e67 100644 --- a/stubs/sagemaker/sagemaker/model_metrics.pyi +++ b/stubs/sagemaker/sagemaker/model_metrics.pyi @@ -32,7 +32,7 @@ class MetricsSource: self, content_type: str | PipelineVariable, s3_uri: str | PipelineVariable, - content_digest: Optional[str | PipelineVariable] = None, + content_digest: str | PipelineVariable | None = None, ) -> None: ... class FileSource: @@ -42,6 +42,6 @@ class FileSource: def __init__( self, s3_uri: str | PipelineVariable, - content_digest: Optional[str | PipelineVariable] = None, - content_type: Optional[str | PipelineVariable] = None, + content_digest: str | PipelineVariable | None = None, + content_type: str | PipelineVariable | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/model_monitor/data_quality_monitoring_config.pyi b/stubs/sagemaker/sagemaker/model_monitor/data_quality_monitoring_config.pyi index eebff2c9040a..3f1ed8c00526 100644 --- a/stubs/sagemaker/sagemaker/model_monitor/data_quality_monitoring_config.pyi +++ b/stubs/sagemaker/sagemaker/model_monitor/data_quality_monitoring_config.pyi @@ -5,7 +5,7 @@ L_INFINITY_METHOD: str class DataQualityDistributionConstraints: categorical_drift_method: Incomplete - def __init__(self, categorical_drift_method: str = None) -> None: ... + def __init__(self, categorical_drift_method: str | None = None) -> None: ... @staticmethod def valid_distribution_constraints(distribution_constraints): ... @staticmethod @@ -13,6 +13,6 @@ class DataQualityDistributionConstraints: class DataQualityMonitoringConfig: distribution_constraints: Incomplete - def __init__(self, distribution_constraints: DataQualityDistributionConstraints = None) -> None: ... + def __init__(self, distribution_constraints: DataQualityDistributionConstraints | None = None) -> None: ... @staticmethod def valid_monitoring_config(monitoring_config): ... diff --git a/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi b/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi index 3354c733b444..916a103623ad 100644 --- a/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi +++ b/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi @@ -100,19 +100,19 @@ class ModelMonitor: def describe_schedule(self): ... def list_executions(self): ... def update_monitoring_alert( - self, monitoring_alert_name: str, data_points_to_alert: Optional[int], evaluation_period: Optional[int] + self, monitoring_alert_name: str, data_points_to_alert: int | None, evaluation_period: int | None ): ... - def list_monitoring_alerts(self, next_token: Optional[str] = None, max_results: Optional[int] = 10): ... + def list_monitoring_alerts(self, next_token: str | None = None, max_results: int | None = 10): ... def list_monitoring_alert_history( self, - monitoring_alert_name: Optional[str] = None, - sort_by: Optional[str] = "CreationTime", - sort_order: Optional[str] = "Descending", - next_token: Optional[str] = None, - max_results: Optional[int] = 10, - creation_time_before: Optional[str] = None, - creation_time_after: Optional[str] = None, - status_equals: Optional[str] = None, + monitoring_alert_name: str | None = None, + sort_by: str | None = "CreationTime", + sort_order: str | None = "Descending", + next_token: str | None = None, + max_results: int | None = 10, + creation_time_before: str | None = None, + creation_time_after: str | None = None, + status_equals: str | None = None, ): ... @classmethod def attach(cls, monitor_schedule_name, sagemaker_session: Incomplete | None = None): ... @@ -367,12 +367,12 @@ class BatchTransformInput(MonitoringInput): dataset_format: MonitoringDatasetFormat, s3_input_mode: str = "File", s3_data_distribution_type: str = "FullyReplicated", - start_time_offset: str = None, - end_time_offset: str = None, - features_attribute: str = None, - inference_attribute: str = None, - probability_attribute: str = None, - probability_threshold_attribute: str = None, + start_time_offset: str | None = None, + end_time_offset: str | None = None, + features_attribute: str | None = None, + inference_attribute: str | None = None, + probability_attribute: str | None = None, + probability_threshold_attribute: str | None = None, ) -> None: ... class MonitoringOutput: diff --git a/stubs/sagemaker/sagemaker/model_uris.pyi b/stubs/sagemaker/sagemaker/model_uris.pyi index b3bcf06d516f..c938769f102d 100644 --- a/stubs/sagemaker/sagemaker/model_uris.pyi +++ b/stubs/sagemaker/sagemaker/model_uris.pyi @@ -4,10 +4,10 @@ from typing import Optional logger: Incomplete def retrieve( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, - model_scope: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, + model_scope: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, ) -> str: ... diff --git a/stubs/sagemaker/sagemaker/multidatamodel.pyi b/stubs/sagemaker/sagemaker/multidatamodel.pyi index 8f44b3357ab8..7e55a0211b16 100644 --- a/stubs/sagemaker/sagemaker/multidatamodel.pyi +++ b/stubs/sagemaker/sagemaker/multidatamodel.pyi @@ -20,10 +20,10 @@ class MultiDataModel(Model): self, name: str, model_data_prefix: str, - model: Optional[Model] = None, - image_uri: Optional[str | PipelineVariable] = None, - role: Optional[str] = None, - sagemaker_session: Optional[Session] = None, + model: Model | None = None, + image_uri: str | PipelineVariable | None = None, + role: str | None = None, + sagemaker_session: Session | None = None, **kwargs, ) -> None: ... def prepare_container_def( diff --git a/stubs/sagemaker/sagemaker/mxnet/estimator.pyi b/stubs/sagemaker/sagemaker/mxnet/estimator.pyi index eec5b37d8a19..dfba054dd5da 100644 --- a/stubs/sagemaker/sagemaker/mxnet/estimator.pyi +++ b/stubs/sagemaker/sagemaker/mxnet/estimator.pyi @@ -12,12 +12,12 @@ class MXNet(Framework): def __init__( self, entry_point: str | PipelineVariable, - framework_version: Optional[str] = None, - py_version: Optional[str] = None, - source_dir: Optional[str | PipelineVariable] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, - image_uri: Optional[str | PipelineVariable] = None, - distribution: Optional[Dict[str, str]] = None, + framework_version: str | None = None, + py_version: str | None = None, + source_dir: str | PipelineVariable | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, + image_uri: str | PipelineVariable | None = None, + distribution: dict[str, str] | None = None, **kwargs, ) -> None: ... def create_model( diff --git a/stubs/sagemaker/sagemaker/mxnet/model.pyi b/stubs/sagemaker/sagemaker/mxnet/model.pyi index 510167b2907a..121ade43a6ec 100644 --- a/stubs/sagemaker/sagemaker/mxnet/model.pyi +++ b/stubs/sagemaker/sagemaker/mxnet/model.pyi @@ -20,39 +20,39 @@ class MXNetModel(FrameworkModel): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - entry_point: Optional[str] = None, + role: str | None = None, + entry_point: str | None = None, framework_version: str = "1.4.0", - py_version: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, + py_version: str | None = None, + image_uri: str | PipelineVariable | None = None, predictor_cls: callable = ..., - model_server_workers: Optional[int | PipelineVariable] = None, + model_server_workers: int | PipelineVariable | None = None, **kwargs, ) -> None: ... image_uri: Incomplete def register( self, - content_types: List[str | PipelineVariable], - response_types: List[str | PipelineVariable], - inference_instances: Optional[List[str | PipelineVariable]] = None, - transform_instances: Optional[List[str | PipelineVariable]] = None, - model_package_name: Optional[str | PipelineVariable] = None, - model_package_group_name: Optional[str | PipelineVariable] = None, - image_uri: Optional[str | PipelineVariable] = None, - model_metrics: Optional[ModelMetrics] = None, - metadata_properties: Optional[MetadataProperties] = None, + content_types: list[str | PipelineVariable], + response_types: list[str | PipelineVariable], + inference_instances: list[str | PipelineVariable] | None = None, + transform_instances: list[str | PipelineVariable] | None = None, + model_package_name: str | PipelineVariable | None = None, + model_package_group_name: str | PipelineVariable | None = None, + image_uri: str | PipelineVariable | None = None, + model_metrics: ModelMetrics | None = None, + metadata_properties: MetadataProperties | None = None, marketplace_cert: bool = False, - approval_status: Optional[str | PipelineVariable] = None, - description: Optional[str] = None, - drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, - domain: Optional[str | PipelineVariable] = None, - sample_payload_url: Optional[str | PipelineVariable] = None, - task: Optional[str | PipelineVariable] = None, - framework: Optional[str | PipelineVariable] = None, - framework_version: Optional[str | PipelineVariable] = None, - nearest_model_name: Optional[str | PipelineVariable] = None, - data_input_configuration: Optional[str | PipelineVariable] = None, + approval_status: str | PipelineVariable | None = None, + description: str | None = None, + drift_check_baselines: DriftCheckBaselines | None = None, + customer_metadata_properties: dict[str, str | PipelineVariable] | None = None, + domain: str | PipelineVariable | None = None, + sample_payload_url: str | PipelineVariable | None = None, + task: str | PipelineVariable | None = None, + framework: str | PipelineVariable | None = None, + framework_version: str | PipelineVariable | None = None, + nearest_model_name: str | PipelineVariable | None = None, + data_input_configuration: str | PipelineVariable | None = None, ): ... def prepare_container_def( self, diff --git a/stubs/sagemaker/sagemaker/mxnet/processing.pyi b/stubs/sagemaker/sagemaker/mxnet/processing.pyi index ca8359d6de5f..6112062df0cd 100644 --- a/stubs/sagemaker/sagemaker/mxnet/processing.pyi +++ b/stubs/sagemaker/sagemaker/mxnet/processing.pyi @@ -11,20 +11,20 @@ class MXNetProcessor(FrameworkProcessor): def __init__( self, framework_version: str, - role: Optional[str | PipelineVariable] = None, - instance_count: int | PipelineVariable = None, - instance_type: str | PipelineVariable = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, py_version: str = "py3", - image_uri: Optional[str | PipelineVariable] = None, - command: Optional[List[str]] = None, + image_uri: str | PipelineVariable | None = None, + command: list[str] | None = None, volume_size_in_gb: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - code_location: Optional[str] = None, - max_runtime_in_seconds: Optional[int | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - network_config: Optional[NetworkConfig] = None, + volume_kms_key: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + code_location: str | None = None, + max_runtime_in_seconds: int | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + network_config: NetworkConfig | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/network.pyi b/stubs/sagemaker/sagemaker/network.pyi index 327771cd1066..f154e6287739 100644 --- a/stubs/sagemaker/sagemaker/network.pyi +++ b/stubs/sagemaker/sagemaker/network.pyi @@ -10,8 +10,8 @@ class NetworkConfig: encrypt_inter_container_traffic: Incomplete def __init__( self, - enable_network_isolation: bool | PipelineVariable = None, - security_group_ids: Optional[List[str | PipelineVariable]] = None, - subnets: Optional[List[str | PipelineVariable]] = None, - encrypt_inter_container_traffic: Optional[bool | PipelineVariable] = None, + enable_network_isolation: bool | PipelineVariable | None = None, + security_group_ids: list[str | PipelineVariable] | None = None, + subnets: list[str | PipelineVariable] | None = None, + encrypt_inter_container_traffic: bool | PipelineVariable | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/pipeline.pyi b/stubs/sagemaker/sagemaker/pipeline.pyi index f3f872e309e5..c2981edd4534 100644 --- a/stubs/sagemaker/sagemaker/pipeline.pyi +++ b/stubs/sagemaker/sagemaker/pipeline.pyi @@ -18,13 +18,13 @@ class PipelineModel: enable_network_isolation: Incomplete def __init__( self, - models: List[Model], - role: str = None, - predictor_cls: Optional[callable] = None, - name: Optional[str] = None, - vpc_config: Optional[Dict[str, List[str | PipelineVariable]]] = None, - sagemaker_session: Optional[Session] = None, - enable_network_isolation: bool | PipelineVariable = None, + models: list[Model], + role: str | None = None, + predictor_cls: callable | None = None, + name: str | None = None, + vpc_config: dict[str, list[str | PipelineVariable]] | None = None, + sagemaker_session: Session | None = None, + enable_network_isolation: bool | PipelineVariable | None = None, ) -> None: ... def pipeline_container_def(self, instance_type: Incomplete | None = None): ... def deploy( @@ -46,27 +46,27 @@ class PipelineModel: def create(self, instance_type: str): ... def register( self, - content_types: List[str | PipelineVariable], - response_types: List[str | PipelineVariable], - inference_instances: Optional[List[str | PipelineVariable]] = None, - transform_instances: Optional[List[str | PipelineVariable]] = None, - model_package_name: Optional[str | PipelineVariable] = None, - model_package_group_name: Optional[str | PipelineVariable] = None, - image_uri: Optional[str | PipelineVariable] = None, - model_metrics: Optional[ModelMetrics] = None, - metadata_properties: Optional[MetadataProperties] = None, + content_types: list[str | PipelineVariable], + response_types: list[str | PipelineVariable], + inference_instances: list[str | PipelineVariable] | None = None, + transform_instances: list[str | PipelineVariable] | None = None, + model_package_name: str | PipelineVariable | None = None, + model_package_group_name: str | PipelineVariable | None = None, + image_uri: str | PipelineVariable | None = None, + model_metrics: ModelMetrics | None = None, + metadata_properties: MetadataProperties | None = None, marketplace_cert: bool = False, - approval_status: Optional[str | PipelineVariable] = None, - description: Optional[str] = None, - drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, - domain: Optional[str | PipelineVariable] = None, - sample_payload_url: Optional[str | PipelineVariable] = None, - task: Optional[str | PipelineVariable] = None, - framework: Optional[str | PipelineVariable] = None, - framework_version: Optional[str | PipelineVariable] = None, - nearest_model_name: Optional[str | PipelineVariable] = None, - data_input_configuration: Optional[str | PipelineVariable] = None, + approval_status: str | PipelineVariable | None = None, + description: str | None = None, + drift_check_baselines: DriftCheckBaselines | None = None, + customer_metadata_properties: dict[str, str | PipelineVariable] | None = None, + domain: str | PipelineVariable | None = None, + sample_payload_url: str | PipelineVariable | None = None, + task: str | PipelineVariable | None = None, + framework: str | PipelineVariable | None = None, + framework_version: str | PipelineVariable | None = None, + nearest_model_name: str | PipelineVariable | None = None, + data_input_configuration: str | PipelineVariable | None = None, ): ... def transformer( self, diff --git a/stubs/sagemaker/sagemaker/predictor.pyi b/stubs/sagemaker/sagemaker/predictor.pyi index 96b2d7f73e7f..b7403f555f28 100644 --- a/stubs/sagemaker/sagemaker/predictor.pyi +++ b/stubs/sagemaker/sagemaker/predictor.pyi @@ -5,10 +5,10 @@ from sagemaker.session import Session def retrieve_default( endpoint_name: str, - sagemaker_session: Optional[Session] = None, - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, + sagemaker_session: Session | None = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, ) -> Predictor: ... diff --git a/stubs/sagemaker/sagemaker/processing.pyi b/stubs/sagemaker/sagemaker/processing.pyi index da9c8cb939e3..1252c191b3fa 100644 --- a/stubs/sagemaker/sagemaker/processing.pyi +++ b/stubs/sagemaker/sagemaker/processing.pyi @@ -31,51 +31,51 @@ class Processor: env: Incomplete def __init__( self, - role: str = None, - image_uri: str | PipelineVariable = None, - instance_count: int | PipelineVariable = None, - instance_type: str | PipelineVariable = None, - entrypoint: Optional[List[str | PipelineVariable]] = None, + role: str | None = None, + image_uri: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + entrypoint: list[str | PipelineVariable] | None = None, volume_size_in_gb: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - max_runtime_in_seconds: Optional[int | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - network_config: Optional[NetworkConfig] = None, + volume_kms_key: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + max_runtime_in_seconds: int | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + network_config: NetworkConfig | None = None, ) -> None: ... def run( self, - inputs: Optional[List["ProcessingInput"]] = None, - outputs: Optional[List["ProcessingOutput"]] = None, - arguments: Optional[List[str | PipelineVariable]] = None, + inputs: list["ProcessingInput"] | None = None, + outputs: list["ProcessingOutput"] | None = None, + arguments: list[str | PipelineVariable] | None = None, wait: bool = True, logs: bool = True, - job_name: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, - kms_key: Optional[str] = None, + job_name: str | None = None, + experiment_config: dict[str, str] | None = None, + kms_key: str | None = None, ): ... class ScriptProcessor(Processor): command: Incomplete def __init__( self, - role: Optional[str | PipelineVariable] = None, - image_uri: str | PipelineVariable = None, - command: List[str] = None, - instance_count: int | PipelineVariable = None, - instance_type: str | PipelineVariable = None, + role: str | PipelineVariable | None = None, + image_uri: str | PipelineVariable | None = None, + command: list[str] | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, volume_size_in_gb: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - max_runtime_in_seconds: Optional[int | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - network_config: Optional[NetworkConfig] = None, + volume_kms_key: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + max_runtime_in_seconds: int | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + network_config: NetworkConfig | None = None, ) -> None: ... def get_run_args( self, code, inputs: Incomplete | None = None, outputs: Incomplete | None = None, arguments: Incomplete | None = None @@ -84,14 +84,14 @@ class ScriptProcessor(Processor): def run( self, code: str, - inputs: Optional[List["ProcessingInput"]] = None, - outputs: Optional[List["ProcessingOutput"]] = None, - arguments: Optional[List[str | PipelineVariable]] = None, + inputs: list["ProcessingInput"] | None = None, + outputs: list["ProcessingOutput"] | None = None, + arguments: list[str | PipelineVariable] | None = None, wait: bool = True, logs: bool = True, - job_name: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, - kms_key: Optional[str] = None, + job_name: str | None = None, + experiment_config: dict[str, str] | None = None, + kms_key: str | None = None, ): ... class ProcessingJob(_Job): @@ -130,15 +130,15 @@ class ProcessingInput: app_managed: Incomplete def __init__( self, - source: Optional[str | PipelineVariable] = None, - destination: Optional[str | PipelineVariable] = None, - input_name: Optional[str | PipelineVariable] = None, + source: str | PipelineVariable | None = None, + destination: str | PipelineVariable | None = None, + input_name: str | PipelineVariable | None = None, s3_data_type: str | PipelineVariable = "S3Prefix", s3_input_mode: str | PipelineVariable = "File", s3_data_distribution_type: str | PipelineVariable = "FullyReplicated", s3_compression_type: str | PipelineVariable = "None", - s3_input: Optional[S3Input] = None, - dataset_definition: Optional[DatasetDefinition] = None, + s3_input: S3Input | None = None, + dataset_definition: DatasetDefinition | None = None, app_managed: bool | PipelineVariable = False, ) -> None: ... @@ -151,9 +151,9 @@ class ProcessingOutput: feature_store_output: Incomplete def __init__( self, - source: Optional[str | PipelineVariable] = None, - destination: Optional[str | PipelineVariable] = None, - output_name: Optional[str | PipelineVariable] = None, + source: str | PipelineVariable | None = None, + destination: str | PipelineVariable | None = None, + output_name: str | PipelineVariable | None = None, s3_upload_mode: str | PipelineVariable = "EndOfJob", app_managed: bool | PipelineVariable = False, feature_store_output: Optional["FeatureStoreOutput"] = None, @@ -185,22 +185,22 @@ class FrameworkProcessor(ScriptProcessor): self, estimator_cls: type, framework_version: str, - role: Optional[str | PipelineVariable] = None, - instance_count: int | PipelineVariable = None, - instance_type: str | PipelineVariable = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, py_version: str = "py3", - image_uri: Optional[str | PipelineVariable] = None, - command: Optional[List[str]] = None, + image_uri: str | PipelineVariable | None = None, + command: list[str] | None = None, volume_size_in_gb: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - code_location: Optional[str] = None, - max_runtime_in_seconds: Optional[int | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - network_config: Optional[NetworkConfig] = None, + volume_kms_key: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + code_location: str | None = None, + max_runtime_in_seconds: int | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + network_config: NetworkConfig | None = None, ) -> None: ... def get_run_args( self, @@ -216,15 +216,15 @@ class FrameworkProcessor(ScriptProcessor): def run( self, code: str, - source_dir: Optional[str] = None, - dependencies: Optional[List[str]] = None, - git_config: Optional[Dict[str, str]] = None, - inputs: Optional[List[ProcessingInput]] = None, - outputs: Optional[List[ProcessingOutput]] = None, - arguments: Optional[List[str | PipelineVariable]] = None, + source_dir: str | None = None, + dependencies: list[str] | None = None, + git_config: dict[str, str] | None = None, + inputs: list[ProcessingInput] | None = None, + outputs: list[ProcessingOutput] | None = None, + arguments: list[str | PipelineVariable] | None = None, wait: bool = True, logs: bool = True, - job_name: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, - kms_key: Optional[str] = None, + job_name: str | None = None, + experiment_config: dict[str, str] | None = None, + kms_key: str | None = None, ): ... diff --git a/stubs/sagemaker/sagemaker/pytorch/estimator.pyi b/stubs/sagemaker/sagemaker/pytorch/estimator.pyi index 3087bbe8be5c..2d8eaa131e93 100644 --- a/stubs/sagemaker/sagemaker/pytorch/estimator.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/estimator.pyi @@ -18,13 +18,13 @@ class PyTorch(Framework): def __init__( self, entry_point: str | PipelineVariable, - framework_version: Optional[str] = None, - py_version: Optional[str] = None, - source_dir: Optional[str | PipelineVariable] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, - image_uri: Optional[str | PipelineVariable] = None, - distribution: Optional[Dict] = None, - compiler_config: Optional[TrainingCompilerConfig] = None, + framework_version: str | None = None, + py_version: str | None = None, + source_dir: str | PipelineVariable | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, + image_uri: str | PipelineVariable | None = None, + distribution: dict | None = None, + compiler_config: TrainingCompilerConfig | None = None, **kwargs, ) -> None: ... def hyperparameters(self): ... diff --git a/stubs/sagemaker/sagemaker/pytorch/model.pyi b/stubs/sagemaker/sagemaker/pytorch/model.pyi index 34009ceb0fe3..fa98cd8175ca 100644 --- a/stubs/sagemaker/sagemaker/pytorch/model.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/model.pyi @@ -20,39 +20,39 @@ class PyTorchModel(FrameworkModel): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - entry_point: Optional[str] = None, + role: str | None = None, + entry_point: str | None = None, framework_version: str = "1.3", - py_version: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, + py_version: str | None = None, + image_uri: str | PipelineVariable | None = None, predictor_cls: callable = ..., - model_server_workers: Optional[int | PipelineVariable] = None, + model_server_workers: int | PipelineVariable | None = None, **kwargs, ) -> None: ... image_uri: Incomplete def register( self, - content_types: List[str | PipelineVariable], - response_types: List[str | PipelineVariable], - inference_instances: Optional[List[str | PipelineVariable]] = None, - transform_instances: Optional[List[str | PipelineVariable]] = None, - model_package_name: Optional[str | PipelineVariable] = None, - model_package_group_name: Optional[str | PipelineVariable] = None, - image_uri: Optional[str | PipelineVariable] = None, - model_metrics: Optional[ModelMetrics] = None, - metadata_properties: Optional[MetadataProperties] = None, + content_types: list[str | PipelineVariable], + response_types: list[str | PipelineVariable], + inference_instances: list[str | PipelineVariable] | None = None, + transform_instances: list[str | PipelineVariable] | None = None, + model_package_name: str | PipelineVariable | None = None, + model_package_group_name: str | PipelineVariable | None = None, + image_uri: str | PipelineVariable | None = None, + model_metrics: ModelMetrics | None = None, + metadata_properties: MetadataProperties | None = None, marketplace_cert: bool = False, - approval_status: Optional[str | PipelineVariable] = None, - description: Optional[str] = None, - drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, - domain: Optional[str | PipelineVariable] = None, - sample_payload_url: Optional[str | PipelineVariable] = None, - task: Optional[str | PipelineVariable] = None, - framework: Optional[str | PipelineVariable] = None, - framework_version: Optional[str | PipelineVariable] = None, - nearest_model_name: Optional[str | PipelineVariable] = None, - data_input_configuration: Optional[str | PipelineVariable] = None, + approval_status: str | PipelineVariable | None = None, + description: str | None = None, + drift_check_baselines: DriftCheckBaselines | None = None, + customer_metadata_properties: dict[str, str | PipelineVariable] | None = None, + domain: str | PipelineVariable | None = None, + sample_payload_url: str | PipelineVariable | None = None, + task: str | PipelineVariable | None = None, + framework: str | PipelineVariable | None = None, + framework_version: str | PipelineVariable | None = None, + nearest_model_name: str | PipelineVariable | None = None, + data_input_configuration: str | PipelineVariable | None = None, ): ... def prepare_container_def( self, diff --git a/stubs/sagemaker/sagemaker/pytorch/processing.pyi b/stubs/sagemaker/sagemaker/pytorch/processing.pyi index 8193fdbb673e..d60e077dfe62 100644 --- a/stubs/sagemaker/sagemaker/pytorch/processing.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/processing.pyi @@ -11,20 +11,20 @@ class PyTorchProcessor(FrameworkProcessor): def __init__( self, framework_version: str, - role: Optional[str | PipelineVariable] = None, - instance_count: int | PipelineVariable = None, - instance_type: str | PipelineVariable = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, py_version: str = "py3", - image_uri: Optional[str | PipelineVariable] = None, - command: Optional[List[str]] = None, + image_uri: str | PipelineVariable | None = None, + command: list[str] | None = None, volume_size_in_gb: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - code_location: Optional[str] = None, - max_runtime_in_seconds: Optional[int | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - network_config: Optional[NetworkConfig] = None, + volume_kms_key: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + code_location: str | None = None, + max_runtime_in_seconds: int | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + network_config: NetworkConfig | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/remote_function/client.pyi b/stubs/sagemaker/sagemaker/remote_function/client.pyi index 6469cf282286..978d5f220f36 100644 --- a/stubs/sagemaker/sagemaker/remote_function/client.pyi +++ b/stubs/sagemaker/sagemaker/remote_function/client.pyi @@ -11,30 +11,30 @@ logger: Incomplete def remote( _func: Incomplete | None = None, *, - dependencies: str = None, - pre_execution_commands: List[str] = None, - pre_execution_script: str = None, - environment_variables: Dict[str, str] = None, - image_uri: str = None, + dependencies: str | None = None, + pre_execution_commands: list[str] | None = None, + pre_execution_script: str | None = None, + environment_variables: dict[str, str] | None = None, + image_uri: str | None = None, include_local_workdir: bool = False, instance_count: int = 1, - instance_type: str = None, - job_conda_env: str = None, - job_name_prefix: str = None, + instance_type: str | None = None, + job_conda_env: str | None = None, + job_name_prefix: str | None = None, keep_alive_period_in_seconds: int = 0, max_retry_attempts: int = 1, max_runtime_in_seconds: int = 86400, - role: str = None, - s3_kms_key: str = None, - s3_root_uri: str = None, - sagemaker_session: Session = None, - security_group_ids: List[str] = None, - subnets: List[str] = None, - tags: List[Tuple[str, str]] = None, - volume_kms_key: str = None, + role: str | None = None, + s3_kms_key: str | None = None, + s3_root_uri: str | None = None, + sagemaker_session: Session | None = None, + security_group_ids: list[str] | None = None, + subnets: list[str] | None = None, + tags: list[tuple[str, str]] | None = None, + volume_kms_key: str | None = None, volume_size: int = 30, - encrypt_inter_container_traffic: bool = None, - spark_config: SparkConfig = None, + encrypt_inter_container_traffic: bool | None = None, + spark_config: SparkConfig | None = None, ): ... class _SubmitRequest: @@ -54,31 +54,31 @@ class RemoteExecutor: def __init__( self, *, - dependencies: str = None, - pre_execution_commands: List[str] = None, - pre_execution_script: str = None, - environment_variables: Dict[str, str] = None, - image_uri: str = None, + dependencies: str | None = None, + pre_execution_commands: list[str] | None = None, + pre_execution_script: str | None = None, + environment_variables: dict[str, str] | None = None, + image_uri: str | None = None, include_local_workdir: bool = False, instance_count: int = 1, - instance_type: str = None, - job_conda_env: str = None, - job_name_prefix: str = None, + instance_type: str | None = None, + job_conda_env: str | None = None, + job_name_prefix: str | None = None, keep_alive_period_in_seconds: int = 0, max_parallel_jobs: int = 1, max_retry_attempts: int = 1, max_runtime_in_seconds: int = 86400, - role: str = None, - s3_kms_key: str = None, - s3_root_uri: str = None, - sagemaker_session: Session = None, - security_group_ids: List[str] = None, - subnets: List[str] = None, - tags: List[Tuple[str, str]] = None, - volume_kms_key: str = None, + role: str | None = None, + s3_kms_key: str | None = None, + s3_root_uri: str | None = None, + sagemaker_session: Session | None = None, + security_group_ids: list[str] | None = None, + subnets: list[str] | None = None, + tags: list[tuple[str, str]] | None = None, + volume_kms_key: str | None = None, volume_size: int = 30, - encrypt_inter_container_traffic: bool = None, - spark_config: SparkConfig = None, + encrypt_inter_container_traffic: bool | None = None, + spark_config: SparkConfig | None = None, ) -> None: ... def submit(self, func, *args, **kwargs): ... def map(self, func, *iterables): ... @@ -90,8 +90,8 @@ class Future: def __init__(self) -> None: ... @staticmethod def from_describe_response(describe_training_job_response, sagemaker_session): ... - def result(self, timeout: float = None) -> Any: ... - def wait(self, timeout: int = None) -> None: ... + def result(self, timeout: float | None = None) -> Any: ... + def wait(self, timeout: int | None = None) -> None: ... def cancel(self) -> bool: ... def running(self) -> bool: ... def cancelled(self) -> bool: ... diff --git a/stubs/sagemaker/sagemaker/remote_function/job.pyi b/stubs/sagemaker/sagemaker/remote_function/job.pyi index b105cb8f7344..16c129b0f36b 100644 --- a/stubs/sagemaker/sagemaker/remote_function/job.pyi +++ b/stubs/sagemaker/sagemaker/remote_function/job.pyi @@ -56,30 +56,30 @@ class _JobSettings: def __init__( self, *, - dependencies: str = None, - pre_execution_commands: List[str] = None, - pre_execution_script: str = None, - environment_variables: Dict[str, str] = None, - image_uri: str = None, - include_local_workdir: bool = None, + dependencies: str | None = None, + pre_execution_commands: list[str] | None = None, + pre_execution_script: str | None = None, + environment_variables: dict[str, str] | None = None, + image_uri: str | None = None, + include_local_workdir: bool | None = None, instance_count: int = 1, - instance_type: str = None, - job_conda_env: str = None, - job_name_prefix: str = None, + instance_type: str | None = None, + job_conda_env: str | None = None, + job_name_prefix: str | None = None, keep_alive_period_in_seconds: int = 0, max_retry_attempts: int = 1, max_runtime_in_seconds: int = 86400, - role: str = None, - s3_kms_key: str = None, - s3_root_uri: str = None, - sagemaker_session: Session = None, - security_group_ids: List[str] = None, - subnets: List[str] = None, - tags: List[Tuple[str, str]] = None, - volume_kms_key: str = None, + role: str | None = None, + s3_kms_key: str | None = None, + s3_root_uri: str | None = None, + sagemaker_session: Session | None = None, + security_group_ids: list[str] | None = None, + subnets: list[str] | None = None, + tags: list[tuple[str, str]] | None = None, + volume_kms_key: str | None = None, volume_size: int = 30, - encrypt_inter_container_traffic: bool = None, - spark_config: SparkConfig = None, + encrypt_inter_container_traffic: bool | None = None, + spark_config: SparkConfig | None = None, ) -> None: ... class _Job: @@ -94,7 +94,7 @@ class _Job: def start(job_settings: _JobSettings, func, func_args, func_kwargs, run_info: Incomplete | None = None): ... def describe(self): ... def stop(self) -> None: ... - def wait(self, timeout: int = None): ... + def wait(self, timeout: int | None = None): ... class _RunInfo: experiment_name: str diff --git a/stubs/sagemaker/sagemaker/remote_function/runtime_environment/runtime_environment_manager.pyi b/stubs/sagemaker/sagemaker/remote_function/runtime_environment/runtime_environment_manager.pyi index bc461b922124..9a6260b7a970 100644 --- a/stubs/sagemaker/sagemaker/remote_function/runtime_environment/runtime_environment_manager.pyi +++ b/stubs/sagemaker/sagemaker/remote_function/runtime_environment/runtime_environment_manager.pyi @@ -10,8 +10,8 @@ def get_logger(): ... logger: Incomplete class RuntimeEnvironmentManager: - def snapshot(self, dependencies: str = None) -> str: ... - def bootstrap(self, local_dependencies_file: str, client_python_version: str, conda_env: str = None): ... + def snapshot(self, dependencies: str | None = None) -> str: ... + def bootstrap(self, local_dependencies_file: str, client_python_version: str, conda_env: str | None = None): ... def run_pre_exec_script(self, pre_exec_script_path: str): ... class RuntimeEnvironmentError(Exception): diff --git a/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi b/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi index 21d36b1c5ac4..78c6417c16ee 100644 --- a/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi +++ b/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi @@ -1,11 +1,11 @@ from typing import Dict, List, Optional class SparkConfig: - submit_jars: Optional[List[str]] - submit_py_files: Optional[List[str]] - submit_files: Optional[List[str]] - configuration: Optional[List[Dict, Dict]] - spark_event_logs_uri: Optional[str] + submit_jars: list[str] | None + submit_py_files: list[str] | None + submit_files: list[str] | None + configuration: list[dict, dict] | None + spark_event_logs_uri: str | None def __init__(self, submit_jars, submit_py_files, submit_files, configuration, spark_event_logs_uri) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/rl/estimator.pyi b/stubs/sagemaker/sagemaker/rl/estimator.pyi index 17d02e4a0bf3..420501cc025f 100644 --- a/stubs/sagemaker/sagemaker/rl/estimator.pyi +++ b/stubs/sagemaker/sagemaker/rl/estimator.pyi @@ -31,13 +31,13 @@ class RLEstimator(Framework): def __init__( self, entry_point: str | PipelineVariable, - toolkit: Optional[RLToolkit] = None, - toolkit_version: Optional[str] = None, - framework: Optional[Framework] = None, - source_dir: Optional[str | PipelineVariable] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, - image_uri: Optional[str | PipelineVariable] = None, - metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, + toolkit: RLToolkit | None = None, + toolkit_version: str | None = None, + framework: Framework | None = None, + source_dir: str | PipelineVariable | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, + image_uri: str | PipelineVariable | None = None, + metric_definitions: list[dict[str, str | PipelineVariable]] | None = None, **kwargs, ) -> None: ... def create_model( diff --git a/stubs/sagemaker/sagemaker/s3_utils.pyi b/stubs/sagemaker/sagemaker/s3_utils.pyi index 18aeb6b8183e..3d58acf2e4e6 100644 --- a/stubs/sagemaker/sagemaker/s3_utils.pyi +++ b/stubs/sagemaker/sagemaker/s3_utils.pyi @@ -6,5 +6,5 @@ logger: Incomplete def parse_s3_url(url): ... def s3_path_join(*args, with_end_slash: bool = False): ... def determine_bucket_and_prefix( - bucket: Optional[str] = None, key_prefix: Optional[str] = None, sagemaker_session: Incomplete | None = None + bucket: str | None = None, key_prefix: str | None = None, sagemaker_session: Incomplete | None = None ): ... diff --git a/stubs/sagemaker/sagemaker/script_uris.pyi b/stubs/sagemaker/sagemaker/script_uris.pyi index e7d49b650701..c7708e0ee9f8 100644 --- a/stubs/sagemaker/sagemaker/script_uris.pyi +++ b/stubs/sagemaker/sagemaker/script_uris.pyi @@ -4,10 +4,10 @@ from typing import Optional logger: Incomplete def retrieve( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, - script_scope: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, + script_scope: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, ) -> str: ... diff --git a/stubs/sagemaker/sagemaker/serializers.pyi b/stubs/sagemaker/sagemaker/serializers.pyi index 2c3776d78e9e..de1ee2357557 100644 --- a/stubs/sagemaker/sagemaker/serializers.pyi +++ b/stubs/sagemaker/sagemaker/serializers.pyi @@ -14,16 +14,16 @@ from sagemaker.base_serializers import ( ) def retrieve_options( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, -) -> List[BaseSerializer]: ... +) -> list[BaseSerializer]: ... def retrieve_default( - region: Optional[str] = None, - model_id: Optional[str] = None, - model_version: Optional[str] = None, + region: str | None = None, + model_id: str | None = None, + model_version: str | None = None, tolerate_vulnerable_model: bool = False, tolerate_deprecated_model: bool = False, ) -> BaseSerializer: ... diff --git a/stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi b/stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi index 211edd548eeb..6ac6a63e7fe5 100644 --- a/stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi +++ b/stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi @@ -6,5 +6,5 @@ class ServerlessInferenceConfig: max_concurrency: Incomplete provisioned_concurrency: Incomplete def __init__( - self, memory_size_in_mb: int = 2048, max_concurrency: int = 5, provisioned_concurrency: Optional[int] = None + self, memory_size_in_mb: int = 2048, max_concurrency: int = 5, provisioned_concurrency: int | None = None ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/session.pyi b/stubs/sagemaker/sagemaker/session.pyi index 64eb768a84c7..56f8db5c3ffd 100644 --- a/stubs/sagemaker/sagemaker/session.pyi +++ b/stubs/sagemaker/sagemaker/session.pyi @@ -1,5 +1,6 @@ from _typeshed import Incomplete -from typing import Any, Dict, List, Optional, Sequence +from typing import Any, Dict, List, Optional +from collections.abc import Sequence from sagemaker.inputs import BatchDataCaptureConfig @@ -31,8 +32,8 @@ class Session: default_bucket: Incomplete | None = None, settings=..., sagemaker_metrics_client: Incomplete | None = None, - sagemaker_config: dict = None, - default_bucket_prefix: str = None, + sagemaker_config: dict | None = None, + default_bucket_prefix: str | None = None, ) -> None: ... @property def boto_region_name(self): ... @@ -74,7 +75,7 @@ class Session: enable_sagemaker_metrics: Incomplete | None = None, profiler_rule_configs: Incomplete | None = None, profiler_config: Incomplete | None = None, - environment: Optional[Dict[str, str]] = None, + environment: dict[str, str] | None = None, retry_strategy: Incomplete | None = None, ): ... def update_training_job( @@ -92,7 +93,7 @@ class Session: resources, stopping_condition, app_specification, - environment: Optional[Dict[str, str]] = None, + environment: dict[str, str] | None = None, network_config: Incomplete | None = None, role_arn: Incomplete | None = None, tags: Incomplete | None = None, @@ -161,20 +162,20 @@ class Session: self, monitoring_schedule_name: str, monitoring_alert_name: str, data_points_to_alert: int, evaluation_period: int ): ... def list_monitoring_alerts( - self, monitoring_schedule_name: str, next_token: Optional[str] = None, max_results: Optional[int] = 10 - ) -> Dict: ... + self, monitoring_schedule_name: str, next_token: str | None = None, max_results: int | None = 10 + ) -> dict: ... def list_monitoring_alert_history( self, - monitoring_schedule_name: Optional[str] = None, - monitoring_alert_name: Optional[str] = None, - sort_by: Optional[str] = "CreationTime", - sort_order: Optional[str] = "Descending", - next_token: Optional[str] = None, - max_results: Optional[int] = 10, - creation_time_before: Optional[str] = None, - creation_time_after: Optional[str] = None, - status_equals: Optional[str] = None, - ) -> Dict: ... + monitoring_schedule_name: str | None = None, + monitoring_alert_name: str | None = None, + sort_by: str | None = "CreationTime", + sort_order: str | None = "Descending", + next_token: str | None = None, + max_results: int | None = 10, + creation_time_before: str | None = None, + creation_time_after: str | None = None, + status_equals: str | None = None, + ) -> dict: ... def was_processing_job_successful(self, job_name): ... def describe_processing_job(self, job_name): ... def stop_processing_job(self, job_name) -> None: ... @@ -286,11 +287,11 @@ class Session: output_config, resource_config, experiment_config, - env: Optional[Dict[str, str]] = None, + env: dict[str, str] | None = None, tags: Incomplete | None = None, data_processing: Incomplete | None = None, model_client_config: Incomplete | None = None, - batch_data_capture_config: BatchDataCaptureConfig = None, + batch_data_capture_config: BatchDataCaptureConfig | None = None, ): ... def create_model( self, @@ -436,15 +437,15 @@ class Session: feature_group_name: str, record_identifier_name: str, event_time_feature_name: str, - feature_definitions: Sequence[Dict[str, str]], - role_arn: str = None, - online_store_config: Dict[str, str] = None, - offline_store_config: Dict[str, str] = None, - description: str = None, - tags: List[Dict[str, str]] = None, - ) -> Dict[str, Any]: ... - def describe_feature_group(self, feature_group_name: str, next_token: str = None) -> Dict[str, Any]: ... - def update_feature_group(self, feature_group_name: str, feature_additions: Sequence[Dict[str, str]]) -> Dict[str, Any]: ... + feature_definitions: Sequence[dict[str, str]], + role_arn: str | None = None, + online_store_config: dict[str, str] | None = None, + offline_store_config: dict[str, str] | None = None, + description: str | None = None, + tags: list[dict[str, str]] | None = None, + ) -> dict[str, Any]: ... + def describe_feature_group(self, feature_group_name: str, next_token: str | None = None) -> dict[str, Any]: ... + def update_feature_group(self, feature_group_name: str, feature_additions: Sequence[dict[str, str]]) -> dict[str, Any]: ... def list_feature_groups( self, name_contains, @@ -456,37 +457,37 @@ class Session: sort_by, max_results, next_token, - ) -> Dict[str, Any]: ... + ) -> dict[str, Any]: ... def update_feature_metadata( self, feature_group_name: str, feature_name: str, - description: str = None, - parameter_additions: Sequence[Dict[str, str]] = None, - parameter_removals: Sequence[str] = None, - ) -> Dict[str, Any]: ... - def describe_feature_metadata(self, feature_group_name: str, feature_name: str) -> Dict[str, Any]: ... + description: str | None = None, + parameter_additions: Sequence[dict[str, str]] | None = None, + parameter_removals: Sequence[str] | None = None, + ) -> dict[str, Any]: ... + def describe_feature_metadata(self, feature_group_name: str, feature_name: str) -> dict[str, Any]: ... def search( self, resource: str, - search_expression: Dict[str, any] = None, - sort_by: str = None, - sort_order: str = None, - next_token: str = None, - max_results: int = None, - ) -> Dict[str, Any]: ... - def put_record(self, feature_group_name: str, record: Sequence[Dict[str, str]]): ... + search_expression: dict[str, any] | None = None, + sort_by: str | None = None, + sort_order: str | None = None, + next_token: str | None = None, + max_results: int | None = None, + ) -> dict[str, Any]: ... + def put_record(self, feature_group_name: str, record: Sequence[dict[str, str]]): ... def delete_record( - self, feature_group_name: str, record_identifier_value_as_string: str, event_time: str, deletion_mode: str = None + self, feature_group_name: str, record_identifier_value_as_string: str, event_time: str, deletion_mode: str | None = None ): ... def get_record( self, record_identifier_value_as_string: str, feature_group_name: str, feature_names: Sequence[str] - ) -> Dict[str, Sequence[Dict[str, str]]]: ... - def batch_get_record(self, identifiers: Sequence[Dict[str, Any]]) -> Dict[str, Any]: ... + ) -> dict[str, Sequence[dict[str, str]]]: ... + def batch_get_record(self, identifiers: Sequence[dict[str, Any]]) -> dict[str, Any]: ... def start_query_execution( - self, catalog: str, database: str, query_string: str, output_location: str, kms_key: str = None, workgroup: str = None - ) -> Dict[str, str]: ... - def get_query_execution(self, query_execution_id: str) -> Dict[str, Any]: ... + self, catalog: str, database: str, query_string: str, output_location: str, kms_key: str | None = None, workgroup: str | None = None + ) -> dict[str, str]: ... + def get_query_execution(self, query_execution_id: str) -> dict[str, Any]: ... def wait_for_athena_query(self, query_execution_id: str, poll: int = 5): ... def download_athena_query_result(self, bucket: str, prefix: str, query_execution_id: str, filename: str): ... def account_id(self) -> str: ... @@ -494,24 +495,24 @@ class Session: self, role: str, sample_payload_url: str, - supported_content_types: List[str], - job_name: str = None, + supported_content_types: list[str], + job_name: str | None = None, job_type: str = "Default", - model_name: str = None, - model_package_version_arn: str = None, - job_duration_in_seconds: int = None, - nearest_model_name: str = None, - supported_instance_types: List[str] = None, - framework: str = None, - framework_version: str = None, - endpoint_configurations: List[Dict[str, any]] = None, - traffic_pattern: Dict[str, any] = None, - stopping_conditions: Dict[str, any] = None, - resource_limit: Dict[str, any] = None, + model_name: str | None = None, + model_package_version_arn: str | None = None, + job_duration_in_seconds: int | None = None, + nearest_model_name: str | None = None, + supported_instance_types: list[str] | None = None, + framework: str | None = None, + framework_version: str | None = None, + endpoint_configurations: list[dict[str, any]] | None = None, + traffic_pattern: dict[str, any] | None = None, + stopping_conditions: dict[str, any] | None = None, + resource_limit: dict[str, any] | None = None, ): ... def wait_for_inference_recommendations_job( self, job_name: str, poll: int = 120, log_level: str = "Verbose" - ) -> Dict[str, Any]: ... + ) -> dict[str, Any]: ... def get_model_package_args( content_types, @@ -557,7 +558,7 @@ def get_create_model_package_request( sample_payload_url: Incomplete | None = None, task: Incomplete | None = None, ): ... -def update_args(args: Dict[str, Any], **kwargs): ... +def update_args(args: dict[str, Any], **kwargs): ... def container_def( image_uri, model_data_url: Incomplete | None = None, diff --git a/stubs/sagemaker/sagemaker/sklearn/estimator.pyi b/stubs/sagemaker/sagemaker/sklearn/estimator.pyi index 58ba3d43bd25..6e316e00f6c6 100644 --- a/stubs/sagemaker/sagemaker/sklearn/estimator.pyi +++ b/stubs/sagemaker/sagemaker/sklearn/estimator.pyi @@ -13,12 +13,12 @@ class SKLearn(Framework): def __init__( self, entry_point: str | PipelineVariable, - framework_version: Optional[str] = None, + framework_version: str | None = None, py_version: str = "py3", - source_dir: Optional[str | PipelineVariable] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, - image_uri: Optional[str | PipelineVariable] = None, - image_uri_region: Optional[str] = None, + source_dir: str | PipelineVariable | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, + image_uri: str | PipelineVariable | None = None, + image_uri_region: str | None = None, **kwargs, ) -> None: ... def create_model( diff --git a/stubs/sagemaker/sagemaker/sklearn/model.pyi b/stubs/sagemaker/sagemaker/sklearn/model.pyi index 61fd7877063b..53eb5accd45b 100644 --- a/stubs/sagemaker/sagemaker/sklearn/model.pyi +++ b/stubs/sagemaker/sagemaker/sklearn/model.pyi @@ -20,39 +20,39 @@ class SKLearnModel(FrameworkModel): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, - entry_point: Optional[str] = None, - framework_version: Optional[str] = None, + role: str | None = None, + entry_point: str | None = None, + framework_version: str | None = None, py_version: str = "py3", - image_uri: Optional[str | PipelineVariable] = None, + image_uri: str | PipelineVariable | None = None, predictor_cls: callable = ..., - model_server_workers: Optional[int | PipelineVariable] = None, + model_server_workers: int | PipelineVariable | None = None, **kwargs, ) -> None: ... image_uri: Incomplete def register( self, - content_types: List[str | PipelineVariable], - response_types: List[str | PipelineVariable], - inference_instances: Optional[List[str | PipelineVariable]] = None, - transform_instances: Optional[List[str | PipelineVariable]] = None, - model_package_name: Optional[str | PipelineVariable] = None, - model_package_group_name: Optional[str | PipelineVariable] = None, - image_uri: Optional[str | PipelineVariable] = None, - model_metrics: Optional[ModelMetrics] = None, - metadata_properties: Optional[MetadataProperties] = None, + content_types: list[str | PipelineVariable], + response_types: list[str | PipelineVariable], + inference_instances: list[str | PipelineVariable] | None = None, + transform_instances: list[str | PipelineVariable] | None = None, + model_package_name: str | PipelineVariable | None = None, + model_package_group_name: str | PipelineVariable | None = None, + image_uri: str | PipelineVariable | None = None, + model_metrics: ModelMetrics | None = None, + metadata_properties: MetadataProperties | None = None, marketplace_cert: bool = False, - approval_status: Optional[str | PipelineVariable] = None, - description: Optional[str] = None, - drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, - domain: Optional[str | PipelineVariable] = None, - sample_payload_url: Optional[str | PipelineVariable] = None, - task: Optional[str | PipelineVariable] = None, - framework: Optional[str | PipelineVariable] = None, - framework_version: Optional[str | PipelineVariable] = None, - nearest_model_name: Optional[str | PipelineVariable] = None, - data_input_configuration: Optional[str | PipelineVariable] = None, + approval_status: str | PipelineVariable | None = None, + description: str | None = None, + drift_check_baselines: DriftCheckBaselines | None = None, + customer_metadata_properties: dict[str, str | PipelineVariable] | None = None, + domain: str | PipelineVariable | None = None, + sample_payload_url: str | PipelineVariable | None = None, + task: str | PipelineVariable | None = None, + framework: str | PipelineVariable | None = None, + framework_version: str | PipelineVariable | None = None, + nearest_model_name: str | PipelineVariable | None = None, + data_input_configuration: str | PipelineVariable | None = None, ): ... def prepare_container_def( self, diff --git a/stubs/sagemaker/sagemaker/sklearn/processing.pyi b/stubs/sagemaker/sagemaker/sklearn/processing.pyi index ec6e9c62b22d..677b03f4adaf 100644 --- a/stubs/sagemaker/sagemaker/sklearn/processing.pyi +++ b/stubs/sagemaker/sagemaker/sklearn/processing.pyi @@ -9,17 +9,17 @@ class SKLearnProcessor(ScriptProcessor): def __init__( self, framework_version: str, - role: Optional[str | PipelineVariable] = None, - instance_count: int | PipelineVariable = None, - instance_type: str | PipelineVariable = None, - command: Optional[List[str]] = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + command: list[str] | None = None, volume_size_in_gb: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - max_runtime_in_seconds: Optional[int | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - network_config: Optional[NetworkConfig] = None, + volume_kms_key: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + max_runtime_in_seconds: int | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + network_config: NetworkConfig | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/spark/processing.pyi b/stubs/sagemaker/sagemaker/spark/processing.pyi index 0020f95cd2fb..4054b6e3ec98 100644 --- a/stubs/sagemaker/sagemaker/spark/processing.pyi +++ b/stubs/sagemaker/sagemaker/spark/processing.pyi @@ -26,8 +26,8 @@ class _SparkProcessorBase(ScriptProcessor): volume_size_in_gb: int = 30, volume_kms_key: Incomplete | None = None, output_kms_key: Incomplete | None = None, - configuration_location: Optional[str] = None, - dependency_location: Optional[str] = None, + configuration_location: str | None = None, + dependency_location: str | None = None, max_runtime_in_seconds: Incomplete | None = None, base_job_name: Incomplete | None = None, sagemaker_session: Incomplete | None = None, @@ -56,24 +56,24 @@ class _SparkProcessorBase(ScriptProcessor): class PySparkProcessor(_SparkProcessorBase): def __init__( self, - role: str = None, - instance_type: str | PipelineVariable = None, - instance_count: int | PipelineVariable = None, - framework_version: Optional[str] = None, - py_version: Optional[str] = None, - container_version: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, + role: str | None = None, + instance_type: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + framework_version: str | None = None, + py_version: str | None = None, + container_version: str | None = None, + image_uri: str | PipelineVariable | None = None, volume_size_in_gb: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - configuration_location: Optional[str] = None, - dependency_location: Optional[str] = None, - max_runtime_in_seconds: Optional[int | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - network_config: Optional[NetworkConfig] = None, + volume_kms_key: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + configuration_location: str | None = None, + dependency_location: str | None = None, + max_runtime_in_seconds: int | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + network_config: NetworkConfig | None = None, ) -> None: ... def get_run_args( self, @@ -91,42 +91,42 @@ class PySparkProcessor(_SparkProcessorBase): def run( self, submit_app: str, - submit_py_files: Optional[List[str | PipelineVariable]] = None, - submit_jars: Optional[List[str | PipelineVariable]] = None, - submit_files: Optional[List[str | PipelineVariable]] = None, - inputs: Optional[List[ProcessingInput]] = None, - outputs: Optional[List[ProcessingOutput]] = None, - arguments: Optional[List[str | PipelineVariable]] = None, + submit_py_files: list[str | PipelineVariable] | None = None, + submit_jars: list[str | PipelineVariable] | None = None, + submit_files: list[str | PipelineVariable] | None = None, + inputs: list[ProcessingInput] | None = None, + outputs: list[ProcessingOutput] | None = None, + arguments: list[str | PipelineVariable] | None = None, wait: bool = True, logs: bool = True, - job_name: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, - configuration: Optional[List[Dict, Dict]] = None, - spark_event_logs_s3_uri: Optional[str | PipelineVariable] = None, - kms_key: Optional[str] = None, + job_name: str | None = None, + experiment_config: dict[str, str] | None = None, + configuration: list[dict, dict] | None = None, + spark_event_logs_s3_uri: str | PipelineVariable | None = None, + kms_key: str | None = None, ): ... class SparkJarProcessor(_SparkProcessorBase): def __init__( self, - role: str = None, - instance_type: str | PipelineVariable = None, - instance_count: int | PipelineVariable = None, - framework_version: Optional[str] = None, - py_version: Optional[str] = None, - container_version: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, + role: str | None = None, + instance_type: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + framework_version: str | None = None, + py_version: str | None = None, + container_version: str | None = None, + image_uri: str | PipelineVariable | None = None, volume_size_in_gb: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - configuration_location: Optional[str] = None, - dependency_location: Optional[str] = None, - max_runtime_in_seconds: Optional[int | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - network_config: Optional[NetworkConfig] = None, + volume_kms_key: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + configuration_location: str | None = None, + dependency_location: str | None = None, + max_runtime_in_seconds: int | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + network_config: NetworkConfig | None = None, ) -> None: ... def get_run_args( self, @@ -145,18 +145,18 @@ class SparkJarProcessor(_SparkProcessorBase): self, submit_app: str, submit_class: str | PipelineVariable, - submit_jars: Optional[List[str | PipelineVariable]] = None, - submit_files: Optional[List[str | PipelineVariable]] = None, - inputs: Optional[List[ProcessingInput]] = None, - outputs: Optional[List[ProcessingOutput]] = None, - arguments: Optional[List[str | PipelineVariable]] = None, + submit_jars: list[str | PipelineVariable] | None = None, + submit_files: list[str | PipelineVariable] | None = None, + inputs: list[ProcessingInput] | None = None, + outputs: list[ProcessingOutput] | None = None, + arguments: list[str | PipelineVariable] | None = None, wait: bool = True, logs: bool = True, - job_name: Optional[str] = None, - experiment_config: Optional[Dict[str, str]] = None, - configuration: Optional[List[Dict, Dict]] = None, - spark_event_logs_s3_uri: Optional[str | PipelineVariable] = None, - kms_key: Optional[str] = None, + job_name: str | None = None, + experiment_config: dict[str, str] | None = None, + configuration: list[dict, dict] | None = None, + spark_event_logs_s3_uri: str | PipelineVariable | None = None, + kms_key: str | None = None, ): ... class _HistoryServer: @@ -177,6 +177,6 @@ class FileType(Enum): class SparkConfigUtils: @staticmethod - def validate_configuration(configuration: Dict): ... + def validate_configuration(configuration: dict): ... @staticmethod def validate_s3_uri(spark_output_s3_path) -> None: ... diff --git a/stubs/sagemaker/sagemaker/sparkml/model.pyi b/stubs/sagemaker/sagemaker/sparkml/model.pyi index 79985c7057e6..85769ef8b73c 100644 --- a/stubs/sagemaker/sagemaker/sparkml/model.pyi +++ b/stubs/sagemaker/sagemaker/sparkml/model.pyi @@ -13,8 +13,8 @@ class SparkMLModel(Model): def __init__( self, model_data: str | PipelineVariable, - role: Optional[str] = None, + role: str | None = None, spark_version: str = "3.3", - sagemaker_session: Optional[Session] = None, + sagemaker_session: Session | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi b/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi index 322cb7a5b5cd..54f5a1ee2030 100644 --- a/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi +++ b/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi @@ -16,12 +16,12 @@ class TensorFlow(Framework): compiler_config: Incomplete def __init__( self, - py_version: Optional[str] = None, - framework_version: Optional[str] = None, - model_dir: Optional[str | PipelineVariable] = None, - image_uri: Optional[str | PipelineVariable] = None, - distribution: Optional[Dict[str, str]] = None, - compiler_config: Optional[TrainingCompilerConfig] = None, + py_version: str | None = None, + framework_version: str | None = None, + model_dir: str | PipelineVariable | None = None, + image_uri: str | PipelineVariable | None = None, + distribution: dict[str, str] | None = None, + compiler_config: TrainingCompilerConfig | None = None, **kwargs, ) -> None: ... def create_model( diff --git a/stubs/sagemaker/sagemaker/tensorflow/model.pyi b/stubs/sagemaker/sagemaker/tensorflow/model.pyi index df35c8a673eb..1862a861de82 100644 --- a/stubs/sagemaker/sagemaker/tensorflow/model.pyi +++ b/stubs/sagemaker/sagemaker/tensorflow/model.pyi @@ -34,38 +34,38 @@ class TensorFlowModel(sagemaker.model.FrameworkModel): def __init__( self, model_data: str | PipelineVariable, - role: str = None, - entry_point: Optional[str] = None, - image_uri: Optional[str | PipelineVariable] = None, - framework_version: Optional[str] = None, - container_log_level: Optional[int] = None, + role: str | None = None, + entry_point: str | None = None, + image_uri: str | PipelineVariable | None = None, + framework_version: str | None = None, + container_log_level: int | None = None, predictor_cls: callable = ..., **kwargs, ) -> None: ... image_uri: Incomplete def register( self, - content_types: List[str | PipelineVariable], - response_types: List[str | PipelineVariable], - inference_instances: Optional[List[str | PipelineVariable]] = None, - transform_instances: Optional[List[str | PipelineVariable]] = None, - model_package_name: Optional[str | PipelineVariable] = None, - model_package_group_name: Optional[str | PipelineVariable] = None, - image_uri: Optional[str | PipelineVariable] = None, - model_metrics: Optional[ModelMetrics] = None, - metadata_properties: Optional[MetadataProperties] = None, + content_types: list[str | PipelineVariable], + response_types: list[str | PipelineVariable], + inference_instances: list[str | PipelineVariable] | None = None, + transform_instances: list[str | PipelineVariable] | None = None, + model_package_name: str | PipelineVariable | None = None, + model_package_group_name: str | PipelineVariable | None = None, + image_uri: str | PipelineVariable | None = None, + model_metrics: ModelMetrics | None = None, + metadata_properties: MetadataProperties | None = None, marketplace_cert: bool = False, - approval_status: Optional[str | PipelineVariable] = None, - description: Optional[str] = None, - drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, - domain: Optional[str | PipelineVariable] = None, - sample_payload_url: Optional[str | PipelineVariable] = None, - task: Optional[str | PipelineVariable] = None, - framework: Optional[str | PipelineVariable] = None, - framework_version: Optional[str | PipelineVariable] = None, - nearest_model_name: Optional[str | PipelineVariable] = None, - data_input_configuration: Optional[str | PipelineVariable] = None, + approval_status: str | PipelineVariable | None = None, + description: str | None = None, + drift_check_baselines: DriftCheckBaselines | None = None, + customer_metadata_properties: dict[str, str | PipelineVariable] | None = None, + domain: str | PipelineVariable | None = None, + sample_payload_url: str | PipelineVariable | None = None, + task: str | PipelineVariable | None = None, + framework: str | PipelineVariable | None = None, + framework_version: str | PipelineVariable | None = None, + nearest_model_name: str | PipelineVariable | None = None, + data_input_configuration: str | PipelineVariable | None = None, ): ... def deploy( self, diff --git a/stubs/sagemaker/sagemaker/tensorflow/processing.pyi b/stubs/sagemaker/sagemaker/tensorflow/processing.pyi index d0fa1560e444..dfbd2cea8de6 100644 --- a/stubs/sagemaker/sagemaker/tensorflow/processing.pyi +++ b/stubs/sagemaker/sagemaker/tensorflow/processing.pyi @@ -11,20 +11,20 @@ class TensorFlowProcessor(FrameworkProcessor): def __init__( self, framework_version: str, - role: Optional[str | PipelineVariable] = None, - instance_count: int | PipelineVariable = None, - instance_type: str | PipelineVariable = None, + role: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, py_version: str = "py3", - image_uri: Optional[str | PipelineVariable] = None, - command: Optional[List[str]] = None, + image_uri: str | PipelineVariable | None = None, + command: list[str] | None = None, volume_size_in_gb: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - code_location: Optional[str] = None, - max_runtime_in_seconds: Optional[int | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - network_config: Optional[NetworkConfig] = None, + volume_kms_key: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + code_location: str | None = None, + max_runtime_in_seconds: int | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + network_config: NetworkConfig | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/transformer.pyi b/stubs/sagemaker/sagemaker/transformer.pyi index 21796c04ff87..169e9ae5eb35 100644 --- a/stubs/sagemaker/sagemaker/transformer.pyi +++ b/stubs/sagemaker/sagemaker/transformer.pyi @@ -30,33 +30,33 @@ class Transformer: model_name: str | PipelineVariable, instance_count: int | PipelineVariable, instance_type: str | PipelineVariable, - strategy: Optional[str | PipelineVariable] = None, - assemble_with: Optional[str | PipelineVariable] = None, - output_path: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - accept: Optional[str | PipelineVariable] = None, - max_concurrent_transforms: Optional[int | PipelineVariable] = None, - max_payload: Optional[int | PipelineVariable] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - base_transform_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - volume_kms_key: Optional[str | PipelineVariable] = None, + strategy: str | PipelineVariable | None = None, + assemble_with: str | PipelineVariable | None = None, + output_path: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + accept: str | PipelineVariable | None = None, + max_concurrent_transforms: int | PipelineVariable | None = None, + max_payload: int | PipelineVariable | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + env: dict[str, str | PipelineVariable] | None = None, + base_transform_job_name: str | None = None, + sagemaker_session: Session | None = None, + volume_kms_key: str | PipelineVariable | None = None, ) -> None: ... def transform( self, data: str | PipelineVariable, data_type: str | PipelineVariable = "S3Prefix", - content_type: Optional[str | PipelineVariable] = None, - compression_type: Optional[str | PipelineVariable] = None, - split_type: Optional[str | PipelineVariable] = None, - job_name: Optional[str] = None, - input_filter: Optional[str | PipelineVariable] = None, - output_filter: Optional[str | PipelineVariable] = None, - join_source: Optional[str | PipelineVariable] = None, - experiment_config: Optional[Dict[str, str]] = None, - model_client_config: Optional[Dict[str, str | PipelineVariable]] = None, - batch_data_capture_config: BatchDataCaptureConfig = None, + content_type: str | PipelineVariable | None = None, + compression_type: str | PipelineVariable | None = None, + split_type: str | PipelineVariable | None = None, + job_name: str | None = None, + input_filter: str | PipelineVariable | None = None, + output_filter: str | PipelineVariable | None = None, + join_source: str | PipelineVariable | None = None, + experiment_config: dict[str, str] | None = None, + model_client_config: dict[str, str | PipelineVariable] | None = None, + batch_data_capture_config: BatchDataCaptureConfig | None = None, wait: bool = True, logs: bool = True, ): ... @@ -66,20 +66,20 @@ class Transformer: monitoring_resource_config, data: str, data_type: str = "S3Prefix", - content_type: str = None, - compression_type: str = None, - split_type: str = None, - input_filter: str = None, - output_filter: str = None, - join_source: str = None, - model_client_config: Dict[str, str] = None, - batch_data_capture_config: BatchDataCaptureConfig = None, + content_type: str | None = None, + compression_type: str | None = None, + split_type: str | None = None, + input_filter: str | None = None, + output_filter: str | None = None, + join_source: str | None = None, + model_client_config: dict[str, str] | None = None, + batch_data_capture_config: BatchDataCaptureConfig | None = None, monitor_before_transform: bool = False, - supplied_baseline_statistics: str = None, - supplied_baseline_constraints: str = None, + supplied_baseline_statistics: str | None = None, + supplied_baseline_constraints: str | None = None, wait: bool = True, - pipeline_name: str = None, - role: str = None, + pipeline_name: str | None = None, + role: str | None = None, ): ... def delete_model(self) -> None: ... def wait(self, logs: bool = True) -> None: ... diff --git a/stubs/sagemaker/sagemaker/tuner.pyi b/stubs/sagemaker/sagemaker/tuner.pyi index 533a4d896887..6b08bb50448e 100644 --- a/stubs/sagemaker/sagemaker/tuner.pyi +++ b/stubs/sagemaker/sagemaker/tuner.pyi @@ -34,7 +34,7 @@ class WarmStartTypes(Enum): class WarmStartConfig: type: Incomplete parents: Incomplete - def __init__(self, warm_start_type: WarmStartTypes, parents: Set[str | PipelineVariable]) -> None: ... + def __init__(self, warm_start_type: WarmStartTypes, parents: set[str | PipelineVariable]) -> None: ... @classmethod def from_job_desc(cls, warm_start_config): ... def to_input_req(self): ... @@ -60,8 +60,8 @@ class InstanceConfig: volume_size: Incomplete def __init__( self, - instance_count: int | PipelineVariable = None, - instance_type: str | PipelineVariable = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, volume_size: int | PipelineVariable = 30, ) -> None: ... @classmethod @@ -74,9 +74,9 @@ class TuningJobCompletionCriteriaConfig: target_objective_metric_value: Incomplete def __init__( self, - max_number_of_training_jobs_not_improving: int = None, - complete_on_convergence: bool = None, - target_objective_metric_value: float = None, + max_number_of_training_jobs_not_improving: int | None = None, + complete_on_convergence: bool | None = None, + target_objective_metric_value: float | None = None, ) -> None: ... @classmethod def from_job_desc(cls, completion_criteria_config): ... @@ -120,31 +120,31 @@ class HyperparameterTuner: self, estimator: EstimatorBase, objective_metric_name: str | PipelineVariable, - hyperparameter_ranges: Dict[str, ParameterRange], - metric_definitions: Optional[List[Dict[str, str | PipelineVariable]]] = None, + hyperparameter_ranges: dict[str, ParameterRange], + metric_definitions: list[dict[str, str | PipelineVariable]] | None = None, strategy: str | PipelineVariable = "Bayesian", objective_type: str | PipelineVariable = "Maximize", - max_jobs: int | PipelineVariable = None, + max_jobs: int | PipelineVariable | None = None, max_parallel_jobs: int | PipelineVariable = 1, - max_runtime_in_seconds: Optional[int | PipelineVariable] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - base_tuning_job_name: Optional[str] = None, - warm_start_config: Optional[WarmStartConfig] = None, - strategy_config: Optional[StrategyConfig] = None, - completion_criteria_config: Optional[TuningJobCompletionCriteriaConfig] = None, + max_runtime_in_seconds: int | PipelineVariable | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + base_tuning_job_name: str | None = None, + warm_start_config: WarmStartConfig | None = None, + strategy_config: StrategyConfig | None = None, + completion_criteria_config: TuningJobCompletionCriteriaConfig | None = None, early_stopping_type: str | PipelineVariable = "Off", - estimator_name: Optional[str] = None, - random_seed: Optional[int] = None, + estimator_name: str | None = None, + random_seed: int | None = None, autotune: bool = False, - hyperparameters_to_keep_static: Optional[List[str]] = None, + hyperparameters_to_keep_static: list[str] | None = None, ) -> None: ... - def override_resource_config(self, instance_configs: List[InstanceConfig, Dict[str, List[InstanceConfig]]]): ... + def override_resource_config(self, instance_configs: list[InstanceConfig, dict[str, list[InstanceConfig]]]): ... def fit( self, - inputs: Optional[str | Dict | List | TrainingInput | FileSystemInput | RecordSet | FileSystemRecordSet] = None, - job_name: Optional[str] = None, - include_cls_metadata: bool | Dict[str | bool] = False, - estimator_kwargs: Optional[Dict[str, dict]] = None, + inputs: str | dict | list | TrainingInput | FileSystemInput | RecordSet | FileSystemRecordSet | None = None, + job_name: str | None = None, + include_cls_metadata: bool | dict[str | bool] = False, + estimator_kwargs: dict[str, dict] | None = None, wait: bool = True, **kwargs, ): ... diff --git a/stubs/sagemaker/sagemaker/utilities/cache.pyi b/stubs/sagemaker/sagemaker/utilities/cache.pyi index 0229648f7e3d..f3db0aaf58b9 100644 --- a/stubs/sagemaker/sagemaker/utilities/cache.pyi +++ b/stubs/sagemaker/sagemaker/utilities/cache.pyi @@ -1,6 +1,7 @@ import datetime from _typeshed import Incomplete -from typing import Callable, Optional, TypeVar +from typing import Optional, TypeVar +from collections.abc import Callable KeyType = TypeVar("KeyType") ValType = TypeVar("ValType") @@ -20,5 +21,5 @@ class LRUCache: def __len__(self) -> int: ... def __contains__(self, key: KeyType) -> bool: ... def clear(self) -> None: ... - def get(self, key: KeyType, data_source_fallback: Optional[bool] = True) -> ValType: ... - def put(self, key: KeyType, value: Optional[ValType] = None) -> None: ... + def get(self, key: KeyType, data_source_fallback: bool | None = True) -> ValType: ... + def put(self, key: KeyType, value: ValType | None = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/utils.pyi b/stubs/sagemaker/sagemaker/utils.pyi index 1cd58a102e0f..24f3baa1fc63 100644 --- a/stubs/sagemaker/sagemaker/utils.pyi +++ b/stubs/sagemaker/sagemaker/utils.pyi @@ -21,8 +21,8 @@ def sagemaker_timestamp(): ... def sagemaker_short_timestamp(): ... def build_dict(key, value): ... def get_config_value(key_path, config): ... -def get_nested_value(dictionary: dict, nested_keys: List[str]): ... -def set_nested_value(dictionary: dict, nested_keys: List[str], value_to_set: object): ... +def get_nested_value(dictionary: dict, nested_keys: list[str]): ... +def set_nested_value(dictionary: dict, nested_keys: list[str], value_to_set: object): ... def get_short_version(framework_version): ... def secondary_training_status_changed(current_job_description, prev_job_description): ... def secondary_training_status_message(job_description, prev_description): ... @@ -71,21 +71,21 @@ def update_container_with_inference_params( container_list: Incomplete | None = None, ): ... def construct_container_object(obj, data_input_configuration, framework, framework_version, nearest_model_name): ... -def pop_out_unused_kwarg(arg_name: str, kwargs: dict, override_val: Optional[str] = None): ... +def pop_out_unused_kwarg(arg_name: str, kwargs: dict, override_val: str | None = None): ... def to_string(obj: object): ... def get_module(module_name): ... -def check_and_get_run_experiment_config(experiment_config: Optional[dict] = None) -> dict: ... +def check_and_get_run_experiment_config(experiment_config: dict | None = None) -> dict: ... def resolve_value_from_config( direct_input: Incomplete | None = None, - config_path: str = None, + config_path: str | None = None, default_value: Incomplete | None = None, sagemaker_session: Incomplete | None = None, - sagemaker_config: dict = None, + sagemaker_config: dict | None = None, ): ... -def get_sagemaker_config_value(sagemaker_session, key, sagemaker_config: dict = None): ... +def get_sagemaker_config_value(sagemaker_session, key, sagemaker_config: dict | None = None): ... def resolve_class_attribute_from_config( - clazz: Optional[type], - instance: Optional[object], + clazz: type | None, + instance: object | None, attribute: str, config_path: str, default_value: Incomplete | None = None, @@ -93,7 +93,7 @@ def resolve_class_attribute_from_config( ): ... def resolve_nested_dict_value_from_config( dictionary: dict, - nested_keys: List[str], + nested_keys: list[str], config_path: str, default_value: object = None, sagemaker_session: Incomplete | None = None, @@ -101,8 +101,8 @@ def resolve_nested_dict_value_from_config( def update_list_of_dicts_with_values_from_config( input_list, config_key_path, - required_key_paths: List[str] = None, - union_key_paths: List[List[str]] = None, + required_key_paths: list[str] | None = None, + union_key_paths: list[list[str]] | None = None, sagemaker_session: Incomplete | None = None, ): ... def update_nested_dictionary_with_values_from_config( diff --git a/stubs/sagemaker/sagemaker/workflow/_utils.pyi b/stubs/sagemaker/sagemaker/workflow/_utils.pyi index 3af080afd34c..e3bcc90526be 100644 --- a/stubs/sagemaker/sagemaker/workflow/_utils.pyi +++ b/stubs/sagemaker/sagemaker/workflow/_utils.pyi @@ -23,12 +23,12 @@ class _RepackModelStep(TrainingStep): role, model_data: str, entry_point: str, - display_name: str = None, - description: str = None, - source_dir: str = None, - dependencies: List = None, - depends_on: Optional[List[str | Step | "StepCollection"]] = None, - retry_policies: List[RetryPolicy] = None, + display_name: str | None = None, + description: str | None = None, + source_dir: str | None = None, + dependencies: list | None = None, + depends_on: list[str | Step | "StepCollection"] | None = None, + retry_policies: list[RetryPolicy] | None = None, subnets: Incomplete | None = None, security_group_ids: Incomplete | None = None, **kwargs, @@ -64,12 +64,12 @@ class _RegisterModelStep(ConfigurableRetryStep): def __init__( self, name: str, - step_args: Optional[dict] = None, - content_types: Optional[list] = None, - response_types: Optional[list] = None, - inference_instances: Optional[list] = None, - transform_instances: Optional[list] = None, - estimator: EstimatorBase = None, + step_args: dict | None = None, + content_types: list | None = None, + response_types: list | None = None, + inference_instances: list | None = None, + transform_instances: list | None = None, + estimator: EstimatorBase | None = None, model_data: Incomplete | None = None, model_package_group_name: Incomplete | None = None, model_metrics: Incomplete | None = None, @@ -77,10 +77,10 @@ class _RegisterModelStep(ConfigurableRetryStep): approval_status: str = "PendingManualApproval", image_uri: Incomplete | None = None, compile_model_family: Incomplete | None = None, - display_name: str = None, + display_name: str | None = None, description: Incomplete | None = None, - depends_on: Optional[List[str | Step | "StepCollection"]] = None, - retry_policies: Optional[List[RetryPolicy]] = None, + depends_on: list[str | Step | "StepCollection"] | None = None, + retry_policies: list[RetryPolicy] | None = None, tags: Incomplete | None = None, container_def_list: Incomplete | None = None, drift_check_baselines: Incomplete | None = None, diff --git a/stubs/sagemaker/sagemaker/workflow/automl_step.pyi b/stubs/sagemaker/sagemaker/workflow/automl_step.pyi index bdfa9e538b4b..354c3a16fe82 100644 --- a/stubs/sagemaker/sagemaker/workflow/automl_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/automl_step.pyi @@ -14,11 +14,11 @@ class AutoMLStep(ConfigurableRetryStep): self, name: str, step_args: _JobStepArguments, - display_name: str = None, - description: str = None, - cache_config: CacheConfig = None, - depends_on: Optional[List[str | Step | StepCollection]] = None, - retry_policies: List[RetryPolicy] = None, + display_name: str | None = None, + description: str | None = None, + cache_config: CacheConfig | None = None, + depends_on: list[str | Step | StepCollection] | None = None, + retry_policies: list[RetryPolicy] | None = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/callback_step.pyi b/stubs/sagemaker/sagemaker/workflow/callback_step.pyi index 8cdac4bd9204..93c67d38cbf6 100644 --- a/stubs/sagemaker/sagemaker/workflow/callback_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/callback_step.pyi @@ -16,7 +16,7 @@ class CallbackOutput: output_name: str output_type: CallbackOutputTypeEnum def to_request(self) -> RequestType: ... - def expr(self, step_name) -> Dict[str, str]: ... + def expr(self, step_name) -> dict[str, str]: ... def __init__(self, output_name, output_type) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -33,11 +33,11 @@ class CallbackStep(Step): name: str, sqs_queue_url: str, inputs: dict, - outputs: List[CallbackOutput], - display_name: str = None, - description: str = None, - cache_config: CacheConfig = None, - depends_on: Optional[List[str | Step | StepCollection]] = None, + outputs: list[CallbackOutput], + display_name: str | None = None, + description: str | None = None, + cache_config: CacheConfig | None = None, + depends_on: list[str | Step | StepCollection] | None = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi b/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi index 27ae9ed9f67f..fc3a8d89ee96 100644 --- a/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi @@ -20,7 +20,7 @@ class ClarifyCheckConfig(ABC): class DataBiasCheckConfig(ClarifyCheckConfig): data_bias_config: BiasConfig - methods: str | List[str] + methods: str | list[str] def __init__(self, data_config, kms_key, monitoring_analysis_config_uri, data_bias_config, methods) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -31,7 +31,7 @@ class ModelBiasCheckConfig(ClarifyCheckConfig): data_bias_config: BiasConfig model_config: ModelConfig model_predicted_label_config: ModelPredictedLabelConfig - methods: str | List[str] + methods: str | list[str] def __init__( self, data_config, @@ -76,12 +76,12 @@ class ClarifyCheckStep(Step): skip_check: bool | PipelineVariable = False, fail_on_violation: bool | PipelineVariable = True, register_new_baseline: bool | PipelineVariable = False, - model_package_group_name: str | PipelineVariable = None, - supplied_baseline_constraints: str | PipelineVariable = None, - display_name: str = None, - description: str = None, - cache_config: CacheConfig = None, - depends_on: Optional[List[str | Step | StepCollection]] = None, + model_package_group_name: str | PipelineVariable | None = None, + supplied_baseline_constraints: str | PipelineVariable | None = None, + display_name: str | None = None, + description: str | None = None, + cache_config: CacheConfig | None = None, + depends_on: list[str | Step | StepCollection] | None = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/condition_step.pyi b/stubs/sagemaker/sagemaker/workflow/condition_step.pyi index 05b6a09d72cd..b9a31b5916f8 100644 --- a/stubs/sagemaker/sagemaker/workflow/condition_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/condition_step.pyi @@ -15,12 +15,12 @@ class ConditionStep(Step): def __init__( self, name: str, - depends_on: Optional[List[str | Step | StepCollection]] = None, - display_name: str = None, - description: str = None, - conditions: List[Condition] = None, - if_steps: List[Step | StepCollection] = None, - else_steps: List[Step | StepCollection] = None, + depends_on: list[str | Step | StepCollection] | None = None, + display_name: str | None = None, + description: str | None = None, + conditions: list[Condition] | None = None, + if_steps: list[Step | StepCollection] | None = None, + else_steps: list[Step | StepCollection] | None = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/conditions.pyi b/stubs/sagemaker/sagemaker/workflow/conditions.pyi index a43ccf05a6e9..af497089df04 100644 --- a/stubs/sagemaker/sagemaker/workflow/conditions.pyi +++ b/stubs/sagemaker/sagemaker/workflow/conditions.pyi @@ -63,7 +63,7 @@ class ConditionIn(Condition): value: Incomplete in_values: Incomplete def __init__( - self, value: ConditionValueType | PrimitiveType, in_values: List[ConditionValueType | PrimitiveType] + self, value: ConditionValueType | PrimitiveType, in_values: list[ConditionValueType | PrimitiveType] ) -> None: ... def to_request(self) -> RequestType: ... @@ -74,9 +74,9 @@ class ConditionNot(Condition): class ConditionOr(Condition): conditions: Incomplete - def __init__(self, conditions: List[Condition] = None) -> None: ... + def __init__(self, conditions: list[Condition] | None = None) -> None: ... def to_request(self) -> RequestType: ... def primitive_or_expr( value: ExecutionVariable | Expression | PrimitiveType | Parameter | Properties, -) -> Dict[str | str, PrimitiveType]: ... +) -> dict[str | str, PrimitiveType]: ... diff --git a/stubs/sagemaker/sagemaker/workflow/emr_step.pyi b/stubs/sagemaker/sagemaker/workflow/emr_step.pyi index 46b355128d8d..ade3ce4a1958 100644 --- a/stubs/sagemaker/sagemaker/workflow/emr_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/emr_step.pyi @@ -10,7 +10,7 @@ class EMRStepConfig: args: Incomplete main_class: Incomplete properties: Incomplete - def __init__(self, jar, args: List[str] = None, main_class: str = None, properties: List[dict] = None) -> None: ... + def __init__(self, jar, args: list[str] | None = None, main_class: str | None = None, properties: list[dict] | None = None) -> None: ... def to_request(self) -> RequestType: ... INSTANCES: str @@ -34,10 +34,10 @@ class EMRStep(Step): description: str, cluster_id: str, step_config: EMRStepConfig, - depends_on: Optional[List[str | Step | StepCollection]] = None, - cache_config: CacheConfig = None, - cluster_config: Dict[str, Any] = None, - execution_role_arn: str = None, + depends_on: list[str | Step | StepCollection] | None = None, + cache_config: CacheConfig | None = None, + cluster_config: dict[str, Any] | None = None, + execution_role_arn: str | None = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/entities.pyi b/stubs/sagemaker/sagemaker/workflow/entities.pyi index 873279ea031a..7e92341711fc 100644 --- a/stubs/sagemaker/sagemaker/workflow/entities.pyi +++ b/stubs/sagemaker/sagemaker/workflow/entities.pyi @@ -4,7 +4,7 @@ from enum import EnumMeta from typing import Any, Dict, List PrimitiveType = str | int | bool | float | None -RequestType = Dict[str | Any, List[Dict[str, Any]]] +RequestType = dict[str | Any, list[dict[str, Any]]] class Entity(abc.ABC, metaclass=abc.ABCMeta): @abc.abstractmethod diff --git a/stubs/sagemaker/sagemaker/workflow/fail_step.pyi b/stubs/sagemaker/sagemaker/workflow/fail_step.pyi index 6b7a3bfeb2a6..8851ff031302 100644 --- a/stubs/sagemaker/sagemaker/workflow/fail_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/fail_step.pyi @@ -10,10 +10,10 @@ class FailStep(Step): def __init__( self, name: str, - error_message: str | PipelineVariable = None, - display_name: str = None, - description: str = None, - depends_on: Optional[List[str | Step | StepCollection]] = None, + error_message: str | PipelineVariable | None = None, + display_name: str | None = None, + description: str | None = None, + depends_on: list[str | Step | StepCollection] | None = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/functions.pyi b/stubs/sagemaker/sagemaker/workflow/functions.pyi index 93f2bb85c9b5..51f6fb198127 100644 --- a/stubs/sagemaker/sagemaker/workflow/functions.pyi +++ b/stubs/sagemaker/sagemaker/workflow/functions.pyi @@ -5,7 +5,7 @@ from sagemaker.workflow.properties import PropertyFile class Join(PipelineVariable): on: str - values: List + values: list def to_string(self) -> PipelineVariable: ... @property def expr(self): ... diff --git a/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi b/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi index c5495579961a..04ebc40ddf08 100644 --- a/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi @@ -17,7 +17,7 @@ class LambdaOutput: output_name: str output_type: LambdaOutputTypeEnum def to_request(self) -> RequestType: ... - def expr(self, step_name) -> Dict[str, str]: ... + def expr(self, step_name) -> dict[str, str]: ... def __init__(self, output_name, output_type) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -33,12 +33,12 @@ class LambdaStep(Step): self, name: str, lambda_func: Lambda, - display_name: str = None, - description: str = None, - inputs: dict = None, - outputs: List[LambdaOutput] = None, - cache_config: CacheConfig = None, - depends_on: Optional[List[str | Step | StepCollection]] = None, + display_name: str | None = None, + description: str | None = None, + inputs: dict | None = None, + outputs: list[LambdaOutput] | None = None, + cache_config: CacheConfig | None = None, + depends_on: list[str | Step | StepCollection] | None = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/model_step.pyi b/stubs/sagemaker/sagemaker/workflow/model_step.pyi index 43c1c82f118f..eacac28c06de 100644 --- a/stubs/sagemaker/sagemaker/workflow/model_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/model_step.pyi @@ -18,8 +18,8 @@ class ModelStep(StepCollection): self, name: str, step_args: _ModelStepArguments, - depends_on: Optional[List[str | Step | StepCollection]] = None, - retry_policies: Optional[List[RetryPolicy, Dict[str, List[RetryPolicy]]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, + depends_on: list[str | Step | StepCollection] | None = None, + retry_policies: list[RetryPolicy, dict[str, list[RetryPolicy]]] | None = None, + display_name: str | None = None, + description: str | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi b/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi index 09355569bd83..861f505aaa6d 100644 --- a/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi @@ -19,8 +19,8 @@ class MonitorBatchTransformStep(StepCollection): check_job_configuration: CheckJobConfig, monitor_before_transform: bool = False, fail_on_violation: bool | PipelineVariable = True, - supplied_baseline_statistics: str | PipelineVariable = None, - supplied_baseline_constraints: str | PipelineVariable = None, - display_name: Optional[str] = None, - description: Optional[str] = None, + supplied_baseline_statistics: str | PipelineVariable | None = None, + supplied_baseline_constraints: str | PipelineVariable | None = None, + display_name: str | None = None, + description: str | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/parameters.pyi b/stubs/sagemaker/sagemaker/workflow/parameters.pyi index c1ded205badc..d7cd12699f60 100644 --- a/stubs/sagemaker/sagemaker/workflow/parameters.pyi +++ b/stubs/sagemaker/sagemaker/workflow/parameters.pyi @@ -16,7 +16,7 @@ class ParameterTypeEnum(Enum, metaclass=DefaultEnumMeta): BOOLEAN: str FLOAT: str @property - def python_type(self) -> Type: ... + def python_type(self) -> type: ... class Parameter(PipelineVariable, Entity): name: str @@ -24,7 +24,7 @@ class Parameter(PipelineVariable, Entity): default_value: PrimitiveType def to_request(self) -> RequestType: ... @property - def expr(self) -> Dict[str, str]: ... + def expr(self) -> dict[str, str]: ... def __init__(self, name, parameter_type, default_value) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -35,13 +35,13 @@ ParameterBoolean: Incomplete class ParameterString(Parameter): enum_values: Incomplete - def __init__(self, name: str, default_value: str = None, enum_values: List[str] = None) -> None: ... + def __init__(self, name: str, default_value: str | None = None, enum_values: list[str] | None = None) -> None: ... def __hash__(self): ... def to_string(self) -> PipelineVariable: ... def to_request(self) -> RequestType: ... class ParameterInteger(Parameter): - def __init__(self, name: str, default_value: int = None) -> None: ... + def __init__(self, name: str, default_value: int | None = None) -> None: ... class ParameterFloat(Parameter): - def __init__(self, name: str, default_value: float = None) -> None: ... + def __init__(self, name: str, default_value: float | None = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/pipeline.pyi b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi index 7b3abf2b45a1..fa9acf68b455 100644 --- a/stubs/sagemaker/sagemaker/workflow/pipeline.pyi +++ b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi @@ -1,5 +1,6 @@ from _typeshed import Incomplete -from typing import Any, Dict, List, Optional, Sequence, Set +from typing import Any, Dict, List, Optional, Set +from collections.abc import Sequence from sagemaker.session import Session from sagemaker.workflow.entities import Entity, RequestType as RequestType @@ -21,49 +22,49 @@ class Pipeline(Entity): def __init__( self, name: str = "", - parameters: Optional[Sequence[Parameter]] = None, - pipeline_experiment_config: Optional[PipelineExperimentConfig] = ..., - steps: Optional[Sequence[Step | StepCollection]] = None, - sagemaker_session: Optional[Session] = None, + parameters: Sequence[Parameter] | None = None, + pipeline_experiment_config: PipelineExperimentConfig | None = ..., + steps: Sequence[Step | StepCollection] | None = None, + sagemaker_session: Session | None = None, ) -> None: ... def to_request(self) -> RequestType: ... def create( self, - role_arn: str = None, - description: str = None, - tags: List[Dict[str, str]] = None, - parallelism_config: ParallelismConfiguration = None, - ) -> Dict[str, Any]: ... - def describe(self) -> Dict[str, Any]: ... + role_arn: str | None = None, + description: str | None = None, + tags: list[dict[str, str]] | None = None, + parallelism_config: ParallelismConfiguration | None = None, + ) -> dict[str, Any]: ... + def describe(self) -> dict[str, Any]: ... def update( - self, role_arn: str = None, description: str = None, parallelism_config: ParallelismConfiguration = None - ) -> Dict[str, Any]: ... + self, role_arn: str | None = None, description: str | None = None, parallelism_config: ParallelismConfiguration | None = None + ) -> dict[str, Any]: ... def upsert( self, - role_arn: str = None, - description: str = None, - tags: List[Dict[str, str]] = None, - parallelism_config: ParallelismConfiguration = None, - ) -> Dict[str, Any]: ... - def delete(self) -> Dict[str, Any]: ... + role_arn: str | None = None, + description: str | None = None, + tags: list[dict[str, str]] | None = None, + parallelism_config: ParallelismConfiguration | None = None, + ) -> dict[str, Any]: ... + def delete(self) -> dict[str, Any]: ... def start( self, - parameters: Dict[str, str | bool | int | float] = None, - execution_display_name: str = None, - execution_description: str = None, - parallelism_config: ParallelismConfiguration = None, - selective_execution_config: SelectiveExecutionConfig = None, + parameters: dict[str, str | bool | int | float] | None = None, + execution_display_name: str | None = None, + execution_description: str | None = None, + parallelism_config: ParallelismConfiguration | None = None, + selective_execution_config: SelectiveExecutionConfig | None = None, ): ... def definition(self) -> str: ... def list_executions( - self, sort_by: str = None, sort_order: str = None, max_results: int = None, next_token: str = None - ) -> Dict[str, Any]: ... + self, sort_by: str | None = None, sort_order: str | None = None, max_results: int | None = None, next_token: str | None = None + ) -> dict[str, Any]: ... -def format_start_parameters(parameters: Dict[str, Any]) -> List[Dict[str, Any]]: ... +def format_start_parameters(parameters: dict[str, Any]) -> list[dict[str, Any]]: ... def interpolate( - request_obj: RequestType, callback_output_to_step_map: Dict[str, str], lambda_output_to_step_map: Dict[str, str] + request_obj: RequestType, callback_output_to_step_map: dict[str, str], lambda_output_to_step_map: dict[str, str] ) -> RequestType: ... -def update_args(args: Dict[str, Any], **kwargs): ... +def update_args(args: dict[str, Any], **kwargs): ... class _PipelineExecution: arn: str @@ -85,7 +86,7 @@ class PipelineGraph: @classmethod def from_pipeline(cls, pipeline: Pipeline): ... def is_cyclic(self) -> bool: ... - def get_steps_in_sub_dag(self, current_step: Step | StepCollection, sub_dag_steps: Set[str] = None) -> Set[str]: ... + def get_steps_in_sub_dag(self, current_step: Step | StepCollection, sub_dag_steps: set[str] | None = None) -> set[str]: ... stack: Incomplete def __iter__(self): ... def __next__(self) -> Step: ... diff --git a/stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi b/stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi index ed5fdb97b84e..76b5342d8c78 100644 --- a/stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi +++ b/stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi @@ -1,5 +1,6 @@ from _typeshed import Incomplete -from typing import Callable, Optional +from typing import Optional +from collections.abc import Callable from sagemaker.local import LocalSession from sagemaker.session import Session @@ -9,7 +10,7 @@ class _StepArguments: func: Incomplete func_args: Incomplete func_kwargs: Incomplete - def __init__(self, caller_name: str = None, func: Callable = None, *func_args, **func_kwargs) -> None: ... + def __init__(self, caller_name: str | None = None, func: Callable | None = None, *func_args, **func_kwargs) -> None: ... class _JobStepArguments(_StepArguments): args: Incomplete @@ -37,13 +38,13 @@ class PipelineSession(Session): sagemaker_client: Incomplete | None = None, default_bucket: Incomplete | None = None, settings=..., - sagemaker_config: dict = None, - default_bucket_prefix: str = None, + sagemaker_config: dict | None = None, + default_bucket_prefix: str | None = None, ) -> None: ... @property def context(self): ... @context.setter - def context(self, value: Optional[_StepArguments] = ...): ... + def context(self, value: _StepArguments | None = ...): ... def init_model_step_arguments(self, model) -> None: ... class LocalPipelineSession(LocalSession, PipelineSession): diff --git a/stubs/sagemaker/sagemaker/workflow/properties.pyi b/stubs/sagemaker/sagemaker/workflow/properties.pyi index 3dc9f41c2740..cd6e88e2a226 100644 --- a/stubs/sagemaker/sagemaker/workflow/properties.pyi +++ b/stubs/sagemaker/sagemaker/workflow/properties.pyi @@ -13,9 +13,9 @@ class Properties(PipelineVariable, metaclass=PropertiesMeta): def __init__( self, step_name: str, - path: str = None, - shape_name: str = None, - shape_names: List[str] = None, + path: str | None = None, + shape_name: str | None = None, + shape_names: list[str] | None = None, service_name: str = "sagemaker", ) -> None: ... @property @@ -24,13 +24,13 @@ class Properties(PipelineVariable, metaclass=PropertiesMeta): class PropertiesList(Properties): shape_name: Incomplete service_name: Incomplete - def __init__(self, step_name: str, path: str, shape_name: str = None, service_name: str = "sagemaker") -> None: ... + def __init__(self, step_name: str, path: str, shape_name: str | None = None, service_name: str = "sagemaker") -> None: ... def __getitem__(self, item: int | str): ... class PropertiesMap(Properties): shape_name: Incomplete service_name: Incomplete - def __init__(self, step_name: str, path: str, shape_name: str = None, service_name: str = "sagemaker") -> None: ... + def __init__(self, step_name: str, path: str, shape_name: str | None = None, service_name: str = "sagemaker") -> None: ... def __getitem__(self, item: int | str): ... class PropertyFile(Expression): @@ -38,7 +38,7 @@ class PropertyFile(Expression): output_name: str path: str @property - def expr(self) -> Dict[str, str]: ... + def expr(self) -> dict[str, str]: ... def __init__(self, name, output_name, path) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... diff --git a/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi b/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi index de6b5321cf23..d3fbc44efabc 100644 --- a/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi @@ -69,13 +69,13 @@ class QualityCheckStep(Step): skip_check: bool | PipelineVariable = False, fail_on_violation: bool | PipelineVariable = True, register_new_baseline: bool | PipelineVariable = False, - model_package_group_name: str | PipelineVariable = None, - supplied_baseline_statistics: str | PipelineVariable = None, - supplied_baseline_constraints: str | PipelineVariable = None, - display_name: str = None, - description: str = None, - cache_config: CacheConfig = None, - depends_on: Optional[List[str | Step | StepCollection]] = None, + model_package_group_name: str | PipelineVariable | None = None, + supplied_baseline_statistics: str | PipelineVariable | None = None, + supplied_baseline_constraints: str | PipelineVariable | None = None, + display_name: str | None = None, + description: str | None = None, + cache_config: CacheConfig | None = None, + depends_on: list[str | Step | StepCollection] | None = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/retry.pyi b/stubs/sagemaker/sagemaker/workflow/retry.pyi index 18c7335be39f..3d302161ce7f 100644 --- a/stubs/sagemaker/sagemaker/workflow/retry.pyi +++ b/stubs/sagemaker/sagemaker/workflow/retry.pyi @@ -38,11 +38,11 @@ class StepRetryPolicy(RetryPolicy): exception_types: Incomplete def __init__( self, - exception_types: List[StepExceptionTypeEnum], + exception_types: list[StepExceptionTypeEnum], backoff_rate: float = 2.0, interval_seconds: int = 1, - max_attempts: int = None, - expire_after_mins: int = None, + max_attempts: int | None = None, + expire_after_mins: int | None = None, ) -> None: ... def to_request(self) -> RequestType: ... def __hash__(self): ... @@ -51,12 +51,12 @@ class SageMakerJobStepRetryPolicy(RetryPolicy): exception_type_list: Incomplete def __init__( self, - exception_types: List[SageMakerJobExceptionTypeEnum] = None, - failure_reason_types: List[SageMakerJobExceptionTypeEnum] = None, + exception_types: list[SageMakerJobExceptionTypeEnum] | None = None, + failure_reason_types: list[SageMakerJobExceptionTypeEnum] | None = None, backoff_rate: float = 2.0, interval_seconds: int = 1, - max_attempts: int = None, - expire_after_mins: int = None, + max_attempts: int | None = None, + expire_after_mins: int | None = None, ) -> None: ... def to_request(self) -> RequestType: ... def __hash__(self): ... diff --git a/stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi b/stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi index f28ce9bdc5e2..7ce7d156a974 100644 --- a/stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi +++ b/stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi @@ -6,5 +6,5 @@ from sagemaker.workflow.entities import RequestType as RequestType class SelectiveExecutionConfig: source_pipeline_execution_arn: Incomplete selected_steps: Incomplete - def __init__(self, selected_steps: List[str], source_pipeline_execution_arn: str = None) -> None: ... + def __init__(self, selected_steps: list[str], source_pipeline_execution_arn: str | None = None) -> None: ... def to_request(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/step_collections.pyi b/stubs/sagemaker/sagemaker/workflow/step_collections.pyi index 1063d4816c2d..51014766097e 100644 --- a/stubs/sagemaker/sagemaker/workflow/step_collections.pyi +++ b/stubs/sagemaker/sagemaker/workflow/step_collections.pyi @@ -10,8 +10,8 @@ from sagemaker.workflow.steps import Step class StepCollection: name: str - steps: List[Step] - def request_dicts(self) -> List[RequestType]: ... + steps: list[Step] + def request_dicts(self) -> list[RequestType]: ... @property def properties(self): ... def __init__(self, name, steps) -> None: ... @@ -32,11 +32,11 @@ class RegisterModel(StepCollection): response_types, inference_instances: Incomplete | None = None, transform_instances: Incomplete | None = None, - estimator: EstimatorBase = None, + estimator: EstimatorBase | None = None, model_data: Incomplete | None = None, - depends_on: Optional[List[str | Step | StepCollection]] = None, - repack_model_step_retry_policies: List[RetryPolicy] = None, - register_model_step_retry_policies: List[RetryPolicy] = None, + depends_on: list[str | Step | StepCollection] | None = None, + repack_model_step_retry_policies: list[RetryPolicy] | None = None, + register_model_step_retry_policies: list[RetryPolicy] | None = None, model_package_group_name: Incomplete | None = None, model_metrics: Incomplete | None = None, approval_status: Incomplete | None = None, @@ -45,7 +45,7 @@ class RegisterModel(StepCollection): display_name: Incomplete | None = None, description: Incomplete | None = None, tags: Incomplete | None = None, - model: Model | PipelineModel = None, + model: Model | PipelineModel | None = None, drift_check_baselines: Incomplete | None = None, customer_metadata_properties: Incomplete | None = None, domain: Incomplete | None = None, @@ -70,8 +70,8 @@ class EstimatorTransformer(StepCollection): instance_count, instance_type, transform_inputs, - description: str = None, - display_name: str = None, + description: str | None = None, + display_name: str | None = None, image_uri: Incomplete | None = None, predictor_cls: Incomplete | None = None, env: Incomplete | None = None, @@ -84,9 +84,9 @@ class EstimatorTransformer(StepCollection): max_payload: Incomplete | None = None, tags: Incomplete | None = None, volume_kms_key: Incomplete | None = None, - depends_on: Optional[List[str | Step | StepCollection]] = None, - repack_model_step_retry_policies: List[RetryPolicy] = None, - model_step_retry_policies: List[RetryPolicy] = None, - transform_step_retry_policies: List[RetryPolicy] = None, + depends_on: list[str | Step | StepCollection] | None = None, + repack_model_step_retry_policies: list[RetryPolicy] | None = None, + model_step_retry_policies: list[RetryPolicy] | None = None, + transform_step_retry_policies: list[RetryPolicy] | None = None, **kwargs, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/steps.pyi b/stubs/sagemaker/sagemaker/workflow/steps.pyi index ac3f5bbfc615..41aa5d7896f6 100644 --- a/stubs/sagemaker/sagemaker/workflow/steps.pyi +++ b/stubs/sagemaker/sagemaker/workflow/steps.pyi @@ -34,10 +34,10 @@ class StepTypeEnum(Enum, metaclass=DefaultEnumMeta): class Step(Entity, metaclass=abc.ABCMeta): name: str - display_name: Optional[str] - description: Optional[str] + display_name: str | None + description: str | None step_type: StepTypeEnum - depends_on: Optional[List[str | "Step" | "StepCollection"]] + depends_on: list[str | "Step" | "StepCollection"] | None @property @abc.abstractmethod def arguments(self) -> RequestType: ... @@ -47,9 +47,9 @@ class Step(Entity, metaclass=abc.ABCMeta): @abc.abstractmethod def properties(self): ... def to_request(self) -> RequestType: ... - def add_depends_on(self, step_names: List[str | "Step" | "StepCollection"]): ... + def add_depends_on(self, step_names: list[str | "Step" | "StepCollection"]): ... @property - def ref(self) -> Dict[str, str]: ... + def ref(self) -> dict[str, str]: ... def __init__(self, name, display_name, description, step_type, depends_on) -> None: ... def __lt__(self, other): ... def __le__(self, other): ... @@ -73,10 +73,10 @@ class ConfigurableRetryStep(Step, metaclass=abc.ABCMeta): self, name: str, step_type: StepTypeEnum, - display_name: str = None, - description: str = None, - depends_on: Optional[List[str | Step | "StepCollection"]] = None, - retry_policies: List[RetryPolicy] = None, + display_name: str | None = None, + description: str | None = None, + depends_on: list[str | Step | "StepCollection"] | None = None, + retry_policies: list[RetryPolicy] | None = None, ) -> None: ... def add_retry_policy(self, retry_policy: RetryPolicy): ... def to_request(self) -> RequestType: ... @@ -90,14 +90,14 @@ class TrainingStep(ConfigurableRetryStep): def __init__( self, name: str, - step_args: _JobStepArguments = None, - estimator: EstimatorBase = None, - display_name: str = None, - description: str = None, - inputs: TrainingInput | dict | str | FileSystemInput = None, - cache_config: CacheConfig = None, - depends_on: Optional[List[str | Step | "StepCollection"]] = None, - retry_policies: List[RetryPolicy] = None, + step_args: _JobStepArguments | None = None, + estimator: EstimatorBase | None = None, + display_name: str | None = None, + description: str | None = None, + inputs: TrainingInput | dict | str | FileSystemInput | None = None, + cache_config: CacheConfig | None = None, + depends_on: list[str | Step | "StepCollection"] | None = None, + retry_policies: list[RetryPolicy] | None = None, ) -> None: ... @property def arguments(self) -> RequestType: ... @@ -112,13 +112,13 @@ class CreateModelStep(ConfigurableRetryStep): def __init__( self, name: str, - step_args: Optional[dict] = None, - model: Optional[Model | PipelineModel] = None, - inputs: Optional[CreateModelInput] = None, - depends_on: Optional[List[str | Step | "StepCollection"]] = None, - retry_policies: Optional[List[RetryPolicy]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, + step_args: dict | None = None, + model: Model | PipelineModel | None = None, + inputs: CreateModelInput | None = None, + depends_on: list[str | Step | "StepCollection"] | None = None, + retry_policies: list[RetryPolicy] | None = None, + display_name: str | None = None, + description: str | None = None, ) -> None: ... @property def arguments(self) -> RequestType: ... @@ -133,14 +133,14 @@ class TransformStep(ConfigurableRetryStep): def __init__( self, name: str, - step_args: _JobStepArguments = None, - transformer: Transformer = None, - inputs: TransformInput = None, - display_name: str = None, - description: str = None, - cache_config: CacheConfig = None, - depends_on: Optional[List[str | Step | "StepCollection"]] = None, - retry_policies: List[RetryPolicy] = None, + step_args: _JobStepArguments | None = None, + transformer: Transformer | None = None, + inputs: TransformInput | None = None, + display_name: str | None = None, + description: str | None = None, + cache_config: CacheConfig | None = None, + depends_on: list[str | Step | "StepCollection"] | None = None, + retry_policies: list[RetryPolicy] | None = None, ) -> None: ... @property def arguments(self) -> RequestType: ... @@ -162,18 +162,18 @@ class ProcessingStep(ConfigurableRetryStep): def __init__( self, name: str, - step_args: _JobStepArguments = None, - processor: Processor = None, - display_name: str = None, - description: str = None, - inputs: List[ProcessingInput] = None, - outputs: List[ProcessingOutput] = None, - job_arguments: List[str] = None, - code: str = None, - property_files: List[PropertyFile] = None, - cache_config: CacheConfig = None, - depends_on: Optional[List[str | Step | "StepCollection"]] = None, - retry_policies: List[RetryPolicy] = None, + step_args: _JobStepArguments | None = None, + processor: Processor | None = None, + display_name: str | None = None, + description: str | None = None, + inputs: list[ProcessingInput] | None = None, + outputs: list[ProcessingOutput] | None = None, + job_arguments: list[str] | None = None, + code: str | None = None, + property_files: list[PropertyFile] | None = None, + cache_config: CacheConfig | None = None, + depends_on: list[str | Step | "StepCollection"] | None = None, + retry_policies: list[RetryPolicy] | None = None, kms_key: Incomplete | None = None, ) -> None: ... @property @@ -191,15 +191,15 @@ class TuningStep(ConfigurableRetryStep): def __init__( self, name: str, - step_args: _JobStepArguments = None, - tuner: HyperparameterTuner = None, - display_name: str = None, - description: str = None, + step_args: _JobStepArguments | None = None, + tuner: HyperparameterTuner | None = None, + display_name: str | None = None, + description: str | None = None, inputs: Incomplete | None = None, - job_arguments: List[str] = None, - cache_config: CacheConfig = None, - depends_on: Optional[List[str | Step | "StepCollection"]] = None, - retry_policies: List[RetryPolicy] = None, + job_arguments: list[str] | None = None, + cache_config: CacheConfig | None = None, + depends_on: list[str | Step | "StepCollection"] | None = None, + retry_policies: list[RetryPolicy] | None = None, ) -> None: ... @property def arguments(self) -> RequestType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/utilities.pyi b/stubs/sagemaker/sagemaker/workflow/utilities.pyi index 67660d63c110..f0e6738af4ac 100644 --- a/stubs/sagemaker/sagemaker/workflow/utilities.pyi +++ b/stubs/sagemaker/sagemaker/workflow/utilities.pyi @@ -1,5 +1,6 @@ from _typeshed import Incomplete -from typing import List, Sequence, Set +from typing import List, Set +from collections.abc import Sequence from sagemaker.workflow.entities import Entity, RequestType as RequestType from sagemaker.workflow.pipeline_context import _StepArguments @@ -7,16 +8,16 @@ from sagemaker.workflow.pipeline_context import _StepArguments logger: Incomplete BUF_SIZE: int -def list_to_request(entities: Sequence[Entity | "StepCollection"]) -> List[RequestType]: ... +def list_to_request(entities: Sequence[Entity | "StepCollection"]) -> list[RequestType]: ... def build_steps(steps: Sequence[Entity], pipeline_name: str): ... def get_code_hash(step: Entity) -> str: ... -def get_processing_dependencies(dependency_args: List[List[str]]) -> List[str]: ... -def get_processing_code_hash(code: str, source_dir: str, dependencies: List[str]) -> str: ... -def get_training_code_hash(entry_point: str, source_dir: str, dependencies: List[str]) -> str: ... +def get_processing_dependencies(dependency_args: list[list[str]]) -> list[str]: ... +def get_processing_code_hash(code: str, source_dir: str, dependencies: list[str]) -> str: ... +def get_training_code_hash(entry_point: str, source_dir: str, dependencies: list[str]) -> str: ... def get_config_hash(step: Entity): ... def hash_object(obj) -> str: ... def hash_file(path: str) -> str: ... -def hash_files_or_dirs(paths: List[str]) -> str: ... -def validate_step_args_input(step_args: _StepArguments, expected_caller: Set[str], error_message: str): ... +def hash_files_or_dirs(paths: list[str]) -> str: ... +def validate_step_args_input(step_args: _StepArguments, expected_caller: set[str], error_message: str): ... def override_pipeline_parameter_var(func): ... def execute_job_functions(step_args: _StepArguments): ... diff --git a/stubs/sagemaker/sagemaker/wrangler/ingestion.pyi b/stubs/sagemaker/sagemaker/wrangler/ingestion.pyi index 796076df0ac0..70e6186712cb 100644 --- a/stubs/sagemaker/sagemaker/wrangler/ingestion.pyi +++ b/stubs/sagemaker/sagemaker/wrangler/ingestion.pyi @@ -8,11 +8,11 @@ def generate_data_ingestion_flow_from_s3_input( s3_content_type: str = "csv", s3_has_header: bool = False, operator_version: str = "0.1", - schema: Dict = None, + schema: dict | None = None, ): ... def generate_data_ingestion_flow_from_athena_dataset_definition( - input_name: str, athena_dataset_definition: AthenaDatasetDefinition, operator_version: str = "0.1", schema: Dict = None + input_name: str, athena_dataset_definition: AthenaDatasetDefinition, operator_version: str = "0.1", schema: dict | None = None ): ... def generate_data_ingestion_flow_from_redshift_dataset_definition( - input_name: str, redshift_dataset_definition: RedshiftDatasetDefinition, operator_version: str = "0.1", schema: Dict = None + input_name: str, redshift_dataset_definition: RedshiftDatasetDefinition, operator_version: str = "0.1", schema: dict | None = None ): ... diff --git a/stubs/sagemaker/sagemaker/wrangler/processing.pyi b/stubs/sagemaker/sagemaker/wrangler/processing.pyi index de9051a81302..87fcd0698946 100644 --- a/stubs/sagemaker/sagemaker/wrangler/processing.pyi +++ b/stubs/sagemaker/sagemaker/wrangler/processing.pyi @@ -10,17 +10,17 @@ class DataWranglerProcessor(Processor): sagemaker_session: Incomplete def __init__( self, - role: str = None, - data_wrangler_flow_source: str = None, - instance_count: int = None, - instance_type: str = None, + role: str | None = None, + data_wrangler_flow_source: str | None = None, + instance_count: int | None = None, + instance_type: str | None = None, volume_size_in_gb: int = 30, - volume_kms_key: str = None, - output_kms_key: str = None, - max_runtime_in_seconds: int = None, - base_job_name: str = None, - sagemaker_session: Session = None, - env: Dict[str, str] = None, - tags: List[dict] = None, - network_config: NetworkConfig = None, + volume_kms_key: str | None = None, + output_kms_key: str | None = None, + max_runtime_in_seconds: int | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str] | None = None, + tags: list[dict] | None = None, + network_config: NetworkConfig | None = None, ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/xgboost/estimator.pyi b/stubs/sagemaker/sagemaker/xgboost/estimator.pyi index 993065991227..cc2c005e35f0 100644 --- a/stubs/sagemaker/sagemaker/xgboost/estimator.pyi +++ b/stubs/sagemaker/sagemaker/xgboost/estimator.pyi @@ -14,11 +14,11 @@ class XGBoost(Framework): self, entry_point: str | PipelineVariable, framework_version: str, - source_dir: Optional[str | PipelineVariable] = None, - hyperparameters: Optional[Dict[str, str | PipelineVariable]] = None, + source_dir: str | PipelineVariable | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, py_version: str = "py3", - image_uri: Optional[str | PipelineVariable] = None, - image_uri_region: Optional[str] = None, + image_uri: str | PipelineVariable | None = None, + image_uri_region: str | None = None, **kwargs, ) -> None: ... def create_model( diff --git a/stubs/sagemaker/sagemaker/xgboost/model.pyi b/stubs/sagemaker/sagemaker/xgboost/model.pyi index f90f3f040839..fe811b899867 100644 --- a/stubs/sagemaker/sagemaker/xgboost/model.pyi +++ b/stubs/sagemaker/sagemaker/xgboost/model.pyi @@ -20,39 +20,39 @@ class XGBoostModel(FrameworkModel): def __init__( self, model_data: str | PipelineVariable, - role: str = None, - entry_point: str = None, - framework_version: str = None, - image_uri: Optional[str | PipelineVariable] = None, + role: str | None = None, + entry_point: str | None = None, + framework_version: str | None = None, + image_uri: str | PipelineVariable | None = None, py_version: str = "py3", predictor_cls: callable = ..., - model_server_workers: Optional[int | PipelineVariable] = None, + model_server_workers: int | PipelineVariable | None = None, **kwargs, ) -> None: ... image_uri: Incomplete def register( self, - content_types: List[str | PipelineVariable], - response_types: List[str | PipelineVariable], - inference_instances: Optional[List[str | PipelineVariable]] = None, - transform_instances: Optional[List[str | PipelineVariable]] = None, - model_package_name: Optional[str | PipelineVariable] = None, - model_package_group_name: Optional[str | PipelineVariable] = None, - image_uri: Optional[str | PipelineVariable] = None, - model_metrics: Optional[ModelMetrics] = None, - metadata_properties: Optional[MetadataProperties] = None, + content_types: list[str | PipelineVariable], + response_types: list[str | PipelineVariable], + inference_instances: list[str | PipelineVariable] | None = None, + transform_instances: list[str | PipelineVariable] | None = None, + model_package_name: str | PipelineVariable | None = None, + model_package_group_name: str | PipelineVariable | None = None, + image_uri: str | PipelineVariable | None = None, + model_metrics: ModelMetrics | None = None, + metadata_properties: MetadataProperties | None = None, marketplace_cert: bool = False, - approval_status: Optional[str | PipelineVariable] = None, - description: Optional[str] = None, - drift_check_baselines: Optional[DriftCheckBaselines] = None, - customer_metadata_properties: Optional[Dict[str, str | PipelineVariable]] = None, - domain: Optional[str | PipelineVariable] = None, - sample_payload_url: Optional[str | PipelineVariable] = None, - task: Optional[str | PipelineVariable] = None, - framework: Optional[str | PipelineVariable] = None, - framework_version: Optional[str | PipelineVariable] = None, - nearest_model_name: Optional[str | PipelineVariable] = None, - data_input_configuration: Optional[str | PipelineVariable] = None, + approval_status: str | PipelineVariable | None = None, + description: str | None = None, + drift_check_baselines: DriftCheckBaselines | None = None, + customer_metadata_properties: dict[str, str | PipelineVariable] | None = None, + domain: str | PipelineVariable | None = None, + sample_payload_url: str | PipelineVariable | None = None, + task: str | PipelineVariable | None = None, + framework: str | PipelineVariable | None = None, + framework_version: str | PipelineVariable | None = None, + nearest_model_name: str | PipelineVariable | None = None, + data_input_configuration: str | PipelineVariable | None = None, ): ... def prepare_container_def( self, diff --git a/stubs/sagemaker/sagemaker/xgboost/processing.pyi b/stubs/sagemaker/sagemaker/xgboost/processing.pyi index e73715950366..28e6e9e19031 100644 --- a/stubs/sagemaker/sagemaker/xgboost/processing.pyi +++ b/stubs/sagemaker/sagemaker/xgboost/processing.pyi @@ -11,20 +11,20 @@ class XGBoostProcessor(FrameworkProcessor): def __init__( self, framework_version: str, - role: str = None, - instance_count: int | PipelineVariable = None, - instance_type: str | PipelineVariable = None, + role: str | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, py_version: str = "py3", - image_uri: Optional[str | PipelineVariable] = None, - command: Optional[List[str]] = None, + image_uri: str | PipelineVariable | None = None, + command: list[str] | None = None, volume_size_in_gb: int | PipelineVariable = 30, - volume_kms_key: Optional[str | PipelineVariable] = None, - output_kms_key: Optional[str | PipelineVariable] = None, - code_location: Optional[str] = None, - max_runtime_in_seconds: Optional[int | PipelineVariable] = None, - base_job_name: Optional[str] = None, - sagemaker_session: Optional[Session] = None, - env: Optional[Dict[str, str | PipelineVariable]] = None, - tags: Optional[List[Dict[str, str | PipelineVariable]]] = None, - network_config: Optional[NetworkConfig] = None, + volume_kms_key: str | PipelineVariable | None = None, + output_kms_key: str | PipelineVariable | None = None, + code_location: str | None = None, + max_runtime_in_seconds: int | PipelineVariable | None = None, + base_job_name: str | None = None, + sagemaker_session: Session | None = None, + env: dict[str, str | PipelineVariable] | None = None, + tags: list[dict[str, str | PipelineVariable]] | None = None, + network_config: NetworkConfig | None = None, ) -> None: ... From 5a7161043260ebb5b40d0fd271ec5acc4a6e324b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 25 Jun 2023 14:08:50 +0000 Subject: [PATCH 06/10] [pre-commit.ci] auto fixes from pre-commit.com hooks --- stubs/sagemaker/sagemaker/accept_types.pyi | 2 -- stubs/sagemaker/sagemaker/algorithm.pyi | 1 - .../sagemaker/sagemaker/amazon/amazon_estimator.pyi | 1 - .../sagemaker/amazon/factorization_machines.pyi | 7 +------ stubs/sagemaker/sagemaker/amazon/ipinsights.pyi | 7 +------ stubs/sagemaker/sagemaker/amazon/kmeans.pyi | 7 +------ stubs/sagemaker/sagemaker/amazon/knn.pyi | 7 +------ stubs/sagemaker/sagemaker/amazon/lda.pyi | 7 +------ stubs/sagemaker/sagemaker/amazon/linear_learner.pyi | 7 +------ stubs/sagemaker/sagemaker/amazon/ntm.pyi | 7 +------ stubs/sagemaker/sagemaker/amazon/object2vec.pyi | 8 +------- stubs/sagemaker/sagemaker/amazon/pca.pyi | 7 +------ .../sagemaker/sagemaker/amazon/randomcutforest.pyi | 7 +------ stubs/sagemaker/sagemaker/automl/automl.pyi | 1 - stubs/sagemaker/sagemaker/base_predictor.pyi | 2 +- stubs/sagemaker/sagemaker/chainer/estimator.pyi | 1 - stubs/sagemaker/sagemaker/chainer/model.pyi | 1 - stubs/sagemaker/sagemaker/clarify.pyi | 2 +- stubs/sagemaker/sagemaker/collection.pyi | 1 - stubs/sagemaker/sagemaker/config/config.pyi | 1 - stubs/sagemaker/sagemaker/content_types.pyi | 2 -- stubs/sagemaker/sagemaker/debugger/debugger.pyi | 1 - .../sagemaker/debugger/profiler_config.pyi | 1 - stubs/sagemaker/sagemaker/deserializers.pyi | 2 -- stubs/sagemaker/sagemaker/djl_inference/model.pyi | 2 +- stubs/sagemaker/sagemaker/drift_check_baselines.pyi | 1 - stubs/sagemaker/sagemaker/environment_variables.pyi | 1 - stubs/sagemaker/sagemaker/estimator.pyi | 1 - stubs/sagemaker/sagemaker/experiments/_utils.pyi | 2 -- stubs/sagemaker/sagemaker/experiments/run.pyi | 2 +- .../explainer/clarify_explainer_config.pyi | 1 - .../sagemaker/explainer/explainer_config.pyi | 1 - .../sagemaker/feature_store/dataset_builder.pyi | 5 +++-- .../sagemaker/feature_store/feature_definition.pyi | 2 +- .../sagemaker/feature_store/feature_group.pyi | 4 ++-- .../sagemaker/feature_store/feature_store.pyi | 2 +- stubs/sagemaker/sagemaker/feature_store/inputs.pyi | 2 +- stubs/sagemaker/sagemaker/fw_utils.pyi | 2 +- stubs/sagemaker/sagemaker/huggingface/estimator.pyi | 1 - stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi | 2 -- stubs/sagemaker/sagemaker/huggingface/model.pyi | 1 - .../sagemaker/sagemaker/huggingface/processing.pyi | 1 - stubs/sagemaker/sagemaker/hyperparameters.pyi | 1 - .../inference_recommender_mixin.pyi | 1 - stubs/sagemaker/sagemaker/inputs.pyi | 6 ++++-- stubs/sagemaker/sagemaker/instance_types.pyi | 1 - stubs/sagemaker/sagemaker/jumpstart/accessors.pyi | 2 +- stubs/sagemaker/sagemaker/jumpstart/cache.pyi | 1 - stubs/sagemaker/sagemaker/jumpstart/constants.pyi | 1 - stubs/sagemaker/sagemaker/jumpstart/estimator.pyi | 1 - stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi | 1 - .../sagemaker/jumpstart/factory/estimator.pyi | 1 - .../sagemaker/sagemaker/jumpstart/factory/model.pyi | 1 - stubs/sagemaker/sagemaker/jumpstart/filters.pyi | 2 +- stubs/sagemaker/sagemaker/jumpstart/model.pyi | 1 - .../sagemaker/jumpstart/notebook_utils.pyi | 2 -- stubs/sagemaker/sagemaker/jumpstart/types.pyi | 2 +- stubs/sagemaker/sagemaker/jumpstart/utils.pyi | 11 +++-------- stubs/sagemaker/sagemaker/jumpstart/validators.pyi | 2 +- stubs/sagemaker/sagemaker/lineage/action.pyi | 3 +-- stubs/sagemaker/sagemaker/lineage/artifact.pyi | 3 +-- stubs/sagemaker/sagemaker/lineage/association.pyi | 9 ++++++--- stubs/sagemaker/sagemaker/lineage/context.pyi | 3 +-- .../sagemaker/lineage/lineage_trial_component.pyi | 1 - stubs/sagemaker/sagemaker/lineage/query.pyi | 6 ++++-- stubs/sagemaker/sagemaker/lineage/visualizer.pyi | 2 -- stubs/sagemaker/sagemaker/local/pipeline.pyi | 1 - stubs/sagemaker/sagemaker/metadata_properties.pyi | 1 - stubs/sagemaker/sagemaker/metric_definitions.pyi | 1 - stubs/sagemaker/sagemaker/model.pyi | 1 - stubs/sagemaker/sagemaker/model_card/helpers.pyi | 2 +- stubs/sagemaker/sagemaker/model_card/model_card.pyi | 10 ++-------- .../sagemaker/model_monitor/model_monitoring.pyi | 1 - stubs/sagemaker/sagemaker/model_uris.pyi | 1 - stubs/sagemaker/sagemaker/multidatamodel.pyi | 1 - stubs/sagemaker/sagemaker/mxnet/estimator.pyi | 1 - stubs/sagemaker/sagemaker/mxnet/model.pyi | 1 - stubs/sagemaker/sagemaker/mxnet/processing.pyi | 2 -- stubs/sagemaker/sagemaker/network.pyi | 1 - stubs/sagemaker/sagemaker/pipeline.pyi | 1 - stubs/sagemaker/sagemaker/predictor.pyi | 2 -- stubs/sagemaker/sagemaker/processing.pyi | 2 +- stubs/sagemaker/sagemaker/pytorch/estimator.pyi | 1 - stubs/sagemaker/sagemaker/pytorch/model.pyi | 1 - stubs/sagemaker/sagemaker/pytorch/processing.pyi | 2 -- .../sagemaker/sagemaker/remote_function/client.pyi | 2 +- stubs/sagemaker/sagemaker/remote_function/job.pyi | 1 - .../sagemaker/remote_function/spark_config.pyi | 2 -- stubs/sagemaker/sagemaker/rl/estimator.pyi | 1 - stubs/sagemaker/sagemaker/s3_utils.pyi | 1 - stubs/sagemaker/sagemaker/script_uris.pyi | 1 - stubs/sagemaker/sagemaker/serializers.pyi | 2 -- .../serverless/serverless_inference_config.pyi | 1 - stubs/sagemaker/sagemaker/session.pyi | 10 ++++++++-- stubs/sagemaker/sagemaker/sklearn/estimator.pyi | 1 - stubs/sagemaker/sagemaker/sklearn/model.pyi | 1 - stubs/sagemaker/sagemaker/sklearn/processing.pyi | 2 -- stubs/sagemaker/sagemaker/spark/processing.pyi | 1 - stubs/sagemaker/sagemaker/sparkml/model.pyi | 1 - stubs/sagemaker/sagemaker/tensorflow/estimator.pyi | 1 - stubs/sagemaker/sagemaker/tensorflow/model.pyi | 1 - stubs/sagemaker/sagemaker/tensorflow/processing.pyi | 2 -- stubs/sagemaker/sagemaker/transformer.pyi | 1 - stubs/sagemaker/sagemaker/tuner.pyi | 1 - stubs/sagemaker/sagemaker/utilities/cache.pyi | 2 +- stubs/sagemaker/sagemaker/utils.pyi | 2 +- stubs/sagemaker/sagemaker/workflow/_utils.pyi | 1 - stubs/sagemaker/sagemaker/workflow/automl_step.pyi | 1 - .../sagemaker/sagemaker/workflow/callback_step.pyi | 1 - .../sagemaker/workflow/clarify_check_step.pyi | 1 - .../sagemaker/sagemaker/workflow/condition_step.pyi | 1 - stubs/sagemaker/sagemaker/workflow/conditions.pyi | 1 - stubs/sagemaker/sagemaker/workflow/emr_step.pyi | 6 ++++-- stubs/sagemaker/sagemaker/workflow/entities.pyi | 2 +- stubs/sagemaker/sagemaker/workflow/fail_step.pyi | 1 - stubs/sagemaker/sagemaker/workflow/functions.pyi | 2 -- stubs/sagemaker/sagemaker/workflow/lambda_step.pyi | 1 - stubs/sagemaker/sagemaker/workflow/model_step.pyi | 1 - .../workflow/monitor_batch_transform_step.pyi | 1 - stubs/sagemaker/sagemaker/workflow/parameters.pyi | 1 - stubs/sagemaker/sagemaker/workflow/pipeline.pyi | 13 ++++++++++--- .../sagemaker/workflow/pipeline_context.pyi | 1 - stubs/sagemaker/sagemaker/workflow/properties.pyi | 1 - .../sagemaker/workflow/quality_check_step.pyi | 1 - stubs/sagemaker/sagemaker/workflow/retry.pyi | 1 - .../workflow/selective_execution_config.pyi | 1 - .../sagemaker/workflow/step_collections.pyi | 1 - stubs/sagemaker/sagemaker/workflow/steps.pyi | 1 - stubs/sagemaker/sagemaker/workflow/utilities.pyi | 1 - stubs/sagemaker/sagemaker/wrangler/ingestion.pyi | 7 ++++--- stubs/sagemaker/sagemaker/wrangler/processing.pyi | 1 - stubs/sagemaker/sagemaker/xgboost/estimator.pyi | 1 - stubs/sagemaker/sagemaker/xgboost/model.pyi | 1 - stubs/sagemaker/sagemaker/xgboost/processing.pyi | 2 -- 134 files changed, 81 insertions(+), 230 deletions(-) diff --git a/stubs/sagemaker/sagemaker/accept_types.pyi b/stubs/sagemaker/sagemaker/accept_types.pyi index 08282c899c31..1ffd31dba8cd 100644 --- a/stubs/sagemaker/sagemaker/accept_types.pyi +++ b/stubs/sagemaker/sagemaker/accept_types.pyi @@ -1,5 +1,3 @@ -from typing import List, Optional - def retrieve_options( region: str | None = None, model_id: str | None = None, diff --git a/stubs/sagemaker/sagemaker/algorithm.pyi b/stubs/sagemaker/sagemaker/algorithm.pyi index e93a25e20479..a88be9ce9d38 100644 --- a/stubs/sagemaker/sagemaker/algorithm.pyi +++ b/stubs/sagemaker/sagemaker/algorithm.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker.estimator import EstimatorBase from sagemaker.inputs import FileSystemInput, TrainingInput diff --git a/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi b/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi index 8f491abc4744..1a6f4c6addf5 100644 --- a/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi +++ b/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi @@ -1,6 +1,5 @@ import abc from _typeshed import Incomplete -from typing import Dict, Optional from sagemaker.amazon.hyperparameter import Hyperparameter as hp from sagemaker.estimator import EstimatorBase diff --git a/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi b/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi index 68668c6f3207..880922adef98 100644 --- a/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi +++ b/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp @@ -73,9 +72,5 @@ class FactorizationMachinesPredictor(Predictor): class FactorizationMachinesModel(Model): def __init__( - self, - model_data: str | PipelineVariable, - role: str | None = None, - sagemaker_session: Session | None = None, - **kwargs, + self, model_data: str | PipelineVariable, role: str | None = None, sagemaker_session: Session | None = None, **kwargs ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi b/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi index a60ceea84067..e0014fcb71d6 100644 --- a/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi +++ b/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp @@ -44,9 +43,5 @@ class IPInsightsPredictor(Predictor): class IPInsightsModel(Model): def __init__( - self, - model_data: str | PipelineVariable, - role: str | None = None, - sagemaker_session: Session | None = None, - **kwargs, + self, model_data: str | PipelineVariable, role: str | None = None, sagemaker_session: Session | None = None, **kwargs ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/kmeans.pyi b/stubs/sagemaker/sagemaker/amazon/kmeans.pyi index c2bd3bed0dac..754d45566d56 100644 --- a/stubs/sagemaker/sagemaker/amazon/kmeans.pyi +++ b/stubs/sagemaker/sagemaker/amazon/kmeans.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp @@ -46,9 +45,5 @@ class KMeansPredictor(Predictor): class KMeansModel(Model): def __init__( - self, - model_data: str | PipelineVariable, - role: str | None = None, - sagemaker_session: Session | None = None, - **kwargs, + self, model_data: str | PipelineVariable, role: str | None = None, sagemaker_session: Session | None = None, **kwargs ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/knn.pyi b/stubs/sagemaker/sagemaker/amazon/knn.pyi index 06636307df50..a783065ef124 100644 --- a/stubs/sagemaker/sagemaker/amazon/knn.pyi +++ b/stubs/sagemaker/sagemaker/amazon/knn.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp @@ -43,9 +42,5 @@ class KNNPredictor(Predictor): class KNNModel(Model): def __init__( - self, - model_data: str | PipelineVariable, - role: str | None = None, - sagemaker_session: Session | None = None, - **kwargs, + self, model_data: str | PipelineVariable, role: str | None = None, sagemaker_session: Session | None = None, **kwargs ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/lda.pyi b/stubs/sagemaker/sagemaker/amazon/lda.pyi index e0b1ec0c7247..864fa5c0232d 100644 --- a/stubs/sagemaker/sagemaker/amazon/lda.pyi +++ b/stubs/sagemaker/sagemaker/amazon/lda.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp @@ -36,9 +35,5 @@ class LDAPredictor(Predictor): class LDAModel(Model): def __init__( - self, - model_data: str | PipelineVariable, - role: str | None = None, - sagemaker_session: Session | None = None, - **kwargs, + self, model_data: str | PipelineVariable, role: str | None = None, sagemaker_session: Session | None = None, **kwargs ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi b/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi index 2011a49f3374..ef89c31f3828 100644 --- a/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi +++ b/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp @@ -112,9 +111,5 @@ class LinearLearnerPredictor(Predictor): class LinearLearnerModel(Model): def __init__( - self, - model_data: str | PipelineVariable, - role: str | None = None, - sagemaker_session: Session | None = None, - **kwargs, + self, model_data: str | PipelineVariable, role: str | None = None, sagemaker_session: Session | None = None, **kwargs ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/ntm.pyi b/stubs/sagemaker/sagemaker/amazon/ntm.pyi index 5d8b67dfa1d1..44c35d295e19 100644 --- a/stubs/sagemaker/sagemaker/amazon/ntm.pyi +++ b/stubs/sagemaker/sagemaker/amazon/ntm.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp @@ -49,9 +48,5 @@ class NTMPredictor(Predictor): class NTMModel(Model): def __init__( - self, - model_data: str | PipelineVariable, - role: str | None = None, - sagemaker_session: Session | None = None, - **kwargs, + self, model_data: str | PipelineVariable, role: str | None = None, sagemaker_session: Session | None = None, **kwargs ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/object2vec.pyi b/stubs/sagemaker/sagemaker/amazon/object2vec.pyi index cdad95098444..7b5cf4eb8e3e 100644 --- a/stubs/sagemaker/sagemaker/amazon/object2vec.pyi +++ b/stubs/sagemaker/sagemaker/amazon/object2vec.pyi @@ -1,5 +1,3 @@ -from typing import Optional - from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp from sagemaker.model import Model @@ -87,9 +85,5 @@ class Object2Vec(AmazonAlgorithmEstimatorBase): class Object2VecModel(Model): def __init__( - self, - model_data: str | PipelineVariable, - role: str | None = None, - sagemaker_session: Session | None = None, - **kwargs, + self, model_data: str | PipelineVariable, role: str | None = None, sagemaker_session: Session | None = None, **kwargs ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/pca.pyi b/stubs/sagemaker/sagemaker/amazon/pca.pyi index 958f083fe6e9..ddc84ebcaa09 100644 --- a/stubs/sagemaker/sagemaker/amazon/pca.pyi +++ b/stubs/sagemaker/sagemaker/amazon/pca.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp @@ -34,9 +33,5 @@ class PCAPredictor(Predictor): class PCAModel(Model): def __init__( - self, - model_data: str | PipelineVariable, - role: str | None = None, - sagemaker_session: Session | None = None, - **kwargs, + self, model_data: str | PipelineVariable, role: str | None = None, sagemaker_session: Session | None = None, **kwargs ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi b/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi index 6d79c22be3c5..356f6b9c1d98 100644 --- a/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi +++ b/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Optional from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase from sagemaker.amazon.hyperparameter import Hyperparameter as hp @@ -33,9 +32,5 @@ class RandomCutForestPredictor(Predictor): class RandomCutForestModel(Model): def __init__( - self, - model_data: str | PipelineVariable, - role: str | None = None, - sagemaker_session: Session | None = None, - **kwargs, + self, model_data: str | PipelineVariable, role: str | None = None, sagemaker_session: Session | None = None, **kwargs ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/automl/automl.pyi b/stubs/sagemaker/sagemaker/automl/automl.pyi index c222638df7d1..d1a2f9df8ecb 100644 --- a/stubs/sagemaker/sagemaker/automl/automl.pyi +++ b/stubs/sagemaker/sagemaker/automl/automl.pyi @@ -1,6 +1,5 @@ import abc from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker.job import _Job from sagemaker.session import Session diff --git a/stubs/sagemaker/sagemaker/base_predictor.pyi b/stubs/sagemaker/sagemaker/base_predictor.pyi index 1f87d79c2cb4..08a936d00526 100644 --- a/stubs/sagemaker/sagemaker/base_predictor.pyi +++ b/stubs/sagemaker/sagemaker/base_predictor.pyi @@ -1,6 +1,6 @@ import abc from _typeshed import Incomplete -from typing import Any, Tuple +from typing import Any from sagemaker.deserializers import StreamDeserializer as StreamDeserializer, StringDeserializer as StringDeserializer diff --git a/stubs/sagemaker/sagemaker/chainer/estimator.pyi b/stubs/sagemaker/sagemaker/chainer/estimator.pyi index 9fefc5424cc0..890c6e4eacfd 100644 --- a/stubs/sagemaker/sagemaker/chainer/estimator.pyi +++ b/stubs/sagemaker/sagemaker/chainer/estimator.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/chainer/model.pyi b/stubs/sagemaker/sagemaker/chainer/model.pyi index a2ea2c387fac..5fe2f8f63409 100644 --- a/stubs/sagemaker/sagemaker/chainer/model.pyi +++ b/stubs/sagemaker/sagemaker/chainer/model.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker import ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/clarify.pyi b/stubs/sagemaker/sagemaker/clarify.pyi index c027928f133a..f696db5b7176 100644 --- a/stubs/sagemaker/sagemaker/clarify.pyi +++ b/stubs/sagemaker/sagemaker/clarify.pyi @@ -2,7 +2,7 @@ import abc from _typeshed import Incomplete from abc import ABC, abstractmethod from enum import Enum -from typing import Any, Dict, List, Optional +from typing import Any from sagemaker.network import NetworkConfig from sagemaker.processing import Processor diff --git a/stubs/sagemaker/sagemaker/collection.pyi b/stubs/sagemaker/sagemaker/collection.pyi index b1704d849589..74e6d31012f3 100644 --- a/stubs/sagemaker/sagemaker/collection.pyi +++ b/stubs/sagemaker/sagemaker/collection.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List class Collection: sagemaker_session: Incomplete diff --git a/stubs/sagemaker/sagemaker/config/config.pyi b/stubs/sagemaker/sagemaker/config/config.pyi index 4d73876d6281..a4317ffe3669 100644 --- a/stubs/sagemaker/sagemaker/config/config.pyi +++ b/stubs/sagemaker/sagemaker/config/config.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List logger: Incomplete ENV_VARIABLE_ADMIN_CONFIG_OVERRIDE: str diff --git a/stubs/sagemaker/sagemaker/content_types.pyi b/stubs/sagemaker/sagemaker/content_types.pyi index 2852ed5e7419..acc1b8503a40 100644 --- a/stubs/sagemaker/sagemaker/content_types.pyi +++ b/stubs/sagemaker/sagemaker/content_types.pyi @@ -1,5 +1,3 @@ -from typing import List, Optional - def retrieve_options( region: str | None = None, model_id: str | None = None, diff --git a/stubs/sagemaker/sagemaker/debugger/debugger.pyi b/stubs/sagemaker/sagemaker/debugger/debugger.pyi index 15fe3c9785b0..0d9280346b8e 100644 --- a/stubs/sagemaker/sagemaker/debugger/debugger.pyi +++ b/stubs/sagemaker/sagemaker/debugger/debugger.pyi @@ -1,6 +1,5 @@ from _typeshed import Incomplete from abc import ABC -from typing import Dict, List, Optional from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi b/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi index 9c56853c61c5..8d28f4bada68 100644 --- a/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi +++ b/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.debugger.framework_profile import FrameworkProfile from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/deserializers.pyi b/stubs/sagemaker/sagemaker/deserializers.pyi index cf345a84a5c3..9d12487cb3c4 100644 --- a/stubs/sagemaker/sagemaker/deserializers.pyi +++ b/stubs/sagemaker/sagemaker/deserializers.pyi @@ -1,5 +1,3 @@ -from typing import List, Optional - from sagemaker.base_deserializers import ( BaseDeserializer, BytesDeserializer as BytesDeserializer, diff --git a/stubs/sagemaker/sagemaker/djl_inference/model.pyi b/stubs/sagemaker/sagemaker/djl_inference/model.pyi index a809237fa7fe..a1626711c2fe 100644 --- a/stubs/sagemaker/sagemaker/djl_inference/model.pyi +++ b/stubs/sagemaker/sagemaker/djl_inference/model.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from enum import Enum -from typing import Any, Dict, Optional +from typing import Any from sagemaker import Predictor from sagemaker.deserializers import BaseDeserializer diff --git a/stubs/sagemaker/sagemaker/drift_check_baselines.pyi b/stubs/sagemaker/sagemaker/drift_check_baselines.pyi index 784fd440a9c9..8746a676c54c 100644 --- a/stubs/sagemaker/sagemaker/drift_check_baselines.pyi +++ b/stubs/sagemaker/sagemaker/drift_check_baselines.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.model_metrics import FileSource, MetricsSource diff --git a/stubs/sagemaker/sagemaker/environment_variables.pyi b/stubs/sagemaker/sagemaker/environment_variables.pyi index 208e30e9e19e..f32ee3020a9c 100644 --- a/stubs/sagemaker/sagemaker/environment_variables.pyi +++ b/stubs/sagemaker/sagemaker/environment_variables.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, Optional logger: Incomplete diff --git a/stubs/sagemaker/sagemaker/estimator.pyi b/stubs/sagemaker/sagemaker/estimator.pyi index 60ba0099b77b..41dc88730b1a 100644 --- a/stubs/sagemaker/sagemaker/estimator.pyi +++ b/stubs/sagemaker/sagemaker/estimator.pyi @@ -1,7 +1,6 @@ import abc from _typeshed import Incomplete from abc import abstractmethod -from typing import Dict, List, Optional from sagemaker.debugger import ( DEBUGGER_FLAG as DEBUGGER_FLAG, diff --git a/stubs/sagemaker/sagemaker/experiments/_utils.pyi b/stubs/sagemaker/sagemaker/experiments/_utils.pyi index 138a27c84289..848ea430c7d3 100644 --- a/stubs/sagemaker/sagemaker/experiments/_utils.pyi +++ b/stubs/sagemaker/sagemaker/experiments/_utils.pyi @@ -1,5 +1,3 @@ -from typing import Optional - from sagemaker import Session from sagemaker.experiments._environment import _RunEnvironment diff --git a/stubs/sagemaker/sagemaker/experiments/run.pyi b/stubs/sagemaker/sagemaker/experiments/run.pyi index 166d2118c78d..cd968b7a5d67 100644 --- a/stubs/sagemaker/sagemaker/experiments/run.pyi +++ b/stubs/sagemaker/sagemaker/experiments/run.pyi @@ -1,7 +1,7 @@ import datetime from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Optional +from typing import Optional from numpy import array from sagemaker import Session diff --git a/stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi b/stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi index fb034b672b25..50e98f45d19a 100644 --- a/stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi +++ b/stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Optional class ClarifyTextConfig: language: Incomplete diff --git a/stubs/sagemaker/sagemaker/explainer/explainer_config.pyi b/stubs/sagemaker/sagemaker/explainer/explainer_config.pyi index fa0b5557d107..3fa08fb928e4 100644 --- a/stubs/sagemaker/sagemaker/explainer/explainer_config.pyi +++ b/stubs/sagemaker/sagemaker/explainer/explainer_config.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.explainer.clarify_explainer_config import ClarifyExplainerConfig diff --git a/stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi b/stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi index 386f2babf3b8..dd2ca1e51cc5 100644 --- a/stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi @@ -1,6 +1,5 @@ import datetime from enum import Enum -from typing import List, Tuple import pandas as pd from sagemaker.feature_store.feature_group import FeatureDefinition, FeatureGroup @@ -99,7 +98,9 @@ class DatasetBuilder: def with_number_of_recent_records_by_record_identifier(self, number_of_recent_records: int): ... def with_number_of_records_from_query_results(self, number_of_records: int): ... def as_of(self, timestamp: datetime.datetime): ... - def with_event_time_range(self, starting_timestamp: datetime.datetime | None = None, ending_timestamp: datetime.datetime | None = None): ... + def with_event_time_range( + self, starting_timestamp: datetime.datetime | None = None, ending_timestamp: datetime.datetime | None = None + ): ... def to_csv_file(self) -> tuple[str, str]: ... def to_dataframe(self) -> tuple[pd.DataFrame, str]: ... def __init__( diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi index 1e5c8ec2ce06..6b6394f84040 100644 --- a/stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi @@ -1,5 +1,5 @@ from enum import Enum -from typing import Any, Dict +from typing import Any from sagemaker.feature_store.inputs import Config diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi index 792918b8e623..6b004b3e84ba 100644 --- a/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi @@ -1,7 +1,7 @@ from _typeshed import Incomplete -from multiprocessing.pool import AsyncResult as AsyncResult -from typing import Any, Dict, List from collections.abc import Sequence +from multiprocessing.pool import AsyncResult as AsyncResult +from typing import Any from botocore.config import Config as Config from pandas import DataFrame as DataFrame diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi index 0b9acd9cd8d1..e516a205c81e 100644 --- a/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi @@ -1,6 +1,6 @@ import datetime -from typing import Any, Dict from collections.abc import Sequence +from typing import Any import pandas as pd from sagemaker import Session diff --git a/stubs/sagemaker/sagemaker/feature_store/inputs.pyi b/stubs/sagemaker/sagemaker/feature_store/inputs.pyi index 8e43957f0750..ca07c76c3135 100644 --- a/stubs/sagemaker/sagemaker/feature_store/inputs.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/inputs.pyi @@ -1,6 +1,6 @@ import abc from enum import Enum -from typing import Any, Dict, List +from typing import Any class Config(abc.ABC, metaclass=abc.ABCMeta): @abc.abstractmethod diff --git a/stubs/sagemaker/sagemaker/fw_utils.pyi b/stubs/sagemaker/sagemaker/fw_utils.pyi index f787c8a4aa5d..35cea594c4f2 100644 --- a/stubs/sagemaker/sagemaker/fw_utils.pyi +++ b/stubs/sagemaker/sagemaker/fw_utils.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, NamedTuple, Optional +from typing import NamedTuple from sagemaker.session_settings import SessionSettings from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/huggingface/estimator.pyi b/stubs/sagemaker/sagemaker/huggingface/estimator.pyi index d76c839250f7..81feb5d86458 100644 --- a/stubs/sagemaker/sagemaker/huggingface/estimator.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/estimator.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig diff --git a/stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi b/stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi index 31eb491b3c67..c3e2460137e4 100644 --- a/stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi @@ -1,5 +1,3 @@ -from typing import Optional - from sagemaker.session import Session def get_huggingface_llm_image_uri( diff --git a/stubs/sagemaker/sagemaker/huggingface/model.pyi b/stubs/sagemaker/sagemaker/huggingface/model.pyi index 6cdf18caf685..1ed4530b0033 100644 --- a/stubs/sagemaker/sagemaker/huggingface/model.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/model.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker import ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/huggingface/processing.pyi b/stubs/sagemaker/sagemaker/huggingface/processing.pyi index 162b9d200351..34bd8c3191dd 100644 --- a/stubs/sagemaker/sagemaker/huggingface/processing.pyi +++ b/stubs/sagemaker/sagemaker/huggingface/processing.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker.huggingface.estimator import HuggingFace from sagemaker.network import NetworkConfig diff --git a/stubs/sagemaker/sagemaker/hyperparameters.pyi b/stubs/sagemaker/sagemaker/hyperparameters.pyi index a808d5b6974d..12af8fd60218 100644 --- a/stubs/sagemaker/sagemaker/hyperparameters.pyi +++ b/stubs/sagemaker/sagemaker/hyperparameters.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, Optional from sagemaker.jumpstart.enums import HyperparameterValidationMode diff --git a/stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi b/stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi index f942f08a9749..0abfb3101b79 100644 --- a/stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi +++ b/stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker.parameter import CategoricalParameter diff --git a/stubs/sagemaker/sagemaker/inputs.pyi b/stubs/sagemaker/sagemaker/inputs.pyi index c61e7ae20dac..72bae95d28fc 100644 --- a/stubs/sagemaker/sagemaker/inputs.pyi +++ b/stubs/sagemaker/sagemaker/inputs.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import List, Optional +from typing import Optional from sagemaker.workflow.entities import PipelineVariable @@ -80,4 +80,6 @@ class BatchDataCaptureConfig: destination_s3_uri: Incomplete kms_key_id: Incomplete generate_inference_id: Incomplete - def __init__(self, destination_s3_uri: str, kms_key_id: str | None = None, generate_inference_id: bool | None = None) -> None: ... + def __init__( + self, destination_s3_uri: str, kms_key_id: str | None = None, generate_inference_id: bool | None = None + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/instance_types.pyi b/stubs/sagemaker/sagemaker/instance_types.pyi index 5d86e5786160..a6b9e5865814 100644 --- a/stubs/sagemaker/sagemaker/instance_types.pyi +++ b/stubs/sagemaker/sagemaker/instance_types.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Optional logger: Incomplete diff --git a/stubs/sagemaker/sagemaker/jumpstart/accessors.pyi b/stubs/sagemaker/sagemaker/jumpstart/accessors.pyi index 097c23013cb9..38535a13b75d 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/accessors.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/accessors.pyi @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Optional +from typing import Any from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartModelSpecs diff --git a/stubs/sagemaker/sagemaker/jumpstart/cache.pyi b/stubs/sagemaker/sagemaker/jumpstart/cache.pyi index ebbbdac1fb56..b8ae6b53fc10 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/cache.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/cache.pyi @@ -1,6 +1,5 @@ import datetime from _typeshed import Incomplete -from typing import List, Optional import botocore from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartModelSpecs diff --git a/stubs/sagemaker/sagemaker/jumpstart/constants.pyi b/stubs/sagemaker/sagemaker/jumpstart/constants.pyi index c095b070d3de..658de1fac609 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/constants.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/constants.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, Set, Type from sagemaker.base_deserializers import BaseDeserializer from sagemaker.base_serializers import BaseSerializer diff --git a/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi b/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi index 041a0ad70610..2fb009cf93f4 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker import session as session from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig diff --git a/stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi b/stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi index 60be283cc732..c9578364fe25 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Optional from sagemaker.jumpstart.constants import JumpStartScriptScope diff --git a/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi b/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi index 630e15be1a1e..5be3db1249ea 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig from sagemaker.base_deserializers import BaseDeserializer diff --git a/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi b/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi index b8e531fdb354..9dbb7cf24825 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig from sagemaker.base_deserializers import BaseDeserializer diff --git a/stubs/sagemaker/sagemaker/jumpstart/filters.pyi b/stubs/sagemaker/sagemaker/jumpstart/filters.pyi index e344a892d92d..3a63e4b32cfd 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/filters.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/filters.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from enum import Enum -from typing import Any, Dict, List +from typing import Any from sagemaker.jumpstart.types import JumpStartDataHolderType diff --git a/stubs/sagemaker/sagemaker/jumpstart/model.pyi b/stubs/sagemaker/sagemaker/jumpstart/model.pyi index 1e7ba70eb65f..6ebffc54c7cb 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/model.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/model.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig from sagemaker.base_deserializers import BaseDeserializer diff --git a/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi b/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi index 5290926abc8d..18680b92f0dd 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi @@ -1,5 +1,3 @@ -from typing import List, Tuple - from sagemaker.jumpstart.filters import Operator def extract_framework_task_model(model_id: str) -> tuple[str, str, str]: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/types.pyi b/stubs/sagemaker/sagemaker/jumpstart/types.pyi index 97f93a5fd1cc..e16da7647ce8 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/types.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/types.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from enum import Enum -from typing import Any, Dict, List, Optional, Set +from typing import Any class JumpStartDataHolderType: def __eq__(self, other: Any) -> bool: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/utils.pyi b/stubs/sagemaker/sagemaker/jumpstart/utils.pyi index d14dae7de992..6868e7585985 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/utils.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/utils.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Any, Dict, List, Optional +from typing import Any from sagemaker.jumpstart import enums from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartModelSpecs, JumpStartVersionedModelId @@ -16,9 +16,7 @@ def is_jumpstart_model_input(model_id: str | None, version: str | None) -> bool: def is_jumpstart_model_uri(uri: str | None) -> bool: ... def tag_key_in_array(tag_key: str, tag_array: list[dict[str, str]]) -> bool: ... def get_tag_value(tag_key: str, tag_array: list[dict[str, str]]) -> str: ... -def add_single_jumpstart_tag( - uri: str, tag_key: enums.JumpStartTag, curr_tags: list[dict[str, str]] | None -) -> list | None: ... +def add_single_jumpstart_tag(uri: str, tag_key: enums.JumpStartTag, curr_tags: list[dict[str, str]] | None) -> list | None: ... def get_jumpstart_base_name_if_jumpstart_model(*uris: str | None) -> str | None: ... def add_jumpstart_tags( tags: list[dict[str, str]] | None = None, @@ -46,8 +44,5 @@ def resolve_estimator_sagemaker_config_field( field_name: str, field_val: Any | None, sagemaker_session: Session, default_value: str | None = None ) -> Any: ... def is_valid_model_id( - model_id: str | None, - region: str | None = None, - model_version: str | None = None, - script: enums.JumpStartScriptScope = ..., + model_id: str | None, region: str | None = None, model_version: str | None = None, script: enums.JumpStartScriptScope = ... ) -> bool: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/validators.pyi b/stubs/sagemaker/sagemaker/jumpstart/validators.pyi index 36c366b6b0c3..1babe6e0594d 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/validators.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/validators.pyi @@ -1,4 +1,4 @@ -from typing import Any, Dict, Optional +from typing import Any from sagemaker.jumpstart.enums import HyperparameterValidationMode diff --git a/stubs/sagemaker/sagemaker/lineage/action.pyi b/stubs/sagemaker/sagemaker/lineage/action.pyi index 537b8254a365..fe92d8d959ba 100644 --- a/stubs/sagemaker/sagemaker/lineage/action.pyi +++ b/stubs/sagemaker/sagemaker/lineage/action.pyi @@ -1,7 +1,6 @@ from _typeshed import Incomplete -from datetime import datetime -from typing import Optional from collections.abc import Iterator +from datetime import datetime from sagemaker.apiutils import _base_types from sagemaker.lineage._api_types import ActionSource, ActionSummary diff --git a/stubs/sagemaker/sagemaker/lineage/artifact.pyi b/stubs/sagemaker/sagemaker/lineage/artifact.pyi index a1bb108e6b72..ca8e10f96ae3 100644 --- a/stubs/sagemaker/sagemaker/lineage/artifact.pyi +++ b/stubs/sagemaker/sagemaker/lineage/artifact.pyi @@ -1,7 +1,6 @@ from _typeshed import Incomplete -from datetime import datetime -from typing import Optional from collections.abc import Iterator +from datetime import datetime from sagemaker.apiutils import _base_types from sagemaker.lineage._api_types import ArtifactSource, ArtifactSummary diff --git a/stubs/sagemaker/sagemaker/lineage/association.pyi b/stubs/sagemaker/sagemaker/lineage/association.pyi index ef758e2fc502..6b0b8d51cf8a 100644 --- a/stubs/sagemaker/sagemaker/lineage/association.pyi +++ b/stubs/sagemaker/sagemaker/lineage/association.pyi @@ -1,7 +1,6 @@ from _typeshed import Incomplete -from datetime import datetime -from typing import Optional from collections.abc import Iterator +from datetime import datetime from sagemaker.apiutils import _base_types from sagemaker.lineage._api_types import AssociationSummary @@ -16,7 +15,11 @@ class Association(_base_types.Record): def set_tags(self, tags: Incomplete | None = None): ... @classmethod def create( - cls, source_arn: str, destination_arn: str, association_type: str | None = None, sagemaker_session: Incomplete | None = None + cls, + source_arn: str, + destination_arn: str, + association_type: str | None = None, + sagemaker_session: Incomplete | None = None, ) -> Association: ... @classmethod def list( diff --git a/stubs/sagemaker/sagemaker/lineage/context.pyi b/stubs/sagemaker/sagemaker/lineage/context.pyi index 180c3e015837..105b82a4801b 100644 --- a/stubs/sagemaker/sagemaker/lineage/context.pyi +++ b/stubs/sagemaker/sagemaker/lineage/context.pyi @@ -1,7 +1,6 @@ from _typeshed import Incomplete -from datetime import datetime -from typing import Optional from collections.abc import Iterator +from datetime import datetime from sagemaker.apiutils import _base_types from sagemaker.lineage import association diff --git a/stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi b/stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi index d78438f5beb6..3af3aaa43f15 100644 --- a/stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi +++ b/stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List from sagemaker.apiutils import _base_types from sagemaker.lineage.artifact import Artifact diff --git a/stubs/sagemaker/sagemaker/lineage/query.pyi b/stubs/sagemaker/sagemaker/lineage/query.pyi index 2e22c56382df..b3c5cda6ec0b 100644 --- a/stubs/sagemaker/sagemaker/lineage/query.pyi +++ b/stubs/sagemaker/sagemaker/lineage/query.pyi @@ -1,7 +1,7 @@ from _typeshed import Incomplete from datetime import datetime from enum import Enum -from typing import Any, Dict, List, Optional +from typing import Any class LineageEntityEnum(Enum): TRIAL: str @@ -57,7 +57,9 @@ class LineageQueryResult: edges: Incomplete vertices: Incomplete startarn: Incomplete - def __init__(self, edges: list[Edge] | None = None, vertices: list[Vertex] | None = None, startarn: list[str] | None = None) -> None: ... + def __init__( + self, edges: list[Edge] | None = None, vertices: list[Vertex] | None = None, startarn: list[str] | None = None + ) -> None: ... def visualize(self, path: str | None = "lineage_graph_pyvis.html", pyvis_options: dict[str, Any] | None = None): ... class LineageFilter: diff --git a/stubs/sagemaker/sagemaker/lineage/visualizer.pyi b/stubs/sagemaker/sagemaker/lineage/visualizer.pyi index f24bb67a39ec..0eb962e8b7ee 100644 --- a/stubs/sagemaker/sagemaker/lineage/visualizer.pyi +++ b/stubs/sagemaker/sagemaker/lineage/visualizer.pyi @@ -1,5 +1,3 @@ -from typing import Optional - from pandas import DataFrame as DataFrame class LineageTableVisualizer: diff --git a/stubs/sagemaker/sagemaker/local/pipeline.pyi b/stubs/sagemaker/sagemaker/local/pipeline.pyi index b1ee84ca1901..f590eb33bc6c 100644 --- a/stubs/sagemaker/sagemaker/local/pipeline.pyi +++ b/stubs/sagemaker/sagemaker/local/pipeline.pyi @@ -1,7 +1,6 @@ import abc from _typeshed import Incomplete from abc import ABC, abstractmethod -from typing import Dict from sagemaker.workflow.steps import Step diff --git a/stubs/sagemaker/sagemaker/metadata_properties.pyi b/stubs/sagemaker/sagemaker/metadata_properties.pyi index c63450183f4c..4821a30d5b94 100644 --- a/stubs/sagemaker/sagemaker/metadata_properties.pyi +++ b/stubs/sagemaker/sagemaker/metadata_properties.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/metric_definitions.pyi b/stubs/sagemaker/sagemaker/metric_definitions.pyi index a0d688803871..a9c77cf52f47 100644 --- a/stubs/sagemaker/sagemaker/metric_definitions.pyi +++ b/stubs/sagemaker/sagemaker/metric_definitions.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional logger: Incomplete diff --git a/stubs/sagemaker/sagemaker/model.pyi b/stubs/sagemaker/sagemaker/model.pyi index 125b98a886f4..6a547b135b88 100644 --- a/stubs/sagemaker/sagemaker/model.pyi +++ b/stubs/sagemaker/sagemaker/model.pyi @@ -1,6 +1,5 @@ import abc from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker.drift_check_baselines import DriftCheckBaselines from sagemaker.inference_recommender.inference_recommender_mixin import InferenceRecommenderMixin diff --git a/stubs/sagemaker/sagemaker/model_card/helpers.pyi b/stubs/sagemaker/sagemaker/model_card/helpers.pyi index 815343a726f4..564be87dab37 100644 --- a/stubs/sagemaker/sagemaker/model_card/helpers.pyi +++ b/stubs/sagemaker/sagemaker/model_card/helpers.pyi @@ -4,7 +4,7 @@ import json from _typeshed import Incomplete from abc import ABC, abstractmethod from enum import Enum -from typing import Any, Optional +from typing import Any from boto3.session import Session as Session diff --git a/stubs/sagemaker/sagemaker/model_card/model_card.pyi b/stubs/sagemaker/sagemaker/model_card/model_card.pyi index 0e23f9645505..4cdf80f012ec 100644 --- a/stubs/sagemaker/sagemaker/model_card/model_card.pyi +++ b/stubs/sagemaker/sagemaker/model_card/model_card.pyi @@ -1,6 +1,5 @@ from _typeshed import Incomplete from datetime import datetime -from typing import List, Optional from boto3.session import Session as boto3_Session from sagemaker.model_card.evaluation_metric_parsers import EvaluationMetricTypeEnum @@ -67,10 +66,7 @@ class BusinessDetails(_DefaultToRequestDict, _DefaultFromDict): business_stakeholders: Incomplete line_of_business: Incomplete def __init__( - self, - business_problem: str | None = None, - business_stakeholders: str | None = None, - line_of_business: str | None = None, + self, business_problem: str | None = None, business_stakeholders: str | None = None, line_of_business: str | None = None ) -> None: ... class Function(_DefaultToRequestDict, _DefaultFromDict): @@ -236,9 +232,7 @@ class ModelCard: def load(cls, name: str, version: int | None = None, sagemaker_session: Session | None = None): ... def update(self, **kwargs): ... def delete(self): ... - def export_pdf( - self, s3_output_path: str, export_job_name: str | None = None, model_card_version: int | None = None - ): ... + def export_pdf(self, s3_output_path: str, export_job_name: str | None = None, model_card_version: int | None = None): ... def list_export_jobs(self, **kwargs): ... def get_version_history(self, **kwargs): ... diff --git a/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi b/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi index 916a103623ad..4de57d55dd71 100644 --- a/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi +++ b/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.model_monitor.dataset_format import MonitoringDatasetFormat from sagemaker.processing import ProcessingJob diff --git a/stubs/sagemaker/sagemaker/model_uris.pyi b/stubs/sagemaker/sagemaker/model_uris.pyi index c938769f102d..968de9f1dd17 100644 --- a/stubs/sagemaker/sagemaker/model_uris.pyi +++ b/stubs/sagemaker/sagemaker/model_uris.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional logger: Incomplete diff --git a/stubs/sagemaker/sagemaker/multidatamodel.pyi b/stubs/sagemaker/sagemaker/multidatamodel.pyi index 7e55a0211b16..756fe1720ab6 100644 --- a/stubs/sagemaker/sagemaker/multidatamodel.pyi +++ b/stubs/sagemaker/sagemaker/multidatamodel.pyi @@ -1,6 +1,5 @@ from _typeshed import Incomplete from collections.abc import Generator -from typing import Optional from sagemaker.model import Model from sagemaker.session import Session diff --git a/stubs/sagemaker/sagemaker/mxnet/estimator.pyi b/stubs/sagemaker/sagemaker/mxnet/estimator.pyi index dfba054dd5da..e257089c0ec1 100644 --- a/stubs/sagemaker/sagemaker/mxnet/estimator.pyi +++ b/stubs/sagemaker/sagemaker/mxnet/estimator.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/mxnet/model.pyi b/stubs/sagemaker/sagemaker/mxnet/model.pyi index 121ade43a6ec..55e0d2bea7a8 100644 --- a/stubs/sagemaker/sagemaker/mxnet/model.pyi +++ b/stubs/sagemaker/sagemaker/mxnet/model.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker import ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/mxnet/processing.pyi b/stubs/sagemaker/sagemaker/mxnet/processing.pyi index 6112062df0cd..35e6d82fdca7 100644 --- a/stubs/sagemaker/sagemaker/mxnet/processing.pyi +++ b/stubs/sagemaker/sagemaker/mxnet/processing.pyi @@ -1,5 +1,3 @@ -from typing import Dict, List, Optional - from sagemaker.mxnet.estimator import MXNet from sagemaker.network import NetworkConfig from sagemaker.processing import FrameworkProcessor diff --git a/stubs/sagemaker/sagemaker/network.pyi b/stubs/sagemaker/sagemaker/network.pyi index f154e6287739..e676e919bede 100644 --- a/stubs/sagemaker/sagemaker/network.pyi +++ b/stubs/sagemaker/sagemaker/network.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Optional from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/pipeline.pyi b/stubs/sagemaker/sagemaker/pipeline.pyi index c2981edd4534..33f727c1e029 100644 --- a/stubs/sagemaker/sagemaker/pipeline.pyi +++ b/stubs/sagemaker/sagemaker/pipeline.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker import Model, ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/predictor.pyi b/stubs/sagemaker/sagemaker/predictor.pyi index b7403f555f28..9d7cce2d60ea 100644 --- a/stubs/sagemaker/sagemaker/predictor.pyi +++ b/stubs/sagemaker/sagemaker/predictor.pyi @@ -1,5 +1,3 @@ -from typing import Optional - from sagemaker.base_predictor import Predictor, PredictorBase as PredictorBase, RealTimePredictor as RealTimePredictor from sagemaker.session import Session diff --git a/stubs/sagemaker/sagemaker/processing.pyi b/stubs/sagemaker/sagemaker/processing.pyi index 1252c191b3fa..a90cef462c56 100644 --- a/stubs/sagemaker/sagemaker/processing.pyi +++ b/stubs/sagemaker/sagemaker/processing.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional +from typing import Optional from sagemaker.apiutils._base_types import ApiObject from sagemaker.dataset_definition.inputs import DatasetDefinition, S3Input diff --git a/stubs/sagemaker/sagemaker/pytorch/estimator.pyi b/stubs/sagemaker/sagemaker/pytorch/estimator.pyi index 2d8eaa131e93..c81e77625b06 100644 --- a/stubs/sagemaker/sagemaker/pytorch/estimator.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/estimator.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.pytorch.training_compiler.config import TrainingCompilerConfig diff --git a/stubs/sagemaker/sagemaker/pytorch/model.pyi b/stubs/sagemaker/sagemaker/pytorch/model.pyi index fa98cd8175ca..b3650cbc0a7b 100644 --- a/stubs/sagemaker/sagemaker/pytorch/model.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/model.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker import ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/pytorch/processing.pyi b/stubs/sagemaker/sagemaker/pytorch/processing.pyi index d60e077dfe62..4dbc8ba470bc 100644 --- a/stubs/sagemaker/sagemaker/pytorch/processing.pyi +++ b/stubs/sagemaker/sagemaker/pytorch/processing.pyi @@ -1,5 +1,3 @@ -from typing import Dict, List, Optional - from sagemaker.network import NetworkConfig from sagemaker.processing import FrameworkProcessor from sagemaker.pytorch.estimator import PyTorch diff --git a/stubs/sagemaker/sagemaker/remote_function/client.pyi b/stubs/sagemaker/sagemaker/remote_function/client.pyi index 978d5f220f36..5e5251d3305e 100644 --- a/stubs/sagemaker/sagemaker/remote_function/client.pyi +++ b/stubs/sagemaker/sagemaker/remote_function/client.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from collections.abc import Generator -from typing import Any, Dict, List, Tuple +from typing import Any from sagemaker.remote_function.job import _JobSettings from sagemaker.remote_function.spark_config import SparkConfig diff --git a/stubs/sagemaker/sagemaker/remote_function/job.pyi b/stubs/sagemaker/sagemaker/remote_function/job.pyi index 16c129b0f36b..c1de023558f8 100644 --- a/stubs/sagemaker/sagemaker/remote_function/job.pyi +++ b/stubs/sagemaker/sagemaker/remote_function/job.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Tuple from sagemaker.remote_function.spark_config import SparkConfig from sagemaker.session import Session diff --git a/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi b/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi index 78c6417c16ee..fdbe2f02db1a 100644 --- a/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi +++ b/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi @@ -1,5 +1,3 @@ -from typing import Dict, List, Optional - class SparkConfig: submit_jars: list[str] | None submit_py_files: list[str] | None diff --git a/stubs/sagemaker/sagemaker/rl/estimator.pyi b/stubs/sagemaker/sagemaker/rl/estimator.pyi index 420501cc025f..227b5fcf3d99 100644 --- a/stubs/sagemaker/sagemaker/rl/estimator.pyi +++ b/stubs/sagemaker/sagemaker/rl/estimator.pyi @@ -1,6 +1,5 @@ import enum from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker.estimator import Framework from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/s3_utils.pyi b/stubs/sagemaker/sagemaker/s3_utils.pyi index 3d58acf2e4e6..7f8bfab41678 100644 --- a/stubs/sagemaker/sagemaker/s3_utils.pyi +++ b/stubs/sagemaker/sagemaker/s3_utils.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional logger: Incomplete diff --git a/stubs/sagemaker/sagemaker/script_uris.pyi b/stubs/sagemaker/sagemaker/script_uris.pyi index c7708e0ee9f8..a10f72891c1d 100644 --- a/stubs/sagemaker/sagemaker/script_uris.pyi +++ b/stubs/sagemaker/sagemaker/script_uris.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional logger: Incomplete diff --git a/stubs/sagemaker/sagemaker/serializers.pyi b/stubs/sagemaker/sagemaker/serializers.pyi index de1ee2357557..cb33e7b60360 100644 --- a/stubs/sagemaker/sagemaker/serializers.pyi +++ b/stubs/sagemaker/sagemaker/serializers.pyi @@ -1,5 +1,3 @@ -from typing import List, Optional - from sagemaker.base_serializers import ( BaseSerializer, CSVSerializer as CSVSerializer, diff --git a/stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi b/stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi index 6ac6a63e7fe5..2c16bb841418 100644 --- a/stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi +++ b/stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional class ServerlessInferenceConfig: memory_size_in_mb: Incomplete diff --git a/stubs/sagemaker/sagemaker/session.pyi b/stubs/sagemaker/sagemaker/session.pyi index 56f8db5c3ffd..615cb191192f 100644 --- a/stubs/sagemaker/sagemaker/session.pyi +++ b/stubs/sagemaker/sagemaker/session.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete -from typing import Any, Dict, List, Optional from collections.abc import Sequence +from typing import Any from sagemaker.inputs import BatchDataCaptureConfig @@ -485,7 +485,13 @@ class Session: ) -> dict[str, Sequence[dict[str, str]]]: ... def batch_get_record(self, identifiers: Sequence[dict[str, Any]]) -> dict[str, Any]: ... def start_query_execution( - self, catalog: str, database: str, query_string: str, output_location: str, kms_key: str | None = None, workgroup: str | None = None + self, + catalog: str, + database: str, + query_string: str, + output_location: str, + kms_key: str | None = None, + workgroup: str | None = None, ) -> dict[str, str]: ... def get_query_execution(self, query_execution_id: str) -> dict[str, Any]: ... def wait_for_athena_query(self, query_execution_id: str, poll: int = 5): ... diff --git a/stubs/sagemaker/sagemaker/sklearn/estimator.pyi b/stubs/sagemaker/sagemaker/sklearn/estimator.pyi index 6e316e00f6c6..3a9f6264ff6c 100644 --- a/stubs/sagemaker/sagemaker/sklearn/estimator.pyi +++ b/stubs/sagemaker/sagemaker/sklearn/estimator.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/sklearn/model.pyi b/stubs/sagemaker/sagemaker/sklearn/model.pyi index 53eb5accd45b..cb00309a9bc9 100644 --- a/stubs/sagemaker/sagemaker/sklearn/model.pyi +++ b/stubs/sagemaker/sagemaker/sklearn/model.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker import ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/sklearn/processing.pyi b/stubs/sagemaker/sagemaker/sklearn/processing.pyi index 677b03f4adaf..e26176dcf0df 100644 --- a/stubs/sagemaker/sagemaker/sklearn/processing.pyi +++ b/stubs/sagemaker/sagemaker/sklearn/processing.pyi @@ -1,5 +1,3 @@ -from typing import Dict, List, Optional - from sagemaker import Session from sagemaker.network import NetworkConfig from sagemaker.processing import ScriptProcessor diff --git a/stubs/sagemaker/sagemaker/spark/processing.pyi b/stubs/sagemaker/sagemaker/spark/processing.pyi index 4054b6e3ec98..37e2e7bde131 100644 --- a/stubs/sagemaker/sagemaker/spark/processing.pyi +++ b/stubs/sagemaker/sagemaker/spark/processing.pyi @@ -1,6 +1,5 @@ from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Optional from sagemaker.network import NetworkConfig from sagemaker.processing import ProcessingInput, ProcessingOutput, ScriptProcessor diff --git a/stubs/sagemaker/sagemaker/sparkml/model.pyi b/stubs/sagemaker/sagemaker/sparkml/model.pyi index 85769ef8b73c..60904f7fc396 100644 --- a/stubs/sagemaker/sagemaker/sparkml/model.pyi +++ b/stubs/sagemaker/sagemaker/sparkml/model.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker import Model, Predictor, Session from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi b/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi index 54f5a1ee2030..4aa5bc44f6b5 100644 --- a/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi +++ b/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.tensorflow.training_compiler.config import TrainingCompilerConfig diff --git a/stubs/sagemaker/sagemaker/tensorflow/model.pyi b/stubs/sagemaker/sagemaker/tensorflow/model.pyi index 1862a861de82..e998f2bcae81 100644 --- a/stubs/sagemaker/sagemaker/tensorflow/model.pyi +++ b/stubs/sagemaker/sagemaker/tensorflow/model.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional import sagemaker from sagemaker import ModelMetrics diff --git a/stubs/sagemaker/sagemaker/tensorflow/processing.pyi b/stubs/sagemaker/sagemaker/tensorflow/processing.pyi index dfbd2cea8de6..d03dcd68a6b4 100644 --- a/stubs/sagemaker/sagemaker/tensorflow/processing.pyi +++ b/stubs/sagemaker/sagemaker/tensorflow/processing.pyi @@ -1,5 +1,3 @@ -from typing import Dict, List, Optional - from sagemaker.network import NetworkConfig from sagemaker.processing import FrameworkProcessor from sagemaker.session import Session diff --git a/stubs/sagemaker/sagemaker/transformer.pyi b/stubs/sagemaker/sagemaker/transformer.pyi index 169e9ae5eb35..36b17a3c1ed7 100644 --- a/stubs/sagemaker/sagemaker/transformer.pyi +++ b/stubs/sagemaker/sagemaker/transformer.pyi @@ -1,6 +1,5 @@ import abc from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker.inputs import BatchDataCaptureConfig from sagemaker.job import _Job diff --git a/stubs/sagemaker/sagemaker/tuner.pyi b/stubs/sagemaker/sagemaker/tuner.pyi index 6b08bb50448e..829ef1790907 100644 --- a/stubs/sagemaker/sagemaker/tuner.pyi +++ b/stubs/sagemaker/sagemaker/tuner.pyi @@ -1,7 +1,6 @@ import abc from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Optional, Set from sagemaker.amazon.amazon_estimator import FileSystemRecordSet, RecordSet from sagemaker.estimator import EstimatorBase diff --git a/stubs/sagemaker/sagemaker/utilities/cache.pyi b/stubs/sagemaker/sagemaker/utilities/cache.pyi index f3db0aaf58b9..28fff0cd59cb 100644 --- a/stubs/sagemaker/sagemaker/utilities/cache.pyi +++ b/stubs/sagemaker/sagemaker/utilities/cache.pyi @@ -1,7 +1,7 @@ import datetime from _typeshed import Incomplete -from typing import Optional, TypeVar from collections.abc import Callable +from typing import TypeVar KeyType = TypeVar("KeyType") ValType = TypeVar("ValType") diff --git a/stubs/sagemaker/sagemaker/utils.pyi b/stubs/sagemaker/sagemaker/utils.pyi index 24f3baa1fc63..7b8765fbc5cd 100644 --- a/stubs/sagemaker/sagemaker/utils.pyi +++ b/stubs/sagemaker/sagemaker/utils.pyi @@ -1,7 +1,7 @@ import abc from _typeshed import Incomplete from collections.abc import Generator -from typing import Any, List, Optional +from typing import Any ECR_URI_PATTERN: str MAX_BUCKET_PATHS_COUNT: int diff --git a/stubs/sagemaker/sagemaker/workflow/_utils.pyi b/stubs/sagemaker/sagemaker/workflow/_utils.pyi index e3bcc90526be..4daa82eb6f1c 100644 --- a/stubs/sagemaker/sagemaker/workflow/_utils.pyi +++ b/stubs/sagemaker/sagemaker/workflow/_utils.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Optional from sagemaker.estimator import EstimatorBase from sagemaker.workflow.entities import RequestType as RequestType diff --git a/stubs/sagemaker/sagemaker/workflow/automl_step.pyi b/stubs/sagemaker/sagemaker/workflow/automl_step.pyi index 354c3a16fe82..1f812bf36e3c 100644 --- a/stubs/sagemaker/sagemaker/workflow/automl_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/automl_step.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Optional from sagemaker.workflow.entities import RequestType as RequestType from sagemaker.workflow.pipeline_context import _JobStepArguments diff --git a/stubs/sagemaker/sagemaker/workflow/callback_step.pyi b/stubs/sagemaker/sagemaker/workflow/callback_step.pyi index 93c67d38cbf6..c2c948176b20 100644 --- a/stubs/sagemaker/sagemaker/workflow/callback_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/callback_step.pyi @@ -1,6 +1,5 @@ from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Optional from sagemaker.workflow.entities import DefaultEnumMeta, RequestType as RequestType from sagemaker.workflow.step_collections import StepCollection diff --git a/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi b/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi index fc3a8d89ee96..4b1a6556e8e5 100644 --- a/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi @@ -1,6 +1,5 @@ from _typeshed import Incomplete from abc import ABC -from typing import List, Optional from sagemaker.clarify import BiasConfig, DataConfig, ModelConfig, ModelPredictedLabelConfig, SHAPConfig from sagemaker.workflow.check_job_config import CheckJobConfig diff --git a/stubs/sagemaker/sagemaker/workflow/condition_step.pyi b/stubs/sagemaker/sagemaker/workflow/condition_step.pyi index b9a31b5916f8..dadc03f3ffaf 100644 --- a/stubs/sagemaker/sagemaker/workflow/condition_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/condition_step.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Optional from sagemaker.workflow.conditions import Condition from sagemaker.workflow.entities import RequestType as RequestType diff --git a/stubs/sagemaker/sagemaker/workflow/conditions.pyi b/stubs/sagemaker/sagemaker/workflow/conditions.pyi index af497089df04..7bc8c7557703 100644 --- a/stubs/sagemaker/sagemaker/workflow/conditions.pyi +++ b/stubs/sagemaker/sagemaker/workflow/conditions.pyi @@ -1,7 +1,6 @@ import abc from _typeshed import Incomplete from enum import Enum -from typing import Dict, List from sagemaker.workflow.entities import ( DefaultEnumMeta, diff --git a/stubs/sagemaker/sagemaker/workflow/emr_step.pyi b/stubs/sagemaker/sagemaker/workflow/emr_step.pyi index ade3ce4a1958..1d9c754ac637 100644 --- a/stubs/sagemaker/sagemaker/workflow/emr_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/emr_step.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Any, Dict, List, Optional +from typing import Any from sagemaker.workflow.entities import RequestType as RequestType from sagemaker.workflow.step_collections import StepCollection @@ -10,7 +10,9 @@ class EMRStepConfig: args: Incomplete main_class: Incomplete properties: Incomplete - def __init__(self, jar, args: list[str] | None = None, main_class: str | None = None, properties: list[dict] | None = None) -> None: ... + def __init__( + self, jar, args: list[str] | None = None, main_class: str | None = None, properties: list[dict] | None = None + ) -> None: ... def to_request(self) -> RequestType: ... INSTANCES: str diff --git a/stubs/sagemaker/sagemaker/workflow/entities.pyi b/stubs/sagemaker/sagemaker/workflow/entities.pyi index 7e92341711fc..5c7767cbf73b 100644 --- a/stubs/sagemaker/sagemaker/workflow/entities.pyi +++ b/stubs/sagemaker/sagemaker/workflow/entities.pyi @@ -1,7 +1,7 @@ import abc from _typeshed import Incomplete from enum import EnumMeta -from typing import Any, Dict, List +from typing import Any PrimitiveType = str | int | bool | float | None RequestType = dict[str | Any, list[dict[str, Any]]] diff --git a/stubs/sagemaker/sagemaker/workflow/fail_step.pyi b/stubs/sagemaker/sagemaker/workflow/fail_step.pyi index 8851ff031302..25d49d60a8ad 100644 --- a/stubs/sagemaker/sagemaker/workflow/fail_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/fail_step.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Optional from sagemaker.workflow.entities import PipelineVariable, RequestType as RequestType from sagemaker.workflow.step_collections import StepCollection diff --git a/stubs/sagemaker/sagemaker/workflow/functions.pyi b/stubs/sagemaker/sagemaker/workflow/functions.pyi index 51f6fb198127..db8bde5182f8 100644 --- a/stubs/sagemaker/sagemaker/workflow/functions.pyi +++ b/stubs/sagemaker/sagemaker/workflow/functions.pyi @@ -1,5 +1,3 @@ -from typing import List - from sagemaker.workflow.entities import PipelineVariable from sagemaker.workflow.properties import PropertyFile diff --git a/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi b/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi index 04ebc40ddf08..c9d64cecc3c8 100644 --- a/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi @@ -1,6 +1,5 @@ from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Optional from sagemaker.lambda_helper import Lambda from sagemaker.workflow.entities import DefaultEnumMeta, RequestType as RequestType diff --git a/stubs/sagemaker/sagemaker/workflow/model_step.pyi b/stubs/sagemaker/sagemaker/workflow/model_step.pyi index eacac28c06de..3338fc8745c2 100644 --- a/stubs/sagemaker/sagemaker/workflow/model_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/model_step.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker.workflow.pipeline_context import _ModelStepArguments from sagemaker.workflow.retry import RetryPolicy diff --git a/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi b/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi index 861f505aaa6d..2453f744c93d 100644 --- a/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.workflow.check_job_config import CheckJobConfig from sagemaker.workflow.clarify_check_step import ClarifyCheckConfig diff --git a/stubs/sagemaker/sagemaker/workflow/parameters.pyi b/stubs/sagemaker/sagemaker/workflow/parameters.pyi index d7cd12699f60..05144e5707c8 100644 --- a/stubs/sagemaker/sagemaker/workflow/parameters.pyi +++ b/stubs/sagemaker/sagemaker/workflow/parameters.pyi @@ -1,6 +1,5 @@ from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Type from sagemaker.workflow.entities import ( DefaultEnumMeta, diff --git a/stubs/sagemaker/sagemaker/workflow/pipeline.pyi b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi index fa9acf68b455..7f7e35d8405e 100644 --- a/stubs/sagemaker/sagemaker/workflow/pipeline.pyi +++ b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete -from typing import Any, Dict, List, Optional, Set from collections.abc import Sequence +from typing import Any from sagemaker.session import Session from sagemaker.workflow.entities import Entity, RequestType as RequestType @@ -37,7 +37,10 @@ class Pipeline(Entity): ) -> dict[str, Any]: ... def describe(self) -> dict[str, Any]: ... def update( - self, role_arn: str | None = None, description: str | None = None, parallelism_config: ParallelismConfiguration | None = None + self, + role_arn: str | None = None, + description: str | None = None, + parallelism_config: ParallelismConfiguration | None = None, ) -> dict[str, Any]: ... def upsert( self, @@ -57,7 +60,11 @@ class Pipeline(Entity): ): ... def definition(self) -> str: ... def list_executions( - self, sort_by: str | None = None, sort_order: str | None = None, max_results: int | None = None, next_token: str | None = None + self, + sort_by: str | None = None, + sort_order: str | None = None, + max_results: int | None = None, + next_token: str | None = None, ) -> dict[str, Any]: ... def format_start_parameters(parameters: dict[str, Any]) -> list[dict[str, Any]]: ... diff --git a/stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi b/stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi index 76b5342d8c78..fff6fb56f835 100644 --- a/stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi +++ b/stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from collections.abc import Callable from sagemaker.local import LocalSession diff --git a/stubs/sagemaker/sagemaker/workflow/properties.pyi b/stubs/sagemaker/sagemaker/workflow/properties.pyi index cd6e88e2a226..cfa979772a78 100644 --- a/stubs/sagemaker/sagemaker/workflow/properties.pyi +++ b/stubs/sagemaker/sagemaker/workflow/properties.pyi @@ -1,6 +1,5 @@ from _typeshed import Incomplete from abc import ABCMeta -from typing import Dict, List from sagemaker.workflow.entities import Expression, PipelineVariable diff --git a/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi b/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi index d3fbc44efabc..e1fa486a6fa1 100644 --- a/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi +++ b/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi @@ -1,6 +1,5 @@ from _typeshed import Incomplete from abc import ABC -from typing import List, Optional from sagemaker.workflow.check_job_config import CheckJobConfig from sagemaker.workflow.entities import PipelineVariable, RequestType as RequestType diff --git a/stubs/sagemaker/sagemaker/workflow/retry.pyi b/stubs/sagemaker/sagemaker/workflow/retry.pyi index 3d302161ce7f..4460ea5a8066 100644 --- a/stubs/sagemaker/sagemaker/workflow/retry.pyi +++ b/stubs/sagemaker/sagemaker/workflow/retry.pyi @@ -1,6 +1,5 @@ from _typeshed import Incomplete from enum import Enum -from typing import List from sagemaker.workflow.entities import DefaultEnumMeta, Entity, RequestType as RequestType diff --git a/stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi b/stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi index 7ce7d156a974..7108b4cc5dde 100644 --- a/stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi +++ b/stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List from sagemaker.workflow.entities import RequestType as RequestType diff --git a/stubs/sagemaker/sagemaker/workflow/step_collections.pyi b/stubs/sagemaker/sagemaker/workflow/step_collections.pyi index 51014766097e..c2351233fdb6 100644 --- a/stubs/sagemaker/sagemaker/workflow/step_collections.pyi +++ b/stubs/sagemaker/sagemaker/workflow/step_collections.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Optional from sagemaker import PipelineModel from sagemaker.estimator import EstimatorBase diff --git a/stubs/sagemaker/sagemaker/workflow/steps.pyi b/stubs/sagemaker/sagemaker/workflow/steps.pyi index 41aa5d7896f6..fb9cc3fff9ba 100644 --- a/stubs/sagemaker/sagemaker/workflow/steps.pyi +++ b/stubs/sagemaker/sagemaker/workflow/steps.pyi @@ -1,7 +1,6 @@ import abc from _typeshed import Incomplete from enum import Enum -from typing import Dict, List, Optional from sagemaker.estimator import EstimatorBase from sagemaker.inputs import CreateModelInput, FileSystemInput, TrainingInput, TransformInput diff --git a/stubs/sagemaker/sagemaker/workflow/utilities.pyi b/stubs/sagemaker/sagemaker/workflow/utilities.pyi index f0e6738af4ac..61d59fd5ce6b 100644 --- a/stubs/sagemaker/sagemaker/workflow/utilities.pyi +++ b/stubs/sagemaker/sagemaker/workflow/utilities.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import List, Set from collections.abc import Sequence from sagemaker.workflow.entities import Entity, RequestType as RequestType diff --git a/stubs/sagemaker/sagemaker/wrangler/ingestion.pyi b/stubs/sagemaker/sagemaker/wrangler/ingestion.pyi index 70e6186712cb..9d075c6f76c1 100644 --- a/stubs/sagemaker/sagemaker/wrangler/ingestion.pyi +++ b/stubs/sagemaker/sagemaker/wrangler/ingestion.pyi @@ -1,5 +1,3 @@ -from typing import Dict - from sagemaker.dataset_definition.inputs import AthenaDatasetDefinition, RedshiftDatasetDefinition def generate_data_ingestion_flow_from_s3_input( @@ -14,5 +12,8 @@ def generate_data_ingestion_flow_from_athena_dataset_definition( input_name: str, athena_dataset_definition: AthenaDatasetDefinition, operator_version: str = "0.1", schema: dict | None = None ): ... def generate_data_ingestion_flow_from_redshift_dataset_definition( - input_name: str, redshift_dataset_definition: RedshiftDatasetDefinition, operator_version: str = "0.1", schema: dict | None = None + input_name: str, + redshift_dataset_definition: RedshiftDatasetDefinition, + operator_version: str = "0.1", + schema: dict | None = None, ): ... diff --git a/stubs/sagemaker/sagemaker/wrangler/processing.pyi b/stubs/sagemaker/sagemaker/wrangler/processing.pyi index 87fcd0698946..14de5a1f2614 100644 --- a/stubs/sagemaker/sagemaker/wrangler/processing.pyi +++ b/stubs/sagemaker/sagemaker/wrangler/processing.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List from sagemaker.network import NetworkConfig from sagemaker.processing import Processor diff --git a/stubs/sagemaker/sagemaker/xgboost/estimator.pyi b/stubs/sagemaker/sagemaker/xgboost/estimator.pyi index cc2c005e35f0..04c78c0e6015 100644 --- a/stubs/sagemaker/sagemaker/xgboost/estimator.pyi +++ b/stubs/sagemaker/sagemaker/xgboost/estimator.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, Optional from sagemaker.estimator import Framework from sagemaker.workflow.entities import PipelineVariable diff --git a/stubs/sagemaker/sagemaker/xgboost/model.pyi b/stubs/sagemaker/sagemaker/xgboost/model.pyi index fe811b899867..77900846aaa2 100644 --- a/stubs/sagemaker/sagemaker/xgboost/model.pyi +++ b/stubs/sagemaker/sagemaker/xgboost/model.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Dict, List, Optional from sagemaker import ModelMetrics from sagemaker.drift_check_baselines import DriftCheckBaselines diff --git a/stubs/sagemaker/sagemaker/xgboost/processing.pyi b/stubs/sagemaker/sagemaker/xgboost/processing.pyi index 28e6e9e19031..2b2b247edd49 100644 --- a/stubs/sagemaker/sagemaker/xgboost/processing.pyi +++ b/stubs/sagemaker/sagemaker/xgboost/processing.pyi @@ -1,5 +1,3 @@ -from typing import Dict, List, Optional - from sagemaker.network import NetworkConfig from sagemaker.processing import FrameworkProcessor from sagemaker.session import Session From 63866d11f534ae067e0a67fd78d407ea28e3d06c Mon Sep 17 00:00:00 2001 From: DemyCode Date: Sun, 25 Jun 2023 18:09:57 +0200 Subject: [PATCH 07/10] fix: fixing flake8 issues --- stubs/sagemaker/sagemaker/clarify.pyi | 15 ++++---- .../sagemaker/sagemaker/debugger/debugger.pyi | 4 +-- .../sagemaker/djl_inference/model.pyi | 7 ++-- stubs/sagemaker/sagemaker/experiments/run.pyi | 36 ++++++++++++------- .../sagemaker/feature_store/feature_group.pyi | 2 +- stubs/sagemaker/sagemaker/inputs.pyi | 8 ++--- .../sagemaker/sagemaker/jumpstart/filters.pyi | 2 +- stubs/sagemaker/sagemaker/jumpstart/types.pyi | 10 +++--- .../sagemaker/model_card/helpers.pyi | 22 ++++++------ .../sagemaker/model_card/model_card.pyi | 21 +++++------ stubs/sagemaker/sagemaker/model_metrics.pyi | 17 +++++---- stubs/sagemaker/sagemaker/parameter.pyi | 4 +-- stubs/sagemaker/sagemaker/processing.pyi | 11 +++--- stubs/sagemaker/sagemaker/workflow/_utils.pyi | 16 +++++---- .../sagemaker/workflow/conditions.pyi | 5 +-- .../sagemaker/sagemaker/workflow/entities.pyi | 6 ++-- .../sagemaker/sagemaker/workflow/pipeline.pyi | 2 +- stubs/sagemaker/sagemaker/workflow/steps.pyi | 22 ++++++------ .../sagemaker/workflow/utilities.pyi | 4 ++- 19 files changed, 117 insertions(+), 97 deletions(-) diff --git a/stubs/sagemaker/sagemaker/clarify.pyi b/stubs/sagemaker/sagemaker/clarify.pyi index f696db5b7176..c641fbb74a85 100644 --- a/stubs/sagemaker/sagemaker/clarify.pyi +++ b/stubs/sagemaker/sagemaker/clarify.pyi @@ -51,7 +51,7 @@ class DataConfig: predicted_label_dataset_uri: str | None = None, predicted_label_headers: list[str] | None = None, predicted_label: str | int | None = None, - excluded_columns: list[int, list[str]] | None = None, + excluded_columns: list[int] | list[str] | None = None, ) -> None: ... def get_config(self): ... @@ -59,9 +59,10 @@ class BiasConfig: analysis_config: Incomplete def __init__( self, - label_values_or_threshold: int | float | str, - facet_name: str | int | list[str, list[int]], - facet_values_or_threshold: int | float | str | None = None, + label_values_or_threshold: float | str, + # Union[str, int, List[str], List[int]], + facet_name: str | int | list[str] | list[int], + facet_values_or_threshold: float | str | None = None, group_name: str | None = None, ) -> None: ... def get_config(self): ... @@ -106,7 +107,7 @@ class ExplainabilityConfig(ABC, metaclass=abc.ABCMeta): class PDPConfig(ExplainabilityConfig): pdp_config: Incomplete - def __init__(self, features: list | None = None, grid_resolution: int = 15, top_k_features: int = 10) -> None: ... + def __init__(self, features: list[Any] | None = None, grid_resolution: int = 15, top_k_features: int = 10) -> None: ... def get_explainability_config(self): ... class TextConfig: @@ -132,7 +133,7 @@ class SHAPConfig(ExplainabilityConfig): shap_config: Incomplete def __init__( self, - baseline: str | list | dict | None = None, + baseline: str | list[Any] | dict[Any, Any] | None = None, num_samples: int | None = None, agg_method: str | None = None, use_logit: bool = False, @@ -207,7 +208,7 @@ class SageMakerClarifyProcessor(Processor): self, data_config: DataConfig, model_config: ModelConfig, - explainability_config: ExplainabilityConfig | list, + explainability_config: ExplainabilityConfig | list[ExplainabilityConfig], model_scores: int | str | ModelPredictedLabelConfig | None = None, wait: bool = True, logs: bool = True, diff --git a/stubs/sagemaker/sagemaker/debugger/debugger.pyi b/stubs/sagemaker/sagemaker/debugger/debugger.pyi index 0d9280346b8e..29a952153239 100644 --- a/stubs/sagemaker/sagemaker/debugger/debugger.pyi +++ b/stubs/sagemaker/sagemaker/debugger/debugger.pyi @@ -65,7 +65,7 @@ class Rule(RuleBase): s3_output_path: str | PipelineVariable | None = None, other_trials_s3_input_paths: list[str | PipelineVariable] | None = None, rule_parameters: dict[str, str | PipelineVariable] | None = None, - collections_to_save: list["CollectionConfig"] | None = None, + collections_to_save: list[CollectionConfig] | None = None, actions: Incomplete | None = None, ): ... def prepare_actions(self, training_job_name) -> None: ... @@ -105,7 +105,7 @@ class DebuggerHookConfig: s3_output_path: str | PipelineVariable | None = None, container_local_output_path: str | PipelineVariable | None = None, hook_parameters: dict[str, str | PipelineVariable] | None = None, - collection_configs: list["CollectionConfig"] | None = None, + collection_configs: list[CollectionConfig] | None = None, ) -> None: ... class TensorBoardOutputConfig: diff --git a/stubs/sagemaker/sagemaker/djl_inference/model.pyi b/stubs/sagemaker/sagemaker/djl_inference/model.pyi index a1626711c2fe..efd6207b1b8f 100644 --- a/stubs/sagemaker/sagemaker/djl_inference/model.pyi +++ b/stubs/sagemaker/sagemaker/djl_inference/model.pyi @@ -1,8 +1,9 @@ from _typeshed import Incomplete from enum import Enum from typing import Any +from collections.abc import Callable +from sagemaker.base_predictor import Predictor -from sagemaker import Predictor from sagemaker.deserializers import BaseDeserializer from sagemaker.model import FrameworkModel from sagemaker.serializers import BaseSerializer @@ -57,7 +58,7 @@ class DJLModel(FrameworkModel): prediction_timeout: int | None = None, entry_point: str | None = None, image_uri: str | PipelineVariable | None = None, - predictor_cls: callable = ..., + predictor_cls: Callable[[str, Session], Predictor] | Predictor = ..., **kwargs, ) -> None: ... def package_for_edge(self, **_) -> None: ... @@ -145,7 +146,7 @@ class HuggingFaceAccelerateModel(DJLModel): role: str, number_of_partitions: int | None = None, device_id: int | None = None, - device_map: str | dict[str | str] | None = None, + device_map: str | dict[str, str] | None = None, load_in_8bit: bool = False, low_cpu_mem_usage: bool = False, **kwargs, diff --git a/stubs/sagemaker/sagemaker/experiments/run.pyi b/stubs/sagemaker/sagemaker/experiments/run.pyi index cd968b7a5d67..29f4a6b38f81 100644 --- a/stubs/sagemaker/sagemaker/experiments/run.pyi +++ b/stubs/sagemaker/sagemaker/experiments/run.pyi @@ -1,9 +1,9 @@ import datetime from _typeshed import Incomplete from enum import Enum -from typing import Optional +from typing import Any -from numpy import array +from numpy import ndarray from sagemaker import Session logger: Incomplete @@ -38,25 +38,35 @@ class Run: experiment_display_name: str | None = None, run_display_name: str | None = None, tags: list[dict[str, str]] | None = None, - sagemaker_session: Optional["Session"] = None, + sagemaker_session: Session | None = None, ) -> None: ... @property - def experiment_config(self) -> dict: ... - def log_parameter(self, name: str, value: str | int | float): ... - def log_parameters(self, parameters: dict[str, str | int | float]): ... + def experiment_config(self) -> dict[Any, Any]: ... + def log_parameter(self, name: str, value: str | float): ... + def log_parameters(self, parameters: dict[str, str | float]): ... def log_metric(self, name: str, value: float, timestamp: datetime.datetime | None = None, step: int | None = None): ... def log_precision_recall( self, - y_true: list | array, - predicted_probabilities: list | array, + y_true: list[Any] | ndarray[Any, Any], + predicted_probabilities: list[Any] | ndarray[Any, Any], positive_label: str | int | None = None, title: str | None = None, is_output: bool = True, no_skill: int | None = None, ): ... - def log_roc_curve(self, y_true: list | array, y_score: list | array, title: str | None = None, is_output: bool = True): ... + def log_roc_curve( + self, + y_true: list[Any] | ndarray[Any, Any], + y_score: list[Any] | ndarray[Any, Any], + title: str | None = None, + is_output: bool = True, + ): ... def log_confusion_matrix( - self, y_true: list | array, y_pred: list | array, title: str | None = None, is_output: bool = True + self, + y_true: list[Any] | ndarray[Any, Any], + y_pred: list[Any] | ndarray[Any, Any], + title: str | None = None, + is_output: bool = True, ): ... def log_artifact(self, name: str, value: str, media_type: str | None = None, is_output: bool = True): ... def log_file(self, file_path: str, name: str | None = None, media_type: str | None = None, is_output: bool = True): ... @@ -65,15 +75,15 @@ class Run: def __exit__(self, exc_type, exc_value, exc_traceback) -> None: ... def load_run( - run_name: str | None = None, experiment_name: str | None = None, sagemaker_session: Optional["Session"] = None + run_name: str | None = None, experiment_name: str | None = None, sagemaker_session: Session | None = None ) -> Run: ... def list_runs( experiment_name: str, created_before: datetime.datetime | None = None, created_after: datetime.datetime | None = None, - sagemaker_session: Optional["Session"] = None, + sagemaker_session: Session | None = None, max_results: int | None = None, next_token: str | None = None, sort_by: SortByType = ..., sort_order: SortOrderType = ..., -) -> list: ... +) -> list[Run]: ... diff --git a/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi index 6b004b3e84ba..8d949c3a4806 100644 --- a/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi +++ b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi @@ -106,7 +106,7 @@ class FeatureGroup: max_workers: int = 1, max_processes: int = 1, wait: bool = True, - timeout: int | float | None = None, + timeout: float | None = None, profile_name: str | None = None, ) -> IngestionManagerPandas: ... def athena_query(self) -> AthenaQuery: ... diff --git a/stubs/sagemaker/sagemaker/inputs.pyi b/stubs/sagemaker/sagemaker/inputs.pyi index 72bae95d28fc..aba4ef0e1f2b 100644 --- a/stubs/sagemaker/sagemaker/inputs.pyi +++ b/stubs/sagemaker/sagemaker/inputs.pyi @@ -1,5 +1,5 @@ from _typeshed import Incomplete -from typing import Optional +from typing import Any, Literal from sagemaker.workflow.entities import PipelineVariable @@ -20,7 +20,7 @@ class TrainingInput: input_mode: str | PipelineVariable | None = None, attribute_names: list[str | PipelineVariable] | None = None, target_attribute_name: str | PipelineVariable | None = None, - shuffle_config: Optional["ShuffleConfig"] = None, + shuffle_config: ShuffleConfig | None = None, ) -> None: ... class ShuffleConfig: @@ -45,8 +45,8 @@ class TransformInput: input_filter: str output_filter: str join_source: str - model_client_config: dict - batch_data_capture_config: dict + model_client_config: dict[Literal['InvocationsTimeoutInSeconds', 'InvocationsMaxRetries'], int] + batch_data_capture_config: dict[Any, Any] def __init__( self, data, diff --git a/stubs/sagemaker/sagemaker/jumpstart/filters.pyi b/stubs/sagemaker/sagemaker/jumpstart/filters.pyi index 3a63e4b32cfd..8c9abccb4fde 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/filters.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/filters.pyi @@ -83,5 +83,5 @@ class ModelFilter(JumpStartDataHolderType): def parse_filter_string(filter_string: str) -> ModelFilter: ... def evaluate_filter_expression( - model_filter: ModelFilter, cached_model_value: str | bool | int | float | dict[str | Any, list[Any]] + model_filter: ModelFilter, cached_model_value: str | bool | float | dict[str, Any] | list[Any] ) -> BooleanValues: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/types.pyi b/stubs/sagemaker/sagemaker/jumpstart/types.pyi index e16da7647ce8..f1415b226abe 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/types.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/types.pyi @@ -1,9 +1,11 @@ from _typeshed import Incomplete from enum import Enum from typing import Any +from collections.abc import Callable +from sagemaker import Predictor class JumpStartDataHolderType: - def __eq__(self, other: Any) -> bool: ... + def __eq__(self, other: object) -> bool: ... def __hash__(self) -> int: ... class JumpStartS3FileType(str, Enum): @@ -168,7 +170,7 @@ class JumpStartModelInitKwargs(JumpStartKwargs): image_uri: str | Any | None = None, model_data: str | Any | None = None, role: str | None = None, - predictor_cls: callable | None = None, + predictor_cls: Callable[..., Predictor] | Predictor | None = None, env: dict[str, str | Any] | None = None, name: str | None = None, vpc_config: dict[str, list[str | Any]] | None = None, @@ -355,7 +357,7 @@ class JumpStartEstimatorFitKwargs(JumpStartKwargs): model_id: str, model_version: str | None = None, region: str | None = None, - inputs: str | dict | Any | Any | None = None, + inputs: str | dict[Any, Any] | None = None, wait: bool | None = None, logs: str | None = None, job_name: str | None = None, @@ -429,7 +431,7 @@ class JumpStartEstimatorDeployKwargs(JumpStartKwargs): explainer_config: Any | None = None, image_uri: str | Any | None = None, role: str | None = None, - predictor_cls: callable | None = None, + predictor_cls: Callable[..., Predictor] | Predictor | None = None, env: dict[str, str | Any] | None = None, model_name: str | None = None, vpc_config: dict[str, list[str | Any]] | None = None, diff --git a/stubs/sagemaker/sagemaker/model_card/helpers.pyi b/stubs/sagemaker/sagemaker/model_card/helpers.pyi index 564be87dab37..35cd0b665aad 100644 --- a/stubs/sagemaker/sagemaker/model_card/helpers.pyi +++ b/stubs/sagemaker/sagemaker/model_card/helpers.pyi @@ -4,7 +4,7 @@ import json from _typeshed import Incomplete from abc import ABC, abstractmethod from enum import Enum -from typing import Any +from typing import Any, Literal from boto3.session import Session as Session @@ -24,9 +24,9 @@ class _DescriptorBase(ABC, metaclass=abc.ABCMeta): @abstractmethod def validate(self, value): ... @abstractmethod - def require_decode(self, value: dict): ... + def require_decode(self, value: dict[Any, Any]): ... @abstractmethod - def decode(self, value: dict): ... + def decode(self, value: dict[Any, Any]): ... class _OneOf(_DescriptorBase): options: Incomplete @@ -41,25 +41,25 @@ class _IsList(_DescriptorBase): item_type: Incomplete max_size: Incomplete def __init__(self, item_type: object, max_size: int | None = None) -> None: ... - def validate(self, value: list): ... - def require_decode(self, value: list): ... - def decode(self, value: list): ... + def validate(self, value: list[Any]): ... + def require_decode(self, value: list[Any]) -> Literal[True]: ... + def decode(self, value: list[Any]) -> _MaxSizeArray: ... class _IsModelCardObject(_DescriptorBase): custom_class: Incomplete def __init__(self, custom_class: object) -> None: ... - def validate(self, value: dict | object): ... - def require_decode(self, value: dict | object): ... - def decode(self, value: dict): ... + def validate(self, value: dict[Any, Any] | object): ... + def require_decode(self, value: dict[Any, Any] | object) -> bool: ... + def decode(self, value: dict[Any, Any]): ... class _MaxSizeArray(collections.abc.MutableSequence): list: Incomplete - def __init__(self, max_size: int, item_type: Any, array: list | None = None) -> None: ... + def __init__(self, max_size: int, item_type: Any, array: list[Any] | None = None) -> None: ... def __len__(self) -> int: ... def __getitem__(self, index): ... def __delitem__(self, index) -> None: ... def __setitem__(self, index: int, value: Any): ... - def __eq__(self, other: Any): ... + def __eq__(self, other: _MaxSizeArray): ... def check(self, value: Any): ... def insert(self, index, value) -> None: ... def to_map(self, key_attribute: str): ... diff --git a/stubs/sagemaker/sagemaker/model_card/model_card.pyi b/stubs/sagemaker/sagemaker/model_card/model_card.pyi index 4cdf80f012ec..15c02813d828 100644 --- a/stubs/sagemaker/sagemaker/model_card/model_card.pyi +++ b/stubs/sagemaker/sagemaker/model_card/model_card.pyi @@ -1,5 +1,6 @@ from _typeshed import Incomplete from datetime import datetime +from typing import Any from boto3.session import Session as boto3_Session from sagemaker.model_card.evaluation_metric_parsers import EvaluationMetricTypeEnum @@ -35,7 +36,7 @@ class ModelOverview(_DefaultToRequestDict, _DefaultFromDict): model_id: str | None = None, model_name: str | None = None, model_description: str | None = None, - model_version: int | float | None = None, + model_version: float | None = None, problem_type: str | None = None, algorithm_type: str | None = None, model_creator: str | None = None, @@ -95,21 +96,21 @@ class Metric(_DefaultToRequestDict, _DefaultFromDict): self, name: str, type: MetricTypeEnum | str, - value: int | float | str | bool | list, + value: float | str | bool | list[Any], notes: str | None = None, - x_axis_name: str | list | None = None, - y_axis_name: str | list | None = None, + x_axis_name: str | list[str] | None = None, + y_axis_name: str | list[str] | None = None, ) -> None: ... @property def value(self): ... @value.setter - def value(self, val: int | float | str | bool | list): ... + def value(self, val: float | str | bool | list[Any]): ... class TrainingMetric(_DefaultToRequestDict, _DefaultFromDict): name: Incomplete value: Incomplete notes: Incomplete - def __init__(self, name: str, value: int | float, notes: str | None = None) -> None: ... + def __init__(self, name: str, value: float, notes: str | None = None) -> None: ... class HyperParameter(_DefaultToRequestDict, _DefaultFromDict): name: Incomplete @@ -171,7 +172,7 @@ class EvaluationJob(_DefaultToRequestDict, _DefaultFromDict): evaluation_observation: str | None = None, evaluation_job_arn: str | None = None, datasets: list[str] | None = None, - metadata: dict | None = None, + metadata: dict[Any, Any] | None = None, metric_groups: list[MetricGroup] | None = None, ) -> None: ... def get_metric_group(self, group_name): ... @@ -187,7 +188,7 @@ class AdditionalInformation(_DefaultToRequestDict, _DefaultFromDict): self, ethical_considerations: str | None = None, caveats_and_recommendations: str | None = None, - custom_details: dict | None = None, + custom_details: dict[Any, Any] | None = None, ) -> None: ... class ModelCard: @@ -216,9 +217,9 @@ class ModelCard: arn: str | None = None, version: int | None = None, created_time: datetime | None = None, - created_by: dict | None = None, + created_by: dict[Any, Any] | None = None, last_modified_time: datetime | None = None, - last_modified_by: dict | None = None, + last_modified_by: dict[Any, Any] | None = None, model_overview: ModelOverview | None = None, intended_uses: IntendedUses | None = None, business_details: BusinessDetails | None = None, diff --git a/stubs/sagemaker/sagemaker/model_metrics.pyi b/stubs/sagemaker/sagemaker/model_metrics.pyi index 0b0f54e13e67..970ab668b03a 100644 --- a/stubs/sagemaker/sagemaker/model_metrics.pyi +++ b/stubs/sagemaker/sagemaker/model_metrics.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.workflow.entities import PipelineVariable @@ -14,14 +13,14 @@ class ModelMetrics: explainability: Incomplete def __init__( self, - model_statistics: Optional["MetricsSource"] = None, - model_constraints: Optional["MetricsSource"] = None, - model_data_statistics: Optional["MetricsSource"] = None, - model_data_constraints: Optional["MetricsSource"] = None, - bias: Optional["MetricsSource"] = None, - explainability: Optional["MetricsSource"] = None, - bias_pre_training: Optional["MetricsSource"] = None, - bias_post_training: Optional["MetricsSource"] = None, + model_statistics: MetricsSource | None = None, + model_constraints: MetricsSource | None = None, + model_data_statistics: MetricsSource | None = None, + model_data_constraints: MetricsSource | None = None, + bias: MetricsSource | None = None, + explainability: MetricsSource | None = None, + bias_pre_training: MetricsSource | None = None, + bias_post_training: MetricsSource | None = None, ) -> None: ... class MetricsSource: diff --git a/stubs/sagemaker/sagemaker/parameter.pyi b/stubs/sagemaker/sagemaker/parameter.pyi index 6b5fe6129022..d57ba98476cc 100644 --- a/stubs/sagemaker/sagemaker/parameter.pyi +++ b/stubs/sagemaker/sagemaker/parameter.pyi @@ -9,8 +9,8 @@ class ParameterRange: scaling_type: Incomplete def __init__( self, - min_value: int | float | PipelineVariable, - max_value: int | float | PipelineVariable, + min_value: float | PipelineVariable, + max_value: float | PipelineVariable, scaling_type: str | PipelineVariable = "Auto", ) -> None: ... def is_valid(self, value): ... diff --git a/stubs/sagemaker/sagemaker/processing.pyi b/stubs/sagemaker/sagemaker/processing.pyi index a90cef462c56..1bc931ad35ec 100644 --- a/stubs/sagemaker/sagemaker/processing.pyi +++ b/stubs/sagemaker/sagemaker/processing.pyi @@ -1,5 +1,4 @@ from _typeshed import Incomplete -from typing import Optional from sagemaker.apiutils._base_types import ApiObject from sagemaker.dataset_definition.inputs import DatasetDefinition, S3Input @@ -48,8 +47,8 @@ class Processor: ) -> None: ... def run( self, - inputs: list["ProcessingInput"] | None = None, - outputs: list["ProcessingOutput"] | None = None, + inputs: ProcessingInput | None = None, + outputs: ProcessingOutput | None = None, arguments: list[str | PipelineVariable] | None = None, wait: bool = True, logs: bool = True, @@ -84,8 +83,8 @@ class ScriptProcessor(Processor): def run( self, code: str, - inputs: list["ProcessingInput"] | None = None, - outputs: list["ProcessingOutput"] | None = None, + inputs: ProcessingInput | None = None, + outputs: ProcessingOutput | None = None, arguments: list[str | PipelineVariable] | None = None, wait: bool = True, logs: bool = True, @@ -156,7 +155,7 @@ class ProcessingOutput: output_name: str | PipelineVariable | None = None, s3_upload_mode: str | PipelineVariable = "EndOfJob", app_managed: bool | PipelineVariable = False, - feature_store_output: Optional["FeatureStoreOutput"] = None, + feature_store_output: FeatureStoreOutput | None = None, ) -> None: ... class RunArgs: diff --git a/stubs/sagemaker/sagemaker/workflow/_utils.pyi b/stubs/sagemaker/sagemaker/workflow/_utils.pyi index 4daa82eb6f1c..dc4e18754d13 100644 --- a/stubs/sagemaker/sagemaker/workflow/_utils.pyi +++ b/stubs/sagemaker/sagemaker/workflow/_utils.pyi @@ -1,8 +1,10 @@ from _typeshed import Incomplete +from typing import Any from sagemaker.estimator import EstimatorBase from sagemaker.workflow.entities import RequestType as RequestType from sagemaker.workflow.retry import RetryPolicy +from sagemaker.workflow.step_collections import StepCollection from sagemaker.workflow.steps import ConfigurableRetryStep, Step, TrainingStep logger: Incomplete @@ -25,8 +27,8 @@ class _RepackModelStep(TrainingStep): display_name: str | None = None, description: str | None = None, source_dir: str | None = None, - dependencies: list | None = None, - depends_on: list[str | Step | "StepCollection"] | None = None, + dependencies: list[str] | None = None, + depends_on: list[str | Step | StepCollection] | None = None, retry_policies: list[RetryPolicy] | None = None, subnets: Incomplete | None = None, security_group_ids: Incomplete | None = None, @@ -63,11 +65,11 @@ class _RegisterModelStep(ConfigurableRetryStep): def __init__( self, name: str, - step_args: dict | None = None, - content_types: list | None = None, - response_types: list | None = None, - inference_instances: list | None = None, - transform_instances: list | None = None, + step_args: dict[Any, Any] | None = None, + content_types: list[Any] | None = None, + response_types: list[Any] | None = None, + inference_instances: list[Any] | None = None, + transform_instances: list[Any] | None = None, estimator: EstimatorBase | None = None, model_data: Incomplete | None = None, model_package_group_name: Incomplete | None = None, diff --git a/stubs/sagemaker/sagemaker/workflow/conditions.pyi b/stubs/sagemaker/sagemaker/workflow/conditions.pyi index 7bc8c7557703..5713f9f23e3f 100644 --- a/stubs/sagemaker/sagemaker/workflow/conditions.pyi +++ b/stubs/sagemaker/sagemaker/workflow/conditions.pyi @@ -1,6 +1,7 @@ import abc from _typeshed import Incomplete from enum import Enum +from typing import TypeAlias from sagemaker.workflow.entities import ( DefaultEnumMeta, @@ -13,7 +14,7 @@ from sagemaker.workflow.execution_variables import ExecutionVariable from sagemaker.workflow.parameters import Parameter from sagemaker.workflow.properties import Properties -ConditionValueType = ExecutionVariable | Parameter | Properties +ConditionValueType: TypeAlias = ExecutionVariable | Parameter | Properties class ConditionTypeEnum(Enum, metaclass=DefaultEnumMeta): EQ: str @@ -78,4 +79,4 @@ class ConditionOr(Condition): def primitive_or_expr( value: ExecutionVariable | Expression | PrimitiveType | Parameter | Properties, -) -> dict[str | str, PrimitiveType]: ... +) -> dict[str, str] | PrimitiveType: ... diff --git a/stubs/sagemaker/sagemaker/workflow/entities.pyi b/stubs/sagemaker/sagemaker/workflow/entities.pyi index 5c7767cbf73b..66ebfce4ea2d 100644 --- a/stubs/sagemaker/sagemaker/workflow/entities.pyi +++ b/stubs/sagemaker/sagemaker/workflow/entities.pyi @@ -1,10 +1,10 @@ import abc from _typeshed import Incomplete from enum import EnumMeta -from typing import Any +from typing import Any, TypeAlias -PrimitiveType = str | int | bool | float | None -RequestType = dict[str | Any, list[dict[str, Any]]] +PrimitiveType: TypeAlias = str | int | bool | float | None +RequestType: TypeAlias = dict[str | Any, list[dict[str, Any]]] class Entity(abc.ABC, metaclass=abc.ABCMeta): @abc.abstractmethod diff --git a/stubs/sagemaker/sagemaker/workflow/pipeline.pyi b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi index 7f7e35d8405e..a5808ff2b4e6 100644 --- a/stubs/sagemaker/sagemaker/workflow/pipeline.pyi +++ b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi @@ -52,7 +52,7 @@ class Pipeline(Entity): def delete(self) -> dict[str, Any]: ... def start( self, - parameters: dict[str, str | bool | int | float] | None = None, + parameters: dict[str, str | bool | float] | None = None, execution_display_name: str | None = None, execution_description: str | None = None, parallelism_config: ParallelismConfiguration | None = None, diff --git a/stubs/sagemaker/sagemaker/workflow/steps.pyi b/stubs/sagemaker/sagemaker/workflow/steps.pyi index fb9cc3fff9ba..692ae4e8f06c 100644 --- a/stubs/sagemaker/sagemaker/workflow/steps.pyi +++ b/stubs/sagemaker/sagemaker/workflow/steps.pyi @@ -1,6 +1,7 @@ import abc from _typeshed import Incomplete from enum import Enum +from typing import Any from sagemaker.estimator import EstimatorBase from sagemaker.inputs import CreateModelInput, FileSystemInput, TrainingInput, TransformInput @@ -14,6 +15,7 @@ from sagemaker.workflow.functions import Join from sagemaker.workflow.pipeline_context import _JobStepArguments from sagemaker.workflow.properties import PropertyFile from sagemaker.workflow.retry import RetryPolicy +from sagemaker.workflow.step_collections import StepCollection class StepTypeEnum(Enum, metaclass=DefaultEnumMeta): CONDITION: str @@ -36,7 +38,7 @@ class Step(Entity, metaclass=abc.ABCMeta): display_name: str | None description: str | None step_type: StepTypeEnum - depends_on: list[str | "Step" | "StepCollection"] | None + depends_on: list[str | Step | StepCollection] | None @property @abc.abstractmethod def arguments(self) -> RequestType: ... @@ -46,7 +48,7 @@ class Step(Entity, metaclass=abc.ABCMeta): @abc.abstractmethod def properties(self): ... def to_request(self) -> RequestType: ... - def add_depends_on(self, step_names: list[str | "Step" | "StepCollection"]): ... + def add_depends_on(self, step_names: list[str | Step | StepCollection]): ... @property def ref(self) -> dict[str, str]: ... def __init__(self, name, display_name, description, step_type, depends_on) -> None: ... @@ -74,7 +76,7 @@ class ConfigurableRetryStep(Step, metaclass=abc.ABCMeta): step_type: StepTypeEnum, display_name: str | None = None, description: str | None = None, - depends_on: list[str | Step | "StepCollection"] | None = None, + depends_on: list[str | Step | StepCollection] | None = None, retry_policies: list[RetryPolicy] | None = None, ) -> None: ... def add_retry_policy(self, retry_policy: RetryPolicy): ... @@ -93,9 +95,9 @@ class TrainingStep(ConfigurableRetryStep): estimator: EstimatorBase | None = None, display_name: str | None = None, description: str | None = None, - inputs: TrainingInput | dict | str | FileSystemInput | None = None, + inputs: TrainingInput | dict[str, str] | dict[str, TrainingInput] | str | FileSystemInput | None = None, cache_config: CacheConfig | None = None, - depends_on: list[str | Step | "StepCollection"] | None = None, + depends_on: list[str | Step | StepCollection] | None = None, retry_policies: list[RetryPolicy] | None = None, ) -> None: ... @property @@ -111,10 +113,10 @@ class CreateModelStep(ConfigurableRetryStep): def __init__( self, name: str, - step_args: dict | None = None, + step_args: dict[Any, Any] | None = None, model: Model | PipelineModel | None = None, inputs: CreateModelInput | None = None, - depends_on: list[str | Step | "StepCollection"] | None = None, + depends_on: list[str | Step | StepCollection] | None = None, retry_policies: list[RetryPolicy] | None = None, display_name: str | None = None, description: str | None = None, @@ -138,7 +140,7 @@ class TransformStep(ConfigurableRetryStep): display_name: str | None = None, description: str | None = None, cache_config: CacheConfig | None = None, - depends_on: list[str | Step | "StepCollection"] | None = None, + depends_on: list[str | Step | StepCollection] | None = None, retry_policies: list[RetryPolicy] | None = None, ) -> None: ... @property @@ -171,7 +173,7 @@ class ProcessingStep(ConfigurableRetryStep): code: str | None = None, property_files: list[PropertyFile] | None = None, cache_config: CacheConfig | None = None, - depends_on: list[str | Step | "StepCollection"] | None = None, + depends_on: list[str | Step | StepCollection] | None = None, retry_policies: list[RetryPolicy] | None = None, kms_key: Incomplete | None = None, ) -> None: ... @@ -197,7 +199,7 @@ class TuningStep(ConfigurableRetryStep): inputs: Incomplete | None = None, job_arguments: list[str] | None = None, cache_config: CacheConfig | None = None, - depends_on: list[str | Step | "StepCollection"] | None = None, + depends_on: list[str | Step | StepCollection] | None = None, retry_policies: list[RetryPolicy] | None = None, ) -> None: ... @property diff --git a/stubs/sagemaker/sagemaker/workflow/utilities.pyi b/stubs/sagemaker/sagemaker/workflow/utilities.pyi index 61d59fd5ce6b..71249491c86f 100644 --- a/stubs/sagemaker/sagemaker/workflow/utilities.pyi +++ b/stubs/sagemaker/sagemaker/workflow/utilities.pyi @@ -3,11 +3,13 @@ from collections.abc import Sequence from sagemaker.workflow.entities import Entity, RequestType as RequestType from sagemaker.workflow.pipeline_context import _StepArguments +from sagemaker.workflow.step_collections import StepCollection + logger: Incomplete BUF_SIZE: int -def list_to_request(entities: Sequence[Entity | "StepCollection"]) -> list[RequestType]: ... +def list_to_request(entities: Sequence[Entity | StepCollection]) -> list[RequestType]: ... def build_steps(steps: Sequence[Entity], pipeline_name: str): ... def get_code_hash(step: Entity) -> str: ... def get_processing_dependencies(dependency_args: list[list[str]]) -> list[str]: ... From ee45b61fb0f3c69ebfc8b5b62219382ba9f25f49 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 25 Jun 2023 16:12:01 +0000 Subject: [PATCH 08/10] [pre-commit.ci] auto fixes from pre-commit.com hooks --- stubs/sagemaker/sagemaker/djl_inference/model.pyi | 4 ++-- stubs/sagemaker/sagemaker/inputs.pyi | 2 +- stubs/sagemaker/sagemaker/jumpstart/types.pyi | 3 ++- stubs/sagemaker/sagemaker/workflow/utilities.pyi | 1 - 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/stubs/sagemaker/sagemaker/djl_inference/model.pyi b/stubs/sagemaker/sagemaker/djl_inference/model.pyi index efd6207b1b8f..173a88f2b67a 100644 --- a/stubs/sagemaker/sagemaker/djl_inference/model.pyi +++ b/stubs/sagemaker/sagemaker/djl_inference/model.pyi @@ -1,9 +1,9 @@ from _typeshed import Incomplete +from collections.abc import Callable from enum import Enum from typing import Any -from collections.abc import Callable -from sagemaker.base_predictor import Predictor +from sagemaker.base_predictor import Predictor from sagemaker.deserializers import BaseDeserializer from sagemaker.model import FrameworkModel from sagemaker.serializers import BaseSerializer diff --git a/stubs/sagemaker/sagemaker/inputs.pyi b/stubs/sagemaker/sagemaker/inputs.pyi index aba4ef0e1f2b..c82718fcd154 100644 --- a/stubs/sagemaker/sagemaker/inputs.pyi +++ b/stubs/sagemaker/sagemaker/inputs.pyi @@ -45,7 +45,7 @@ class TransformInput: input_filter: str output_filter: str join_source: str - model_client_config: dict[Literal['InvocationsTimeoutInSeconds', 'InvocationsMaxRetries'], int] + model_client_config: dict[Literal["InvocationsTimeoutInSeconds", "InvocationsMaxRetries"], int] batch_data_capture_config: dict[Any, Any] def __init__( self, diff --git a/stubs/sagemaker/sagemaker/jumpstart/types.pyi b/stubs/sagemaker/sagemaker/jumpstart/types.pyi index f1415b226abe..3227d171d8d5 100644 --- a/stubs/sagemaker/sagemaker/jumpstart/types.pyi +++ b/stubs/sagemaker/sagemaker/jumpstart/types.pyi @@ -1,7 +1,8 @@ from _typeshed import Incomplete +from collections.abc import Callable from enum import Enum from typing import Any -from collections.abc import Callable + from sagemaker import Predictor class JumpStartDataHolderType: diff --git a/stubs/sagemaker/sagemaker/workflow/utilities.pyi b/stubs/sagemaker/sagemaker/workflow/utilities.pyi index 71249491c86f..863b46da20a9 100644 --- a/stubs/sagemaker/sagemaker/workflow/utilities.pyi +++ b/stubs/sagemaker/sagemaker/workflow/utilities.pyi @@ -5,7 +5,6 @@ from sagemaker.workflow.entities import Entity, RequestType as RequestType from sagemaker.workflow.pipeline_context import _StepArguments from sagemaker.workflow.step_collections import StepCollection - logger: Incomplete BUF_SIZE: int From c29657f76285b1f3a6edb550a972ade59bd70695 Mon Sep 17 00:00:00 2001 From: DemyCode Date: Sun, 25 Jun 2023 18:29:17 +0200 Subject: [PATCH 09/10] fix: trying renaming typevar with underscore --- stubs/sagemaker/sagemaker/utilities/cache.pyi | 14 +++++++------- stubs/sagemaker/sagemaker/workflow/_utils.pyi | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/stubs/sagemaker/sagemaker/utilities/cache.pyi b/stubs/sagemaker/sagemaker/utilities/cache.pyi index 28fff0cd59cb..b0e1973cfe67 100644 --- a/stubs/sagemaker/sagemaker/utilities/cache.pyi +++ b/stubs/sagemaker/sagemaker/utilities/cache.pyi @@ -3,23 +3,23 @@ from _typeshed import Incomplete from collections.abc import Callable from typing import TypeVar -KeyType = TypeVar("KeyType") -ValType = TypeVar("ValType") +_KeyType = TypeVar("_KeyType") +_ValType = TypeVar("_ValType") class LRUCache: class Element: value: Incomplete creation_time: Incomplete - def __init__(self, value: ValType, creation_time: datetime.datetime) -> None: ... + def __init__(self, value: _ValType, creation_time: datetime.datetime) -> None: ... def __init__( self, max_cache_items: int, expiration_horizon: datetime.timedelta, - retrieval_function: Callable[[KeyType, ValType], ValType], + retrieval_function: Callable[[_KeyType, _ValType], _ValType], ) -> None: ... def __len__(self) -> int: ... - def __contains__(self, key: KeyType) -> bool: ... + def __contains__(self, key: _KeyType) -> bool: ... def clear(self) -> None: ... - def get(self, key: KeyType, data_source_fallback: bool | None = True) -> ValType: ... - def put(self, key: KeyType, value: ValType | None = None) -> None: ... + def get(self, key: _KeyType, data_source_fallback: bool | None = True) -> _ValType: ... + def put(self, key: _KeyType, value: _ValType | None = None) -> None: ... diff --git a/stubs/sagemaker/sagemaker/workflow/_utils.pyi b/stubs/sagemaker/sagemaker/workflow/_utils.pyi index dc4e18754d13..fafb6cab4687 100644 --- a/stubs/sagemaker/sagemaker/workflow/_utils.pyi +++ b/stubs/sagemaker/sagemaker/workflow/_utils.pyi @@ -80,7 +80,7 @@ class _RegisterModelStep(ConfigurableRetryStep): compile_model_family: Incomplete | None = None, display_name: str | None = None, description: Incomplete | None = None, - depends_on: list[str | Step | "StepCollection"] | None = None, + depends_on: list[str | Step | StepCollection] | None = None, retry_policies: list[RetryPolicy] | None = None, tags: Incomplete | None = None, container_def_list: Incomplete | None = None, From a06680d5a32f7086792c571d4612bb46e2bad13c Mon Sep 17 00:00:00 2001 From: Alex Waygood Date: Sat, 22 Jul 2023 19:18:29 +0100 Subject: [PATCH 10/10] add upstream repo --- stubs/sagemaker/METADATA.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/stubs/sagemaker/METADATA.toml b/stubs/sagemaker/METADATA.toml index 4398389b5581..0b3e0b7b46f3 100644 --- a/stubs/sagemaker/METADATA.toml +++ b/stubs/sagemaker/METADATA.toml @@ -1 +1,2 @@ version = "2.168.*" +upstream_repository = "https://github.com/aws/sagemaker-python-sdk/"