11
11
# ANY KIND, either express or implied. See the License for the specific
12
12
# language governing permissions and limitations under the License.
13
13
from __future__ import absolute_import
14
+ from __future__ import absolute_import
14
15
15
16
import numpy as np
16
17
import os
22
23
from sagemaker .tensorflow import TensorFlow
23
24
from six .moves .urllib .parse import urlparse
24
25
from sagemaker .utils import unique_name_from_base
25
- import tests .integ as integ
26
- from tests .integ import kms_utils
27
- import tests .integ .timeout as timeout
26
+
27
+ import tests .integ
28
28
29
29
ROLE = 'SageMakerRole'
30
30
35
35
TAGS = [{'Key' : 'some-key' , 'Value' : 'some-value' }]
36
36
37
37
38
- @pytest .fixture (scope = 'session' , params = ['ml.c5.xlarge' , 'ml.p2.xlarge' ])
38
+ @pytest .fixture (scope = 'session' , params = [
39
+ 'ml.c5.xlarge' ,
40
+ pytest .param ('ml.p2.xlarge' ,
41
+ marks = pytest .mark .skipif (
42
+ tests .integ .test_region () in tests .integ .HOSTING_NO_P2_REGIONS ,
43
+ reason = 'no ml.p2 instances in this region' ))])
39
44
def instance_type (request ):
40
45
return request .param
41
46
42
47
43
- @pytest .mark .skipif (integ .test_region () in integ .HOSTING_NO_P2_REGIONS ,
44
- reason = 'no ml.p2 instances in these regions' )
45
- @pytest .mark .skipif (integ .PYTHON_VERSION != 'py3' , reason = "Script Mode tests are only configured to run with Python 3" )
48
+ @pytest .mark .skipif (tests .integ .PYTHON_VERSION != 'py3' ,
49
+ reason = "Script Mode tests are only configured to run with Python 3" )
46
50
def test_mnist (sagemaker_session , instance_type ):
47
51
estimator = TensorFlow (entry_point = SCRIPT ,
48
52
role = 'SageMakerRole' ,
@@ -51,26 +55,26 @@ def test_mnist(sagemaker_session, instance_type):
51
55
sagemaker_session = sagemaker_session ,
52
56
py_version = 'py3' ,
53
57
framework_version = TensorFlow .LATEST_VERSION ,
54
- metric_definitions = [{'Name' : 'train:global_steps' , 'Regex' : r'global_step\/sec:\s(.*)' }])
58
+ metric_definitions = [
59
+ {'Name' : 'train:global_steps' , 'Regex' : r'global_step\/sec:\s(.*)' }])
55
60
inputs = estimator .sagemaker_session .upload_data (
56
61
path = os .path .join (RESOURCE_PATH , 'data' ),
57
62
key_prefix = 'scriptmode/mnist' )
58
63
59
- with timeout .timeout (minutes = integ .TRAINING_DEFAULT_TIMEOUT_MINUTES ):
64
+ with tests . integ . timeout .timeout (minutes = tests . integ .TRAINING_DEFAULT_TIMEOUT_MINUTES ):
60
65
estimator .fit (inputs = inputs , job_name = unique_name_from_base ('test-tf-sm-mnist' ))
61
66
_assert_s3_files_exist (estimator .model_dir ,
62
67
['graph.pbtxt' , 'model.ckpt-0.index' , 'model.ckpt-0.meta' ])
63
68
df = estimator .training_job_analytics .dataframe ()
64
- print (df )
65
69
assert df .size > 0
66
70
67
71
68
72
def test_server_side_encryption (sagemaker_session ):
69
-
70
73
boto_session = sagemaker_session .boto_session
71
- with kms_utils .bucket_with_encryption (boto_session , ROLE ) as (bucket_with_kms , kms_key ):
72
-
73
- output_path = os .path .join (bucket_with_kms , 'test-server-side-encryption' , time .strftime ('%y%m%d-%H%M' ))
74
+ with tests .integ .kms_utils .bucket_with_encryption (boto_session , ROLE ) as (
75
+ bucket_with_kms , kms_key ):
76
+ output_path = os .path .join (bucket_with_kms , 'test-server-side-encryption' ,
77
+ time .strftime ('%y%m%d-%H%M' ))
74
78
75
79
estimator = TensorFlow (entry_point = SCRIPT ,
76
80
role = ROLE ,
@@ -88,28 +92,29 @@ def test_server_side_encryption(sagemaker_session):
88
92
path = os .path .join (RESOURCE_PATH , 'data' ),
89
93
key_prefix = 'scriptmode/mnist' )
90
94
91
- with timeout .timeout (minutes = integ .TRAINING_DEFAULT_TIMEOUT_MINUTES ):
92
- estimator .fit (inputs = inputs , job_name = unique_name_from_base ('test-server-side-encryption' ))
95
+ with tests .integ .timeout .timeout (minutes = tests .integ .TRAINING_DEFAULT_TIMEOUT_MINUTES ):
96
+ estimator .fit (inputs = inputs ,
97
+ job_name = unique_name_from_base ('test-server-side-encryption' ))
93
98
94
99
95
100
@pytest .mark .canary_quick
96
- @pytest .mark .skipif (integ .PYTHON_VERSION != 'py3' , reason = "Script Mode tests are only configured to run with Python 3" )
101
+ @pytest .mark .skipif (tests .integ .PYTHON_VERSION != 'py3' ,
102
+ reason = "Script Mode tests are only configured to run with Python 3" )
97
103
def test_mnist_distributed (sagemaker_session , instance_type ):
98
104
estimator = TensorFlow (entry_point = SCRIPT ,
99
105
role = ROLE ,
100
106
train_instance_count = 2 ,
101
- # TODO: change train_instance_type to instance_type once the test is passing consistently
102
- train_instance_type = 'ml.c5.xlarge' ,
107
+ train_instance_type = instance_type ,
103
108
sagemaker_session = sagemaker_session ,
104
- py_version = integ .PYTHON_VERSION ,
109
+ py_version = tests . integ .PYTHON_VERSION ,
105
110
script_mode = True ,
106
111
framework_version = TensorFlow .LATEST_VERSION ,
107
112
distributions = PARAMETER_SERVER_DISTRIBUTION )
108
113
inputs = estimator .sagemaker_session .upload_data (
109
114
path = os .path .join (RESOURCE_PATH , 'data' ),
110
115
key_prefix = 'scriptmode/distributed_mnist' )
111
116
112
- with timeout .timeout (minutes = integ .TRAINING_DEFAULT_TIMEOUT_MINUTES ):
117
+ with tests . integ . timeout .timeout (minutes = tests . integ .TRAINING_DEFAULT_TIMEOUT_MINUTES ):
113
118
estimator .fit (inputs = inputs , job_name = unique_name_from_base ('test-tf-sm-distributed' ))
114
119
_assert_s3_files_exist (estimator .model_dir ,
115
120
['graph.pbtxt' , 'model.ckpt-0.index' , 'model.ckpt-0.meta' ])
@@ -131,22 +136,26 @@ def test_mnist_async(sagemaker_session):
131
136
training_job_name = estimator .latest_training_job .name
132
137
time .sleep (20 )
133
138
endpoint_name = training_job_name
134
- _assert_training_job_tags_match (sagemaker_session .sagemaker_client , estimator .latest_training_job .name , TAGS )
135
- with timeout .timeout_and_delete_endpoint_by_name (endpoint_name , sagemaker_session ):
136
- estimator = TensorFlow .attach (training_job_name = training_job_name , sagemaker_session = sagemaker_session )
139
+ _assert_training_job_tags_match (sagemaker_session .sagemaker_client ,
140
+ estimator .latest_training_job .name , TAGS )
141
+ with tests .integ .timeout .timeout_and_delete_endpoint_by_name (endpoint_name , sagemaker_session ):
142
+ estimator = TensorFlow .attach (training_job_name = training_job_name ,
143
+ sagemaker_session = sagemaker_session )
137
144
predictor = estimator .deploy (initial_instance_count = 1 , instance_type = 'ml.c4.xlarge' ,
138
145
endpoint_name = endpoint_name )
139
146
140
147
result = predictor .predict (np .zeros (784 ))
141
148
print ('predict result: {}' .format (result ))
142
149
_assert_endpoint_tags_match (sagemaker_session .sagemaker_client , predictor .endpoint , TAGS )
143
- _assert_model_tags_match (sagemaker_session .sagemaker_client , estimator .latest_training_job .name , TAGS )
150
+ _assert_model_tags_match (sagemaker_session .sagemaker_client ,
151
+ estimator .latest_training_job .name , TAGS )
144
152
145
153
146
154
def _assert_s3_files_exist (s3_url , files ):
147
155
parsed_url = urlparse (s3_url )
148
156
s3 = boto3 .client ('s3' )
149
- contents = s3 .list_objects_v2 (Bucket = parsed_url .netloc , Prefix = parsed_url .path .lstrip ('/' ))["Contents" ]
157
+ contents = s3 .list_objects_v2 (Bucket = parsed_url .netloc , Prefix = parsed_url .path .lstrip ('/' ))[
158
+ "Contents" ]
150
159
for f in files :
151
160
found = [x ['Key' ] for x in contents if x ['Key' ].endswith (f )]
152
161
if not found :
@@ -169,5 +178,6 @@ def _assert_endpoint_tags_match(sagemaker_client, endpoint_name, tags):
169
178
170
179
171
180
def _assert_training_job_tags_match (sagemaker_client , training_job_name , tags ):
172
- training_job_description = sagemaker_client .describe_training_job (TrainingJobName = training_job_name )
181
+ training_job_description = sagemaker_client .describe_training_job (
182
+ TrainingJobName = training_job_name )
173
183
_assert_tags_match (sagemaker_client , training_job_description ['TrainingJobArn' ], tags )
0 commit comments