Toy Machine Learning Pipeline

About

This is a toy example of a standalone ML pipeline written entirely in Python. No external tools are incorporated into the master branch. I built this for two reasons:

1. To experiment with my own ideas for MLOps tools, as it is hard to develop devtools in a vacuum :)
2. To have something to integrate existing MLOps tools with so I can have real opinions

The following diagram describes the pipeline at a high level. The README describes it in more detail.

Getting started

This pipeline is broken down into several components, described in a high level by the directories in this repository. See the Makefile for various commands you can run, but to serve the inference API locally, you can do the following:

1. git clone the repository
2. In the root directory of the repo, run make serve
3. [OPTIONAL] In a new tab, run make inference to ping the API with some sample records

All Python dependencies and virtual environment creation is handled by the Makefile. See setup.py to see the packages installed into the virtual environment, which mainly consist of basic Python packages such as pandas or sklearn.

ML task description and evaluation procedure

We train a model to predict whether a passenger in a NYC taxicab ride will give the driver a large tip. This is a binary classification task. A large tip is arbitrarily defined as greater than 20% of the total fare (before tip). To evaluate the model or measure the efficacy of the model, we measure the F1 score.

The current best model is an instance of sklearn.ensemble.RandomForestClassifier with max_depth of 10 and other default parameters. The test set F1 score is 0.716. I explored this toy task earlier in my debugging ML talk.

Dataset description

We use the yellow taxicab trip records from the NYC Taxi & Limousine Comission public dataset, which is stored in a public aws S3 bucket. The data dictionary can be found here and is also shown below:

Repository structure

The pipeline contains multiple components, each organized into the following high-level subdirectories:

• etl
• training
• inference

Pipeline components

Any applied ML pipeline is essentially a series of functions applied one after the other, such as data transformations, models, and output transformations. This pipeline was initially built in a lightweight fashion to run on a regular laptop with around 8 GB of RAM. The logic in these components is a first pass; there is a lot of room to improve.

The following table describes the components of this pipeline, in order:

Data storage

The inputs and outputs for the pipeline components, as well as other artifacts, are stored in a public S3 bucket named toy-applied-ml-pipeline located in us-west-1. Read access is universal and doesn't require special permissions. Write access is limited to those with credentials. If you are interested in contributing and want write access, please contact me directly describing how you would like to be involved, and I can send you keys.

The bucket has a scratch folder, where random scratch files live. These random scratch files were likely generated by the write_file function in utils.io. The bulk of the bucket lies in the dev directory, or s3://toy-applied-ml-pipeline/dev.

The dev directory's subdirectories represent the components in the pipeline. These subdirectories contain the outputs of each component respectively, where the outputs are versioned with the timestamp the component was run. The utils.io library contains helper functions to write outputs and load the latest component output as input to another component. To inspect the filesystem structure further, you can call io.list_files(dirname), which returns the immediate files in dirname.

If you have write permissions, store your keys/ids in an .env file, and the Makefile will automatically pick it up. If you do not have write permissions, you will run into an error if you try to write to the S3 bucket.

Utils documentation

The utils directory contains helper functions and abstractions for expanding upon the current pipeline. Tests are in utils/tests.py. Note that only the io functions are tested as of now.

io

utils/io.py contains various helper functions to interface with S3. The two most useful functions are:

def load_output_df(component: str, dev: bool = True, version: str = None) -> pd.DataFrame:
  """
    This function loads the latest version of data that was produced by a component.
    Args:
        component (str): component name that we want to get the output from
        dev (bool): whether this is run in development or "production" mode
        version (str, optional): specified version of the data
    Returns:
        df (pd.DataFrame): dataframe corresponding to the data in the latest version of the output for the specified component
    """
    ...

def save_output_df(df: pd.DataFrame, component: str, dev: bool = True, overwrite: bool = False, version: str = None) -> str:
    """
    This function writes the output of a pipeline component (a dataframe) to a parquet file.
    Args:
        df (pd.DataFrame): dataframe representing the output
        component (str): name of the component that produced the output (ex: clean)
        dev (bool, optional): whether this is run in development or "production" mode
        overwrite (bool, optional): whether to overwrite a file with the same name
        version (str, optional): optional version for the output. If not specified, the function will create the version number.
    Returns:
        path (str): Full path that the file can be accessed at
    """
    ...

Note that save_output_df's default parameters are set such that you cannot overwrite an existing file. You can change this by setting overwrite = True.

Feature generators

utils.feature_generators.py contains the lightweight abstraction for a feature generator to make it easy for someone to create a new feature. The abstraction is as follows:

class FeatureGenerator(ABC):
    """Abstract class for a feature generator."""

    def __init__(self, name: str, required_columns: typing.List[str]):
        """Constructor stores the name of the feature and columns required in a df to construct that feature."""
        self.name = name
        self.required_columns = required_columns

    @abstractmethod
    def compute(self):
        pass

    @abstractmethod
    def schema(self):
        pass

See utils.feature_generators.py for examples on how to create specific feature types and etl/featuregen.py for an example on how to create the actual instances of the features themselves.

Models

utils/models.py contains the ModelWrapper abstraction. This abstraction is essentially a wrapper around a model and consists of:

• the model binary
• pointer to dataset(s)
• metric values

To use this abstraction, you must create a subclass of ModelWrapper and implement the preprocess, train, predict, and score methods. The base class also provides methods to save and load the ModelWrapper object. It will fail to save if the client has not added data paths and metrics to the object.

An example of a subclass of ModelWrapper is the RandomForestModelWrapper, which is also found in utils/models.py. The RandomForestModelWrapper client usage example is in training/train.py and is partially shown below:

from utils import models

# Create and train model
mw = models.RandomForestModelWrapper(
    feature_columns=feature_columns, model_params=model_params)
mw.train(train_df, label_column)

# Score model
train_score = mw.score(train_df, label_column)
test_score = mw.score(test_df, label_column)

mw.add_data_path('train_df', train_file_path)
mw.add_data_path('test_df', test_file_path)
mw.add_metric('train_f1', train_score)
mw.add_metric('test_f1', test_score)

# Save model
print(mw.save('training/models'))

# Load latest model version
reloaded_mw = models.RandomForestModelWrapper.load('training/models')
test_preds = reloaded_mw.predict(test_df)

Roadmap

See the open issues for tickets corresponding to feature ideas. The issues in this repo are mainly tagged either data science or engineering.

Contributing

Having a toy example of an ML pipeline isn't just nice to have for people experimenting with MLOps tools. ML beginners or data science enthusiasts looking to understand how to build pipelines around ML models can also benefit from this repository.

Anyone is welcome to contribute, and your contribution is greatly appreciated! Feel free to either create issues or pull requests to address issues.

1. Fork the repo
2. Create your branch (git checkout -b YOUR_GITHUB_USERNAME/somefeature)
3. Make changes and add files to the commit (git add .)
4. Commit your changes (git commit -m 'Add something')
5. Push to your branch (git push origin YOUR_GITHUB_USERNAME/somefeature)
6. Make a pull request

Contact

Original author: Shreya Shankar

Email: [email protected]