Roman Parykin

Managing nested models with Bubble Tea

Wed, 09 Jul 2025 20:00:02 +0200

I don't get Bubbletea

r/golang, Reddit

Judging by other comments and similar posts, this experience is not unique.

In this post, I’d like to share an intuition behind model composition using the Bubble Tea TUI framework for Go, with the main focus on model nesting. While I assume a basic understanding of the Bubble Tea framework, let’s start by going over what a model is.

The Elm Architecture

The Bubble Tea is based on The Elm Architecture (The acronym TEA, short for The Elm Architecture, is what gives Bubble Tea its name.), where the model is the central piece of application state, representing the entire state of the application at any given point in time. The other two components are:

The Update function that modifies the model in response to messages.
The View function that uses the model to determine what to show to the user.

In Elm, the model itself isn’t a formal concept, but a common pattern used to structure different states of the application. In practice, it’s a data structure that holds everything needed to render itself.

Nested models

And that’s where the official tutorial ends, without even mentioning how to handle nested states, or, in other words, nested models. Let’s take a look at an example called Composable Views from Bubble Tea’s examples folder. The model, which composes two other models (timer and spinner), looks like this:

type sessionState uint

type mainModel struct {
	state   sessionState
	timer   timer.Model
	spinner spinner.Model
	//...
}

The two models, timer.Model and spinner.Model (All xyz.Models are part of the Charm bubbles project.), are nested within the mainModel. The sessionState is used to determine which model is currently active. Both the Update and View functions then branch their logic based on the current sessionState, either routing messages to the appropriate model(s) or rendering the active model, respectively.

View the diagrams on the website

As the number of nested models grows, we may want to move some of them into separate models, thereby creating a tree model structure.

type subModel struct {
	time    timer.Model
	spinner spinner.Model
}

type mainModel struct {
	state    sessionState
	list     list.Model
	subModel subModel
}

View the diagrams on the website

No matter the size of the model tree, the mainModel remains the top-level model, responsible for routing messages to the child models and view rendering with the content of the child’s View methods. For example, this is how updating the timer.Model looks like:

View the diagrams on the website

There is nothing wrong with this approach. In fact, this is the recommended way to handle the nested models, as demonstrated in more detail in the “Bubble Tea Nested Models” video. It makes perfect sense if the child models are part of the current view layout (Generalizing the modeling of components (or bubbles) with the current Go (1.25) type system is challenging. Although referred to as models, bubbles aren’t true models in the strict sense - they don’t implement the tea.Model interface and they provide additional API. One possible workaround would be to make their API event-driven, but this would come at the cost of event loop performance.). However, not all application models are like that. Some models represent entirely new screen views and have nothing to do with the main or previous model. Does this change the handling of nested models, as described above? Not at all. The main model is still responsible for tracking the current model and rendering the appropriate view in place of its own.

To recap, the top-down approach is characterized by the following attributes:

Keeping track of the application state.
Routing messages to the appropriate models in the Update function based on the current state.
Composing or delegating the view in the View function based on the current state.

Independent nested models

Let’s bring in one more example to the picture, such as the Credit Card Form. The main model will consist of list.Model and two submodels: SubModel1 for Composable Views and SubModel2 for the Credit Card Form. The main view might show a list of available examples (two, in this case). When the user selects an example, the corresponding view is shown instead of the list. Both SubModel1 and SubModel2 are completely independent of each other and of the main model. In this scenario, the MainModel’s Update and View functions should handle at least three states. Even if the submodel’s view is self-contained, the MainModel’s view should still be updated.

View the diagrams on the website

An independent model is one that contains all the information necessary to render a screen and handle user interactions. Examples of independent models include the main model and both submodels. In contrast, models like list.Model or textinput.Model are usually not independent, as they require additional models to manage user interaction. Some models can serve both roles simultaneously, such as both submodels. If a submodel is shown upon item selection, it is considered independent. However, if it’s shown alongside the list, it becomes part of the main model’s view and is therefore not independent.

The steps for adding an independent model are closely tied to the top-down approach:

Add a new state to the main model.
Extend the main model’s Update function to handle the new state and route messages to the new model.
Update the main model’s View function to render the new model’s view based on the updated state.

Let’s see how to get rid of all these steps with the Model Stack architecture.

Model Stack architecture

The concept behind the Model Stack architecture is to restructure a tree of nested models into a stack of independent ones. In the example above, the model tree can be decomposed into three independent models, stacked on top of one another as needed.

View the diagrams on the website

Initially, the model stack contains only the entry model. Each model is responsible for managing its own state and knows how to render itself - it behaves just like any other component. How is another model displayed? It’s pushed onto the stack, becoming the current model. The model stack controller handles model initialization, routes messages to the current model, and displays the current model’s view.

View the diagrams on the website

The model stack controller takes over routing for nested models - replacing the need for coordination required in the top-down approach within the main model’s Update and View functions - thus eliminating the steps for adding an independent model. The controller plays the role of the main model and intercepts all commands, but it does not need session state. All that’s left for models is to notify the controller.

Bubblon

Let’s see the benefits of this approach using the bubblon library, which implements the Model Stack architecture for the Bubble Tea framework.

Event-driven

Since the models are supposed to be independent, they cannot be accessed directly because the stack is encapsulated by the controller. The only way to interact with the controller is by sending commands. This allows any model to communicate with the controller without holding a direct reference to it.

Minimalistic

At its core, the Model Stack architecture is about simplification. For example, the bubblon.Open command will initialize the provided model, make it the current one, and redirect messages to it. All of this is accomplished in a single line of code (Internally, there is a stack of models, but semantic API - using commands like Open and Close - makes more sense in the context of TUI apps.):

func (m MainModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	// ...
	return m, bubblon.Open(SubModel1.New())
}

When we’re done with the submodel and want to return to the main model:

func (m SubModel1) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	// ...
	return m, bubblon.Close
}

Following its event-driven design, bubblon will send the MainModel the bubblon.Closed message, indicating that the submodel has been closed.

The API remains unobtrusive by introducing no new interfaces or concepts. Models are expected to implement the existing tea.Model interface, and all interactions are performed using Bubble Tea commands. The controller itself is a tea.Model and can be used like any other model.

Modular

One advantage of using independent models is dynamic composition. Any model can become the current one. Model A can open model B, which in turn can open model C - without model A knowing anything about it. Or C → B → A. No state, no cry. While Go prevents circular dependencies at the package level, it’s still possible to pass a constructed instance or a function to initialize a parent model.

Resource-efficient

By default, models are released and made available to the garbage collector once closed by the controller, rather than being kept indefinitely in the main model’s internal state. In addition, more efficient commands are available for replacing the top model - or all models - with a new one, without relying on a Close/Open sequence: bubblon.Replace and bubblon.ReplaceAll (The full API reference can be found on go.dev.).

Since the controller is not responsible for model creation, it is not limited to dynamic composition. Otherwise it would be a waste of resources to recreate a model that can be instantiated just once. This is where a hybrid approach to model caching comes in. A submodel that is better suited for one-time creation can be embedded for repeated use, with all the other benefits of the Model Stack architecture preserved.

type model struct {
	submodel tea.Model
}

// somewhere
// m := model{&submodel{}}

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	// Maybe reset submodel's state first
	return m, bubblon.Open(m.submodel)
}

Without any session state or command coordination, the submodel is cached by the main model and can be reused at any time. The goal of the Model Stack architecture is to enable both dynamic composition and model caching with minimal code changes.

To summarize, there are at least three progressive levels for building Bubble Tea applications using nested models with bubblon:

Level 1: Use top-down approach to build a model from individual components.
Level 2: Use the Model Stack architecture to compose an application from independent models.
Level 3: Use a hybrid approach with the Model Stack architecture to maximize flexibility and performance.

Check out bubblon on GitHub for more practical examples.

Stretchy architecture

Sun, 23 Jul 2023 16:33:52 +0300

Stretchy architecture is an architectural approach where each dimension is chosen in a way that makes it easier to move in one direction on the decision spectrum than the other. An important condition is to move back to the starting point if possible.

