Skip to content

Commit 0b762fc

Browse files
committed
feat(serve): support fine-tuned models in deployment-config API
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 the SAME shape the base response returns (DeploymentConfigName plus a nested DeploymentArgs block with ImageUri/InstanceType/Environment/ ComputeResourceRequirements and the other base DeploymentArgs keys, present as None when a recipe does not provide them), so a caller can iterate results from either pathway identically. BenchmarkMetrics and AccelerationConfigs are None (recipes publish none today); an additive IsDefault flag marks the recipe's Default config. set_deployment_config(instance_type=...) applies the whole matching config (image, environment, compute requirements) at build time and raises on an unpublished or ambiguous instance type. list_deployment_ configs(instance_type=...) filters to configs offering that instance. Configs published with only DefaultInstanceType are matched consistently end to end so they are not silently dropped at build. 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 Testing: python3 -m pytest sagemaker-serve/tests/unit/test_recipe_hosting_config_selection.py (31 tests) plus the base-model regression suites (coverage boost, instance type inference, nova hosting config, compute requirements) — 102 passed. Includes a shape-parity test tied to the base DeploymentArgs / DeploymentConfigMetadata dataclass slots so it fails if base adds a field the fine-tuned shape does not mirror.
1 parent 3f810c6 commit 0b762fc

2 files changed

Lines changed: 817 additions & 18 deletions

File tree

‎sagemaker-serve/src/sagemaker/serve/model_builder.py‎

Lines changed: 284 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -393,6 +393,7 @@ class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuild
393393
_tags: Optional[Tags] = field(default=None, init=False)
394394
_optimizing: bool = field(default=False, init=False)
395395
_deployment_config: Optional[Dict[str, Any]] = field(default=None, init=False)
396+
_selected_hosting_config: Optional[Dict[str, Any]] = field(default=None, init=False)
396397

