Skip to content

Commit 864310a

Browse files
ajaykarpurajaykarpur
authored
fix: Append serving to model framework name for PyTorch, MXNet, and TensorFlow (#1247)
1 parent 26acec7 commit 864310a

File tree

5 files changed

+36
-38
lines changed

5 files changed

+36
-38
lines changed

‎src/sagemaker/mxnet/model.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,7 @@ def serving_image_uri(self, region_name, instance_type):
183183
"""
184184
return fw_utils.create_image_uri(
185185
region_name,
186-
self.__framework_name__,
186+
"-".join([self.__framework_name__, "serving"]),
187187
instance_type,
188188
self.framework_version,
189189
self.py_version,

‎src/sagemaker/pytorch/model.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,7 @@ def serving_image_uri(self, region_name, instance_type):
183183
"""
184184
return fw_utils.create_image_uri(
185185
region_name,
186-
self.__framework_name__,
186+
"-".join([self.__framework_name__, "serving"]),
187187
instance_type,
188188
self.framework_version,
189189
self.py_version,

‎src/sagemaker/tensorflow/model.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ def serving_image_uri(self, region_name, instance_type):
174174
"""
175175
return fw_utils.create_image_uri(
176176
region_name,
177-
self.__framework_name__,
177+
"-".join([self.__framework_name__, "serving"]),
178178
instance_type,
179179
self.framework_version,
180180
self.py_version,

‎tests/integ/test_airflow_config.py

Lines changed: 23 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -620,39 +620,37 @@ def test_pytorch_airflow_config_uploads_data_source_to_s3_when_inputs_not_provid
620620
hyperparameters={"epochs": 6, "backend": "gloo"},
621621
)
622622

623-
train_config = sm_airflow.training_config(estimator=estimator)
624-
625-
uploaded_s3_data = train_config["HyperParameters"]["sagemaker_submit_directory"].strip('"')
626-
627-
transform_config = sm_airflow.transform_config_from_estimator(
628-
estimator=estimator,
629-
task_id="transform_config",
630-
task_type="training",
631-
instance_count=SINGLE_INSTANCE_COUNT,
632-
instance_type=cpu_instance_type,
633-
data=uploaded_s3_data,
634-
content_type="text/csv",
623+
training_config = _build_airflow_workflow(
624+
estimator=estimator, instance_type=cpu_instance_type
635625
)
636626

637-
default_args = {
638-
"owner": "airflow",
639-
"start_date": airflow.utils.dates.days_ago(2),
640-
"provide_context": True,
641-
}
627+
_assert_that_s3_url_contains_data(
628+
sagemaker_session,
629+
training_config["HyperParameters"]["sagemaker_submit_directory"].strip('"'),
630+
)
642631

643-
dag = DAG("tensorflow_example", default_args=default_args, schedule_interval="@once")
644632

645-
train_op = SageMakerTrainingOperator(
646-
task_id="tf_training", config=train_config, wait_for_completion=True, dag=dag
633+
def test_pytorch_12_airflow_config_uploads_data_source_to_s3_when_inputs_not_provided(
634+
sagemaker_session, cpu_instance_type
635+
):
636+
with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
637+
estimator = PyTorch(
638+
entry_point=PYTORCH_MNIST_SCRIPT,
639+
role=ROLE,
640+
framework_version="1.2.0",
641+
train_instance_count=2,
642+
train_instance_type=cpu_instance_type,
643+
hyperparameters={"epochs": 6, "backend": "gloo"},
647644
)
648645

649-
transform_op = SageMakerTransformOperator(
650-
task_id="transform_operator", config=transform_config, wait_for_completion=True, dag=dag
646+
training_config = _build_airflow_workflow(
647+
estimator=estimator, instance_type=cpu_instance_type
651648
)
652649

653-
transform_op.set_upstream(train_op)
654-
655-
_assert_that_s3_url_contains_data(sagemaker_session, uploaded_s3_data)
650+
_assert_that_s3_url_contains_data(
651+
sagemaker_session,
652+
training_config["HyperParameters"]["sagemaker_submit_directory"].strip('"'),
653+
)
656654

657655

658656
def _assert_that_s3_url_contains_data(sagemaker_session, s3_url):

‎tests/unit/test_airflow.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1051,9 +1051,9 @@ def test_model_config_from_framework_estimator(sagemaker_session):
10511051
task_type="training",
10521052
)
10531053
expected_config = {
1054-
"ModelName": "sagemaker-mxnet-%s" % TIME_STAMP,
1054+
"ModelName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
10551055
"PrimaryContainer": {
1056-
"Image": "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.3.0-cpu-py3",
1056+
"Image": "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet-serving:1.3.0-cpu-py3",
10571057
"Environment": {
10581058
"SAGEMAKER_PROGRAM": "{{ entry_point }}",
10591059
"SAGEMAKER_SUBMIT_DIRECTORY": "s3://output/{{ ti.xcom_pull(task_ids='task_id')['Training']"
@@ -1209,9 +1209,9 @@ def test_transform_config_from_framework_estimator(sagemaker_session):
12091209
)
12101210
expected_config = {
12111211
"Model": {
1212-
"ModelName": "sagemaker-mxnet-%s" % TIME_STAMP,
1212+
"ModelName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
12131213
"PrimaryContainer": {
1214-
"Image": "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.3.0-gpu-py3",
1214+
"Image": "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet-serving:1.3.0-gpu-py3",
12151215
"Environment": {
12161216
"SAGEMAKER_PROGRAM": "{{ entry_point }}",
12171217
"SAGEMAKER_SUBMIT_DIRECTORY": "s3://output/{{ ti.xcom_pull(task_ids='task_id')"
@@ -1228,7 +1228,7 @@ def test_transform_config_from_framework_estimator(sagemaker_session):
12281228
},
12291229
"Transform": {
12301230
"TransformJobName": "{{ base_job_name }}-%s" % TIME_STAMP,
1231-
"ModelName": "sagemaker-mxnet-%s" % TIME_STAMP,
1231+
"ModelName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
12321232
"TransformInput": {
12331233
"DataSource": {
12341234
"S3DataSource": {"S3DataType": "S3Prefix", "S3Uri": "{{ transform_data }}"}
@@ -1449,9 +1449,9 @@ def test_deploy_config_from_framework_estimator(sagemaker_session):
14491449
)
14501450
expected_config = {
14511451
"Model": {
1452-
"ModelName": "sagemaker-mxnet-%s" % TIME_STAMP,
1452+
"ModelName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
14531453
"PrimaryContainer": {
1454-
"Image": "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.3.0-cpu-py3",
1454+
"Image": "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet-serving:1.3.0-cpu-py3",
14551455
"Environment": {
14561456
"SAGEMAKER_PROGRAM": "{{ entry_point }}",
14571457
"SAGEMAKER_SUBMIT_DIRECTORY": "s3://output/{{ ti.xcom_pull(task_ids='task_id')['Training']"
@@ -1466,20 +1466,20 @@ def test_deploy_config_from_framework_estimator(sagemaker_session):
14661466
"ExecutionRoleArn": "{{ role }}",
14671467
},
14681468
"EndpointConfig": {
1469-
"EndpointConfigName": "sagemaker-mxnet-%s" % TIME_STAMP,
1469+
"EndpointConfigName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
14701470
"ProductionVariants": [
14711471
{
14721472
"InstanceType": "ml.c4.large",
14731473
"InitialInstanceCount": "{{ instance_count}}",
1474-
"ModelName": "sagemaker-mxnet-%s" % TIME_STAMP,
1474+
"ModelName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
14751475
"VariantName": "AllTraffic",
14761476
"InitialVariantWeight": 1,
14771477
}
14781478
],
14791479
},
14801480
"Endpoint": {
14811481
"EndpointName": "mxnet-endpoint",
1482-
"EndpointConfigName": "sagemaker-mxnet-%s" % TIME_STAMP,
1482+
"EndpointConfigName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
14831483
},
14841484
}
14851485

0 commit comments

Comments
 (0)