feature: Add more flexibility to docker based tests so that we can pa…

…ss different dockerfiles for different image types (aws#65)

Authored-by: Balaji Sankar <[email protected]>

pytest.ini

@@ -1,6 +1,7 @@
 [pytest]
 markers =
-    slow: For running tests on docker images to verify package installations.
+    gpu: For running tests on GPU docker images to verify package installations.
+    cpu: For running tests on CPU docker images to verify package installations.
     unit: For running unit tests
 pythonpath = src/
 addopts = -n 2

src/main.py

@@ -160,7 +160,8 @@ def _push_images_upstream(image_versions_to_push: list[dict[str, str]], region:
 def _test_local_images(image_ids_to_test: list[str]):
     assert len(image_ids_to_test) == len(_image_generator_configs)
     for (image_id, config) in zip(image_ids_to_test, _image_generator_configs):
-        exit_code = pytest.main(['-n', '2', '-m', 'slow', '--local-image-id', image_id, *config['pytest_flags']])
+        exit_code = pytest.main(['-n', '2', '-m', config['image_type'], '--local-image-id',
+                                 image_id, *config['pytest_flags']])
 
         assert exit_code == 0, f'Tests failed with exit code: {exit_code} against: {image_id}'
 

test/test_dockerfile_based_harness.py

@@ -3,12 +3,25 @@
 import docker
 import pytest
 from docker.errors import BuildError, ContainerError
-pytestmark = pytest.mark.slow
-
 
 _docker_client = docker.from_env()
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("dockerfile_path", ["keras.test.Dockerfile",
+                                             "matplotlib.test.Dockerfile",
+                                             "scipy.test.Dockerfile",
+                                             "numpy.test.Dockerfile",
+                                             "boto3.test.Dockerfile",
+                                             "pandas.test.Dockerfile",
+                                             "sm-python-sdk.test.Dockerfile",
+                                             "pytorch.examples.Dockerfile",
+                                             "tensorflow.examples.Dockerfile"])
+def test_dockerfiles_for_cpu(dockerfile_path: str, local_image_id: str, use_gpu: bool):
+    _validate_docker_images(dockerfile_path, local_image_id, use_gpu)
+
+
+@pytest.mark.gpu
 @pytest.mark.parametrize("dockerfile_path", ["keras.test.Dockerfile",
                                              "matplotlib.test.Dockerfile",
                                              "scipy.test.Dockerfile",
@@ -18,7 +31,24 @@
                                              "sm-python-sdk.test.Dockerfile",
                                              "pytorch.examples.Dockerfile",
                                              "tensorflow.examples.Dockerfile"])
-def test_dockerfiles(dockerfile_path: str, local_image_id: str, use_gpu: bool):
+def test_dockerfiles_for_gpu(dockerfile_path: str, local_image_id: str, use_gpu: bool):
+    _validate_docker_images(dockerfile_path, local_image_id, use_gpu)
+
+
+
+# The following is a simple function to check whether the local machine has at least 1 GPU and some Nvidia driver
+# version.
+def _is_nvidia_drivers_available() -> bool:
+    exitcode, output = subprocess.getstatusoutput("nvidia-smi --query-gpu=driver_version --format=csv,noheader --id=0")
+    if exitcode == 0:
+        print(f'Found Nvidia driver version: {output}')
+    else:
+        print(f'No Nvidia drivers found on the machine. Error output: {output}')
+
+    return exitcode == 0
+
+
+def _validate_docker_images(dockerfile_path: str, local_image_id: str, use_gpu: bool):
     print(f'Will start running test for: {dockerfile_path} against: {local_image_id}')
     try:
         image, _ = _docker_client.images.build(path='test/test_artifacts',
@@ -35,7 +65,7 @@ def test_dockerfiles(dockerfile_path: str, local_image_id: str, use_gpu: bool):
     # that the image would have supplied the right entrypoint.
 
     device_requests = []
-    if use_gpu and is_nvidia_drivers_available():
+    if use_gpu and _is_nvidia_drivers_available():
         # Pass all available GPUs, if available.
         device_requests.append(docker.types.DeviceRequest(count=-1, capabilities=[['gpu']]))
 
@@ -53,14 +83,3 @@ def test_dockerfiles(dockerfile_path: str, local_image_id: str, use_gpu: bool):
         # Remove the test docker image after running the test.
         _docker_client.images.remove(image=image.id, force=True)
 
-
-# The following is a simple function to check whether the local machine has at least 1 GPU and some Nvidia driver
-# version.
-def is_nvidia_drivers_available() -> bool:
-    exitcode, output = subprocess.getstatusoutput("nvidia-smi --query-gpu=driver_version --format=csv,noheader --id=0")
-    if exitcode == 0:
-        print(f'Found Nvidia driver version: {output}')
-    else:
-        print(f'No Nvidia drivers found on the machine. Error output: {output}')
-
-    return exitcode == 0