checkstyle

src/config.py

@@ -76,5 +76,5 @@
             "pytest_flags": [],
             "image_type": "cpu",
         },
-    ]
+    ],
 }

src/main.py

@@ -53,7 +53,8 @@ def create_and_get_semver_dir(version: Version, exist_ok: bool = False):
 
 def _delete_all_files_except_additional_packages_input_files(base_version_dir, version: Version):
     additional_package_env_in_files = [
-        image_generator_config["additional_packages_env_in_file"] for image_generator_config in _image_generator_configs[version.major]
+        image_generator_config["additional_packages_env_in_file"]
+        for image_generator_config in _image_generator_configs[version.major]
     ]
     for filename in os.listdir(base_version_dir):
         if filename not in additional_package_env_in_files:
@@ -236,6 +237,7 @@ def _get_config_for_image(target_version_dir: str, image_generator_config, force
     config_for_image["build_args"].pop("ARG_BASED_ENV_IN_FILENAME", None)
     return config_for_image
 
+
 # Returns a tuple of: 1/ list of actual images generated; 2/ list of tagged images. A given image can be tagged by
 # multiple different strings - for e.g., a CPU image can be tagged as '1.3.2-cpu', '1.3-cpu', '1-cpu' and/or
 # 'latest-cpu'. Therefore, (1) is strictly a subset of (2).

test/test_dockerfile_based_harness.py

@@ -117,7 +117,9 @@ def _check_required_package_constraints(target_version: Version, required_packag
         pytest.skip(f"Skipping test because {target_version_dir} does not exist.")
     # fetch the env.out file for this image_type
     env_out_file_name = next(
-        config["env_out_filename"] for config in _image_generator_configs[target_version.major] if config["image_type"] == image_type
+        config["env_out_filename"]
+        for config in _image_generator_configs[target_version.major]
+        if config["image_type"] == image_type
     )
     env_out_path = f"{target_version_dir}/{env_out_file_name}"
     if not os.path.exists(env_out_path):
@@ -136,7 +138,9 @@ def _validate_docker_images(
     _check_docker_file_existence(dockerfile_path, test_artifacts_path)
     _check_required_package_constraints(target_version, required_packages, image_type)
     image_tag_generator_from_config = next(
-        config["image_tag_generator"] for config in _image_generator_configs[target_version.major] if config["image_type"] == image_type
+        config["image_tag_generator"]
+        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_identifier = f"localhost/sagemaker-distribution:{docker_image_tag}"

test/test_main.py

@@ -100,10 +100,11 @@ def _create_prev_docker_file(file_path):
         FROM mambaorg / micromamba:$TAG_FOR_BASE_MICROMAMBA_IMAGE\nprevious_dockerfile\n"""
         )
 
+
 def _create_gpu_cuda_config_file(file_path):
     gpu_cuda_config_context = {
         "TAG_FOR_BASE_MICROMAMBA_IMAGE": "jammy-cuda-test-version",
-        "CUDA_MAJOR_MINOR_VERSION": "test-major-minor-version"
+        "CUDA_MAJOR_MINOR_VERSION": "test-major-minor-version",
     }
 
     with open(file_path, "w") as gpu_cuda_config:
@@ -121,7 +122,7 @@ def mock_get_dir_for_version(base_version):
     input_version = get_semver(version)
     # Create directory for base version
     input_version_dir = create_and_get_semver_dir(input_version)
-    print('input_version_dir', input_version_dir)
+    print("input_version_dir", input_version_dir)
     # Create env.in and env.out for base version
     _create_docker_cpu_env_in_file(input_version_dir + "/cpu.env.in")
     _create_docker_gpu_env_in_file(input_version_dir + "/gpu.env.in")