gokart is a wrapper for the luigi data pipeline library designed specifically for Machine Learning pipelines. It focuses on solving reproducibility, managing task dependencies, enforcing good code constraints, and providing ease of use for creating ML batches.

Key features include:

• Reproducibility: Automatically stores task output data, module versions, processing time, random seeds, logs, and task parameters in separate .pkl files indexed by hash.
• Automatic Reruns: If a task parameter changes, the hash changes, triggering spontaneous reruns of the task and all its dependents.
• Memory Efficiency: Intermediate data is exchanged via files (supporting GCS or S3) rather than keeping everything in memory.
• Data Integrity: Provides pandas.DataFrame type and column checking during I/O.
• Concurrency Control: Tasks are locked via Redis to prevent conflicts during parallel execution.
• Automatic Seeding: Automatically fixes seeds for numpy and random.

The Mypy plugin provides two main type-safety benefits for gokart tasks:

1. Parameter Validation: It ensures that when instantiating a task, the arguments passed match the annotated types of the luigi.Parameter or luigi.IntParameter fields. Note: All parameters must be annotated for this to work.
2. Generic Type Inference: It checks gokart.TaskOnKart[T] generic types. When you use self.load(task_instance_parameter), Mypy correctly infers the return type based on the generic type of the parameter (e.g., TaskOnKart[int] results in an int being returned by .load()). It also validates that the tasks passed to TaskInstanceParameter match the expected generic type.

import gokart
import luigi

class Foo(gokart.TaskOnKart):
    # All parameters must be annotated
    foo: int = luigi.IntParameter(default=1)
    bar: str = luigi.Parameter()

Foo(foo=1, bar='2')   # OK
Foo(foo='1')          # Mypy error: foo is not int and bar is missing

class SampleTask(gokart.TaskOnKart):
    str_task: gokart.TaskOnKart[str] = gokart.TaskInstanceParameter()
    int_task: gokart.TaskOnKart[int] = gokart.TaskInstanceParameter()

    def requires(self):
        return dict(str=self.str_task, int=self.int_task)

    def run(self):
        s = self.load(self.str_task)  # Inferred as str
        i = self.load(self.int_task)  # Inferred as int

# Mypy validates that the tasks provided match the generic requirements
SampleTask(
    str_task=StrTask(),  # OK
    int_task=StrTask(),  # Mypy error: expected TaskOnKart[int]
) 

TaskOnKart is the primary class for defining tasks in gokart. It inherits from luigi.Task and provides type-safe methods for loading and dumping data. You define a task by specifying its type (e.g., TaskOnKart[str]), its requirements via requires(), its output location via output(), and its logic in run().

By default, the output file path is inferred from the script's __name__, so the output() method can often be omitted.

import luigi
import gokart

class TaskA(gokart.TaskOnKart[str]):
    param: luigi.Parameter = luigi.Parameter()

    def output(self):
        return self.make_target('output_of_task_a.pkl')

    def run(self):
        results = f'param={self.param}'
        self.dump(results)

class TaskB(gokart.TaskOnKart[str]):
    param: luigi.Parameter = luigi.Parameter()

    def requires(self):
        return TaskA(param='world')

    def output(self):
        return self.make_target('output_of_task_b.pkl')

    def run(self):
        output_of_task_a = self.load()
        results = f'Task A: {output_of_task_a}\nTaskB: param={self.param}'
        self.dump(results)

if __name__ == '__main__':
    print(gokart.build([TaskB(param='Hello')]))

There are two primary ways to force a task to rerun in gokart, depending on whether you want to preserve the output file hash or generate a new one.

1. Using the rerun parameter

You can set rerun=True on a task instance. This forces the task to execute even if the output exists, but it does not change the hash value of the output file. This is useful for quick re-executions without changing the downstream dependency chain.

class Main(GokartTask):
    def requires(self):
        return StringToSplit(task=Sample(rerun=True), rerun=True)

2. Changing task parameters

By adding or changing a parameter (e.g., a luigi.Parameter) on a task, you change its hash value. This causes the task to generate a new output file with a different name, which in turn triggers all dependent tasks to rerun because their inputs have changed.

import luigi

class Sample(GokartTask):
    version: luigi.IntParameter = luigi.IntParameter(default=1)

    def run(self):
        self.dump(f'sample output version {self.version}')

Summary of differences:

Gokart allows mixing pandas and Polars tasks in the same pipeline. This is useful for gradual migrations. When a Polars task requires a pandas task, use pl.from_pandas() to convert the loaded data.

import pandas as pd
import polars as pl
from gokart import TaskOnKart
from gokart.file_processor import FeatherFileProcessor

class PandasTask(TaskOnKart[pd.DataFrame]):
    def output(self):
        return self.make_target(
            'path/to/pandas_output.feather',
            processor=FeatherFileProcessor(store_index_in_feather=False, dataframe_type='pandas')
        )

    def run(self):
        df = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})
        self.dump(df)

class PolarsTask(TaskOnKart[pl.DataFrame]):
    def requires(self):
        return PandasTask()

    def output(self):
        return self.make_target(
            'path/to/polars_output.feather',
            processor=FeatherFileProcessor(store_index_in_feather=False, dataframe_type='polars')
        )

    def run(self):
        # Load pandas DataFrame and convert to Polars
        pandas_df = self.load()  # Returns pandas DataFrame
        polars_df = pl.from_pandas(pandas_df)

        # Process with Polars
        result = polars_df.with_columns((pl.col('a') * 2).alias('a_doubled'))
        self.dump(result)