397398
shared_libs: List[str] = field(
398399
default_factory=list,
@@ -964,6 +965,169 @@ def _is_gpu_instance(self, instance_type: str) -> bool:
964965
gpu_count = self._infer_accelerator_count_from_instance_type(instance_type)
965966
return gpu_count is not None and gpu_count > 0
966967

968+
def _resolve_recipe_hosting_configs(self) -> List[Dict[str, Any]]:
969+
"""Return the raw ``HostingConfigs`` list published for this model's recipe.
970+
971+
A fine-tuned (model-customization) model serves the base model under its recipe, so the
972+
authoritative hosting configurations live in the JumpStart hub document under
973+
``RecipeCollection[<recipe>].HostingConfigs``. Falls back to the top-level
974+
``HostingConfigs`` when the recipe entry does not carry its own (some models publish
975+
there instead). Raises when neither is present.
976+
"""
977+
hub_document = self._fetch_hub_document_for_custom_model()
978+
model_package = self._fetch_model_package()
979+
containers = getattr(
980+
getattr(model_package, "inference_specification", None), "containers", None
981+
)
982+
if not containers:
983+
raise ValueError(
984+
"Model is not supported for deployment: the model package has no inference "
985+
"container to resolve a recipe from."
986+
)
987+
recipe_name = getattr(containers[0].base_model, "recipe_name", None) or ""
988+
989+
for recipe in hub_document.get("RecipeCollection", []):
990+
if recipe.get("Name") == recipe_name:
991+
configs = recipe.get("HostingConfigs")
992+
if configs:
993+
return configs
994+
break
995+
996+
configs = hub_document.get("HostingConfigs")
997+
if configs:
998+
return configs
999+
1000+
detail = (
1001+
f"recipe '{recipe_name}'" if recipe_name else "this model's recipe"
1002+
)
1003+
raise ValueError(
1004+
f"Model is not supported for deployment: {detail} does not publish any "
1005+
"hosting configurations."
1006+
)
1007+
1008+
@staticmethod
1009+
def _raw_config_instance_type(cfg: Dict[str, Any]) -> Optional[str]:
1010+
"""Instance type of a raw recipe ``HostingConfigs`` entry.
1011+
1012+
Recipe entries publish the instance type under ``InstanceType`` or, for some, only
1013+
``DefaultInstanceType``. Every place that reads a raw entry's instance type MUST use this
1014+
so normalization, selection, and build-time matching stay consistent (otherwise a
1015+
``DefaultInstanceType``-only config is listable/selectable but silently fails to re-match
1016+
at build and falls back to Default).
1017+
"""
1018+
return cfg.get("InstanceType") or cfg.get("DefaultInstanceType")
1019+
1020+
@staticmethod
1021+
def _normalize_hosting_config(cfg: Dict[str, Any]) -> Dict[str, Any]:
1022+
"""Normalize a raw recipe ``HostingConfigs`` entry into the SAME shape the base/JumpStart
1023+
``list_deployment_configs`` response uses, so a caller can iterate results from either
1024+
pathway identically.
1025+
1026+
The serving fields live under ``DeploymentArgs`` with the SAME keys the base response
1027+
nests (``ImageUri``, ``InstanceType``, ``Environment``, ``ComputeResourceRequirements``,
1028+
``ModelData``, ``ModelPackageArn``, ``ModelDataDownloadTimeout``,
1029+
``ContainerStartupHealthCheckTimeout``, ``AdditionalDataSources``). Keys a recipe does not
1030+
provide are present with a ``None`` value so a caller can index the same keys on either
1031+
pathway without a ``KeyError``. ``BenchmarkMetrics``/``AccelerationConfigs`` are ``None``
1032+
(recipes don't publish these today; base emits ``None`` when absent too — extensible:
1033+
populate when recipes do). ``IsDefault`` is an additive convenience the recipe uniquely
1034+
provides (it marks one config with ``Profile == "Default"``) and is a superset of the base
1035+
response, which carries no per-config default flag.
1036+
"""
1037+
profile = cfg.get("Profile")
1038+
instance_type = ModelBuilder._raw_config_instance_type(cfg)
1039+
return {
1040+
# Configs are mostly unnamed (only "Default" has a Profile); fall back to the instance
1041+
# type as a stable, human-meaningful identifier. Same key the base configs use.
1042+
"DeploymentConfigName": profile or instance_type,
1043+
"DeploymentArgs": {
1044+
"ImageUri": cfg.get("EcrAddress"),
1045+
"InstanceType": instance_type,
1046+
"Environment": cfg.get("Environment", {}),
1047+
"ComputeResourceRequirements": cfg.get("ComputeResourceRequirements", {}),
1048+
# Keys the base DeploymentArgs carries that recipe configs don't populate — present
1049+
# as None so both pathways expose an identical DeploymentArgs key set.
1050+
"ModelData": None,
1051+
"ModelPackageArn": None,
1052+
"ModelDataDownloadTimeout": None,
1053+
"ContainerStartupHealthCheckTimeout": None,
1054+
"AdditionalDataSources": None,
1055+
},
1056+
"BenchmarkMetrics": None,
1057+
"AccelerationConfigs": None,
1058+
"IsDefault": profile == "Default",
1059+
}
1060+
1061+
def _select_recipe_hosting_config(self, hosting_configs: list) -> dict:
1062+
"""Select which recipe hosting configuration to apply.
1063+
1064+
A recipe publishes several pre-benchmarked hosting configurations (its
1065+
``HostingConfigs``); each is a self-consistent bundle of an instance type,
1066+
serving container image (``EcrAddress``), container ``Environment`` (e.g.
1067+
``OPTION_TENSOR_PARALLEL_DEGREE``), and ``ComputeResourceRequirements``.
1068+
1069+
Selection precedence:
1070+
1071+
1. A config explicitly chosen via :meth:`set_deployment_config`
1072+
(``_selected_hosting_config``).
1073+
2. The config whose ``InstanceType`` matches a caller-supplied ``instance_type`` — so its
1074+
image, environment, and compute requirements are all applied together and consistently,
1075+
rather than pinning the caller's instance onto the Default config's environment (which
1076+
would, for example, leave tensor-parallel degree at the Default value on a differently
1077+
sized instance). Mirrors the Nova path's ``_select_nova_hosting_config_entry``.
1078+
3. The ``Default`` config (or the first entry) — the prior behavior — when no instance type
1079+
was supplied, or when the supplied instance type matches no published config.
1080+
1081+
Args:
1082+
hosting_configs: The recipe's ``HostingConfigs`` list.
1083+
1084+
Returns:
1085+
The selected hosting config dict.
1086+
"""
1087+
if not hosting_configs:
1088+
raise ValueError(
1089+
"Model is not supported for deployment: no hosting configurations to select from."
1090+
)
1091+
if self._selected_hosting_config is not None:
1092+
selected_instance = self._config_instance_type(self._selected_hosting_config)
1093+
match = next(
1094+
(
1095+
cfg
1096+
for cfg in hosting_configs
1097+
if self._raw_config_instance_type(cfg) == selected_instance
1098+
),
1099+
None,
1100+
)
1101+
if match is not None:
1102+
return match
1103+
1104+
if self._user_provided_instance_type and self.instance_type:
1105+
match = next(
1106+
(
1107+
cfg
1108+
for cfg in hosting_configs
1109+
if self._raw_config_instance_type(cfg) == self.instance_type
1110+
),
1111+
None,
1112+
)
1113+
if match is not None:
1114+
logger.info(
1115+
"Selected recipe hosting config for instance type %s.",
1116+
self.instance_type,
1117+
)
1118+
return match
1119+
logger.warning(
1120+
"Instance type %s does not match any published hosting config for this "
1121+
"recipe; falling back to the Default hosting config's container image, "
1122+
"environment, and compute requirements. Published instance types: %s",
1123+
self.instance_type,
1124+
[self._raw_config_instance_type(cfg) for cfg in hosting_configs],
1125+
)
1126+
return next(
1127+
(cfg for cfg in hosting_configs if cfg.get("Profile") == "Default"),
1128+
hosting_configs[0],
1129+
)
1130+
9671131
def _fetch_and_cache_recipe_config(self):
9681132
"""Fetch and cache image URI, compute requirements, and s3_upload_path from recipe during build."""
9691133
hub_document = self._fetch_hub_document_for_custom_model()
@@ -981,10 +1145,7 @@ def _fetch_and_cache_recipe_config(self):
9811145
if recipe.get("Name") == recipe_name:
9821146
hosting_configs = recipe.get("HostingConfigs", [])
9831147
if hosting_configs:
984-
config = next(
985-
(cfg for cfg in hosting_configs if cfg.get("Profile") == "Default"),
986-
hosting_configs[0],
987-
)
1148+
config = self._select_recipe_hosting_config(hosting_configs)
9881149
if not self.image_uri:
9891150
self.image_uri = config.get("EcrAddress")
9901151

@@ -4035,8 +4196,53 @@ def display_benchmark_metrics(self, **kwargs) -> None:
40354196
@_telemetry_emitter(
40364197
feature=Feature.MODEL_CUSTOMIZATION, func_name="model_builder.set_deployment_config"
40374198
)
4038-
def set_deployment_config(self, config_name: str, instance_type: str) -> None:
4039-
"""Sets the deployment config to apply to the model."""
4199+
def set_deployment_config(
4200+
self, config_name: Optional[str] = None, instance_type: Optional[str] = None
4201+
) -> None:
4202+
"""Select the deployment config to apply to the model.
4203+
4204+
One API for both model types. For base/JumpStart models, ``config_name`` chooses a named
4205+
deployment config (with ``instance_type``). For fine-tuned models, the configs are largely
4206+
unnamed, so selection is by ``instance_type`` — pass the ``InstanceType`` of a config from
4207+
:meth:`list_deployment_configs`. The whole matching config (container image, environment,
4208+
and compute requirements) is applied at build/deploy time.
4209+
"""
4210+
if self._is_model_customization():
4211+
if not instance_type:
4212+
raise ValueError(
4213+
"Select a deployment config by instance_type for this model "
4214+
"(see list_deployment_configs())."
4215+
)
4216+
# Resolve the recipe's configs ONCE and derive both the match and the available-list
4217+
# from it (avoids re-resolving/re-fetching the hub document per branch).
4218+
configs = [
4219+
self._normalize_hosting_config(cfg)
4220+
for cfg in self._resolve_recipe_hosting_configs()
4221+
]
4222+
matches = [c for c in configs if self._config_instance_type(c) == instance_type]
4223+
if not matches:
4224+
available = sorted(
4225+
{
4226+
self._config_instance_type(c)
4227+
for c in configs
4228+
if self._config_instance_type(c)
4229+
}
4230+
)
4231+
raise ValueError(
4232+
f"No deployment config published for instance type '{instance_type}'. "
4233+
f"Available instance types: {available}"
4234+
)
4235+
if len(matches) > 1:
4236+
raise ValueError(
4237+
f"Instance type '{instance_type}' matches {len(matches)} deployment configs; "
4238+
"selection by instance type is ambiguous for this model."
4239+
)
4240+
self._selected_hosting_config = matches[0]
4241+
self.instance_type = instance_type
4242+
self._user_provided_instance_type = True
4243+
logger.info("Selected deployment config for instance type %s.", instance_type)
4244+
return
4245+
40404246
if not isinstance(self.model, str):
40414247
raise ValueError(
40424248
"Deployment config is only supported for JumpStart or HuggingFace models"
@@ -4045,6 +4251,9 @@ def set_deployment_config(self, config_name: str, instance_type: str) -> None:
40454251
if not (self._is_jumpstart_model_id() or self._use_jumpstart_equivalent()):
40464252
raise ValueError(f"The deployment config {config_name} cannot be set on this model")
40474253

4254+
if config_name is None:
4255+
raise ValueError("config_name is required to set a deployment config on this model.")
4256+
40484257
self.config_name = config_name
40494258
self.instance_type = instance_type
40504259

@@ -4068,7 +4277,24 @@ def set_deployment_config(self, config_name: str, instance_type: str) -> None:
40684277
feature=Feature.MODEL_CUSTOMIZATION, func_name="model_builder.get_deployment_config"
40694278
)
40704279
def get_deployment_config(self) -> Optional[Dict[str, Any]]:
4071-
"""Gets the deployment config to apply to the model."""
4280+
"""Get the deployment config that will be applied to the model.
4281+
4282+
One API for both model types. For fine-tuned models, returns the config chosen via
4283+
:meth:`set_deployment_config` if one was set, otherwise the config that would be selected
4284+
by default at build time (matching a provided instance type, else the default config); it
4285+
raises if the model publishes no hosting configs. NOTE: this differs from the base/
4286+
JumpStart path, which returns ``None`` until a config is explicitly set — the fine-tuned
4287+
path always has a well-defined default (its recipe's Default config), so it reports that.
4288+
"""
4289+
if self._is_model_customization():
4290+
if self._selected_hosting_config is not None:
4291+
return self._selected_hosting_config
4292+
# _resolve_recipe_hosting_configs raises when no configs are published, so a non-empty
4293+
# list is guaranteed here.
4294+
return self._normalize_hosting_config(
4295+
self._select_recipe_hosting_config(self._resolve_recipe_hosting_configs())
4296+
)
4297+
40724298
if not isinstance(self.model, str):
40734299
raise ValueError(
40744300
"Deployment config is only supported for JumpStart or HuggingFace models"
@@ -4092,19 +4318,59 @@ def get_deployment_config(self) -> Optional[Dict[str, Any]]:
40924318
@_telemetry_emitter(
40934319
feature=Feature.MODEL_CUSTOMIZATION, func_name="model_builder.list_deployment_configs"
40944320
)
4095-
def list_deployment_configs(self) -> List[Dict[str, Any]]:
4096-
"""List deployment configs for the model in the current region."""
4097-
if not isinstance(self.model, str):
4098-
raise ValueError(
4099-
"Deployment config is only supported for JumpStart or HuggingFace models"
4100-
)
4321+
def list_deployment_configs(
4322+
self, instance_type: Optional[str] = None
4323+
) -> List[Dict[str, Any]]:
4324+
"""List the deployment configs available for the model in the current region.
4325+
4326+
One API for both model types, returning the SAME dict shape so callers can iterate
4327+
results from either pathway identically: ``DeploymentConfigName`` plus a
4328+
``DeploymentArgs`` block (``ImageUri``, ``InstanceType``, ``Environment``,
4329+
``ComputeResourceRequirements``) and ``BenchmarkMetrics`` / ``AccelerationConfigs``. For
4330+
fine-tuned models the recipe's published configs are normalized into this shape;
4331+
``BenchmarkMetrics``/``AccelerationConfigs`` are empty (recipes don't publish them yet) and
4332+
an additive ``IsDefault`` flag marks the recipe's Default config. When ``instance_type`` is
4333+
provided, only configs offering that instance type are returned — the supported way to
4334+
narrow the alternatives before selecting.
4335+
"""
4336+
if self._is_model_customization():
4337+
configs = [
4338+
self._normalize_hosting_config(cfg)
4339+
for cfg in self._resolve_recipe_hosting_configs()
4340+
]
4341+
else:
4342+
if not isinstance(self.model, str):
4343+
raise ValueError(
4344+
"Deployment config is only supported for JumpStart or HuggingFace models"
4345+
)
41014346

4102-
if not (self._is_jumpstart_model_id() or self._use_jumpstart_equivalent()):
4103-
raise ValueError("Deployment config is only supported for JumpStart models")
4347+
if not (self._is_jumpstart_model_id() or self._use_jumpstart_equivalent()):
4348+
raise ValueError("Deployment config is only supported for JumpStart models")
4349+
4350+
configs = self.deployment_config_response_data(
4351+
self._get_deployment_configs(self.config_name, self.instance_type)
4352+
) # Delegate to JumpStart builder
41044353

4105-
return self.deployment_config_response_data(
4106-
self._get_deployment_configs(self.config_name, self.instance_type)
4107-
) # Delegate to JumpStart builder
4354+
if instance_type is not None:
4355+
configs = [
4356+
cfg for cfg in configs if self._config_instance_type(cfg) == instance_type
4357+
]
4358+
return configs
4359+
4360+
@staticmethod
4361+
def _config_instance_type(config: Dict[str, Any]) -> Optional[str]:
4362+
"""Extract the instance type from a normalized deployment config.
4363+
4364+
Both pathways nest the instance type under ``DeploymentArgs`` (fine-tuned configs are
4365+
normalized to the base shape by :meth:`_normalize_hosting_config`). A top-level
4366+
``InstanceType`` is accepted as a defensive fallback.
4367+
"""
4368+
deployment_args = config.get("DeploymentArgs") or {}
4369+
return (
4370+
deployment_args.get("InstanceType")
4371+
or deployment_args.get("instance_type")
4372+
or config.get("InstanceType")
4373+
)
41084374

41094375
@_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="model_builder.optimize")
41104376
# Add these methods to the current V3 ModelBuilder class:

0 commit comments

Comments
 (0)