Proposal for KubeFlow + Pachyderm Integration

Motivation

From the KubeFlow Perspective

KubeFlow is providing some great standards for the deployment and management of distributed ML/AI tooling on Kubernetes (distributed TensorFlow, TF serving, JupyterHub, etc.). However, as of now, KubeFlow does not have a built in solution for orchestrating ML/AI workflows that utilize these tools. That is, what are users to do if they want to string together pre-processing, feature engineering, training, inference, post-processing, etc. into an automated data pipeline that is sustainable, scalable, and tracked over time.

In addition, vanilla KubeFlow leaves the data management solution up to the individual users and relies on users to integrate their own ingress/egress/tracking/versioning of data in TFJobs, notebooks, etc. This is good in that it provides flexibility for users, but many users will want a scalable and unified way of managing and versioning data sets and connecting those to corresponding processing (e.g., with TF).

From the Pachyderm perspective

Pachyderm is already allowing users to build scalable, tracked, and manageable ML/AI pipelines on Kubernetes. However, for the most part (outside of GPU usage and some custom solutions), these pipelines rely on single node (non-distributed) training (although Pachyderm can automatically parallelize map/reduce style processing for these stages).

Pachyderm is managing production model development workflows on GPU instances, but its users frequently ask how they integrate distributed model training with things like TF distributed or Spark. Thus, Pachyderm needs a way to enable these distributed frameworks to run as stages of a Pachyderm pipeline, thus giving its users access to accelerated training paired with scalable data management and pipelining.

Moreover, anything that can be done to ease the burden of deploying these frameworks and tools on Kubernetes will benefit the Pachyderm ecosystem and its users. For example, easily spinning up TF serving would be useful to Pachyderm users, because they could automatically egress models from Pachyderm to TF serving and manage/version that workflow in a unified way.

Pachyderm + KubeFlow

Pachyderm + KubeFlow will satisfy these needs from both the KubeFlow and Pachyderm perspective. Pachyderm can provide the pipeline/orchestration/data management pieces missing with KubeFlow, and KubeFlow can provide the ease of deployment and distributed framework integrations that are missing in Pachyderm.

Use Cases

I have tested out distributed TF training in a Jupyter notebook, but now I want to hook that training into my ETL pipeline and/or integrate large scale data with its corresponding pre-processing stages.
I have my TF model development workflow running as a pipeline in Pachyderm, but I need to scale my training using TF distributed.
I want to unify my management of model training and deployment/inference in a way that maintains versions of models and provides data provenance for my results.
I want to both easily deploy the various data science, ML, AI frameworks I need to use on k8s and stitch those together in a data pipeline.
I need to automatically trigger model training and/or inference based on new data.
I want to utilize the framework deployed with KubeFlow to also manage large scale data on the same cluster, such that I can ingress/egress data to/from KubeFlow components.

Goals

Trigger TFJobs from Pachyderm
- Allow users to trigger these jobs without modifying their TF code
- Support both distributed and non-distributed training
Implement a mechanism to transfer data to/from TFJobs that are launched by Pachyderm.
Implement a mechanism to egress models trained and versioned in Pachyderm to TF serving, as deployed by KubeFlow.
Deploy Pachyderm in a manner consistent with the rest of the KubeFlow offerings.

Workflow

Fig. 1 Distributed training as part of a Pachyderm pipeline

As shown in Fig. 1, users can manage their data sets in Pachyderm distributed file system, and have various stages of processing (e.g., pre-processing) that run as normal Pachyderm pipelines stages. However, the model training stage of the Pachyderm pipeline would trigger a TFJob via KubeFlow that would run distributed (or not) training via TF. The Pachyderm pipeline stage would wait for this Job to complete and then version any output data from the TFJob (e..g, the serialize *.pb file from TF) as output in Pachyderm. This back and forth of data and the corresponding mechanisms are discussed below. In addition, users could pull in other external data (e.g., from databases, HDFS, etc.) in the TFJob if they wanted (similar to how this could be done in any other Pachyderm pipeline stage).

Fig. 2 Training + Model versioning export

As shown in Fig. 2, users could train model using the mechanism discussed wrt to Fig. 1 (or another mechanism). After a model is trained and versioned, a Pachyderm pipeline stage could be utilized to export this newly trained model to TF serving. The has some huge benefits around automation, as models will be kept up to date in TF serving, and around tracking, as models can be automatically versioned in pachyderm. Models utilized in TF serving can also be tracked back to the exact state of pre-processing, raw data, etc. that produced them.

Components

Fig. 3 Pachyderm + KubeFlow Integration Components (proposed components shown in green)

The design has several components (some of which are pictured above, and some of which are configuration/templates):

A Pachyderm KubeFlow worker image, utilized as the image for Pachyderm pipelines that utilize KubeFlow components
A Pachyderm TFJob Sidecar
A Pachyderm ksonnet package
Pachyderm template pipeline specs
- A template spec for utilizing the KubeFlow worker image to do training
- A template spec for utilizing the KubeFlow worker image to do export a model for serving

KubeFlow Worker Image

The KubeFlow Worker Image will be a Docker image including a binary for launching TFJobs and exporting models to TF serving. This binary will be included in Pachyderm’s source repo and the image will be integrated into Pachyderm CI/CD.

