diff --git a/pyrightconfig.stricter.json b/pyrightconfig.stricter.json index 341b43b90883..912434a701d5 100644 --- a/pyrightconfig.stricter.json +++ b/pyrightconfig.stricter.json @@ -68,6 +68,7 @@ "stubs/qrcode", "stubs/redis", "stubs/requests", + "stubs/sagemaker", "stubs/setuptools", "stubs/stripe", "stubs/tqdm", diff --git a/stubs/sagemaker/METADATA.toml b/stubs/sagemaker/METADATA.toml new file mode 100644 index 000000000000..0b3e0b7b46f3 --- /dev/null +++ b/stubs/sagemaker/METADATA.toml @@ -0,0 +1,2 @@ +version = "2.168.*" +upstream_repository = "https://github.com/aws/sagemaker-python-sdk/" 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..1ffd31dba8cd --- /dev/null +++ b/stubs/sagemaker/sagemaker/accept_types.pyi @@ -0,0 +1,14 @@ +def retrieve_options( + 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]: ... +def retrieve_default( + 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 new file mode 100644 index 000000000000..a88be9ce9d38 --- /dev/null +++ b/stubs/sagemaker/sagemaker/algorithm.pyi @@ -0,0 +1,74 @@ +from _typeshed import Incomplete + +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 = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + volume_size: int | PipelineVariable = 30, + volume_kms_key: str | PipelineVariable | None = None, + max_run: int | PipelineVariable = 86400, + input_mode: str | PipelineVariable = "File", + 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: list[dict[str, str | PipelineVariable]] | None = None, + encrypt_inter_container_traffic: bool | PipelineVariable = False, + use_spot_instances: bool | PipelineVariable = False, + max_wait: int | PipelineVariable | None = 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: str | dict | TrainingInput | FileSystemInput | None = None, + wait: bool = True, + logs: bool = True, + job_name: str | None = 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..1a6f4c6addf5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/amazon_estimator.pyi @@ -0,0 +1,83 @@ +import abc +from _typeshed import Incomplete + +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: str | None + repo_version: str | None + DEFAULT_MINI_BATCH_SIZE: int | None + def __init__( + self, + 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: ... + 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: int | None = None, + wait: bool = True, + logs: bool = True, + 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): ... + +class RecordSet: + s3_data: Incomplete + feature_dim: Incomplete + num_records: Incomplete + s3_data_type: Incomplete + channel: Incomplete + def __init__( + self, + s3_data: str | PipelineVariable, + num_records: int, + feature_dim: int, + s3_data_type: str | PipelineVariable = "ManifestFile", + channel: 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..880922adef98 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/factorization_machines.pyi @@ -0,0 +1,76 @@ +from _typeshed import Incomplete + +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: 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): ... + +class FactorizationMachinesPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class FactorizationMachinesModel(Model): + def __init__( + self, model_data: str | PipelineVariable, role: str | None = None, sagemaker_session: Session | None = 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..e0014fcb71d6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/ipinsights.pyi @@ -0,0 +1,47 @@ +from _typeshed import Incomplete + +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: 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): ... + +class IPInsightsPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class IPInsightsModel(Model): + def __init__( + 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 new file mode 100644 index 000000000000..754d45566d56 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/kmeans.pyi @@ -0,0 +1,49 @@ +from _typeshed import Incomplete + +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: 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): ... + 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: 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 new file mode 100644 index 000000000000..a783065ef124 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/knn.pyi @@ -0,0 +1,46 @@ +from _typeshed import Incomplete + +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: 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): ... + +class KNNPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class KNNModel(Model): + def __init__( + 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 new file mode 100644 index 000000000000..864fa5c0232d --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/lda.pyi @@ -0,0 +1,39 @@ +from _typeshed import Incomplete + +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: 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): ... + +class LDAPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class LDAModel(Model): + def __init__( + 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 new file mode 100644 index 000000000000..ef89c31f3828 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/linear_learner.pyi @@ -0,0 +1,115 @@ +from _typeshed import Incomplete + +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: 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): ... + +class LinearLearnerPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class LinearLearnerModel(Model): + def __init__( + 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 new file mode 100644 index 000000000000..44c35d295e19 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/ntm.pyi @@ -0,0 +1,52 @@ +from _typeshed import Incomplete + +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: 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): ... + +class NTMPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class NTMModel(Model): + def __init__( + 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 new file mode 100644 index 000000000000..7b5cf4eb8e3e --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/object2vec.pyi @@ -0,0 +1,89 @@ +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: 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): ... + +class Object2VecModel(Model): + def __init__( + 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 new file mode 100644 index 000000000000..ddc84ebcaa09 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/pca.pyi @@ -0,0 +1,37 @@ +from _typeshed import Incomplete + +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: 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): ... + +class PCAPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class PCAModel(Model): + def __init__( + 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 new file mode 100644 index 000000000000..356f6b9c1d98 --- /dev/null +++ b/stubs/sagemaker/sagemaker/amazon/randomcutforest.pyi @@ -0,0 +1,36 @@ +from _typeshed import Incomplete + +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: 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): ... + +class RandomCutForestPredictor(Predictor): + def __init__(self, endpoint_name, sagemaker_session: Incomplete | None = None, serializer=..., deserializer=...) -> None: ... + +class RandomCutForestModel(Model): + def __init__( + self, model_data: str | PipelineVariable, role: str | None = None, sagemaker_session: Session | None = 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..d1a2f9df8ecb --- /dev/null +++ b/stubs/sagemaker/sagemaker/automl/automl.pyi @@ -0,0 +1,144 @@ +import abc +from _typeshed import Incomplete + +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: 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( + 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..08a936d00526 --- /dev/null +++ b/stubs/sagemaker/sagemaker/base_predictor.pyi @@ -0,0 +1,71 @@ +import abc +from _typeshed import Incomplete +from typing import Any + +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..890c6e4eacfd --- /dev/null +++ b/stubs/sagemaker/sagemaker/chainer/estimator.pyi @@ -0,0 +1,39 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + 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): ... + 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..5fe2f8f63409 --- /dev/null +++ b/stubs/sagemaker/sagemaker/chainer/model.pyi @@ -0,0 +1,68 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + 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: int | PipelineVariable | None = None, + **kwargs, + ) -> None: ... + image_uri: Incomplete + def register( + self, + 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: 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, + 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..c641fbb74a85 --- /dev/null +++ b/stubs/sagemaker/sagemaker/clarify.pyi @@ -0,0 +1,282 @@ +import abc +from _typeshed import Incomplete +from abc import ABC, abstractmethod +from enum import Enum +from typing import Any + +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: 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: 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): ... + +class BiasConfig: + analysis_config: Incomplete + def __init__( + self, + 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): ... + +class ModelConfig: + predictor_config: Incomplete + def __init__( + self, + 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): ... + +class ModelPredictedLabelConfig: + label: Incomplete + probability: Incomplete + probability_threshold: Incomplete + label_headers: Incomplete + predictor_config: Incomplete + def __init__( + self, + 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): ... + +class ExplainabilityConfig(ABC, metaclass=abc.ABCMeta): + @abstractmethod + def get_explainability_config(self): ... + +class PDPConfig(ExplainabilityConfig): + pdp_config: Incomplete + 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: + 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: 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): ... + +class SHAPConfig(ExplainabilityConfig): + shap_config: Incomplete + def __init__( + self, + baseline: str | list[Any] | dict[Any, Any] | None = None, + num_samples: int | None = None, + agg_method: str | None = None, + use_logit: bool = False, + save_local_shap_values: bool = True, + 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): ... + +class SageMakerClarifyProcessor(Processor): + job_name_prefix: Incomplete + skip_early_validation: Incomplete + def __init__( + self, + role: str | None = None, + instance_count: int | None = None, + instance_type: str | None = None, + volume_size_in_gb: int = 30, + 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: ... + def run_pre_training_bias( + self, + data_config: DataConfig, + data_bias_config: BiasConfig, + methods: str | list[str] = "all", + wait: bool = True, + logs: bool = True, + 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: ModelConfig | None = None, + model_predicted_label_config: ModelPredictedLabelConfig | None = None, + methods: str | list[str] = "all", + wait: bool = True, + logs: bool = True, + 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: 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: 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[ExplainabilityConfig], + model_scores: int | str | ModelPredictedLabelConfig | None = None, + wait: bool = True, + logs: bool = True, + 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], + bias_config: BiasConfig, + 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, + 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: ExplainabilityConfig | list[ExplainabilityConfig], + bias_config: BiasConfig, + pre_training_methods: str | list[str] = "all", + post_training_methods: str | list[str] = "all", + ): ... + @classmethod + def explainability( + cls, + data_config: DataConfig, + model_config: ModelConfig, + model_predicted_label_config: ModelPredictedLabelConfig, + explainability_config: ExplainabilityConfig | list[ExplainabilityConfig], + ): ... + @classmethod + 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], + model_config: ModelConfig, + ): ... + @classmethod + def bias( + cls, + data_config: DataConfig, + 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", + ): ... + +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..74e6d31012f3 --- /dev/null +++ b/stubs/sagemaker/sagemaker/collection.pyi @@ -0,0 +1,11 @@ +from _typeshed import Incomplete + +class Collection: + sagemaker_session: Incomplete + def __init__(self, sagemaker_session) -> None: ... + 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 = 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..a4317ffe3669 --- /dev/null +++ b/stubs/sagemaker/sagemaker/config/config.pyi @@ -0,0 +1,9 @@ +from _typeshed import Incomplete + +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 = None, s3_resource: Incomplete | None = None) -> dict: ... +def validate_sagemaker_config(sagemaker_config: dict | None = 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..acc1b8503a40 --- /dev/null +++ b/stubs/sagemaker/sagemaker/content_types.pyi @@ -0,0 +1,19 @@ +def retrieve_options( + 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]: ... +def retrieve_default( + 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: ... + +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..29a952153239 --- /dev/null +++ b/stubs/sagemaker/sagemaker/debugger/debugger.pyi @@ -0,0 +1,124 @@ +from _typeshed import Incomplete +from abc import ABC + +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: str | PipelineVariable, + instance_type: str | PipelineVariable, + volume_size_in_gb: int | PipelineVariable, + 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: ... + 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: 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: str | PipelineVariable | None = None + ) -> None: ... + +class CollectionConfig: + name: Incomplete + parameters: Incomplete + 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/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..8d28f4bada68 --- /dev/null +++ b/stubs/sagemaker/sagemaker/debugger/profiler_config.pyi @@ -0,0 +1,19 @@ +from _typeshed import Incomplete + +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: 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/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..9d12487cb3c4 --- /dev/null +++ b/stubs/sagemaker/sagemaker/deserializers.pyi @@ -0,0 +1,28 @@ +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: 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]: ... +def retrieve_default( + 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/__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..173a88f2b67a --- /dev/null +++ b/stubs/sagemaker/sagemaker/djl_inference/model.pyi @@ -0,0 +1,169 @@ +from _typeshed import Incomplete +from collections.abc import Callable +from enum import Enum +from typing import Any + +from sagemaker.base_predictor 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 = 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: str | None = None, + task: str | None = None, + dtype: str = "fp32", + 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: int | None = None, + prediction_timeout: int | None = None, + entry_point: str | None = None, + image_uri: str | PipelineVariable | None = None, + predictor_cls: Callable[[str, Session], Predictor] | Predictor = ..., + **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 = 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 = 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: 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 partition( + self, + instance_type: str, + 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 = 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: 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 partition( + self, + instance_type: str, + 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 = None, + enable_network_isolation: bool = False, + ): ... + +class FasterTransformerModel(DJLModel): + number_of_partitions: Incomplete + 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 new file mode 100644 index 000000000000..8746a676c54c --- /dev/null +++ b/stubs/sagemaker/sagemaker/drift_check_baselines.pyi @@ -0,0 +1,26 @@ +from _typeshed import Incomplete + +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: 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 new file mode 100644 index 000000000000..f32ee3020a9c --- /dev/null +++ b/stubs/sagemaker/sagemaker/environment_variables.pyi @@ -0,0 +1,12 @@ +from _typeshed import Incomplete + +logger: Incomplete + +def retrieve_default( + 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]: ... diff --git a/stubs/sagemaker/sagemaker/estimator.pyi b/stubs/sagemaker/sagemaker/estimator.pyi new file mode 100644 index 000000000000..41dc88730b1a --- /dev/null +++ b/stubs/sagemaker/sagemaker/estimator.pyi @@ -0,0 +1,392 @@ +import abc +from _typeshed import Incomplete +from abc import abstractmethod + +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 = 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: str | PipelineVariable | None = None, + max_run: int | PipelineVariable = 86400, + input_mode: str | PipelineVariable = "File", + 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: list[dict[str, str | PipelineVariable]] | None = None, + encrypt_inter_container_traffic: bool | PipelineVariable | None = None, + use_spot_instances: bool | PipelineVariable = False, + 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: 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: ... + @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: str | dict | TrainingInput | FileSystemInput | None = None, + wait: bool = True, + logs: str = "All", + job_name: str | None = None, + experiment_config: dict[str, str] | None = 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: str | PipelineVariable, + 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: str | PipelineVariable | None = None, + max_run: int | PipelineVariable = 86400, + input_mode: str | PipelineVariable = "File", + 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: list[dict[str, str | PipelineVariable]] | None = None, + encrypt_inter_container_traffic: bool | PipelineVariable | None = None, + use_spot_instances: bool | PipelineVariable = False, + 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: 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: 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: ... + 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: str | PipelineVariable, + source_dir: str | PipelineVariable | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, + container_log_level: int | PipelineVariable = 20, + 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: ... + 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..848ea430c7d3 --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/_utils.pyi @@ -0,0 +1,11 @@ +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: 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/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..29f4a6b38f81 --- /dev/null +++ b/stubs/sagemaker/sagemaker/experiments/run.pyi @@ -0,0 +1,89 @@ +import datetime +from _typeshed import Incomplete +from enum import Enum +from typing import Any + +from numpy import ndarray +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: str | None = None, + experiment_display_name: str | None = None, + run_display_name: str | None = None, + tags: list[dict[str, str]] | None = None, + sagemaker_session: Session | None = None, + ) -> None: ... + @property + 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[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[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[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): ... + def close(self) -> None: ... + def __enter__(self): ... + 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: Session | None = None +) -> Run: ... +def list_runs( + experiment_name: str, + created_before: datetime.datetime | None = None, + created_after: datetime.datetime | None = None, + sagemaker_session: Session | None = None, + max_results: int | None = None, + next_token: str | None = None, + sort_by: SortByType = ..., + sort_order: SortOrderType = ..., +) -> list[Run]: ... 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..50e98f45d19a --- /dev/null +++ b/stubs/sagemaker/sagemaker/explainer/clarify_explainer_config.pyi @@ -0,0 +1,67 @@ +from _typeshed import Incomplete + +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: str | None = "text/csv", shap_baseline: str | None = None, shap_baseline_uri: str | None = 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: int | None = None, + seed: int | None = None, + use_logit: bool | None = False, + text_config: ClarifyTextConfig | None = 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: 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: + enable_explanations: Incomplete + shap_config: Incomplete + inference_config: Incomplete + def __init__( + self, + shap_config: ClarifyShapConfig, + 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 new file mode 100644 index 000000000000..3fa08fb928e4 --- /dev/null +++ b/stubs/sagemaker/sagemaker/explainer/explainer_config.pyi @@ -0,0 +1,7 @@ +from _typeshed import Incomplete + +from sagemaker.explainer.clarify_explainer_config import ClarifyExplainerConfig + +class ExplainerConfig: + clarify_explainer_config: Incomplete + def __init__(self, clarify_explainer_config: ClarifyExplainerConfig | None = 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..dd2ca1e51cc5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/feature_store/dataset_builder.pyi @@ -0,0 +1,129 @@ +import datetime +from enum import Enum + +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 = None, + feature_name_in_target: str | None = 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 = None, + included_feature_names: list[str] | None = None, + feature_name_in_target: str | None = 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 = 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, + 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..6b6394f84040 --- /dev/null +++ b/stubs/sagemaker/sagemaker/feature_store/feature_definition.pyi @@ -0,0 +1,28 @@ +from enum import Enum +from typing import Any + +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..8d949c3a4806 --- /dev/null +++ b/stubs/sagemaker/sagemaker/feature_store/feature_group.pyi @@ -0,0 +1,118 @@ +from _typeshed import Incomplete +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 +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 = None, workgroup: str | None = 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: str | bool, + record_identifier_name: str, + event_time_feature_name: str, + 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 = None, + disable_glue_table_creation: bool = False, + 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 = 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 = 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 = 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: 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 = 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..e516a205c81e --- /dev/null +++ b/stubs/sagemaker/sagemaker/feature_store/feature_store.pyi @@ -0,0 +1,49 @@ +import datetime +from collections.abc import Sequence +from typing import Any + +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: FeatureGroup | pd.DataFrame, + output_path: str, + 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 = 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 = 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): ... + 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..ca07c76c3135 --- /dev/null +++ b/stubs/sagemaker/sagemaker/feature_store/inputs.pyi @@ -0,0 +1,135 @@ +import abc +from enum import Enum +from typing import Any + +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..35cea594c4f2 --- /dev/null +++ b/stubs/sagemaker/sagemaker/fw_utils.pyi @@ -0,0 +1,74 @@ +from _typeshed import Incomplete +from typing import NamedTuple + +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: 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): ... +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: SessionSettings | None = 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..81feb5d86458 --- /dev/null +++ b/stubs/sagemaker/sagemaker/huggingface/estimator.pyi @@ -0,0 +1,45 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + 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): ... + 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..c3e2460137e4 --- /dev/null +++ b/stubs/sagemaker/sagemaker/huggingface/llm_utils.pyi @@ -0,0 +1,5 @@ +from sagemaker.session import Session + +def get_huggingface_llm_image_uri( + 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 new file mode 100644 index 000000000000..1ed4530b0033 --- /dev/null +++ b/stubs/sagemaker/sagemaker/huggingface/model.pyi @@ -0,0 +1,98 @@ +from _typeshed import Incomplete + +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: 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: int | PipelineVariable | None = 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[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: 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, + 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..34bd8c3191dd --- /dev/null +++ b/stubs/sagemaker/sagemaker/huggingface/processing.pyi @@ -0,0 +1,34 @@ +from _typeshed import Incomplete + +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: 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: str | PipelineVariable | None = None, + command: list[str] | None = None, + volume_size_in_gb: int | PipelineVariable = 30, + 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/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..7d5f4a15d459 --- /dev/null +++ b/stubs/sagemaker/sagemaker/huggingface/training_compiler/config.pyi @@ -0,0 +1,13 @@ +from _typeshed import Incomplete + +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: bool | PipelineVariable = True, debug: 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..12af8fd60218 --- /dev/null +++ b/stubs/sagemaker/sagemaker/hyperparameters.pyi @@ -0,0 +1,21 @@ +from _typeshed import Incomplete + +from sagemaker.jumpstart.enums import HyperparameterValidationMode + +logger: Incomplete + +def retrieve_default( + 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]: ... +def validate( + 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/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..0abfb3101b79 --- /dev/null +++ b/stubs/sagemaker/sagemaker/inference_recommender/inference_recommender_mixin.pyi @@ -0,0 +1,35 @@ +from _typeshed import Incomplete + +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 = 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 new file mode 100644 index 000000000000..c82718fcd154 --- /dev/null +++ b/stubs/sagemaker/sagemaker/inputs.pyi @@ -0,0 +1,85 @@ +from _typeshed import Incomplete +from typing import Any, Literal + +from sagemaker.workflow.entities import PipelineVariable + +FILE_SYSTEM_TYPES: Incomplete +FILE_SYSTEM_ACCESS_MODES: Incomplete + +class TrainingInput: + config: Incomplete + def __init__( + self, + s3_data: str | PipelineVariable, + 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: 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: ShuffleConfig | None = 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[Literal["InvocationsTimeoutInSeconds", "InvocationsMaxRetries"], int] + batch_data_capture_config: dict[Any, Any] + 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 = None, generate_inference_id: bool | None = 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..a6b9e5865814 --- /dev/null +++ b/stubs/sagemaker/sagemaker/instance_types.pyi @@ -0,0 +1,20 @@ +from _typeshed import Incomplete + +logger: Incomplete + +def retrieve_default( + 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: 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]: ... 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..38535a13b75d --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/accessors.pyi @@ -0,0 +1,21 @@ +from typing import Any + +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) -> None: ... + @staticmethod + def reset_cache(cache_kwargs: dict[str, Any] | None = None, region: str | None = None) -> None: ... + @staticmethod + def get_manifest(cache_kwargs: dict[str, Any] | None = ..., region: str | None = ...) -> 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..b8ae6b53fc10 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/cache.pyi @@ -0,0 +1,29 @@ +import datetime +from _typeshed import Incomplete + +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: str | None = None, + s3_client_config: botocore.config.Config | None = 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..658de1fac609 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/constants.pyi @@ -0,0 +1,27 @@ +from _typeshed import Incomplete + +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..2fb009cf93f4 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/estimator.pyi @@ -0,0 +1,126 @@ +from _typeshed import Incomplete + +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: 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: 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: 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 new file mode 100644 index 000000000000..c9578364fe25 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/exceptions.pyi @@ -0,0 +1,25 @@ +from _typeshed import Incomplete + +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: str | None = None) -> None: ... + +class VulnerableJumpStartModelError(ValueError): + message: Incomplete + def __init__( + self, + 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: str | None = None, version: str | None = None, message: str | None = 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..5be3db1249ea --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/factory/estimator.pyi @@ -0,0 +1,121 @@ +from _typeshed import Incomplete + +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: 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: 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: 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 new file mode 100644 index 000000000000..9dbb7cf24825 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/factory/model.pyi @@ -0,0 +1,73 @@ +from _typeshed import Incomplete + +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: 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: 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 new file mode 100644 index 000000000000..8c9abccb4fde --- /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 + +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: 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: Operand | str) -> None: ... + def __iter__(self) -> Any: ... + resolved_value: Incomplete + def eval(self) -> None: ... + +class Or(Operator): + operands: Incomplete + 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: 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: str | bool | 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..6ebffc54c7cb --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/model.pyi @@ -0,0 +1,69 @@ +from _typeshed import Incomplete + +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: 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: 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 new file mode 100644 index 000000000000..18680b92f0dd --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/notebook_utils.pyi @@ -0,0 +1,14 @@ +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 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]]]: ... +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..3227d171d8d5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/types.pyi @@ -0,0 +1,452 @@ +from _typeshed import Incomplete +from collections.abc import Callable +from enum import Enum +from typing import Any + +from sagemaker import Predictor + +class JumpStartDataHolderType: + def __eq__(self, other: object) -> 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: 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: dict[str, Any] | None) -> 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: dict[JumpStartVersionedModelId | JumpStartModelHeader, JumpStartModelSpecs], + md5_hash: str | None = 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: 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[..., Predictor] | Predictor | 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: 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): + 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: 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): + 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: 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: 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: 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): + 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: 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): + 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: 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[..., 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, + sagemaker_session: Any | None = None, + enable_network_isolation: bool | Any = 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 new file mode 100644 index 000000000000..6868e7585985 --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/utils.pyi @@ -0,0 +1,48 @@ +from _typeshed import Incomplete +from typing import Any + +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: 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 get_jumpstart_base_name_if_jumpstart_model(*uris: str | None) -> str | None: ... +def add_jumpstart_tags( + 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: 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: 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: Any | None, sagemaker_session: Session, default_value: str | None = None +) -> Any: ... +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 = ... +) -> bool: ... diff --git a/stubs/sagemaker/sagemaker/jumpstart/validators.pyi b/stubs/sagemaker/sagemaker/jumpstart/validators.pyi new file mode 100644 index 000000000000..1babe6e0594d --- /dev/null +++ b/stubs/sagemaker/sagemaker/jumpstart/validators.pyi @@ -0,0 +1,11 @@ +from typing import Any + +from sagemaker.jumpstart.enums import HyperparameterValidationMode + +def validate_hyperparameters( + model_id: str, + model_version: str, + hyperparameters: dict[str, Any], + validation_mode: HyperparameterValidationMode = ..., + region: str | None = "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..7dee30dd812a --- /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 = 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 = None, + environment: dict | None = None, + layers: list | None = 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..fe92d8d959ba --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/action.pyi @@ -0,0 +1,62 @@ +from _typeshed import Incomplete +from collections.abc import Iterator +from datetime import datetime + +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 = 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: 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]: ... + +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..ca8e10f96ae3 --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/artifact.pyi @@ -0,0 +1,73 @@ +from _typeshed import Incomplete +from collections.abc import Iterator +from datetime import datetime + +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: 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: 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 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..6b0b8d51cf8a --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/association.pyi @@ -0,0 +1,39 @@ +from _typeshed import Incomplete +from collections.abc import Iterator +from datetime import datetime + +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 = None, + sagemaker_session: Incomplete | None = None, + ) -> Association: ... + @classmethod + def list( + cls, + 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 new file mode 100644 index 000000000000..105b82a4801b --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/context.pyi @@ -0,0 +1,67 @@ +from _typeshed import Incomplete +from collections.abc import Iterator +from datetime import datetime + +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 = 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: 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]: ... + +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..3af3aaa43f15 --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/lineage_trial_component.pyi @@ -0,0 +1,33 @@ +from _typeshed import Incomplete + +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..b3c5cda6ec0b --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/query.pyi @@ -0,0 +1,93 @@ +from _typeshed import Incomplete +from datetime import datetime +from enum import Enum +from typing import Any + +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: 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 = 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 + sources: Incomplete + created_before: Incomplete + created_after: Incomplete + modified_before: Incomplete + modified_after: Incomplete + properties: Incomplete + def __init__( + self, + 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], + direction: LineageQueryDirectionEnum = ..., + include_edges: bool = True, + 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 new file mode 100644 index 000000000000..0eb962e8b7ee --- /dev/null +++ b/stubs/sagemaker/sagemaker/lineage/visualizer.pyi @@ -0,0 +1,16 @@ +from pandas import DataFrame as DataFrame + +class LineageTableVisualizer: + def __init__(self, sagemaker_session) -> None: ... + def show( + self, + 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/__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..3175cff45bb6 --- /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 = 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..f590eb33bc6c --- /dev/null +++ b/stubs/sagemaker/sagemaker/local/pipeline.pyi @@ -0,0 +1,48 @@ +import abc +from _typeshed import Incomplete +from abc import ABC, abstractmethod + +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..4821a30d5b94 --- /dev/null +++ b/stubs/sagemaker/sagemaker/metadata_properties.pyi @@ -0,0 +1,16 @@ +from _typeshed import Incomplete + +from sagemaker.workflow.entities import PipelineVariable + +class MetadataProperties: + commit_id: Incomplete + repository: Incomplete + generated_by: Incomplete + project_id: Incomplete + def __init__( + self, + 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 new file mode 100644 index 000000000000..a9c77cf52f47 --- /dev/null +++ b/stubs/sagemaker/sagemaker/metric_definitions.pyi @@ -0,0 +1,11 @@ +from _typeshed import Incomplete + +logger: Incomplete + +def retrieve_default( + 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[dict[str, str]] | None: ... diff --git a/stubs/sagemaker/sagemaker/model.pyi b/stubs/sagemaker/sagemaker/model.pyi new file mode 100644 index 000000000000..6a547b135b88 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model.pyi @@ -0,0 +1,210 @@ +import abc +from _typeshed import Incomplete + +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: str | PipelineVariable, + 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: 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: 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: 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: 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, + 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: str | PipelineVariable, + image_uri: str | PipelineVariable, + 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: str | None = None, + sagemaker_session: Session | None = None, + dependencies: list[str] | None = None, + git_config: dict[str, str] | None = 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..35cd0b665aad --- /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, Literal + +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 = None): ... + def __set__(self, obj: object, value: Any): ... + @abstractmethod + def validate(self, value): ... + @abstractmethod + def require_decode(self, value: dict[Any, Any]): ... + @abstractmethod + def decode(self, value: dict[Any, Any]): ... + +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: Enum | str): ... + def decode(self, value: str): ... + +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[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[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[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: _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 new file mode 100644 index 000000000000..15c02813d828 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_card/model_card.pyi @@ -0,0 +1,267 @@ +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 +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: str | None = None, + model_name: str | None = None, + model_description: str | None = None, + model_version: 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 = 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: 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): + business_problem: Incomplete + 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 + ) -> None: ... + +class Function(_DefaultToRequestDict, _DefaultFromDict): + function: Incomplete + facet: Incomplete + condition: Incomplete + def __init__( + self, + 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: str | None = 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: MetricTypeEnum | str, + value: float | str | bool | list[Any], + notes: str | 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: float | str | bool | list[Any]): ... + +class TrainingMetric(_DefaultToRequestDict, _DefaultFromDict): + name: Incomplete + value: Incomplete + notes: Incomplete + def __init__(self, name: str, value: float, notes: str | None = 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: 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): + objective_function: Incomplete + training_job_details: Incomplete + training_observations: Incomplete + def __init__( + self, + 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 = None, **kwargs): ... + @classmethod + 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: list[Metric] | None = 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: str | None = None, + evaluation_job_arn: str | None = None, + datasets: list[str] | None = None, + metadata: dict[Any, Any] | None = None, + metric_groups: list[MetricGroup] | None = 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: str | None = None, + caveats_and_recommendations: str | None = None, + custom_details: dict[Any, Any] | None = 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: ModelCardStatusEnum | str | None = ..., + arn: str | None = None, + version: int | None = None, + created_time: datetime | None = None, + created_by: dict[Any, Any] | None = None, + last_modified_time: datetime | 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, + 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: 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 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: 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 = None): ... + @staticmethod + def list_export_jobs(model_card_name: str, sagemaker_session: Session | None = 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..970ab668b03a --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_metrics.pyi @@ -0,0 +1,46 @@ +from _typeshed import Incomplete + +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: 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: + content_type: Incomplete + s3_uri: Incomplete + content_digest: Incomplete + def __init__( + self, + content_type: str | PipelineVariable, + s3_uri: str | PipelineVariable, + content_digest: str | PipelineVariable | None = None, + ) -> None: ... + +class FileSource: + content_type: Incomplete + s3_uri: Incomplete + content_digest: Incomplete + def __init__( + self, + s3_uri: str | PipelineVariable, + content_digest: str | PipelineVariable | None = None, + content_type: str | PipelineVariable | None = 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..3f1ed8c00526 --- /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) -> 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) -> 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..4de57d55dd71 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_monitor/model_monitoring.pyi @@ -0,0 +1,381 @@ +from _typeshed import Incomplete + +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: int | None, evaluation_period: int | None + ): ... + def list_monitoring_alerts(self, next_token: str | None = None, max_results: int | None = 10): ... + def list_monitoring_alert_history( + self, + 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): ... + @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: 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 = 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: + 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..968de9f1dd17 --- /dev/null +++ b/stubs/sagemaker/sagemaker/model_uris.pyi @@ -0,0 +1,12 @@ +from _typeshed import Incomplete + +logger: Incomplete + +def retrieve( + 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 new file mode 100644 index 000000000000..756fe1720ab6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/multidatamodel.pyi @@ -0,0 +1,49 @@ +from _typeshed import Incomplete +from collections.abc import Generator + +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: Model | None = None, + image_uri: str | PipelineVariable | None = None, + role: str | None = None, + sagemaker_session: Session | None = 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..e257089c0ec1 --- /dev/null +++ b/stubs/sagemaker/sagemaker/mxnet/estimator.pyi @@ -0,0 +1,32 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + 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( + 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..55e0d2bea7a8 --- /dev/null +++ b/stubs/sagemaker/sagemaker/mxnet/model.pyi @@ -0,0 +1,68 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + role: str | None = None, + entry_point: str | None = None, + framework_version: str = "1.4.0", + py_version: str | None = None, + image_uri: str | PipelineVariable | None = None, + predictor_cls: callable = ..., + 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: 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: 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, + 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..35e6d82fdca7 --- /dev/null +++ b/stubs/sagemaker/sagemaker/mxnet/processing.pyi @@ -0,0 +1,28 @@ +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: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + py_version: str = "py3", + image_uri: str | PipelineVariable | None = None, + command: list[str] | None = None, + volume_size_in_gb: int | PipelineVariable = 30, + 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 new file mode 100644 index 000000000000..e676e919bede --- /dev/null +++ b/stubs/sagemaker/sagemaker/network.pyi @@ -0,0 +1,16 @@ +from _typeshed import Incomplete + +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: 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/parameter.pyi b/stubs/sagemaker/sagemaker/parameter.pyi new file mode 100644 index 000000000000..d57ba98476cc --- /dev/null +++ b/stubs/sagemaker/sagemaker/parameter.pyi @@ -0,0 +1,39 @@ +from _typeshed import Incomplete + +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: float | PipelineVariable, + max_value: float | PipelineVariable, + scaling_type: 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..33f727c1e029 --- /dev/null +++ b/stubs/sagemaker/sagemaker/pipeline.pyi @@ -0,0 +1,85 @@ +from _typeshed import Incomplete + +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 = 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( + 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[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: 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, + 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..9d7cce2d60ea --- /dev/null +++ b/stubs/sagemaker/sagemaker/predictor.pyi @@ -0,0 +1,12 @@ +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: 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/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..1bc931ad35ec --- /dev/null +++ b/stubs/sagemaker/sagemaker/processing.pyi @@ -0,0 +1,229 @@ +from _typeshed import Incomplete + +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 = 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: 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: ProcessingInput | None = None, + outputs: ProcessingOutput | None = None, + arguments: list[str | PipelineVariable] | None = None, + wait: bool = True, + logs: bool = True, + 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: 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: 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 + ): ... + latest_job: Incomplete + def run( + self, + code: str, + inputs: ProcessingInput | None = None, + outputs: ProcessingOutput | None = None, + arguments: list[str | PipelineVariable] | None = None, + wait: bool = True, + logs: bool = True, + job_name: str | None = None, + experiment_config: dict[str, str] | None = None, + kms_key: str | None = 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: 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: S3Input | None = None, + dataset_definition: DatasetDefinition | None = None, + app_managed: 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: 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: FeatureStoreOutput | None = 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: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + py_version: str = "py3", + image_uri: str | PipelineVariable | None = None, + command: list[str] | None = None, + volume_size_in_gb: int | PipelineVariable = 30, + 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, + 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: 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: str | None = None, + experiment_config: dict[str, str] | None = None, + kms_key: str | None = 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..c81e77625b06 --- /dev/null +++ b/stubs/sagemaker/sagemaker/pytorch/estimator.pyi @@ -0,0 +1,39 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + 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): ... + 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..b3650cbc0a7b --- /dev/null +++ b/stubs/sagemaker/sagemaker/pytorch/model.pyi @@ -0,0 +1,68 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + role: str | None = None, + entry_point: str | None = None, + framework_version: str = "1.3", + py_version: str | None = None, + image_uri: str | PipelineVariable | None = None, + predictor_cls: callable = ..., + 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: 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: 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, + 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..4dbc8ba470bc --- /dev/null +++ b/stubs/sagemaker/sagemaker/pytorch/processing.pyi @@ -0,0 +1,28 @@ +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: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + py_version: str = "py3", + image_uri: str | PipelineVariable | None = None, + command: list[str] | None = None, + volume_size_in_gb: int | PipelineVariable = 30, + 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/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..7d5f4a15d459 --- /dev/null +++ b/stubs/sagemaker/sagemaker/pytorch/training_compiler/config.pyi @@ -0,0 +1,13 @@ +from _typeshed import Incomplete + +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: bool | PipelineVariable = True, debug: 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..5e5251d3305e --- /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 + +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 = 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 = 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 = 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 = None, + spark_config: SparkConfig | None = 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 = 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 = 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 = 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 = None, + spark_config: SparkConfig | None = 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 = None) -> Any: ... + def wait(self, timeout: int | None = 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..c1de023558f8 --- /dev/null +++ b/stubs/sagemaker/sagemaker/remote_function/job.pyi @@ -0,0 +1,101 @@ +from _typeshed import Incomplete + +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 = 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 = 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 = 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 = None, + spark_config: SparkConfig | None = 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 = 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..9a6260b7a970 --- /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 = 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): + 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..fdbe2f02db1a --- /dev/null +++ b/stubs/sagemaker/sagemaker/remote_function/spark_config.pyi @@ -0,0 +1,11 @@ +class SparkConfig: + 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): ... + 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..227b5fcf3d99 --- /dev/null +++ b/stubs/sagemaker/sagemaker/rl/estimator.pyi @@ -0,0 +1,54 @@ +import enum +from _typeshed import Incomplete + +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: str | PipelineVariable, + 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( + 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..45ecce1fd8fe --- /dev/null +++ b/stubs/sagemaker/sagemaker/s3.pyi @@ -0,0 +1,31 @@ +import io +from _typeshed import Incomplete + +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: 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..7f8bfab41678 --- /dev/null +++ b/stubs/sagemaker/sagemaker/s3_utils.pyi @@ -0,0 +1,9 @@ +from _typeshed import Incomplete + +logger: Incomplete + +def parse_s3_url(url): ... +def s3_path_join(*args, with_end_slash: bool = False): ... +def determine_bucket_and_prefix( + 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 new file mode 100644 index 000000000000..a10f72891c1d --- /dev/null +++ b/stubs/sagemaker/sagemaker/script_uris.pyi @@ -0,0 +1,12 @@ +from _typeshed import Incomplete + +logger: Incomplete + +def retrieve( + 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 new file mode 100644 index 000000000000..cb33e7b60360 --- /dev/null +++ b/stubs/sagemaker/sagemaker/serializers.pyi @@ -0,0 +1,27 @@ +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: 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]: ... +def retrieve_default( + 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/__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..2c16bb841418 --- /dev/null +++ b/stubs/sagemaker/sagemaker/serverless/serverless_inference_config.pyi @@ -0,0 +1,9 @@ +from _typeshed import Incomplete + +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: int | None = None + ) -> None: ... diff --git a/stubs/sagemaker/sagemaker/session.pyi b/stubs/sagemaker/sagemaker/session.pyi new file mode 100644 index 000000000000..615cb191192f --- /dev/null +++ b/stubs/sagemaker/sagemaker/session.pyi @@ -0,0 +1,592 @@ +from _typeshed import Incomplete +from collections.abc import Sequence +from typing import Any + +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 = None, + default_bucket_prefix: str | None = 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: dict[str, str] | None = 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: dict[str, str] | None = 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: str | None = None, max_results: int | None = 10 + ) -> dict: ... + def list_monitoring_alert_history( + self, + 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: ... + 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: 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 = 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 = 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, + 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 = 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 = 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 = 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 = 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: ... + def create_inference_recommendations_job( + self, + role: str, + sample_payload_url: str, + supported_content_types: list[str], + job_name: str | None = None, + job_type: str = "Default", + 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]: ... + +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..3a9f6264ff6c --- /dev/null +++ b/stubs/sagemaker/sagemaker/sklearn/estimator.pyi @@ -0,0 +1,32 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + framework_version: str | None = None, + py_version: str = "py3", + 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( + 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..cb00309a9bc9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/sklearn/model.pyi @@ -0,0 +1,62 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + role: str | None = None, + entry_point: str | None = None, + framework_version: str | None = None, + py_version: str = "py3", + image_uri: str | PipelineVariable | None = None, + predictor_cls: callable = ..., + 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: 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: 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, + 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..e26176dcf0df --- /dev/null +++ b/stubs/sagemaker/sagemaker/sklearn/processing.pyi @@ -0,0 +1,23 @@ +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: 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: 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/__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..37e2e7bde131 --- /dev/null +++ b/stubs/sagemaker/sagemaker/spark/processing.pyi @@ -0,0 +1,181 @@ +from _typeshed import Incomplete +from enum import Enum + +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: 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, + 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 = 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: 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, + 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: 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: 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 = 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: 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, + 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: str | PipelineVariable, + 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: 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: + 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..60904f7fc396 --- /dev/null +++ b/stubs/sagemaker/sagemaker/sparkml/model.pyi @@ -0,0 +1,19 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + role: str | None = None, + spark_version: str = "3.3", + sagemaker_session: Session | None = 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..4aa5bc44f6b5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/tensorflow/estimator.pyi @@ -0,0 +1,55 @@ +from _typeshed import Incomplete + +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: 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( + 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..e998f2bcae81 --- /dev/null +++ b/stubs/sagemaker/sagemaker/tensorflow/model.pyi @@ -0,0 +1,102 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + 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: 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: 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, + 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..d03dcd68a6b4 --- /dev/null +++ b/stubs/sagemaker/sagemaker/tensorflow/processing.pyi @@ -0,0 +1,28 @@ +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: str | PipelineVariable | None = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + py_version: str = "py3", + image_uri: str | PipelineVariable | None = None, + command: list[str] | None = None, + volume_size_in_gb: int | PipelineVariable = 30, + 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/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..36b17a3c1ed7 --- /dev/null +++ b/stubs/sagemaker/sagemaker/transformer.pyi @@ -0,0 +1,107 @@ +import abc +from _typeshed import Incomplete + +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: str | PipelineVariable, + instance_count: int | PipelineVariable, + instance_type: str | PipelineVariable, + 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: 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, + ): ... + def transform_with_monitoring( + self, + monitoring_config, + monitoring_resource_config, + data: str, + data_type: str = "S3Prefix", + 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 = None, + supplied_baseline_constraints: str | None = None, + wait: bool = True, + pipeline_name: str | None = None, + role: str | None = 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..829ef1790907 --- /dev/null +++ b/stubs/sagemaker/sagemaker/tuner.pyi @@ -0,0 +1,222 @@ +import abc +from _typeshed import Incomplete +from enum import Enum + +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[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: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + volume_size: 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 = None, + complete_on_convergence: bool | None = None, + target_objective_metric_value: float | None = 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: str | PipelineVariable, + 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 = None, + max_parallel_jobs: int | PipelineVariable = 1, + 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: str | None = None, + random_seed: int | None = None, + autotune: bool = False, + hyperparameters_to_keep_static: list[str] | None = None, + ) -> None: ... + def override_resource_config(self, instance_configs: list[InstanceConfig, dict[str, list[InstanceConfig]]]): ... + def fit( + self, + 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, + ): ... + @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..b0e1973cfe67 --- /dev/null +++ b/stubs/sagemaker/sagemaker/utilities/cache.pyi @@ -0,0 +1,25 @@ +import datetime +from _typeshed import Incomplete +from collections.abc import Callable +from typing import 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: bool | None = True) -> _ValType: ... + def put(self, key: _KeyType, value: _ValType | None = 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..7b8765fbc5cd --- /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 + +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: str | None = None): ... +def to_string(obj: object): ... +def get_module(module_name): ... +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 = None, + default_value: Incomplete | None = None, + sagemaker_session: Incomplete | None = None, + sagemaker_config: dict | None = None, +): ... +def get_sagemaker_config_value(sagemaker_session, key, sagemaker_config: dict | None = None): ... +def resolve_class_attribute_from_config( + clazz: type | None, + instance: object | None, + 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 = None, + union_key_paths: list[list[str]] | None = 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..fafb6cab4687 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/_utils.pyi @@ -0,0 +1,97 @@ +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 +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 = None, + description: str | None = None, + source_dir: str | 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, + **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: 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, + 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 = None, + description: Incomplete | 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, + 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..1f812bf36e3c --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/automl_step.pyi @@ -0,0 +1,27 @@ +from _typeshed import Incomplete + +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 = 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: ... + @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..c2c948176b20 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/callback_step.pyi @@ -0,0 +1,45 @@ +from _typeshed import Incomplete +from enum import Enum + +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 = None, + description: str | None = None, + cache_config: CacheConfig | None = None, + depends_on: list[str | Step | StepCollection] | None = 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..4b1a6556e8e5 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/clarify_check_step.pyi @@ -0,0 +1,89 @@ +from _typeshed import Incomplete +from abc import ABC + +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: 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: 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: 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: bool | PipelineVariable = False, + fail_on_violation: bool | PipelineVariable = True, + register_new_baseline: bool | PipelineVariable = False, + 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: ... + @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..dadc03f3ffaf --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/condition_step.pyi @@ -0,0 +1,32 @@ +from _typeshed import Incomplete + +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: 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: ... + @property + def step_only_arguments(self): ... + @property + def properties(self): ... + +class JsonGet(NewJsonGet): + 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 new file mode 100644 index 000000000000..5713f9f23e3f --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/conditions.pyi @@ -0,0 +1,82 @@ +import abc +from _typeshed import Incomplete +from enum import Enum +from typing import TypeAlias + +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: TypeAlias = 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: ConditionValueType | PrimitiveType + right: 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: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType) -> None: ... + +class ConditionGreaterThan(ConditionComparison): + def __init__(self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType) -> None: ... + +class ConditionGreaterThanOrEqualTo(ConditionComparison): + def __init__(self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType) -> None: ... + +class ConditionLessThan(ConditionComparison): + def __init__(self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType) -> None: ... + +class ConditionLessThanOrEqualTo(ConditionComparison): + def __init__(self, left: ConditionValueType | PrimitiveType, right: ConditionValueType | PrimitiveType) -> None: ... + +class ConditionIn(Condition): + value: Incomplete + in_values: Incomplete + def __init__( + self, value: ConditionValueType | PrimitiveType, in_values: list[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) -> None: ... + def to_request(self) -> RequestType: ... + +def primitive_or_expr( + 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 new file mode 100644 index 000000000000..1d9c754ac637 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/emr_step.pyi @@ -0,0 +1,48 @@ +from _typeshed import Incomplete +from typing import Any + +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 = None, main_class: str | None = None, properties: list[dict] | None = 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: 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: ... + @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..66ebfce4ea2d --- /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, TypeAlias + +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 + 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: 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..25d49d60a8ad --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/fail_step.pyi @@ -0,0 +1,20 @@ +from _typeshed import Incomplete + +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: 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: ... + @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..db8bde5182f8 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/functions.pyi @@ -0,0 +1,26 @@ +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: 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..c9d64cecc3c8 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/lambda_step.pyi @@ -0,0 +1,46 @@ +from _typeshed import Incomplete +from enum import Enum + +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 = 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: ... + @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..3338fc8745c2 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/model_step.pyi @@ -0,0 +1,24 @@ +from _typeshed import Incomplete + +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: 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 new file mode 100644 index 000000000000..2453f744c93d --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/monitor_batch_transform_step.pyi @@ -0,0 +1,25 @@ +from _typeshed import Incomplete + +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: QualityCheckConfig | ClarifyCheckConfig, + check_job_configuration: CheckJobConfig, + monitor_before_transform: bool = False, + fail_on_violation: bool | PipelineVariable = True, + 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/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..05144e5707c8 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/parameters.pyi @@ -0,0 +1,46 @@ +from _typeshed import Incomplete +from enum import Enum + +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 = 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) -> None: ... + +class ParameterFloat(Parameter): + 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 new file mode 100644 index 000000000000..a5808ff2b4e6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/pipeline.pyi @@ -0,0 +1,99 @@ +from _typeshed import Incomplete +from collections.abc import Sequence +from typing import Any + +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: 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 = 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 = None, + description: str | None = None, + parallelism_config: ParallelismConfiguration | None = None, + ) -> dict[str, Any]: ... + def upsert( + self, + 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 | 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 = 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 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[Step | StepCollection]) -> None: ... + @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 = 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..fff6fb56f835 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/pipeline_context.pyi @@ -0,0 +1,60 @@ +from _typeshed import Incomplete +from collections.abc import Callable + +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 = None, func: Callable | None = 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 = None, + default_bucket_prefix: str | None = None, + ) -> None: ... + @property + def context(self): ... + @context.setter + def context(self, value: _StepArguments | None = ...): ... + 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..ff69e202de8c --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/pipeline_experiment_config.pyi @@ -0,0 +1,25 @@ +from _typeshed import Incomplete + +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: str | Parameter | ExecutionVariable | Expression, + trial_name: 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..cfa979772a78 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/properties.pyi @@ -0,0 +1,45 @@ +from _typeshed import Incomplete +from abc import ABCMeta + +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 = None, + shape_name: str | None = None, + shape_names: list[str] | None = 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 = 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 = None, service_name: str = "sagemaker") -> None: ... + def __getitem__(self, item: 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..e1fa486a6fa1 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/quality_check_step.pyi @@ -0,0 +1,83 @@ +from _typeshed import Incomplete +from abc import ABC + +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: str | PipelineVariable + dataset_format: dict + 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): ... + 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: 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, + 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: bool | PipelineVariable = False, + fail_on_violation: bool | PipelineVariable = True, + register_new_baseline: bool | PipelineVariable = False, + 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: ... + @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..4460ea5a8066 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/retry.pyi @@ -0,0 +1,61 @@ +from _typeshed import Incomplete +from enum import Enum + +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 = None, + expire_after_mins: int | None = 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 = None, + failure_reason_types: list[SageMakerJobExceptionTypeEnum] | None = None, + backoff_rate: float = 2.0, + interval_seconds: int = 1, + 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 new file mode 100644 index 000000000000..7108b4cc5dde --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/selective_execution_config.pyi @@ -0,0 +1,9 @@ +from _typeshed import Incomplete + +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) -> 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..c2351233fdb6 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/step_collections.pyi @@ -0,0 +1,91 @@ +from _typeshed import Incomplete + +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 = None, + model_data: Incomplete | None = 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, + 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: Model | PipelineModel | 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, + ) -> 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 = None, + display_name: str | None = 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: 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 new file mode 100644 index 000000000000..692ae4e8f06c --- /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 Any + +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: str | None + description: str | None + step_type: StepTypeEnum + depends_on: list[str | Step | StepCollection] | None + @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[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 = 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: ... + +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 = None, + estimator: EstimatorBase | None = None, + display_name: str | None = None, + description: str | 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, + retry_policies: list[RetryPolicy] | None = 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: dict[Any, Any] | 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: ... + @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 = 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: ... + @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 = 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 + 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 = None, + tuner: HyperparameterTuner | None = None, + display_name: str | None = None, + description: str | None = None, + inputs: Incomplete | None = 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: ... + @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..863b46da20a9 --- /dev/null +++ b/stubs/sagemaker/sagemaker/workflow/utilities.pyi @@ -0,0 +1,23 @@ +from _typeshed import Incomplete +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 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..9d075c6f76c1 --- /dev/null +++ b/stubs/sagemaker/sagemaker/wrangler/ingestion.pyi @@ -0,0 +1,19 @@ +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 = 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 = 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, +): ... diff --git a/stubs/sagemaker/sagemaker/wrangler/processing.pyi b/stubs/sagemaker/sagemaker/wrangler/processing.pyi new file mode 100644 index 000000000000..14de5a1f2614 --- /dev/null +++ b/stubs/sagemaker/sagemaker/wrangler/processing.pyi @@ -0,0 +1,25 @@ +from _typeshed import Incomplete + +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 = 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 = 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/__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..04c78c0e6015 --- /dev/null +++ b/stubs/sagemaker/sagemaker/xgboost/estimator.pyi @@ -0,0 +1,34 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + framework_version: str, + source_dir: str | PipelineVariable | None = None, + hyperparameters: dict[str, str | PipelineVariable] | None = None, + py_version: str = "py3", + image_uri: str | PipelineVariable | None = None, + image_uri_region: str | None = 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..77900846aaa2 --- /dev/null +++ b/stubs/sagemaker/sagemaker/xgboost/model.pyi @@ -0,0 +1,62 @@ +from _typeshed import Incomplete + +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: str | PipelineVariable, + 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: int | PipelineVariable | None = None, + **kwargs, + ) -> None: ... + image_uri: Incomplete + def register( + self, + 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: 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, + 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..2b2b247edd49 --- /dev/null +++ b/stubs/sagemaker/sagemaker/xgboost/processing.pyi @@ -0,0 +1,28 @@ +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 = None, + instance_count: int | PipelineVariable | None = None, + instance_type: str | PipelineVariable | None = None, + py_version: str = "py3", + image_uri: str | PipelineVariable | None = None, + command: list[str] | None = None, + volume_size_in_gb: int | PipelineVariable = 30, + 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/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: ...