How to write Pytest code for reliably testing your data science pipeline

Testing software modules enhances robustness and trust in the final product/ service. The importance of high-quality testing cannot be emphasized enough in any software development.

The same argument goes for your data science pipeline too. Even if you are developing a data science codebase mainly for prototyping and research, it’s a good idea to know how to write basic testing modules to check if the functions and classes you are developing are working as they are expected to.

In fact, I discussed this in detail in one of my previous articles. There I showed that it’s not about checking the input data type. It’s about checking whether your data science pipeline can handle that. It’s not just randomly throwing out-of-range variables at the function. It’s about checking whether the response is as expected.

PyTest for Machine Learning — a simple example-based tutorial

In this short article, we will look at an example of a typical data science pipeline (consisting of small dedicated functions), instead of a monolith (as shown in the article above), and see how to write a Pytest module for that. We will also look at why writing test modules for data science can be slightly different from that what software engineers or Quality Assurance (QA) folks do every day.

The Data Science Pipeline

Our pipeline is going to be short and sweet.

Preprocess

Here is the code for preprocessing. It just takes in a Pandas DataFrame and returns training features and targets as Numpy arrays extracted from that DataFrame. We assume that the last column of the DataFrame contains the target values and the rest of the columns are features.

def preprocess(df, 
               test_size=0.25,
               random_state=101):
    """
    Precrocesses data
    Loads from a dataframe, test/train split
    Returns as Numpy arrays
    """
    X = df[df.columns[:-1]]
    y = df[df.columns[-1]]
    
    X_train, X_test, y_train, y_test = train_test_split(X, y, 
                         test_size=test_size, 
                         random_state=random_state)
    
    X_train = np.array(X_train)
    y_train = np.array(y_train)
    X_test = np.array(X_test)
    y_test = np.array(y_test)
    return (X_train, X_test, y_train, y_test)

Model Fit

The model fitting code is even shorter. But it does have a try-except block and a couple of assertion checks to make sure that the input type is a Numpy array. We will see how Pytest uses these checks.

def training(X_train,y_train):
    """
    Trains a model
    """
    try:
        assert isinstance(X_train,np.ndarray),"Features must be a Numpy array"
        assert isinstance(y_train,np.ndarray),"Targets must be a Numpy array"
        lm = LinearRegression()
        lm.fit(X_train,y_train)
        return lm
    except AssertionError as msg:
        return msg

Model Save

Then we have a model-saving code that saves the fitted model to a file using the joblib module. It also adds date-time versioning to the filename.

def save_model(model, fname = 'linear-model'):
    """
    Saves the model with datetime adder (versioning)
    """
    try:
        now = datetime.datetime.now()
        dt_adder = str(now.month)+str(now.day)+str(now.year)+'_'+\
            str(now.hour)+str(now.minute)
        fname = fname+'_'+dt_adder+'.sav'
        dump(model, fname)
        return 1
    except:
        return 0

Load and Predict

Finally, we have a short function for loading the joblib object (from a given filename) and making predictions for some given data. It checks the data dimension to make sure that the model can work with that i.e. the dimensions are compatible.

def predict(model,data):
    """
    Predicts for some data using the model (give as the filename)
    """
    try:
        lm = load(model)
    except:
        return 0
    assert len(lm.coef_) == data.shape[1],"Data dimension is not suitable for this model"
    preds = lm.predict(data)
    
    return preds

That’s it for the actual pipeline code. We put all of these functions into a single file called pipeline.py and move on to write the actual test code.

Existing ML applications may surprise you — watch our interview with GE Healthcare’s Vignesh Shetty to learn how his team is using ML in the healthcare setting.

The Test Code for the Pipeline

We will use a lightweight but powerful Python testing library, Pytest, for this demo. You can read more about it here.

The way Pytest works is that we create a test script (a Python file) and run it with the following command:

pytest test_pipeline.py

Here, we assumed that our test script is named test_pipeline.py for the obvious reason that it is testing the code in pipeline.py.

The types and variety of the test are only limited by your imagination and business needs. We just show some simple examples with the first two modules of the data pipeline here. An illustration is shown below:

The fundamental principle to remember is that you should write test code only for the specific function that you want to test i.e. test what that function is supposed to do. This means for the preprocess function — as it is written today — do not try to check its input data type yet. That should be handled by the function code and not the external test code.

So, our test_pipeline.py starts by loading the main script and pytest, and then it imports other libraries as needed.

from pipeline import *
import pytest
import...

Testing the Preprocess Function

Here is the code for testing the preprocess function only. Note that the test functions must have a name starting with test_.

We start by generating some random data and using the preprocess function on this data. We also have a test case for a specific test-set (related to test/train split, not Pytest) fraction of 0.5.

def test_preprocess():
    """
    Tests the returned object from preprocessing
    """
    # Generate a random dataset
    df = random_data_constructor()
    # Use the preprocess function and get the returned object
    preprocess_return = preprocess(df)
    # Same as above with a specific test_size of 0.5
    preprocess_return_test_size = preprocess(df,test_size=0.5)

Then, we write the actual check code. It checks

• That the returned data type is a tuple
• The tuple has 4 objects
• Each of those objects is NumPy array type

# Type check
assert isinstance(preprocess_return, tuple)
# Length check
assert len(preprocess_return) == 4
# Tuple element type check
assert isinstance(preprocess_return[0],np.ndarray)
assert isinstance(preprocess_return[1],np.ndarray)
assert isinstance(preprocess_return[2],np.ndarray)
assert isinstance(preprocess_return[3],np.ndarray)

We also check if the test/train splits dimensions match up properly.

# Test size fraction check
    preprocess_return_test_size[0].shape == preprocess_return_test_size[2].shape
    preprocess_return_test_size[1].shape == preprocess_return_test_size[3].shape

The next step is interesting. Pytest provides context managers that can accept exceptions like ValueError or ZeroDiviosnError and check if the function actually raised that. In this case, we pass invalid test set fractions (-0.5 or 1.5) and expect the Scikit-learn function to raise a ValueError. We test for that.

# (Invalid) test size exception check
    with pytest.raises(ValueError):
        preprocess_return_invalid_test_size = preprocess(df,test_size=-0.5)
        preprocess_return_invalid_test_size = preprocess(df,test_size=1.5)

This is where testing for data science code needs data science knowledge. As a data scientist, you need to know how to break the function and also what specific kind of exception/error will be generated. To pass the Pytest screen, you have to incorporate that in the test code.

Testing the Model Training Function

Again, we start by generating fresh random data, preprocessing it, and getting a trained model back by passing the preprocessed data into the training function.

def test_training():
    """
    Tests the training function
    """
    # Generate a random dataset
    df = random_data_constructor()
    # Use the preprocess function
    X_train, X_test, y_train, y_test = preprocess(df)
    # Get back the trained model
    trained_model = training(X_train,y_train)

Then, the primary tests we do are type checking. First, type check for the returned object.

# Type check for the returned object
assert isinstance(trained_model, sklearn.linear_model._base.LinearRegression)

Then, we pass on an invalid data type for the X input. The function is supposed to check this internally and return a specific AssertionError as message. Note how we are checking the content of that message here using Pytest.

# AssertionError check for wrong input type
msg = training('X',y_train)
assert isinstance(msg, AssertionError)
assert msg.args[0] == "Features must be a Numpy array"

We do the exact same thing for the y vector.

msg = training(X_train,'y')
assert isinstance(msg, AssertionError)
assert msg.args[0] == "Targets must be a Numpy array"

One thing is to be noted with care: Pytest did not check for the X or y data type, the main training function did. Pytest is only checking whether the training function is doing its job properly by checking the returned message.

Running the Tests and Next Steps

As we said before, we run these tests on the CLI.

pytest test_pipeline.py

And, if everything goes well, we may get this. There were two functions starting with test_

• test_preprocess
• test_training

And they passed — all the assert statements inside them yielded True.

The Rest is Up to You

You may be tempted to test the training function for all kinds of predictions with all kinds of data. Don’t do that here. Do those tests in the predict function, because that’s where the prediction is coming into the picture. This is how you maintain the modularity and compartmentalization of your code and software.

Loved the article? Become a Medium member to continue learning without limits. I’ll receive a portion of your membership fee if you use the following link, with no extra cost to you.

Join Medium with my referral link - Tirthajyoti Sarkar

Editor’s Note: Heartbeat is a contributor-driven online publication and community dedicated to providing premier educational resources for data science, machine learning, and deep learning practitioners. We’re committed to supporting and inspiring developers and engineers from all walks of life.

Editorially independent, Heartbeat is sponsored and published by Comet, an MLOps platform that enables data scientists & ML teams to track, compare, explain, & optimize their experiments. We pay our contributors, and we don’t sell ads.

If you’d like to contribute, head on over to our call for contributors. You can also sign up to receive our weekly newsletters (Deep Learning Weekly and the Comet Newsletter), join us on Slack, and follow Comet on Twitter and LinkedIn for resources, events, and much more that will help you build better ML models, faster.

How to Write Test Code for Data Science Pipeline was originally published in Heartbeat on Medium, where people are continuing the conversation by highlighting and responding to this story.

As a data scientist, a common question to expect from your engineering/ platform team is “how much memory footprint does your model/code have?” or “what is the execution time of the various data science (DS) modules when running with some given data load?”

This is quite a natural query and there could be a number of reasons behind it:

• Hardware resources may be limited.
• A single ML or DS module should not hog all the memory or CPU resources of the system.
• The stability of the overall platform may depend on how the individual DS modules are using the CPU resources.
• Input data characteristics may change frequently that may render DS modules less performant e.g., a simple data cleaning routine scaling badly when the data velocity goes from gigabytes/hour to terabytes/hour.

Some of these concerns are particularly critical for edge computing scenarios where the ML app may be running inside a virtualized container on an industrial PC (and with no cloud-supported auto-scaling of memory or dynamic allocation).

Also, your model may just be one of the hundreds of ML models running on that piece of hardware or platform. Therefore, you must have some idea about the peak memory/ CPU usage of the model because if multiple models peak in their usage at the same time, that can crash the whole system.

All of these models do not necessarily come from the same data scientist either. Various teams might have developed them over time. It makes sense to have a common mechanism of measuring or gauging the memory usage (peak and average) and CPU execution profile of all those models (and data science code in general).

Therefore, in this article, we will learn how to time-profile a typical piece of data science code using a built-in Python module cProfile. This knowledge will help data scientists in all sorts of projects and tasks where they have to measure the computational load of their code or modules. For a team, such a tool can help compare and benchmark the optimal code/function that satisfies some timing requirements and find out inefficiencies along the way.

Memory Profile Using “Scalene”

In a previous article, I discussed in detail memory profiling a typical ML workflow code using a nifty little library called Scalene.

How Much Memory is your ML Code Consuming?

In that article, I showed how the memory profile may vary widely between ML model types and how they scale with data size. It was shown that external I/O may take up more memory than the ML model if the parametric size of the model is modest. In the current article, we focus on basic CPU time profiles (as taken up by individual functions or even lines of code) and how they can be gauged using a built-in Python profiler.

Timing Profile with cProfile

In one of my previous articles, I showed basic tricks and techniques to measure the execution time of simple code blocks or a function with a timing decorator.

Now we will discuss a built-in Python library cProfile that can give us a detailed timing profile about the various parts of our data science code with a simple command. The advantage is that we don’t have to insert code snippets like time.time() in various places and track them manually.

Basic Usage

The cProfile library comes with the default Python installation so there is nothing to install. Let us see the basic usage with a simple code where we are adding two Numpy arrays.

import numpy as np
import cProfile

# Two Numpy arrays
SIZE = 10_000_000
a = np.arange(SIZE)
b = np.random.normal(size=SIZE)

# Profling array addition
cProfile.run('a+b')

So, the main thing to notice is that we have to wrap the code within a string and pass it on to the cProfile.run function. Here the code is simply ‘a+b’. The output may look something like the following (the exact time will vary, of course, depending on the underlying hardware).

The interesting thing to remember is that the only piece of code that was measured for timing is the snippet — “a+b”. The array creation is not being measured here, only the addition.

If we wanted to measure the timing profile of all the steps, we could write:

code = 
"""
SIZE = 10_000_000
a = np.arange(SIZE)
b = np.random.normal(size=SIZE)
a+b
"""

cProfile.run(code)

Here we put all the code inside the string variable code and then pass that on to the cProfile.run function. The output looks different, as expected.

Note the extra array creation operations resulted in a total of five function calls, as opposed to three for the basic addition code.

With a Function as the Argument

We could, of course, create a standalone function and pass the name of that object to cProfile.run function for the same task.

def add():
    SIZE = 10_000_000
    a = np.arange(SIZE)
    b = np.random.normal(size=SIZE)
    c=a+b

cProfile.run('add()')

The output is similar to the figure above, but an additional function call is registered that comes from the construction of the add function itself.

