使用Kubeflow构建机器学习流水线

在此前的文章中，我已经向你介绍了Kubeflow，这是一个为团队设置的机器学习平台，需要构建机器学习流水线。
在本文中，我们将了解如何采用现有的机器学习详细并将其变成Kubeflow的机器学习流水线，进而可以部署在Kubernetes上。在进行本次练习的时候，请考虑你该如何将现有的机器学习项目转换到Kubeflow上。
我将使用Fashion MNIST作为例子，因为在本次练习中模型的复杂性并不是我们需要解决的主要目标。对于这一简单的例子，我将流水线分为3个阶段：

• Git clone代码库
• 下载并重新处理训练和测试数据
• 训练评估

当然，你可以根据自己的用例将流水线以任意形式拆分，并且可以随意扩展流水线。

% git clone https://github.com/benjamintanweihao/kubeflow-mnist.git

以下是我们用来创建流水线的完整清单。实际上，你的代码很可能跨多个库和文件。在我们的例子中，我们将代码分为两个脚本， preprocessing.py
和 train.py
。

from tensorflow import keras

import argparse

import os

import pickle







def preprocess(data_dir: str):

    fashion_mnist = keras.datasets.fashion_mnist

    (train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()




    train_images = train_images / 255.0

    test_images = test_images / 255.0




    os.makedirs(data_dir, exist_ok=True)




    with open(os.path.join(data_dir, 'train_images.pickle'), 'wb') as f:

  pickle.dump(train_images, f)




    with open(os.path.join(data_dir, 'train_labels.pickle'), 'wb') as f:

  pickle.dump(train_labels, f)




    with open(os.path.join(data_dir, 'test_images.pickle'), 'wb') as f:

        pickle.dump(test_images, f)




    with open(os.path.join(data_dir, 'test_labels.pickle'), 'wb') as f:

        pickle.dump(test_labels, f)




if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Kubeflow MNIST training script')

    parser.add_argument('--data_dir', help='path to images and labels.')

    args = parser.parse_args()




    preprocess(data_dir=args.data_dir)

处理脚本采用单个参数 data_dir
。它下载并预处理数据，并将 pickled
版本保存在 data_dir
中。在生产代码中，这可能是TFRecords的存储目录。

train.py

import calendar

import os

import time




import tensorflow as tf

import pickle

import argparse




from tensorflow import keras

from constants import PROJECT_ROOT







def train(data_dir: str):

    # Training

    model = keras.Sequential([

          keras.layers.Flatten(input_shape=(28, 28)),

          keras.layers.Dense(128, activation='relu'),

          keras.layers.Dense(10)])




    model.compile(optimizer='adam',

              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),

                  metrics=['accuracy'])




    with open(os.path.join(data_dir, 'train_images.pickle'), 'rb') as f:

        train_images = pickle.load(f)




    with open(os.path.join(data_dir, 'train_labels.pickle'), 'rb') as f:

        train_labels = pickle.load(f)




    model.fit(train_images, train_labels, epochs=10)




    with open(os.path.join(data_dir, 'test_images.pickle'), 'rb') as f:

        test_images = pickle.load(f)




    with open(os.path.join(data_dir, 'test_labels.pickle'), 'rb') as f:

        test_labels = pickle.load(f)




    # Evaluation

    test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=2)




    print(f'Test Loss: {test_loss}')

    print(f'Test Acc: {test_acc}')




    # Save model

    ts = calendar.timegm(time.gmtime())

    model_path = os.path.join(PROJECT_ROOT, f'mnist-{ts}.h5')

    tf.saved_model.save(model, model_path)




    with open(os.path.join(PROJECT_ROOT, 'output.txt'), 'w') as f:

        f.write(model_path)

        print(f'Model written to: {model_path}')







if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Kubeflow FMNIST training script')

    parser.add_argument('--data_dir', help='path to images and labels.')

    args = parser.parse_args()




    train(data_dir=args.data_dir)

在 train.py
中，将建立模型，并使用 data_dir
指定训练和测试数据的位置。模型训练完毕并且开始执行评估后，将模型写入带有时间戳的路径。请注意，该路径也已写入 output.txt
。稍后将对此进行引用。

