diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 59c1fd2aa9..471141db55 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -428,6 +428,7 @@ def __post_init__(self) -> None: logger.setLevel(self.log_level) self._warn_about_deprecated_parameters(warnings) + self._initialize_region() self._initialize_compute_config() self._initialize_network_config() self._initialize_defaults() @@ -496,17 +497,12 @@ def _initialize_network_config(self) -> None: if not hasattr(self, "_enable_network_isolation"): self._enable_network_isolation = False - def _initialize_defaults(self) -> None: - """Initialize default values for unset parameters.""" - if not hasattr(self, "model_name") or self.model_name is None: - self.model_name = "model-" + str(uuid.uuid4())[:8] - - if not hasattr(self, "mode") or self.mode is None: - self.mode = Mode.SAGEMAKER_ENDPOINT - - if not hasattr(self, "env_vars") or self.env_vars is None: - self.env_vars = {} + def _initialize_region(self) -> None: + """Resolve the region before any spec lookups that depend on it. + Must run before _initialize_compute_config: JumpStart default + instance type resolution requires self.region to exist. + """ # Set region with priority: user input > sagemaker session > AWS account region > default if not hasattr(self, "region") or not self.region: if self.sagemaker_session and self.sagemaker_session.boto_region_name: @@ -520,6 +516,21 @@ def _initialize_defaults(self) -> None: except Exception: self.region = None # Default fallback + def _initialize_defaults(self) -> None: + """Initialize default values for unset parameters.""" + if not hasattr(self, "model_name") or self.model_name is None: + self.model_name = "model-" + str(uuid.uuid4())[:8] + + if not hasattr(self, "mode") or self.mode is None: + self.mode = Mode.SAGEMAKER_ENDPOINT + + if not hasattr(self, "env_vars") or self.env_vars is None: + self.env_vars = {} + + # Region is resolved earlier in _initialize_region(); re-run for safety + # in case callers invoke _initialize_defaults directly. + self._initialize_region() + # At construction, only resolve a default role when none was supplied (so # building a ModelBuilder does no IAM work when a role is given). The # resolved role is validated for serving permissions at the actual diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py index e58ea4d7ad..4a3da090d1 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py @@ -322,8 +322,13 @@ def _get_jumpstart_recommended_instance_type(self) -> Optional[str]: if hasattr(deploy_kwargs, "instance_type") and deploy_kwargs.instance_type: return deploy_kwargs.instance_type - except Exception: - pass + except Exception as e: + logger.warning( + "Failed to retrieve JumpStart default instance type for model '%s': %s. " + "Falling back to generic instance type detection.", + self.model, + e, + ) return None diff --git a/sagemaker-serve/tests/unit/data/jumpstart_base_spec.json b/sagemaker-serve/tests/unit/data/jumpstart_base_spec.json new file mode 100644 index 0000000000..ef8fa262d8 --- /dev/null +++ b/sagemaker-serve/tests/unit/data/jumpstart_base_spec.json @@ -0,0 +1,841 @@ +{ + "hosting_resource_requirements": { + "num_accelerators": 1, + "min_memory_mb": 34360 + }, + "inference_volume_size": 123, + "training_volume_size": 456, + "dynamic_container_deployment_supported": true, + "model_id": "pytorch-ic-mobilenet-v2", + "url": "https://pytorch.org/hub/pytorch_vision_mobilenet_v2/", + "version": "3.0.6", + "min_sdk_version": "2.189.0", + "incremental_training_supported": true, + "hosting_ecr_specs": { + "framework": "pytorch", + "framework_version": "1.10.0", + "py_version": "py38" + }, + "hosting_artifact_uri": null, + "hosting_artifact_key": "pytorch-ic/pytorch-ic-mobilenet-v2/artifacts/inference/v2.0.0/", + "hosting_script_key": "source-directory-tarballs/pytorch/inference/ic/v2.0.0/sourcedir.tar.gz", + "training_supported": true, + "training_ecr_specs": { + "framework": "pytorch", + "framework_version": "1.10.0", + "py_version": "py38" + }, + "training_artifact_key": "pytorch-training/v2.0.0/train-pytorch-ic-mobilenet-v2.tar.gz", + "training_script_key": "source-directory-tarballs/pytorch/transfer_learning/ic/v2.3.0/sourcedir.tar.gz", + "hyperparameters": [ + { + "name": "train_only_top_layer", + "type": "text", + "options": [ + "True", + "False" + ], + "default": "True", + "scope": "algorithm" + }, + { + "name": "epochs", + "type": "int", + "default": 5, + "scope": "algorithm", + "min": 1, + "max": 1000 + }, + { + "name": "learning_rate", + "type": "float", + "default": 0.001, + "scope": "algorithm", + "min": 1e-08, + "max": 1 + }, + { + "name": "batch_size", + "type": "int", + "default": 4, + "scope": "algorithm", + "min": 1, + "max": 1024 + }, + { + "name": "reinitialize_top_layer", + "type": "text", + "options": [ + "Auto", + "True", + "False" + ], + "default": "Auto", + "scope": "algorithm" + }, + { + "name": "sagemaker_submit_directory", + "type": "text", + "default": "/opt/ml/input/data/code/sourcedir.tar.gz", + "scope": "container" + }, + { + "name": "sagemaker_program", + "type": "text", + "default": "transfer_learning.py", + "scope": "container" + }, + { + "name": "sagemaker_container_log_level", + "type": "text", + "default": "20", + "scope": "container" + } + ], + "inference_environment_variables": [ + { + "name": "SAGEMAKER_PROGRAM", + "type": "text", + "default": "inference.py", + "scope": "container", + "required_for_model_class": true + }, + { + "name": "SAGEMAKER_SUBMIT_DIRECTORY", + "type": "text", + "default": "/opt/ml/model/code", + "scope": "container", + "required_for_model_class": false + }, + { + "name": "SAGEMAKER_CONTAINER_LOG_LEVEL", + "type": "text", + "default": "20", + "scope": "container", + "required_for_model_class": false + }, + { + "name": "SAGEMAKER_MODEL_SERVER_TIMEOUT", + "type": "text", + "default": "3600", + "scope": "container", + "required_for_model_class": false + }, + { + "name": "ENDPOINT_SERVER_TIMEOUT", + "type": "int", + "default": 3600, + "scope": "container", + "required_for_model_class": true + }, + { + "name": "MODEL_CACHE_ROOT", + "type": "text", + "default": "/opt/ml/model", + "scope": "container", + "required_for_model_class": true + }, + { + "name": "SAGEMAKER_ENV", + "type": "text", + "default": "1", + "scope": "container", + "required_for_model_class": true + }, + { + "name": "SAGEMAKER_MODEL_SERVER_WORKERS", + "type": "int", + "default": 1, + "scope": "container", + "required_for_model_class": true + } + ], + "inference_vulnerable": false, + "inference_dependencies": [], + "inference_vulnerabilities": [], + "training_vulnerable": false, + "training_dependencies": [ + "sagemaker_jumpstart_prepack_script_utilities==1.0.0" + ], + "training_vulnerabilities": [], + "deprecated": false, + "usage_info_message": null, + "deprecated_message": null, + "deprecate_warn_message": null, + "default_inference_instance_type": "ml.m5.large", + "supported_inference_instance_types": [ + "ml.m5.large", + "ml.m5.xlarge", + "ml.c5.xlarge", + "ml.c5.2xlarge", + "ml.m4.large", + "ml.m4.xlarge" + ], + "default_training_instance_type": "ml.m5.xlarge", + "supported_training_instance_types": [ + "ml.m5.xlarge", + "ml.c5.2xlarge", + "ml.m4.xlarge" + ], + "metrics": [ + { + "Name": "pytorch-ic:val-accuracy", + "Regex": "val_accuracy: ([0-9\\.]+)" + } + ], + "training_prepacked_script_key": "source-directory-tarballs/pytorch/transfer_learning/ic/prepack/v1.1.0/sourcedir.tar.gz", + "hosting_prepacked_artifact_key": "pytorch-ic/pytorch-ic-mobilenet-v2/artifacts/inference-prepack/v1.0.0/", + "model_kwargs": {}, + "deploy_kwargs": {}, + "estimator_kwargs": { + "encrypt_inter_container_traffic": true, + "max_run": 360000 + }, + "fit_kwargs": {}, + "predictor_specs": { + "default_content_type": "application/x-image", + "supported_content_types": [ + "application/x-image" + ], + "default_accept_type": "application/json", + "supported_accept_types": [ + "application/json;verbose", + "application/json" + ] + }, + "inference_enable_network_isolation": true, + "training_enable_network_isolation": true, + "default_training_dataset_uri": null, + "default_training_dataset_key": "training-datasets/tf_flowers/", + "resource_name_base": "pt-ic-mobilenet-v2", + "hosting_eula_key": null, + "hosting_model_package_arns": {}, + "training_model_package_artifact_uris": null, + "hosting_use_script_uri": false, + "hosting_instance_type_variants": { + "regional_aliases": { + "af-south-1": { + "cpu_ecr_uri_1": "626614931356.dkr.ecr.af-south-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "626614931356.dkr.ecr.af-south-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "ap-east-1": { + "cpu_ecr_uri_1": "871362719292.dkr.ecr.ap-east-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "871362719292.dkr.ecr.ap-east-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "ap-northeast-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "ap-northeast-2": { + "cpu_ecr_uri_1": "763104351884.dkr.ecr.ap-northeast-2.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.ap-northeast-2.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "ap-northeast-3": { + "cpu_ecr_uri_1": "364406365360.dkr.ecr.ap-northeast-3.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "364406365360.dkr.ecr.ap-northeast-3.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "ap-south-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.ap-south-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.ap-south-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.ap-south-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "ap-southeast-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.ap-southeast-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.ap-southeast-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.ap-southeast-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "ap-southeast-2": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.ap-southeast-2.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.ap-southeast-2.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.ap-southeast-2.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "ap-southeast-3": { + "cpu_ecr_uri_1": "907027046896.dkr.ecr.ap-southeast-3.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "907027046896.dkr.ecr.ap-southeast-3.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "ap-southeast-5": { + "cpu_ecr_uri_1": "550225433462.dkr.ecr.ap-southeast-5.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "550225433462.dkr.ecr.ap-southeast-5.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "ca-central-1": { + "cpu_ecr_uri_1": "763104351884.dkr.ecr.ca-central-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.ca-central-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "cn-north-1": { + "cpu_ecr_uri_1": "727897471807.dkr.ecr.cn-north-1.amazonaws.com.cn/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "727897471807.dkr.ecr.cn-north-1.amazonaws.com.cn/pytorch-inference:1.10.0-gpu-py38" + }, + "cn-northwest-1": { + "cpu_ecr_uri_1": "727897471807.dkr.ecr.cn-northwest-1.amazonaws.com.cn/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "727897471807.dkr.ecr.cn-northwest-1.amazonaws.com.cn/pytorch-inference:1.10.0-gpu-py38" + }, + "eu-central-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.eu-central-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.eu-central-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.eu-central-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "eu-central-2": { + "cpu_ecr_uri_1": "380420809688.dkr.ecr.eu-central-2.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "380420809688.dkr.ecr.eu-central-2.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "eu-north-1": { + "cpu_ecr_uri_1": "763104351884.dkr.ecr.eu-north-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.eu-north-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "eu-south-1": { + "cpu_ecr_uri_1": "692866216735.dkr.ecr.eu-south-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "692866216735.dkr.ecr.eu-south-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "eu-west-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.eu-west-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.eu-west-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.eu-west-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "eu-west-2": { + "cpu_ecr_uri_1": "763104351884.dkr.ecr.eu-west-2.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.eu-west-2.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "eu-west-3": { + "cpu_ecr_uri_1": "763104351884.dkr.ecr.eu-west-3.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.eu-west-3.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "il-central-1": { + "cpu_ecr_uri_1": "780543022126.dkr.ecr.il-central-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "780543022126.dkr.ecr.il-central-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "me-central-1": { + "cpu_ecr_uri_1": "914824155844.dkr.ecr.me-central-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "914824155844.dkr.ecr.me-central-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "me-south-1": { + "cpu_ecr_uri_1": "217643126080.dkr.ecr.me-south-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "217643126080.dkr.ecr.me-south-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "sa-east-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.sa-east-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.sa-east-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.sa-east-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "us-east-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "us-east-2": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "us-gov-east-1": { + "cpu_ecr_uri_1": "446045086412.dkr.ecr.us-gov-east-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "446045086412.dkr.ecr.us-gov-east-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "us-gov-west-1": { + "cpu_ecr_uri_1": "442386744353.dkr.ecr.us-gov-west-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "442386744353.dkr.ecr.us-gov-west-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "us-west-1": { + "cpu_ecr_uri_1": "763104351884.dkr.ecr.us-west-1.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.us-west-1.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + }, + "us-west-2": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.10.0-gpu-py38" + } + }, + "aliases": null, + "variants": { + "c4": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "c5": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "c5d": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "c5n": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "c6i": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "c6id": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "c7i": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "g4dn": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "g5": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "g6": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "g6e": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "local": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "local_gpu": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "m4": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "m5": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "m5d": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "m6i": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "m6id": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "m7i": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "p2": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "p3": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "p3dn": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "p4d": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "p4de": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "p5": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "r5": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "r5d": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "r6i": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "r6id": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "r7i": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "t2": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "t3": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "trn1": { + "regional_properties": { + "image_uri": "$alias_ecr_uri_3" + } + }, + "trn1n": { + "regional_properties": { + "image_uri": "$alias_ecr_uri_3" + } + } + } + }, + "training_instance_type_variants": { + "regional_aliases": { + "af-south-1": { + "cpu_ecr_uri_1": "626614931356.dkr.ecr.af-south-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "626614931356.dkr.ecr.af-south-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "ap-east-1": { + "cpu_ecr_uri_1": "871362719292.dkr.ecr.ap-east-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "871362719292.dkr.ecr.ap-east-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "ap-northeast-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "ap-northeast-2": { + "cpu_ecr_uri_1": "763104351884.dkr.ecr.ap-northeast-2.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.ap-northeast-2.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "ap-northeast-3": { + "cpu_ecr_uri_1": "364406365360.dkr.ecr.ap-northeast-3.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "364406365360.dkr.ecr.ap-northeast-3.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "ap-south-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.ap-south-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.ap-south-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.ap-south-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "ap-southeast-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.ap-southeast-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.ap-southeast-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.ap-southeast-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "ap-southeast-2": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.ap-southeast-2.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.ap-southeast-2.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.ap-southeast-2.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "ap-southeast-3": { + "cpu_ecr_uri_1": "907027046896.dkr.ecr.ap-southeast-3.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "907027046896.dkr.ecr.ap-southeast-3.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "ap-southeast-5": { + "cpu_ecr_uri_1": "550225433462.dkr.ecr.ap-southeast-5.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "550225433462.dkr.ecr.ap-southeast-5.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "ca-central-1": { + "cpu_ecr_uri_1": "763104351884.dkr.ecr.ca-central-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.ca-central-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "cn-north-1": { + "cpu_ecr_uri_1": "727897471807.dkr.ecr.cn-north-1.amazonaws.com.cn/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "727897471807.dkr.ecr.cn-north-1.amazonaws.com.cn/pytorch-training:1.10.0-gpu-py38" + }, + "cn-northwest-1": { + "cpu_ecr_uri_1": "727897471807.dkr.ecr.cn-northwest-1.amazonaws.com.cn/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "727897471807.dkr.ecr.cn-northwest-1.amazonaws.com.cn/pytorch-training:1.10.0-gpu-py38" + }, + "eu-central-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.eu-central-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.eu-central-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.eu-central-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "eu-central-2": { + "cpu_ecr_uri_1": "380420809688.dkr.ecr.eu-central-2.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "380420809688.dkr.ecr.eu-central-2.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "eu-north-1": { + "cpu_ecr_uri_1": "763104351884.dkr.ecr.eu-north-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.eu-north-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "eu-south-1": { + "cpu_ecr_uri_1": "692866216735.dkr.ecr.eu-south-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "692866216735.dkr.ecr.eu-south-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "eu-west-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.eu-west-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.eu-west-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.eu-west-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "eu-west-2": { + "cpu_ecr_uri_1": "763104351884.dkr.ecr.eu-west-2.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.eu-west-2.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "eu-west-3": { + "cpu_ecr_uri_1": "763104351884.dkr.ecr.eu-west-3.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.eu-west-3.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "il-central-1": { + "cpu_ecr_uri_1": "780543022126.dkr.ecr.il-central-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "780543022126.dkr.ecr.il-central-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "me-central-1": { + "cpu_ecr_uri_1": "914824155844.dkr.ecr.me-central-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "914824155844.dkr.ecr.me-central-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "me-south-1": { + "cpu_ecr_uri_1": "217643126080.dkr.ecr.me-south-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "217643126080.dkr.ecr.me-south-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "sa-east-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.sa-east-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.sa-east-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.sa-east-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "us-east-1": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "us-east-2": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "us-gov-east-1": { + "cpu_ecr_uri_1": "446045086412.dkr.ecr.us-gov-east-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "446045086412.dkr.ecr.us-gov-east-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "us-gov-west-1": { + "cpu_ecr_uri_1": "442386744353.dkr.ecr.us-gov-west-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "442386744353.dkr.ecr.us-gov-west-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "us-west-1": { + "cpu_ecr_uri_1": "763104351884.dkr.ecr.us-west-1.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.us-west-1.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + }, + "us-west-2": { + "alias_ecr_uri_3": "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.4.0-ubuntu20.04", + "cpu_ecr_uri_1": "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.10.0-cpu-py38", + "gpu_ecr_uri_2": "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.10.0-gpu-py38" + } + }, + "aliases": null, + "variants": { + "c4": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "c5": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "c5d": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "c5n": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "c6i": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "c6id": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "c7i": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "g4dn": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "g5": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "g6": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "g6e": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "local": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "local_gpu": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "m4": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "m5": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "m5d": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "m6i": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "m6id": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "m7i": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "p2": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "p3": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "p3dn": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "p4d": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "p4de": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "p5": { + "regional_properties": { + "image_uri": "$gpu_ecr_uri_2" + } + }, + "r5": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "r5d": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "r6i": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "r6id": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "r7i": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "t2": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "t3": { + "regional_properties": { + "image_uri": "$cpu_ecr_uri_1" + } + }, + "trn1": { + "regional_properties": { + "image_uri": "$alias_ecr_uri_3" + } + }, + "trn1n": { + "regional_properties": { + "image_uri": "$alias_ecr_uri_3" + } + } + } + }, + "default_payloads": null, + "gated_bucket": false, + "model_subscription_link": null, + "hosting_additional_data_sources": null, + "hosting_neuron_model_id": null, + "hosting_neuron_model_version": null, + "inference_configs": null, + "inference_config_components": null, + "inference_config_rankings": null, + "training_configs": null, + "training_config_components": null, + "training_config_rankings": null +} diff --git a/sagemaker-serve/tests/unit/test_jumpstart_default_instance_type.py b/sagemaker-serve/tests/unit/test_jumpstart_default_instance_type.py new file mode 100644 index 0000000000..e73fe84795 --- /dev/null +++ b/sagemaker-serve/tests/unit/test_jumpstart_default_instance_type.py @@ -0,0 +1,361 @@ +"""Unit tests for JumpStart default instance type resolution in ModelBuilder. + +Regression tests for the bug where ``__post_init__`` ran +``_initialize_compute_config()`` before ``self.region`` was assigned, +causing the JumpStart spec-default lookup to fail with a silently +swallowed AttributeError and every JumpStart model to default to +``ml.m5.large`` instead of the spec's default instance type. +""" + +import json +import pathlib +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from sagemaker.serve.model_builder import ModelBuilder + +JS_MODEL_ID = "openai-reasoning-gpt-oss-120b" +SPEC_DEFAULT_INSTANCE_TYPE = "ml.g7e.2xlarge" +TEST_ROLE = "arn:aws:iam::123456789012:role/TestRole" + +# Production-shaped JumpStart model spec, extracted from the master-v2 branch's +# tests/unit/sagemaker/jumpstart/constants.py BASE_SPEC fixture. +_SPEC_FIXTURE_PATH = pathlib.Path(__file__).parent / "data" / "jumpstart_base_spec.json" +# BASE_SPEC's real default is ml.m5.large, which collides with the buggy code's +# hardcoded fallback and cannot discriminate pass from fail. Override the +# discriminating field (and keep it in the supported list) so a resolved spec +# default is distinguishable from the fallback. +_FIXTURE_DEFAULT_INSTANCE_TYPE = "ml.g5.12xlarge" + + +def _load_fixture_spec_dict(): + spec = json.loads(_SPEC_FIXTURE_PATH.read_text()) + spec["default_inference_instance_type"] = _FIXTURE_DEFAULT_INSTANCE_TYPE + # Exclude ml.m5.large (the buggy code's hardcoded fallback) from the + # supported list so membership assertions discriminate red from green. + spec["supported_inference_instance_types"] = [_FIXTURE_DEFAULT_INSTANCE_TYPE] + [ + t + for t in spec.get("supported_inference_instance_types", []) + if t != "ml.m5.large" + ] + return spec + + +def _get_fixture_spec(*args, **kwargs): + """side_effect for JumpStartModelsAccessor.get_model_specs, mirroring the + master-v2 get_spec_from_base_spec test convention.""" + from sagemaker.core.jumpstart.types import JumpStartModelSpecs + + return JumpStartModelSpecs(_load_fixture_spec_dict()) + + + +def _mock_session(): + """Mock sagemaker session following the package's shared-setup convention.""" + mock_session = Mock() + mock_session.boto_region_name = "us-west-2" + mock_session.config = {} + mock_session.sagemaker_config = {} + mock_session.boto_session = Mock() + mock_session.boto_session.region_name = "us-west-2" + return mock_session + + +def _build_jumpstart_builder(get_deploy_kwargs_mock=None, **builder_kwargs): + """Construct a ModelBuilder for a JumpStart model with the spec lookup mocked. + + Returns (builder, get_deploy_kwargs_mock) so tests can assert on the lookup call. + """ + if get_deploy_kwargs_mock is None: + get_deploy_kwargs_mock = Mock( + return_value=SimpleNamespace(instance_type=SPEC_DEFAULT_INSTANCE_TYPE) + ) + + with patch( + "sagemaker.serve.model_builder_utils._ModelBuilderUtils._is_jumpstart_model_id", + return_value=True, + ), patch( + "sagemaker.serve.model_builder_utils.get_deploy_kwargs", + get_deploy_kwargs_mock, + ), patch.object( + ModelBuilder, "_initialize_jumpstart_config" + ): + builder = ModelBuilder( + model=JS_MODEL_ID, + role_arn=TEST_ROLE, + sagemaker_session=_mock_session(), + **builder_kwargs, + ) + return builder, get_deploy_kwargs_mock + + +class TestJumpStartDefaultInstanceTypeResolution(unittest.TestCase): + """Construction must resolve the JumpStart spec default instance type.""" + + def test_construction_resolves_spec_default_instance_type(self): + """No instance_type passed: builder must adopt the spec default, not a hardcoded guess.""" + builder, _ = _build_jumpstart_builder() + + self.assertEqual(builder.instance_type, SPEC_DEFAULT_INSTANCE_TYPE) + self.assertFalse(builder._user_provided_instance_type) + + def test_spec_lookup_receives_resolved_region(self): + """Region must be resolved before the spec lookup runs (init-order regression guard).""" + builder, get_deploy_kwargs_mock = _build_jumpstart_builder() + + get_deploy_kwargs_mock.assert_called_once() + _, called_kwargs = get_deploy_kwargs_mock.call_args + self.assertEqual(called_kwargs.get("region"), "us-west-2") + self.assertEqual(called_kwargs.get("model_id"), JS_MODEL_ID) + + def test_region_attribute_exists_before_compute_config(self): + """self.region must exist by the time instance-type detection runs.""" + seen = {} + original = ModelBuilder._get_default_instance_type + + def spy(self): + seen["has_region"] = hasattr(self, "region") + seen["region"] = getattr(self, "region", None) + return original(self) + + with patch.object(ModelBuilder, "_get_default_instance_type", spy): + _build_jumpstart_builder() + + self.assertTrue(seen.get("has_region"), "self.region missing during instance-type detection") + self.assertEqual(seen.get("region"), "us-west-2") + + def test_user_provided_instance_type_wins_and_skips_lookup(self): + """Explicit instance_type must be honored without consulting the spec.""" + builder, get_deploy_kwargs_mock = _build_jumpstart_builder( + instance_type="ml.p5.48xlarge" + ) + + self.assertEqual(builder.instance_type, "ml.p5.48xlarge") + self.assertTrue(builder._user_provided_instance_type) + get_deploy_kwargs_mock.assert_not_called() + + def test_spec_lookup_failure_logs_warning_and_falls_back(self): + """A genuine lookup failure must be logged, not silently swallowed.""" + failing = Mock(side_effect=RuntimeError("spec fetch exploded")) + + with self.assertLogs("sagemaker.core.utils.utils", level="WARNING") as logs: + builder, _ = _build_jumpstart_builder(get_deploy_kwargs_mock=failing) + + self.assertEqual(builder.instance_type, "ml.m5.large") # documented fallback + self.assertTrue( + any("spec fetch exploded" in line for line in logs.output), + f"expected swallowed exception in warning log, got: {logs.output}", + ) + + +class TestDefaultInstanceTypePropagatesToDeploy(unittest.TestCase): + """The resolved spec default must flow into the deploy chain unchanged.""" + + def test_deploy_passes_spec_default_to_internal_deploy(self): + """deploy() with no instance_type must hand the spec default to _deploy.""" + builder, _ = _build_jumpstart_builder() + builder.built_model = Mock() + + with patch.object(ModelBuilder, "_is_model_customization", return_value=False), patch.object( + ModelBuilder, "_deploy", return_value=Mock() + ) as mock_deploy: + builder.deploy(endpoint_name="test-endpoint") + + mock_deploy.assert_called_once() + _, called_kwargs = mock_deploy.call_args + self.assertEqual(called_kwargs.get("instance_type"), SPEC_DEFAULT_INSTANCE_TYPE) + + def test_deploy_core_endpoint_passes_instance_type_to_production_variant(self): + """_deploy_core_endpoint must place the instance type into the ProductionVariant + used for CreateEndpointConfig/CreateEndpoint.""" + builder, _ = _build_jumpstart_builder() + builder.built_model = Mock() + builder.built_model.model_name = "test-model" + builder.model_name = "test-model" + builder.sagemaker_session.endpoint_in_service_or_not = Mock(return_value=False) + + with patch( + "sagemaker.serve.model_builder.session_helper.production_variant", + return_value={"VariantName": "AllTraffic"}, + ) as mock_pv: + try: + builder._deploy_core_endpoint( + instance_type=builder.instance_type, + initial_instance_count=1, + endpoint_name="test-endpoint", + wait=False, + ) + except Exception: + # Downstream endpoint creation uses mocks; we only assert the + # ProductionVariant handoff below. + pass + + mock_pv.assert_called_once() + called_args, called_kwargs = mock_pv.call_args + passed_instance_type = ( + called_kwargs.get("instance_type") + if "instance_type" in called_kwargs + else called_args[1] + ) + self.assertEqual(passed_instance_type, SPEC_DEFAULT_INSTANCE_TYPE) + + +class TestSpecFixtureDrivenResolution(unittest.TestCase): + """End-to-end resolution against a production-shaped spec fixture. + + Unlike the classes above, these tests do NOT mock get_deploy_kwargs. + They patch only the spec source (JumpStartModelsAccessor.get_model_specs, + the master-v2 test convention) with a realistic spec document, so the + real chain runs: _get_default_instance_type -> + get_deploy_kwargs -> _add_instance_type_to_kwargs -> + instance_types.retrieve_default -> spec.default_inference_instance_type. + """ + + def _build(self, **builder_kwargs): + with patch( + "sagemaker.core.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs", + side_effect=_get_fixture_spec, + ), patch( + "sagemaker.serve.model_builder_utils._ModelBuilderUtils._is_jumpstart_model_id", + return_value=True, + ), patch.object( + ModelBuilder, "_initialize_jumpstart_config" + ): + return ModelBuilder( + model="pytorch-ic-mobilenet-v2", # model_id inside the fixture + role_arn=TEST_ROLE, + sagemaker_session=_mock_session(), + **builder_kwargs, + ) + + def test_spec_document_default_resolves_through_real_machinery(self): + """The fixture's default_inference_instance_type must be resolved + by the real retrieve_default chain, not any hardcoded value.""" + builder = self._build() + + self.assertEqual(builder.instance_type, _FIXTURE_DEFAULT_INSTANCE_TYPE) + + def test_resolved_default_is_in_spec_supported_list(self): + """The resolved type must come from the spec's supported instance types.""" + builder = self._build() + + spec = _load_fixture_spec_dict() + self.assertIn(builder.instance_type, spec["supported_inference_instance_types"]) + + def test_explicit_instance_type_still_wins_over_spec_document(self): + builder = self._build(instance_type="ml.p4d.24xlarge") + + self.assertEqual(builder.instance_type, "ml.p4d.24xlarge") + + +class TestUserOverrideFlowsDownstream(unittest.TestCase): + """A user-selected instance type must not only win at construction time, + it must be the value that drives all downstream behavior: the internal + deploy call, the ProductionVariant sent to CreateEndpointConfig, and the + JumpStart init-kwargs resolution that selects instance-type-specific + image URIs and environment variables. + """ + + USER_INSTANCE_TYPE = "ml.p5.48xlarge" + + def test_deploy_passes_user_override_to_internal_deploy(self): + """deploy() must hand the user's instance type (not the spec default) + to _deploy.""" + builder, _ = _build_jumpstart_builder(instance_type=self.USER_INSTANCE_TYPE) + builder.built_model = Mock() + + with patch.object(ModelBuilder, "_is_model_customization", return_value=False), patch.object( + ModelBuilder, "_deploy", return_value=Mock() + ) as mock_deploy: + builder.deploy(endpoint_name="test-endpoint") + + mock_deploy.assert_called_once() + _, called_kwargs = mock_deploy.call_args + self.assertEqual(called_kwargs.get("instance_type"), self.USER_INSTANCE_TYPE) + self.assertNotEqual(called_kwargs.get("instance_type"), SPEC_DEFAULT_INSTANCE_TYPE) + + def test_deploy_core_endpoint_passes_user_override_to_production_variant(self): + """The ProductionVariant for CreateEndpointConfig must carry the + user's instance type.""" + builder, _ = _build_jumpstart_builder(instance_type=self.USER_INSTANCE_TYPE) + builder.built_model = Mock() + builder.built_model.model_name = "test-model" + builder.model_name = "test-model" + builder.sagemaker_session.endpoint_in_service_or_not = Mock(return_value=False) + + with patch( + "sagemaker.serve.model_builder.session_helper.production_variant", + return_value={"VariantName": "AllTraffic"}, + ) as mock_pv: + try: + builder._deploy_core_endpoint( + instance_type=builder.instance_type, + initial_instance_count=1, + endpoint_name="test-endpoint", + wait=False, + ) + except Exception: + # Downstream endpoint creation uses mocks; we only assert the + # ProductionVariant handoff below. + pass + + mock_pv.assert_called_once() + called_args, called_kwargs = mock_pv.call_args + passed_instance_type = ( + called_kwargs.get("instance_type") + if "instance_type" in called_kwargs + else called_args[1] + ) + self.assertEqual(passed_instance_type, self.USER_INSTANCE_TYPE) + + def test_user_override_forwarded_to_jumpstart_init_kwargs(self): + """_build_for_jumpstart must resolve image URI and env vars against the + user's instance type by forwarding it to get_init_kwargs.""" + builder, _ = _build_jumpstart_builder(instance_type=self.USER_INSTANCE_TYPE) + builder._optimizing = False + builder.model_version = "*" + + mock_init_kwargs = Mock() + mock_init_kwargs.image_uri = "123456789012.dkr.ecr.us-west-2.amazonaws.com/img:tag" + mock_init_kwargs.env = {} + mock_init_kwargs.model_data = "s3://bucket/model.tar.gz" + mock_init_kwargs.enable_network_isolation = None + mock_init_kwargs.model_reference_arn = None + + with patch( + "sagemaker.core.jumpstart.utils.validate_model_id_and_get_type", + return_value=None, + ), patch( + "sagemaker.core.jumpstart.factory.utils.get_init_kwargs", + return_value=mock_init_kwargs, + ) as mock_get_kwargs, patch( + "sagemaker.serve.model_builder.ModelBuilder._create_model", + return_value=Mock(), + ), patch( + "sagemaker.serve.model_builder.ModelBuilder._prepare_for_mode" + ): + builder._build_for_jumpstart() + + mock_get_kwargs.assert_called_once() + self.assertEqual( + mock_get_kwargs.call_args.kwargs.get("instance_type"), + self.USER_INSTANCE_TYPE, + ) + + def test_compute_object_instance_type_wins_and_skips_lookup(self): + """instance_type provided via the Compute object must behave the same + as a direct instance_type: honored verbatim, spec lookup skipped.""" + from sagemaker.core.training.configs import Compute + + builder, get_deploy_kwargs_mock = _build_jumpstart_builder( + compute=Compute(instance_type=self.USER_INSTANCE_TYPE) + ) + + self.assertEqual(builder.instance_type, self.USER_INSTANCE_TYPE) + self.assertTrue(builder._user_provided_instance_type) + get_deploy_kwargs_mock.assert_not_called() + + +if __name__ == "__main__": + unittest.main()