The function that is passed on to cProfile can have any arguments as well. In many cases, we can change the arguments and see the impact on the profile results. This is one of the most obvious use-cases of the library. We rewrite the add function to accept a size argument.

def add(size):
    a = np.arange(size)
    b = np.random.normal(size=size)
    c=a+b

Siloed workflows no longer cut it. An MLOps strategy can bring your team together and get models out the door faster. Watch our webinar with industry experts to learn more about the power of MLOps.

Then, we can use that to profile the array operations with 10-million elements.

SIZE = 10_000_000
cProfile.run('add(SIZE)')

We get this as the output:

When we change the number of elements to 20 million, it reflects immediately:

Using the Profiler Class

cProfile has a special Profiler class that stores all the important information and it can be enabled/disabled programmatically. We can also use the pstats library and pass this Profiler object to it for printing and extracting data. Here is the code to measure and print the total execution time:

import cProfile, pstats

profiler = cProfile.Profile()

# Enable profiler
profiler.enable()

# Function execution
add(SIZE)

# Disable profiler
profiler.disable()

# pstats
stats = pstats.Stats(profiler)

# Print the total time and function calls
print("Total function calls:", stats.total_calls)
print("Total time (seconds):", stats.total_tt)

The result will look something like this:

>> Total function calls: 48
>> Total time (seconds): 1.1797575999999999

Here we get the execution time from the total_tt attribute of the stats object and the number of calls from the total_calls attribute.

Consequently, this opens up the possibility to programmatically control the profiling and storing of the information as needed. For example, we can profile the execution of the same add function over a range of array sizes.

size = [int(i*1e6) for i in range(5,26,5)]
total_tt = []

for s in size:
    profiler = cProfile.Profile()
    profiler.enable()
    add(s)
    profiler.disable()
    stats = pstats.Stats(profiler)
    total_tt.append(round(stats.total_tt,3))

So, the timings are stored in the total_tt array. When plotted, it shows the expected pattern.

Considerations for DS Workflow

While measuring the execution time of these small standalone functions serves as a basic demonstration of the usage of these profilers, the real utility is realized when they are used in a large-scale data science workflow.

Final Validation can be Useful

For low-resource situations, it could be a good idea to host a validation environment/server which will accept a given modeling code (when developed) and run it through such a memory profiler to create runtime statistics. If it passes pre-determined criteria of memory footprint, only then the modeling code will be accepted for further deployment.

Measure all Modules

Such a workflow has a variety of modules and functions and we can set up profiling for all of them if necessary. The output may be logged into a database or even be fed into a monitoring system that will track the performance of the modules over time and take action if needed (e.g., a function has poorly performed and took too much time in a certain run or for a certain input data).

Summary

In this article we discussed the importance of memory and time profiling your data science/machine learning code for smooth and easy interfacing with the platform/engineering team. This is expected to ease the deployment of the code on a service/machine.

We showed the basic usage of the built-in Python timing profiler cProfile and how it can be used with raw code or function modules. We demonstrated how to extract the total execution time or the number of function calls using this library for Numpy operations. The idea can be extended to any data science workflow that consists of many stages and modules.

Note: The basic code is contained in this Jupyter Notebook on my Github repo.

Every data science team or organization may have its own style for measuring code and module efficiency and memory footprint. The motivation for this article was to introduce the readers to their importance and show some hands-on examples so that they can explore further to be ready for such implementation.

Loved the article? Become a Medium member to continue learning without limits. I’ll receive a portion of your membership fee if you use the following link, with no extra cost to you.

Join Medium with my referral link - Tirthajyoti Sarkar

Editor’s Note: Heartbeat is a contributor-driven online publication and community dedicated to providing premier educational resources for data science, machine learning, and deep learning practitioners. We’re committed to supporting and inspiring developers and engineers from all walks of life.

Editorially independent, Heartbeat is sponsored and published by Comet, an MLOps platform that enables data scientists & ML teams to track, compare, explain, & optimize their experiments. We pay our contributors, and we don’t sell ads.

If you’d like to contribute, head on over to our call for contributors. You can also sign up to receive our weekly newsletters (Deep Learning Weekly and the Comet Newsletter), join us on Slack, and follow Comet on Twitter and LinkedIn for resources, events, and much more that will help you build better ML models, faster.

Time-Profiling Data Science Code was originally published in Heartbeat on Medium, where people are continuing the conversation by highlighting and responding to this story.