Skip to content

Commit 7304bfe

Browse files
beniericbenieric
authored
Merge branch 'master' into enable-github-actions-ci
2 parents 96fb7ed + 615a8ad commit 7304bfe

File tree

7 files changed

+259
-31
lines changed

7 files changed

+259
-31
lines changed

‎src/sagemaker/huggingface/llm_utils.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,8 @@ def get_huggingface_model_metadata(model_id: str, hf_hub_token: Optional[str] =
8181
Returns:
8282
dict: The model metadata retrieved with the HuggingFace API
8383
"""
84-
84+
if not model_id:
85+
raise ValueError("Model ID is empty. Please provide a valid Model ID.")
8586
hf_model_metadata_url = f"https://huggingface.co/api/models/{model_id}"
8687
hf_model_metadata_json = None
8788
try:

‎src/sagemaker/model.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -766,8 +766,8 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
766766

767767
def _script_mode_env_vars(self):
768768
"""Returns a mapping of environment variables for script mode execution"""
769-
script_name = None
770-
dir_name = None
769+
script_name = self.env.get(SCRIPT_PARAM_NAME.upper(), "")
770+
dir_name = self.env.get(DIR_PARAM_NAME.upper(), "")
771771
if self.uploaded_code:
772772
script_name = self.uploaded_code.script_name
773773
if self.repacked_model_data or self.enable_network_isolation():
@@ -783,8 +783,8 @@ def _script_mode_env_vars(self):
783783
else "file://" + self.source_dir
784784
)
785785
return {
786-
SCRIPT_PARAM_NAME.upper(): script_name or str(),
787-
DIR_PARAM_NAME.upper(): dir_name or str(),
786+
SCRIPT_PARAM_NAME.upper(): script_name,
787+
DIR_PARAM_NAME.upper(): dir_name,
788788
CONTAINER_LOG_LEVEL_PARAM_NAME.upper(): to_string(self.container_log_level),
789789
SAGEMAKER_REGION_PARAM_NAME.upper(): self.sagemaker_session.boto_region_name,
790790
}

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

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -124,8 +124,8 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
124124
into a stream. All translations between the server and the client are handled
125125
automatically with the specified input and output.
126126
model (Optional[Union[object, str]): Model object (with ``predict`` method to perform
127-
inference) or a HuggingFace/JumpStart Model ID. Either ``model`` or
128-
``inference_spec`` is required for the model builder to build the artifact.
127+
inference) or a HuggingFace/JumpStart Model ID. Either ``model`` or ``inference_spec``
128+
is required for the model builder to build the artifact.
129129
inference_spec (InferenceSpec): The inference spec file with your customized
130130
``invoke`` and ``load`` functions.
131131
image_uri (Optional[str]): The container image uri (which is derived from a
@@ -145,6 +145,8 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
145145
to the model server). Possible values for this argument are
146146
``TORCHSERVE``, ``MMS``, ``TENSORFLOW_SERVING``, ``DJL_SERVING``,
147147
``TRITON``, and``TGI``.
148+
model_metadata (Optional[Dict[str, Any]): Dictionary used to override the HuggingFace
149+
model metadata. Currently ``HF_TASK`` is overridable.
148150
"""
149151

150152
model_path: Optional[str] = field(
@@ -241,6 +243,10 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
241243
model_server: Optional[ModelServer] = field(
242244
default=None, metadata={"help": "Define the model server to deploy to."}
243245
)
246+
model_metadata: Optional[Dict[str, Any]] = field(
247+
default=None,
248+
metadata={"help": "Define the model metadata to override, currently supports `HF_TASK`"},
249+
)
244250

245251
def _build_validations(self):
246252
"""Placeholder docstring"""
@@ -616,6 +622,9 @@ def build( # pylint: disable=R0911
616622
self._is_custom_image_uri = self.image_uri is not None
617623

618624
if isinstance(self.model, str):
625+
model_task = None
626+
if self.model_metadata:
627+
model_task = self.model_metadata.get("HF_TASK")
619628
if self._is_jumpstart_model_id():
620629
return self._build_for_jumpstart()
621630
if self._is_djl(): # pylint: disable=R1705
@@ -625,10 +634,10 @@ def build( # pylint: disable=R0911
625634
self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
626635
)
627636

628-
model_task = hf_model_md.get("pipeline_tag")
629-
if self.schema_builder is None and model_task:
637+
if model_task is None:
638+
model_task = hf_model_md.get("pipeline_tag")
639+
if self.schema_builder is None and model_task is not None:
630640
self._schema_builder_init(model_task)
631-
632641
if model_task == "text-generation": # pylint: disable=R1705
633642
return self._build_for_tgi()
634643
elif self._can_fit_on_single_gpu():

‎src/sagemaker/serve/schema/task.json

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,28 @@
11
{
22
"fill-mask": {
3-
"sample_inputs": {
3+
"sample_inputs": {
44
"properties": {
55
"inputs": "Paris is the [MASK] of France.",
66
"parameters": {}
77
}
8-
},
9-
"sample_outputs": {
8+
},
9+
"sample_outputs": {
1010
"properties": [
1111
{
1212
"sequence": "Paris is the capital of France.",
1313
"score": 0.7
1414
}
1515
]
1616
}
17-
},
17+
},
1818
"question-answering": {
19-
"sample_inputs": {
19+
"sample_inputs": {
2020
"properties": {
2121
"context": "I have a German Shepherd dog, named Coco.",
2222
"question": "What is my dog's breed?"
2323
}
24-
},
25-
"sample_outputs": {
24+
},
25+
"sample_outputs": {
2626
"properties": [
2727
{
2828
"answer": "German Shepherd",
@@ -32,36 +32,36 @@
3232
}
3333
]
3434
}
35-
},
35+
},
3636
"text-classification": {
37-
"sample_inputs": {
37+
"sample_inputs": {
3838
"properties": {
3939
"inputs": "Where is the capital of France?, Paris is the capital of France.",
4040
"parameters": {}
4141
}
42-
},
43-
"sample_outputs": {
42+
},
43+
"sample_outputs": {
4444
"properties": [
4545
{
4646
"label": "entailment",
4747
"score": 0.997
4848
}
4949
]
5050
}
51-
},
52-
"text-generation": {
53-
"sample_inputs": {
51+
},
52+
"text-generation": {
53+
"sample_inputs": {
5454
"properties": {
5555
"inputs": "Hello, I'm a language model",
5656
"parameters": {}
5757
}
58-
},
59-
"sample_outputs": {
58+
},
59+
"sample_outputs": {
6060
"properties": [
61-
{
62-
"generated_text": "Hello, I'm a language modeler. So while writing this, when I went out to meet my wife or come home she told me that my"
63-
}
61+
{
62+
"generated_text": "Hello, I'm a language modeler. So while writing this, when I went out to meet my wife or come home she told me that my"
63+
}
6464
]
6565
}
66-
}
66+
}
6767
}

‎tests/integ/sagemaker/serve/test_schema_builder.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,3 +99,69 @@ def test_model_builder_negative_path(sagemaker_session):
9999
match="Error Message: Schema builder for text-to-image could not be found.",
100100
):
101101
model_builder.build(sagemaker_session=sagemaker_session)
102+
103+
104+
@pytest.mark.skipif(
105+
PYTHON_VERSION_IS_NOT_310,
106+
reason="Testing Schema Builder Simplification feature",
107+
)
108+
@pytest.mark.parametrize(
109+
"model_id, task_provided",
110+
[
111+
("bert-base-uncased", "fill-mask"),
112+
("bert-large-uncased-whole-word-masking-finetuned-squad", "question-answering"),
113+
],
114+
)
115+
def test_model_builder_happy_path_with_task_provided(
116+
model_id, task_provided, sagemaker_session, gpu_instance_type
117+
):
118+
model_builder = ModelBuilder(model=model_id, model_metadata={"HF_TASK": task_provided})
119+
120+
model = model_builder.build(sagemaker_session=sagemaker_session)
121+
122+
assert model is not None
123+
assert model_builder.schema_builder is not None
124+
125+
inputs, outputs = task.retrieve_local_schemas(task_provided)
126+
assert model_builder.schema_builder.sample_input == inputs
127+
assert model_builder.schema_builder.sample_output == outputs
128+
129+
with timeout(minutes=SERVE_SAGEMAKER_ENDPOINT_TIMEOUT):
130+
caught_ex = None
131+
try:
132+
iam_client = sagemaker_session.boto_session.client("iam")
133+
role_arn = iam_client.get_role(RoleName="SageMakerRole")["Role"]["Arn"]
134+
135+
logger.info("Deploying and predicting in SAGEMAKER_ENDPOINT mode...")
136+
predictor = model.deploy(
137+
role=role_arn, instance_count=1, instance_type=gpu_instance_type
138+
)
139+
140+
predicted_outputs = predictor.predict(inputs)
141+
assert predicted_outputs is not None
142+
143+
except Exception as e:
144+
caught_ex = e
145+
finally:
146+
cleanup_model_resources(
147+
sagemaker_session=model_builder.sagemaker_session,
148+
model_name=model.name,
149+
endpoint_name=model.endpoint_name,
150+
)
151+
if caught_ex:
152+
logger.exception(caught_ex)
153+
assert (
154+
False
155+
), f"{caught_ex} was thrown when running transformers sagemaker endpoint test"
156+
157+
158+
def test_model_builder_negative_path_with_invalid_task(sagemaker_session):
159+
model_builder = ModelBuilder(
160+
model="bert-base-uncased", model_metadata={"HF_TASK": "invalid-task"}
161+
)
162+
163+
with pytest.raises(
164+
TaskNotFoundException,
165+
match="Error Message: Schema builder for invalid-task could not be found.",
166+
):
167+
model_builder.build(sagemaker_session=sagemaker_session)

‎tests/unit/sagemaker/serve/builder/test_model_builder.py

Lines changed: 124 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1076,7 +1076,7 @@ def test_build_negative_path_when_schema_builder_not_present(
10761076

10771077
model_builder = ModelBuilder(model="CompVis/stable-diffusion-v1-4")
10781078

1079-
self.assertRaisesRegexp(
1079+
self.assertRaisesRegex(
10801080
TaskNotFoundException,
10811081
"Error Message: Schema builder for text-to-image could not be found.",
10821082
lambda: model_builder.build(sagemaker_session=mock_session),
@@ -1593,3 +1593,126 @@ def test_total_inference_model_size_mib_throws(
15931593
model_builder.build(sagemaker_session=mock_session)
15941594

15951595
self.assertEqual(model_builder._can_fit_on_single_gpu(), False)
1596+
1597+
@patch("sagemaker.serve.builder.tgi_builder.HuggingFaceModel")
1598+
@patch("sagemaker.image_uris.retrieve")
1599+
@patch("sagemaker.djl_inference.model.urllib")
1600+
@patch("sagemaker.djl_inference.model.json")
1601+
@patch("sagemaker.huggingface.llm_utils.urllib")
1602+
@patch("sagemaker.huggingface.llm_utils.json")
1603+
@patch("sagemaker.model_uris.retrieve")
1604+
@patch("sagemaker.serve.builder.model_builder._ServeSettings")
1605+
def test_build_happy_path_override_with_task_provided(
1606+
self,
1607+
mock_serveSettings,
1608+
mock_model_uris_retrieve,
1609+
mock_llm_utils_json,
1610+
mock_llm_utils_urllib,
1611+
mock_model_json,
1612+
mock_model_urllib,
1613+
mock_image_uris_retrieve,
1614+
mock_hf_model,
1615+
):
1616+
# Setup mocks
1617+
1618+
mock_setting_object = mock_serveSettings.return_value
1619+
mock_setting_object.role_arn = mock_role_arn
1620+
mock_setting_object.s3_model_data_url = mock_s3_model_data_url
1621+
1622+
# HF Pipeline Tag
1623+
mock_model_uris_retrieve.side_effect = KeyError
1624+
mock_llm_utils_json.load.return_value = {"pipeline_tag": "fill-mask"}
1625+
mock_llm_utils_urllib.request.Request.side_effect = Mock()
1626+
1627+
# HF Model config
1628+
mock_model_json.load.return_value = {"some": "config"}
1629+
mock_model_urllib.request.Request.side_effect = Mock()
1630+
1631+
mock_image_uris_retrieve.return_value = "https://some-image-uri"
1632+
1633+
model_builder = ModelBuilder(
1634+
model="bert-base-uncased", model_metadata={"HF_TASK": "text-generation"}
1635+
)
1636+
model_builder.build(sagemaker_session=mock_session)
1637+
1638+
self.assertIsNotNone(model_builder.schema_builder)
1639+
sample_inputs, sample_outputs = task.retrieve_local_schemas("text-generation")
1640+
self.assertEqual(
1641+
sample_inputs["inputs"], model_builder.schema_builder.sample_input["inputs"]
1642+
)
1643+
self.assertEqual(sample_outputs, model_builder.schema_builder.sample_output)
1644+
1645+
@patch("sagemaker.image_uris.retrieve")
1646+
@patch("sagemaker.djl_inference.model.urllib")
1647+
@patch("sagemaker.djl_inference.model.json")
1648+
@patch("sagemaker.huggingface.llm_utils.urllib")
1649+
@patch("sagemaker.huggingface.llm_utils.json")
1650+
@patch("sagemaker.model_uris.retrieve")
1651+
@patch("sagemaker.serve.builder.model_builder._ServeSettings")
1652+
def test_build_task_override_with_invalid_task_provided(
1653+
self,
1654+
mock_serveSettings,
1655+
mock_model_uris_retrieve,
1656+
mock_llm_utils_json,
1657+
mock_llm_utils_urllib,
1658+
mock_model_json,
1659+
mock_model_urllib,
1660+
mock_image_uris_retrieve,
1661+
):
1662+
# Setup mocks
1663+
1664+
mock_setting_object = mock_serveSettings.return_value
1665+
mock_setting_object.role_arn = mock_role_arn
1666+
mock_setting_object.s3_model_data_url = mock_s3_model_data_url
1667+
1668+
# HF Pipeline Tag
1669+
mock_model_uris_retrieve.side_effect = KeyError
1670+
mock_llm_utils_json.load.return_value = {"pipeline_tag": "fill-mask"}
1671+
mock_llm_utils_urllib.request.Request.side_effect = Mock()
1672+
1673+
# HF Model config
1674+
mock_model_json.load.return_value = {"some": "config"}
1675+
mock_model_urllib.request.Request.side_effect = Mock()
1676+
1677+
mock_image_uris_retrieve.return_value = "https://some-image-uri"
1678+
model_ids_with_invalid_task = {
1679+
"bert-base-uncased": "invalid-task",
1680+
"bert-large-uncased-whole-word-masking-finetuned-squad": "",
1681+
}
1682+
for model_id in model_ids_with_invalid_task:
1683+
provided_task = model_ids_with_invalid_task[model_id]
1684+
model_builder = ModelBuilder(model=model_id, model_metadata={"HF_TASK": provided_task})
1685+
1686+
self.assertRaisesRegex(
1687+
TaskNotFoundException,
1688+
f"Error Message: Schema builder for {provided_task} could not be found.",
1689+
lambda: model_builder.build(sagemaker_session=mock_session),
1690+
)
1691+
1692+
@patch("sagemaker.image_uris.retrieve")
1693+
@patch("sagemaker.model_uris.retrieve")
1694+
@patch("sagemaker.serve.builder.model_builder._ServeSettings")
1695+
def test_build_task_override_with_invalid_model_provided(
1696+
self,
1697+
mock_serveSettings,
1698+
mock_model_uris_retrieve,
1699+
mock_image_uris_retrieve,
1700+
):
1701+
# Setup mocks
1702+
1703+
mock_setting_object = mock_serveSettings.return_value
1704+
mock_setting_object.role_arn = mock_role_arn
1705+
mock_setting_object.s3_model_data_url = mock_s3_model_data_url
1706+
1707+
# HF Pipeline Tag
1708+
mock_model_uris_retrieve.side_effect = KeyError
1709+
1710+
mock_image_uris_retrieve.return_value = "https://some-image-uri"
1711+
invalid_model_id = ""
1712+
provided_task = "fill-mask"
1713+
1714+
model_builder = ModelBuilder(
1715+
model=invalid_model_id, model_metadata={"HF_TASK": provided_task}
1716+
)
1717+
with self.assertRaises(Exception):
1718+
model_builder.build(sagemaker_session=mock_session)

0 commit comments

Comments
 (0)