Skip to content

Commit 15c5c6a

Browse files
feature: Allow users to send custom attributes to the model endpoint (#2198)
1 parent 334f942 commit 15c5c6a

File tree

2 files changed

+16
-0
lines changed

2 files changed

+16
-0
lines changed

src/sagemaker/clarify.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,7 @@ def __init__(
122122
accept_type=None,
123123
content_type=None,
124124
content_template=None,
125+
custom_attributes=None,
125126
):
126127
"""Initializes a configuration of a model and the endpoint to be created for it.
127128
@@ -141,6 +142,15 @@ def __init__(
141142
"application/jsonlines". The template should have one and only one placeholder
142143
$features which will be replaced by a features list for to form the model inference
143144
input.
145+
custom_attributes (str): Provides additional information about a request for an
146+
inference submitted to a model hosted at an Amazon SageMaker endpoint. The
147+
information is an opaque value that is forwarded verbatim. You could use this
148+
value, for example, to provide an ID that you can use to track a request or to
149+
provide other metadata that a service endpoint was programmed to process. The value
150+
must consist of no more than 1024 visible US-ASCII characters as specified in
151+
Section 3.3.6. Field Value Components (
152+
https://tools.ietf.org/html/rfc7230#section-3.2.6) of the Hypertext Transfer
153+
Protocol (HTTP/1.1).
144154
"""
145155
self.predictor_config = {
146156
"model_name": model_name,
@@ -169,6 +179,9 @@ def __init__(
169179
)
170180
self.predictor_config["content_template"] = content_template
171181

182+
if custom_attributes is not None:
183+
self.predictor_config["custom_attributes"] = custom_attributes
184+
172185
def get_predictor_config(self):
173186
"""Returns part of the predictor dictionary of the analysis config."""
174187
return copy.deepcopy(self.predictor_config)

tests/unit/test_clarify.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,19 +91,22 @@ def test_model_config():
9191
instance_count = 1
9292
accept_type = "text/csv"
9393
content_type = "application/jsonlines"
94+
custom_attributes = "c000b4f9-df62-4c85-a0bf-7c525f9104a4"
9495
model_config = ModelConfig(
9596
model_name=model_name,
9697
instance_type=instance_type,
9798
instance_count=instance_count,
9899
accept_type=accept_type,
99100
content_type=content_type,
101+
custom_attributes=custom_attributes,
100102
)
101103
expected_config = {
102104
"model_name": model_name,
103105
"instance_type": instance_type,
104106
"initial_instance_count": instance_count,
105107
"accept_type": accept_type,
106108
"content_type": content_type,
109+
"custom_attributes": custom_attributes,
107110
}
108111
assert expected_config == model_config.get_predictor_config()
109112

0 commit comments

Comments
 (0)