kissml

Keep It Simple Stupid Tools for Machine Learning

A Python library providing simple, powerful tools for ML workflows with minimal boilerplate.

I made this because:

• Most data science services are notebook based, but notebooks are difficult to debug
• Most frameworks (flyte, metaflow) focus on extending to the cloud. This is great, but for local iteration all we really need is reproducible pipeline steps.

Installation

pip install kissml

Steps

The @step decorator provides:

• execution tracking
• persistent disk-based caching for your functions
• post-run execution (i.e., after effects) for the return value -- useful to visualize data or log stats.

Basic Usage

from kissml import step, CacheConfig
import logging

# Simple execution time logging
@step(log_level=logging.INFO)
def process_data(data):
    # Your processing logic here
    return result

# With persistent caching
@step(
    log_level=logging.INFO,
    cache=CacheConfig(version=1)
)
def expensive_computation(data):
    # This will only run once per unique input
    # Subsequent calls return cached results
    return result

Key Features

Execution Time Tracking: Log how long your functions take to run

@step(log_level=logging.INFO)
def train_model(X, y):
    # Logs: "train_model completed in 45.2341 seconds"
    return model

Persistent Disk Caching: Cache results to disk and reuse them across runs

@step(cache=CacheConfig(version=1))
def load_and_preprocess(filepath):
    # Expensive preprocessing runs once
    # Subsequent calls load from cache in milliseconds
    return processed_data

Version-Based Invalidation: Bump the version to invalidate old cache

# Old implementation
@step(cache=CacheConfig(version=1))
def feature_engineering(df):
    return old_features(df)

# Updated implementation - cache automatically invalidated
@step(cache=CacheConfig(version=2))
def feature_engineering(df):
    return new_improved_features(df)

Smart Serialization: Efficient storage for pandas DataFrames and nested collections

import pandas as pd

@step(cache=CacheConfig(version=1))
def analyze_data(df: pd.DataFrame) -> pd.DataFrame:
    # DataFrames cached as Parquet files (requires pyarrow)
    # Much more efficient than pickle
    return processed_df

@step(cache=CacheConfig(version=1))
def complex_pipeline(data) -> dict:
    # Returns dict with DataFrames, lists, etc.
    # Each type uses optimal serialization
    return {
        "results": some_dataframe,
        "metrics": [metric1, metric2],
        "metadata": {"key": "value"}
    }

Cache Configuration

Control cache behavior with CacheConfig:

from kissml import step, CacheConfig, EvictionPolicy

# No eviction (default) - cache grows forever
@step(cache=CacheConfig(version=1, eviction_policy=EvictionPolicy.NONE))
def permanent_cache(x):
    return x

# Least Recently Used - evicts oldest accessed items
@step(cache=CacheConfig(version=1, eviction_policy=EvictionPolicy.LEAST_RECENTLY_USED))
def lru_cache(x):
    return x

# Least Recently Stored - evicts oldest stored items
@step(cache=CacheConfig(version=1, eviction_policy=EvictionPolicy.LEAST_RECENTLY_STORED))
def lrs_cache(x):
    return x

# Least Frequently Used - evicts least accessed items
@step(cache=CacheConfig(version=1, eviction_policy=EvictionPolicy.LEAST_FREQUENTLY_USED))
def lfu_cache(x):
    return x

AfterEffects

AfterEffects allow you to automatically execute side effects (like visualization, logging, or validation) after a step completes, whether the result was cached or freshly computed.

from typing import Annotated
from kissml import step, AfterEffect, CacheConfig
import mlflow

# Define a custom AfterEffect
class HTMLVisualizer(AfterEffect):
    def __init__(self, max_rows=100):
        self.max_rows = max_rows
    
    def __call__(self, result, was_cached, func_name, execution_time):
        # Create HTML preview
        html = result.head(self.max_rows).to_html()
        html = f"<h3>{func_name} - {execution_time:.2f}s {'(cached)' if was_cached else ''}</h3>" + html
        
        # Log to MLflow
        with open(f"{func_name}.html", "w") as f:
            f.write(html)
        mlflow.log_artifact(f"{func_name}.html")

# Use it with type annotations
@step(cache=CacheConfig(version=1))
def load_data() -> Annotated[pd.DataFrame, HTMLVisualizer(max_rows=200)]:
    return pd.read_csv("data.csv")

