@@ -60,54 +60,52 @@ def tar_dir(directory, tmpdir):
60
60
61
61
62
62
@pytest .fixture
63
- def tfs_predictor_with_model_and_entry_point_same_tar (instance_type ,
64
- sagemaker_session ,
63
+ def tfs_predictor_with_model_and_entry_point_same_tar (sagemaker_local_session ,
65
64
tf_full_version ,
66
65
tmpdir ):
67
66
endpoint_name = sagemaker .utils .unique_name_from_base ('sagemaker-tensorflow-serving' )
68
67
69
68
model_tar = tar_dir (os .path .join (tests .integ .DATA_DIR , 'tfs/tfs-test-model-with-inference' ),
70
69
tmpdir )
71
70
72
- model_data = sagemaker_session .upload_data (
73
- path = model_tar ,
74
- key_prefix = 'tensorflow-serving/models' )
71
+ model = Model (model_data = 'file://' + model_tar ,
72
+ role = 'SageMakerRole' ,
73
+ framework_version = tf_full_version ,
74
+ sagemaker_session = sagemaker_local_session )
75
+ predictor = model .deploy (1 , 'local' , endpoint_name = endpoint_name )
75
76
76
- with tests .integ .timeout .timeout_and_delete_endpoint_by_name (endpoint_name ,
77
- sagemaker_session ):
78
- model = Model (model_data = model_data ,
79
- role = 'SageMakerRole' ,
80
- framework_version = tf_full_version ,
81
- sagemaker_session = sagemaker_session )
82
- predictor = model .deploy (1 , instance_type , endpoint_name = endpoint_name )
77
+ try :
83
78
yield predictor
79
+ finally :
80
+ predictor .delete_endpoint ()
84
81
85
82
86
83
@pytest .fixture (scope = 'module' )
87
- def tfs_predictor_with_model_and_entry_point_and_dependencies (instance_type ,
88
- sagemaker_session , tf_full_version ):
84
+ def tfs_predictor_with_model_and_entry_point_and_dependencies (sagemaker_local_session ,
85
+ tf_full_version ):
89
86
endpoint_name = sagemaker .utils .unique_name_from_base ('sagemaker-tensorflow-serving' )
90
87
91
- model_data = sagemaker_session .upload_data (
92
- path = os .path .join (tests .integ .DATA_DIR ,
93
- 'tensorflow-serving-test-model.tar.gz' ),
94
- key_prefix = 'tensorflow-serving/models' )
88
+ entry_point = os .path .join (tests .integ .DATA_DIR ,
89
+ 'tfs/tfs-test-entrypoint-and-dependencies/inference.py' )
90
+ dependencies = [os .path .join (tests .integ .DATA_DIR ,
91
+ 'tfs/tfs-test-entrypoint-and-dependencies/dependency.py' )]
92
+
93
+ model_data = 'file://' + os .path .join (tests .integ .DATA_DIR ,
94
+ 'tensorflow-serving-test-model.tar.gz' )
95
+
96
+ model = Model (entry_point = entry_point ,
97
+ model_data = model_data ,
98
+ role = 'SageMakerRole' ,
99
+ dependencies = dependencies ,
100
+ framework_version = tf_full_version ,
101
+ sagemaker_session = sagemaker_local_session )
102
+
103
+ predictor = model .deploy (1 , 'local' , endpoint_name = endpoint_name )
104
+ try :
95
105
96
- with tests .integ .timeout .timeout_and_delete_endpoint_by_name (endpoint_name ,
97
- sagemaker_session ):
98
- entry_point = os .path .join (tests .integ .DATA_DIR ,
99
- 'tfs/tfs-test-entrypoint-and-dependencies/inference.py' )
100
- dependencies = [os .path .join (tests .integ .DATA_DIR ,
101
- 'tfs/tfs-test-entrypoint-and-dependencies/dependency.py' )]
102
-
103
- model = Model (entry_point = entry_point ,
104
- model_data = model_data ,
105
- role = 'SageMakerRole' ,
106
- dependencies = dependencies ,
107
- framework_version = tf_full_version ,
108
- sagemaker_session = sagemaker_session )
109
- predictor = model .deploy (1 , instance_type , endpoint_name = endpoint_name )
110
106
yield predictor
107
+ finally :
108
+ predictor .delete_endpoint ()
111
109
112
110
113
111
@pytest .fixture (scope = 'module' )
0 commit comments