为了开始创建Kubeflow流水线，我们需要拉取一些依赖项。我准备了一个 environment.yml
，其中包括了 kfp 0.5.0
、 tensorflow
以及其他所需的依赖项。
你需要安装Conda，然后执行以下步骤：

% conda env create -f environment.yml

% source activate kubeflow-mnist

% python preprocessing.py --data_dir=/path/to/data

% python train.py --data_dir=/path/to/data

现在我们来回顾一下我们流水线中的几个步骤：

• Git clone代码库
• 下载并预处理训练和测试数据
• 训练并进行评估

在我们开始写代码之前，需要从宏观上了解Kubeflow流水线。

流水线由连接组件构成。一个组件的输出成为另一个组件的输入，每个组件实际上都在容器中执行（在本例中为Docker）。将发生的情况是，我们会执行一个我们稍后将要指定的Docker镜像，它包含了我们运行 preprocessing.py
和 train.py
所需的一切。当然，这两个阶段会有它们的组件。

我们还需要额外的一个镜像以 git clone
项目。我们需要将项目bake到Docker镜像，但在实际项目中，这可能会导致Docker镜像的大小膨胀。
说到Docker镜像，我们应该先创建一个。

FROM tensorflow/tensorflow:1.14.0-gpu-py3

LABEL MAINTAINER "Benjamin Tan "

SHELL ["/bin/bash", "-c"]




# Set the locale

RUN echo 'Acquire {http::Pipeline-Depth "0";};' >> /etc/apt/apt.conf

RUN DEBIAN_FRONTEND="noninteractive"

RUN apt-get update  && apt-get -y install --no-install-recommends locales && locale-gen en_US.UTF-8

ENV LANG en_US.UTF-8

ENV LANGUAGE en_US:en

ENV LC_ALL en_US.UTF-8




