DEV Community: eblocha

Parallel Matrix Multiplication in Rust

eblocha — Sat, 01 Oct 2022 12:56:52 +0000

Much of linear algebra revolves around matrix multiplications. Since this is an O(n^3) operation, other techniques are needed to improve its performance. Today, I will show how I parallelized matrix multiplication, improving speed by 4x on my 12-thread machine.

In my examples, I am using the nalgebra crate to handle matrix storage (which stores dynamically sized matrices as a Vec in column-major order). This isn't too important, the only thing that really matters is the ability to index the matrix by row and column index.

Single-Threaded Implementation

In order to multiply matrices, we need to iterate over the columns of the rhs, then for each column, iterate over the rows of the lhs. We then zip the lhs row with the rhs column, taking the dot product of these two arrays.

let l_shape = lhs.shape();
let r_shape = rhs.shape();

// check for shape compatibility here...

// the multiplication
let result: Vec<f64> = (0..r_shape.1).flat_map(move |rj| {
    (0..l_shape.0).flat_map(move |li| {
        (0..r_shape.0)
            .zip(0..l_shape.1)
            .map(move |(ri, lj)| {
                lhs.index((li, lj)) * rhs.index((ri, rj))
            })
            .sum::<f64>()
    })
})
.collect();

// result is a vec in column-major order

Parallelizing

I will now use rayon to parallelize this operation. It's way simpler than you may think!

let result: Vec<f64> = (0..r_shape.1).into_par_iter().flat_map(move |rj| {
    (0..l_shape.0).into_par_iter().flat_map(move |li| {
        (0..r_shape.0)
            .zip(0..l_shape.1)
            .map(move |(ri, lj)| {
                lhs.index((li, lj)) * rhs.index((ri, rj))
            })
            .sum::<f64>()
    })
})
.collect();

That's it! Just adding into_par_iter after the ranges is all that is needed. I benchmarked this on an Intel i5-10600K (12) @ 4.800GHz, multiplying a 1000x1000 matrix with a 1000x1 vector, and the average execution time went from 8 ms down to 2 ms. That is a 75% improvement in speed.

Logging in Python Command Line Applications

eblocha — Sun, 03 Apr 2022 20:24:58 +0000

Introduction

Logging is a necessary part of any serious application. It allows you and your users to effectively debug issues. In this article, I will showcase a good way to configure logging for command-line python apps. You can see the finished product as a github gist here.

Goal

The goal of our logger will have multiple parts:

Make the console output pretty and easy to read by a human
Save additional output to a log file that can be parsed by a computer
Save all traceback information for debugging later
Allow users to change the verbosity of the console and log file output

We are going to make the console output look like this:

And the log file look like this:



DEBUG:2022-04-03 15:41:17,920:root:A debug message
INFO:2022-04-03 15:41:17,920:root:An info message
WARNING:2022-04-03 15:41:17,920:root:A warning message
ERROR:2022-04-03 15:41:17,920:root:An error message
CRITICAL:2022-04-03 15:41:17,920:root:A critical message from an exception
    Traceback (most recent call last):
        /home/eb/projects/py-scratch/color-log.py  <module>  327: raise ValueError("A critical message from an exception")
    ValueError: A critical message from an exception

We've got a few tricky things going on here.

We are using different colors and formats based on the log level
We are using different formats for the console and file output
We are changing the traceback format

Formatter

The python logging formatter does not allow different format strings based on the log level, so we will need to implement our own formatter.



import typing as t

class MultiFormatter(PrettyExceptionFormatter):
    """Format log messages differently for each log level"""

    def __init__(self, formats: t.Dict[int, str] = None, **kwargs):
        base_format = kwargs.pop("fmt", None)
        super().__init__(base_format, **kwargs)

        formats = formats or default_formats

        self.formatters = {
            level: PrettyExceptionFormatter(fmt, **kwargs)
            for level, fmt in formats.items()
        }

    def format(self, record: logging.LogRecord):
        formatter = self.formatters.get(record.levelno)

        if formatter is None:
            return super().format(record)

        return formatter.format(record)

In our MultiFormatter class, we take a mapping of log level to a format string, then create different formatters for each level. In .format(), we dispatch to the formatter for the level logged.

Now, what is PrettyExceptionFormatter? We also need to implement that. It will format the traceback and exception information when it is included in the log record.



from textwrap import indent
from pretty_traceback.formatting import exc_to_traceback_str