The purpose of the KubeFlow Worker Image, as enabled by the included binary (likely a Go program) is to:

Launch and manage TFJobs via the Kubernetes API
- Create a TFJob based on a user supplied manifest
- Wait for the TFJob to finish or fail
- Gather results and logs from the TFJob, which will be made available and will be versioned in Pachyderm.
Export TF models to TF Serving

By having this officially supported KubeFlow worker, users of Pachyderm will be able to:

Develop TFJobs as documented in KubeFlow documentation
Trigger those TFJobs from Pachyderm, without having to modify the TF code (e.g., by injecting references to Pachyderm)
Version and manage the output of TFJobs via Pachyderm data versioning
Both version models and easily export them to TF serving, without having to write extra code to enable that export step

An example of the functionality required in this worker has been created here. In addition to that creating, waiting, and gathering functionality, this worker image will, for training jobs, likely also need to read and slightly modify the user supplied TFJob manifest, to include:

The Pachyderm TFJob Sidecar Container (further discussed below) in each of the TF master, worker, and ps pods.
Environmental variables for the sidecar containers that reference the input data repository that is being processed. That way the sidecar containers can know what data to shim into the TFJob pods. The environment variable will include input repo name and branch at the minimum.

Pachyderm TFJob Sidecar Container

There are a couple of use cases in terms of the data that would be required to run a TFJob:

The training data comes from a Pachyderm data repository that is input to the KubeFlow worker pipeline stage.
The training data or other data is pull from data stores outside of Pachyderm (HDFS, databases, object store buckets, etc.)

In either case, the TFJob can be triggered the same way. The KubeFlow worker image will create the TFJob using a modified version of the user manifest that includes the Pachyderm TFJob sidecar containers in each of the master, worker, and ps pods. The Pachyderm TFJob sidecar will then read in the environmental variable set by the KubeFlow worker (Pachyderm input repo name and branch) and then shim in the data from the input repo via a Pachyderm language client to the pod.

In the first case above (1), this will allow the TFJob to have access to input data from Pachyderm, without having to use any Pachyderm components explicitly. In the second case above (2), the data that is shimmed in is likely irrelevant and the user can just ignore it. This is consistent with the general Pachyderm philosophy of data in/out. It’s up to the user whether they want to pull data from the Pachyderm input or not.

If the node is the master TF node, the Pachyderm sidecar will also take care of collecting output from the TFJob that is meant to go back to Pachyderm. This will be done by writing to a specified directory (e.g., /pfs/out as in a normal Pachyderm pipeline). Once the TF code completes and the user container is exiting, the Pachyderm sidecar can gather whatever is in that directory and:

use the Pachyderm client to commit this output back into Pachyderm as a “partial result” in the Kubeflow worker pipeline’s output repo.
the Kubeflow worker would then, by default, gather all partial results in the output repo, including the ones from the sidecar as the pipeline’s final output (which is what happens now when Pachyderm is running multiple workers in a job).

(Note an alternative or future direction here might be to exploit k8s volumes to facilitate this interaction)

Pachyderm ksonnet package

KubeFlow users should be able to deploy Pachyderm as they would other components of the KubeFlow ecosystem. This would look like:

$ ks pkg install kubeflow/pachyderm

Possible with some arguments or setting a cloud/nocloud option. Some inspiration can be taken from the Helm chart for Pachyderm.

Pachyderm pipeline templates

We should also provide pipeline templates for KubeFlow/TFJob/TF serving, so users can get up and running quickly. This would look like the following for training (where pachkf is the example binary included in the KubeFlow Worker, further discussed above):

{
  "pipeline": {
    "name": "model_training"
  },
  "transform": {
    "image": "pachyderm/kubeflow-worker",
    "cmd": [ "pachkf", “tfjob”, “/pfs/tfjob/tfjob.yaml” ]
  },
  "input": {
    "cross": [
      {
        "atom": {
          "repo": "tfjob",
          "glob": "/"
        }
      },
      {
        "atom": {
          "repo": "training_data",
          "glob": "/"
        }
      }
    ]
  }
}

(note, the TFJob manifest could also be shim in directly from a git repo using Pachyderm git repo inputs)

The pipeline spec for serving would look like:

{
  "pipeline": {
    "name": "model_export"
  },
  "transform": {
    "image": "pachyderm/kubeflow-worker",
    "cmd": [ "pachkf", “tfserving”, “/pfs/model/frozen_model.pb” ]
  },
  "input": {
    "atom": {
      "repo": "model",
      "glob": "/"
    }
  }
}

Moving forward and discussion

We look forward to further discussion and fine tuning of this proposal here. @dwhitena is also in KubeFlow Slack, and Pachyderm has their own public Slack for discussion: http://slack.pachyderm.io/

We also propose that, to expedite this integration, the Pachyderm team starts working on the KubeFlow worker image and pipeline templates, and the KubeFlow team starts working on some of the functionality for the Pachyderm sidecar (based on their TFJob expertise). Then we can take care of the ksonnet package collaboratively.

Pachyderm and Kubeflow integration #151

Description

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions