[v1.9.x] Backport docker cache fixes (#20537)

* [v1.x] Fix docker cache build pipeline (#20385)

* libturbojpeg0-dev is only for ubuntu 18.04, so don't install on other ubuntu versions.

* Change logic to not install on 16.04 (works on 18.04 and 20.04.)

* Fix docker registry variable - was inadvertently changed to DOCKER_CACHE_REGISTRY from DOCKER_ECR_REGISTRY.

* Use our local ECR registry for docker-compose images as well.

* Revert changing DOCKER_CACHE_REGISTRY variable.

* Override DOCKER_CACHE_REGISTRY for docker-compose in docker cache pipeline.

* Login to ECR if using ECR for docker registry.

* Import docker_cache when needed.

* Add support for awscli v2, use docker credentials already saved from docker_cache.py run. (#20458)

* [v1.x] Docker cache enhancements (#20525)

* Properly initialize set so that docker cache builds exclude the centos7 dockerfile.

* Change the number of parallel builds to 1 for easier debugging and bandwidth control, also parameterize the number of build retries and reduce from 10 to 1 for easier debugging.

* Remove old dockerfiles.

* Update docker-compose config to use new ECR-based naming scheme.

* Refactor build.py and docker_cache.py to support building and pushing containers for the docker-compose based platforms. Add options for selecting architecture (defaults to x86) since aarch64 container builds will fail on x86 systems.

* Remove old commands for building docker-compose based containers since they are not built in docker_cache.py.

* Fix running of docker-compose based containers via build.py.

Author: josephevans

ci/Jenkinsfile_docker_cache

@@ -38,7 +38,6 @@ core_logic: {
         timeout(time: total_timeout, unit: 'MINUTES') {
           utils.init_git()
           sh "ci/docker_cache.py --docker-registry ${env.DOCKER_ECR_REGISTRY}"
-          sh "cd ci && python3 ./docker_login.py --secret-name ${env.DOCKERHUB_SECRET_NAME} && docker-compose -f docker/docker-compose.yml build --parallel && docker-compose -f docker/docker-compose.yml push && docker logout"
         }
       }
     }

ci/Jenkinsfile_utils.groovy

@@ -148,7 +148,7 @@ def collect_test_results_windows(original_file_name, new_file_name) {
 
 
 def docker_run(platform, function_name, use_nvidia, shared_mem = '500m', env_vars = "", build_args = "") {
-  def command = "ci/build.py %ENV_VARS% %BUILD_ARGS% --docker-registry ${env.DOCKER_CACHE_REGISTRY} %USE_NVIDIA% --platform %PLATFORM% --docker-build-retries 3 --shm-size %SHARED_MEM% /work/runtime_functions.sh %FUNCTION_NAME%"
+  def command = "ci/build.py %ENV_VARS% %BUILD_ARGS% --docker-registry ${env.DOCKER_ECR_REGISTRY} %USE_NVIDIA% --platform %PLATFORM% --docker-build-retries 3 --shm-size %SHARED_MEM% /work/runtime_functions.sh %FUNCTION_NAME%"
   command = command.replaceAll('%ENV_VARS%', env_vars.length() > 0 ? "-e ${env_vars}" : '')
   command = command.replaceAll('%BUILD_ARGS%', env_vars.length() > 0 ? "${build_args}" : '')
   command = command.replaceAll('%USE_NVIDIA%', use_nvidia ? '--nvidiadocker' : '')

ci/build.py

@@ -42,24 +42,36 @@
 
 from util import *
 
-DOCKER_COMPOSE_WHITELIST = ('centos7_cpu', 'centos7_gpu_cu92', 'centos7_gpu_cu100',
-                            'centos7_gpu_cu101', 'centos7_gpu_cu102', 'centos7_gpu_cu110',
-                            'centos7_gpu_cu112')
-
 # Files for docker compose
-DOCKER_COMPOSE_FILES = set(('docker/build.centos7'))
+DOCKER_COMPOSE_FILES = set(['docker/build.centos7'])
+
+# keywords to identify arm-based dockerfiles
+AARCH_FILE_KEYWORDS = ['armv', 'aarch64']
 
 def get_dockerfiles_path():
     return "docker"
 
+def get_docker_compose_platforms(path: str = get_dockerfiles_path()):
+    platforms = set()
+    with open(os.path.join(path, "docker-compose.yml"), "r") as f:
+        compose_config = yaml.load(f.read(), yaml.SafeLoader)
+        for platform in compose_config["services"]:
+            platforms.add(platform)
+    return platforms
 
-def get_platforms(path: str = get_dockerfiles_path(), legacy_only=False) -> List[str]:
-    """Get a list of architectures given our dockerfiles"""
+
+def get_platforms(path: str = get_dockerfiles_path(), arch='x86') -> List[str]:
+    """Get a list of platforms given our dockerfiles"""
     dockerfiles = glob.glob(os.path.join(path, "Dockerfile.*"))
     dockerfiles = set(filter(lambda x: x[-1] != '~', dockerfiles))
     files = set(map(lambda x: re.sub(r"Dockerfile.(.*)", r"\1", x), dockerfiles))
-    if legacy_only:
-        files = files - DOCKER_COMPOSE_FILES
+    files = files - DOCKER_COMPOSE_FILES
+    files.update(["build."+x for x in get_docker_compose_platforms()])
+    arm_files = set(filter(lambda x: any(y in x for y in AARCH_FILE_KEYWORDS), files))
+    if arch == 'x86':
+        files = files - arm_files
+    elif arch == 'aarch64':
+        files = arm_files
     platforms = list(map(lambda x: os.path.split(x)[1], sorted(files)))
     return platforms
 
@@ -87,14 +99,21 @@ def _hash_file(ctx, filename):
                 break
             ctx.update(d)
 
+def is_docker_compose(platform: str) -> bool:
+    """:return: boolean whether specified platform container uses docker-compose"""
+    platlist = get_docker_compose_platforms()
+    platform = platform.split(".")[1] if any(x in platform for x in ['build.', 'publish.']) else platform
+    return platform in platlist
+
+
 def get_docker_tag(platform: str, registry: str) -> str:
     """:return: docker tag to be used for the container"""
-    if platform in DOCKER_COMPOSE_WHITELIST:
+    platform = platform if any(x in platform for x in ['build.', 'publish.']) else 'build.{}'.format(platform)
+    if is_docker_compose(platform):
         with open("docker/docker-compose.yml", "r") as f:
             compose_config = yaml.load(f.read(), yaml.SafeLoader)
-            return compose_config["services"][platform]["image"].replace('${DOCKER_CACHE_REGISTRY}', registry)
+            return compose_config["services"][platform.split(".")[1]]["image"].replace('${DOCKER_CACHE_REGISTRY}', registry)
 
-    platform = platform if any(x in platform for x in ['build.', 'publish.']) else 'build.{}'.format(platform)
     if not registry:
         registry = "mxnet_local"
     dockerfile = get_dockerfile(platform)
@@ -121,9 +140,9 @@ def build_docker(platform: str, registry: str, num_retries: int, no_cache: bool,
     :return: Id of the top level image
     """
     tag = get_docker_tag(platform=platform, registry=registry)
-    
     # docker-compose
-    if platform in DOCKER_COMPOSE_WHITELIST:
+    if is_docker_compose(platform):
+        docker_compose_platform = platform.split(".")[1] if any(x in platform for x in ['build.', 'publish.']) else platform
         logging.info('Building docker container tagged \'%s\' based on ci/docker/docker-compose.yml', tag)
         # We add a user with the same group as the executing non-root user so files created in the
         # container match permissions of the local user. Same for the group.
@@ -132,7 +151,7 @@ def build_docker(platform: str, registry: str, num_retries: int, no_cache: bool,
                "--build-arg", "GROUP_ID={}".format(os.getgid())]
         if cache_intermediate:
             cmd.append('--no-rm')
-        cmd.append(platform)
+        cmd.append(docker_compose_platform)
     else:
         logging.info("Building docker container tagged '%s'", tag)
         #
@@ -286,17 +305,24 @@ def docker_run_cmd(cmd):
     return 0
 
 
-def list_platforms() -> str:
-    return "\nSupported platforms:\n{}".format('\n'.join(get_platforms()))
+def list_platforms(arch='x86') -> str:
+    return "\nSupported platforms:\n{}".format('\n'.join(get_platforms(arch=arch)))
 
 
 def load_docker_cache(platform, tag, docker_registry) -> None:
     """Imports tagged container from the given docker registry"""
     if docker_registry:
-        if platform in DOCKER_COMPOSE_WHITELIST:
+        if is_docker_compose(platform):
+            docker_compose_platform = platform.split(".")[1] if any(x in platform for x in ['build.', 'publish.']) else platform
             env = os.environ.copy()
             env["DOCKER_CACHE_REGISTRY"] = docker_registry
-            cmd = ['docker-compose', '-f', 'docker/docker-compose.yml', 'pull', platform]
+            if "dkr.ecr" in docker_registry:
+                try:
+                    import docker_cache
+                    docker_cache._ecr_login(docker_registry)
+                except Exception:
+                    logging.exception('Unable to login to ECR...')
+            cmd = ['docker-compose', '-f', 'docker/docker-compose.yml', 'pull', docker_compose_platform]
             logging.info("Running command: 'DOCKER_CACHE_REGISTRY=%s %s'", docker_registry, ' '.join(cmd))
             check_call(cmd, env=env)
             return
@@ -335,6 +361,11 @@ def main() -> int:
                         help="platform",
                         type=str)
 
+    parser.add_argument("-A", "--architecture",
+                        help="Architecture of images to build (x86 or aarch64). Default is x86.",
+                        default='x86',
+                        dest='architecture')
+
     parser.add_argument("-b", "--build-only",
                         help="Only build the container, don't build the project",
                         action='store_true')
@@ -401,7 +432,7 @@ def main() -> int:
                         for e in args.environment])
 
     if args.list:
-        print(list_platforms())
+        print(list_platforms(arch=args.architecture))
     elif args.platform:
         platform = args.platform
         tag = get_docker_tag(platform=platform, registry=args.docker_registry)
@@ -445,9 +476,9 @@ def main() -> int:
             return ret
 
     elif args.all:
-        platforms = get_platforms()
+        platforms = get_platforms(arch=args.architecture)
         platforms = [platform for platform in platforms if 'build.' in platform]
-        logging.info("Building for all architectures: %s", platforms)
+        logging.info("Building for all platforms: %s", platforms)
         logging.info("Artifacts will be produced in the build/ directory.")
         for platform in platforms:
             tag = get_docker_tag(platform=platform, registry=args.docker_registry)
@@ -474,7 +505,7 @@ def main() -> int:
 
     else:
         parser.print_help()
-        list_platforms()
+        list_platforms(arch=args.architecture)
         print("""
 Examples: