From ed67e3d881ce556780e78551c540a067a410e157 Mon Sep 17 00:00:00 2001 From: Loki Ravi Date: Fri, 17 Jul 2026 03:07:59 +0000 Subject: [PATCH] feat(serve): support fine-tuned models in deployment-config API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make ModelBuilder.list_deployment_configs / get_deployment_config / set_deployment_config work for fine-tuned (model-customization) models in addition to base/JumpStart models, dispatching internally on _is_model_customization(). For fine-tuned models the recipe's published HostingConfigs are the source of truth; selection is by instance type (recipe configs are largely unnamed). The base "deployment config" vs recipe "hosting config" distinction is internal and does not surface: fine-tuned configs are normalized to a shape compatible with the base response (DeploymentConfigName plus a nested DeploymentArgs block, plus BenchmarkMetrics/AccelerationConfigs and an additive IsDefault flag). The normalized fine-tuned shape always populates the DeploymentArgs keys (None when unset); the base response may omit unset keys (its serializer drops empty slots), so the fine-tuned shape is a superset — callers consuming both pathways should use .get() for the optional keys (documented on the methods). set_deployment_config applies the whole matching config (image, environment, compute requirements) at build time. A caller-provided image_uri still takes precedence over the config's ImageUri (documented). Both branches fail fast on bad input: the fine-tuned branch requires instance_type and raises on an unpublished or ambiguous instance; the base branch requires config_name AND instance_type and now also rejects an unpublished config name or an instance the config does not support (validated against JumpStart metadata) instead of silently recording a no-op selection. Instance-type matching honors the config's full offered set, not only its default. A base config is a multi-instance bundle: list_deployment_ configs(instance_type=X) filters against each config's supported-instance metadata (supported_inference_instance_types) and materializes matched configs FOR X, so a config that supports X but defaults to another instance is neither discarded nor materialized at the wrong instance. Recipe hosting configs are per-instance bundles by contract, but SupportedInstanceTypes, when present, is honored for selection and filtering. list/get/set agree for the same selection: get_deployment_ config() materializes the pinned instance (matching what list returns and what build deploys), including when the pinned instance came from SupportedInstanceTypes (differs from the config's default). The pinned instance is preserved through build. Configs published with only DefaultInstanceType are matched consistently end to end. The build env merge tolerates a config that publishes an explicit null Environment. Config discovery (recipe-level with a top-level HostingConfigs fallback) is centralized in _extract_hosting_configs_from_hub() and used by BOTH the selection API and the build path, so a top-level config that is listable/selectable is also applied at build. Nova models are routed to their dedicated build path first (per-tier SMI validation + Nova env precedence), so the top-level fallback never diverts them. An explicit selection stores a deep copy of the raw config, is applied exactly at build (or raises if no longer published), and returned configs are copied so caller mutation cannot corrupt internal state. This is the SDK half of a paired change; the SageMaker agent model-deployment skill consumes this unified API. Internal CR: https://code.amazon.com/reviews/CR-289903643 --- .../src/sagemaker/serve/model_builder.py | 550 ++++++- .../test_recipe_hosting_config_selection.py | 1363 +++++++++++++++++ 2 files changed, 1865 insertions(+), 48 deletions(-) create mode 100644 sagemaker-serve/tests/unit/test_recipe_hosting_config_selection.py diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 275b09707e..f3a3b6eeeb 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -393,6 +393,7 @@ class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuild _tags: Optional[Tags] = field(default=None, init=False) _optimizing: bool = field(default=False, init=False) _deployment_config: Optional[Dict[str, Any]] = field(default=None, init=False) + _selected_hosting_config: Optional[Dict[str, Any]] = field(default=None, init=False) shared_libs: List[str] = field( default_factory=list, @@ -964,6 +965,273 @@ def _is_gpu_instance(self, instance_type: str) -> bool: gpu_count = self._infer_accelerator_count_from_instance_type(instance_type) return gpu_count is not None and gpu_count > 0 + @staticmethod + def _extract_hosting_configs_from_hub( + hub_document: Dict[str, Any], recipe_name: str + ) -> List[Dict[str, Any]]: + """Discover the raw ``HostingConfigs`` for a recipe from an already-fetched hub document. + + The authoritative hosting configurations live under + ``RecipeCollection[].HostingConfigs``. Falls back to the top-level + ``HostingConfigs`` when the matching recipe entry does not carry its own (some models + publish there instead). Returns ``[]`` when neither is present — callers decide whether an + empty result is an error (fine-tuned) or a signal to try another path (e.g. Nova). + + This is the single source of truth for config discovery so that + :meth:`_resolve_recipe_hosting_configs` (used by list/get/set_deployment_config) and + :meth:`_fetch_and_cache_recipe_config` (used by build) never diverge — in particular the + top-level fallback is honored on BOTH the selection and the build paths. + """ + for recipe in hub_document.get("RecipeCollection", []): + if recipe.get("Name") == recipe_name: + configs = recipe.get("HostingConfigs") + if configs: + return configs + break + return hub_document.get("HostingConfigs") or [] + + def _resolve_recipe_hosting_configs(self) -> List[Dict[str, Any]]: + """Return the raw ``HostingConfigs`` list published for this model's recipe. + + A fine-tuned (model-customization) model serves the base model under its recipe, so the + authoritative hosting configurations live in the JumpStart hub document (recipe-level, with + a top-level fallback — see :meth:`_extract_hosting_configs_from_hub`). Raises when neither + is present. + """ + hub_document = self._fetch_hub_document_for_custom_model() + model_package = self._fetch_model_package() + containers = getattr( + getattr(model_package, "inference_specification", None), "containers", None + ) + if not containers: + raise ValueError( + "Model is not supported for deployment: the model package has no inference " + "container to resolve a recipe from." + ) + recipe_name = getattr(containers[0].base_model, "recipe_name", None) or "" + + configs = self._extract_hosting_configs_from_hub(hub_document, recipe_name) + if configs: + return configs + + detail = f"recipe '{recipe_name}'" if recipe_name else "this model's recipe" + raise ValueError( + f"Model is not supported for deployment: {detail} does not publish any " + "hosting configurations." + ) + + @staticmethod + def _raw_config_instance_type(cfg: Dict[str, Any]) -> Optional[str]: + """Instance type of a raw recipe ``HostingConfigs`` entry. + + Recipe entries publish the instance type under ``InstanceType`` or, for some, only + ``DefaultInstanceType``. Every place that reads a raw entry's instance type MUST use this + so normalization, selection, and build-time matching stay consistent (otherwise a + ``DefaultInstanceType``-only config is listable/selectable but silently fails to re-match + at build and falls back to Default). + + This returns the config's PRIMARY/default instance (the one used as the normalized + identifier and materialization target). A config may additionally advertise more instances + via ``SupportedInstanceTypes`` — use :meth:`_raw_config_offered_instances` for matching. + """ + return cfg.get("InstanceType") or cfg.get("DefaultInstanceType") + + @staticmethod + def _raw_config_offered_instances(cfg: Dict[str, Any]) -> set: + """The full set of instance types a raw recipe ``HostingConfigs`` entry offers. + + Recipe hosting configs are, by contract, per-instance bundles: each entry publishes a + single instance type via ``InstanceType``/``DefaultInstanceType`` (this is the observed + recipe shape and what :meth:`_normalize_hosting_config` surfaces as one config == one + instance). A config MAY, however, also advertise additional instances under + ``SupportedInstanceTypes``; when present, those are honored for MATCHING (filter and + select) so a config is findable by any instance it offers — not only its default. The + normalized output still carries the primary/default instance as its identifier. + """ + offered = set(cfg.get("SupportedInstanceTypes") or []) + primary = ModelBuilder._raw_config_instance_type(cfg) + if primary: + offered.add(primary) + return offered + + def _base_config_supported_instances(self, config_name: str) -> set: + """Instance types a base/JumpStart deployment config supports, from JumpStart metadata. + + A base config is a multi-instance bundle: its JumpStart metadata lists + ``supported_inference_instance_types`` plus a ``default_inference_instance_type``. Unlike a + recipe's per-instance bundle, ONE base config can be deployed on several instances, so + matching by the single materialized instance would wrongly discard a config that supports + the requested instance but defaults to another. Return the full supported set (union of the + supported list and the default) so filtering is against what the config actually offers. + """ + meta = (self._metadata_configs or {}).get(config_name) + resolved = getattr(meta, "resolved_config", None) or {} + supported = set(resolved.get("supported_inference_instance_types") or []) + default = resolved.get("default_inference_instance_type") + if default: + supported.add(default) + return supported + + @staticmethod + def _normalize_hosting_config(cfg: Dict[str, Any]) -> Dict[str, Any]: + """Normalize a raw recipe ``HostingConfigs`` entry into the SAME shape the base/JumpStart + ``list_deployment_configs`` response uses, so a caller can iterate results from either + pathway identically. + + The serving fields live under ``DeploymentArgs`` with the SAME keys the base response + nests (``ImageUri``, ``InstanceType``, ``Environment``, ``ComputeResourceRequirements``, + ``ModelData``, ``ModelPackageArn``, ``ModelDataDownloadTimeout``, + ``ContainerStartupHealthCheckTimeout``, ``AdditionalDataSources``). This normalized + (fine-tuned) shape always populates all of these keys, using ``None`` for values a recipe + does not provide. NOTE: the base/JumpStart response is built from a slotted dataclass whose + ``to_json()`` OMITS unset slots, so a base config may be MISSING some of these keys + entirely. The normalized shape is therefore a superset of the base shape, not byte-identical + to it — a caller iterating results from both pathways should use ``.get()`` for the optional + ``DeploymentArgs`` keys rather than direct indexing. ``BenchmarkMetrics``/ + ``AccelerationConfigs`` are ``None`` (recipes don't publish these today; extensible — + populate when recipes do). ``IsDefault`` is an additive convenience the recipe uniquely + provides (it marks one config with ``Profile == "Default"``); base responses carry no such + flag, so callers of the base pathway must not assume it is present. + """ + profile = cfg.get("Profile") + instance_type = ModelBuilder._raw_config_instance_type(cfg) + # Nested dicts are copied so a caller mutating a returned config cannot corrupt the raw + # recipe config (or the stored selection, which get_deployment_config normalizes). + return { + # Configs are mostly unnamed (only "Default" has a Profile); fall back to the instance + # type. This occupies the base API's config-name slot for readability, but it is NOT a + # guaranteed-unique identifier — selection is by instance type (multiple configs may + # share an instance type; the setter rejects such ambiguity). Do not treat this as a + # stable ID. + "DeploymentConfigName": profile or instance_type, + "DeploymentArgs": { + "ImageUri": cfg.get("EcrAddress"), + "InstanceType": instance_type, + "Environment": dict(cfg.get("Environment") or {}), + "ComputeResourceRequirements": dict(cfg.get("ComputeResourceRequirements") or {}), + # Keys the base DeploymentArgs carries that recipe configs don't populate — present + # as None so this fine-tuned shape is a compatible SUPERSET of the base + # DeploymentArgs (the base to_json() omits unset slots, so base configs may lack + # these keys; callers should .get() them, not index). + "ModelData": None, + "ModelPackageArn": None, + "ModelDataDownloadTimeout": None, + "ContainerStartupHealthCheckTimeout": None, + "AdditionalDataSources": None, + }, + "BenchmarkMetrics": None, + "AccelerationConfigs": None, + "IsDefault": profile == "Default", + } + + @staticmethod + def _materialize_normalized_for_instance( + norm: Dict[str, Any], instance_type: str + ) -> Dict[str, Any]: + """Surface ``instance_type`` as the instance of an already-normalized config. + + Used when a config is returned in response to an instance-type filter: the returned data + should be materialized for the requested instance. For the common per-instance bundle this + is a no-op (the config's instance already equals the request). For a config that offered + the instance only via ``SupportedInstanceTypes`` (default differs), this rewrites the + nested ``InstanceType`` and, for an unnamed config (whose name defaulted to its instance + type), the ``DeploymentConfigName`` too — so the caller sees the config it asked for. + """ + if norm["DeploymentArgs"].get("InstanceType") == instance_type: + return norm + materialized = copy.deepcopy(norm) + old_instance = materialized["DeploymentArgs"].get("InstanceType") + materialized["DeploymentArgs"]["InstanceType"] = instance_type + # The unnamed-config identifier is its instance type; keep it in sync. A real profile name + # (e.g. "Default") is left untouched. + if not materialized.get("IsDefault") and materialized["DeploymentConfigName"] == old_instance: + materialized["DeploymentConfigName"] = instance_type + return materialized + + def _select_recipe_hosting_config(self, hosting_configs: list) -> dict: + """Select which recipe hosting configuration to apply. + + A recipe publishes several pre-benchmarked hosting configurations (its + ``HostingConfigs``); each is a self-consistent bundle of an instance type, + serving container image (``EcrAddress``), container ``Environment`` (e.g. + ``OPTION_TENSOR_PARALLEL_DEGREE``), and ``ComputeResourceRequirements``. + + Selection precedence — note the two input mechanisms have deliberately different + no-match semantics: + + 1. EXPLICIT selection via :meth:`set_deployment_config` (``_selected_hosting_config``): + the stored raw config is applied exactly. If its instance type is no longer published + at build time, this RAISES — an explicit selection is never silently downgraded to + Default. + 2. CONSTRUCTOR-supplied ``instance_type`` (``_user_provided_instance_type`` with no explicit + selection): the config whose ``InstanceType`` matches is applied so its image, + environment, and compute requirements travel together. If it matches no published + config, this WARNS and falls back to Default (precedence 3) — preserving prior behavior + for direct callers who pin an instance without selecting a published bundle. Prefer + :meth:`set_deployment_config` for fail-fast behavior. + 3. The ``Default`` config (or the first entry) — when no instance type was supplied, or the + constructor-supplied instance type matched no published config. + + Args: + hosting_configs: The recipe's ``HostingConfigs`` list. + + Returns: + The selected hosting config dict. + """ + if not hosting_configs: + raise ValueError( + "Model is not supported for deployment: no hosting configurations to select from." + ) + if self._selected_hosting_config is not None: + # An explicit set_deployment_config() selection is applied EXACTLY (the stored raw + # config), never a re-derived one — so a hub-document change between selection and + # build cannot silently swap in different image/environment/compute values. If the + # selected instance type is no longer published at build time, fail loudly rather than + # fall back to Default (an explicit selection must be honored or surfaced). + selected_instance = self._raw_config_instance_type(self._selected_hosting_config) + selected_offered = self._raw_config_offered_instances(self._selected_hosting_config) + # Still published if any fresh config shares an offered instance with the stored + # selection (matching the supported-set semantics used for selection). + still_published = any( + self._raw_config_offered_instances(cfg) & selected_offered + for cfg in hosting_configs + ) + if not still_published: + raise ValueError( + "The deployment config previously selected via set_deployment_config() for " + f"instance type '{selected_instance}' is no longer published by this model's " + "recipe. Re-select an available config with " + "set_deployment_config(instance_type=...)." + ) + return self._selected_hosting_config + + if self._user_provided_instance_type and self.instance_type: + match = next( + ( + cfg + for cfg in hosting_configs + if self.instance_type in self._raw_config_offered_instances(cfg) + ), + None, + ) + if match is not None: + logger.info( + "Selected recipe hosting config for instance type %s.", + self.instance_type, + ) + return match + logger.warning( + "Instance type %s does not match any published hosting config for this " + "recipe; falling back to the Default hosting config's container image, " + "environment, and compute requirements. Published instance types: %s", + self.instance_type, + [self._raw_config_instance_type(cfg) for cfg in hosting_configs], + ) + return next( + (cfg for cfg in hosting_configs if cfg.get("Profile") == "Default"), + hosting_configs[0], + ) + def _fetch_and_cache_recipe_config(self): """Fetch and cache image URI, compute requirements, and s3_upload_path from recipe during build.""" hub_document = self._fetch_hub_document_for_custom_model() @@ -977,46 +1245,11 @@ def _fetch_and_cache_recipe_config(self): if s3_uri: self.s3_upload_path = s3_uri - for recipe in hub_document.get("RecipeCollection", []): - if recipe.get("Name") == recipe_name: - hosting_configs = recipe.get("HostingConfigs", []) - if hosting_configs: - config = next( - (cfg for cfg in hosting_configs if cfg.get("Profile") == "Default"), - hosting_configs[0], - ) - if not self.image_uri: - self.image_uri = config.get("EcrAddress") - - # Cache environment variables from recipe config - if self.env_vars: - self.env_vars.update(config.get("Environment", {})) - else: - self.env_vars = config.get("Environment", {}) - - # Infer instance type from JumpStart metadata if not provided by user - if not self._user_provided_instance_type: - recipe_instance_type = config.get("InstanceType") or config.get( - "DefaultInstanceType" - ) - if recipe_instance_type: - self.instance_type = recipe_instance_type - elif not self.instance_type: - self.instance_type = self._infer_instance_type_from_jumpstart() - - # Resolve compute requirements using the already-fetched hub document - # This ensures requirements are determined through public methods - # and properly merged with any user-provided configuration - self._cached_compute_requirements = ( - self._resolve_compute_requirements_from_config( - instance_type=self.instance_type, - config=config, - user_resource_requirements=None, # No user config at build time - ) - ) - return - - # Fallback: Nova recipes don't have hosting configs in the hub document + # Nova models resolve through their dedicated path (per-tier SMI validation + Nova env + # precedence, where user overrides are applied LAST), regardless of where their hosting + # configs live in the hub document. This MUST run before the generic recipe branch below — + # otherwise a Nova model that publishes top-level HostingConfigs would be diverted to the + # generic path and skip _validate_nova_smi_config / invert env precedence. if self._is_nova_model(): nova_config = self._get_nova_hosting_config(instance_type=self.instance_type) if not self.image_uri: @@ -1032,6 +1265,50 @@ def _fetch_and_cache_recipe_config(self): self._validate_nova_smi_config() return + # Use the SAME discovery logic as the selection API (recipe-level + top-level fallback) so + # a config that was listable/selectable via list/set_deployment_config is also applied here + # at build time. Previously this walked only recipe-level HostingConfigs, so a selected + # top-level config was silently dropped at build. + hosting_configs = self._extract_hosting_configs_from_hub(hub_document, recipe_name) + if hosting_configs: + config = self._select_recipe_hosting_config(hosting_configs) + if not self.image_uri: + self.image_uri = config.get("EcrAddress") + + # Cache environment variables from recipe config. Use `or {}` (not a `{}` default) so a + # config that publishes an explicit "Environment": null does not blow up dict()/update. + recipe_env = config.get("Environment") or {} + if self.env_vars: + self.env_vars.update(recipe_env) + else: + self.env_vars = dict(recipe_env) + + if self._selected_hosting_config is not None: + # An explicit selection is a self-consistent bundle. Keep the pinned instance if it + # is one the config offers (the instance chosen at set time — a config may offer + # several via SupportedInstanceTypes); otherwise snap to the config's primary so the + # endpoint matches the applied image/env/compute even if the caller reassigned + # instance_type after selecting. + if self.instance_type not in self._raw_config_offered_instances(config): + self.instance_type = self._raw_config_instance_type(config) + elif not self._user_provided_instance_type: + # Infer instance type from JumpStart metadata if not provided by user + recipe_instance_type = self._raw_config_instance_type(config) + if recipe_instance_type: + self.instance_type = recipe_instance_type + elif not self.instance_type: + self.instance_type = self._infer_instance_type_from_jumpstart() + + # Resolve compute requirements using the already-fetched hub document + # This ensures requirements are determined through public methods + # and properly merged with any user-provided configuration + self._cached_compute_requirements = self._resolve_compute_requirements_from_config( + instance_type=self.instance_type, + config=config, + user_resource_requirements=None, # No user config at build time + ) + return + raise ValueError( f"Model with recipe '{recipe_name}' is not supported for deployment. " f"The recipe does not have hosting configuration. " @@ -1090,6 +1367,13 @@ def _is_nova_model(self) -> bool: return False recipe_name = getattr(base_model, "recipe_name", "") or "" hub_content_name = getattr(base_model, "hub_content_name", "") or "" + # Coerce to str defensively: these attributes are normally strings, but a partially + # populated model package can leave them as None or a non-string, and "nova" in + # raises TypeError. Treat any non-string as absent (not Nova). + if not isinstance(recipe_name, str): + recipe_name = "" + if not isinstance(hub_content_name, str): + hub_content_name = "" return "nova" in recipe_name.lower() or "nova" in hub_content_name.lower() def _is_nova_model_for_telemetry(self) -> bool: @@ -4035,8 +4319,63 @@ def display_benchmark_metrics(self, **kwargs) -> None: @_telemetry_emitter( feature=Feature.MODEL_CUSTOMIZATION, func_name="model_builder.set_deployment_config" ) - def set_deployment_config(self, config_name: str, instance_type: str) -> None: - """Sets the deployment config to apply to the model.""" + def set_deployment_config( + self, config_name: Optional[str] = None, instance_type: Optional[str] = None + ) -> None: + """Select the deployment config to apply to the model. + + One API for both model types. For base/JumpStart models, ``config_name`` chooses a named + deployment config (with ``instance_type``). For fine-tuned models, the configs are largely + unnamed, so selection is by ``instance_type`` — pass the ``InstanceType`` of a config from + :meth:`list_deployment_configs`. The whole matching config (container image, environment, + and compute requirements) is applied at build/deploy time. + + Image override: a caller-provided ``image_uri`` on the builder takes precedence over the + selected config's ``ImageUri`` — the config's environment and compute requirements still + apply, so the result is a hybrid of the caller's image with the recipe config's other + fields. Leave ``image_uri`` unset to apply the published config's image verbatim. + """ + if self._is_model_customization(): + if not instance_type: + raise ValueError( + "Select a deployment config by instance_type for this model " + "(see list_deployment_configs())." + ) + # Resolve the recipe's raw configs ONCE and match on instance type. Store a deep copy + # of the matched RAW config so the exact selection (image/env/compute) is applied at + # build time — independent of later hub-document changes or caller mutation of any + # returned config. get_deployment_config() normalizes this raw config for output. + raw_configs = self._resolve_recipe_hosting_configs() + # Match on the config's full offered set (its instance plus any SupportedInstanceTypes), + # so a config is selectable by any instance it offers, not only its default. + matches = [ + c + for c in raw_configs + if instance_type in self._raw_config_offered_instances(c) + ] + if not matches: + available = sorted( + { + inst + for c in raw_configs + for inst in self._raw_config_offered_instances(c) + } + ) + raise ValueError( + f"No deployment config published for instance type '{instance_type}'. " + f"Available instance types: {available}" + ) + if len(matches) > 1: + raise ValueError( + f"Instance type '{instance_type}' matches {len(matches)} deployment configs; " + "selection by instance type is ambiguous for this model." + ) + self._selected_hosting_config = copy.deepcopy(matches[0]) + self.instance_type = instance_type + self._user_provided_instance_type = True + logger.info("Selected deployment config for instance type %s.", instance_type) + return + if not isinstance(self.model, str): raise ValueError( "Deployment config is only supported for JumpStart or HuggingFace models" @@ -4045,6 +4384,30 @@ def set_deployment_config(self, config_name: str, instance_type: str) -> None: if not (self._is_jumpstart_model_id() or self._use_jumpstart_equivalent()): raise ValueError(f"The deployment config {config_name} cannot be set on this model") + if not config_name: + raise ValueError("config_name is required to set a deployment config on this model.") + + if not instance_type: + raise ValueError("instance_type is required to set a deployment config on this model.") + + # Fail fast on an unpublished config name or an instance the config does not support, + # mirroring the fine-tuned branch — otherwise a bogus selection is recorded silently and + # simply applies nothing at build time. Only enforced when JumpStart metadata is available + # (the authoritative source); if it can't be resolved we defer to the downstream lookup. + self._ensure_metadata_configs() + if self._metadata_configs: + if config_name not in self._metadata_configs: + raise ValueError( + f"No deployment config named '{config_name}' is published for this model. " + f"Available config names: {sorted(self._metadata_configs)}" + ) + supported = self._base_config_supported_instances(config_name) + if supported and instance_type not in supported: + raise ValueError( + f"Deployment config '{config_name}' does not support instance type " + f"'{instance_type}'. Supported instance types: {sorted(supported)}" + ) + self.config_name = config_name self.instance_type = instance_type @@ -4068,7 +4431,42 @@ def set_deployment_config(self, config_name: str, instance_type: str) -> None: feature=Feature.MODEL_CUSTOMIZATION, func_name="model_builder.get_deployment_config" ) def get_deployment_config(self) -> Optional[Dict[str, Any]]: - """Gets the deployment config to apply to the model.""" + """Get the deployment config that will be applied to the model. + + One API for both model types. For fine-tuned models, returns the config chosen via + :meth:`set_deployment_config` if one was set, otherwise the config that would be selected + by default at build time (matching a provided instance type, else the default config); it + raises if the model publishes no hosting configs. NOTE: this differs from the base/ + JumpStart path, which returns ``None`` until a config is explicitly set — the fine-tuned + path always has a well-defined default (its recipe's Default config), so it reports that. + + This reports the RECORDED selection without re-validating it against the live recipe; + build (:meth:`_select_recipe_hosting_config`) is the source of truth and raises if an + explicitly selected config is no longer published. So after a recipe changes, this getter + may still return a selection that build would reject — call it before build, not as a + freshness check. + """ + if self._is_model_customization(): + # The chosen RAW config: an explicit selection, else the config build would pick. + # _resolve_recipe_hosting_configs raises when no configs are published, so a non-empty + # list is guaranteed on the fallback path. + chosen = ( + self._selected_hosting_config + if self._selected_hosting_config is not None + else self._select_recipe_hosting_config(self._resolve_recipe_hosting_configs()) + ) + norm = self._normalize_hosting_config(chosen) + # Report the config materialized for the EFFECTIVE instance (what build will deploy), + # so this agrees with list_deployment_configs(instance_type=...) and set_deployment_ + # config() for the same selection — including when the pinned instance came from the + # config's SupportedInstanceTypes (differs from its default). Only rewrite when the + # pinned instance is one the config offers; otherwise the normalized default stands. + if self.instance_type and self.instance_type in self._raw_config_offered_instances( + chosen + ): + norm = self._materialize_normalized_for_instance(norm, self.instance_type) + return norm + if not isinstance(self.model, str): raise ValueError( "Deployment config is only supported for JumpStart or HuggingFace models" @@ -4092,8 +4490,43 @@ def get_deployment_config(self) -> Optional[Dict[str, Any]]: @_telemetry_emitter( feature=Feature.MODEL_CUSTOMIZATION, func_name="model_builder.list_deployment_configs" ) - def list_deployment_configs(self) -> List[Dict[str, Any]]: - """List deployment configs for the model in the current region.""" + def list_deployment_configs( + self, instance_type: Optional[str] = None + ) -> List[Dict[str, Any]]: + """List the deployment configs available for the model in the current region. + + One API for both model types, returning a compatible dict shape so callers can iterate + results from either pathway: ``DeploymentConfigName`` plus a ``DeploymentArgs`` block + (``ImageUri``, ``InstanceType``, ``Environment``, ``ComputeResourceRequirements``) and + ``BenchmarkMetrics`` / ``AccelerationConfigs``. The base/JumpStart response may OMIT unset + ``DeploymentArgs`` keys (its serializer drops empty slots), while the fine-tuned shape + always populates them (``None`` when unset) — a superset. Use ``.get()`` for the optional + keys when consuming both pathways. For fine-tuned models the recipe's published configs are + normalized into this shape; + ``BenchmarkMetrics``/``AccelerationConfigs`` are empty (recipes don't publish them yet) and + an additive ``IsDefault`` flag marks the recipe's Default config. When ``instance_type`` is + provided, only configs offering that instance type are returned — the supported way to + narrow the alternatives before selecting. A config "offers" an instance if that instance is + its published/default instance OR appears in the config's supported-instance metadata; a + matched config is materialized for the requested instance in the returned data. + """ + if self._is_model_customization(): + configs = [] + for raw in self._resolve_recipe_hosting_configs(): + if instance_type is not None and instance_type not in ( + self._raw_config_offered_instances(raw) + ): + continue + norm = self._normalize_hosting_config(raw) + if instance_type is not None: + # Materialize the returned config FOR the requested instance. Common case: the + # config's single instance already equals it (no-op). Superset case: the config + # offered the instance via SupportedInstanceTypes but defaults to another — + # surface the requested instance so the result is materialized for it. + norm = self._materialize_normalized_for_instance(norm, instance_type) + configs.append(norm) + return configs + if not isinstance(self.model, str): raise ValueError( "Deployment config is only supported for JumpStart or HuggingFace models" @@ -4102,9 +4535,30 @@ def list_deployment_configs(self) -> List[Dict[str, Any]]: if not (self._is_jumpstart_model_id() or self._use_jumpstart_equivalent()): raise ValueError("Deployment config is only supported for JumpStart models") - return self.deployment_config_response_data( - self._get_deployment_configs(self.config_name, self.instance_type) - ) # Delegate to JumpStart builder + if instance_type is None: + # No narrowing: return every config materialized at its default (original base API). + return self.deployment_config_response_data( + self._get_deployment_configs(self.config_name, self.instance_type) + ) + + # A base config is a MULTI-instance bundle. Filtering on the single materialized instance + # would wrongly discard a config that supports the requested instance but defaults to + # another, so filter against each config's supported-instance metadata and materialize the + # matched configs FOR the requested instance (pass it as the selected instance so + # JumpStart resolves image/env/compute for it, not the default). + self._ensure_metadata_configs() + configs = [] + for config_name in self._metadata_configs or {}: + if instance_type not in self._base_config_supported_instances(config_name): + continue + configs.extend( + cfg + for cfg in self.deployment_config_response_data( + self._get_deployment_configs(config_name, instance_type) + ) + if cfg.get("DeploymentConfigName") == config_name + ) + return configs @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="model_builder.optimize") # Add these methods to the current V3 ModelBuilder class: diff --git a/sagemaker-serve/tests/unit/test_recipe_hosting_config_selection.py b/sagemaker-serve/tests/unit/test_recipe_hosting_config_selection.py new file mode 100644 index 0000000000..e725e2d314 --- /dev/null +++ b/sagemaker-serve/tests/unit/test_recipe_hosting_config_selection.py @@ -0,0 +1,1363 @@ +"""Unit tests for ModelBuilder._select_recipe_hosting_config. + +A recipe publishes several pre-benchmarked hosting configurations. When the caller supplies an +instance type that matches one of them, ModelBuilder must select THAT config (its image, +environment, and compute requirements) rather than pinning the caller's instance onto the +Default config's environment. These tests pin that selection behavior. +""" + +import unittest +from unittest.mock import Mock, patch + +from sagemaker.serve.model_builder import ModelBuilder +from sagemaker.serve.mode.function_pointers import Mode + +# Mirrors a real recipe's HostingConfigs (e.g. llmft_qwen3_0_dot_6b_seq4k_gpu_sft_lora): +# a 1-GPU Default plus a same-topology alternative and two 8-GPU / TP=8 alternatives. +HOSTING_CONFIGS = [ + { + "Profile": "Default", + "InstanceType": "ml.g6.4xlarge", + "EcrAddress": "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.34.0-lmi16", + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "1"}, + "ComputeResourceRequirements": { + "MinMemoryRequiredInMb": 32768, + "NumberOfAcceleratorDevicesRequired": 1, + "NumberOfCpuCoresRequired": 12, + }, + }, + { + "InstanceType": "ml.g5.4xlarge", + "EcrAddress": "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.34.0-lmi16", + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "1"}, + }, + { + "InstanceType": "ml.g6e.48xlarge", + "EcrAddress": "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.34.0-lmi16", + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "8"}, + }, + { + "InstanceType": "ml.p5.48xlarge", + "EcrAddress": "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.34.0-lmi16", + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "8"}, + }, +] + + +def _make_mock_session(): + """A minimally-populated SageMaker session mock sufficient to construct a ModelBuilder.""" + mock_session = Mock() + mock_session.boto_region_name = "us-west-2" + mock_session.default_bucket.return_value = "test-bucket" + mock_session.default_bucket_prefix = "test-prefix" + mock_session.config = {} + mock_session.sagemaker_config = {} + mock_session.settings = Mock() + mock_session.settings.include_jumpstart_tags = False + creds = Mock() + creds.access_key = "k" + creds.secret_key = "s" + creds.token = None + mock_session.boto_session = Mock() + mock_session.boto_session.get_credentials.return_value = creds + mock_session.boto_session.region_name = "us-west-2" + return mock_session + + +class TestSelectRecipeHostingConfig(unittest.TestCase): + def setUp(self): + self.mock_session = _make_mock_session() + + def _builder(self, instance_type): + return ModelBuilder( + model="huggingface-reasoning-qwen3-06b", + model_metadata={ + "CUSTOM_MODEL_ID": "huggingface-reasoning-qwen3-06b", + "CUSTOM_MODEL_VERSION": "3.9.0", + }, + mode=Mode.SAGEMAKER_ENDPOINT, + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + instance_type=instance_type, + ) + + def test_matching_alternative_instance_selects_that_config(self): + # A different-topology alternative (8 GPU / TP=8) must be selected in full, not Default. + b = self._builder("ml.g6e.48xlarge") + cfg = b._select_recipe_hosting_config(HOSTING_CONFIGS) + self.assertEqual(cfg["InstanceType"], "ml.g6e.48xlarge") + self.assertEqual(cfg["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"], "8") + + def test_matching_same_topology_alternative_selects_that_config(self): + b = self._builder("ml.g5.4xlarge") + cfg = b._select_recipe_hosting_config(HOSTING_CONFIGS) + self.assertEqual(cfg["InstanceType"], "ml.g5.4xlarge") + + def test_default_instance_selects_default_config(self): + b = self._builder("ml.g6.4xlarge") + cfg = b._select_recipe_hosting_config(HOSTING_CONFIGS) + self.assertEqual(cfg["Profile"], "Default") + self.assertEqual(cfg["InstanceType"], "ml.g6.4xlarge") + + def test_no_user_instance_type_falls_back_to_default(self): + b = self._builder("ml.g6e.48xlarge") + # Simulate "user did not pin an instance type". + b._user_provided_instance_type = False + cfg = b._select_recipe_hosting_config(HOSTING_CONFIGS) + self.assertEqual(cfg["Profile"], "Default") + + def test_unmatched_instance_type_falls_back_to_default(self): + # A constructor-supplied instance not published for the recipe keeps prior behavior: WARN + # and fall back to Default (contrast with an explicit selection, which raises). + b = self._builder("ml.g6.xlarge") + with self.assertLogs(level="WARNING") as logs: + cfg = b._select_recipe_hosting_config(HOSTING_CONFIGS) + self.assertEqual(cfg["Profile"], "Default") + self.assertTrue( + any("does not match any published hosting config" in m for m in logs.output) + ) + + def test_no_default_profile_falls_back_to_first_entry(self): + b = self._builder("ml.g6.xlarge") # unmatched -> fallback path + configs = [c for c in HOSTING_CONFIGS if c.get("Profile") != "Default"] + cfg = b._select_recipe_hosting_config(configs) + self.assertEqual(cfg, configs[0]) + + def test_stale_explicit_selection_raises_not_silent_fallback(self): + # An EXPLICIT selection (_selected_hosting_config) whose instance type is no longer present + # at build time must RAISE — never silently fall back to Default. (A constructor-provided + # instance type, by contrast, warns and falls back; see the unmatched test above.) + b = self._builder("ml.g6.4xlarge") + b._selected_hosting_config = {"InstanceType": "ml.not-published.xlarge", "Environment": {}} + with self.assertRaises(ValueError) as ctx: + b._select_recipe_hosting_config(HOSTING_CONFIGS) + msg = str(ctx.exception) + self.assertIn("no longer published", msg) + self.assertIn("ml.not-published.xlarge", msg) + + def test_stale_selection_still_published_via_supported_instance(self): + # The stored selection offers [g5.2xlarge, g5.12xlarge]; the fresh list has a config that + # offers g5.12xlarge (as its primary). They share an offered instance -> NOT stale. + b = self._builder("ml.g5.12xlarge") + b._selected_hosting_config = { + "DefaultInstanceType": "ml.g5.2xlarge", + "SupportedInstanceTypes": ["ml.g5.2xlarge", "ml.g5.12xlarge"], + "Environment": {}, + } + fresh = [{"InstanceType": "ml.g5.12xlarge", "Environment": {}}] + # Must not raise; returns the stored selection exactly. + self.assertIs(b._select_recipe_hosting_config(fresh), b._selected_hosting_config) + + def test_stale_selection_raises_when_no_offered_instance_shared(self): + # Stored selection offers only g5.2xlarge/g5.12xlarge; fresh list shares none -> stale. + b = self._builder("ml.g5.12xlarge") + b._selected_hosting_config = { + "DefaultInstanceType": "ml.g5.2xlarge", + "SupportedInstanceTypes": ["ml.g5.2xlarge", "ml.g5.12xlarge"], + "Environment": {}, + } + fresh = [{"InstanceType": "ml.p5.48xlarge", "Environment": {}}] + with self.assertRaises(ValueError) as ctx: + b._select_recipe_hosting_config(fresh) + self.assertIn("no longer published", str(ctx.exception)) + + def test_explicit_selection_applied_exactly(self): + # An explicit selection is applied EXACTLY (the stored raw config), not a re-derived one — + # even if the freshly-resolved list has a same-instance entry with different contents. + b = self._builder("ml.g6e.48xlarge") + b._selected_hosting_config = { + "InstanceType": "ml.g6e.48xlarge", + "EcrAddress": "stored-image", + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "8"}, + } + # Fresh list has the same instance type but different (drifted) contents. + fresh = [ + {"Profile": "Default", "InstanceType": "ml.g6.4xlarge", "Environment": {}}, + {"InstanceType": "ml.g6e.48xlarge", "EcrAddress": "drifted-image", "Environment": {}}, + ] + cfg = b._select_recipe_hosting_config(fresh) + self.assertEqual(cfg["EcrAddress"], "stored-image") + self.assertEqual(cfg["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"], "8") + + +class TestDeploymentConfigForFineTunedModels(unittest.TestCase): + """The unified deployment-config API (list/get/set) applied to fine-tuned models. + + Same public methods as base/JumpStart models — the base "deployment config" vs recipe + "hosting config" split is internal and never surfaces to the caller. + """ + + def setUp(self): + self.mock_session = _make_mock_session() + + def _builder(self, instance_type="ml.g6.4xlarge"): + return ModelBuilder( + model="huggingface-reasoning-qwen3-06b", + model_metadata={ + "CUSTOM_MODEL_ID": "huggingface-reasoning-qwen3-06b", + "CUSTOM_MODEL_VERSION": "3.9.0", + }, + mode=Mode.SAGEMAKER_ENDPOINT, + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + instance_type=instance_type, + ) + + def _customization_builder(self, instance_type="ml.g6.4xlarge"): + """A builder patched to look like a fine-tuned model with the fixture configs.""" + b = self._builder(instance_type) + patcher_cust = patch.object(ModelBuilder, "_is_model_customization", return_value=True) + patcher_cfgs = patch.object( + ModelBuilder, "_resolve_recipe_hosting_configs", return_value=HOSTING_CONFIGS + ) + patcher_cust.start() + patcher_cfgs.start() + self.addCleanup(patcher_cust.stop) + self.addCleanup(patcher_cfgs.stop) + return b + + def test_list_deployment_configs_returns_base_shape(self): + # Fine-tuned configs are normalized into the SAME shape as the base/JumpStart response: + # DeploymentConfigName + a nested DeploymentArgs block + BenchmarkMetrics, plus the + # additive IsDefault flag. A unified consumer reads DeploymentArgs["InstanceType"] etc. + # identically for base and fine-tuned. + b = self._customization_builder() + configs = b.list_deployment_configs() + self.assertEqual(len(configs), 4) + default = next(c for c in configs if c["IsDefault"]) + self.assertEqual(default["DeploymentConfigName"], "Default") + self.assertEqual(default["DeploymentArgs"]["InstanceType"], "ml.g6.4xlarge") + self.assertEqual(default["DeploymentArgs"]["ImageUri"].split(":")[-1], "0.34.0-lmi16") + self.assertEqual( + default["DeploymentArgs"]["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"], "1" + ) + # BenchmarkMetrics present but None (recipes publish none today; base emits None when + # absent, so None matches the base sentinel). + self.assertIsNone(default["BenchmarkMetrics"]) + # Unnamed configs get their instance type as a stable identifier. + g6e = next(c for c in configs if c["DeploymentArgs"]["InstanceType"] == "ml.g6e.48xlarge") + self.assertEqual(g6e["DeploymentConfigName"], "ml.g6e.48xlarge") + self.assertFalse(g6e["IsDefault"]) + self.assertEqual(g6e["DeploymentArgs"]["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"], "8") + + def test_list_deployment_configs_filters_by_instance_type(self): + b = self._customization_builder() + configs = b.list_deployment_configs(instance_type="ml.g6e.48xlarge") + self.assertEqual(len(configs), 1) + self.assertEqual(configs[0]["DeploymentArgs"]["InstanceType"], "ml.g6e.48xlarge") + # A non-published instance filters to empty. + self.assertEqual(b.list_deployment_configs(instance_type="ml.g6.xlarge"), []) + + def test_set_deployment_config_selects_and_pins_instance(self): + b = self._customization_builder() + b.set_deployment_config(instance_type="ml.g6e.48xlarge") + self.assertEqual(b.instance_type, "ml.g6e.48xlarge") + self.assertTrue(b._user_provided_instance_type) + # _selected_hosting_config stores the RAW recipe config (not the normalized shape). + self.assertEqual(b._raw_config_instance_type(b._selected_hosting_config), "ml.g6e.48xlarge") + # The selection flows into the internal build-time selector (which reads raw configs). + chosen = b._select_recipe_hosting_config(HOSTING_CONFIGS) + self.assertEqual(chosen["InstanceType"], "ml.g6e.48xlarge") + self.assertEqual(chosen["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"], "8") + + def test_set_deployment_config_requires_instance_type(self): + b = self._customization_builder() + with self.assertRaises(ValueError) as ctx: + b.set_deployment_config() + self.assertIn("instance_type", str(ctx.exception)) + + def test_set_deployment_config_unknown_instance_raises_with_available(self): + b = self._customization_builder() + with self.assertRaises(ValueError) as ctx: + b.set_deployment_config(instance_type="ml.g6.xlarge") + msg = str(ctx.exception) + self.assertIn("ml.g6.4xlarge", msg) # lists available instance types + self.assertIn("ml.p5.48xlarge", msg) + + def test_set_deployment_config_ambiguous_raises(self): + b = self._builder() + dup = [ + {"InstanceType": "ml.g6.4xlarge", "Environment": {}}, + {"InstanceType": "ml.g6.4xlarge", "Environment": {}}, + ] + with patch.object(ModelBuilder, "_is_model_customization", return_value=True), patch.object( + ModelBuilder, "_resolve_recipe_hosting_configs", return_value=dup + ): + with self.assertRaises(ValueError) as ctx: + b.set_deployment_config(instance_type="ml.g6.4xlarge") + self.assertIn("ambiguous", str(ctx.exception)) + + def test_get_deployment_config_returns_selected_then_default(self): + b = self._customization_builder() + # Before any selection -> the Default config. + self.assertEqual(b.get_deployment_config()["DeploymentConfigName"], "Default") + # After selecting an alternative -> that config. + b.set_deployment_config(instance_type="ml.p5.48xlarge") + got = b.get_deployment_config() + self.assertEqual(got["DeploymentArgs"]["InstanceType"], "ml.p5.48xlarge") + self.assertFalse(got["IsDefault"]) + + def test_set_deployment_config_base_requires_config_name(self): + # A base/JumpStart model going through the SAME setter must still require config_name — the + # instance_type-only form is fine-tuned-only. Guards the config_name-None base branch. + b = self._builder() + with patch.object( + ModelBuilder, "_is_model_customization", return_value=False + ), patch.object(ModelBuilder, "_is_jumpstart_model_id", return_value=True): + with self.assertRaises(ValueError) as ctx: + b.set_deployment_config(instance_type="ml.g5.2xlarge") + self.assertIn("config_name is required", str(ctx.exception)) + + def test_get_deployment_config_base_returns_none_until_set(self): + # Base/JumpStart contract (distinct from fine-tuned): get returns None until a config is + # explicitly set. Pins the documented semantic difference from the fine-tuned path. + b = self._builder() + with patch.object( + ModelBuilder, "_is_model_customization", return_value=False + ), patch.object(ModelBuilder, "_is_jumpstart_model_id", return_value=True): + b.config_name = None + self.assertIsNone(b.get_deployment_config()) + + def test_set_deployment_config_base_success_sets_name_and_instance(self): + # Base happy path: config_name + instance_type are recorded (contrast with the fine-tuned + # instance-type-only selection). Metadata confirms the config publishes the instance. + b = self._builder() + b.additional_model_data_sources = None # normally populated by the build path + meta = { + "lmi": Mock( + resolved_config={ + "default_inference_instance_type": "ml.g5.2xlarge", + "supported_inference_instance_types": ["ml.g5.2xlarge"], + } + ) + } + with patch.object( + ModelBuilder, "_is_model_customization", return_value=False + ), patch.object(ModelBuilder, "_is_jumpstart_model_id", return_value=True), patch.object( + ModelBuilder, "_ensure_metadata_configs" + ), patch.object(ModelBuilder, "get_deployment_config", return_value=None): + b._metadata_configs = meta + b.set_deployment_config(config_name="lmi", instance_type="ml.g5.2xlarge") + self.assertEqual(b.config_name, "lmi") + self.assertEqual(b.instance_type, "ml.g5.2xlarge") + + def test_set_deployment_config_base_rejects_unknown_config_name(self): + # Base setter must fail fast on an unpublished config name (mirrors the fine-tuned branch), + # not silently record a bogus selection that applies nothing at build. + b = self._builder() + meta = {"lmi": Mock(resolved_config={"default_inference_instance_type": "ml.g5.2xlarge"})} + with patch.object( + ModelBuilder, "_is_model_customization", return_value=False + ), patch.object(ModelBuilder, "_is_jumpstart_model_id", return_value=True), patch.object( + ModelBuilder, "_ensure_metadata_configs" + ): + b._metadata_configs = meta + with self.assertRaises(ValueError) as ctx: + b.set_deployment_config(config_name="bogus", instance_type="ml.g5.2xlarge") + self.assertIn("bogus", str(ctx.exception)) + + def test_set_deployment_config_base_rejects_unsupported_instance(self): + # Base setter must fail fast when the named config does not support the requested instance. + b = self._builder() + meta = { + "lmi": Mock( + resolved_config={ + "default_inference_instance_type": "ml.g5.2xlarge", + "supported_inference_instance_types": ["ml.g5.2xlarge"], + } + ) + } + with patch.object( + ModelBuilder, "_is_model_customization", return_value=False + ), patch.object(ModelBuilder, "_is_jumpstart_model_id", return_value=True), patch.object( + ModelBuilder, "_ensure_metadata_configs" + ): + b._metadata_configs = meta + with self.assertRaises(ValueError) as ctx: + b.set_deployment_config(config_name="lmi", instance_type="ml.bogus.xlarge") + self.assertIn("does not support", str(ctx.exception)) + + def test_repeated_selection_applies_latest(self): + # Selecting A then B must fully replace A — no stale env/compute from the first selection. + b = self._customization_builder() + b.set_deployment_config(instance_type="ml.g5.4xlarge") # TP=1 alternative + b.set_deployment_config(instance_type="ml.g6e.48xlarge") # TP=8 alternative + self.assertEqual(b._raw_config_instance_type(b._selected_hosting_config), "ml.g6e.48xlarge") + got = b.get_deployment_config() + self.assertEqual(got["DeploymentArgs"]["InstanceType"], "ml.g6e.48xlarge") + self.assertEqual(got["DeploymentArgs"]["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"], "8") + + def test_returned_config_mutation_does_not_corrupt_selection(self): + # Mutating a dict returned by get_deployment_config() must not alter internal state. + b = self._customization_builder() + b.set_deployment_config(instance_type="ml.g6e.48xlarge") + got = b.get_deployment_config() + got["DeploymentArgs"]["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"] = "999" + got["DeploymentConfigName"] = "mutated" + again = b.get_deployment_config() + self.assertEqual( + again["DeploymentArgs"]["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"], "8" + ) + self.assertEqual(again["DeploymentConfigName"], "ml.g6e.48xlarge") + + def test_list_config_mutation_does_not_corrupt_next_list(self): + b = self._customization_builder() + first = b.list_deployment_configs() + first[0]["DeploymentArgs"]["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"] = "999" + second = b.list_deployment_configs() + self.assertNotEqual( + second[0]["DeploymentArgs"]["Environment"].get("OPTION_TENSOR_PARALLEL_DEGREE"), "999" + ) + + # ---- SupportedInstanceTypes: a config offered under a multi-instance list ---- + + def _supported_types_builder(self): + # A recipe config whose DEFAULT is ml.g5.2xlarge but which also OFFERS ml.g5.12xlarge via + # SupportedInstanceTypes. Recipe configs are per-instance bundles by contract, but the API + # honors SupportedInstanceTypes for matching so such a config is still findable/selectable. + configs = [ + { + "Profile": "Default", + "InstanceType": "ml.g6.4xlarge", + "EcrAddress": "img:default", + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "1"}, + }, + { + "DefaultInstanceType": "ml.g5.2xlarge", + "SupportedInstanceTypes": ["ml.g5.2xlarge", "ml.g5.12xlarge"], + "EcrAddress": "img:multi", + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "4"}, + }, + ] + b = self._builder() + p1 = patch.object(ModelBuilder, "_is_model_customization", return_value=True) + p2 = patch.object( + ModelBuilder, "_resolve_recipe_hosting_configs", return_value=configs + ) + p1.start() + p2.start() + self.addCleanup(p1.stop) + self.addCleanup(p2.stop) + return b + + def test_list_filter_finds_config_by_supported_non_default_instance(self): + # Filtering for a supported non-default instance returns the config, materialized FOR it. + b = self._supported_types_builder() + configs = b.list_deployment_configs(instance_type="ml.g5.12xlarge") + self.assertEqual(len(configs), 1) + self.assertEqual(configs[0]["DeploymentArgs"]["InstanceType"], "ml.g5.12xlarge") + # Unnamed config's identifier tracks the materialized instance. + self.assertEqual(configs[0]["DeploymentConfigName"], "ml.g5.12xlarge") + # The config's own env travels with it. + self.assertEqual( + configs[0]["DeploymentArgs"]["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"], "4" + ) + + def test_set_selects_config_by_supported_non_default_instance(self): + # Selecting a supported non-default instance matches the config and pins that instance. + b = self._supported_types_builder() + b.set_deployment_config(instance_type="ml.g5.12xlarge") + self.assertEqual(b.instance_type, "ml.g5.12xlarge") + self.assertEqual(b._selected_hosting_config["EcrAddress"], "img:multi") + # Available-instances error message includes the supported (non-default) instance too. + with self.assertRaises(ValueError) as ctx: + b.set_deployment_config(instance_type="ml.nonexistent.xlarge") + self.assertIn("ml.g5.12xlarge", str(ctx.exception)) + + def test_get_after_set_agrees_with_list_for_supported_non_default_instance(self): + # Round-trip consistency: for a config selected on a SupportedInstanceTypes entry that + # DIFFERS from its default, list / set / get must all report the SAME instance and config + # identity — get_deployment_config() must materialize the pinned instance, not the default. + b = self._supported_types_builder() + listed = b.list_deployment_configs(instance_type="ml.g5.12xlarge")[0] + b.set_deployment_config(instance_type="ml.g5.12xlarge") + got = b.get_deployment_config() + self.assertEqual(got["DeploymentArgs"]["InstanceType"], "ml.g5.12xlarge") + self.assertEqual(got["DeploymentConfigName"], listed["DeploymentConfigName"]) + self.assertEqual( + got["DeploymentArgs"]["InstanceType"], listed["DeploymentArgs"]["InstanceType"] + ) + + def test_get_default_before_selection_reports_offered_instance(self): + # With no explicit selection but a constructor instance that a config offers via + # SupportedInstanceTypes, get_deployment_config() reports that pinned instance. + b = self._builder(instance_type="ml.g5.12xlarge") + configs = [ + { + "DefaultInstanceType": "ml.g5.2xlarge", + "SupportedInstanceTypes": ["ml.g5.2xlarge", "ml.g5.12xlarge"], + "EcrAddress": "img:multi", + "Environment": {}, + }, + ] + with patch.object(ModelBuilder, "_is_model_customization", return_value=True), patch.object( + ModelBuilder, "_resolve_recipe_hosting_configs", return_value=configs + ): + got = b.get_deployment_config() + self.assertEqual(got["DeploymentArgs"]["InstanceType"], "ml.g5.12xlarge") + + def test_list_returns_both_configs_that_offer_same_instance(self): + # Listing does not dedup: two distinct configs both offering the requested instance are + # both returned (materialized for it). Selection, by contrast, rejects this as ambiguous. + configs = [ + {"InstanceType": "ml.g5.2xlarge", "SupportedInstanceTypes": ["ml.g5.12xlarge"]}, + {"InstanceType": "ml.g6.4xlarge", "SupportedInstanceTypes": ["ml.g5.12xlarge"]}, + ] + b = self._builder() + with patch.object(ModelBuilder, "_is_model_customization", return_value=True), patch.object( + ModelBuilder, "_resolve_recipe_hosting_configs", return_value=configs + ): + listed = b.list_deployment_configs(instance_type="ml.g5.12xlarge") + with self.assertRaises(ValueError) as ctx: + b.set_deployment_config(instance_type="ml.g5.12xlarge") + self.assertEqual(len(listed), 2) + self.assertIn("ambiguous", str(ctx.exception)) + + +class TestBuildAppliesSelectedConfig(unittest.TestCase): + """The point of the feature: the SELECTED hosting config's container image, environment, and + compute requirements must actually reach the built artifact via _fetch_and_cache_recipe_config + — not just be returned by the selector. Also pins the DefaultInstanceType round-trip so a + config published with only DefaultInstanceType is not silently dropped at build time. + """ + + def setUp(self): + self.mock_session = _make_mock_session() + + def _builder(self, instance_type): + return ModelBuilder( + model="huggingface-reasoning-qwen3-06b", + model_metadata={ + "CUSTOM_MODEL_ID": "huggingface-reasoning-qwen3-06b", + "CUSTOM_MODEL_VERSION": "3.9.0", + }, + mode=Mode.SAGEMAKER_ENDPOINT, + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + instance_type=instance_type, + ) + + def _patch_recipe(self, b, hosting_configs, recipe_name="recipe-1"): + hub = {"RecipeCollection": [{"Name": recipe_name, "HostingConfigs": hosting_configs}]} + mp = Mock() + container = Mock() + container.base_model.recipe_name = recipe_name + # Non-Nova names so _is_nova_model() (now consulted first at build) returns False. + container.base_model.hub_content_name = "test-hub-content" + mp.inference_specification.containers = [container] + b._fetch_hub_document_for_custom_model = Mock(return_value=hub) + b._fetch_model_package = Mock(return_value=mp) + b.s3_upload_path = "s3://bucket/model" # skip the s3-uri resolution block + b._resolve_compute_requirements_from_config = Mock(return_value={}) + + def test_build_applies_selected_alternative_not_default(self): + # User picked the 8-GPU / TP=8 alternative; build must apply THAT config's env/image, + # not the 1-GPU Default's. + b = self._builder("ml.g6e.48xlarge") + self._patch_recipe(b, HOSTING_CONFIGS) + b.image_uri = None + b.env_vars = None + b._fetch_and_cache_recipe_config() + self.assertEqual(b.image_uri.split(":")[-1], "0.34.0-lmi16") + # The decisive assertion: TP degree came from the selected alternative (8), not Default (1). + self.assertEqual(b.env_vars["OPTION_TENSOR_PARALLEL_DEGREE"], "8") + self.assertEqual(b.instance_type, "ml.g6e.48xlarge") + + def test_build_merges_selected_env_over_user_env(self): + b = self._builder("ml.g6e.48xlarge") + self._patch_recipe(b, HOSTING_CONFIGS) + b.image_uri = None + b.env_vars = {"MY_FLAG": "1"} + b._fetch_and_cache_recipe_config() + self.assertEqual(b.env_vars["MY_FLAG"], "1") # user env preserved + self.assertEqual(b.env_vars["OPTION_TENSOR_PARALLEL_DEGREE"], "8") # recipe env applied + + def test_defaultinstancetype_only_config_round_trips(self): + # A config published with only DefaultInstanceType (no InstanceType) must be selectable AND + # re-match at build time — otherwise it silently falls back to Default (the M1 bug). + img = "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.34.0-lmi16" + raw = [ + { + "Profile": "Default", + "InstanceType": "ml.g6.4xlarge", + "EcrAddress": img, + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "1"}, + }, + { + "DefaultInstanceType": "ml.g6e.48xlarge", + "EcrAddress": img, + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "8"}, + }, + ] + b = self._builder("ml.g6.4xlarge") + with patch.object(ModelBuilder, "_is_model_customization", return_value=True), patch.object( + ModelBuilder, "_resolve_recipe_hosting_configs", return_value=raw + ): + b.set_deployment_config(instance_type="ml.g6e.48xlarge") + # Stored selection is the RAW config (only DefaultInstanceType, no InstanceType key). + self.assertEqual( + b._raw_config_instance_type(b._selected_hosting_config), "ml.g6e.48xlarge" + ) + # Build-time selector must re-match the DefaultInstanceType-only entry, not fall back. + chosen = b._select_recipe_hosting_config(raw) + self.assertEqual(chosen.get("DefaultInstanceType"), "ml.g6e.48xlarge") + self.assertEqual(chosen["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"], "8") + + def test_supported_instance_selection_preserved_at_build(self): + # A config offered via SupportedInstanceTypes (default differs) selected on a NON-default + # supported instance must, at build, apply that config's image/env AND keep the pinned + # instance — not snap back to the config's default instance. + img = "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.34.0-lmi16" + raw = [ + { + "Profile": "Default", + "InstanceType": "ml.g6.4xlarge", + "EcrAddress": img, + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "1"}, + }, + { + "DefaultInstanceType": "ml.g5.2xlarge", + "SupportedInstanceTypes": ["ml.g5.2xlarge", "ml.g5.12xlarge"], + "EcrAddress": img, + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "4"}, + }, + ] + b = self._builder("ml.g6.4xlarge") + self._patch_recipe(b, raw) + b.image_uri = None + b.env_vars = None + with patch.object(ModelBuilder, "_is_model_customization", return_value=True): + b.set_deployment_config(instance_type="ml.g5.12xlarge") + b._fetch_and_cache_recipe_config() + # The multi-instance config's env is applied... + self.assertEqual(b.env_vars["OPTION_TENSOR_PARALLEL_DEGREE"], "4") + # ...and the pinned non-default supported instance is preserved (not snapped to g5.2xlarge). + self.assertEqual(b.instance_type, "ml.g5.12xlarge") + + def test_build_tolerates_explicit_null_environment(self): + # A config that publishes an explicit "Environment": null (key present, value None) must + # not blow up the build's env merge (dict(None)/update(None) would raise TypeError). + img = "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.34.0-lmi16" + raw = [ + { + "Profile": "Default", + "InstanceType": "ml.g6.4xlarge", + "EcrAddress": img, + "Environment": None, + }, + ] + b = self._builder("ml.g6.4xlarge") + self._patch_recipe(b, raw) + b.image_uri = None + b.env_vars = None + b._fetch_and_cache_recipe_config() + self.assertEqual(b.env_vars, {}) + self.assertEqual(b.image_uri, img) + + def test_top_level_config_selected_and_applied_at_build(self): + # P1 (blocking): a config that lives ONLY at the top level (recipe entry publishes no + # HostingConfigs) must be both selectable AND applied at build. Previously the build path + # walked only recipe-level configs, so a selected top-level config was silently dropped. + img = "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.34.0-lmi16" + top = [ + { + "Profile": "Default", + "InstanceType": "ml.g6.4xlarge", + "EcrAddress": img, + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "1"}, + }, + { + "InstanceType": "ml.g6e.48xlarge", + "EcrAddress": img, + "Environment": {"OPTION_TENSOR_PARALLEL_DEGREE": "8"}, + }, + ] + # Recipe entry exists but carries NO HostingConfigs -> top-level fallback. + hub = {"RecipeCollection": [{"Name": "recipe-1"}], "HostingConfigs": top} + b = self._builder("ml.g6e.48xlarge") + mp = Mock() + container = Mock() + container.base_model.recipe_name = "recipe-1" + container.base_model.hub_content_name = "test-hub-content" # non-Nova + mp.inference_specification.containers = [container] + b._fetch_hub_document_for_custom_model = Mock(return_value=hub) + b._fetch_model_package = Mock(return_value=mp) + b.s3_upload_path = "s3://bucket/model" + b._resolve_compute_requirements_from_config = Mock(return_value={}) + b.image_uri = None + b.env_vars = None + with patch.object(ModelBuilder, "_is_model_customization", return_value=True): + # Select via the public API (resolves through the top-level fallback)... + b.set_deployment_config(instance_type="ml.g6e.48xlarge") + # ...and the build path must apply that same top-level config. + b._fetch_and_cache_recipe_config() + self.assertEqual(b.image_uri.split(":")[-1], "0.34.0-lmi16") + self.assertEqual(b.env_vars["OPTION_TENSOR_PARALLEL_DEGREE"], "8") + + def test_user_image_uri_not_overridden_by_selected_config(self): + # A caller-provided image_uri wins; the recipe env is still applied on top. + b = self._builder("ml.g6e.48xlarge") + self._patch_recipe(b, HOSTING_CONFIGS) + b.image_uri = "user-supplied-image" + b.env_vars = None + b._fetch_and_cache_recipe_config() + self.assertEqual(b.image_uri, "user-supplied-image") + self.assertEqual(b.env_vars["OPTION_TENSOR_PARALLEL_DEGREE"], "8") + + def test_constructor_instance_duplicate_matches_first(self): + # Constructor-supplied instance type (no explicit selection) with duplicate matches: the + # first published config wins. (set_deployment_config, by contrast, rejects ambiguity.) + b = self._builder("ml.dup.xlarge") + dup = [ + {"InstanceType": "ml.dup.xlarge", "EcrAddress": "first", "Environment": {}}, + {"InstanceType": "ml.dup.xlarge", "EcrAddress": "second", "Environment": {}}, + ] + cfg = b._select_recipe_hosting_config(dup) + self.assertEqual(cfg["EcrAddress"], "first") + + def test_build_nova_model_uses_nova_path_not_generic(self): + # A Nova model that publishes top-level HostingConfigs must still resolve via the Nova path + # (Nova env precedence + SMI validation), NOT the generic recipe branch. Guards the + # regression where the top-level fallback would divert Nova to the generic path. + b = self._builder("ml.p5.48xlarge") + hub = { + "HostingConfigs": [ + { + "InstanceType": "ml.p5.48xlarge", + "EcrAddress": "generic-image", + "Environment": {"GENERIC": "1"}, + } + ] + } + mp = Mock() + container = Mock() + container.base_model.recipe_name = "nova-lite-recipe" # matches _is_nova_model + container.base_model.hub_content_name = "nova-textgeneration-lite" + mp.inference_specification.containers = [container] + b._fetch_hub_document_for_custom_model = Mock(return_value=hub) + b._fetch_model_package = Mock(return_value=mp) + b.s3_upload_path = "s3://bucket/model" + b.image_uri = None + b.env_vars = {"USER_FLAG": "keep"} + b._get_nova_hosting_config = Mock( + return_value={ + "image_uri": "nova-image", + "env_vars": {"CONTEXT_LENGTH": "8000", "USER_FLAG": "recipe-default"}, + "instance_type": "ml.p5.48xlarge", + } + ) + validate = Mock() + b._validate_nova_smi_config = validate + b._fetch_and_cache_recipe_config() + # Nova image applied (not the generic top-level EcrAddress). + self.assertEqual(b.image_uri, "nova-image") + # Nova env precedence: user override wins over the recipe default. + self.assertEqual(b.env_vars["USER_FLAG"], "keep") + self.assertEqual(b.env_vars["CONTEXT_LENGTH"], "8000") + validate.assert_called_once() + + def test_build_resyncs_instance_type_to_explicit_selection(self): + # If a caller reassigns instance_type AFTER an explicit selection, build re-syncs to the + # selected bundle's instance type so the endpoint instance matches the applied image/env. + b = self._builder("ml.g6.4xlarge") + self._patch_recipe(b, HOSTING_CONFIGS) + with patch.object(ModelBuilder, "_is_model_customization", return_value=True): + b.set_deployment_config(instance_type="ml.g6e.48xlarge") + b.instance_type = "ml.g5.4xlarge" # contradictory direct reassignment + b.image_uri = None + b.env_vars = None + b._fetch_and_cache_recipe_config() + self.assertEqual(b.instance_type, "ml.g6e.48xlarge") # re-synced to the selection + self.assertEqual(b.env_vars["OPTION_TENSOR_PARALLEL_DEGREE"], "8") + + def test_build_passes_selected_config_to_compute_requirements(self): + # The SELECTED config (not Default) must reach _resolve_compute_requirements_from_config, + # and its result is cached — proving compute requirements travel with the selection. + b = self._builder("ml.g6e.48xlarge") + self._patch_recipe(b, HOSTING_CONFIGS) + calls = {} + + def recorder(instance_type, config, user_resource_requirements): + calls["instance_type"] = instance_type + calls["config"] = config + return {"NumberOfAcceleratorDevicesRequired": 8} + + b._resolve_compute_requirements_from_config = recorder + b.image_uri = None + b.env_vars = None + with patch.object(ModelBuilder, "_is_model_customization", return_value=True): + b.set_deployment_config(instance_type="ml.g6e.48xlarge") + b._fetch_and_cache_recipe_config() + self.assertEqual(calls["instance_type"], "ml.g6e.48xlarge") + self.assertEqual(calls["config"]["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"], "8") + self.assertEqual(b._cached_compute_requirements["NumberOfAcceleratorDevicesRequired"], 8) + + def test_build_no_configs_non_nova_raises(self): + # No hosting configs anywhere and not a Nova model -> clear unsupported error. + b = self._builder("ml.g6.4xlarge") + hub = {"RecipeCollection": [{"Name": "r1"}]} # no recipe-level or top-level configs + mp = Mock() + container = Mock() + container.base_model.recipe_name = "r1" + container.base_model.hub_content_name = "not-nova" + mp.inference_specification.containers = [container] + b._fetch_hub_document_for_custom_model = Mock(return_value=hub) + b._fetch_model_package = Mock(return_value=mp) + b.s3_upload_path = "s3://bucket/model" + with self.assertRaises(ValueError) as ctx: + b._fetch_and_cache_recipe_config() + self.assertIn("not supported for deployment", str(ctx.exception)) + + def test_returned_config_mutation_does_not_corrupt_build(self): + # Cross-path isolation: mutating a config returned by get/list must not change what build + # applies. + b = self._builder("ml.g6e.48xlarge") + self._patch_recipe(b, HOSTING_CONFIGS) + with patch.object(ModelBuilder, "_is_model_customization", return_value=True): + b.set_deployment_config(instance_type="ml.g6e.48xlarge") + got = b.get_deployment_config() + got["DeploymentArgs"]["Environment"]["OPTION_TENSOR_PARALLEL_DEGREE"] = "999" + b.image_uri = None + b.env_vars = None + b._fetch_and_cache_recipe_config() + self.assertEqual(b.env_vars["OPTION_TENSOR_PARALLEL_DEGREE"], "8") + + +class TestRecipeHostingConfigHelpers(unittest.TestCase): + """Direct coverage of the internal helpers behind the unified deployment-config API, plus a + guard that a base/JumpStart model still routes through the same public entry point — the + base-vs-fine-tuned split is chosen only by ``_is_model_customization()``, never surfaced. + """ + + def setUp(self): + self.mock_session = _make_mock_session() + + def _builder(self, model="huggingface-reasoning-qwen3-06b"): + return ModelBuilder( + model=model, + model_metadata={ + "CUSTOM_MODEL_ID": "huggingface-reasoning-qwen3-06b", + "CUSTOM_MODEL_VERSION": "3.9.0", + }, + mode=Mode.SAGEMAKER_ENDPOINT, + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + # Pin an instance type so construction stays hermetic — without it the constructor + # auto-detects the instance type via a live JumpStart S3 lookup. + instance_type="ml.g6.4xlarge", + ) + + # ---- _normalize_hosting_config: stable list/get-friendly shape ---- + + def test_normalize_named_default_base_shape(self): + norm = ModelBuilder._normalize_hosting_config(HOSTING_CONFIGS[0]) + self.assertEqual(norm["DeploymentConfigName"], "Default") + self.assertTrue(norm["IsDefault"]) + args = norm["DeploymentArgs"] + self.assertEqual(args["InstanceType"], "ml.g6.4xlarge") + self.assertEqual(args["ImageUri"].split(":")[-1], "0.34.0-lmi16") + self.assertEqual( + args["ComputeResourceRequirements"]["NumberOfAcceleratorDevicesRequired"], 1 + ) + # BenchmarkMetrics/AccelerationConfigs present but None, matching the base sentinel. + self.assertIsNone(norm["BenchmarkMetrics"]) + self.assertIsNone(norm["AccelerationConfigs"]) + + def test_normalize_unnamed_uses_instance_type_identifier(self): + # Unnamed configs (no Profile) get their instance type as a stable identifier. + norm = ModelBuilder._normalize_hosting_config(HOSTING_CONFIGS[2]) + self.assertEqual(norm["DeploymentConfigName"], "ml.g6e.48xlarge") + self.assertFalse(norm["IsDefault"]) + self.assertEqual(norm["DeploymentArgs"]["InstanceType"], "ml.g6e.48xlarge") + + def test_normalize_defaults_missing_fields(self): + norm = ModelBuilder._normalize_hosting_config({"InstanceType": "ml.g6.4xlarge"}) + args = norm["DeploymentArgs"] + self.assertEqual(args["Environment"], {}) + self.assertEqual(args["ComputeResourceRequirements"], {}) + self.assertIsNone(args["ImageUri"]) + + # ---- _resolve_recipe_hosting_configs: recipe match, top-level fallback, no-configs raise ---- + + def _patch_recipe(self, b, hub_document, recipe_name): + mp = Mock() + container = Mock() + container.base_model.recipe_name = recipe_name + mp.inference_specification.containers = [container] + b._fetch_hub_document_for_custom_model = Mock(return_value=hub_document) + b._fetch_model_package = Mock(return_value=mp) + + def test_resolve_from_recipe_collection(self): + b = self._builder() + hub = {"RecipeCollection": [{"Name": "r1", "HostingConfigs": HOSTING_CONFIGS}]} + self._patch_recipe(b, hub, "r1") + self.assertEqual(b._resolve_recipe_hosting_configs(), HOSTING_CONFIGS) + + def test_resolve_falls_back_to_top_level(self): + # Recipe entry matches but publishes no HostingConfigs (missing key) -> fall back to top. + b = self._builder() + hub = {"RecipeCollection": [{"Name": "r1"}], "HostingConfigs": HOSTING_CONFIGS} + self._patch_recipe(b, hub, "r1") + self.assertEqual(b._resolve_recipe_hosting_configs(), HOSTING_CONFIGS) + + def test_resolve_recipe_empty_list_falls_back_to_top_level(self): + # Recipe entry matches with an EMPTY HostingConfigs list (not just a missing key) -> top. + b = self._builder() + hub = { + "RecipeCollection": [{"Name": "r1", "HostingConfigs": []}], + "HostingConfigs": HOSTING_CONFIGS, + } + self._patch_recipe(b, hub, "r1") + self.assertEqual(b._resolve_recipe_hosting_configs(), HOSTING_CONFIGS) + + def test_resolve_unmatched_recipe_name_falls_back_to_top_level(self): + # The model's recipe_name matches no RecipeCollection entry -> top-level fallback (does not + # accidentally pick another recipe's configs). + b = self._builder() + other = [{"InstanceType": "ml.other.xlarge", "Environment": {}}] + hub = { + "RecipeCollection": [{"Name": "some-other-recipe", "HostingConfigs": other}], + "HostingConfigs": HOSTING_CONFIGS, + } + self._patch_recipe(b, hub, "r1") # model's recipe_name is r1, unmatched + self.assertEqual(b._resolve_recipe_hosting_configs(), HOSTING_CONFIGS) + + def test_extract_hosting_configs_from_hub_precedence(self): + # Direct coverage of the single-source-of-truth discovery helper used by BOTH the selection + # API and the build path. + rc = [{"InstanceType": "ml.recipe.xlarge"}] + top = [{"InstanceType": "ml.top.xlarge"}] + # recipe-level present -> recipe-level wins + self.assertEqual( + ModelBuilder._extract_hosting_configs_from_hub( + {"RecipeCollection": [{"Name": "r1", "HostingConfigs": rc}], "HostingConfigs": top}, + "r1", + ), + rc, + ) + # recipe matched but empty -> top-level + self.assertEqual( + ModelBuilder._extract_hosting_configs_from_hub( + {"RecipeCollection": [{"Name": "r1", "HostingConfigs": []}], "HostingConfigs": top}, + "r1", + ), + top, + ) + # recipe unmatched -> top-level + self.assertEqual( + ModelBuilder._extract_hosting_configs_from_hub( + { + "RecipeCollection": [{"Name": "other", "HostingConfigs": rc}], + "HostingConfigs": top, + }, + "r1", + ), + top, + ) + # neither -> empty + self.assertEqual( + ModelBuilder._extract_hosting_configs_from_hub({"RecipeCollection": []}, "r1"), [] + ) + + def test_resolve_no_configs_raises(self): + b = self._builder() + self._patch_recipe(b, {"RecipeCollection": []}, "r1") + with self.assertRaises(ValueError) as ctx: + b._resolve_recipe_hosting_configs() + self.assertIn("does not publish any hosting configurations", str(ctx.exception)) + + def test_resolve_no_container_raises_clean_error(self): + # A model package with no inference container must raise a clear ValueError, not a raw + # AttributeError/IndexError from containers[0]. + b = self._builder() + b._fetch_hub_document_for_custom_model = Mock(return_value={"RecipeCollection": []}) + mp = Mock() + mp.inference_specification.containers = [] + b._fetch_model_package = Mock(return_value=mp) + with self.assertRaises(ValueError) as ctx: + b._resolve_recipe_hosting_configs() + self.assertIn("no inference container", str(ctx.exception)) + + def test_normalized_shape_covers_base_deployment_config_keys(self): + # SHAPE PARITY (the whole point of normalization): the normalized fine-tuned config must + # expose at least every key the real base/JumpStart response emits, so a caller can index + # the same keys on either pathway. Tie the assertion to the actual base dataclasses so it + # fails loudly if base ever adds a field we don't mirror. + from sagemaker.core.jumpstart.types import DeploymentArgs, DeploymentConfigMetadata + + def pascal(s): + return s.replace("_", " ").title().replace(" ", "") + + base_top_keys = {pascal(s) for s in DeploymentConfigMetadata.__slots__} + base_args_keys = {pascal(s) for s in DeploymentArgs.__slots__} + + norm = ModelBuilder._normalize_hosting_config(HOSTING_CONFIGS[0]) + self.assertTrue( + base_top_keys <= set(norm), + f"normalized top-level missing base keys: {base_top_keys - set(norm)}", + ) + self.assertTrue( + base_args_keys <= set(norm["DeploymentArgs"]), + f"normalized DeploymentArgs missing base keys: " + f"{base_args_keys - set(norm['DeploymentArgs'])}", + ) + + # ---- base-model routing through the SAME unified entry point ---- + + def test_list_deployment_configs_routes_to_base_for_string_model(self): + # For a base/JumpStart string model, list_deployment_configs must take the base branch + # (never the recipe resolver). With no instance_type it returns every config materialized + # at its default (the original base API). + b = self._builder() + base_configs = [ + {"DeploymentConfigName": "lmi", "DeploymentArgs": {"InstanceType": "ml.g5.2xlarge"}}, + {"DeploymentConfigName": "tgi", "DeploymentArgs": {"InstanceType": "ml.g6.4xlarge"}}, + ] + with patch.object( + ModelBuilder, "_is_model_customization", return_value=False + ), patch.object(ModelBuilder, "_is_jumpstart_model_id", return_value=True), patch.object( + ModelBuilder, "_get_deployment_configs", return_value=["sentinel"] + ), patch.object( + ModelBuilder, "deployment_config_response_data", return_value=base_configs + ), patch.object( + ModelBuilder, "_resolve_recipe_hosting_configs" + ) as resolve_recipe: + self.assertEqual(len(b.list_deployment_configs()), 2) + # The fine-tuned resolver is never touched on the base path. + resolve_recipe.assert_not_called() + + def test_list_deployment_configs_base_filter_keeps_supported_non_default_instance(self): + # Regression for the base-filter bug (reviewer P1): a base config is a MULTI-instance + # bundle. list_deployment_configs(instance_type=X) must return a config that SUPPORTS X even + # when X is not its default — materialized FOR X — instead of discarding it. Exercised via + # the supported-instance metadata + the real per-config materialization loop. + b = self._builder() + # "big" supports the requested ml.g5.12xlarge but DEFAULTS to ml.g5.2xlarge; "small" does + # not support it at all and must be filtered out. + meta_big = Mock( + resolved_config={ + "default_inference_instance_type": "ml.g5.2xlarge", + "supported_inference_instance_types": ["ml.g5.2xlarge", "ml.g5.12xlarge"], + } + ) + meta_small = Mock( + resolved_config={ + "default_inference_instance_type": "ml.m5.large", + "supported_inference_instance_types": ["ml.m5.large"], + } + ) + b._metadata_configs = {"big": meta_big, "small": meta_small} + + def _fake_get_configs(selected_config_name, selected_instance_type): + # Mimic _get_deployment_configs: the SELECTED config is materialized at the requested + # instance; every config is returned. + return [ + { + "DeploymentConfigName": name, + "DeploymentArgs": { + "InstanceType": ( + selected_instance_type + if name == selected_config_name + else "DEFAULT" + ) + }, + } + for name in ("big", "small") + ] + + with patch.object( + ModelBuilder, "_is_model_customization", return_value=False + ), patch.object(ModelBuilder, "_is_jumpstart_model_id", return_value=True), patch.object( + ModelBuilder, "_ensure_metadata_configs" + ), patch.object( + ModelBuilder, "_get_deployment_configs", side_effect=_fake_get_configs + ), patch.object( + ModelBuilder, + "deployment_config_response_data", + side_effect=lambda configs: configs, + ): + result = b.list_deployment_configs(instance_type="ml.g5.12xlarge") + + # Only the supporting config is returned, materialized FOR the requested instance. + self.assertEqual(len(result), 1) + self.assertEqual(result[0]["DeploymentConfigName"], "big") + self.assertEqual(result[0]["DeploymentArgs"]["InstanceType"], "ml.g5.12xlarge") + + def test_get_deployment_configs_honors_requested_instance_for_selected(self): + # Exercises the REAL _get_deployment_configs path (not a mocked deployment_config_response_ + # data). The SELECTED config must be materialized with the REQUESTED instance type — even + # when that differs from its default — while other configs keep their defaults, none dropped. + from sagemaker.serve import model_builder_utils as mbu + + b = self._builder() + b.image_uri = "img" + meta_lmi = Mock( + benchmark_metrics=None, + config_components={}, + resolved_config={"default_inference_instance_type": "ml.g5.2xlarge"}, + ) + meta_tgi = Mock( + benchmark_metrics=None, + config_components={}, + resolved_config={"default_inference_instance_type": "ml.g5.2xlarge"}, + ) + b._metadata_configs = {"lmi": meta_lmi, "tgi": meta_tgi} + + captured = {} + + def _capture_init_kwargs(**kwargs): + captured[kwargs["config_name"]] = kwargs["instance_type"] + return Mock() + + with patch.object(mbu, "get_init_kwargs", side_effect=_capture_init_kwargs), patch.object( + mbu, "get_deploy_kwargs", return_value=Mock() + ), patch.object(mbu, "DeploymentConfigMetadata", side_effect=lambda *a, **k: Mock()): + configs = b._get_deployment_configs("lmi", "ml.g5.12xlarge") + + # Both configs materialized (none dropped). + self.assertEqual(len(configs), 2) + # Selected config uses the REQUESTED instance; the other keeps its default. + self.assertEqual(captured["lmi"], "ml.g5.12xlarge") + self.assertEqual(captured["tgi"], "ml.g5.2xlarge") + + def test_set_deployment_config_base_requires_instance_type(self): + # Base/JumpStart setter: config_name alone is not enough — instance_type is still required. + # Guards against the optional-signature (added for the fine-tuned path) silently forwarding + # None into the base selection machinery instead of failing fast. + b = self._builder() + with patch.object( + ModelBuilder, "_is_model_customization", return_value=False + ), patch.object(ModelBuilder, "_is_jumpstart_model_id", return_value=True): + with self.assertRaises(ValueError) as ctx: + b.set_deployment_config(config_name="lmi") + self.assertIn("instance_type is required", str(ctx.exception)) + + def test_list_deployment_configs_base_non_jumpstart_raises(self): + # A non-customization, non-JumpStart string model still hits the base guard. + b = self._builder() + with patch.object( + ModelBuilder, "_is_model_customization", return_value=False + ), patch.object(ModelBuilder, "_is_jumpstart_model_id", return_value=False), patch.object( + ModelBuilder, "_use_jumpstart_equivalent", return_value=False + ): + with self.assertRaises(ValueError) as ctx: + b.list_deployment_configs() + self.assertIn("JumpStart", str(ctx.exception)) + + def test_is_nova_model_tolerates_non_string_attrs(self): + # Regression: _is_nova_model (consulted first on the model-customization build path) must + # not raise when a partially-populated model package leaves recipe_name/hub_content_name as + # a non-string (e.g. an unset Mock attribute). "nova" in would raise TypeError. + b = self._builder() + pkg = Mock() + container = Mock() + container.base_model.recipe_name = "test-recipe" # non-Nova string + # hub_content_name intentionally left as an auto Mock (not a string). + pkg.inference_specification.containers = [container] + with patch.object(ModelBuilder, "_fetch_model_package", return_value=pkg): + self.assertFalse(b._is_nova_model()) # returns False instead of raising + + +# A matrix of ADVERSARIAL raw HostingConfigs sets that deliberately covers the shapes the earlier +# hand-picked fixture never did — most importantly configs whose default instance DIFFERS from an +# instance they offer via SupportedInstanceTypes (the exact case that hid the get() bug), plus +# DefaultInstanceType-only entries, ambiguous overlaps, and a Default profile with a supported +# superset. The invariant tests below run EVERY scenario against ALL of list/set/get so a +# disagreement between the methods surfaces — coverage that can't be tuned to the implementation. +_ADVERSARIAL_CONFIG_SETS = { + "single_instance_bundles": [ + {"Profile": "Default", "InstanceType": "ml.g6.4xlarge", "Environment": {"T": "d"}}, + {"InstanceType": "ml.g6e.48xlarge", "Environment": {"T": "a"}}, + ], + "default_instance_type_only": [ + {"DefaultInstanceType": "ml.g6.4xlarge", "Environment": {"T": "d"}}, + {"DefaultInstanceType": "ml.g6e.48xlarge", "Environment": {"T": "a"}}, + ], + "supported_superset_differs_from_default": [ + {"Profile": "Default", "InstanceType": "ml.g6.4xlarge", "Environment": {"T": "d"}}, + { + "DefaultInstanceType": "ml.g5.2xlarge", + "SupportedInstanceTypes": ["ml.g5.2xlarge", "ml.g5.12xlarge"], + "Environment": {"T": "m"}, + }, + ], + "supported_not_including_primary": [ + { + "InstanceType": "ml.g6.4xlarge", + "SupportedInstanceTypes": ["ml.p5.48xlarge"], + "Environment": {"T": "s"}, + }, + ], + "default_profile_with_supported_superset": [ + { + "Profile": "Default", + "InstanceType": "ml.g6.4xlarge", + "SupportedInstanceTypes": ["ml.g6.4xlarge", "ml.g5.12xlarge"], + "Environment": {"T": "d"}, + }, + {"InstanceType": "ml.p5.48xlarge", "Environment": {"T": "a"}}, + ], + "ambiguous_overlap": [ + {"InstanceType": "ml.g6.4xlarge", "SupportedInstanceTypes": ["ml.g5.12xlarge"]}, + {"InstanceType": "ml.p5.48xlarge", "SupportedInstanceTypes": ["ml.g5.12xlarge"]}, + ], + "many_mixed": [ + {"Profile": "Default", "InstanceType": "ml.g6.4xlarge", "Environment": {"T": "d"}}, + {"DefaultInstanceType": "ml.g5.2xlarge", "Environment": {"T": "a"}}, + { + "InstanceType": "ml.g6e.48xlarge", + "SupportedInstanceTypes": ["ml.g6e.48xlarge", "ml.p5.48xlarge"], + "Environment": {"T": "b"}, + }, + ], +} + + +class TestDeploymentConfigInvariants(unittest.TestCase): + """Cross-method invariants for the unified deployment-config API, run over an adversarial matrix. + + These exist because the earlier suite was superficial in two ways that let a real bug ship: + (1) every method was tested in ISOLATION, so a DISAGREEMENT between list/set/get could not be + seen; and (2) the single hand-picked fixture had one instance per config (default == the + requested instance), so the materialization step was always a no-op and a missing + materialization in get_deployment_config() was invisible. + + Each test asserts an invariant that must hold for EVERY scenario in `_ADVERSARIAL_CONFIG_SETS` + — including `supported_superset_differs_from_default`, the exact shape the shipped get() bug + got wrong. The round-trip test fails against that bug; the earlier per-method tests could not. + (No external test dependency — deterministic table-driven checks; hypothesis is not a declared + dependency of this package.) + """ + + def setUp(self): + self.mock_session = _make_mock_session() + + def _builder(self, configs): + # A fine-tuned-looking builder over `configs`. Patches stay active for the caller (started + # here, stopped via addCleanup) so list/get/set all see the same configs. + b = ModelBuilder( + model="huggingface-reasoning-qwen3-06b", + model_metadata={ + "CUSTOM_MODEL_ID": "huggingface-reasoning-qwen3-06b", + "CUSTOM_MODEL_VERSION": "3.9.0", + }, + mode=Mode.SAGEMAKER_ENDPOINT, + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + instance_type="ml.g6.4xlarge", + ) + p1 = patch.object(ModelBuilder, "_is_model_customization", return_value=True) + p2 = patch.object(ModelBuilder, "_resolve_recipe_hosting_configs", return_value=configs) + p1.start() + p2.start() + self.addCleanup(p1.stop) + self.addCleanup(p2.stop) + return b + + @staticmethod + def _offered_counts(configs): + counts = {} + for cfg in configs: + for inst in ModelBuilder._raw_config_offered_instances(cfg): + counts[inst] = counts.get(inst, 0) + 1 + return counts + + _DEPLOYMENT_ARGS_KEYS = { + "ImageUri", + "InstanceType", + "Environment", + "ComputeResourceRequirements", + "ModelData", + "ModelPackageArn", + "ModelDataDownloadTimeout", + "ContainerStartupHealthCheckTimeout", + "AdditionalDataSources", + } + + def test_list_shape_invariant(self): + # INVARIANT: every listed config exposes the full normalized shape, and exactly one config + # is flagged IsDefault iff a "Default" profile exists. + for name, configs in _ADVERSARIAL_CONFIG_SETS.items(): + with self.subTest(scenario=name): + b = self._builder(configs) + listed = b.list_deployment_configs() + self.assertEqual(len(listed), len(configs)) + for c in listed: + self.assertIn("DeploymentConfigName", c) + self.assertIn("BenchmarkMetrics", c) + self.assertIn("AccelerationConfigs", c) + self.assertIn("IsDefault", c) + self.assertEqual(set(c["DeploymentArgs"]), self._DEPLOYMENT_ARGS_KEYS) + num_default = sum(1 for c in listed if c["IsDefault"]) + expected_default = sum(1 for cfg in configs if cfg.get("Profile") == "Default") + self.assertEqual(num_default, expected_default) + + def test_list_filter_sound_and_complete(self): + # INVARIANT: for every offered instance X, list(instance_type=X) returns EXACTLY the configs + # that offer X (completeness), each materialized FOR X (soundness). This is the invariant the + # base-filter bug violated (a supported non-default config was dropped). + for name, configs in _ADVERSARIAL_CONFIG_SETS.items(): + with self.subTest(scenario=name): + b = self._builder(configs) + for x, count in self._offered_counts(configs).items(): + filtered = b.list_deployment_configs(instance_type=x) + self.assertEqual(len(filtered), count, f"{name}/{x}") + for c in filtered: + self.assertEqual(c["DeploymentArgs"]["InstanceType"], x) + self.assertEqual( + b.list_deployment_configs(instance_type="ml.not-offered.xlarge"), [] + ) + + def test_list_set_get_round_trip(self): + # INVARIANT (the one that would have caught the shipped bug): for any UNAMBIGUOUS offered + # instance X, set(instance_type=X) then get() must agree with list(instance_type=X) on BOTH + # DeploymentArgs.InstanceType (== X) and DeploymentConfigName. Fresh builder per X. + for name, configs in _ADVERSARIAL_CONFIG_SETS.items(): + counts = self._offered_counts(configs) + for x, count in counts.items(): + if count != 1: + continue # ambiguous instances covered by the ambiguity invariant + with self.subTest(scenario=name, instance=x): + b = self._builder(configs) + listed = b.list_deployment_configs(instance_type=x) + self.assertEqual(len(listed), 1) + b.set_deployment_config(instance_type=x) + got = b.get_deployment_config() + self.assertEqual(got["DeploymentArgs"]["InstanceType"], x) + self.assertEqual( + got["DeploymentConfigName"], listed[0]["DeploymentConfigName"] + ) + + def test_ambiguous_instance_rejected_by_set(self): + # INVARIANT: an instance offered by MORE THAN ONE config must be rejected by set() rather + # than silently picking one. + for name, configs in _ADVERSARIAL_CONFIG_SETS.items(): + counts = self._offered_counts(configs) + for x, count in counts.items(): + if count < 2: + continue + with self.subTest(scenario=name, instance=x): + b = self._builder(configs) + with self.assertRaises(ValueError) as ctx: + b.set_deployment_config(instance_type=x) + self.assertIn("ambiguous", str(ctx.exception)) + + def test_returned_configs_are_isolated_from_internal_state(self): + # INVARIANT: mutating any dict returned by list()/get() must not affect a later call — + # returned data is always freshly built / copied. + for name, configs in _ADVERSARIAL_CONFIG_SETS.items(): + with self.subTest(scenario=name): + b = self._builder(configs) + first = b.list_deployment_configs() + for c in first: + c["DeploymentConfigName"] = "MUTATED" + c["DeploymentArgs"]["Environment"]["T"] = "MUTATED" + second = b.list_deployment_configs() + self.assertTrue(all(c["DeploymentConfigName"] != "MUTATED" for c in second)) + self.assertTrue( + all(c["DeploymentArgs"]["Environment"].get("T") != "MUTATED" for c in second) + ) + + +if __name__ == "__main__": + unittest.main()