File tree 2 files changed +6
-0
lines changed
2 files changed +6
-0
lines changed Original file line number Diff line number Diff line change @@ -185,6 +185,7 @@ def __init__(
185
185
custom_attributes = None ,
186
186
accelerator_type = None ,
187
187
endpoint_name_prefix = None ,
188
+ target_model = None ,
188
189
):
189
190
r"""Initializes a configuration of a model and the endpoint to be created for it.
190
191
@@ -218,6 +219,8 @@ def __init__(
218
219
https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html.
219
220
endpoint_name_prefix (str): The endpoint name prefix of a new endpoint. Must follow
220
221
pattern "^[a-zA-Z0-9](-\*[a-zA-Z0-9]".
222
+ target_model (str): Sets the target model name when using a multi-model endpoint. For more information
223
+ about multi-model endpoints
221
224
"""
222
225
self .predictor_config = {
223
226
"model_name" : model_name ,
@@ -261,6 +264,7 @@ def __init__(
261
264
self .predictor_config ["content_template" ] = content_template
262
265
_set (custom_attributes , "custom_attributes" , self .predictor_config )
263
266
_set (accelerator_type , "accelerator_type" , self .predictor_config )
267
+ _set (target_model , "target_model" , self .predictor_config )
264
268
265
269
def get_predictor_config (self ):
266
270
"""Returns part of the predictor dictionary of the analysis config."""
Original file line number Diff line number Diff line change @@ -240,6 +240,7 @@ def test_model_config():
240
240
accept_type = "text/csv"
241
241
content_type = "application/jsonlines"
242
242
custom_attributes = "c000b4f9-df62-4c85-a0bf-7c525f9104a4"
243
+ target_model = "target_model_name"
243
244
accelerator_type = "ml.eia1.medium"
244
245
model_config = ModelConfig (
245
246
model_name = model_name ,
@@ -258,6 +259,7 @@ def test_model_config():
258
259
"content_type" : content_type ,
259
260
"custom_attributes" : custom_attributes ,
260
261
"accelerator_type" : accelerator_type ,
262
+ "target_model_name" : target_model ,
261
263
}
262
264
assert expected_config == model_config .get_predictor_config ()
263
265
You can’t perform that action at this time.
0 commit comments