Add offline apt package cache in kube runtime (#4226)

* squash commit

* Update package_cache.md

* fix & lint

* add fix for nfs, samba, azurefile

* fix for 32-bit os

* fix

* use dpkg -V instead of dpkg -l

* fix

* fix

* skip some uts

src/kube-runtime/build/kube-runtime.k8s.dockerfile

@@ -15,6 +15,24 @@
 # DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+# Package Cache Data Layer Starts
+
+FROM ubuntu:16.04 as ubuntu_16_04_cache
+
+WORKDIR /src
+COPY src/package_cache ./
+RUN chmod -R +x ./
+RUN /bin/bash ubuntu_build.sh package_cache_info ubuntu16.04
+
+
+FROM ubuntu:18.04 as ubuntu_18_04_cache
+
+WORKDIR /src
+COPY src/package_cache ./
+RUN chmod -R +x ./
+RUN /bin/bash ubuntu_build.sh package_cache_info ubuntu18.04
+
+# Package Cache Data Layer Ends
 
 FROM golang:1.12.6-alpine as builder
 
@@ -29,6 +47,10 @@ RUN ${PROJECT_DIR}/build/runtime/go-build.sh && \
 
 FROM python:3.7-alpine
 
+RUN mkdir -p /opt/package_cache
+COPY --from=ubuntu_16_04_cache /package_cache /opt/package_cache/
+COPY --from=ubuntu_18_04_cache /package_cache /opt/package_cache/
+
 RUN pip install kubernetes pyyaml requests jinja2 pystache
 
 ENV INSTALL_DIR=/opt/kube-runtime

src/kube-runtime/package_cache.md

@@ -0,0 +1,54 @@
+## Package Cache
+
+Some runtime plugins use `apt update` and `apt install ...` to install packages. If the network connection from pai cluster to the apt repository is poor, `apt install` may take a long period of time or fail.
+
+Since we use kube runtime to initialize the environment before job container, it is possible to use the runtime container as a "cache" of some frequently-needed packages.
+
+Note: `Package Cache` only stores packages for systems of arch `x64`.
+
+## How to enable cache for your plugin
+
+**1. Add packages you want to store cache for in the file [`package_cache_info`](src/package_cache/package_cache_info)**:
+
+```
+# group_name, os, packages(space-gapped)
+# "#" can be used for comments
+ssh,ubuntu16.04,openssh-client openssh-server
+ssh,ubuntu18.04,openssh-client openssh-server
+nfs,ubuntu16.04,nfs-common
+nfs,ubuntu18.04,nfs-common
+```
+
+The first column is `group_name`. One group can contain multiple packages. The second column stands for the OS type. Currently only `ubuntu16.04` and `ubuntu18.04` are supported. The third column is the packages you want to add for the group. The last column is the precommands, which will be executed before gathering packages and it can be left empty.
+
+**2. In `init.py` of each plugin:**
+
+```python
+from plugins.plugin_utils import try_to_install_by_cache
+command = [
+    try_to_install_by_cache('<group_name>', fallback_cmds=[
+        '<fallback commands>',
+    ])
+]    
+```
+
+`try_to_install_by_cache(group_name, fallback_cmds)` will generate a script to install all packages of a certain group name. It guarantees:
+
+- If it returns 0, all the packages are installed successfully.
+- If it has a non-zero exit code, the package installation has failed. Reasons could be that the required cache is not found or other internal problems. In such case, plugin will executes the `fallback_cmds` instead. You can use `apt-get`, `yum` or other commands to install the packages.
+
+Here is an example for the `ssh` plugin:
+
+```python
+command = [
+    try_to_install_by_cache('ssh', fallback_cmds=[
+        'apt-get update',
+        'apt-get install -y openssh-client openssh-server',
+    ]),
+    '......'
+]
+```
+
+## Optimization (TBD)
+
+1. Use hash of package url to save storage space: Packages with the same url can be saved together.

src/kube-runtime/src/package_cache/package_cache_info

@@ -0,0 +1,10 @@
+# group_name, os, packages(space-gapped)
+# "#" can be used for comments
+ssh,ubuntu16.04,openssh-client openssh-server
+ssh,ubuntu18.04,openssh-client openssh-server
+nfs,ubuntu16.04,nfs-common
+nfs,ubuntu18.04,nfs-common
+samba,ubuntu16.04,cifs-utils
+samba,ubuntu18.04,cifs-utils
+azurefile,ubuntu16.04,cifs-utils sshpass
+azurefile,ubuntu18.04,cifs-utils sshpass

src/kube-runtime/src/package_cache/ubuntu_build.sh

@@ -0,0 +1,72 @@
+#!/bin/bash
+
+# Copyright (c) Microsoft Corporation
+# All rights reserved.
+#
+# MIT License
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+# documentation files (the "Software"), to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and
+# to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
+# BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+if [ $# -ne 2 ]; then
+  echo "Usage: bash -x <this-script.sh> <path-to-package-cache-info> <os-type>"
+  exit 1
+else
+  package_cache_info=$1
+  os_type=$2
+fi
+
+apt-get update
+ROOT_DIR=/package_cache
+mkdir -p $ROOT_DIR
+
+i=0
+package_dirs=()
+
+while IFS= read -r line || [[ -n "$line" ]] ;
+do
+  start_char=`echo $line | cut -b 1`
+  if [ ! "$start_char" == "#" ]; then
+    name=`echo $line | cut -d , -f 1`
+    os=`echo $line | cut -d , -f 2`
+    packages=`echo $line | cut -d , -f 3`
+    if [ "$os" = "$os_type" ]; then
+      echo "name: ${name} os: ${os} packages: ${packages}"
+      package_dir=$ROOT_DIR"/${name}-${os}"
+      package_dirs[$i]=$package_dir
+      let i++
+      mkdir -p $package_dir
+      cd $package_dir
+      echo $packages > ./packages
+      apt-get -y install --print-uris ${packages} | cut -d " " -f 1-2 | grep http:// > /aptinfo && \
+      cat /aptinfo | cut -d\' -f 2 > ./urls && \
+      apt-get -y install ${packages} --dry-run &> /dry_run_log && \
+      cat /dry_run_log  | grep Conf | cut -d " " -f 2 > ./order
+      if [ $? -ne 0 ]; then
+        echo 'There is an error during package collection.'
+        exit 1
+      fi
+    fi
+  fi
+done < $package_cache_info  
+
+apt-get -y install wget
+
+for package_dir in ${package_dirs[@]};do
+  cd $package_dir
+  wget -i ./urls --tries 3 -P ./ && \
+  ls -la *.deb | awk '{print $9}' | while read filename; do mv $filename `echo $filename | cut -d "_" -f1`".deb"; done;
+  if [ $? -ne 0 ]; then
+    echo 'There is an error during package collection.'
+    exit 1
+  fi
+done

src/kube-runtime/src/plugins/plugin_utils.py

@@ -18,6 +18,8 @@
 
 import logging
 import argparse
+import os
+import shutil
 import yaml
 
 from common.utils import init_logger
@@ -54,3 +56,22 @@ def plugin_init():
     plugin_config = yaml.safe_load(args.plugin_config)
 
     return [plugin_config, args.pre_script, args.post_script]
+
+
+PAI_WORK_DIR = "/usr/local/pai"
+
+
+def try_to_install_by_cache(group_name: str, fallback_cmds: list):
+    source_folder = "/opt/package_cache"
+    target_folder = os.path.join(PAI_WORK_DIR, "package_cache")
+    if not os.path.exists(source_folder):
+        return "{};".format(";".join(fallback_cmds))
+    exists_group_names = os.listdir(source_folder)
+    needed_group_names = [name for name in exists_group_names if name.startswith(group_name + '-')]
+    for name in needed_group_names:
+        name_source_folder = os.path.join(source_folder, name)
+        name_target_folder = os.path.join(target_folder, name)
+        if not os.path.exists(name_target_folder):  # avoid duplicate copy
+            shutil.copytree(name_source_folder, name_target_folder)
+    cached_cmd = "/bin/bash {}/runtime.d/install_group.sh ".format(PAI_WORK_DIR) + group_name
+    return "{} || {{ {}; }}".format(cached_cmd, ";".join(fallback_cmds))

src/kube-runtime/src/plugins/ssh/init.py

@@ -22,7 +22,7 @@
 
 sys.path.append(
     os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
-from plugins.plugin_utils import plugin_init, PluginHelper  #pylint: disable=wrong-import-position
+from plugins.plugin_utils import plugin_init, PluginHelper, try_to_install_by_cache  #pylint: disable=wrong-import-position
 
 LOGGER = logging.getLogger(__name__)
 
@@ -58,6 +58,10 @@ def main():
         LOGGER.info("Skip sshd script since neither jobssh or userssh is set")
     else:
         command = [
+            try_to_install_by_cache("ssh", fallback_cmds=[
+                "apt-get update",
+                "apt-get install -y openssh-client openssh-server",
+            ]),
             "{}/sshd.sh {}\n".format(os.path.dirname(os.path.abspath(__file__)),
                                      " ".join(cmd_params))
         ]

src/kube-runtime/src/plugins/ssh/sshd.sh

@@ -102,13 +102,6 @@ function start_ssh()
   service ssh restart
 }
 
-# Try to install openssh if sshd is not found
-if [ ! -f /usr/sbin/sshd ] ; then
-  apt-get update
-  apt-get install -y openssh-client openssh-server
-fi
-
-
 if [ -f /usr/sbin/sshd ] ; then
     if [ -z "$PAI_CONTAINER_SSH_PORT" ] ; then
         echo "no ssh port provided" >&2

src/kube-runtime/src/plugins/teamwise_storage/storage_command_generator.py

@@ -33,7 +33,7 @@
 
 KUBE_TOKEN_FILE = "/var/run/secrets/kubernetes.io/serviceaccount/token"
 
-STORAGE_PRE_COMMAND = ["apt-get update", "umask 000"]
+STORAGE_PRE_COMMAND = ["umask 000"]
 
 
 def _covert_secret_to_server_config(secret) -> dict:

src/kube-runtime/src/plugins/teamwise_storage/storage_helper.py

@@ -17,6 +17,13 @@
 
 import posixpath
 import re
+import sys
+import os
+
+#pylint: disable=wrong-import-position
+sys.path.append(
+    os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
+from plugins.plugin_utils import try_to_install_by_cache
 
 
 class StorageHelper():
@@ -103,7 +110,10 @@ def _get_nfs_setup_commands(self, server_config, mount_point,
         if phrase == "pre_mount":
             return [
                 "mkdir --parents {}".format(mount_point),
-                "apt-get install --assume-yes nfs-common",
+                try_to_install_by_cache("nfs", fallback_cmds=[
+                    "apt-get update",
+                    "apt-get install --assume-yes nfs-common",
+                ])
             ]
         if phrase in ("tmp_mount", "real_mount"):
             server_data = server_config["data"]
@@ -126,7 +136,10 @@ def _get_samba_setup_commands(self, server_config, mount_point,
         if phrase == "pre_mount":
             return [
                 "mkdir --parents {}".format(mount_point),
-                "apt-get install --assume-yes cifs-utils",
+                try_to_install_by_cache("samba", fallback_cmds=[
+                    "apt-get update",
+                    "apt-get install --assume-yes cifs-utils",
+                ])
             ]
         if phrase in ("tmp_mount", "real_mount"):
             server_data = server_config["data"]
@@ -154,10 +167,12 @@ def _get_azurefile_setup_commands(self, server_config, mount_point,
         if phrase == "pre_mount":
             ret = [
                 "mkdir --parents {}".format(mount_point),
-                "apt-get install --assume-yes cifs-utils",
+                try_to_install_by_cache("azurefile", fallback_cmds=[
+                    "apt-get update",
+                    "apt-get install --assume-yes cifs-utils sshpass",
+                ])
             ]
             if "proxy" in server_data and len(server_data["proxy"]) == 2:
-                ret.append("apt-get install --assume-yes sshpass")
                 proxy_info: str = server_data["proxy"][0]
                 proxy_password: str = server_data["proxy"][1]
                 proxy_ip = proxy_info if proxy_info.find(
@@ -201,6 +216,7 @@ def _get_azureblob_setup_commands(self, server_config, mount_point,
         cfg_file = "/{}.cfg".format(server_name)
         if phrase == "pre_mount":
             return [
+                "apt-get update",
                 "apt-get install --assume-yes wget curl lsb-release apt-transport-https",
                 "valid_release=('14.04' '15.10' '16.04' '16.10' '17.04' '17.10' '18.04' '18.10' '19.04')",
                 "release=`lsb_release -r | cut -f 2`",

src/kube-runtime/src/runtime.d/install_group.sh

@@ -0,0 +1,78 @@
+#!/bin/bash
+
+# Copyright (c) Microsoft Corporation
+# All rights reserved.
+#
+# MIT License
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+# documentation files (the "Software"), to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and
+# to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
+# BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+PAI_WORK_DIR=/usr/local/pai
+CACHE_ROOT_DIR=${PAI_WORK_DIR}/package_cache
+
+function is_ubuntu_package_installed(){
+  for package in $1
+  do
+    dpkg -V $package &> /dev/null
+    if [ $? -ne 0 ]; then
+      return 1
+    fi
+  done
+  return 0
+}
+
+if [ $# -ne 1 ]; then
+  echo "Usage: bash -x install_group.sh <group_name>"
+  exit 1
+else
+  name=$1
+fi
+
+if cat /etc/issue | grep "Ubuntu 16.04" &> /dev/null ; then
+  os='ubuntu16.04'
+elif cat /etc/issue | grep "Ubuntu 18.04" &> /dev/null ; then
+  os='ubuntu18.04'
+else
+  echo "[package_cache] This os doesn't support package cache!"
+  exit 1
+fi
+if [ -d $CACHE_ROOT_DIR"/${name}-${os}" ]; then
+  package_dir=$CACHE_ROOT_DIR"/${name}-${os}"
+  packages=`cat ${package_dir}"/packages"`
+  is_ubuntu_package_installed "${packages}"
+  if [ $? -eq 0 ]; then
+    echo "[package_cache] Skip installation of group ${name}."
+    exit 0
+  fi
+  if [ `getconf LONG_BIT` -eq 64 ]; then
+    echo "[package_cache] Install group ${name} from cache ${package_dir}."
+    cat ${package_dir}"/order" | while read file; do dpkg -i ${package_dir}"/"$file".deb"; done;
+    apt-get install -f
+    # check if packages are installed
+    is_ubuntu_package_installed "${packages}"
+    if [ $? -eq 0 ]; then
+      echo "[package_cache] Install group ${name} from cache ${package_dir} succeeded!"
+    else
+      echo "[package_cache] Install group ${name} from cache ${package_dir} failed. Fallback to apt-get."
+      apt-get update
+      apt-get install -y ${packages}
+    fi
+  else
+    echo "[package_cache] 32-Bit OS is not supported! Fallback to apt-get."
+    apt-get update
+    apt-get install -y ${packages}
+  fi
+else
+  echo "Cannot find dependency ${name}-${os}!"
+  exit 1
+fi