Docker load ydbd_slice (#12797)

3 files changed, 108 insertions, 12 deletions

diff --git a/ydb/tools/ydbd_slice/__init__.py b/ydb/tools/ydbd_slice/__init__.py
index cccee7041f..2d5d211c65 100644
--- a/ydb/tools/ydbd_slice/__init__.py
+++ b/ydb/tools/ydbd_slice/__init__.py
@@ -857,26 +857,36 @@ def add_sample_config_mode(modes):
 
 #
 # docker and kube scenarios
-def build_and_push_docker_image(build_args, docker_package, build_ydbd, image, force_rebuild):
+def build_docker_image(build_args, docker_package, build_ydbd, image, force_rebuild):
     if docker_package is None:
         docker_package = docker.DOCKER_IMAGE_YDBD_PACKAGE_SPEC
 
     logger.debug(f'using docker package spec: {docker_package}')
 
     image_details = docker.docker_inspect(image)
+    output_path = docker.get_image_output_path(image)
 
     if image_details is None:
         logger.debug('ydb image %s is not present on host, building', image)
         root = arcadia_root()
         ya_package_docker(root, build_args, docker_package, image)
+        docker.docker_image_save(image, output_path, True)
     elif force_rebuild:
         logger.debug('ydb image %s is already present on host, rebuilding', image)
         root = arcadia_root()
         ya_package_docker(root, build_args, docker_package, image)
+        docker.docker_image_save(image, output_path, True)
     else:
         logger.debug('ydb image %s is already present on host, using existing image', image)
+        docker.docker_image_save(image, output_path, False)
 
-    docker.docker_push(image)
+
+def push_docker_image(image):
+    image_details = docker.docker_inspect(image)
+    if image_details is not None:
+        docker.docker_push(image)
+    else:
+        logger.error('ydb image %s is not present on host, skip', image)
 
 
 def add_arguments_docker_build_with_remainder(mode, add_force_rebuild=False):
@@ -908,13 +918,24 @@ def add_arguments_docker_build_with_remainder(mode, add_force_rebuild=False):
     )
 
 
+def add_arguments_docker_push_with_remainder(mode):
+    group = mode.add_argument_group('docker push options')
+    group.add_argument(
+        '-i', '--image',
+        help='Optional: docker image name and tag to push. Conflicts with "-t" argument.',
+    )
+    group.add_argument(
+        '-t', '--tag',
+        help='Optional: docker image tag to push. Conflicts with "-i" argument. Default is {user}-latest.',
+    )
+
+
 def add_docker_build_mode(modes):
     def _run(args):
         logger.debug("starting docker-build cmd with args '%s'", args)
         try:
             image = docker.get_image_from_args(args)
-            build_and_push_docker_image(args.build_args, args.docker_package, False, image, force_rebuild=True)
-
+            build_docker_image(args.build_args, args.docker_package, False, image, True)
             logger.info('docker-build finished')
         except RuntimeError as e:
             logger.error(e.args[0])
@@ -929,6 +950,26 @@ def add_docker_build_mode(modes):
     mode.set_defaults(handler=_run)
 
 
+def add_docker_push_mode(modes):
+    def _run(args):
+        logger.debug("starting docker-push cmd with args '%s'", args)
+        try:
+            image = docker.get_image_from_args(args)
+            push_docker_image(image)
+            logger.info('docker-push finished')
+        except RuntimeError as e:
+            logger.error(e.args[0])
+            sys.exit(1)
+
+    mode = modes.add_parser(
+        "docker-push",
+        parents=[],
+        description="Push YDB docker image."
+    )
+    add_arguments_docker_push_with_remainder(mode)
+    mode.set_defaults(handler=_run)
+
+
 def add_kube_generate_mode(modes):
     def _run(args):
         logger.debug("starting kube-generate cmd with args '%s'", args)
@@ -989,11 +1030,11 @@ def add_kube_install_mode(modes):
         try:
             image = docker.get_image_from_args(args)
             if not args.use_prebuilt_image:
-                build_and_push_docker_image(args.build_args, args.docker_package, False, image, force_rebuild=args.force_rebuild)
+                build_docker_image(args.build_args, args.docker_package, False, image, args.force_rebuild)
 
             manifests = kube_handlers.get_all_manifests(args.path)
             kube_handlers.manifests_ydb_set_image(args.path, manifests, image)
-            kube_handlers.slice_install(args.path, manifests, args.wait_ready, args.dynamic_config_type)
+            kube_handlers.slice_install(args.path, manifests, args.wait_ready, args.dynamic_config_type, image, args.use_prebuilt_image)
 
             logger.info('kube-install finished')
         except RuntimeError as e:
@@ -1036,12 +1077,12 @@ def add_kube_update_mode(modes):
         try:
             image = docker.get_image_from_args(args)
             if not args.use_prebuilt_image:
-                build_and_push_docker_image(args.build_args, args.docker_package, False, image, force_rebuild=args.force_rebuild)
+                build_docker_image(args.build_args, args.docker_package, False, image, args.force_rebuild)
 
             manifests = kube_handlers.get_all_manifests(args.path)
             manifests = kube_handlers.manifests_ydb_filter_components(args.path, manifests, args.components)
             kube_handlers.manifests_ydb_set_image(args.path, manifests, image)
-            kube_handlers.slice_update(args.path, manifests, args.wait_ready, args.dynamic_config_type)
+            kube_handlers.slice_update(args.path, manifests, args.wait_ready, args.dynamic_config_type, image, args.use_prebuilt_image)
 
             logger.info('kube-update finished')
         except RuntimeError as e:
@@ -1368,6 +1409,7 @@ def main(walle_provider=None):
         add_sample_config_mode(modes)
 
         add_docker_build_mode(modes)
+        add_docker_push_mode(modes)
         add_kube_generate_mode(modes)
         add_kube_install_mode(modes)
         add_kube_update_mode(modes)
diff --git a/ydb/tools/ydbd_slice/kube/docker.py b/ydb/tools/ydbd_slice/kube/docker.py
index 1c3530a2d2..db3b8a9f5d 100644
--- a/ydb/tools/ydbd_slice/kube/docker.py
+++ b/ydb/tools/ydbd_slice/kube/docker.py
@@ -43,6 +43,16 @@ def get_image_from_args(args):
         return "%s:%s" % (DOCKER_IMAGE_FULL_NAME, tag)
 
 
