diff --git a/test/test_dockerfile_based_harness.py b/test/test_dockerfile_based_harness.py index b8172a05..39ef2f83 100644 --- a/test/test_dockerfile_based_harness.py +++ b/test/test_dockerfile_based_harness.py @@ -73,9 +73,7 @@ def test_dockerfiles_for_cpu( local_image_version: str, use_gpu: bool, ): - _validate_docker_images( - dockerfile_path, required_packages, local_image_version, use_gpu, "cpu" - ) + _validate_docker_images(dockerfile_path, required_packages, local_image_version, use_gpu, "cpu") @pytest.mark.gpu @@ -139,17 +137,13 @@ def test_dockerfiles_for_gpu( local_image_version: str, use_gpu: bool, ): - _validate_docker_images( - dockerfile_path, required_packages, local_image_version, use_gpu, "gpu" - ) + _validate_docker_images(dockerfile_path, required_packages, local_image_version, use_gpu, "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" - ) + 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: @@ -163,9 +157,7 @@ def _check_docker_file_existence(dockerfile_name: str, test_artifacts_path: str) pytest.skip(f"Skipping test because {dockerfile_name} does not exist.") -def _check_required_package_constraints( - target_version: Version, required_packages: List[str], image_type: str -): +def _check_required_package_constraints(target_version: Version, required_packages: List[str], image_type: str): target_version_dir = get_dir_for_version(target_version) if not os.path.exists(target_version_dir): pytest.skip(f"Skipping test because {target_version_dir} does not exist.") @@ -181,9 +173,7 @@ def _check_required_package_constraints( target_match_spec_out = get_match_specs(env_out_path) for required_package in required_packages: if required_package not in target_match_spec_out: - pytest.skip( - f"Skipping test because {required_package} is not present in {env_out_file_name}" - ) + pytest.skip(f"Skipping test because {required_package} is not present in {env_out_file_name}") def _validate_docker_images( @@ -202,13 +192,9 @@ def _validate_docker_images( for config in _image_generator_configs[target_version.major] if config["image_type"] == image_type ) - docker_image_tag = image_tag_generator_from_config.format( - image_version=local_image_version - ) + docker_image_tag = image_tag_generator_from_config.format(image_version=local_image_version) docker_image_identifier = f"localhost/sagemaker-distribution:{docker_image_tag}" - print( - f"Will start running test for: {dockerfile_path} against: {docker_image_identifier}" - ) + print(f"Will start running test for: {dockerfile_path} against: {docker_image_identifier}") try: image, _ = _docker_client.images.build( @@ -232,9 +218,7 @@ def _validate_docker_images( device_requests = [] if use_gpu and _is_nvidia_drivers_available(): # Pass all available GPUs, if available. - device_requests.append( - docker.types.DeviceRequest(count=-1, capabilities=[["gpu"]]) - ) + device_requests.append(docker.types.DeviceRequest(count=-1, capabilities=[["gpu"]])) # We assume that the image above would have supplied the right entrypoint, so we just run it as is. If the container # didn't execute successfully, the Docker client below will throw an error and fail the test. @@ -310,6 +294,4 @@ def _wait_for_logs(container, search_string, timeout=5, poll_interval=1): if search_string in logs: return True time.sleep(poll_interval) - raise TimeoutError( - f"Container did not log '{search_string}' within {timeout} seconds." - ) + raise TimeoutError(f"Container did not log '{search_string}' within {timeout} seconds.")