class PrettyExceptionFormatter(logging.Formatter):
    """Uses pretty-traceback to format exceptions"""

    def __init__(self, *args, color=True, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.color = color

    def formatException(self, ei):
        _, exc_value, traceback = ei
        return exc_to_traceback_str(exc_value, traceback, color=self.color)

    def format(self, record: logging.LogRecord):
        record.message = record.getMessage()

        if self.usesTime():
            record.asctime = self.formatTime(record, self.datefmt)

        s = self.formatMessage(record)

        if record.exc_info:
            # Don't assign to exc_text here, since we don't want to inject color all the time
            if s[-1:] != "\n":
                s += "\n"
            # Add indent to indicate the traceback is part of the previous message
            text = indent(self.formatException(record.exc_info), " " * 4)
            s += text

        return s

We're using the awesome pretty-traceback package here. Since the default behavior of logging.Formatter is to modify the record.exc_text with the output of .formatException(), we need to override that behavior. This is because we are adding ANSI color, and don't want to see that in log files.

In the standard logging.Formatter implementation, the record is modified when formatting exceptions (as of python 3.10.2):



def format(self, record):
    ...
    # exc_text is MODIFIED, which propagates to other formatters for other handlers
    record.exc_text = self.formatException(record.exc_info)
    ...
    return s

The MultiFormatter class takes an argument to change the format string per-level, with the default strings as follows:



default_formats = {
    logging.DEBUG: style("DEBUG", fg="cyan") + " | " + style("%(message)s", fg="cyan"),
    logging.INFO: "%(message)s",
    logging.WARNING: style("WARN ", fg="yellow") + " | " + style("%(message)s", fg="yellow"),
    logging.ERROR: style("ERROR", fg="red") + " | " + style("%(message)s", fg="red"),
    logging.CRITICAL: style("FATAL", fg="white", bg="red", bold=True) + " | " + style("%(message)s", fg="red", bold=True),
}

This adds a vertical line between the level name and message, except for info messages, which are passed through as plain messages. The style function here is a direct copy of the click.style utility.

Context Managers

The end goal here is to simply call with cli_log_config():, and get beautiful output. We will need some context managers. Starting with a logging context, straight from the python docs:



class LoggingContext:
    def __init__(
        self,
        logger: logging.Logger = None,
        level: int = None,
        handler: logging.Handler = None,
        close: bool = True,
    ):
        self.logger = logger or logging.root
        self.level = level
        self.handler = handler
        self.close = close

    def __enter__(self):
        if self.level is not None:
            self.old_level = self.logger.level
            self.logger.setLevel(self.level)

        if self.handler:
            self.logger.addHandler(self.handler)

    def __exit__(self, *exc_info):
        if self.level is not None:
            self.logger.setLevel(self.old_level)

        if self.handler:
            self.logger.removeHandler(self.handler)

        if self.handler and self.close:
            self.handler.close()

Next, we create a special context manager to combine multiple context managers dynamically:



class MultiContext:
    """Can be used to dynamically combine context managers"""

    def __init__(self, *contexts) -> None:
        self.contexts = contexts

    def __enter__(self):
        return tuple(ctx.__enter__() for ctx in self.contexts)

    def __exit__(self, *exc_info):
        for ctx in self.contexts:
            ctx.__exit__(*exc_info)

Finally, we can combine everything we've done so far into a single context manager:



def cli_log_config(
    logger: logging.Logger = None,
    verbose: int = 2,
    filename: str = None,
    file_verbose: int = None,
):
    """
    Use a logging configuration for a CLI application.
    This will prettify log messages for the console, and show more info in a log file.

    Parameters
    ----------
    logger : logging.Logger, default None
        The logger to configure. If None, configures the root logger
    verbose : int from 0 to 3, default 2
        Sets the output verbosity.
        Verbosity 0 shows critical errors
        Verbosity 1 shows warnings and above
        Verbosity 2 shows info and above
        Verbosity 3 and above shows debug and above
    filename : str, default None
        The file name of the log file to log to. If None, no log file is generated.
    file_verbose : int from 0 to 3, default None
        Set a different verbosity for the log file. If None, is set to `verbose`.
        This has no effect if `filename` is None.

    Returns
    -------
    A context manager that will configure the logger, and reset to the previous configuration afterwards.
    """

    if file_verbose is None:
        file_verbose = verbose

    verbosities = {
        0: logging.CRITICAL,
        1: logging.WARNING,
        2: logging.INFO,
        3: logging.DEBUG,
    }

    console_level = verbosities.get(verbose, logging.DEBUG)
    file_level = verbosities.get(file_verbose, logging.DEBUG)

    # This configuration will print pretty tracebacks with color to the console,
    # and log pretty tracebacks without color to the log file.

    console_handler = logging.StreamHandler()
    console_handler.setFormatter(MultiFormatter())
    console_handler.setLevel(console_level)

    contexts = [
        LoggingContext(logger=logger, level=min(console_level, file_level)),
        LoggingContext(logger=logger, handler=console_handler, close=False),
    ]

    if filename:
        file_handler = logging.FileHandler(filename)
        file_handler.setFormatter(
            PrettyExceptionFormatter(
                "%(levelname)s:%(asctime)s:%(name)s:%(message)s", color=False
            )
        )
        file_handler.setLevel(file_level)
        contexts.append(LoggingContext(logger=logger, handler=file_handler))

    return MultiContext(*contexts)

We now have the option to specify a verbosity level, a log file, and a different verbosity for the file. Try out the example:



with cli_log_config(verbose=3, filename="test.log"):

    try:

        logging.debug("A debug message")

        logging.info("An info message")

        logging.warning("A warning message")

        logging.error("An error message")

        raise ValueError("A critical message from an exception")

    except Exception as exc:

        logging.critical(str(exc), exc_info=True)

Conculsion

In this article, we:

Implemented a custom logging formatter to dynamically format messages based on log level
Added color to console log output
Prettified exceptions in log messages
Wrapped everything into a reusable context manager

I hope you enjoyed, and drew some inspiration for making your CLI apps more user-friendly and robust!

Using Dataclasses for Configuration in Python

eblocha — Sat, 02 Apr 2022 15:17:49 +0000

Introduction

Defining configuration schemes for python apps can be tricky. In this article, I will showcase a simple and effective way to make handling configuration easy, with only the standard library.

Format

For this method, any file format will work, as long as you can parse it into a python dict.

A Simple Example

We'll make a simple example with this made-up config from wikipedia:

---
receipt:     Oz-Ware Purchase Invoice
date:        2012-08-06
customer:
    first_name:   Dorothy
    family_name:  Gale

items:
    - part_no:   A4786
      descrip:   Water Bucket (Filled)
      price:     1.47
      quantity:  4

    - part_no:   E1628
      descrip:   High Heeled "Ruby" Slippers
      size:      8
      price:     133.7
      quantity:  1

bill-to:  &id001
    street: |
            123 Tornado Alley
            Suite 16
    city:   East Centerville
    state:  KS

ship-to:  *id001

specialDelivery:  >
    Follow the Yellow Brick
    Road to the Emerald City.
    Pay no attention to the
    man behind the curtain.

We can parse this into python to get a dict:

config = {
    "reciept": "Oz-Ware Purchase Invoice",
    "date": "2012-08-06",
    "customer": {
        "first_name": "Dorothy",
        "family_name": "Gale",
    },
    "items": [
        {
            "part_no": "A4786",
            "descrip": "Water Bucket (Filled)",
            "price": 1.47,
            "quantity": 4,
        },
        {
            "part_no": "E1628",
            "descrip": "High Heeled \"Ruby\" Slippers",
            "size": 8,
            "price": 133.7,
            "quantity": 1,
        },
    ],
    "bill-to": {
        "street" : "123 Tornado Alley\nSuite 16",
        "city": "East Centerville",
        "state": "KS",
    },
    "ship-to": {
        "street" : "123 Tornado Alley\nSuite 16",
        "city": "East Centerville",
        "state": "KS",
    },
    "specialDelivery": "Follow the Yellow Brick Road to the Emerald City. Pay no attention to the man behind the curtain.",
}

However, using this in the code is cumbersome. config["customer"]["first_name"] is prone to error, and difficult to refactor.

Dataclasses

Dataclasses will make our life much easier. We define the config properties, sub-properties, and types in a config.py file:

import typing as t
from dataclasses import dataclass
from datetime import date

@dataclass
class Item:
    part_no: str
    description: str
    price: float
    quantity: int
    size: int = None

    def __post_init__(self):
        # Do some validation
        if self.quantity <= 0:
            raise ValueError("quantity must be greater than zero")

    @classmethod
    def from_dict(cls: t.Type["Item"], obj: dict):
        return cls(
            part_no=obj["part_no"],
            description=obj["descrip"],
            price=obj["price"],
            quantity=obj["quantity"],
            size=obj.get("size"),
        )

@dataclass
class Customer:
    first_name: str
    family_name: str

    @classmethod
    def from_dict(cls: t.Type["Customer"], obj: dict):
        return cls(
            first_name=obj["first_name"],
            last_name=obj["family_name"],
        )

@dataclass
class Address:
    street: str
    city: str
    state: str

    @classmethod
    def from_dict(cls: t.Type["Address"], obj: dict):
        return cls(
            street=obj["street"],
            city=obj["city"],
            state=obj["state"],
        )

@dataclass
class Order:
    reciept: str
    date: date
    customer: Customer
    items: t.Sequence[Item]
    bill_to: Address
    ship_to: Address
    special_delivery: str = None

    @classmethod
    def from_dict(cls: t.Type["Order"], obj: dict):
        return cls(
            receipt=obj["reciept"],
            date=date(obj["date"]),
            customer=Customer.from_dict(obj["customer"]),
            items=[Item.from_dict(item) for item in obj["items"]),
            bill_to=Address.from_dict(obj["bill-to"]),
            ship_to=Address.from_dict(obj["ship-to"]),
            special_delivery=obj.get("specialDelivery"),
        )

Now, when we want to use the config in our application, we can simply do:

raw_config = {...}

config = Order.from_dict(raw_config)

config.customer.first_name

This method has a ton of benefits:

We get code completion and type hints in the editor
It's easier to maintain, since you only have to change a config property name in one place
Can implement version reconciliation in the from_dict method
Refactoring is a breeze, since editors can auto-refactor class property names
Allows you to define configurations with python code, since you can instantiate the dataclasses directly in a settings.py file, for example
It's testable:

import unittest
from .config import Order

class TestOrderConfig(unittest.TestCase):
    def test_example_config(self):
        raw_config = {...}
        expected = Order(
            customer=Customer(...),
            ...
        )

        self.assertEqual(Order.from_dict(raw_config), expected)

Hopefully you found this useful, and can use this method to clean up some of your projects!