1818import platform
1919import pytest
2020import shutil
21- import sys
2221import tempfile
2322
2423from sagemaker import LocalSession , Session
25- from sagemaker .pytorch import PyTorch
2624
27- from test . utils import image_utils
25+ from utils import image_utils
2826
2927logger = logging .getLogger (__name__ )
3028logging .getLogger ('boto' ).setLevel (logging .INFO )
4442
4543
4644def pytest_addoption (parser ):
47- parser .addoption ('--build-image' , '-D' , action = 'store_true' )
48- parser .addoption ('--build-base-image' , '-B' , action = 'store_true' )
49- parser .addoption ('--aws-id' )
45+ parser .addoption ('--build-image' , '-B' , action = 'store_true' )
46+ parser .addoption ('--push-image' , '-P' , action = 'store_true' )
47+ parser .addoption ('--dockerfile-type' , '-T' ,
48+ choices = ['dlc.cpu' , 'dlc.gpu' , 'dlc.eia' , 'pytorch' ],
49+ default = 'pytorch' )
50+ parser .addoption ('--dockerfile' , '-D' , default = None )
51+ parser .addoption ('--aws-id' , default = None )
5052 parser .addoption ('--instance-type' )
51- parser .addoption ('--docker-base-name' , default = 'pytorch' )
53+ parser .addoption ('--accelerator-type' )
54+ parser .addoption ('--docker-base-name' , default = 'sagemaker-pytorch-inference' )
5255 parser .addoption ('--region' , default = 'us-west-2' )
53- parser .addoption ('--framework-version' , default = PyTorch .LATEST_VERSION )
54- parser .addoption ('--py-version' , choices = ['2' , '3' ], default = str (sys .version_info .major ))
56+ parser .addoption ('--framework-version' , default = "1.4.0" )
57+ parser .addoption ('--py-version' , choices = ['2' , '3' ], default = '3' )
58+ # Processor is still "cpu" for EIA tests
5559 parser .addoption ('--processor' , choices = ['gpu' , 'cpu' ], default = 'cpu' )
5660 # If not specified, will default to {framework-version}-{processor}-py{py-version}
5761 parser .addoption ('--tag' , default = None )
5862
5963
64+ @pytest .fixture (scope = 'session' , name = 'dockerfile_type' )
65+ def fixture_dockerfile_type (request ):
66+ return request .config .getoption ('--dockerfile-type' )
67+
68+
69+ @pytest .fixture (scope = 'session' , name = 'dockerfile' )
70+ def fixture_dockerfile (request , dockerfile_type ):
71+ dockerfile = request .config .getoption ('--dockerfile' )
72+ return dockerfile if dockerfile else 'Dockerfile.{}' .format (dockerfile_type )
73+
74+
6075@pytest .fixture (scope = 'session' , name = 'docker_base_name' )
6176def fixture_docker_base_name (request ):
6277 return request .config .getoption ('--docker-base-name' )
@@ -89,11 +104,6 @@ def fixture_tag(request, framework_version, processor, py_version):
89104 return provided_tag if provided_tag else default_tag
90105
91106
92- @pytest .fixture (scope = 'session' , name = 'docker_image' )
93- def fixture_docker_image (docker_base_name , tag ):
94- return '{}:{}' .format (docker_base_name , tag )
95-
96-
97107@pytest .fixture
98108def opt_ml ():
99109 tmp = tempfile .mkdtemp ()
@@ -112,32 +122,25 @@ def fixture_use_gpu(processor):
112122 return processor == 'gpu'
113123
114124
115- @pytest .fixture (scope = 'session' , name = 'build_base_image' , autouse = True )
116- def fixture_build_base_image (request , framework_version , py_version , processor , tag , docker_base_name ):
117- build_base_image = request .config .getoption ('--build-base-image' )
118- if build_base_image :
119- return image_utils .build_base_image (framework_name = docker_base_name ,
120- framework_version = framework_version ,
121- py_version = py_version ,
122- base_image_tag = tag ,
123- processor = processor ,
124- cwd = os .path .join (dir_path , '..' ))
125-
126- return tag
127-
128-
129125@pytest .fixture (scope = 'session' , name = 'build_image' , autouse = True )
130- def fixture_build_image (request , framework_version , py_version , processor , tag , docker_base_name ):
126+ def fixture_build_image (request , framework_version , dockerfile , image_uri , region ):
131127 build_image = request .config .getoption ('--build-image' )
132128 if build_image :
133- return image_utils .build_image (framework_name = docker_base_name ,
134- framework_version = framework_version ,
135- py_version = py_version ,
136- processor = processor ,
137- tag = tag ,
129+ return image_utils .build_image (framework_version = framework_version ,
130+ dockerfile = dockerfile ,
131+ image_uri = image_uri ,
132+ region = region ,
138133 cwd = os .path .join (dir_path , '..' ))
139134
140- return tag
135+ return image_uri
136+
137+
138+ @pytest .fixture (scope = 'session' , name = 'push_image' , autouse = True )
139+ def fixture_push_image (request , image_uri , region , aws_id ):
140+ push_image = request .config .getoption ('--push-image' )
141+ if push_image :
142+ return image_utils .push_image (image_uri , region , aws_id )
143+ return None
141144
142145
143146@pytest .fixture (scope = 'session' , name = 'sagemaker_session' )
@@ -162,32 +165,80 @@ def fixture_instance_type(request, processor):
162165 return provided_instance_type or default_instance_type
163166
164167
168+ @pytest .fixture (name = 'accelerator_type' , scope = 'session' )
169+ def fixture_accelerator_type (request ):
170+ return request .config .getoption ('--accelerator-type' )
171+
172+
165173@pytest .fixture (name = 'docker_registry' , scope = 'session' )
166174def fixture_docker_registry (aws_id , region ):
167- return '{}.dkr.ecr.{}.amazonaws.com' .format (aws_id , region )
175+ return '{}.dkr.ecr.{}.amazonaws.com' .format (aws_id , region ) if aws_id else None
176+
177+
178+ @pytest .fixture (name = 'image_uri' , scope = 'session' )
179+ def fixture_image_uri (docker_registry , docker_base_name , tag ):
180+ if docker_registry :
181+ return '{}/{}:{}' .format (docker_registry , docker_base_name , tag )
182+ return '{}:{}' .format (docker_base_name , tag )
183+
184+
185+ @pytest .fixture (scope = 'session' , name = 'dist_cpu_backend' , params = ['gloo' ])
186+ def fixture_dist_cpu_backend (request ):
187+ return request .param
168188
169189
170- @pytest .fixture (name = 'ecr_image ' , scope = 'session' )
171- def fixture_ecr_image ( docker_registry , docker_base_name , tag ):
172- return '{}/{}:{}' . format ( docker_registry , docker_base_name , tag )
190+ @pytest .fixture (scope = 'session' , name = 'dist_gpu_backend ' , params = [ 'gloo' , 'nccl' ] )
191+ def fixture_dist_gpu_backend ( request ):
192+ return request . param
173193
174194
175195@pytest .fixture (autouse = True )
176- def skip_by_device_type (request , use_gpu , instance_type ):
196+ def skip_by_device_type (request , use_gpu , instance_type , accelerator_type ):
177197 is_gpu = use_gpu or instance_type [3 ] in ['g' , 'p' ]
178- if (request .node .get_closest_marker ('skip_gpu' ) and is_gpu ) or \
179- (request .node .get_closest_marker ('skip_cpu' ) and not is_gpu ):
198+ is_eia = accelerator_type is not None
199+
200+ # Separate out cases for clearer logic.
201+ # When running GPU test, skip CPU test. When running CPU test, skip GPU test.
202+ if (request .node .get_closest_marker ('gpu_test' ) and not is_gpu ) or \
203+ (request .node .get_closest_marker ('cpu_test' ) and is_gpu ):
204+ pytest .skip ('Skipping because running on \' {}\' instance' .format (instance_type ))
205+
206+ # When running EIA test, skip the CPU and GPU functions
207+ elif (request .node .get_closest_marker ('gpu_test' ) or request .node .get_closest_marker ('cpu_test' )) and is_eia :
208+ pytest .skip ('Skipping because running on \' {}\' instance' .format (instance_type ))
209+
210+ # When running CPU or GPU test, skip EIA test.
211+ elif request .node .get_closest_marker ('eia_test' ) and not is_eia :
180212 pytest .skip ('Skipping because running on \' {}\' instance' .format (instance_type ))
181213
182214
183215@pytest .fixture (autouse = True )
184216def skip_by_py_version (request , py_version ):
217+ """
218+ This will cause tests to be skipped w/ py3 containers if "py-version" flag is not set
219+ and pytest is running from py2. We can rely on this when py2 is deprecated, but for now
220+ we must use "skip_py2_containers"
221+ """
185222 if request .node .get_closest_marker ('skip_py2' ) and py_version != 'py3' :
186223 pytest .skip ('Skipping the test because Python 2 is not supported.' )
187224
188225
226+ @pytest .fixture (autouse = True )
227+ def skip_test_in_region (request , region ):
228+ if request .node .get_closest_marker ('skip_test_in_region' ):
229+ if region == 'me-south-1' :
230+ pytest .skip ('Skipping SageMaker test in region {}' .format (region ))
231+
232+
189233@pytest .fixture (autouse = True )
190234def skip_gpu_instance_restricted_regions (region , instance_type ):
191- if (region in NO_P2_REGIONS and instance_type .startswith ('ml.p2' )) \
192- or (region in NO_P3_REGIONS and instance_type .startswith ('ml.p3' )):
235+ if (( region in NO_P2_REGIONS and instance_type .startswith ('ml.p2' ))
236+ or (region in NO_P3_REGIONS and instance_type .startswith ('ml.p3' ))) :
193237 pytest .skip ('Skipping GPU test in region {}' .format (region ))
238+
239+
240+ @pytest .fixture (autouse = True )
241+ def skip_py2_containers (request , tag ):
242+ if request .node .get_closest_marker ('skip_py2_containers' ):
243+ if 'py2' in tag :
244+ pytest .skip ('Skipping python2 container with tag {}' .format (tag ))
0 commit comments