Asynchronous vs synchronous communication. Usually presented as a trade-off, in the stretchy architecture style, moving from async to sync is fine and, in practice, it’s much easier then the other way round. Starting with async communication by default allows switching to sync and back anytime if needed. Switching from sync to async is not only more difficult but might also be impossible, which locks (see?) you into the sync style.

Eventual vs strong consistency. Stick to eventual consistency (or any other weaker consistency model), and any part can be made atomic without changing much of the rest. Suggesting to replace the existing strong consistency with weaker models might be dangerous for the messenger. This friction is what makes some people argue about which style is better. Neither is; it’s just easier to start with a weaker model and move towards stronger ones if new requirements dictate that.

Choreography vs orchestrated coordination. Besides the tight coupling introduced by the mediator, orchestration implementations often lead to synchronous communication, where components must wait for instructions from the orchestrator. On the other hand, a fully choreographed system offers the flexibility to add an orchestrator to facilitate coordination among components if the need arises.

Immutability vs mutability. Given a piece of immutable code, it’s fine to introduce mutability in certain parts for performance reasons, as long as it remains isolated from the rest of the codebase. It’s a perfect illustration of “stretchiness”, where we begin with immutable design, stretch it a bit to benefit from lower-level optimizations, and then return back to immutability.

Stateless vs stateful protocol. The recurring pattern of increased friction of movement in one direction more than the other is a defining characteristic of the stretchy architecture. Can be confirmed by anyone who has tried to transition from a stateful protocol to a stateless one.

Schema vs schemaless modeling. Having a schema for data modeling condones hacks like treating JSON objects as text or bytes. When everything is text, where do you go from there?

Static vs dynamic typing. Similar to data modeling but applied to code, statically typed languages provide more guarantees out of the box, and most of them are capable of loosening data representation. One of the useful patterns is to parse data as soon as possible (for example, when reading a file) and work with typed data afterwards. With a map of any values by any keys, the code is destined to become increasingly problematic.

Thinking in terms of the stretchy architecture results in responsive, highly-available, scalable, fault-tolerant, decoupled, and maintainable systems - essentially for free.

Hugging Face pipeline inference optimization

Sun, 19 Feb 2023 14:43:36 +0200

The goal of this post is to show how to apply a few practical optimizations to improve inference performance of 🤗 Transformers pipelines on a single GPU. Compatibility with pipeline API is the driving factor behind the selection of approaches for inference optimization. This is a practical guide to optimizing inference of 🤗 Transformers pipelines based on my personal experience. For more methods on how to make transformer inference more efficient, I recommend checking out Lilian Weng’s blog (Weng, Lilian. (Jan 2023). Large Transformer Model Inference Optimization. Lil’Log.).

All the code is run on NVIDIA RTX 4090 24G using Python 3.10 and PyTorch 1.13.1. Other dependencies are:

datasets==2.9.0
evaluate[evaluator]==0.4.0
optimum[onnxruntime-gpu]==1.6.4
transformers==4.26.1

To start off, let’s establish the baseline and evaluation method, which will remain consistent across all approaches.

Baseline

As a baseline, I’m going to use the RoBERTa base model fine-tuned on the SQuAD dataset for extractive question answering. The training is done using the scripts from the transformers examples (All training examples are based on the v4.26.1 tag.) for PyTorch.

First, let’s fine-tune the roberta-base model on the squad dataset for two epochs with the following parameters (to match the examples from the transformers):

python run_qa.py \
  --model_name_or_path roberta-base \
  --dataset_name rajpurkar/squad \
  --do_train \
  --per_device_train_batch_size 32 \
  --learning_rate 3e-5 \
  --num_train_epochs 2 \
  --save_strategy no \
  --output_dir /workspace/roberta-base-squad

Pipeline Evaluation

The following code is going to be used for the 🤗 Transformers pipeline (The official tutorial on how to use pipelines for inference can be found here.) evaluation:

from datasets import load_dataset
from evaluate import evaluator
from transformers import pipeline

Metrics = dict[str, float]

batch_size = 8
data = load_dataset('rajpurkar/squad', split='validation')
qa_evaluator = evaluator('question-answering')
def qa_eval(model, tokenizer) -> Metrics:
    pipe = pipeline(
        task='question-answering',
        model=model,
        tokenizer=tokenizer,
        batch_size=batch_size,
        device=0,
    )
    metrics = qa_evaluator.compute(
        model_or_pipeline=pipe,
        data=data,
        metric='squad',
        device=0,
    )
    return {k: round(v, 4) for k, v in metrics.items()}

With the fine-tuned model and evaluation method in place, we can establish the baseline as follows:

from transformers import AutoModelForQuestionAnswering, AutoTokenizer

model_name = '/workspace/roberta-base-squad'
model = AutoModelForQuestionAnswering.from_pretrained(model_name).to('cuda')
tokenizer = AutoTokenizer.from_pretrained(model_name)

qa_eval(model=model, tokenizer=tokenizer)

{'exact_match': 85.6481, 'f1': 92.1859, 'samples_per_second': 293.3532}

The accuracy metrics are exact_match and f1, while the throughput is reported as samples_per_second. The important thing to note is that the numbers are for the pipeline and not the model itself, as the pipeline has extra logic for computing the best answer. Additionally, there is overhead caused by the evaluation. According to which, the pipeline baseline is indicated by an f1 score of 92.1859 and a throughput of 293 samples per second. The exact numbers are not that significant as relative performance improvements compared to the baseline.

With the baseline and evaluation out of the way, let’s see how to apply the first optimization technique—knowledge distillation.

Knowledge Distillation

Knowledge distillation is a training technique for transferring knowledge from a pre-trained model (“teacher”) to a less complex one (“student”).

In some cases, task-specific knowledge distillation refers to the process of fine-tuning already distilled language model such as DistilBERT. Let’s see what we’ll get by fine-tuning the DistilRoBERTa base model on the SQuAD with the same parameters:

python run_qa.py \
  --model_name_or_path distilroberta-base \
  --dataset_name rajpurkar/squad \
  --do_train \
  --per_device_train_batch_size 32 \
  --learning_rate 3e-5 \
  --num_train_epochs 2 \
  --save_strategy no \
  --output_dir /workspace/distilroberta-base-squad

Evaluation:

from transformers import AutoModelForQuestionAnswering, AutoTokenizer

model_name = '/workspace/distilroberta-base-squad'
model = AutoModelForQuestionAnswering.from_pretrained(model_name).to('cuda')
tokenizer = AutoTokenizer.from_pretrained(model_name)
qa_eval(model=model, tokenizer=tokenizer)

{'exact_match': 80.5109, 'f1': 87.7095, 'samples_per_second': 381.8372}

Although a 5% drop in pipeline accuracy might be the biggest trade-off to make, a 30% increase in speed is still quite appealing, especially when combined with other techniques. Let’s see if we can do better with two-step knowledge distillation.

Two-step distillation

We also studied whether we could add another step of distillation during the adaptation phase by fine-tuning DistilBERT on SQuAD using a BERT model previously fine-tuned on SQuAD as a teacher for an additional term in the loss (knowledge distillation). In this setting, there are thus two successive steps of distillation, one during the pre-training phase and one during the adaptation phase. In this case, we were able to reach interesting performances given the size of the model: 79.8 F1 and 70.4 EM, i.e. within 3 points of the full model.

Sanh, V.; Debut, L.; Chaumond, J. & Wolf, T. (2019), 'DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter'

Let me show how to implement two-step knowledge distillation to improve accuracy while keeping the throughput. In order to do this, we need to add one file to the same folder where the run_qa.py script is (examples/pytorch/question-answering):

distil_trainer_qa.py

from dataclasses import dataclass, field
from typing import Optional

import torch
import torch.nn.functional as F
from transformers import AutoModelForQuestionAnswering, TrainingArguments

from trainer_qa import QuestionAnsweringTrainer


@dataclass
class DistilTrainingArguments(TrainingArguments):
    alpha: float = field(
        default=0.25,
        metadata={"help": "Controls the relative strength of each loss"}
    )

    temperature: int = field(
        default=3,
        metadata={"help": "Scaling factor to soften probabilities"}
    )

    teacher_name_or_path: Optional[str] = field(
        default=None,
        metadata={"help": "Teacher model name or path"}
    )


class DistilQuestionAnsweringTrainer(QuestionAnsweringTrainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.teacher = None
        if self.args.teacher_name_or_path:
            self.teacher = AutoModelForQuestionAnswering.from_pretrained(
                self.args.teacher_name_or_path
            )
            self.teacher.to(self.args.device).eval()
        self.loss_fn = torch.nn.KLDivLoss(reduction="batchmean")
        self.alpha = self.args.alpha
        self.T = self.args.temperature
        self.T2 = self.T ** 2

    def compute_loss(self, model, inputs, return_outputs=False):
        if self.teacher:
            outputs_stu = model(**inputs)
            with torch.no_grad():
                outputs_tea = self.teacher(**inputs)
            loss_start = self.T2 * self.loss_fn(
                F.log_softmax(outputs_stu.start_logits / self.T, dim=-1),
                F.softmax(outputs_tea.start_logits / self.T, dim=-1)
            )
            loss_end = self.T2 * self.loss_fn(
                F.log_softmax(outputs_stu.end_logits / self.T, dim=-1),
                F.softmax(outputs_tea.end_logits / self.T, dim=-1)
            )
            loss_kd = (loss_start + loss_end) / 2.0
            # Overall loss as a weighted sum
            loss = self.alpha * outputs_stu.loss + (1.0 - self.alpha) * loss_kd
            return (loss, outputs_stu) if return_outputs else loss
        else:
            return super().compute_loss(model, inputs, return_outputs)

It extends the TrainingArguments by adding three new parameters:

alpha: Controls how much weight is put on the student-teacher loss relative to the student loss alone, with higher value giving more weight to the student loss.
temperature: Controls the level of smoothing of the teacher probability distribution, with higher temperature leading to a softer distribution.
teacher_name_or_path: Name or path to the pre-trained teacher model.

The default values for alpha and temperature are set as a result of a dozen of Optuna trials (Hyperparameter space: 𝛼 ∈ [0, 1], T ∈ [1, 10].) for this specific task. In general, for many NLP tasks a good starting point would be alpha < 0.3 and temperature > 1.

The DistilQuestionAnsweringTrainer class extends the QuestionAnsweringTrainer from the trainer_qa.py file by overriding the compute_loss function. The new loss function involves logits from the teacher to calculate the overall loss as a weighted sum of the student’s loss and the distillation loss using the new alpha and temperature parameters.

In addition to this new file, the run_qa.py itself should be modified to support the new distillation trainer:

run_qa.py

from utils_qa import postprocess_qa_predictions
+ from distil_trainer_qa import DistilTrainingArguments, DistilQuestionAnsweringTrainer

- parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)
+ parser = HfArgumentParser((ModelArguments, DataTrainingArguments, DistilTrainingArguments))

- trainer = QuestionAnsweringTrainer(
+ trainer = DistilQuestionAnsweringTrainer(

If the changes have been applied correctly, training with two-step knowledge distillation should be as easy as adding the --teacher parameter (when leaving the alpha and the temperature parameters as default):

python run_qa.py \
  --model_name_or_path distilroberta-base \
  --dataset_name rajpurkar/squad \
  --do_train \
  --per_device_train_batch_size 32 \
  --learning_rate 3e-5 \
  --num_train_epochs 2 \
  --save_strategy no \
  --output_dir /workspace/distilroberta-base-distilsquad \
  --teacher /workspace/roberta-base-squad

If there is any justice in the world, bringing teacher into the equation should move accuracy closer to our baseline.

from transformers import AutoModelForQuestionAnswering, AutoTokenizer

model_name = '/workspace/distilroberta-base-distilsquad'
model = AutoModelForQuestionAnswering.from_pretrained(model_name).to('cuda')
tokenizer = AutoTokenizer.from_pretrained(model_name)
qa_eval(model=model, tokenizer=tokenizer)

{'exact_match': 81.6746, 'f1': 88.7431, 'samples_per_second': 384.2462}

The two-step distillation process maintains a throughput of 31% over the baseline, retaining 96% of the baseline accuracy. Increasing the number of epochs (and playing with other hyperparameters) will bump this number even further (The DistilBERT paper reports “within 3 points of the full model”.) without any negative effect on throughput. I will use this distilled model for the upcoming set of methods.

Automatic Mixed Precision

Automatic Mixed Precision (AMP) allows to use a mix of torch.float32 and half-precision (torch.float16) floating point datatypes during inference, thereby reducing the memory footprint and improving performance while maintaining accuracy.

import torch
from transformers import AutoModelForQuestionAnswering, AutoTokenizer

model_name = '/workspace/distilroberta-base-distilsquad'
model = AutoModelForQuestionAnswering.from_pretrained(model_name).to('cuda')
tokenizer = AutoTokenizer.from_pretrained(model_name)

# autocast context manager allows code regions to run in mixed precision
with torch.cuda.amp.autocast(dtype=torch.float16):
    print(qa_eval(model=model, tokenizer=tokenizer))

{'exact_match': 81.7029, 'f1': 88.7576, 'samples_per_second': 416.5399}

In addition to the performance improvement by 8.4% compared to the original model (Double that for the non-distilled baseline model.), there is also a nearly 0.03% increase in accuracy as well. In general, it’s a safe approach to use the AMP autocast with language models. Later on, we’ll see that, even in situations where it does not impact the performance, it typically does not cause any harm.

TorchScript

TorchScript is a way to create serializable and optimizable models from PyTorch code (From PyTorch documentation on TorchScript.).

There are two ways to use TorchScript models with pipelines:

Tracing an existing model.
Loading a TorchScript module to use for inference.

We’ll now explore how to trace an existing model to use for pipeline inference, as loading a TorchScript module will be straightforward with tricks used for tracing.

There are three steps to trace a model to use it with pipeline:

Prepare example inputs that will be passed to the model.
Trace the model to capture its structure by evaluating it using example inputs.
Patch existing model to use TorchScript for forward pass.

from datasets import load_dataset
import torch
from transformers import AutoModelForQuestionAnswering, AutoTokenizer
from transformers.modeling_outputs import QuestionAnsweringModelOutput

model_name = '/workspace/distilroberta-base-distilsquad'
model = AutoModelForQuestionAnswering.from_pretrained(model_name, torchscript=True).to('cuda')
tokenizer = AutoTokenizer.from_pretrained(model_name)

def encode(examples):
    return tokenizer(
        examples['question'],
        examples['context'],
        truncation='only_second',
        max_length=384,
        stride=128,
        padding='max_length',
    )

# 1. Prepare example inputs
squad = load_dataset('rajpurkar/squad', split=f'validation[:{batch_size}]')
example_batch = squad.map(
    encode,
    batched=True,
    remove_columns=squad.column_names
).with_format('torch')

# 2. Trace the model
jit_inputs = (
    example_batch['input_ids'].to('cuda'),
    example_batch['attention_mask'].to('cuda')
)
jit_model = torch.jit.trace(model, jit_inputs)
jit_model = torch.jit.freeze(jit_model).to('cuda').eval()

# 3. Patch the model forward function to be compatible with pipeline
def forward_qa_wrap(**inputs):
    with torch.jit.optimized_execution(False):
        start_logits, end_logits = jit_model(**inputs)
    return QuestionAnsweringModelOutput(
        start_logits=start_logits,
        end_logits=end_logits
    )

model.forward = forward_qa_wrap

with torch.cuda.amp.autocast(dtype=torch.float16):
    print(qa_eval(model=model, tokenizer=tokenizer))

{'exact_match': 81.7029, 'f1': 88.7576, 'samples_per_second': 446.3015}

The performance increase is 16% compared to the original distilled model, which is almost twice the improvement achieved by autocast alone.

There are several important tricks to get there:

TorchScript is optimized using just-in-time (JIT) compilation based on inputs. When inputs are different (and they are in the case of question answering), JIT recompilation will slow down the pipeline performance eventually. To keep the performance stable, it’s essential to run TorchScript inside the torch.jit.optimized_execution(False) context manager.
The pipeline for question answering expects a model to return an instance of the QuestionAnsweringModelOutput class, while TorchScript will return a tuple of logits. To make it compatible with the pipeline API, it should be wrapped in a function mimicking the forward call. The original model’s forward function is then replaced with the new wrapper (This is a more general technique that could be used to wrap the NVIDIA Apex for example.).
Using TorchScript along with autocast is crucial for inference optimization. In my experience, utilizing autocast during inference instead of tracing itself results in a slight improvement in performance.

While this approach may result in less noticeable performance improvements on some GPUs, such as the RTX 30 series, it is more stable from run to run than AMP alone.

DeepSpeed

Setting up DeepSpeed integration requires more effort if there are missing system dependencies, such as Python shared libs and header files (More about DeepSpeed installation here.). If you are using DeepSpeed for training, you can leverage DeepSpeed-Inference to perform inference as well. Despite the fact that DeepSpeed-Inference shares its name with DeepSpeed, it does not use ZeRO (Zero Redundancy Optimizer) technology.

Using DeepSpeed-Inference with the pipeline API is straightforward (DeepSpeed-Inference tutorial.):

import deepspeed
import torch
from transformers import AutoModelForQuestionAnswering, AutoTokenizer

model_name = '/workspace/distilroberta-base-distilsquad'
model = AutoModelForQuestionAnswering.from_pretrained(model_name).to('cuda')
tokenizer = AutoTokenizer.from_pretrained(model_name)

engine = deepspeed.init_inference(
    model,
    dtype=torch.float16,
    replace_with_kernel_inject=True,
)

with torch.cuda.amp.autocast(dtype=torch.float16):
    print(qa_eval(model=engine.module, tokenizer=tokenizer))

Although I don’t have the exact numbers for the same GPU used in the other methods, it consistently outperforms the previous approach (TorchScript with AMP) by 5–7% on the RTX 3090 24G.

Optimum

🤗 Optimum is an extension of 🤗 Transformers, providing a set of performance optimization tools with unified API to train and run models on targeted hardware with maximum efficiency (More about 🤗 Optimum can be found in the official documentation.). Optimum can be used for accelerated inference with built-in support for transformers pipelines, making it an ideal candidate for pipeline optimization techniques. I will demonstrate how to apply graph optimization to accelerate inference with ONNX Runtime.

ONNX Runtime (ORT) is a cross-platform, high-performance engine for Open Neural Network Exchange (ONNX) models used to accelerate inference and training of machine learning models.

The code below outlines every step required for 🤗 Optimum inference using ONNX Runtime with relevant comments:

from optimum.onnxruntime import ORTModelForQuestionAnswering
from optimum.onnxruntime import ORTOptimizer
from optimum.onnxruntime.configuration import OptimizationConfig
from transformers import AutoTokenizer

model_name = '/workspace/distilroberta-base-distilsquad'
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Load a model and export it to ONNX format
model = ORTModelForQuestionAnswering.from_pretrained(model_name, from_transformers=True)

# Define optimization configuration
opt_config = OptimizationConfig(
    optimization_level=2,
    optimize_for_gpu=True,
    fp16=True,
)

# Create optimizer
optimizer = ORTOptimizer.from_pretrained(model)

# Optimize the model applying defined optimization strategy
optimizer.optimize(
    optimization_config=opt_config,
    save_dir=f'{model_name}/onnx',
    file_suffix=None,
)

# Load the optimized model for inference
model = ORTModelForQuestionAnswering.from_pretrained(model_name, file_name='onnx/model.onnx')

qa_eval(model=model, tokenizer=tokenizer)

{'exact_match': 81.7029, 'f1': 88.7575, 'samples_per_second': 477.7468}

The throughput of the optimized model improved by 24% compared to the initial distilled model.

A few notes on the example configuration.

Since our focus is on GPU inference, setting the optimize_for_gpu to True is crucial. There is no specific reason to go with the optimization_level=2 other than it being a good starting point. If mixed precision is enabled (fp16=True), there is no need to use autocast during inference. However, if the optimization configuration being used is unknown, it’s ok to use both autocast and optimized models together.

Numerology

Ok, these are some of the optimization techniques that work well with 🤗 Transformers pipelines in my experience. As most of the examples are based on the distilled model, let’s see the results with the same techniques being applied to the baseline model too. The table can be mentally broken down into two parts—the first part without distillation, and the second part containing optimizations combined with distillation.

Approach	Exact Match	F1	F1 diff	Samples/sec	Speedup
Baseline	85.65	92.19	100%	293.4	1.00
Automatic Mixed Precision (AMP)	85.64	92.19	100%	347.2	1.27
TorchScript + AMP	85.64	92.19	100%	379.8	1.29
Optimum (ORT)	85.7	92.21	100%	429.6	1.46
Knowledge Distillation	81.67	88.74	96%	384.2	1.31
Distillation + AMP	81.7	88.76	96%	416.5	1.42
Distillation + TorchScript + AMP	81.7	88.76	96%	446.3	1.52
Distillation + Optimum (ORT)	81.7	88.76	96%	477.7	1.63

Although the difference is negligible, using 🤗 Optimum with ONNX Runtime results in slightly higher accuracy compared to the baseline model. What’s even better is that all the described methods have no negative impact on accuracy, except for knowledge distillation. It looks like that, in order to surpass the 1.5x throughput speedup, one may have to sacrifice some accuracy. While this threshold is specific to the particular hardware, the general trend should be similar for different GPUs. So it’s not surprising that the highest throughput is achieved by using 🤗 Optimum in combination with knowledge distillation, as these are the two most significant contributing factors.

Adaptive Reranker Weight

Sat, 15 Oct 2022 11:38:30 +0200

A method for improving the search results (e.g. for question answering retrieval) is known as Retrieve & Re-rank. The Retrieve & Re-rank pipeline consists of at least two components: the retriever and the re-ranker. The retriever is responsible for retrieving a number of documents given a search query. The relevancy of the retrieved documents can be further improved by the re-ranker, which allows retrieving more documents than required for the next pipeline stage.

Let’s see an example of the output at each stage.

query = 'what is re-ranking?' # Search query
n = 10 # Number of documents to retrieve

# Assume this function returns a list of documents with score
docs = retrieve(query=query, n=n)
scores = [doc['score'] for doc in docs]

# [0.9782995053726794,
#  0.9504939500760989,
#  0.8765814146070106,
#  0.8623934128019434,
#  0.842523354483268,
#  0.7736853461402741,
#  0.7713904667955406,
#  0.6740331628686816,
#  0.6378117863548827,
#  0.5634670917387724]

# And this function re-ranks the documents based on the
# relevancy to the query and returns the new scores
reranked = rerank(query=query, docs=docs)

# [0.8958727100108653,
#  0.9704265468563152,
#  0.8037856351531634,
#  0.4605732745735953,
#  0.9991750843646917,
#  0.7299899568668072,
#  0.6836966943663378,
#  0.6294383998509153,
#  0.5605524792499585,
#  0.41810846856511075]

Now, given two sets of scores per document, what is the total score? Calculating the average is one of the simplest approaches that could work.

total_scores = [(s + r) / 2 for s, r in zip(scores, reranked)]

# [0.9370861076917724,
#  0.960460248466207,
#  0.840183524880087,
#  0.6614833436877694,
#  0.9208492194239799,
#  0.7518376515035406,
#  0.7275435805809392,
#  0.6517357813597985,
#  0.5991821328024206,
#  0.4907877801519416]

After combining both scores, the second document has the highest score now. If we want to assign different weights to the retriever and re-ranker scores, we can do so by multiplying the scores with their respective weight parameters:

# The weight values are completely arbitrary
retriever_weight = 1.2
reranker_weight = 1.5 # Boost reranker score a little bit

scoring = lambda s, r: (s * retriever_weight + r * reranker_weight) / 2
total_scores = [scoring(s, r) for s, r in zip(scores, reranked)]

Now the question is, what value should be used for the re-ranker weight? Looking at the total scores, it appears that the positions of the first and second documents have been swapped, with the fifth document occupying third place and so on. There is nothing wrong with keeping the re-ranker weight static. But what if, instead swapping the first two places, the first document ends up becoming the last one? The same weight might not be sufficient for so many search queries and numbers of documents.

The intuition behind the idea is that the more spread out the documents are after re-ranking, the greater the weight. What makes it intuitive is that it’s about the positional difference, rather than the score. Although the total score difference might be insignificant, moving several positions down from the top can be detrimental to the pipeline performance. It’s not a big deal if we retrieve 10 documents and show/process them all. In case when we retrieve 100 documents to return top-5 of them, any document ranked higher than 5 will be filtered out.

If we view the document positions before and after re-ranking as observed and predicted values, we can use any deviation or error function that is used to measure the difference between values. One such function is Root Mean Square Error:

$$RMSE = \sqrt{\frac{1}{n}\sum_{i=1}^{n}(y_{i}-x_{i})^{2}}$$

Using the scores from above, it can be calculated as follows:

import numpy as np

retrieved_pos = np.argsort(scores)
reranked_pos = np.argsort(reranked)

np.sqrt(np.mean(np.square(reranked_pos - retrieved_pos)))

# 2.23606797749979

Another example is Mean Absolute Error:

$$MAE = \frac{1}{n}\sum_{i=1}^{n}|y_{i}-x_{i}|$$

np.mean(np.abs(reranked_pos - retrieved_pos))

# 1.6

Using different error functions will result in different values given the same magnitude of errors. In this case, the error of the MAE is smaller than the RMSE. We cannot compare the errors of different functions, but using the same magnitude of numbers allows us to choose any of them based on how large we want the resulting value to be.

Adaptive reranker weight is the measure between the retrieved and reranked positions. In other words, the adaptive weight is determined by how much the retrieved documents are rearranged by the re-ranker.

$$ARW = max(e, w), \qquad \forall e,w\in\mathbb{R}$$ , where $e$ is error value and $w$ is minimum weight.

With minimum weight set to zero, the adaptive reranker weight is just the error:

$$ARW = max(RMSE, 0) = RMSE = \sqrt{\frac{1}{n}\sum_{i=1}^{n}(y_{i}-x_{i})^{2}}$$

The adaptive reranker weight with MAE error and minimum weight of 1:

$$ARW = max(\frac{1}{n}\sum_{i=1}^{n}|y_{i}-x_{i}|, 1)$$

With scores changing per query, the re-ranker weight is adapted based on the difference between document positions. This approach can be further extended by incorporating the scores or by using different error functions. Logging the adaptive weight can also serve the purpose of monitoring the scoring deviations. The same idea can be applied down the pipeline to the components that calculate new sets of scores, thus affecting the document ranks.

Compositional streaming of services with FS2

Sat, 23 May 2020 12:56:19 +0200

The intention of this post is to demonstrate a blueprint for composing multiple services as streams within the scope of one process with FS2. The idea is that, when dealing with unbounded I/O streams, modeling everything as streams leads to better composition.

The libraries that will be used for all the examples:

"co.fs2" %% "fs2-core" % "2.3.0",
// Configuration and validation
"is.cir" %% "ciris-refined" % "1.0.4"

And the imports:

import cats.effect._
import cats.implicits._
// Configuration and validation
import ciris._
import ciris.refined._
import eu.timepit.refined.W
import eu.timepit.refined.api.Refined
import eu.timepit.refined.string.MatchesRegex
import eu.timepit.refined.types.net.PortNumber
import eu.timepit.refined.types.string.NonEmptyString

Configuration

Let’s start with configuration. Unlike the rest of the examples, I’d like to pause and show the awesomeness of the Ciris configuration library. The docs cover most of its features, so let’s jump to defining a sample config for Elasticsearch. The basic config would probably contain the host and port number, at least. If it’s deployed on AWS, it could require the AWS region as well (along with many other parameters skipped for brevity). The region names follow a particular pattern so let’s start by validating it first:

object AWS {
  type Region = String Refined MatchesRegex[W.`"([a-z]{2,}-)+[0-9]"`.T]
}

We’re using the refined library here to specify the Region type, which checks if a string matches the region name regex.

The full config will look like this:

final case class ElasticsearchConfig(
  host: NonEmptyString,
  port: PortNumber,
  region: AWS.Region
)

Notice how every field is validated using special types from refined. The host field is a string but cannot be empty, while the port number is guaranteed to be in the correct range by using the PortNumber type. As a result, there are no Strings and Ints all over the place with custom validation on top. Ok, it’s just a model, time to specify how to populate it with values.

object ElasticsearchConfig {
  val config: ConfigValue[ElasticsearchConfig] =
    (
      env("ELASTICSEARCH_HOST").default("localhost").as[NonEmptyString],
      env("ELASTICSEARCH_PORT").as[PortNumber].default(PortNumber(9200)),
      env("ELASTICSEARCH_REGION").as[AWS.Region]
    ).parMapN(apply)
}

The values are loaded from environment variables with defaults, validated, and combined in parallel. The automatic support for refined is provided by the ciris-refined module. The same approach is used to compose multiple configurations into larger ones. For example given a separate Kafka config

final case class KafkaConfig(bootstrapServers: NonEmptyString, ...)

both configs can be accumulated into one main config:

final case class Config(es: ElasticsearchConfig, kafka: KafkaConfig)

object Config {
  val config: ConfigValue[Config] = (
    ElasticsearchConfig.config,
    KafkaConfig.config
  ).parMapN(apply)
}

To load the configuration, we have to specify an effect type, which in this case will be cats-effect’s IO type:

val config: IO[Config] = Config.config.load[IO]

Dependencies

An initial stream to work with can be built by evaluating the IO effect to emit the config value (effectively going from IO[A] to Stream[IO, A]):

val stream: fs2.Stream[IO, Config] = fs2.Stream.eval(config)

We can go back any time with stream.compile.lastOrError: IO[Config]. The main reason not to, is to model an effectful computation as a stream, allowing it to be composed with a streaming service later on. In other words, instead of “wrapping” a dependency inside a stream, this modeling approach represents it as a stream itself.

Another dependency is an Elasticsearch client. Here’s a complete example of the client interface with a dummy implementation:

trait Event

trait ElasticsearchClient[F[_]] {
  def index(event: Event): F[Unit]
}

object ElasticsearchClient {
  def create[F[_]](config: ElasticsearchConfig)(
    implicit F: Sync[F]
  ): ElasticsearchClient[F] =
    new ElasticsearchClient[F] {
      def index(event: Event): F[Unit] =
        F.delay(println(s"Indexing $event to ${config.host}:${config.port}"))
    }
}

The config is required to create a new client. There is already a stream emitting the config while evaluating the effect. The basic composition should be familiar:

stream.map(config => ElasticsearchClient.create[IO](config.es))
// fs2.Stream[IO, ElasticsearchClient[IO]] = Stream(..)

Although it’s completely valid, the idea is to stick with modeling everything as streams. Additionally, the config could be needed downstream, making both the config and the client dependencies.

val stream: fs2.Stream[IO, ElasticsearchClient[IO]] = for {
  config <- fs2.Stream.eval(Config.config.load[IO])
  client <- fs2.Stream.emit(ElasticsearchClient.create[IO](config.es))
} yield client

At this step, configuration loading and client creation emit one value, so they can remain in IO without any problem. The composition would be similar. When FS2 shines, though, is in dealing with I/O computations in constant memory.

Composition

It’s common to run both an HTTP server and a Kafka (or any other) consumer concurrently within the scope of one microservice. Continuing this example, an HTTP server could serve data from an Elasticsearch index, while a consumer would index data from a Kafka topic(s) (The code will be more abstract but two good libraries with FS2 support are http4s and fs2-kafka respectively.).

A hypothetical API for a consumer that requires its own config and an Elasticsearch client producing a stream:

def eventStream[F[_]](config: KafkaConfig, client: ElasticsearchClient[F])(
  implicit F: ConcurrentEffect[F]
): fs2.Stream[F, Unit] = ???

And an HTTP server API that is modeled as cats-effect’s Resource:

trait Server[F[_]]

def server[F[_]](implicit F: ConcurrentEffect[F]): Resource[F, Server[F]] = ???

Abstracting over effect type (all these F[_]) following the least powerful typeclass rule is crucial for stream composition and testing. There is no need to constrain an effect with ConcurrentEffect when Sync or Applicative would be enough.

Given that the consumer and the server are independent services, they cannot be composed like dependencies, as they are supposed to run concurrently. Instead of “sequential” flatMap, we can use parJoin, which will nondeterministically merge a stream of streams into a single stream. To run multiple streams concurrently with FS2:

val stream: fs2.Stream[IO, Unit] = for {
  config <- fs2.Stream.eval(Config.config.load[IO])
  client <- fs2.Stream.emit(ElasticsearchClient.create[IO](config.es))
  app <- fs2.Stream(
    eventStream(config.kafka, client),
    fs2.Stream.resource(server[IO]) // Resource[F, Server[F]] → Stream[F, Server[F]]
  ).parJoinUnbounded
} yield app

And that’s it! More information about parJoin and parJoinUnbounded (For this example parJoin(2) and parJoinUnbounded are interchangeable but when adding one more service it’s easy to forget to bump the maxOpen arguments of parJoin.), as well as many other composition primitives such as concurrently (when the services depend on each other), can be found here. The complete snippet of the main entry point with basic error handling would look like this:

object Main extends IOApp {
  def run(args: List[String]): IO[ExitCode] = {
    val stream: fs2.Stream[IO, Unit] = for {
      config <- fs2.Stream.eval(Config.config.load[IO])
      client <- fs2.Stream.emit(ElasticsearchClient.create[IO](config.es))
      app <- fs2.Stream(
        eventStream(config.kafka, client),
        fs2.Stream.resource(server[IO])
      ).parJoinUnbounded
    } yield app

    stream
      .compile
      .drain
      .onError {
        case e: Throwable => IO(println(e.getMessage))
      }
      .as(ExitCode.Success)
  }
  // ...
}

Not only is the flow concise, but every line packs a punch. The part with eval includes type-safe configuration loading semantics, performs it in parallel with validation, and evaluates an effect. Then, any value can be represented as a stream, be it effectful (with eval) or not (with emit, where a sequence of values is handled with emits). The main part comprises an unbounded stream of messages and a server with built-in resource safety (resource) that run concurrently (parJoin, where service independence can be controlled by changing the concurrency operation). Adding a third service to the example (let say a message producer along with the consumer) won’t change the structure at all. To finish things off, there is a centralized place for error handling and straightforward support for testing. Mocking every dependency and service in unit tests can be avoided by replacing the IO effect with another. Streaming semantics are already tested by the FS2 library itself, which allows focusing on service logic.

Why not use compositional streaming modeling for applications with streaming core?

Scaling effects in Scala

Sun, 10 Feb 2019 15:32:58 +0200

42. A value. Different from 37 but both are numbers.

Effect is to type what type is to value.

🄰s type allows to abstract over value, effect allows to abstract over type.

42: Int
"text": String

def opt[A](value: A): Option[A]
opt(42): Option[Int]
opt("text"): Option[String]

There are two types of effect scaling—horizontal ⇄ and vertical ⇅.

Horizontal Scaling

The simplest scaling approach that replaces an effect with another one.

Built-in effects: Option[A] → Future[A].

Third-party effects: Future[A] → Task[A]. (https://monix.io)

import scala.concurrent.Future
import monix.eval.Task
import monix.execution.Scheduler.Implicits.global

val future = Future(42)
val task = Task.deferFuture(future) // Future[A] → Task[A]
task.runToFuture: Future[Int]       // Task[A] → Future[A]

val io = task.toIO                  // Task[A] → IO[A]
io.to[Task]: Task[Int]              // IO[A] → Task[A]

// Future[A] → Task[A] → IO[A] → Task[A] → Future[A]
Task.fromFuture(Future(42)).toIO.to[Task].runToFuture: Future[Int]

Horizontal scaling is bidirectional if types are isomorphic, and unidirectional otherwise.

Vertical Scaling

⼐ With ability to abstract over types, the next step is to abstract over effects.

Algebra is to effect what effect is to type.

One way to abstract over effects is the tagless final encoding.

import java.util.UUID

final case class Snapshot(tenantId: UUID, id: UUID)

trait SnapshotRepo[F[_]] {
  def current(tenantId: UUID): F[Option[Snapshot]]
}

Another way is using Free monads. (https://typelevel.org/cats/datatypes/freemonad.html)

import cats.~>
import cats.free.Free

sealed trait ObjectStoreAlgebra[A]
final case class Put(path: String, bytes: Array[Byte]) extends ObjectStoreAlgebra[Unit]

object ObjectStore {
  type ObjectStore[A] = Free[ObjectStoreAlgebra, A]

  def put(path: String, bytes: Array[Byte]): ObjectStore[Unit] =
    Free.liftF[ObjectStoreAlgebra, Unit](Put(path, bytes))
}

def compiler[F[_]]: ObjectStoreAlgebra ~> F = ???

⇄ Algebras can be scaled horizontally by replacing one approach with another.

Extensible effects are to algebras what algebras are to effects. (https://atnos-org.github.io/eff)

import cats.Monad
import cats.data.Writer
import org.atnos.eff._
import org.atnos.eff.all._

sealed trait Event
final case class ObjectStored(path: String) extends Event

trait PayloadAlgebra[F[_]] {
  type _effect[R] = F |= R                    // Expect to deal with F effect
  type _writer[R] = Writer[String, ?] |= R    // And the Cats' Writer for logging

  def process[R: _option: _effect: _writer](  // Expect to see some Options
    tenantId:     UUID,
    payload:      Array[Byte],
    snapshotRepo: SnapshotRepo[F],            // Pass tagless final interpreter
    compiler:     ObjectStoreAlgebra ~> F     // Pass free monad compiler
  )(implicit M: Monad[F]): Eff[R, Event] = for {
    _ ← tell("Starting processing")                                 // Cats data type
    maybeSnapshot ← Eff.send(snapshotRepo.current(tenantId))        // Tagless final algebra effect
    snapshot ← fromOption(maybeSnapshot)                            // Built-in Option type
    path = s"${tenantId}/${snapshot.id}"                            // A value
    _ ← Eff.send(ObjectStore.put(path, payload).foldMap(compiler))  // Free monad effect
  } yield ObjectStored(path)
}

⇅ Vertical scaling implies both scaling up and scaling down.

☰ Different effects, along with algebra interpreters, are computed in one place. (https://typelevel.org/cats/datatypes/writer.html)

↓ The API is scaled down being encoded as tagless final trait PayloadAlgebra[F[_]].

↓ Further downscaling is achieved by running the computation.

type Stack = Fx.fx3[Option, Task, Writer[String, ?]]
val algebra = new PayloadAlgebra[Task] {}

val computation: Task[Option[Event]] = algebra
  .process[Stack](tenantId, payload, snapshotRepo, compiler)
  .runOption
  .runWriterUnsafe[String](println)
  .runAsync

⧇ Using the abstract interface, the computation results in nested effects.

⟷ The effect can be scaled horizontally as at the beginning.

View the diagrams on the website

↻

Identify tables to cache with Spark SQL Analyzer

Wed, 28 Feb 2018 19:56:54 +0300

Consider having dozens of Cassandra tables as data sources/sinks, with Spark SQL transformations in between. Sometimes, loading the whole table or a data subset by multiple keys is inevitable. Every optimization applied to the underlying data source and intermediary dataframes goes a long way. One of such optimizations is caching. As the number of transformations grows, the DAG can become unreadable very quickly. It would be nice to be able to check if there are any candidates for caching without having to dig deep into the code. Let’s take a look at some examples.

Use Cases

Basically, any table that is used more than once can be a candidate for caching. Let’s create a test table named t1 (all the code is tested on Spark 2.3.0):

spark.sql("select 1").createOrReplaceTempView("t1")

Now let’s create another table that is based on t1:

spark.sql("select * from t1").createOrReplaceTempView("t2")

When a Spark action is applied to t2 table, t1 is read only once, so it should be fine. What if we create yet another table that is using t1?

spark.sql("select * from t1").createOrReplaceTempView("t3")

t1 is now used twice, by both t2 and t3, so it could potentially benefit from caching. The pitfall is that to determine this, all of the transformations should be analyzed.

Let’s create yet another transformation that does not require caching:

spark.sql("select 2").createOrReplaceTempView("t4")
spark.sql("select * from t4").createOrReplaceTempView("t5")

There are some corner cases that would be nice to avoid. In the following example

spark.sql("select * from (select * from t5) as t1").createOrReplaceTempView("t6")

t1 is an alias, but there is a table with exactly the same name. Therefore, it should not be mistaken for the actual table.

If a table is already cached, it should be excluded from consideration:

spark.sql("select 3").createOrReplaceTempView("t7")
spark.table("t7").cache()
spark.sql("select * from t7").createOrReplaceTempView("t8")
spark.sql("select * from t7").createOrReplaceTempView("t9")

In Action

Following is the code covering the aforementioned use cases:

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.plans.logical.SubqueryAlias

object Analyzer {
  type PlanCache = Map[String, LogicalPlan]
  case class Result(parents: Map[String, List[String]], planCache: PlanCache)
    
  def parents(implicit spark: SparkSession): Map[String, List[String]] = {
    def updateCache(table: String, cache: PlanCache): PlanCache =
      if (cache.contains(table)) cache else cache + (table -> analyzedPlan(table))
      
    def subquery(plan: SubqueryAlias): Boolean =
      plan.child.nodeName != "SubqueryAlias" &&
      spark.catalog.tableExists(plan.alias) &&
      !spark.catalog.isCached(plan.alias)

    @annotation.tailrec
    def parentAliases(
      nodes: Seq[LogicalPlan],
      names: List[String],
      planCache: PlanCache
    ): (List[String], PlanCache) = nodes match {
      case (plan @ SubqueryAlias(alias, child)) +: plans if subquery(plan) =>
        val newCache = updateCache(alias, planCache)
        if (plan.sameResult(newCache(alias))) {
          parentAliases(plans, alias :: names, newCache)
        } else {
          parentAliases(plans ++ subplans(plan), names, planCache)
        }
      case plan +: plans =>
        parentAliases(plans ++ subplans(plan), names, planCache)
      case _ => (names, planCache)
    }

    val tables = spark.catalog.listTables.collect
    tables.foldLeft(Result(Map(), Map())) { case (result, table) =>
      val cache = updateCache(table.name, result.planCache)
      val nodes = cache(table.name).children
      val (aliases, newCache) = parentAliases(nodes, Nil, cache)
      val parents = aliases.foldLeft(result.parents) { case (parents, alias) =>
        val children = table.name :: parents.getOrElse(alias, Nil)
        parents + (alias -> children)
      }
      Result(parents, newCache)
    }.parents
  }
  
  def analyzedPlan(table: String)(implicit spark: SparkSession): LogicalPlan =
    spark.table(table).queryExecution.analyzed
    
  def subplans(plan: LogicalPlan): Seq[LogicalPlan] =
    plan.children ++ plan.subqueries

  def showReused(implicit spark: SparkSession): Unit =
    parents.foreach { case (parent, children) =>
      if (children.size > 1)
        println(s"""$parent is used by ${children.mkString(", ")}""")
    }
}

To ensure efficiency, the code has been optimized to process every registered table only once. It is worth noting that most of the cache references in the code refer to the internal cache rather than the Spark cache. The idea is this:

The analyzed logical plan generated by Spark SQL Analyzer is acquired for every table spark.table(table).queryExecution.analyzed.
The child nodes of a query plan are processed recursively to search for SubqueryAlias nodes.
When a SubqueryAlias node is found, checks are performed to verify that the table exists, is not cached, is not an alias, etc.
The goal is to collect all the tables that are used by any other table in order to identify common parent tables.

If we run the code

scala> Analyzer.showReused(spark)
t1 is used by t3, t2

we can see that:

t5 is not a valid candidate and is therefore missing from the results, as expected;
t6, which uses an alias, is also not included in the results;
t7 is already cached, so there is no need to consider it for caching.

Thus, only table t1 is correctly identified as being used by both t2 and t3, making it a suitable candidate for caching.

While Catalyst is primarily used for logical optimizations, its initial analysis phase can also be useful in cases like this, particularly when migrating large volumes of SQL business logic to Spark SQL. In practice, it proved helpful in identifying a few tables that did, in fact, need to be cached.

Functional builder pattern

Sat, 11 Mar 2017 15:54:04 +0200

The builder pattern is used for constructing objects as a series of steps.

Let’s use a database query builder as an example.

import scala.util.{Failure, Success, Try}

case class Query(
  table: String,
  keyspace: Option[String] = None,
  select: List[String] = Nil,
  where: List[String] = Nil,
  limit: Option[Int] = None
) {
  def build: Try[String] =
    for {
      table <- renderTable
      select <- renderSelect
      where <- renderWhere
      limit <- renderLimit
    } yield s"select $select from $table $where $limit".trim

  def keyspace(name: String): Query =
    this.copy(keyspace = Option(name).filterNot(_.isEmpty))

  def select(col: String): Query =
    this.copy(select = col :: select)

  def where(clause: String): Query =
    this.copy(where = clause :: where)

  def limit(l: Int): Query =
    this.copy(limit = Option(l))

  private def renderTable: Try[String] =
    if (table.isEmpty) Failure(error("Table name cannot be empty"))
    else Success(keyspace.filterNot(_.isEmpty).fold(table)(k => s"$k.$table"))

  private def renderSelect: Try[String] =
    Success(if (select.isEmpty) "*" else select.mkString(", "))

  private def renderWhere: Try[String] =
    Success(if (where.isEmpty) "" else s"""where ${where.mkString(" and ")}""")

  private def renderLimit: Try[String] =
    limit.fold(Try("")) { l =>
      if (l < 1) Failure(error("Limit must be strictly positive"))
      else Success(s"limit $l")
    }

  private def error(msg: String): IllegalArgumentException =
    new IllegalArgumentException(msg)
}

With some basic validation, there are strings everywhere for the sake of simplicity, as it’s not about type-safety now. The query can be constructed as follows:

val query = Query(table = "test")
  .select("col")
  .where("col2 = 1")
  .limit(1)
  .keyspace("stg")
  .build

scala> scala.util.Try[String] = Success(select col from stg.test where col2 = 1 limit 1)

Functional approach

The gist of the functional approach is making each transformation step autonomous by disentangling it from the builder. In other words, each step becomes a first-class citizen function.

object Query {
  def keyspace(name: String): Query => Query = _.keyspace(name)
  def select(col: String):    Query => Query = _.select(col)
  def where(clause: String):  Query => Query = _.where(clause)
  def limit(value: Int):      Query => Query = _.limit(value)
}

Each step is an explicit Query => Query function that delegates the construction logic to the builder. The benefits of such wrapping include:

The original builder can be used without any changes.
The construction can be composed in multiple steps (e.g. query optimization based on client request).
Conditional steps become possible (try to add the where step conditionally in the example above).
Parts of the construction can be done elsewhere and passed for composition.
New steps can be created without the original builder being aware of them.

The initial example can be written like this:

import Query._

val baseQuery = Query(table = "test")

val steps = select("col") andThen where("col2 = 1") andThen limit(1) andThen keyspace("stg")
steps(baseQuery).build

scala> scala.util.Try[String] = Success(select col from stg.test where col2 = 1 limit 1)

One of the potential drawbacks is chain ergonomics with all these andThens. It can be improved by adding a shorter alias to the Function1 type:

implicit class FunctionAlias[A, B](f: Function1[A, B]) {
  def ~[C](g: B => C): A => C = f andThen g
}

val steps = select("col") ~ where("col2 = 1") ~ limit(1) ~ keyspace("stg")
steps(baseQuery).build

scala> scala.util.Try[String] = Success(select col from stg.test where col2 = 1 limit 1)

Another example of function composition is creating a step from a list of columns:

val identity = (q: Query) => q
val cols = List("col1", "col2", "col3").foldLeft(identity)(_ ~ select(_))
val steps = keyspace("namespace") ~ cols
steps(baseQuery).build

scala> scala.util.Try[String] = Success(select col3, col2, col1 from namespace.test)

Arrows

If you already use Cats or need additional functionality, the same can be accomplished with the Arrow type class.

import cats.implicits._

val identity = cats.arrow.Arrow[Function1].id[Query]
val cols = List("col1", "col2", "col3").foldLeft(identity)(_ >>> select(_))
val steps = keyspace("namespace") >>> limit(1) >>> select("col") >>> cols
steps(baseQuery).build

scala> scala.util.Try[String] = Success(select col3, col2, col1, col from namespace.test  limit 1)

Note that >>> is identical to ours ~ as it’s just an alias for andThen.

The whole construction can be built from any collection(s) of steps:

List(
  keyspace("namespace"),
  limit(1),
  select("col"),
  cols
).foldLeft(identity)(_ >>> _)(baseQuery).build

The unobtrusive functional layer on top of any builder is what makes it an elegant approach in my opinion.

Akka persistent views

Fri, 10 Apr 2015 22:09:55 +0300

In this post, I’d like to talk about certain aspects of Akka persistence, which is still experimental at the time of writing. The official documentation does a good job of explaining the concept behind Akka persistence and how to make an actor persistent. However, I noticed one thing that can be confusing to many people when they start dealing with akka-persistence module—Persistent Views. Before we discuss it, let’s start with the pattern that will help to understand persistent views better.

CQRS (Command Query Responsibility Segregation)

The idea of the CQRS pattern is to use separate models for reading and updating data, rather than the more classical way of performing CRUD (create, read, update, delete) operations at the same time and place. In terms of CRUD, create, update, and delete are treated as Commands, while read is considered to be a Query. One of the architectural patterns that CRQS fits with is Event Sourcing (you can read about how Akka persistence implements this idea in the official documentation). Before we delve into persistent views, let’s first take a quick look at non-persistent actors.

CQRS with Non Persistent Actors

Let’s imagine that there is an actor that keeps transactions in its internal state. We can use the same actor to store new transactions and read them. One of the issues with such a design is the need for having completely different representations for display. For instance, we may want to display all invalid transactions and the current balance for a particular account on the administration panel. If the actor is in the middle of updating its state, it cannot be queried, which makes the system less responsive. To separate commands and queries, we will most likely end up with more than one actor, each with its own state reflecting changes from the main actor. This setup will work pretty well until the main actor, with all the transactions, goes down. Once this happens, the state is lost forever. The other actors can still function but the data will stay out of sync.

Persistent Views and CQRS

To make it clear, Persistent View represents the Query part of CQRS. Persistent actors are responsible for managing state (persisting events, deleting messages, restoring from the journal, saving the snapshots, etc.). Persistent views poll messages directly from the journal of a persistent actor, without being coupled to the actor itself. All they have to know is the identifier of the corresponding persistent actor. If that persistent actor crashes, it does not affect persistent views as they can still serve data from the journal. Once the persistent actor is recovered, persistent views will eventually become consistent with it again. Moreover, a persistent view can utilize other data sources to optimize data presentation. Basically, it comes down to this: use persistent actors to handle commands and persistent views to handle queries.

Persistent Views in Action

The working demo project can be found here. The source code consists of three files: one for the persistent actor and the other two for the persistent views. The TransactionActor is a persistent actor responsible for persisting transactions. All the machinery in the source code is primarily focused on maintaining and recovering state. The InvalidTransactionsView is a persistent view that shows all invalid transactions. You can see how it’s decoupled from the TransactionActor. Another persistent view is the BalanceView. Instead of keeping a list of transactions, it keeps a map where it stores the current amount per account. In a traditional approach, both persistent views would be methods of the same object. With separate views, even if the persistent actor and one of the views die, the remaining view will still be able to process queries.

I think Persistent Views play an important role when being used with event sourcing. They help to follow the CQRS pattern more easily than using regular actors would. I would even argue that they are just as important as persistent actors. The official documentation would be more helpful if it provided more information about them. I hope that the analogy with CQRS can help develop a better understanding of their role.

Scala Function object to the rescue

Wed, 01 Apr 2015 20:43:55 +0300

The Function object has been there since Scala version 1.0. It provides a bunch of utility methods for dealing with higher-order functions. Despite its simplicity and usefulness, I have found that not only do many novice developers not use them, but they are also unaware of their existence. In this post, I would like to bring to your attention some of the functions I find helpful in certain scenarios.

`chain`

To get started, let’s introduce two extremely simple functions from Int to Int. One of them, let’s call it inc, will increase a number by one, while the other one, double, with multiply a number by two. (We could use a shorter definition form such as val inc: Int => Int = _ + 1 for both functions.)

scala> val inc: Int => Int = x => x + 1
scala> val double: Int => Int = x => x * 2

So what do we do when we have a sequence of such functions that we want to combine? One of the most popular options in my practice is the following:

scala> List(inc, double, inc) reduce (_ andThen _)

This piece of code takes a sequence of functions and combines them using the andThen method starting with the first one in the list and returns the resulting function from Int to Int. It is by no means a bad code. Let’s see how we can simplify it by using the chain function from the Function object:

scala> Function.chain(List(inc, double, inc))

It does exactly the same as the previous snippet but it might be more intuitive. If we first import all the functions from the Function object, it will be dead simple:

scala> import Function._
scala> chain(List(inc, double, inc))

If you want to use the compose function instead of the andThen, you can still achieve the same with the chain by reversing the sequence first:

scala> chain(List(inc, double, double).reverse)

`tupled`

Imagine there is a tuple of two elements, id and name, representing a user wrapped into Option:

scala> val user: Option[(Int, String)] = Some(1, "Bob")

And a function, let’s call it auth, which accepts id and name as two separate arguments:

scala> val auth: (Int, String) => Boolean = (id, name) => id == 1 && name == "Bob"

We can’t just map the auth function over the user because the former expects two arguments while our user is a tuple with two elements. One option would be to extract the id and name and pass them as separate arguments to the auth function but it requires some boilerplate to write. That’s when we can use the tupled function from the Function object. It takes a function with, let’s say, two arguments, and converts it into a function that takes a tuple of two elements with the same types as the original arguments. That’s exactly what we need to map the auth function over the user:

scala> user map auth.tupled
res4: Option[Boolean] = Some(true)

There are tupled functions defined for functions with arity from two to five inclusive. I find it convenient as I don’t use tuples of more than two or three elements that often.

Technically, in the previous example, the tupled was called on the function itself (as it’s defined both in the Function object and the Function* traits). Another “trick” that can be useful is mapping over a hash map in what many people would say a natural way. (It also demonstrates the usage of the tupled function from the Function object and not defined in the Function* trait.) In Scala, you can’t write code like this:

scala> val m = Map(1 -> "first", 2 -> "second")
scala> m map { (k, v) => s"$k:$v" }

<console>:12: error: missing parameter type
Note: The expected type requires a one-argument function accepting a 2-Tuple.
      Consider a pattern matching anonymous function, `{ case (k, v) =>  ... }`
              m map { (k, v) => s"$k:$v" }
                       ^
<console>:12: error: missing parameter type
              m map { (k, v) => s"$k:$v" }

What many developers would do is something like this:

scala> m map { case (k, v) => s"$k:$v" }
res5: scala.collection.immutable.Iterable[String] = List(1:first, 2:second)

You can achieve the same using the tupled function:

scala> m map tupled { (k, v) => s"$k:$v" }
res6: scala.collection.immutable.Iterable[String] = List(1:first, 2:second)

Maybe not so useful but it helps to understand its application.

`unlift`

To illustrate the usage of the unlift, let’s write a function that takes an Int and returns Some(x) if x is equal to or greater than zero, and None otherwise:

scala> val f: Int => Option[Int] = Option(_) filter (_ >= 0)

The unlift function turns an A => Option[B] function into a PartialFunction[A, B]. It allows us to use our function in any place where a partial function is required. To make it clear, this how we can use our function to filter out negative integers from a list:

scala> import Function._
scala> List(-1,0,1) collect unlift(f)
res7: List[Int] = List(0, 1)

I don’t use this on a daily basis but there are cases like this where it comes very handy.

Bonus: there is an opposite function defined in the PartialFunction called lift. To see how it is related to the unlift function, the following equation is always true:

scala> f == unlift(f).lift
res8: Boolean = true

`uncurried/untupled`

These two functions are the opposite to the curried and the tupled functions, respectively. I haven’t seen them being used as frequently as the ones described above.

Although there is nothing new, I believe that we often overlook some of the useful API provided by the Scala standard library. I hope that this refresher is on time and can save you a few lines of code every now and then.