Skip to content

Commit 99b7b48

Browse files
authored
change: enable Neo integ tests (aws#1309)
This change also fixes the logic for generating a model name when compiling a model if self.name is not already set.
1 parent ab58a4c commit 99b7b48

File tree

4 files changed

+22
-35
lines changed

4 files changed

+22
-35
lines changed

src/sagemaker/model.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -304,8 +304,8 @@ def compile(
304304
https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html.
305305
input_shape (dict): Specifies the name and shape of the expected
306306
inputs for your trained model in json dictionary form, for
307-
example: {'data':[1,3,1024,1024]}, or {'var1': [1,1,28,28],
308-
'var2':[1,1,28,28]}
307+
example: {'data': [1,3,1024,1024]}, or {'var1': [1,1,28,28],
308+
'var2': [1,1,28,28]}
309309
output_path (str): Specifies where to store the compiled model
310310
role (str): Execution role
311311
tags (list[dict]): List of tags for labeling a compilation job. For
@@ -436,7 +436,8 @@ def deploy(
436436

437437
compiled_model_suffix = "-".join(instance_type.split(".")[:-1])
438438
if self._is_compiled_model:
439-
self.name += compiled_model_suffix
439+
name_prefix = self.name or utils.name_from_image(self.image)
440+
self.name = "{}{}".format(name_prefix, compiled_model_suffix)
440441

441442
self._create_sagemaker_model(instance_type, accelerator_type, tags)
442443
production_variant = sagemaker.production_variant(

tests/conftest.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -286,7 +286,7 @@ def alternative_cpu_instance_type(sagemaker_session, request):
286286

287287
@pytest.fixture(scope="session")
288288
def cpu_instance_family(cpu_instance_type):
289-
"_".join(cpu_instance_type.split(".")[0:2])
289+
return "_".join(cpu_instance_type.split(".")[0:2])
290290

291291

292292
def pytest_generate_tests(metafunc):

tests/data/mxnet_mnist/mnist_neo.py

Lines changed: 5 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,6 @@
2020
import mxnet as mx
2121
import numpy as np
2222

23-
from sagemaker_mxnet_container.training_utils import scheduler_host
24-
2523

2624
def load_data(path):
2725
with gzip.open(find_file(path, "labels.gz")) as flbl:
@@ -112,11 +110,10 @@ def neo_preprocess(payload, content_type):
112110
if content_type != "application/vnd+python.numpy+binary":
113111
raise RuntimeError("Content type must be application/vnd+python.numpy+binary")
114112

115-
f = io.BytesIO(payload)
116-
return np.load(f)
113+
return np.asarray(json.loads(payload.decode("utf-8")))
117114

118115

119-
### NOTE: this function cannot use MXNet
116+
# NOTE: this function cannot use MXNet
120117
def neo_postprocess(result):
121118
logging.info("Invoking user-defined post-processing function")
122119

@@ -131,19 +128,10 @@ def neo_postprocess(result):
131128
return response_body, content_type
132129

133130

134-
def save(model_dir, model):
135-
model.symbol.save(os.path.join(model_dir, "model-symbol.json"))
136-
model.save_params(os.path.join(model_dir, "model-0000.params"))
137-
138-
signature = [
139-
{"name": data_desc.name, "shape": [dim for dim in data_desc.shape]}
140-
for data_desc in model.data_shapes
141-
]
142-
with open(os.path.join(model_dir, "model-shapes.json"), "w") as f:
143-
json.dump(signature, f)
144-
145-
146131
if __name__ == "__main__":
132+
# Import here to prevent import during serving
133+
from sagemaker_mxnet_container.training_utils import scheduler_host, save
134+
147135
parser = argparse.ArgumentParser()
148136

149137
parser.add_argument("--batch-size", type=int, default=100)

tests/integ/test_neo_mxnet.py

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,27 +12,30 @@
1212
# language governing permissions and limitations under the License.
1313
from __future__ import absolute_import
1414

15-
import numpy
1615
import os
16+
17+
import numpy
1718
import pytest
19+
1820
from sagemaker.mxnet.estimator import MXNet
1921
from sagemaker.mxnet.model import MXNetModel
20-
from sagemaker.utils import sagemaker_timestamp
22+
from sagemaker.utils import unique_name_from_base
2123
from tests.integ import DATA_DIR, PYTHON_VERSION, TRAINING_DEFAULT_TIMEOUT_MINUTES
2224
from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
23-
import time
25+
26+
NEO_MXNET_VERSION = "1.4.1" # Neo doesn't support MXNet 1.6 yet.
2427

2528

2629
@pytest.fixture(scope="module")
27-
def mxnet_training_job(sagemaker_session, mxnet_full_version, cpu_instance_type):
30+
def mxnet_training_job(sagemaker_session, cpu_instance_type):
2831
with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
2932
script_path = os.path.join(DATA_DIR, "mxnet_mnist", "mnist_neo.py")
3033
data_path = os.path.join(DATA_DIR, "mxnet_mnist")
3134

3235
mx = MXNet(
3336
entry_point=script_path,
3437
role="SageMakerRole",
35-
framework_version=mxnet_full_version,
38+
framework_version=NEO_MXNET_VERSION,
3639
py_version=PYTHON_VERSION,
3740
train_instance_count=1,
3841
train_instance_type=cpu_instance_type,
@@ -52,13 +55,10 @@ def mxnet_training_job(sagemaker_session, mxnet_full_version, cpu_instance_type)
5255

5356
@pytest.mark.canary_quick
5457
@pytest.mark.regional_testing
55-
@pytest.mark.skip(
56-
reason="This should be enabled along with the Boto SDK release for Neo API changes"
57-
)
5858
def test_attach_deploy(
5959
mxnet_training_job, sagemaker_session, cpu_instance_type, cpu_instance_family
6060
):
61-
endpoint_name = "test-mxnet-attach-deploy-{}".format(sagemaker_timestamp())
61+
endpoint_name = unique_name_from_base("test-neo-attach-deploy")
6262

6363
with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
6464
estimator = MXNet.attach(mxnet_training_job, sagemaker_session=sagemaker_session)
@@ -77,13 +77,10 @@ def test_attach_deploy(
7777
predictor.predict(data)
7878

7979

80-
@pytest.mark.skip(
81-
reason="This should be enabled along with the Boto SDK release for Neo API changes"
82-
)
8380
def test_deploy_model(
8481
mxnet_training_job, sagemaker_session, cpu_instance_type, cpu_instance_family
8582
):
86-
endpoint_name = "test-mxnet-deploy-model-{}".format(sagemaker_timestamp())
83+
endpoint_name = unique_name_from_base("test-neo-deploy-model")
8784

8885
with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
8986
desc = sagemaker_session.sagemaker_client.describe_training_job(
@@ -97,14 +94,15 @@ def test_deploy_model(
9794
role,
9895
entry_point=script_path,
9996
py_version=PYTHON_VERSION,
97+
framework_version=NEO_MXNET_VERSION,
10098
sagemaker_session=sagemaker_session,
10199
)
102100

103101
model.compile(
104102
target_instance_family=cpu_instance_family,
105103
input_shape={"data": [1, 1, 28, 28]},
106104
role=role,
107-
job_name="test-deploy-model-compilation-job-{}".format(int(time.time())),
105+
job_name=unique_name_from_base("test-deploy-model-compilation-job"),
108106
output_path="/".join(model_data.split("/")[:-1]),
109107
)
110108
predictor = model.deploy(1, cpu_instance_type, endpoint_name=endpoint_name)

0 commit comments

Comments
 (0)