RUN apt-get install -y --no-install-recommends \

    wget \

    git \

    python3-pip \

    openssh-client \

    python3-setuptools \

    google-perftools && \

    rm -rf /var/lib/apt/lists/*




# install conda

WORKDIR /tmp

RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-4.7.12-Linux-x86_64.sh -O ~/miniconda.sh && \

    /bin/bash ~/miniconda.sh -b -p /opt/conda && \

    rm ~/miniconda.sh && \

    ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \

    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc




# build conda environments

COPY environment.yml /tmp/kubeflow-mnist/conda/

RUN /opt/conda/bin/conda update -n base -c defaults conda

RUN /opt/conda/bin/conda env create -f /tmp/kubeflow-mnist/conda/environment.yml

RUN /opt/conda/bin/conda clean -afy




# Cleanup

RUN rm -rf /workspace/{nvidia,docker}-examples && rm -rf /usr/local/nvidia-examples && \

    rm /tmp/kubeflow-mnist/conda/environment.yml




# switch to the conda environment

RUN echo "conda activate kubeflow-mnist" >> ~/.bashrc

ENV PATH /opt/conda/envs/kubeflow-mnist/bin:$PATH

RUN /opt/conda/bin/activate kubeflow-mnist




# make /bin/sh symlink to bash instead of dash:

RUN echo "dash dash/sh boolean false" | debconf-set-selections && \

    DEBIAN_FRONTEND=noninteractive dpkg-reconfigure dash




# Set the new Allocator

ENV LD_PRELOAD /usr/lib/x86_64-linux-gnu/libtcmalloc.so.

关于Dockerfile值得关注的重要一点是Conda环境是否设置完成并准备就绪。要构建镜像：

% docker build -t your-user-name/kubeflow-mnist . -f Dockerfile

% docker push your-user-name/kubeflow-mnist

那么，现在让我们来创建第一个组件！

在 pipeline.py
中可以找到以下代码片段。

def git_clone_darkrai_op(repo_url: str):




    volume_op = dsl.VolumeOp(

        name="create pipeline volume",

        resource_name="pipeline-pvc",

        modes=["ReadWriteOnce"],

        size="3Gi"

    )




    image = 'alpine/git:latest'




    commands = [

        "mkdir ~/.ssh",

        "cp /etc/ssh-key/id_rsa ~/.ssh/id_rsa",

        "chmod 600 ~/.ssh/id_rsa",

        "ssh-keyscan bitbucket.org >> ~/.ssh/known_hosts",

        f"git clone {repo_url} {PROJECT_ROOT}",

        f"cd {PROJECT_ROOT}"]




    op = dsl.ContainerOp(

        name='git clone',

        image=image,

        command=['sh'],

        arguments=['-c', ' && '.join(commands)],

        container_kwargs={'image_pull_policy': 'IfNotPresent'},

        pvolumes={"/workspace": volume_op.volume}

    )




    # Mount Git Secrets

    op.add_volume(V1Volume(name='ssh-key-volume',

                           secret=V1SecretVolumeSource(secret_name='ssh-key-secret')))

    op.add_volume_mount(V1VolumeMount(mount_path='/etc/ssh-key', name='ssh-key-volume', read_only=True))




    return op

首先，创建一个Kubernetes volume，预定义大小为3Gi。其次，将 image
变量指定为我们将要使用的 alpine/git
Docker镜像。之后是在Docker容器中执行的命令列表。这些命令实质上是设置SSH密钥的，以便于流水线可以从私有仓库 git clone
，或者使用 git://URL
来代替  https://
。

该函数的核心是下面一行，返回一个 dsl.ContainerOp
。

command
和 arguments
指定了执行镜像之后需要执行的命令。

最后一个变量十分有趣，是 pvolumes
，它是Pipeline Volumes简称。它创建一个Kubernetes volume并允许流水线组件来共享单个存储。该volume被挂载在 /workspace
上。那么这个组件要做的就是把仓库 git clone
到 /workspace
中。

流水线volume在哪里创建呢？当我们将所有组件都整合到一个流水线中时，就会看到创建好的volume。我们在 /etc/ssh-key/
上安装secrets：

op.add_volume_mount(V1VolumeMount(mount_path='/etc/ssh-key', name='ssh-key-volume', read_only=True))

请记得我们将secret命名为 ssh-key-secret
：

op.add_volume(V1Volume(name='ssh-key-volume',

                           secret=V1SecretVolumeSource(secret_name='ssh-key-secret')))

通过使用相同的volume名称 ssh-key-volume
，我们可以把一切绑定在一起。

def train_and_eval_op(image: str, pvolume: PipelineVolume, data_dir: str, ):

    return dsl.ContainerOp(

        name='training and evaluation',

        image=image,

        command=[CONDA_PYTHON_CMD, f"{PROJECT_ROOT}/train.py"],

        arguments=["--data_dir", data_dir],

        file_outputs={'output': f'{PROJECT_ROOT}/output.txt'},

        container_kwargs={'image_pull_policy': 'IfNotPresent'},

        pvolumes={"/workspace": pvolume}

    )

最后，是时候进行训练和评估这一步骤。这一步唯一的区别在于 file_outputs
变量。如果我们再次查看 train.py
，则有以下代码段：

    with open(os.path.join(PROJECT_ROOT, 'output.txt'), 'w') as f:

        f.write(model_path)

        print(f'Model written to: {model_path}')

我们正在将模型路径写入名为 output.txt
的文本文件中。通常，可以将其发送到下一个流水线组件，在这种情况下，该参数将包含模型的路径。

要指定流水线，你需要使用 dsl.pipeline
来注释流水线功能：

@dsl.pipeline(

    name='Fashion MNIST Training Pipeline',

    description='Fashion MNIST Training Pipeline to be executed on KubeFlow.'

)

def training_pipeline(image: str = 'benjamintanweihao/kubeflow-mnist',

                      repo_url: str = 'https://github.com/benjamintanweihao/kubeflow-mnist.git',

                      data_dir: str = '/workspace'):

    git_clone = git_clone_darkrai_op(repo_url=repo_url)




    preprocess_data = preprocess_op(image=image,

                                    pvolume=git_clone.pvolume,

                                    data_dir=data_dir)




    _training_and_eval = train_and_eval_op(image=image,

                                           pvolume=preprocess_data.pvolume,

                                           data_dir=data_dir)




if __name__ == '__main__':

    import kfp.compiler as compiler

    compiler.Compiler().compile(training_pipeline, __file__ + '.tar.gz')

还记得流水线组件的输出是另一个组件的输入吗？在这里， git clone
、 container_op
的 pvolume
将传递到 preprocess_cp
。

最后一部分将 pipeline.py
转换为可执行脚本。最后一步是编译流水线：

% dsl-compile --py pipeline.py --output pipeline.tar.gz