# Multiple effects run left-to-right
class DatasetLogger(AfterEffect):
    def __call__(self, result, was_cached, func_name, execution_time):
        if not was_cached:  # Only log once
            mlflow.log_metric(f"{func_name}_rows", len(result))

@step(cache=CacheConfig(version=1))
def process() -> Annotated[pd.DataFrame, DatasetLogger(), HTMLVisualizer()]:
    # Both effects run automatically after the function completes
    return load_data()

Error Handling: Control whether AfterEffect failures stop execution:

# Default: errors are logged but don't stop execution
@step(cache=CacheConfig(version=1))
def safe_pipeline() -> Annotated[pd.DataFrame, MyVisualizer()]:
    return data

# Strict mode: effect errors raise exceptions
@step(cache=CacheConfig(version=1), error_on_affect_failure=True)
def strict_pipeline() -> Annotated[pd.DataFrame, MyVisualizer()]:
    return data

Global AfterEffects: Register an AfterEffect once and have it fire after every @step call — no per-step annotation required. Useful for cross-cutting concerns like logging, persistence, or experiment tracking.

import logging
from kissml import settings, step, AfterEffect

class StepTimingLogger(AfterEffect):
    """Log every step's name, runtime, and cache status."""

    def __call__(self, result, was_cached, func_name, execution_time):
        status = "cached" if was_cached else "fresh"
        logging.info(
            f"{func_name} finished in {execution_time:.3f}s ({status})"
        )

# Register once — fires for every @step call from now on
settings.global_after_effects.append(StepTimingLogger())

@step()
def load_data() -> pd.DataFrame:
    return pd.read_csv("data.csv")  # StepTimingLogger runs after this returns

@step()
def transform(df: pd.DataFrame) -> pd.DataFrame:
    return df.dropna()              # StepTimingLogger runs here too

Per-step effects (declared in the return annotation) fire first, then global effects. Both honor the error_on_effect_failure flag on the step.

Subpipelines

The @subpipeline decorator composes steps into an uncached pipeline: it calls other (typically cached) @step functions and returns their results, doing no computation of its own. It never caches -- caching a composition would skip its body on a hit, silently skipping the inner steps too.

Because nothing is ever cached, a subpipeline's body runs on every call, so any AfterEffects declared on its return type are guaranteed to fire every time -- unlike a plain, undecorated function, where Annotated metadata is inert and those effects would never run.

from typing import Annotated
from kissml import step, subpipeline, AfterEffect, CacheConfig

class RowCountLogger(AfterEffect):
    def __call__(self, result, was_cached, func_name, execution_time):
        print(f"{func_name}: {len(result)} rows")

@step(cache=CacheConfig(version=1))
def load_data() -> pd.DataFrame:
    return pd.read_csv("data.csv")

@step(cache=CacheConfig(version=1))
def clean_data(df: pd.DataFrame) -> pd.DataFrame:
    return df.dropna()

@subpipeline()
def prepare_data() -> Annotated[pd.DataFrame, RowCountLogger()]:
    return clean_data(load_data())

prepare_data()
# load_data/clean_data may each hit cache, but prepare_data's own
# body -- and RowCountLogger -- run every time.

subpipeline() has no cache parameter (caching is deliberately unsupported) but accepts the same log_level and error_on_effect_failure options as step(). Both decorators stamp __kissml_kind__ on the wrapped function ("step" or "subpipeline") so tooling can classify decorated functions by introspection.

Configuration

Configure the cache directory via environment variable or settings:

from kissml import settings
from pathlib import Path

# Set cache directory
settings.cache_directory = Path("/path/to/cache")

# Or use environment variable
# export KISSML_CACHE_DIRECTORY=/path/to/cache

Custom Serialization

Register custom serializers for your types:

from kissml.settings import settings
from kissml.types import Serializer
from typing import Any, BinaryIO

class MyCustomSerializer(Serializer):
    def serialize(self, value: Any, out: BinaryIO) -> None:
        # Your serialization logic
        pass

    def deserialize(self, input: BinaryIO) -> Any:
        # Your deserialization logic
        pass

# Register the serializer
settings.serialize_by_type[MyCustomType] = MyCustomSerializer()

# Register a hash function for cache keys
settings.hash_by_type[MyCustomType] = lambda obj: str(hash(obj))

License

Licensed under CC BY-NC-ND 4.0 (Attribution-NonCommercial-NoDerivatives). This is a non-commercial license - see the LICENSE file for full details.

For commercial use, please contact the author.