Data pipelines with Spotify's Luigi
Tags: python, datascience, nlp, luigi
Published: 2021-12-22 17:06:05

Hi!

At [Wonderflow](https://www.wonderflow.ai/) we're doing a lot of ML / NLP using Python and recently we are enjoying writing data pipelines using Spotify's [Luigi](https://github.com/spotify/luigi).

It's quite easy to find online [lots](https://www.integrate.io/blog/airflow-vs-luigi/) [of](https://towardsdatascience.com/data-pipelines-luigi-airflow-everything-you-need-to-know-18dc741449b7) [posts](https://www.datarevenue.com/en-blog/airflow-vs-luigi-vs-argo-vs-mlflow-vs-kubeflow) about Luigi, especially comparisons with other tools like [Apache AirFlow](https://airflow.apache.org/).

In this post I'll try to explain how we're using Luigi at Wonderflow, what issues we found and how we handled them.

## Stateless data pipelines are bad

When it comes to do computationally expensive analysis on data, state management is very important. 

If you have to do a generic computation wrapped in a method like the following:


```python
def compute():

    # step 1
    data = fetch_data(data, params)

    # step 2
    data_1 = heavy_computation_1(data, a=1, b=2, c=3)

    # step 3
    data_2 = heavy_computation_2(data_1, x=1, y=2, z=3)

    return data_2
```

There are 3 main steps here and 2 of them are quite heavy to compute. 

Imagine now that you are tuning step 3 `x`, `y` and `z` parameters. And you find yourself running this method over and over with different parameter values and check every time if the output suits you.

This might seem the fastest way to do it, but you're actually wasting a lot of time: to check `heavy_computation_2` result you're continuously run `fetch_data` and `heavy_computation_1` on (probably) the same data. 

Although you can apply workarounds, even if you end up with a complete pipeline scheduled by an orchestrator, you might end up doing many useless things (i.e. doing the same things on the same data many times)

We ran into this kind of issues and we need to find a clean workaround.

## Luigi to the rescue

As you may already found out, Luigi has been designed to solve issues like the one above. 

Let's rewrite the task as a Luigi pipeline:

```python
import luigi
from luigi.local_target import LocalTarget
from fetching import fetch_data
from processing import heavy_computation_1, heavy_computation_2
from json import dumps, loads

class FetchData(luigi.Task):
    def run(self):
        data = fetch_data()

        with self.output().open("w") as outfile:
            outfile.write(dumps(data))

    def output(self):
        return LocalTarget("output/data.json")

class HeavyComputationOne(luigi.Task):
    a = luigi.IntParameter(default=42)
    b = luigi.IntParameter(default=42)
    c = luigi.IntParameter(default=42)

    def requires(self):
        return FetchData()

    def run(self):
        with self.input().open('r') as infile:
            input_data = loads(infile.read())

        data = heavy_computation_1(input_data, a=self.a, b=self.b, c=self.c)

        with self.output().open("w") as outfile:
            outfile.write(dumps(data))

    def output(self):
        return LocalTarget(f"output/processed_data_1-{self.a}-{self.b}-{self.c}.json")


class HeavyComputationTwo(luigi.Task):
    x = luigi.IntParameter(default=42)
    y = luigi.IntParameter(default=42)
    z = luigi.IntParameter(default=42)

    def requires(self):
        return HeavyComputationOne(a=1, b=2, c=3)

    def run(self):
        with self.input().open('r') as infile:
            input_data = loads(infile.read())

        data = heavy_computation_2(input_data, x=self.x, y=self.y, z=self.z)

        with self.output().open("w") as outfile:
            outfile.write(dumps(data))

    def output(self):
        return LocalTarget(f"output/processed_data_2-{self.x}-{self.y}-{self.z}.json")
```

Luigi [Task](https://luigi.readthedocs.io/en/stable/tasks.html) structure is quite simple, you simply write a Python (sub)class with 3 methods:

- `requires()` which defines task requirements (if any)
- `run()` which contains the task business logic
- `output()` which defines the kind of output (in this case files on local filesystem, but they could be e.g. files on S3 [or many more things](https://luigi.readthedocs.io/en/stable/api/luigi.contrib.html))

As you can see, each task has a **defined output file** which is reflecting input parameters. This is needed since in Luigi each task is **idempotent**. 

It means that Luigi **will not run the same task with same output twice**.

Now you can run the `HeavyComputationTwo` from CLI (yes, [Luigi will generate automatically the CLI for you](https://luigi.readthedocs.io/en/stable/running_luigi.html)) task N times, using different values of `x`, `y` and `z` parameters. Luigi will run the requirements (`HeavyComputationOne` and `FetchData`) only once.

So no more wasting time running unneeded tasks :-) And you can focus on tuning your task's logic.

Moreover, this makes the pipeline _stateful_, fixing the issue we had with the `compute()` method.

After trying Luigi a bit, we initially tough we found the silver bullet for our needs, but as everything it has pros and cons. 

Let's go through some of them.

## 1. Luigi pipelines are hard to test

You may have noticed that Luigi is heavily coupled with your code (e.g. the main class task is a subclass of a Luigi's class). Mainly for this reason, lots of people are saying that Luigi pipelines are hard to test.

Since we're doing [TDD](https://en.wikipedia.org/wiki/Test-driven_development) on almost any kind of software in Wonderflow, we wanted to make sure that TDD is possible even using Luigi.

I think it's true that Luigi is heavily coupled with your code, **but I don't think this makes pipelines or tasks hard to test**. 

Here's the testing pattern we're using.

When you usually write a test, you follow the [Given-When-Then](https://en.wikipedia.org/wiki/Given-When-Then) structure:

```python
def test_something():

# Given an initial situation

# When something happens

# Then there's an outcome
```

Imagine we want to test the `FetchData` task (ideally we write the test first!):

```python
import luigi
from tasks import FetchData
from json impo

def setup_db():
   # here we setup a local test DB with test data
   pass

def test_fetch_data_task():
    setup_db()
    task = FetchData()

    assert luigi.build([task], local_scheduler=True)
    assert task.output().exists()

    data = loads(task.output().open("r").read())
    
    assert len(data) > 0
    # other assertions (if needed)
```

To recap:

1. We setup a local database and instantiate the FetchData task
2. We run the task and ensure the output exists
3. We read the task's output and make needed assertions on it

This is quite effective and nice to read, perfectly following the Given-When-Then pattern.

You may have noticed that this is an **integration test**, which is indirectly covering the `fetch_data()` method.

As a best practice, you need to do unit test on all the methods you're using inside the main `run()` one, since they also can be used outside a Luigi task. 

So, since on `FetchData` you're calling `fetch_data()` you need to unit-testing it as well :-)

If all the methods called inside the task are unit tested, and you test the whole task (or pipeline, if task has requirements) as described above, you're basically doing a nice unit + integration test combo, which is usually enough to avoid code regression.

## 2. Chaining tasks leads to parameters repetition

Considering the pipeline above:

`FetchData` -> `HeavyComputationOne` -> `HeavyComputationTwo`

You may have noticed that I've hardcoded `a`, `b` and `c` params on `HeavyComputationTwo` task requirements. 

So, if you have to run `HeavyComputationTwo` but you need e.g. to specify different `a`, `b` or `c` values other than `x`, `y` and `z` ones, do you need to declare again the `HeavyComputationOne` parameters again?

```python
class HeavyComputationTwo(luigi.Task):
    a = luigi.IntParameter(default=42)
    b = luigi.IntParameter(default=42)
    c = luigi.IntParameter(default=42)

    x = luigi.IntParameter(default=42)
    y = luigi.IntParameter(default=42)
    z = luigi.IntParameter(default=42)

    # ...
```

The answer is: **of course not**. If you do that, the code will be a lot less readable and if params changes, you have to change them in every chained task.

Fortunately, Luigi comes with a very handy [inherits](https://luigi.readthedocs.io/en/stable/api/luigi.util.html) method from its `util` module.

So you simply have to decorate the `HeavyComputationTwo` method specifying that it inherits from `HeavyComputationOne`, in this way:

```python
from tasks import HeavyComputationOne
from luigi.util import inherits

@inherits(HeavyComputationOne)
class HeavyComputationTwo(luigi.Task):
    x = luigi.IntParameter(default=42)
    y = luigi.IntParameter(default=42)
    z = luigi.IntParameter(default=42)

    # ...
```

So, `a`, `b` and `c` are now inherited from `HeavyComputationOne`, and there's no need to re-declare them anymore.

## 3. Luigi doesn't scale

As described [in the docs](https://luigi.readthedocs.io/en/stable/design_and_limitations.html), there are some limitations by design.

According to Luigi's [execution model](https://luigi.readthedocs.io/en/stable/execution_model.html), authors say that Luigi scheduler could hang if you scheduled more than a few thousands of jobs at once.

It's also true that Luigi is designed to handle well _batch processing_ jobs. So if you need to run _lots_ of pipelines on an almost real-time situation, it might be not the best tool to use. Better to read carefully the docs first!

We are running only batch processing pipelines, mainly related to ML / NLP tasks and it suits us quite well.

If you really need to launch many tasks at once, there's also a [Kubernetes Job wrapper](https://luigi.readthedocs.io/en/stable/api/luigi.contrib.kubernetes.html) plugin available.

## 4. There's no in-task parallelism

If having N different workers does not suits you and the specific task is very complex in term of computation, you may delegate the required parallelism to other tools (if `multiprocessing` is not enough.

We have tasks which actually require lots of different [Spacy language models](https://spacy.io) to be loaded at once, and we load them on many processes at once.

To do that, we are efficiently using [Dask](https://dask.org/), simply creating on-demand local (or remote) clusters on task `run()` method:

```python
import luigi
import dask.bag as db

class HeavyTask(luigi.Task):
    def run():
    # ...

    with create_local_cluster(memory_limit="4GB", n_workers=6) as cluster, Client(cluster):
        bag = db.from_sequence(data)

        output +=bag.map_partitions(heavy_method, arg1=arg1)
            .compute()
        )
```

Where `create_local_cluster()` simply wraps the related Dask method:

```python
from dask.distributed import LocalCluster

def create_local_cluster(n_workers: int = 4, **kwargs) -> LocalCluster:
    return LocalCluster(
        n_workers=n_workers,
        worker_dashboard_address=None,
        **kwargs,
    )
```

With Dask you can create [local, SSH, or Kubernetes-based clusters](https://docs.dask.org/en/stable/how-to/index.html) according to your needs. It may require a bit of tuning but it's more than enough to handle parallelism in our common use cases.

## 5. There are no triggers for pipelines

This is true, but I would not necessarily see this as a drawback. There are many options here like:

- API
- Crontab
- SSH (e.g. a Jenkins pipeline which fires a Luigi pipeline via SSH)

## Conclusions

If you have a small team which is already proficient in Python, Luigi's relatively short learning curve and benefits could be ideal for your data pipelines, especially if you mainly run batch processing jobs.

Moreover, configure and deploy the Luigi's Scheduler on a server / pod for production use is easy, while it might be not for other similar tools like [Apache AirFlow](https://airflow.apache.org/).