+def get_image_output_path(image):
+    if ":" in image:
+        image_name, tag = image.split(":")
+    else:
+        image_name, tag = image, "latest"
+    image_base_name = image_name.split("/")[-1]
+
+    return f"{image_base_name}.{tag}.tar"
+
+
 def docker_tag(old, new):
     proc = subprocess.Popen(['docker', 'tag', old, new], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
     try:
@@ -70,6 +80,20 @@ def docker_inspect(obj):
         return json.loads(stdout)
 
 
+def docker_image_save(image, output_path, overwrite):
+    if os.path.exists(output_path) and not overwrite:
+        logger.info(f"Compressed file '{output_path}' already exists, using existing archive")
+        return
+
+    logger.info(f"execute command 'docker save image' for '{image}' to output path '{output_path}'")
+    proc = subprocess.Popen(['docker', 'image', 'save', image, '-o', output_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    _, stderr = proc.communicate()
+    if proc.returncode != 0:
+        raise RuntimeError("docker image save: failed with code %d, error: %s" % (proc.returncode, stderr))
+
+    logger.info(f"Docker image '{image}' saved and compressed successfully to '{output_path}'")
+
+
 def docker_push(image):
     proc = subprocess.Popen(['docker', 'push', image], text=True)
     proc.communicate()
diff --git a/ydb/tools/ydbd_slice/kube/handlers.py b/ydb/tools/ydbd_slice/kube/handlers.py
index 597afdaa38..cf9a916d5d 100644
--- a/ydb/tools/ydbd_slice/kube/handlers.py
+++ b/ydb/tools/ydbd_slice/kube/handlers.py
@@ -7,7 +7,7 @@ from collections import defaultdict
 from kubernetes.client import Configuration
 
 from ydb.tools.ydbd_slice import nodes
-from ydb.tools.ydbd_slice.kube import api, kubectl, yaml, generate, cms, dynconfig
+from ydb.tools.ydbd_slice.kube import api, kubectl, yaml, generate, cms, dynconfig, docker
 
 
 logger = logging.getLogger(__name__)
@@ -114,7 +114,7 @@ def update_image(data, image):
         data['spec'].pop('version')
     image_data = data['spec'].setdefault('image', {})
     image_data['name'] = image
-    image_data['pullPolicy'] = 'Always'
+    image_data['pullPolicy'] = 'IfNotPresent'
 
 
 def update_manifest(path, data):
@@ -184,8 +184,34 @@ def manifests_ydb_filter_components(project_path, manifests, update_components):
     return result
 
 
+def slice_docker_load(api_client, project_path, manifests, image, use_prebuilt_image):
+    if use_prebuilt_image:
+        logger.info('arg use_prebuilt_image found, nothing to load.')
+        return
+
+    node_list = get_nodes(api_client, project_path, manifests)
+    if len(node_list) == 0:
+        logger.info('no nodes found, nothing to load.')
+        return
+
+    built_image_path = docker.get_image_output_path(image)
+    remote_path = f"/Berkanavt/kikimr/{os.path.basename(built_image_path)}"
+
+    image_details = docker.docker_inspect(image)
+    local_digest = image_details[0]['Id']
+
+    node_list = nodes.Nodes(node_list)
+    node_list.copy(built_image_path, remote_path)
+
+    cmd = 'remote_digest=$(sudo crictl inspecti -o json {_image} | jq -r ".status.id"); if [[ "{_local_digest}" != "$remote_digest" ]]; then sudo ctr image import {_remote_path}; fi'.format(
+        _image=image, _local_digest=local_digest, _remote_path=remote_path
+    )
+    node_list.execute_async(cmd)
+
 #
 # macro level nodeclaim functions
+
+
 def slice_namespace_apply(api_client, project_path, manifests):
     for (path, _, kind, _, _, data) in manifests:
         if kind != 'namespace':
@@ -269,6 +295,8 @@ def wait_for_storage(api_client, project_path, manifests):
 
 
 # macro level ydb functions
+
+
 def slice_ydb_apply(api_client, project_path, manifests, dynamic_config_type):
     # process storages first
     for (path, api_version, kind, namespace, name, data) in manifests:
@@ -462,7 +490,7 @@ def slice_generate(project_path, user, slice_name, template, template_vars):
         sys.exit(f'Slice template {template} not implemented.')
 
 
-def slice_install(project_path, manifests, wait_ready, dynamic_config_type):
+def slice_install(project_path, manifests, wait_ready, dynamic_config_type, image, use_prebuilt_image):
     with api.ApiClient() as api_client:
         slice_namespace_apply(api_client, project_path, manifests)
         slice_nodeclaim_apply(api_client, project_path, manifests)
@@ -470,14 +498,16 @@ def slice_install(project_path, manifests, wait_ready, dynamic_config_type):
         slice_ydb_delete(api_client, project_path, manifests)
         slice_ydb_storage_wait_pods_deleted(api_client, project_path, manifests)
         slice_nodeclaim_format(api_client, project_path, manifests)
+        slice_docker_load(api_client, project_path, manifests, image, use_prebuilt_image)
         slice_ydb_apply(api_client, project_path, manifests, dynamic_config_type)
         slice_ydb_wait_ready(api_client, project_path, manifests, wait_ready)
 
 
-def slice_update(project_path, manifests, wait_ready, dynamic_config_type):
+def slice_update(project_path, manifests, wait_ready, dynamic_config_type, image, use_prebuilt_image):
     with api.ApiClient() as api_client:
         slice_nodeclaim_apply(api_client, project_path, manifests)
         slice_nodeclaim_wait_ready(api_client, project_path, manifests)
+        slice_docker_load(api_client, project_path, manifests, image, use_prebuilt_image)
         slice_ydb_apply(api_client, project_path, manifests, dynamic_config_type)
         slice_ydb_restart(api_client, project_path, manifests)
         slice_ydb_wait_ready(api_client, project_path, manifests, wait_ready)