diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 9df761e2e7..cddff4ca2b 100755 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -500,7 +500,7 @@ def get_config_build_dir(build_dir, config): return os.path.join(build_dir, config) -def run_subprocess(args, cwd=None, capture=False, dll_path=None, +def run_subprocess(args, cwd=None, capture_stdout=False, dll_path=None, shell=False, env={}): if isinstance(args, str): raise ValueError("args should be a sequence of strings, not a string") @@ -517,7 +517,7 @@ def run_subprocess(args, cwd=None, capture=False, dll_path=None, my_env.update(env) - return run(*args, cwd=cwd, capture=capture, shell=shell, env=my_env) + return run(*args, cwd=cwd, capture_stdout=capture_stdout, shell=shell, env=my_env) def update_submodules(source_dir): @@ -540,7 +540,7 @@ def is_sudo(): def install_apt_package(package): have = package in str(run_subprocess( - ["apt", "list", "--installed", package], capture=True).stdout) + ["apt", "list", "--installed", package], capture_stdout=True).stdout) if not have: if is_sudo(): run_subprocess(['apt-get', 'install', '-y', package]) diff --git a/tools/ci_build/clean_docker_image_cache.py b/tools/ci_build/clean_docker_image_cache.py new file mode 100755 index 0000000000..16b2c82dd6 --- /dev/null +++ b/tools/ci_build/clean_docker_image_cache.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import argparse +import collections +import datetime +import json +import os +import re +import sys +import tempfile +from logger import get_logger + + +SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__)) +REPO_DIR = os.path.normpath(os.path.join(SCRIPT_DIR, "..", "..")) + +sys.path.append(os.path.join(REPO_DIR, "tools", "python")) + + +from util import run # noqa: E402 + + +log = get_logger("clean_docker_image_cache") + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Cleans the docker image cache container registry. " + "This assumes a fairly specific setup - an Azure container registry " + "and a storage account that receives " + "ContainerRegistryRepositoryEvents logs from that registry. " + "The logs are searched to see whether an image was used (pushed or " + "pulled) recently enough in order to determine whether the image " + "should be kept or cleaned up.") + + parser.add_argument( + "--container-registry", required=True, + help="The container registry name.") + + parser.add_argument( + "--log-storage-account", required=True, + help="The storage account name.") + parser.add_argument( + "--log-storage-container", required=True, + help="The storage account container name.") + parser.add_argument( + "--log-storage-path-pattern", default="*.json", + help="The log path pattern in the storage account container.") + + parser.add_argument( + "--cache-lifetime-days", type=int, default=7, + help="How long an image can be cached without being used, in days.") + + parser.add_argument( + "--dry-run", action="store_true", + help="Do a dry-run and don't actually clean anything up.") + + parser.add_argument( + "--az-path", default="az", help="Path to the az client.") + + return parser.parse_args() + + +def az(*args, parse_output=True, az_path): + proc = run(az_path, *args, "--output", "json", capture_stdout=parse_output) + if parse_output: + return json.loads(proc.stdout.decode()) + return None + + +def download_logs(storage_account, container, log_path_pattern, target_dir, az_path): + log_paths = az( + "storage", "blob", "download-batch", + "--destination", target_dir, + "--source", container, + "--account-name", storage_account, + "--pattern", log_path_pattern, + az_path=az_path) + return [os.path.join(target_dir, log_path) for log_path in log_paths] + + +ImageInfo = collections.namedtuple("ImageInfo", ["repository", "digest"]) + + +def get_image_name(image_info): + return "{}@{}".format(image_info.repository, image_info.digest) + + +timestamp_pattern = re.compile( + r"^(?P\d+)-(?P\d+)-(?P\d+)T(?P\d+):(?P\d+):(?P\d+)") + + +def parse_timestamp(timestamp_str): + match = timestamp_pattern.match(timestamp_str) + if match is None: + return None + + return datetime.datetime( + year=int(match['year']), month=int(match['month']), day=int(match['day']), + hour=int(match['hour']), minute=int(match['minute']), second=int(match['second']), + tzinfo=datetime.timezone.utc) + + +def parse_log_line(line, min_datetime): + entry = json.loads(line) + + for field_name, expected_value in [ + ("category", "ContainerRegistryRepositoryEvents"), + ("resultType", "HttpStatusCode"), + ("resultDescription", "200")]: + if entry.get(field_name) != expected_value: + return None + + timestamp = parse_timestamp(entry.get("time", "")) + if timestamp is None or timestamp < min_datetime: + return None + + if entry.get("operationName") not in ["Pull", "Push"]: + return None + + props = entry.get("properties", {}) + repo, digest = props.get("repository"), props.get("digest") + + if repo is None or digest is None: + return None + + return ImageInfo(repo, digest) + + +def get_valid_images_from_logs(log_paths, min_datetime): + valid_images = set() # set of ImageInfo + + for log_path in log_paths: + log.debug("Processing log file: {}".format(log_path)) + with open(log_path, mode="r") as log_file: + for line in log_file: + image_info = parse_log_line(line, min_datetime) + if image_info is not None: + valid_images.add(image_info) + + return valid_images + + +def get_registry_images(container_registry, az_path): + registry_images = set() # set of ImageInfo + + repositories = az( + "acr", "repository", "list", "--name", container_registry, + az_path=az_path) + + for repository in repositories: + digests = az( + "acr", "repository", "show-manifests", + "--repository", repository, "--name", container_registry, + "--query", "[*].digest", + az_path=az_path) + + registry_images.update( + [ImageInfo(repository, digest) for digest in digests]) + + return registry_images + + +def clean_images(container_registry, image_names, az_path): + for image_name in image_names: + az("acr", "repository", "delete", "--name", container_registry, + "--image", image_name, "--yes", + az_path=az_path, + parse_output=False) + + +def main(): + args = parse_args() + + valid_images = set() + + with tempfile.TemporaryDirectory() as tmp_dir: + log_paths = download_logs( + args.log_storage_account, + args.log_storage_container, + args.log_storage_path_pattern, + tmp_dir, + args.az_path) + + cache_lifetime = datetime.timedelta(days=args.cache_lifetime_days) + + min_timestamp = \ + datetime.datetime.now(tz=datetime.timezone.utc) - cache_lifetime + + valid_images = get_valid_images_from_logs(log_paths, min_timestamp) + + all_images = get_registry_images(args.container_registry, args.az_path) + + def sorted_image_names(image_infos): + return sorted([get_image_name(image_info) for image_info in image_infos]) + + log.debug("All images:\n{}".format( + "\n".join(sorted_image_names(all_images)))) + log.debug("Valid images:\n{}".format( + "\n".join(sorted_image_names(valid_images)))) + + images_to_clean = all_images - valid_images + image_names_to_clean = sorted_image_names(images_to_clean) + + log.info("Images to clean:\n{}".format( + "\n".join(image_names_to_clean))) + + if args.dry_run: + log.info("Dry run, no images will be cleaned.") + return 0 + + clean_images(args.container_registry, image_names_to_clean, args.az_path) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/ci_build/github/azure-pipelines/clean-build-docker-image-cache-pipeline.yml b/tools/ci_build/github/azure-pipelines/clean-build-docker-image-cache-pipeline.yml new file mode 100644 index 0000000000..e65394c7c6 --- /dev/null +++ b/tools/ci_build/github/azure-pipelines/clean-build-docker-image-cache-pipeline.yml @@ -0,0 +1,19 @@ +jobs: +- job: Clean_Build_Docker_Image_Cache + + pool: + vmImage: 'ubuntu-18.04' + + timeoutInMinutes: 10 + + steps: + - task: AzureCLI@2 + inputs: + azureSubscription: 'AIInfraBuild' + scriptType: 'bash' + scriptLocation: 'inlineScript' + inlineScript: | + tools/ci_build/clean_docker_image_cache.py \ + --container-registry $(buildcache-container-registry) \ + --log-storage-account $(buildcache-log-storage-account) \ + --log-storage-container $(buildcache-log-storage-account-container) diff --git a/tools/python/util/run.py b/tools/python/util/run.py index 5afe630097..1d6f145954 100644 --- a/tools/python/util/run.py +++ b/tools/python/util/run.py @@ -9,14 +9,15 @@ import subprocess _log = logging.getLogger("util.run") -def run(*args, cwd=None, capture=False, shell=False, env=None, check=True, - quiet=False): +def run(*args, cwd=None, capture_stdout=False, capture_stderr=False, + shell=False, env=None, check=True, quiet=False): """Runs a subprocess. Args: *args: The subprocess arguments. cwd: The working directory. If None, specifies the current directory. - capture: Whether to capture stdout and stderr. + capture_stdout: Whether to capture stdout. + capture_stderr: Whether to capture stderr. shell: Whether to run using the shell. env: The environment variables as a dict. If None, inherits the current environment. @@ -31,11 +32,14 @@ def run(*args, cwd=None, capture=False, shell=False, env=None, check=True, _log.info("Running subprocess in '{0}'\n{1}".format( cwd or os.getcwd(), cmd)) - output = \ - subprocess.PIPE if capture else (subprocess.DEVNULL if quiet else None) + def output(is_stream_captured): + return subprocess.PIPE if is_stream_captured else \ + (subprocess.DEVNULL if quiet else None) + completed_process = subprocess.run( - cmd, cwd=cwd, check=check, stdout=output, stderr=output, env=env, - shell=shell) + cmd, cwd=cwd, check=check, + stdout=output(capture_stdout), stderr=output(capture_stderr), + env=env, shell=shell) _log.debug("Subprocess completed. Return code: {}".format( completed_process.returncode))