Terence Z. Liu

The Wonders of RocksDB in Rust (Part III - MapReduce Group-By)

2026-02-12T00:00:00+00:00

You can check out the code here.

This is the third and final part of the RocksDB in Rust series. In Part I, we looked at how to interact with RocksDB in Rust. In Part II, we looked at how to do parallel two-pointer set operations. In this part, we’ll look at how to do MapReduce for out-of-core aggregations.

Let’s look at the famed algorithm, MapReduce (Wikipedia). MapReduce is a programming model for processing large datasets in parallel originally targeting a big cluster of machines. It is a key component of the Hadoop ecosystem, and is still widely used today. MapReduce is a three-phase process:

Map: The input is split into independent chunks, and each chunk is processed in parallel by a map function.
Shuffle: The output of the map phase is shuffled and sorted by the grouping key.
Reduce: The output of the shuffle phase is aggregated by a reduce function.

From the original paper abstract:

Our implementation of MapReduce runs on a large cluster of commodity machines and is highly scalable: a typical MapReduce computation processes many terabytes of data on thousands of machines.

Back in 2004, TBs of data on HDDs were considered large. In 2026, TBs of data can all fit on one machine with SSDs in RAID0. For example, on AWS EC2, you can get 120TB of NVMe drives on a single i7ie.48xlarge/i8ge.48xlarge instance. We can use one machine to run MapReduce to solve our Big Data problems.

Suppose we have a table with the primary key being hex strings simulating hashes. There is one other column that contains 3 hex characters. In practice, you can store anything into the value column with protobuf or other serialization formats. We want to aggregate the keys by the value, to turn this:

(5d46a9cbdcd7b1ee, 24f)
(41c35e778b34002b, 53a)
(554c67ebe57cbef2, 24f)

into this:

(24f, [5d46a9cbdcd7b1ee, 554c67ebe57cbef2])
(53a, [41c35e778b34002b])

For tabular data, this operation is also called group-by. DuckDB and Polars can do this, but they start to stress out beyond a certain size (several billions of entries), even with disk-spilling, because a lot of their query execution plans are in-memory. We can use RocksDB and Rust to do this, and it’s a lot more efficient without much memory usage at all.

See the full code.

Map Phase

We iterate through the database in parallel, and use the value with some unique identifier as the new key, and the original key as the new value to write into an intermediate database. One of the best ways to generate a unique identifier is to use the hex of the original key. Generating a hash from the key, or a random hash with sufficient bytes works, but there is a theoretical chance of collision.

We need a separator between the grouping key and the hex suffix. Using . is simple but has a pitfall: if one grouping key is a prefix of another, lexicographic order interleaves them. It is not the case for our specific example, because the values are simply 3 hex characters. However, real data may have all kinds of characters in the value. For example keys 24f[.0bc], 24f.A[.abc], 24f[.bbc] ([] only for illustration) sort correctly — so group 24f is split by 24f.A, and the Reduce phase (which assumes all entries for the same group are consecutive) will emit duplicate and incomplete groups, overwriting entries. To fix, we can use a reserved byte that cannot appear in the grouping key (e.g. \0). Then the grouping key sorts contiguously and you parse by splitting on that byte.

(key, value) -> (value + '\0' + hex(key), key)

let prefixes = generate_consecutive_hex_strings(3);
let pb = make_progress_bar(Some(prefixes.len() as u64));

let count = prefixes
    .into_par_iter()
    .map(|prefix| {
        let prefix = prefix.as_bytes();
        let mut db_iter =
            db.full_iterator(IteratorMode::From(prefix, Direction::Forward));
        let mut count = 0;
        let mut write_batch = rust_rocksdb::WriteBatch::default();
        while let Some(item) = db_iter.next() {
            let (key, value) = item.unwrap();
            if &key[..prefix.len()] != prefix {
                break;
            }

            let key_hex = hex::encode(key.as_ref());
            let new_key: Vec<u8> = value
                .iter()
                .chain(std::iter::once(&0u8))
                .chain(key_hex.as_bytes())
                .cloned()
                .collect();
            let new_value = key;

            write_batch.put(&new_key, &new_value);
            count += 1;
        }
        output_db.write_without_wal(&write_batch).unwrap();
        pb.inc(1);
        count
    })
    .reduce(|| 0_usize, |acc, c| acc + c);

output_db.flush()?;

pb.finish_with_message("done");

Shuffle Phase

Shuffle in RocksDB is done by simply running a full compaction after bulk loading. The sorted string table (SST) files are already sorted by the key, so the compaction will naturally shuffle the data by the new key.

let mut compaction_opts = rust_rocksdb::CompactOptions::default();
compaction_opts.set_exclusive_manual_compaction(true);
compaction_opts.set_change_level(true);
compaction_opts.set_target_level(ROCKSDB_NUM_LEVELS - 1);
compaction_opts
    .set_bottommost_level_compaction(rust_rocksdb::BottommostLevelCompaction::ForceOptimized);
output_db.compact_range_opt(None::<&[u8]>, None::<&[u8]>, &compaction_opts);

Reduce Phase

We take the new key value + '\0' + hex(key), split on the first \0 to get the grouping key (original value), and group the original keys with ‘|’ to get the new value. This is a simple way to serialize the new value, and you can use any other serialization format you want.

(value + '\0' + hex(key), key) -> (value, key1|key2|...|keyN)

let prefixes = generate_consecutive_hex_strings(3);
let pb = make_progress_bar(Some(prefixes.len() as u64));

let counts = prefixes
    .into_par_iter()
    .map(|prefix| {
        let prefix = prefix.as_bytes();
        let mut db_iter =
            db.full_iterator(IteratorMode::From(prefix, Direction::Forward));
        let mut write_batch = rust_rocksdb::WriteBatch::default();
        let mut count = 0;
        let mut count_grouped = 0;
        let mut prev_key = Vec::<u8>::new();
        let mut blobs_vec: Vec<Vec<u8>> = vec![];
        while let Some(item) = db_iter.next() {
            let (key, value) = item.unwrap();
            if &key[..prefix.len()] != prefix {
                break;
            }

            // key is value + '\0' + key_hex; group by value = everything before first '\0'
            let sep = key.iter().position(|&b| b == 0).unwrap_or_else(|| {
                panic!("Invalid key: {}", String::from_utf8_lossy(&key))
            });
            let new_key = key[..sep].to_vec();

            if new_key != prev_key {
                if !prev_key.is_empty() {
                    // concatenate with '|'
                    // can use protobuf or anything else to serialize
                    let new_value: Vec<u8> = blobs_vec.join(&b"|"[..]);
                    write_batch.put(prev_key, new_value);
                    count_grouped += 1;
                }
                blobs_vec = vec![];
                prev_key = new_key;
            }

            blobs_vec.push(value.to_vec());
            count += 1;
        }

        if !blobs_vec.is_empty() {
            let new_value: Vec<u8> = blobs_vec.join(&b"|"[..]);
            write_batch.put(prev_key, new_value);
            count_grouped += 1;
        }
        output_db.write_without_wal(&write_batch).unwrap();
        pb.inc(1);
        (count, count_grouped)
    })
    .reduce(
        || (0_usize, 0_usize),
        |accs, counts| (accs.0 + counts.0, accs.1 + counts.1),
    );

output_db.flush()?;

pb.finish_with_message("done");

Then you get the final results shown in the image at the beginning of this post.

The Wonders of RocksDB in Rust (Part II - Set Operations)

2026-02-11T00:00:00+00:00

You can check out the code here.

This is the second part of the RocksDB in Rust series. In Part I, we looked at how to interact with RocksDB in Rust. In this part, we’ll look at how to do parallel two-pointer set operations.

How does one get the intersection of two sorted lists in memory?

One could simply convert each list into a hash set to get the intersection. The time complexity is O(n + m) for the conversion, and O(min(n, m)) for the intersection, with n and m being the lengths of the two lists. The hash set lookup is constant. The space complexity is O(n + m) for the conversion.

def get_intersection(list1: list[int], list2: list[int]) -> set:
    set1 = set(list1)
    set2 = set(list2)
    return set1 & set2

You could get away with converting only one list, then iterate through the other list to check if the element is in the set. The time complexity is still O(n + m) and the space complexity is O(n). This is how DuckDB and Polars do it. For equi-joins (match by equality), both use the same idea: hash join — build a hash table on one side (usually the smaller), then probe with the other. So it’s the same O(n + m) time complexity and O(n) space complexity, just at columnar scale with SIMD and careful memory layout. When the build side doesn’t fit in memory, DuckDB switches to a partitioned out-of-core hash join (spill to disk, then merge). Polars’ streaming engine processes in batches but does not spill join state to disk — joins can still be memory-heavy or fall back to the in-memory engine, so very large joins may OOM.

A Tale of Two Pointers

What if you have way too much data, and don’t want the conversion overhead and disk spilling? Consider two tables with several TBs each. Turns out you can do a lot with the sortedness property. One could iterate through the shorter list and check if the element is in the longer list by binary search. The time complexity is O(n log m), n being the length of the shorter list. This can be worse than the hash set approach if the two list lengths are similar. A better approach in that case is to use the two-pointer technique to get the intersection in one pass. You can in fact get all the set logic counts in one pass. The time complexity is again O(n + m). These two approaches have no additional space complexity.

To keep it simple, let’s still use two lists in memory for the example.

count_left, count_right, count_intersection = 0, 0, 0
i, j = 0, 0

while i < len(list1) and j < len(list2):
    if list1[i] == list2[j]:
        count_left += 1
        count_right += 1
        count_intersection += 1
        i += 1
        j += 1
    elif list1[i] < list2[j]:
        count_left += 1
        i += 1
    else:
        count_right += 1
        j += 1

# drain remaining elements
while i < len(list1):
    count_left += 1
    i += 1

while j < len(list2):
    count_right += 1
    j += 1

Two RocksDBs, One Pass

Now let’s say you have two RocksDBs with similar data. The keys are hex strings emulating hashes. You want to draw the Venn diagram of the two sets of keys. The two-pointer approach is perfect for this, because the sorted string table (SST) files are already sorted. So you can do something similar (full code):

let db_left = open_rocksdb_for_read_only(&args.db_dir_left, true)?;
let db_right = open_rocksdb_for_read_only(&args.db_dir_right, true)?;

let mut db_iter_left = db_left.full_iterator(IteratorMode::Start);
let mut db_iter_right = db_right.full_iterator(IteratorMode::Start);

let mut count_left = 0;
let mut count_right = 0;
let mut count_intersection = 0;
let mut item_left = db_iter_left.next();
let mut item_right = db_iter_right.next();

// Don't use take() — we must keep the item we don't advance so it's compared again next iteration.
while let (Some(Ok((blob_left, _))), Some(Ok((blob_right, _)))) =
    (item_left.as_ref(), item_right.as_ref())
{
    if blob_left == blob_right {
        count_left += 1;
        count_right += 1;
        count_intersection += 1;
        item_left = db_iter_left.next();
        item_right = db_iter_right.next();
    } else if blob_left < blob_right {
        count_left += 1;
        item_left = db_iter_left.next();
    } else {
        count_right += 1;
        item_right = db_iter_right.next();
    }
}

while item_left.is_some() {
    count_left += 1;
    item_left = db_iter_left.next();
}

while item_right.is_some() {
    count_right += 1;
    item_right = db_iter_right.next();
}

Run in Parallel

We can adapt the two-pointer technique to run in parallel with the parallel scan pattern in Part I (full code):

let db_left = open_rocksdb_for_read_only(&args.db_dir_left, true)?;
let db_right = open_rocksdb_for_read_only(&args.db_dir_right, true)?;

let prefixes = generate_consecutive_hex_strings(3);
let pb = make_progress_bar(Some(prefixes.len() as u64));

let counts = prefixes
    .into_par_iter()
    .map(|prefix_str| {
        let prefix = prefix_str.as_bytes();

        let mut db_iter_left =
            db_left.full_iterator(IteratorMode::From(prefix, Direction::Forward));
        let mut db_iter_right =
            db_right.full_iterator(IteratorMode::From(prefix, Direction::Forward));

        // two pointers
        let mut count_left = 0;
        let mut count_right = 0;
        let mut count_intersection = 0;
        let mut item_left = db_iter_left.next();
        let mut item_right = db_iter_right.next();

        // Don't use take() — keep the item we don't advance for the next comparison.
        while let (Some(Ok((blob_left, _))), Some(Ok((blob_right, _)))) =
            (item_left.as_ref(), item_right.as_ref())
        {
            if &blob_left[..prefix.len()] != prefix || &blob_right[..prefix.len()] != prefix {
                break;
            }

            if blob_left == blob_right {
                count_left += 1;
                count_right += 1;
                count_intersection += 1;
                item_left = db_iter_left.next();
                item_right = db_iter_right.next();
            } else if blob_left < blob_right {
                count_left += 1;
                item_left = db_iter_left.next();
            } else {
                count_right += 1;
                item_right = db_iter_right.next();
            }
        }

        while let Some(Ok((blob_left, _))) = item_left.as_ref() {
            if &blob_left[..prefix.len()] != prefix {
                break;
            }
            count_left += 1;
            item_left = db_iter_left.next();
        }

        while let Some(Ok((blob_right, _))) = item_right.as_ref() {
            if &blob_right[..prefix.len()] != prefix {
                break;
            }
            count_right += 1;
            item_right = db_iter_right.next();
        }

        pb.inc(1);
        Counts {
            count_left,
            count_right,
            count_intersection,
        }
    })
    .reduce(
        || Counts {
            count_left: 0,
            count_right: 0,
            count_intersection: 0,
        },
        |accs, counts| Counts {
            count_left: accs.count_left + counts.count_left,
            count_right: accs.count_right + counts.count_right,
            count_intersection: accs.count_intersection + counts.count_intersection,
        },
    );

pb.finish_with_message("done");

This is a very powerful pattern that can be used to solve a lot of problems. It can be used to get the intersection, union, difference, and symmetric difference of two sets of data, and apply arbitrary business logic to those parts of data.

Two-pointer over sorted data gives you Venn-style set operations in one linear pass with no extra space. RocksDB’s sorted keys let you do the same across two DBs, in parallel by splitting on key prefix as in Part I. In Part III, we’ll look at how to do MapReduce for out-of-core aggregations.

The Wonders of RocksDB in Rust (Part I - Intro)

2026-02-10T00:00:00+00:00

You can check out the code here.

Squarely two years ago (Jan 2024), I needed to interact with an in-house K-V store at work, and the underlying DB layer was RocksDB. I had been using and tuning RocksDB in C++ land since 2022 to implement the on-disk Faiss Inverted File index to support our deca-B use case, and spent a good part of 2023 digging my heels into the fast-maturing Rust ecosystem. The natural question was - can I do my data engineering with RocksDB in Rust?

Looking back, I got so much out of that setup, and I still frequently reuse and share with people some of the techniques from that time. This enabled my one-big-machine-style data processing pattern, turning Big Data into smoldata, accelerated by flexibly picking from a fleet of EC2 instances. It’s gotten to a point where a mini blog series is warranted.

Historically, if you want to deal with data in a structured way, you have two more or less isolated options. You either go high-level with packaged database services (OLTP/OLAP, SQL/NoSQL), get locked into their provided transformation functions, and pay the interface overhead. One step down, there are embedded OLTP (SQLite) and OLAP (DuckDB), but they are purpose-built for SQL-centric workloads. On the other end, you roll up your sleeves and reinvent a good part of the wheel from on-disk data representation all the way up. RocksDB filled this gap by providing a strong embedded K-V store for other database services to build upon. But much of the development was confined to the database community, rather than the data engineering community. It was very hard to imagine using systems languages like C/C++ to work on data engineering problems, because one bad pointer may corrupt your data silently, and data outlives software!

Rust changed this. It is a high-level, memory-safe language for your data engineering problems, with performance matching other low-level systems languages. What’s been historically overlooked is now possible, and can give one great leverage. You can actually use Rust to get data in and out of RocksDB as stable and long-term stores!

This Part I just scratches the surface with the state of the ecosystem, basic patterns, parallel reads and writes with performance in mind. It doesn’t go deep into topics that showcase the mojo. We save them for the Part II.

RocksDB and its Bindings

At its core, RocksDB is a C++ project. It comes with Java bindings in the original repo, and you can find bindings for other languages. As of this post’s writing, for Go, linxGnu/gorocksdb is the most up-to-date with the underlying C++ code. For Rust, the rust-rocksdb crate from Zaidoon leads the pack for freshness. The original rocksdb crate from Tyler Neely (spacejam) updates less frequently.

Per Rust ecosystem’s convention, when wrapping code from other languages, it’s split into the high-level and low-level counterparts, two separate crates. The high-level crate takes the more regular name, rust-rocksdb. The low-level crate takes a suffix -sys, in this case, rust-librocksdb-sys. Usually, the low-level crate contains the original code so there is no external dependency to chase outside of cargo. This leads to long compilation times, because the original code needs to compile too. Sometimes, it gives you the option to specify an existing binary distribution on your computer.

How do you know what version of RocksDB it’s wrapping? If you navigate to the Dependencies tab of a certain version (e.g. 0.46.0) of rust-rocksdb, you will find the corresponding version of rust-librocksdb-sys with version string 0.42.0+10.10.1. The main versions, 0.46.0 and 0.42.0 for the two crates don’t tell you anything about the wrapped C++ RocksDB’s version. However, this part after the + sign, 10.10.1, does. Then you can go to facebook/rocksdb’s corresponding tag to find the meat. The wrapper, when it does its job well, should faithfully reflect the features of the wrapped code.

Basic Patterns

This is the basic write and point-lookup pattern (full code):

db.put(key.as_bytes(), val.as_bytes())?;

let value = db.get(key.as_bytes())?;
if let Some(value) = value {
    println!("val: {}", std::str::from_utf8(&value)?);
} else {
    println!("key not found");
}

Since the underlying data are sorted in an LSM-Tree, we can iterate consecutively (full code):

let mut db_iter = db.full_iterator(IteratorMode::Start);
while let Some(Ok((key, value))) = db_iter.next() {
    println!(
        "key: {} value: {}",
        String::from_utf8_lossy(&key),
        String::from_utf8_lossy(&value)
    );
}

Parallel Scan

There are a lot of CPU cores on your computer. If you need to scan a substantial portion or the entirety of the database, do this (full code):

let db = open_rocksdb_for_read_only(&args.db_dir, true)?;

let prefixes = generate_consecutive_hex_strings(4);
let count = prefixes
    .into_par_iter()
    .map(|prefix| {
        let prefix = prefix.as_bytes();
        let mut db_iter = db.full_iterator(IteratorMode::From(prefix, Direction::Forward));
        let mut count = 0;
        while let Some(Ok((key, _value))) = db_iter.next() {
            if &key[..prefix.len()] != prefix {
                break;
            }
            count += 1;
        }
        count
    })
    .reduce(|| 0_usize, |acc, c| acc + c);

The example here has more or less evenly distributed hex strings (emulating hashes) as keys, so one could create consecutive short hex strings with a configurable number of digits, e.g. 4 (generate_consecutive_hex_strings(4)). They look like 0000, 0001, …, f000, …, ffff. Your data might use different formats for keys, and they might not be evenly distributed. However, hashing usually redistributes lexically clustered keys evenly. So consider using them as keys and attaching the original keys as part of the values. Good, fast hashing algorithms for this purpose are murmur3, xxhash.

The parallelization is provided by the excellent rayon crate. .into_par_iter() turns a Vec into a parallel iterator that distributes work across a number of threads in Rayon’s thread pool.

You can put a lot more complex logic in the inner loop, and treat each prefix as the unit of work. Depending on your workload, select the number of digits so that each unit of work is small enough to not cause resource starvation (OOM) or run too long, and big enough to avoid scheduling overhead.

Parallel Bulk Load

If you have many writes, create a WriteBatch and write once (full code):

let mut write_batch = WriteBatch::default();

for _ in 0..ENTRIES {
    let key = generate_random_hex_string(RAND_BYTES_LEN);
    let val = generate_random_hex_string(RAND_BYTES_LEN);
    write_batch.put(key.as_bytes(), val.as_bytes());
}

db.write(&write_batch)?;

Often, you have TBs of data to ingest, and you should leverage those CPU cores. So let’s write the write batches in parallel (full code):

let db = open_rocksdb_for_bulk_ingestion(&args.db_dir, Some(7), None)?;

rayon::ThreadPoolBuilder::new()
    .num_threads(NUM_THREADS)
    .build_global()?;

(0..NUM_THREADS).into_par_iter().for_each(|_| {
    let mut write_batch = WriteBatch::default();

    for _ in 0..ENTRIES_PER_THREAD {
        let key = generate_random_hex_string(RAND_BYTES_LEN);
        let val = generate_random_hex_string(RAND_BYTES_LEN);
        write_batch.put(key.as_bytes(), val.as_bytes());
    }

    db.write_without_wal(&write_batch).unwrap();
});

db.flush()?;

There are a few things here. We configure each thread to write all data entries into a write batch before writing the batch into the DB. This is just an example writing random hex strings. In real work, you should choose a small but chunky parallelization granularity to accumulate the write batches.

We also write without recording entries in the Write-Ahead Log (WAL), so they only go into the memtables. We flush the memtables at the very end. This avoids the extra work of writing to the WAL, at the risk of the bulk load job crashing in the middle. Usually, it’s not worth resuming during a bulk load, and you can just redo completely.

When opening the DB, there are many options to control the behavior of the session while the DB is open. Some options affect reads over existing data in sorted string table (SST) files, while some options affect the new files written. RocksDB is known to be production-ready, highly tunable, and highly backward-compatible. Generally, newer versions can open DBs created from older versions. Some older versions can open DBs created from newer versions. To learn more about this and become an expert, it’s important to read through the Wiki, the FAQ, and read the header files options.h (and table.h, advanced_options.h) / Rust bindings docs.

If you want a head start, look at the various open_rocksdb_*() functions in my rocksdb_utils.rs. It gives you a very strong baseline on modern SSDs.

Compaction After Bulk Load

After a bulk load without compaction, you need to manually compact the DB (full code):

let mut compaction_opts = rust_rocksdb::CompactOptions::default();
compaction_opts.set_exclusive_manual_compaction(true);
compaction_opts.set_change_level(true);
compaction_opts.set_target_level(ROCKSDB_NUM_LEVELS - 1);
compaction_opts
    .set_bottommost_level_compaction(rust_rocksdb::BottommostLevelCompaction::ForceOptimized);
db.compact_range_opt(None::<&[u8]>, None::<&[u8]>, &compaction_opts);

This takes a while, and roughly doubles your disk usage. If you use zstd as the bottommost level, it will be less than doubling. It gradually approaches that peak disk usage, and very quickly drops to the final size by deleting old files. If anything fails in this process, such as running out of disk space, abrupt cancellations, or server crashes, the DB is still OK. RocksDB is used as the embedded layer for many production DB services, and is built rock-solid in its default configuration without turning on risky knobs. It has checksums all over the place to detect errors.

Why is RocksDB in Rust a strong setup for data engineering? It’s an embedded, production-grade K-V store without the overhead of a full database service, coupled with a memory-safe language. With the techniques above, you can already move a lot of data in and out of RocksDB. In Part II and Part III, we’ll look at some very cool patterns, including parallel two-pointer set operations, and MapReduce for out-of-core aggregations.

Notes on People Management

2026-02-07T00:00:00+00:00

Why do you need people management?

Let’s think from first principles. Suppose you have a vision, about a certain shape of a product that is wanted by some people. It involves many functions pushing together to make it into a stable business. The two most fundamental ones are, Product/Engineering that puts in the money to build it, and Sales that gets money from selling it - creating a virtuous loop. Raising money from equity or debt gives the former more lead time, but eventually, the product has to sell. As the business grows, you improve the effectiveness of this loop by adding Customer Success that keeps customers engaged, Marketing that reaches more customers, Legal that protects the business, and Finance that does the accounting, budgeting, and making the business more appealing to investors. There are more support functions like HR, Security, Compliance that may get different emphases in different industries.

Someone needs to be directly responsible for all of the functions. This is the role of the CEO. The CEO and the function heads constitute the executive suite - the first team. All members of the exec suite need to think strategically, and be aligned with maximal information bandwidth. Since they more or less self-organize, my notion of people management isn’t primarily focused on this level.

In this post, I would like to focus more on people management inside of each function, with a particular emphasis on Engineering. Typically, the finer the granularity of responsibilities, the less direct the link between individual performance and overall company success. So the performance evaluation gauges, practices, and culture find other anchors to lean on, e.g. engineering principles, rigor, and teamwork.

Let’s go from an Engineering team of one. This usually is frontend, backend, data loading/migrations, infra, CI/CD together. You start with platforms and tools that automate most of them, and instead focus on the core business logic, the special sauce. As business needs become more complex and bigger scale, and reliability requirements become more demanding, it naturally overwhelms the one-person capacity. It’s time to hire out of the situation.

Assuming each person is of similar intellect, you can aim for a best case of Nx factor to the one-person raw productivity, N being the number of team members. This is linear scaling. Of course, if the team is coordinated well, the business impact can be a lot higher than that, because high execution velocity realizes the vision/strategy. But coordination has communication cost, so in reality, scaling people is usually sub-linear in terms of raw productivity. It’s OK to be sub-linear, because again, businesses demand speed of execution, so as long as more people don’t lead to negative productivity, it is worth it for a high-growth company to keep hiring. This is the business/macro-level consideration.

On the individual/micro-level, a manager needs to optimize towards a higher scaling factor, by focusing on prioritization, unblocking, collaboration, and fostering a culture of learning, so team members can skill up. They also need to identify, convince, and hire more talented/experienced candidates onto their team within the budget, to boost up the per-capita productivity. Together, the team’s output is maximized. The management quality is determined by these two axes, in how strongly they can guarantee their team’s high productivity and success within the budget.

A manager needs to develop an active personal loop with every team member, from when they were still a candidate before hire. The idea is, you need to understand their capabilities and drive, so that you can maintain alignment and let them push themselves the most. You need to mature yourself faster than the team, so you could offer guidance, mentorship, and make the right calls. This creates mutual respect and trust, a critical component in leadership dynamics.

When getting new people onboarded, a manager needs to foster opportunities for them to absorb the instutitional knowledge. Talk to them often, and encourage them to sync up with existing team members. A good start forms the baseline expectations. Let high-bandwidth communication be part of that baseline.

A manager needs to make the productivity pressure gently present as a background hum. This is the glue between the airy “team culture” and formal annual evals. In-office work, open floor plan, etc. have this by default - you see how everyone works. Remote work, on the other hand, loses this baseline, and it takes a lot of brushing up to maintain it without getting tedious. It tends to slip into a “nothing happens every day and it’s ok” kind of complacence. You look away, people stop working. Quarterly meetups, proactively showing work progress every day without awkward standups, keeping the communication channels active and by default open can all help. Remote work is great for discovering great talent, but it is inherently disadvantaged in this regard, therefore best only for highly self-motivated people.

Moreover, all managers in adjacent teams have to apply similar productivity pressure, otherwise it creates unfairness and cynicism. If the infra team is driven hard while other user-facing teams are stalling, it makes everyone doubt company priorities. All managers need to feel the pressure in a stronger way than their team members, and drive them towards progress. Otherwise, the productivity pressure effectively does not exist.

There is a very human aspect in this management process. A manager’s effectiveness fundamentally depends on if they can rebuild the team in their vision. A limited version of this is, if they see a deep problem with someone they manage, they need to be able to find someone else who can do a better and adequate job, so they can uphold the productivity guarantee needed by the business. A strong manager can project this position and move the team forward.

You want to avoid the team from getting into a state, where there are too many things going on, and too few hands to get through them. To draw an analogy from computing, too many threads on too few processors cause thrashing. The context switching overhead becomes too big, and you lose the ability to fairly evaluate your team members’ work. Everyone gets a plausible excuse of “too many things”, and very little ends up getting done, and apparently no one has dropped the ball. Without a reliable signal and feedback, the team becomes effectively unmanaged. Sometimes this is inevitable, and you need to deepen the trust between yourself and team members. However, you need to consider scaling the team when there is a chance, and push this signal upward to your leadership.

You also want to avoid having only one person in charge of an entire field of work, without enough trust-building. Absent a convincing backup option to hold the bar, you are at the whim of that one person, for better or worse. It turns an employment dynamic into a partnership dynamic. But the formal power and employment structure don’t reflect that, so you will get struggles and disappointments. This is a common dilemma, and I’m sure it will resonate with a few readers. There is no easy way to get out of this, and the risk is high. The fundamental solution is developing a deeper understanding of the business, the bar (how things ought to be), the industry consensus, your options in the talent market, and the raw human nature. Being able to make a clear case to your team members, your peers, and your leadership allows you to push forward.

To conclude, people management is ultimately about making the business’s productivity guarantee credible. That means hiring, unblocking, and developing people proactively, and being willing to reshape the team when needed. The hardest situations - thrashing or a single “indispensable” owner - force you to rely on a deeper understanding of everything around you, and human nature. That understanding is what lets you move forward when formal authority alone isn’t enough.

Recommended Reading:

The Manager’s Path - Camille Fournier
Staff Engineer - Will Larson
Why I Ignore The Spotlight as a Staff Engineer - Lalit Maganti

Case Study: turbopuffer ANN v3

2026-02-03T00:00:00+00:00

turbopuffer is all the rage in vector dbs! In the past, its sizing guide recommended against using turbopuffer if the direct search corpus exceeded a certain number of vectors, likely 1B. In other words, it was built more for multi-tenanted/multi-namespaced hosting use cases, such as its current customers, Cursor, Notion, where each tenant/namespace (end user) would only search a very small portion, filtered down to their own data. Lately, it published an article, explaining searching against a 100B-vector corpus at a whopping 1k qps throughput. There were a lot of admirable first-principles estimates that guided their design choices.

I would like to write down a few observations, wiggle a few parameters, and do some cost analyses.

Vector search is a punishing data access pattern - everything is larger. One field is as big as DIM x DTYPE_SIZE bytes, instead of a single INT64, and the entirety of the field is needed in the query hot path. This makes everything downstream constrained by the sheer data quantity, and bandwidth-bound. Traditional databases put smallish indexes in-mem, and do most of the DB operations on disk. In contrast, vector indexes are big, and often need to reside in DRAM entirely to work properly due to graph-like access patterns. Unless filtered down to a brute-force scannable segment by some natural predicate, it usually means a lot of bandwidth and a lot of compute.

Multi-tenanted use cases provide a great opportunity for optimization by amortization, as any large infra/platform/insurance providers would do - pooling uncorrelated demands/risks lets you provision for the average rather than the peak. For example, S3 specialized into a low-cost storage paradigm, based on the right observation that most people store a lot of data, but very infrequently query some of them. If they spread out data in many HDDs, one user’s bursty workload can be spread across many of them through prefixes, thus delivering a lot of throughput, and a decent amount of IOPS. Overall the ecosystem benefits from this specialization.

Vector search over a single large corpus is an even more punishing data access pattern - in theory, the entirety of your corpus needs to be considered. You can’t use the above amortization optimization anymore. As a platform provider, turbopuffer entered into a different race with a new customer profile that wants to search through a single namespace without additional filtering. Now they need to pound-for-pound optimize for those specific customers and make sure their demands are expressed with the best tradeoffs. These customers could make different choices because they are the consumer of a database that ideally fully serves their needs, instead of the provider of a database that needs to find common ground from a lot of customers. These customers could ask for a lot of optimizations uniquely afforded by them, given their own insights, and have better definitions of acceptable tradeoffs.

Cost at 1k qps

Let’s calculate how much it costs for the final setup to reach 1k qps. At the beginning of the article:

100 billion vectors, 1024 dimensions per vector, 2 bytes per dimension (f16). This is vector search over 200TiB of dense vector data.

And almost at the end:

The lowest level of the quantized ANN tree is stored in DRAM, as was the case before. However, because these vectors are compressed, they require less memory bandwidth to access.

This requires 100B x 1024/8 bytes = 12.5TiB of DRAM!

Looking around EC2 instance options, getting 13 x i4i.32xlarge instances (1TiB each) can fulfill the DRAM requirement. Each instance comes with a total of 29TiB NVMe drive space with a RAID0 setup as the disk cache. r6id.32xlarge has the same CPU/DRAM specs, but falls short at only 7.4TiB disk space. i7is are very similar to i4is with better CPUs, but come with a premium. They both support AVX-512. So let’s just use the less aggressive estimate.

With on-demand pricing, it amounts to 13 x $10.98/hr = $142.8/hr ($103k/mo)! With a 3-year savings plan, 13 x $5.08/hr = $66.04/hr ($48k/mo). More than half a mil annually.

But Wait…

In their final setup,

See how the DRAM and NVMe SSDs have so much unused capacity. They explained that the scheme became completely compute-bound with much smaller (f16 -> bit, 1/16) vector sizes. That was true, but that was due to a choice that put all 12.5TiB RaBitQ-quantized vectors into DRAM, and left SSDs only responsible for 1% of the leftover reranking work! There has to be a middle ground somewhere…

I believe there was something off about the estimation method. In the scale chart,

the cloud object storage bandwidth was in line with my past experience on AWS. In fact, AWS would typically rate-limit one IP in the same region from downloading more than 50 Gbps from a single bucket (my Rust-based sulfite downloader would initially go over, but got capped quickly). The NVMe SSD bandwidth was estimated close to what a single i4i.32xlarge with 29TiB would deliver, ~15GB/s. To prove it, I started one and tested it:

--> sudo fio --name=seqread --ioengine=libaio --rw=read --bs=16k --size=1G --numjobs=1024 --runtime=30 --time_based --direct=1 --group_reporting --filename=/mnt/raid/test
seqread: (g=0): rw=read, bs=(R) 16.0KiB-16.0KiB, (W) 16.0KiB-16.0KiB, (T) 16.0KiB-16.0KiB, ioengine=libaio, iodepth=1
...
fio-3.28
Starting 1024 processes
Jobs: 1024 (f=1024): [R(1024)][100.0%][r=14.4GiB/s][r=944k IOPS][eta 00m:00s]
seqread: (groupid=0, jobs=1024): err= 0: pid=18324: Wed Feb  4 23:30:40 2026
  read: IOPS=919k, BW=14.0GiB/s (15.1GB/s)(421GiB/30005msec)
    slat (usec): min=2, max=2766, avg= 7.44, stdev= 5.94
    clat (usec): min=86, max=51355, avg=1103.34, stdev=1026.37
     lat (usec): min=108, max=51368, avg=1110.91, stdev=1027.03

Show full fio output

--> sudo fio --name=seqread --ioengine=libaio --rw=read --bs=16k --size=1G --numjobs=1024 --runtime=30 --time_based --direct=1 --group_reporting --filename=/mnt/raid/test
seqread: (g=0): rw=read, bs=(R) 16.0KiB-16.0KiB, (W) 16.0KiB-16.0KiB, (T) 16.0KiB-16.0KiB, ioengine=libaio, iodepth=1
...
fio-3.28
Starting 1024 processes
Jobs: 1024 (f=1024): [R(1024)][100.0%][r=14.4GiB/s][r=944k IOPS][eta 00m:00s]
seqread: (groupid=0, jobs=1024): err= 0: pid=18324: Wed Feb  4 23:30:40 2026
  read: IOPS=919k, BW=14.0GiB/s (15.1GB/s)(421GiB/30005msec)
    slat (usec): min=2, max=2766, avg= 7.44, stdev= 5.94
    clat (usec): min=86, max=51355, avg=1103.34, stdev=1026.37
     lat (usec): min=108, max=51368, avg=1110.91, stdev=1027.03
    clat percentiles (usec):
     |  1.00th=[  204],  5.00th=[  297], 10.00th=[  379], 20.00th=[  490],
     | 30.00th=[  603], 40.00th=[  742], 50.00th=[  881], 60.00th=[ 1037],
     | 70.00th=[ 1270], 80.00th=[ 1565], 90.00th=[ 2073], 95.00th=[ 2507],
     | 99.00th=[ 3556], 99.50th=[ 4293], 99.90th=[14353], 99.95th=[19006],
     | 99.99th=[28705]
   bw (  MiB/s): min= 2841, max=17049, per=100.00%, avg=14370.15, stdev= 2.09, samples=60416
   iops        : min=181572, max=1090402, avg=919380.14, stdev=134.08, samples=60416
  lat (usec)   : 100=0.01%, 250=2.61%, 500=18.59%, 750=19.32%, 1000=17.26%
  lat (msec)   : 2=31.16%, 4=10.44%, 10=0.44%, 20=0.13%, 50=0.04%
  lat (msec)   : 100=0.01%
  cpu          : usr=0.44%, sys=1.20%, ctx=27584225, majf=0, minf=20030
  IO depths    : 1=100.0%, 2=0.0%, 4=0.0%, 8=0.0%, 16=0.0%, 32=0.0%, >=64=0.0%
     submit    : 0=0.0%, 4=100.0%, 8=0.0%, 16=0.0%, 32=0.0%, 64=0.0%, >=64=0.0%
     complete  : 0=0.0%, 4=100.0%, 8=0.0%, 16=0.0%, 32=0.0%, 64=0.0%, >=64=0.0%
     issued rwts: total=27581027,0,0,0 short=0,0,0,0 dropped=0,0,0,0
     latency   : target=0, window=0, percentile=100.00%, depth=1

Run status group 0 (all jobs):
   READ: bw=14.0GiB/s (15.1GB/s), 14.0GiB/s-14.0GiB/s (15.1GB/s-15.1GB/s), io=421GiB (452GB), run=30005-30005msec

Disk stats (read/write):
    md0: ios=27540231/1, merge=0/0, ticks=28944453/0, in_queue=28944453, util=95.65%, aggrios=3447628/0, aggrmerge=0/0, aggrticks=3583264/0, aggrin_queue=3583264, aggrutil=95.15%
  nvme3n1: ios=3462858/0, merge=0/0, ticks=3868512/0, in_queue=3868512, util=95.15%
  nvme6n1: ios=3450333/0, merge=0/0, ticks=4140371/0, in_queue=4140372, util=94.52%
  nvme2n1: ios=3432066/1, merge=0/0, ticks=3686629/0, in_queue=3686630, util=94.43%
  nvme5n1: ios=3453552/0, merge=0/0, ticks=3129646/0, in_queue=3129646, util=94.54%
  nvme8n1: ios=3441541/0, merge=0/0, ticks=3170668/0, in_queue=3170669, util=94.49%
  nvme1n1: ios=3434711/0, merge=0/0, ticks=3979054/0, in_queue=3979055, util=94.45%
  nvme4n1: ios=3458483/0, merge=0/0, ticks=3662792/0, in_queue=3662792, util=94.58%
  nvme7n1: ios=3447483/0, merge=0/0, ticks=3028440/0, in_queue=3028440, util=94.50%

However, with 13 nodes, they cumulatively should stack up to 195GB/s! A balanced SPTAG index should pretty evenly distribute its data to all the SSDs. Now the SSDs look wayyy more over-provisioned. I wouldn’t be surprised if the DRAM and CPU cache numbers needed to be adjusted due to multiplication - dependent upon how their implementation branches off the non-root level centroid vectors.

With all this extra capacity from SSDs, why don’t we just put the RaBitQ-quantized vectors back onto them? A total of 12.5TiB is just chump change compared to the original 200TiB of f16 vectors. With this new estimate, we can use fewer instances. Let’s stick with i4i.32xlarge, and use SSDs instead of DRAM as the constraint, (200TiB + 12.5TiB) / 29TiB = 7.3 -> 8 instances. This just brought down the cost by ~40%. Using the same bandwidth limit estimate for SSDs, we got 8 x 15GB/s = 120GB/s, 120GB/(1x6MB) = 20,000 qps. This is more than enough, still compute-bound.

It even makes their initial setup without RaBitQ quantization work:

At fp16, 120GB/(1x100MB) = 1,200 qps, just enough to match the target. However, fp16 usually is a lot more compute-intensive than binary calculations. If they got compute-bound 1k qps doing RaBitQ, compute-bound fp16 is going to be a lot slower.

Of course, we miss out on the extra CPUs from the other 5 nodes, and DRAM gets woefully under-utilized. i-family instances have r-family level CPU/DRAM ratio, and over-provision on DRAM. Is there a more CPU heavy instance type with a lot of disks? c6id.32xlarge comes to mind with 128 vCPUs (same as i4i.32xlarge) and 256GB DRAM, instead of 1024GB. But like r6id.32xlarge, they only carry 7.4TiB disk space. (200TiB + 12.5TiB) / 7.4 = 28.7 -> 29 instances! This amounts to a 3-yr savings plan 29 x $2.696/hr = $78/hr, $78 / $66 = 1.2x the original budget. However, it is a net more balanced configuration, because you now get to have 29 / 13 = 2.23x the original CPUs to solve your compute-bound problem, so the overall qps should multiply as much. This might be what some bigger customers need.

Re-examining the Original Vectors

What’s up with fp16 for original vectors? The industry finally accepted fp32 was too much for vector data, and fp16 wouldn’t cause any significant accuracy loss. Most models are even trained with fp16 autocast on most layers and ops, and the fp16 final embeddings are more than fine for L2 distance/cosine similarity calculations. But what about INT8? In practice, it leads to less than 0.01 accuracy loss. For a well-tuned embedding model (texts or images), the similarity range is wide and clear, <0.4 for dissimilar, 0.4~0.6 for some relevance, and >0.6 for similar. You can totally use INT8 quantization. Best, to preserve more accuracy, use non-uniform scalar quantization so that each dimension gets its own dynamic range trained by representative data. This cuts down your 200TiB original vectors by half! So (100TiB + 12.5TiB) / 7.4 = 15.2 -> 16 c6id.32xlarge instances. 16 x $2.696/hr = $43/hr, $43 / $66 = 65% the original budget, and 16 / 13 = 1.23x the original CPUs to run RaBitQ. INT8 also runs faster than fp16 for distance/similarity calculations.

If you go down further to non-uniform INT4, in my setup, the accuracy loss is at max 0.03 by comparing the unquantized query vector to quantized corpus vectors, and 0.06 by comparing two quantized vectors. It might be unacceptable for some, but a great budget option for others. Accepting the loss, INT4 is even faster than INT8. INT6 is a weird one. It’s usually a better trade-off midpoint between INT8 and INT4, but it’s not byte-aligned, so usually slower than INT8. However (many turns in this line of reasoning), if you use RaBitQ, and only use 1% of original vectors to rerank, speed is not an issue here, so just pick your acceptable accuracy-compression tradeoff point. If you do go down to INT4, or somehow find an INT2, there might be no need to keep the RaBitQ quantizer.

Also, maybe your embedding model could be retrained or fine-tuned to use a smaller output dimension with negligible accuracy loss. This is where it helps to optimize the entire ML<>Data as a system, because a decision during training that helped hit leaderboards might come back to haunt you if that means doubling a huge amount of infra cost to store. Study and document it, and you might just get another 1/2 factor! Vector search is full of these 1/2 (or 2x) factors, and when you stack them up, suddenly you get an order of magnitude of differences!

The 100 qps Target

What if I don’t need 1k qps? Let’s re-evaluate the 100 qps target given what we just went through. At this scale, the extra CPUs are extra cost to avoid, so we need to choose a high SSD/CPU ratio instance type. i3en.24xlarge with 96 vCPUs, 768GB DRAM, and 58.5TiB SSDs is great for this. (100TiB + 12.5TiB) / 58.5TiB = 1.9 -> 2 instances. 2 x $4.72/hr = $9.44/hr, $9.44 / $66 = 14% the original budget, and 96 * 2 / 13 / 128 = 12% the original CPUs. Sounds about right.

The Budget Option for 100B

What if you only need 1~10 qps? This is the true budget option. At this scale, the extra speed NVMe SSDs offer is even extra cost to avoid. Should you appreciate the wonder of gp3 EBS drives? Take an r6in.4xlarge (n for high-throughput EBS access), and pair with 2 x 64TB drives, and pay for some extra IOPS and throughput (not counted): $0.565/hr + $0.08/GB/mo * 128000GB / (30 * 24 hrs/mo) = $14.8/hr. That’s already more expensive than the 100 qps target.

In their original blog article, turbopuffer said they built an LSM-tree directly on object storage. That might bring the instance cost down to a fraction of what you pay for object storage, which is roughly (100TiB + 12.5TiB) x $20/mo/TiB / (30 x 24 hrs/mo) = $3.1/hr on AWS S3.

Conclusion

You should be able to get more bang for your buck, by putting quantized vectors onto SSDs, choosing more compute-optimized instance types, and re-examining the original vector precision tradeoffs. Moreover, scaling down from the mighty 1k qps is graceful and elastic. turbopuffer is really cool.

Individual Drive and Organizational Layout in the Age of AI Agents

2026-02-01T00:00:00+00:00

Many people are talking about the power of newer and newer AI agents, and what’s SOTA today is expiring tomorrow. However, many have also observed that AI in automating knowledge work has crossed some sort of boundary of late e.g. Karpathy, Cherny (creator of Claude Code), like the uncanny valley has been crossed. Maybe it was just Opus 4.5, which materially improved beyond the last frontier. Maybe all of those models have improved due to post-training agentic calibration towards practical tasks. From a technological progress standpoint, this may just be one small step towards the eventuality that everyone has been working towards for years. But it signals some kind of maturation, so it could finally be highly depended upon en masse. Most individuals, and most companies, even the ones actively adopting AI into their people’s daily work lives, have not fully embraced or benefited from this new paradigm. I’m trying to draw some metaphysical points from firsthand.

A brief history. My work has been mostly in the backend, ML, and data. Back in 2022/2023 when Copilot and Cursor were released to the public, people were transitioning from more intelligent auto-complete in control to getting a lot of shock and awe from AI coding out entire UIs. I started learning Rust in early 2023, but AI wasn’t very helpful, given its limited exposure to the language on the Internet and limited reasoning abilities. My first big project in Rust was at the end of 2023, and it lasted a couple of months. I remember when approaching mid-year 2024 that AI got a lot more helpful. My two primary modes of AI usage were still quite limited then - I would have Q&As on questions that could help my own learning in Claude.ai, or I would use the increasingly more powerful auto-complete in the IDE. I only started using Cursor at the beginning of 2025 for its more powerful (and trigger-happy) auto-complete, and still did not appreciate its agent mode. Coding agents moved my code and files too much without precision, and the unwanted creativity wasn’t worth giving up control for. In other words, I was a Big Tabber.

Somewhere in 2025, my product-focused colleagues were incorporating AI agents more into their workflows in experimental ways, creating guiding prompts and plans, and sometimes slowing down noticeably because they ended up waiting on AI to delegate their tasks to, and each iteration took a lot of “thinking”. It was also a time when people started to complain that vibe-coded code was getting a little overwhelming at MR/PR review time¹.

The scale really got tipped at the end of 2025. I noticed myself using AI agents in Cursor a lot more, and started using Claude Code, because

they were able to make more pinpointed changes, keeping me in control, and
when they made bigger and cross-file changes, they tended to be more often correct, thorough, and insightful.

The boundary of trust has been crossed. I would ask an agent to do a few things:

generate me some NEON SIMD code to match the AVX2 and AVX-512 counterpart and reason through the lines (and let Claude blind-check then cross-check it).
find me all references of some Terraform minSize in different formats across the repo that were hard to replace-all.
recreate some commits preserving timestamps and authors by digging into the reflogs.
quickly examine all the Alembic heads that got migrated in a botched migration, look at the migration files, figure out the damage, and decide whether to rollback or go forward.

Once they showed that they were working through the problems and cross-checking, I knew I needed to do less of it. And man, they were tenacious (programmed to be, but uncannily alive).

The intuition from using AI in their knowledge work for some time is, if you give it the right context (enough, but not too distracting), the right directive, in a finite form with the possibility to iterate in a close feedback loop, AI would do a great job. AI is both a much better generalist (we knew that for a while now) than people, and a much better specialist (more recent refinements) in highly symbolic and localized settings (like juggling with each hardware vendor’s SIMD functions). Just in terms of automating sophisticated knowledge work, we blasted through the stochastic parrot era, and gave LLMs expressions and tools for thinking and reasoning, unique to their nature. World models are still not here, but in the semantic space of human & machine languages, LLMs flow like fish in the water. This was one of the key insights that allowed NLP to accelerate so much faster than computer vision - compared to the physical world with strict laws and coherence rules (checked by the universe in real time), the world of languages is a flatter projection. Humans are not even that good at them! We make mistakes talking and writing, and certainly need to spend time with programming to get logical flows internalized. In the evolutionary chain of events, mastering languages was the defining character of humans. However, we can only take pride in that when compared to animals, not this new AI machinery, because we made them, and distilled the refined versions of ourselves into them. So now they can write, talk, and do a lot more complicated things following the explicit and latent flow of languages.

Think about this, and think about what individuals now can use AI agents for. In this article, the author documented the lengths agents would go to accomplish the directive. Excerpt:

So Claude Code got creative. It downloaded a Python audio analysis package and essentially listened to its own output, dissecting the waveform.

So much exploratory work dealing with unknowns and unknown unknowns is compressed and accelerated into a session with AI & agents²! There is no excuse anymore for hand-waving “this is an unknown risk” (thus taking more time to figure out). You just didn’t get into a chat with your AI friend, or an exploratory session with agents. The bottleneck is on YOU. Specifically, it is on your driving capacity, and mental plasticity. Just like modern cloud machine vertical scaling and memory-safe languages are accelerating data engineering within an individual’s reach, AI is accelerating all knowledge work.

The best kind of person this benefits is someone who is driven, deliberate, and introspective with the bigger picture. The drive now can be amplified to an unlimited ceiling. The deliberation means they have specific goals with nuanced needs, that require meticulousness and precision. The introspection with the bigger picture means they keep organizing and creating enough empty space in their mind to accommodate new knowledge as AI shoves them in, and convert it to make it their cohesive own. AI encourages, enables, and fosters the entrepreneurial side of people.

Now let’s consider another angle. You have someone new to your company, and want to assign them work, either as their manager, or as their senior mentor/peer. What’s the best way to do it? To start, they need a lot of context. They would spend days and weeks to get onboarded and gain deeper familiarity with your company’s stack, absorb that institutional knowledge, and start contributing to a small thing, then a medium thing, then hopefully bigger things. They would spend a good few hours just taking in the density of one interesting tutorial, RFC, or blog post. They get tired quickly after just a few hours, and have to manage the stress talking to others, asking questions trying not to sound like an idiot, or bringing up proposals trying not to get someone mad. You would need to sit in with them, or better, talk through lunch and dinner many times to point and advise so they can do as good a job, if not better than yourself. This is of course all enjoyable if you enjoy it, but tedious if you do too much just to scale the company ASAP.

Longer term, you would need to create an org structure that reflects the boundary constraints of human attention and communication limitations. Your company might just look like a big Hadoop cluster, with being bandwidth-bound as the usual struggle case. It is helpful on the margin if you really need to get +1 thing done, but scaling is extremely sublinear, and at some point, offers negative benefits. Scalability! But at what COST?

Vertical scaling solves this. If you could consolidate more resources per individual, boost up their capabilities, communication overhead would decrease accordingly. It improves locality, and avoids passing big contexts around. The ZIRP era is closing out, and even if it makes a comeback, the fundamental unit of productivity is drifting towards higher capabilities of individuals, instead of merely more of them.

The best kind of organization this benefits is one that identifies, mentors, and empowers individuals to take whole-cloth scopes, be it functional or divisional. These will be inert from AI automation until robots become first-class citizens, because fundamentally, you want to find people that take care of things, bear the responsibilities, so you don’t have to worry about them. Functional roles like engineering, product, sales, marketing need at least one person to drive them for the long haul, do what the AI can’t and get the job done. At a human coordination level, there is additional context hard for AI agents to pull in. Big divisions target categorical customer segments, also for a stable amount of time with domain-specific requirements, which also have additional context. In other words, cut the org boundaries at how much of the physical world the tasks demand, appoint people, and enable them to scale AI agents to fit the newly broadened individual capability boundaries.

You should totally put Claude Code into your MR/PR review process. It sits in the right context (the diff, description & comments), and helps with reviewer’s understanding. Done well, it’s not just a cynical game of AI reviewing AI. ↩
There is an aside on agent plan of action’s serial nature, clashing with GPU chips’ high throughput but high latency nature. This led to NVIDIA’s choice on licensing Groq. There is also the “acquiring the competitor” angle, as has been lately revealed. ↩

What do Streaming Databases Solve Uniquely?

2026-01-30T00:00:00+00:00

This post started from researching recent stream processing solutions like RisingWave and Materialize. There are some meta-points worth bringing up, so I’m writing them down.

Why do we need data stream processing in the first place? That’s easy - we need to get a hold of fresh data that keep coming into the system. Fresh data are naturally more interesting to us in many aspects of life. “Fresh” is a moving target in terms of absolute timestamps. However, if we anchor it against “now”, and express it in “past x hours/days”, it is a stable target, only with changing content.

This conceptually makes sense, but why do I need a dedicated stream processing solution? I can simply do on-demand batch queries against my OLAP data warehouse or data lake. They have indexes built on timestamps and other fields, right? This certainly works, and if I don’t do too many such queries, it is a viable approach. After all, the type of data that OLAPs would host tend to be billions of rows, tens of TBs, and keeping them as cold (different S3 tiers) as you can tolerate it is a first-order cost concern. Only start building more indexes and caching data to hotter media like SSDs and RAM if you need more frequent and flexible access. So if you keep making a certain kind of queries, OLAPs may rely on some cache and eviction rules to prioritize them. It’s the system dynamically responding to your user intent.

However, if one already knows the fresh part of data is more important to them, they can make this intent explicit, and build solutions that target it. All the optimizations and tradeoffs can orient themselves around this stable intent as a set of invariants. Solutions and products in the software ecosystem (and arguably the society at large) follow this rule closely to specialize, from compilers optimizing code, to OLTP vs OLAP, and stream processing is no exception. These invariants serve as anchors and interfaces for people to collaborate towards, and every participant can do a better job with clear specs. From a system perspective, the better we can articulate the invariants as boundaries of something that we want, the better we can fill in the middle to make it fit perfectly. Every once in a while, the messy amalgamation of disjoint components get “reaped” into just a few elegant integrated systems that reflect the user intents more directly. That’s a win for the system and a natural evolutionary journey.

Stream processing solutions often materialize the “past x hours/days” queries into views, which formalize the intent that recent data has different value and access patterns than historical data. By explicitly materializing fresh data in a fast-access layer, we can

Keep hot data in memory or SSD-backed stores
Push older data into columnar formats optimized for compression and batch queries
Transition data through tiers as it ages (memory -> SSD -> object storage)

Further, when someone creates a materialized view for “last 24 hours of user activity,” they’re telling the system: “I will query this repeatedly, optimize accordingly.” The system can

Allocate resources proportional to declared importance
Skip materialization of unspecified aggregations entirely
Make different storage/indexing decisions based on the view definition

This allows the additional work to map perfectly to user intents. Contrast this with a pure OLAP store relying on query-driven caching - it’s reactive, learns slowly, and wastes resources caching one-off exploratory queries.

Now, what are streaming databases and what do they have on top of stream processing frameworks? RisingWave itself has a comparison with Flink that demonstrated its design principles well - the database is the sink where these fresh data materialized views go into, and it’s often beneficial to let the very system that writes to it manage it. The invariant from user intent is that they always need a database to host the data, so might as well provide a good one. The additional invariant from the data system is what writes to the database knows the most about the database, and can do further optimizations upholding its promises with a better context.

There is another fine point to the choice of SQL as the user-interface language for both RisingWave and Materialize. Some past stream processing frameworks would use procedural interfaces, DSLs. However, the industry consensus has nominated SQL as the declarative and expressive language to work with tabular data. This is an invariant! Yet again it shows how systems and market players organize themselves around invariants.

Some AWS EC2 Instance Choices

2026-01-23T00:00:00+00:00

There are many aspects of AWS EC2 instances. AWS offers a whole fleet of different varieties, targeting a wide range of use cases. It’s a big topic, and one should probably have a conversation with their AI friend and learn as they experiment with a few different scenarios if they want to get a better grounded sense.

Over the years, I have used a few families of instances at length. So here are some of my personal notes.

The `t` Family

They are burstable instances. The smaller ones are often offered for free as the starter plan for your EC2 journey. You can host bursty web servers on them. I don’t use them as much. They are about 14% cheaper compared to same-gen/spec m instances. e.g. t3.large vs m5.large (both Intel Skylake), and t4g.large vs m6g.large (both AWS Graviton2). The generation numbers are simply zigzagged.

The `c`/`m`/`r` Families

They are the workhorse instances, with CPU/RAM ratio of 1:2, 1:4, and 1:8. You’d usually use c when you have compute-heavy workloads, with native languages C/C++/Rust. You’d use m when you use interpreted languages such as Python for web servers, and want to be on the safe side. You’d use r when you want to host a database that requires a good amount of caching, or in-memory data structures, like vector DBs.

It’s important to think about CPU and RAM differently as resources. CPU cores accelerate things, and with more, you get to do things faster, but they are more or less elastic. On the other hand, RAM is usually a rigid requirement - your application might just crash due to OOM if there is not enough memory. Swapping on fast local NVMe drives could alleviate some hot spots occasionally, but the data structures that work well on RAM usually don’t work that well on drives, even NVMe drives, due to random access patterns (it’s what RAM is literally named after). Because of this, when you host databases on r instances, and your traffic is low, you might observe the CPUs are just idling at 5% - kind of a waste.

You could consider the lower-power processor options, like choose r6a (AMD Zen 3) at a 10% discount, or r6g (AWS Graviton2) at a 20% discount compared to same-gen r6i (Intel Ice Lake). Their lower prices reflect the lower performance pretty accurately, so you get what you paid for. Note that *6a instances have AMD Zen 3 chips that don’t come with later SIMD instruction sets (notably AVX-512) but just AVX2. In comparison, *6i instances have Intel Ice Lake chips that do have AVX-512 for your hashing/vector calculations. You do have to compile your code specifically to support these SIMD instructions, or trust some pre-compiled binaries. I would ask you to migrate away from *5 instances, because they cost exactly the same as *6i instances and have Intel Skylake chips that also don’t have AVX-512, only AVX2.

Arm chips require your code/containers to be compiled for them and able to run on them. Nowadays it’s increasingly not a problem in containerized environments. In terms of SIMD support, AWS Graviton2 (ARMv8.2-A) and beyond chips all have NEON instruction sets (typically 128-bit width), but they are a lot weaker than AVX2 (256-bit width)/AVX-512 (512-bit width). Later Arm chips like AWS Graviton3 (ARMv8.4-A) and Graviton4 (ARMv9-A) have wider bit width SIMD instructions (256-bit SVE and 512-bit SVE2 respectively). These are actively becoming more competitive, but the code that uses these advanced instruction sets is only getting started.

Newer Generations

The storytelling and price performance shift quite significantly at *7* and *8* generations in recent years. This is where AMD and AWS’s own Gravitons really took the lead. For AMD, *7a (Zen 4) and *8a (Zen 5) added AVX-512, and (drum roll…) removed Hyper Threading (HT). So each vCPU is pound per pound a real core, instead of *5- and *6*-era Intel/AMD HT. This resolved a minor but annoying quirk when you watched your compute-heavy workloads starting to thrash as CPU utilization reached just 50% with HT… No more! As a result, the nominal sticker prices for AMD chips are ~15% higher than same-gen Intel chips, e.g. c7a.large vs c7i.large. But you get double the physical cores! This is actually insane! *7* Graviton3 chips narrowed the price difference (18%) to same-gen Intel chips, and *8* Graviton4 chips did that further (15%). I personally ran some workloads on big c8a.48xlarge and c8g.48xlarge boxes. I gotta tell ya, they are nice. Note that, because of HT, Intel could put out a c8i.96xlarge, but at high compute load they are effectively 192 instead of 384 physical cores.

Variations

There are these suffixes to the instance names, like d for extra and smallish ephemeral NVMe drives attached if your workloads need them, and n for much higher networking allowance. They could come in combos, e.g. m6i, m6id, m6in, m6idn. However, with the *7* and *8* generations, AWS moved away from Intel chips to prefer their own Gravitons to be the bearer of these varieties. They are also not evenly distributed in the c/m/r spectrum. In *7g, c7g has d and n, but m7g/r7g only have d. *8g has a much fuller line-up, and introduced another variation b for higher EBS volume allowance. I would like to note that n (and aforementioned b) gives you much better EBS volume communication capacity (IOPS and throughput), because EBS volumes are network-connected drives, so they follow similar rules. You need to be aware of this when your database needs larger than usual communication bandwidth than what regular instances have to offer. So upgrade with n and b accordingly.

The `x` Family

Historically you don’t see these often. They are at the right side of the r family with even bigger CPU:RAM ratios 1:16 and 1:32. There were x1*/x2*, and they jumped to x8 to align the generation numbers. Here is a list.

The `i` Family

They come with big numbers of ephemeral NVMe drives (“storage optimized”) than d-suffixed instances. They are very helpful if you have a 50TB database shard to operate on. i4i/i3en used to be the Intel choices, and i4g/im4gn/is4gn used to be the Graviton2 choices. And then things moved strategically, jumping generations. It’s a little inconsistent to explain, so here I link a list down here.

Notably, with these latest i7ie.48xlarge and i8ge.48xlarge instances, you get a whopping 120TB worth of NVMe drives! You can RAID0 them to do large-scale data processing within a single instance. Your Big Data problems suddenly became smoldata problems that just need to be paired with a fast programming language. I wrote a big post about it - Big Data Engineering in the 2020s - One Big Machine. I hope you enjoy that read too.

Another point to make is, compare r7i.24xlarge ($6.35/hr on demand), i7i.24xlarge ($9.061/hr with 22.5TB NVMe), and i7ie.24xlarge ($12.475/hr with 60TB NVMe). They all have the same basic CPU/RAM specs. If you take similar EBS gp3 drives and amp up the IOPS and throughput to max (80k IOPS and 2GB/s throughput since Sep 2025), 22.5TB is around $2000/mo ($2.8/hr), and 60TB is around $5000/mo ($7/hr). Of course, EBS gp3 volumes are way slower than local NVMe drives, so they are not really a fair comparison, but often a practical one in terms of deployment. And see, the extra disk cost almost add up exactly from the r7i.24xlarge base! I thought it’s an interesting fact to share.

The thing about these instances and local ephemeral NVMe drives, is that they get wiped and reprovisioned to someone else once you power down the machine. They are fine if you restart them though. If you intend to run actual monstrous database workloads on them (like with ScyllaDB), you need to stay on top of data backup & restore, and have a run book, because they are less on durable media than EBS drives such as gp3 and io2.

That said, the other thing about these instances and local ephemeral NVMe drives… is that they are monstrously fast:

People that get into databases, ML/AI, and systems programming on AWS often get tripped by the slow disks and get throttled hard. I hope you try these instances.

`g4dn`/`g5`/`g6`/`g6f`

They are the family of NVIDIA inference-oriented cards. g4dn has NVIDIA T4, g5 has A10g, and g6/g6f has L4. Interestingly, g6 (L4) is less powerful than g5 (A10g), and priced lower. g6f is a new savers option featuring fractionals of a L4 GPU, benefiting from NVIDIA’s GPU partitioning.

`p3` (going away)/`p4d`/`p5`/`p6`

They are the family of NVIDIA training-oriented cards. p3 had NVIDIA V100, p4d has A100, p5 has H100, and p6 has B200/B300/GB200. If you are not a large company and have signed a big deal with them, or 3y savings plans, don’t use them. Go somewhere else. This is a hot topic, and deserves its own post. I gotta say for a lot of non-LLM workloads, there are A100s around in various market segments, and the new RTX Pro 6000 Blackwell cards are awesome. This is a fast-moving and fast-expiring situation, so we can look back in a year.

`g6e`/`g7e`

These are in fact in the middle between training and inference, targeting GPU (virtual) workstations. g6e has L40S, and g7e in fact has RTX Pro 6000 Blackwell. I didn’t know about them until writing this post.

There are other families, like u for extreme CPU/RAM ratio, hpc* for HPC parallel scientific workloads, mac* for MacOS CI/CD purposes, inf* and trn* for AWS’s own Inferentia and Trainium chips, and more. I have not directly used them yet. Although, the dl1.24xlarge (Intel’s Habana Gaudi training chip) on AWS Marketpace paved my way to directly contact the Habana team to use their Gaudi2 chips for a few months, and got great training results after jumping through a few hurdles.

EBS News

One last thing of note - back in Sep 2025, AWS announced their gp3 EBS volumes got a lot better. It was fantastic news to me for some workloads that I had to run on expensive io2 drives. Buying IOPS for io2 is extremely expensive, but very cheap for gp3.

Just going to end it here…

People are Computers

2025-12-24T00:00:00+00:00

There are a lot of analogies between people in a team setting, and computers. Calling out some of these analogies helps tremendously when it comes to scaling a team.

raw intelligence <> chip processor clock frequency (and turbo when drinking coffee)
learning-based skill-up <> advanced instruction set / specialization
unblocking each other <> embarrassingly parallel patterns / minimizing work in the critical region / indirection / atomicity
work prioritization <> priority queue / nice values
work validation and feedback loops <> compile-time type checking / schema validation / fail-fast design
work planning around existing interfaces/checkpoints & migration <> compiler optimization upholding invariants / pipeline parallelism

WIP

What Makes a Good Technical Leader?

2025-12-15T00:00:00+00:00

Jotting down the notes.

Technical leadership - databases, distributed systems, computing fundamentals, Cloud security, scaling and tradeoffs, build vs buy, etc.
Product sense & strategy - how to derive intuition from first-principles thinking & market research, and how to validate ideas before committing the team
Leadership style - org design, leveling, talent retention and hiring, how to drive the team, how to evaluate them and apply productivity pressure
Technical communication - story-telling, simplification of systems, use of analogies
Living the values & principles to build rapport and trust with their first team - increasingly more important as the business grows bigger

What's a "Senior Engineer"?

2025-10-28T00:00:00+00:00

What’s a “Senior Engineer”? A senior engineer

can independently take apart an architectural design spec, and independently drive the parts to completion.
clarify logical boundaries, inputs/outputs, and vocalize when unclear or blocked, so that project meets minimal blockage during collaboration.
anticipate, identify bottlenecks during implementation, and find the appropriate algorithmic solutions, to meet latency and throughput requirements.
has enough fundamental working knowledge and experience, and is able to dive into specifics of core technologies in their stack given a reasonable amount of time, and surgically solve the issue according to well established patterns without them reappearing as compounding complexities.

In some ways, it’s easier to check/evaluate a senior person, because there are fewer convergence points that practically every senior person will encounter in their professional lives. For example, a senior systems engineer has to be able to explain and implement a simple hashmap, and explain the various concepts and behaviors. There is no way around it, because without that knowledge, they can’t solve a class of problems.

The other side of the coin is, it’s harder to evaluate a pre-senior person, because it comes a lot down to individual drive and passion. They shine on a team when they are motivated and can do many things to push the projects forward. It’s hard to find those knowledge convergence points to test them during the interviewing process.

Here are some fundamental working knowledge and/or core technologies in different roles.

Senior Full Stack Engineer

JavaScript fundamentals including event loop behavior, closures, and how asynchronous patterns (promises, async/await) affect application performance and user experience
Proficiency with modern framework ecosystems and state management, understanding when client-side vs server-side rendering is appropriate for performance
Database session/transaction management, e.g. debugging Sqlalchemy session issues, the difference between flush and commit, and how to manage the connection pool
Database query optimization and indexing strategies, ORM relationship mapping loading (Sqlalchemy join/select/selectin)
RESTful API design principles and practical experience with authentication flows (OAuth, JWT) including token refresh and session security
Caching strategies across the full stack - browser caching, CDN usage, Redis/Memcached patterns, and when to cache at different layers
Production debugging skills using browser dev tools, server logs, and database query analysis to identify bottlenecks across the request lifecycle

Senior Backend/Platform Engineer

Common concurrency/async patterns, such as green threads, asyncio with async/await, what could block the executor, and when to choose each pattern based on I/O vs CPU-bound workloads
Common parallelization techniques, such as Python thread/process pool/executors, and when process pools are necessary vs when threading suffices
Python GIL, and how it exemplifies itself in various applications, web servers, job workers, how to architect applications to work around GIL limitations in web servers and job workers
When to choose specific data formats (MessagePack for speed, Protobuf for schema evolution, Arrow for analytics, Parquet for storage) based on use case requirements
How to saturate hardware resources and identify bottlenecks, such as CPU cores, disk IO, networking IO, and perform reliable benchmarks for client/server setups
Load balancing strategies and when to scale horizontally vs vertically based on bottleneck analysis

Senior Data & Systems Engineer

Proficiency with systems languages (C/C++, Rust), and computing fundamentals. Knowledge of when to choose systems languages (C/C++, Rust) vs higher-level languages based on performance requirements and development velocity tradeoffs
Strong experience with high-performance databases, e.g. RocksDB, LMDB, ScyllaDB, and mastery of common query patterns, such as point lookup and parallel iterator scans
Strong familiarity with common data formats, e.g. MessagePack, Protobuf, Arrow, Parquet, etc., and when to choose one based on use case requirements
Job queue architecture decisions: when to use Redis vs MQ vs Kafka based on nature of work, durability, throughput, and failure handling requirements
Familiarity with the Python ecosystem: when to use NumPy, and when to use Pandas, Polars, etc.

Senior ML Engineer (Computer Vision)

Model architecture selection: when to use CNNs vs Vision Transformers vs hybrid approaches based on dataset size, computational constraints, and accuracy requirements
Data preparation skills: how to store data at different stages of the ETL process, and how to use common approaches and libraries to dedupe & merge
Training pipeline optimization: when to use mixed precision, gradient accumulation, and distributed training strategies based on hardware limitations and model size
Deployment optimization: when to use model quantization, pruning, or knowledge distillation vs when to scale infrastructure, based on latency and throughput constraints
Evaluation methodology: how to design proper train/validation/test splits for computer vision tasks and when standard metrics are insufficient for production requirements
Production monitoring: how to detect model drift, data quality issues, and performance degradation in deployed CV systems, and when to trigger model retraining
Hardware utilization: how to profile GPU memory usage, optimize batch sizes, and identify bottlenecks in data loading vs compute vs memory transfer

Let's Work Better

2025-06-23T00:00:00+00:00

There aren’t many ways to become better at work as a team. Below is a retrace from personal experience working at different companies and roles, and taking ownership of different projects. I’m writing these down to lay the backdrop as to what a high mark looks like.

Prioritization. Some tasks are more time-sensitive. Some become dependencies of bigger plans. This applies to balancing your own projects & tasks, as well as coordinating with your manager & other teams on the same project. A great outcome is that the most important projects move forward at full speed according to the team plan.
Work through and around blockage by external factors & other team members. Everyone is bound to be blocked by something or someone at times. Don’t just passively wait. Either proactively engage with someone who is blocking you, or vocalize to your manager who can help you get unstuck. Communicate immediately when you realize what you want to do cannot make progress unless something else happens. Your manager can only effectively coordinate something when the thing is communicated. If you have to wait on someone, communicate, and shift gears onto something else in the meantime, be it a backup project, or something exploratory, or technical learning that furthers your skills. Having something as a backup to work on deepens your technical expertise, broadens your scope, and creates more ownership and opportunities to grow. It marks you as exemplary, and the effort will not get overlooked.
Make progress on mastering domain-critical knowledge and skills, and use them at work. It creates a culture of learning, spanning a longer time window, and is the only effective way to work smarter and better. There are some subjects in everyone’s role that once mastered, make the implementations simpler, more scalable, and less maintenance-prone. Identify them and be good at them. Find tasks that allow you to flex those skills on them. This will not go unnoticed.

Productivity Pressure in Management

2025-05-26T00:00:00+00:00

Note: this doc’s audience is leadership & managers. Managers have an elevated role because they sit at the neck of operations, and if they are not on the same page with leadership on defining the team culture and how it exemplifies itself, the team culture effectively does not exist for people on the ground.

People with the trite gripe of management usually take the frame of “true productivity” vs “hours of work”. But this framing is misguided. Given a comfortable work schedule, putting in the necessary hours is a necessary but not sufficient condition for true productivity. It’s one thing to claim staying the hours does not mean being productive, but it’s quite another thing to claim not having enough time somehow makes one more productive.

Look. John Carmack agrees!

The current standard of work schedule is quite comfortable. If one really does the math, counting the wall clock 9-5, 8 hours, subtracting 1 hour of lunch time, and 1 hour of ramping up & dilly dallying, one should aim for 6 core hours of productivity on a work day. It’s a worthwhile internal goal, and it’s OK to not be like that every day. If people have off days, personal obligations, appointments, they should put an OUT code on the calendar. They don’t have to make up for them if employment policy does not warrant that, but it’s far better to have visibility across the team, because it provides clarity and improves coordination and morale. You don’t want a 996 culture, but you don’t want to encourage eating croissants and watching videos all day. You want to take the frame of “we have big responsibilities to this world and ourselves - our work saves lives, and is highly technical and cool”. People can take these skills anywhere, so on a fundamental level, they get it.

Moreover, people should be able to accomplish one or two meaningful things per day. They should be able to articulate why they are meaningful, and describe the challenges. It closes the goal-setting loop.

This productivity pressure has to be gently present throughout the management chain as a background hum. The executives need to feel the strongest pressure, to a personal level, and drive their teams towards progress. They need to get their lieutenants on board with this line of thinking and evaluate how faithfully they are carrying it out. They need to have a center and a good sense of what timely progress feels like, and optimize towards the objectives. They also need to always look for personnel alternatives and backups, so that they can realize their missions. It hurts the most when people with bigger organizational responsibilities are lax. In the absence of this pressure and drive, organizations will not scale, sometimes even negatively.

Migrating 50TB Data From a Home-Grown Database to ScyllaDB, Fast

2025-03-04T00:00:00+00:00

Terence shares how Clearview AI’s infra needs evolved and why they chose ScyllaDB after first-principles research. From fast ingestion to production queries, the talk explores their journey with embedded DB readers and the ScyllaDB Rust driver, plus config tips for bulk ingestion and achieving data parity.

Original link.

Slides:

The Rust Invasion: New Possibilities in the Python Ecosystem

2024-10-10T00:00:00+00:00

I presented at the local MadPy meetup group on the new possibilities in the Python ecosystem in light of recent development in the Rust world. Spent some time curating those links, so they may be worth a shuffle!

In case the embedded slides are not loading, here is the Google Slides.

Big Data Engineering in the 2020s - One Big Machine

2024-06-30T00:00:00+00:00

I wanted to write this post after juggling PBs of (non-tabular) data in the past year at work, and having experienced first-hand how amazing the Big Data landscape is morphing into in the mid-2020s. The ground is shifting under our feet, and individuals & small teams are immensely empowered with recent advances in hardware capabilities, accessible infrastructure, and modern software tooling. I would like to reaffirm that, as data practitioners we totally could, and should do serious work on one single big machine on a good cloud platform, and weigh the pros and cons of adopting abstractions carefully in case of scaling out. We don’t have to be impeded by the idea that Big Data has to be solved by a cluster of machines through distributed computing, thus out of reach for most developers. I will also lay out a minimally viable way to linearly scale out the occasional data processing needs, which can be particularly helpful for non-tabular data, where unified frameworks are still lacking.

Note: The workload described in this post is non-taublar data, usually revolving around big bespoke formats, raw media files, and ML records with binary data and tensors/arrays. Unlike tabular data with tens of popular cluster solutions, working with non-tabular data is more ad hoc. If you work with tabular data, and need to frequently do aggregation or training with columns of data, existing data warehouse and data lake solutions are quite advantageous.

In the past 10 years, a couple of major changes took place in computing:

Computing power per machine continued to grow, not through higher clock frequencies, but by inventing advanced instruction sets (e.g. SIMD), and packing more CPU cores and memory chips. It’s commonplace to see machines with 100+ cores and TBs of memory. What used to be achievable only by a cluster of machines now can be accessed by one single node. Similar things happened to accelerated computing, e.g. GPUs.
Cloud platforms have become commodities, with higher accessibility to a variety of instance types, resources, and predictable & ever-lowering costs. Users also have easy access to high network bandwidth, and near-infinite reliable cloud object storage that speaks S3 with high throughput. For example, you could get an instance from major cloud providers with 128 cores, 1 TB of memory, 30 TB of NVMe low-latency local storage, and 75 Gibps of network bandwidth for $10/hr.
Memory-safe systems programming languages became more mature and gained wide adoption, notably Rust. This enabled close-to-metal performance and utilization of modern processors’ parallel architectures. There emerged a proliferation of native tooling to tackle Big Data problems.
An important side effect of the wider systems languages adoption was high-level dynamic languages benefiting substantially from native extensions written in those languages. A Python programmer can easily access extreme performance otherwise not accessible before. Since Python has become the de facto language for ML/AI used by millions of people, this “upcast” dramatically extended those projects’ reach, as well as developers’ abilities. Examples: for tabular data, DuckDB, Polars, PyArrow, Delta Lake, for vector data, Faiss, hnswlib. New promising format: LanceDB.

Use One Big Machine!

Now, let us look at the benefits of using one single machine to do data work:

It is easy to spin up and down, scale up and down with different combinations of resources by switching VMs on the same cloud platform, attaching and detaching the same volume that hosts your code and key data. This keeps the cost to roughly just what you need.
It is easy to stay organized for code and key data, without needing to track multiple copies and sync between them. This results in less mental load and operational overhead. Data engineering is a very interactive type of work, which requires constant feedback running code on a portion of the actual data, handling errors that pop up, and observing live logs to ensure final success. As you move data, your code moves with them, oftentimes at a higher frequency due to tweaks to user interfaces, logging messages, etc. It’s paramount to stay sane and sanitize your memory of the steps you take and events that take place.
It has less complexity and is easier to program than a cluster of machines. The unit substrate of programming is typically one machine. So we get the full support from the ecosystem around us. Deployment is easier, without paying for containerization when you don’t need it. It’s also easier to plan resource utilization within the single-node model for CPUs, memory, storage, and networking. Observability is also much easier, just a few Linux commands away (htop, iostat, nload, etc.)
You get organic data locality for passing between processes & services through the loopback interface on the same host, without networking overhead. You are also encouraged to use threading that shares memory, or use shared memory (zero-copy memory layout, like a simple array, Arrow, dlpack for GPU tensors) between processes to avoid unnecessary copies and de-/serialization, thus getting peak speed.

With one single machine, you are bound by the CPU, memory, storage IOPS/throughput and network limits of the largest instance you could get on a given cloud platform. However, empowered by systems languages for native performance, you could lift a tremendous amount of data under those limits in a short duration of time, typically faster than the time for you to do the actual programming! Aim at 5~10 GB/s for storage/network throughput saturation, you could go through 400~800 TB of data a day!

Another important consequence of this high ceiling is it shortens the duration of your average stage of work in a long sequence down to minutes and hours, instead of days and weeks. Minutes and hours are within our attention span and context window, while days and weeks are not. Tasks that used to need to explicitly plan for, now are a matter of a lunch break, or as long as a night of sleep during off hours. It is an incredible productivity boost, and puts control and value back onto the developers, instead of machines.

You are also bound by immediate storage limits, but this is not a problem if your data could be chunked, and synced to/from cloud object storage. Of course, you need a toolkit that could saturate network bandwidth to prosper (e.g. AWS Rust SDK, official in Nov 2023). For data that can’t be chunked, usually you put them in a cluster of database nodes that are well-optimized with established interfaces.

Scaling Out

What about the times you actually need to scale out work to multiple machines? There are many cases for this: compute heavy workload, accelerated chips (GPUs) needed, or simply bargaining with some per-IP address rate/connection limiting mechanism from external services. Here is an approachable recipe to scale out linearly, assuming auth & permission control is not an issue (some environments are insulated from prod, thus satisfying this assumption):

Don’t pay orchestration and abstraction costs if you don’t need it. This means you don’t have to start with Kubernetes, or even containerization. It’s just you, your one single node, a code repository somewhere, and some networking setup.
Create a repository, with some control scripts that set up the remote machines through ssh, rsync, and idempotent scripts containing standard Linux command line instructions, that get executed remotely. If you are doing distributed computing, you are going to mess with system resources configuration anyway, and those abstractions are usually incomplete (custom drivers, swap space, ephemeral disk mounts, to name a few), let alone a whole separate thing to learn. Just use your one single node as the “head node”, and set the worker nodes up. Make sure you can run the setup script multiple times without ill effects. Make sure you sync whatever you need to the workers, anticipating worker VM’s OS varieties. Do some systemd-fu so you can remote control or auto-start worker processes.
You could opt in containerization, and set up a container registry to sync over the container images. This makes deployment a lot more predictable, but you pay the cost of pushing a new version with every incremental change. You should weigh the cost and benefit.
See if your data could be chunked in large units. When units of work are big, like batches of tens of thousands of smaller units, it’s easier and more efficient to store, list, and transfer. Your downstream derivative data will also map nicely to those bigger chunks, maintaining logical consistency and sanity. Lots of data sources that need intensive computation could satisfy this criterion.
You could have a manifest of those big chunks of data on your head node, set up the worker nodes, and dispatch only slim “pointer-like” file/object names/paths through a simple Redis list, or Stream (could be on your head node). Workers download from object storage, do work, and upload to a different location, or send back through a different Redis list. Tasks indicate success by uploading/sending their results to the final destination. Your head node can control and monitor those workers, and check work progress. You check the manifest, and resend the diffs for failed tasks, or automatically harvest the failed tasks until complete success.
By doing the above, you make each task a stateless unit of work, and thus enable a fail-able design pattern for batch processing. This also lets you use cheaper preemptable/spot instances, which usually are 1/3 of their full prices. It can scale out linearly, until bottlenecked by external factors, like total bandwidth of your physical worker group, or source/destination cloud storage account-level allowable throughput, if any. The sky is the limit.
This is the smallest amount of cost you need to pay to blast through PBs of data in 5~10 days, and you or your small team can coordinate this without much overhead. There is no unreliable abstraction out of your control, that could trap you or make you submit some ticket or issue to get handled by companies that built those products. Some tricky OS-level config that requires a bespoke setup? Easy - just add a line using command line tools, like how you’d do it on a single machine.
If your major programming language is Python, or you could readily create native extensions and integrate with Python, you could also opt in Ray, or Dask. It takes care of a lot of coordination for you, backed by a sound and efficient distributed computing design.

The progress in computing and data is incredible to witness. What will happen next?

Comparing Value and Reference Semantics across Languages

2024-06-01T00:00:00+00:00

Started learning Mojo, and watched this deep dive on ownership when it came out:

Just summarizing the value and reference semantics across three langauges: C++, Rust, and Mojo:

C++:

copy: `func(string s)`, caller `func(s)`
move: `func(string&& s)`, caller `func(std::move(s))`
immutable reference: `func(const string& s)`, caller `func(s)`
mutable reference: `func(string& s)`, caller `func(s)`

Rust:

copy: `func(s: String)`, caller `func(s.clone())`
move: `func(s: String)`, caller `func(s)`
immutable reference: `func(s: &String)`, caller `func(&s)`
mutable reference: `func(s: &mut String)`, caller `func(&mut s)`

Mojo:

copy: `func(owned s: String)`, caller `func(s)`
move: `func(owned s: String)`, caller `func(s^)`
immutable reference: `func(borrowed s: String)`, caller `func(s)`
mutable reference: `func(inout s: String)`, caller `func(s)`

Stateless Pipeline with Async Stream

2024-01-06T00:00:00+00:00

You can check out the code here.

After the series on pools and pipeline (Part I, Part II) was published, a friend graciously pointed me to another approach that would be simpler for a lot of occasions with less boilerplate code. It uses the async Stream trait and its iterator-like adaptors in StreamExt. That Stack Overflow post was from 2018, but up till today (Jan 2024), the Stream trait is still not stablized into the Rust standard library. However, what worked back then still works today. It is so useful, and has become my go-to option in many cases. So let’s look at it together.

There are some good guides on streams¹ ² ³. Note that both the futures and Tokio crates have their own StreamExt utility traits and methods. As of today, futures’ implementation has a broader selection. We will be using some of them here.

At first, I was having a little hard time wrapping my head around streams. A Stream is the async-equivalent to an Iterator. Iterators are lazy, and they will get executed at the end of the method chain by methods like collect() and for_each(). Streams are lazier, compounded by the fact that Rust futures are lazy. This introduces another degree of freedom and complexity into the richness of Rust functional programming that iterators champion. Because of this, we are able to express concurrency and parallelism directly in the stream method chaining. In the sync land, you’d need Rayon to achieve parallelism (we’ll compare them a bit later).

Comment: Little did I know when writing this post, that the relationship between iterators and futures in Rust has a long history⁴, both of which are fundamentally driven by what defines the modern Rust programming language with a need for zero-cost abstraction. In addition, there exhibits some elegant symmetry between them⁵ ⁶ (Maxwell’s Equations everyone?) These referenced articles were posted very recently (fresh off the Boats, sorry, can’t help it) and they are very relevant and insightful. So I would recommend a read if you are inclined to dig deeper.

Stateless Pipelines

Let’s look at the example in that post. I altered it a bit for clarity.

use anyhow::Result
use futures::{stream, StreamExt};
use reqwest::Client;

const N_TASKS: usize = 8;
const CONCURRENT_REQUESTS: usize = 4;

#[tokio::main]
async fn main() -> Result<()> {
    let client = Client::new();

    let urls = vec!["https://api.ipify.org"; N_TASKS];

    stream::iter(urls)
        .enumerate()
        .map(|(idx, url)| {
            let client = &client;
            async move {
                let t = std::time::Instant::now();
                println!("{idx}: Fetching {}...", url);
                let resp = client.get(url).send().await?;
                let bytes = resp.bytes().await?;
                println!("{idx}: Took {:.2}ms", t.elapsed().as_millis());
                anyhow::Ok((idx, bytes))
            }
        })
        .buffered(CONCURRENT_REQUESTS)
        // .buffer_unordered(CONCURRENT_REQUESTS)
        .for_each(|res| async {
            match res {
                Ok((idx, b)) => println!("{idx}: Got {} bytes", b.len()),
                Err(e) => eprintln!("Got an error: {}", e),
            }
        })
        .await;

    Ok(())
}

We define a vector urls of 8 entries. With some iterator-like syntax, we map() and for_each() to get it started. Because we are doing async here, we need to .await at the end. It all makes sense. But what are we returning in map() and for_each(), and what is buffered()/buffer_unordered()?? Let’s withhold the burning curiosity and run this example:

Fetching https://api.ipify.org...
Fetching https://api.ipify.org...
Fetching https://api.ipify.org...
Fetching https://api.ipify.org...
Took 398ms
Got 13 bytes
Fetching https://api.ipify.org...
Took 424ms
Got 13 bytes
Took 424ms
Fetching https://api.ipify.org...
Took 424ms
Got 13 bytes
Got 13 bytes
Fetching https://api.ipify.org...
Fetching https://api.ipify.org...
Took 89ms
Got 13 bytes
Took 110ms
Took 110ms
Got 13 bytes
Got 13 bytes
Took 110ms
Got 13 bytes

Four tasks in a batch are indeed running concurrently, followed by another batch! So let’s draw a comparison with the actor model here. Is there an actor defined? No. Is there a resource? Sort of - the reqwest::Client. When we use the client, we only need an immutable reference (or a cheap clone with tokio::spawn()). This is unlike a database connection or an ML model. We are using the client as a stateless bundle of code, and passing it around to tasks to download from links. It is in fact a fairly typical use case, and this approach with async streams is perfect for it. Additionally, in the Stack Overflow post, it was also shown how to use tokio::spawn() to get multi-threaded/multi-core parallelism.

Type Reasoning

So let’s get back to the previous question about map(), for_each(), and buffered()/buffer_unordered(). Let’s compare Iterator with Stream:

pub trait Iterator {
    type Item;

    // Provided methods
    fn map<B, F>(self, f: F) -> Map<Self, F>
       where Self: Sized,
             F: FnMut(Self::Item) -> B { ... }

    fn for_each<F>(self, f: F)
       where Self: Sized,
             F: FnMut(Self::Item) { ... }
}

pub trait StreamExt: Stream {
    // Provided methods
    fn map<T, F>(self, f: F) -> Map<Self, F>
       where F: FnMut(Self::Item) -> T,
             Self: Sized { ... }

    fn for_each<Fut, F>(self, f: F) -> ForEach<Self, Fut, F>
       where F: FnMut(Self::Item) -> Fut,
             Fut: Future<Output = ()>,
             Self: Sized { ... }

    fn buffered(self, n: usize) -> Buffered<Self>
       where Self::Item: Future,
             Self: Sized { ... }

    fn buffer_unordered(self, n: usize) -> BufferUnordered<Self>
       where Self::Item: Future,
             Self: Sized { ... }
}

It is curious. The map() signatures are quite similar, but it’s quite different for for_each(). For Iterator, it’s a straightforward closure that doesn’t return anything. For Stream, there is a trait bound F: FnMut(Self::Item) -> Fut, which means the closure needs to return a future (which is bound by Fut: Future<Output = ()>, a future that doesn’t return anything)!

If we dig a little deeper into async Rust, we will see that this is what an async function/closure does. In the RFC:

An async fn foo(args..) -> T is a function of the type fn(args..) -> impl Future<Output = T>. The return type is an anonymous type generated by the compiler.

So Stream’s for_each() needs an async closure, which is a closure that returns a future. That’s really what the type system tells us. Let’s look at what we gave to for_each():

        .for_each(|res| async {
            match res {
                Ok((idx, b)) => println!("{idx}: Got {} bytes", b.len()),
                Err(e) => eprintln!("Got an error: {}", e),
            }
        })

You see that async? That’s exactly it.

This is a good opportunity to share some notes of my Rust learning journey. Coming from a less typed background, it surely was an uphill battle to reason with Rust types in such intensity, and adding async to it was truly icing (salt) on the cake (wound)! On one hand, as a Rust programmer, you kind of have to understand this to get through the beginner’s phase, and this is where a lot of people give up. On the other hand, after these concepts begin to take shape in your head, Rust starts to become magical, and (many of) you might get really addicted - well, until you reach a certain point as a library developer, with some thorny fights with the compiler (here is a rebuttal).

What about `map()`?

Wait, so… Why are the two map()’s signatures so similar? You could certainly pass a regular closure to Stream’s map(). In fact, the doc uses it as an example:

Note that this function consumes the stream passed into it and returns a wrapped version of it, similar to the existing map methods in the standard library.
use futures::stream::{self, StreamExt};

let stream = stream::iter(1..=3);
let stream = stream.map(|x| x + 3);

assert_eq!(vec![4, 5, 6], stream.collect::<Vec<_>>().await);

If we used map() this way, there wouldn’t be any async magic happening. What is important yet left implied, is that you can pass it a closure that returns a future. After all, the trait bound F: FnMut(Self::Item) -> T does not impose what the closure is supposed to return (just a generic T). So it can certainly be a future!

This is in fact how exactly we enabled async concurrency from the stream in our example. Look:

        .map(|(idx, url)| {
            let client = &client;
            async move {
                let t = std::time::Instant::now();
                println!("{idx}: Fetching {}...", url);
                let resp = client.get(url).send().await?;
                let bytes = resp.bytes().await?;
                println!("{idx}: Took {:.2}ms", t.elapsed().as_millis());
                anyhow::Ok((idx, bytes))
            }
        })

The closure we gave to map() returns an async move {}. What is this? It’s an async block, a.k.a. a future! You could have written something like with for_each(), or below to use the async closure syntax, because, an async closure is a closure that returns a future:

// won't compile, due to the need to move `client` vs the `FnMut` bound in `F: FnMut(Self::Item) -> T`
        .map(|(idx, url)| async move {
            let t = std::time::Instant::now();
            println!("{idx}: Fetching {}...", url);
            let resp = client.get(url).send().await?;
            let bytes = resp.bytes().await?;
            println!("{idx}: Took {:.2}ms", t.elapsed().as_millis());
            anyhow::Ok((idx, bytes))
        })

// or
// still won't compile, due to capturing of `idx` and `url` by reference
        .map(|(idx, url)| async {
            let t = std::time::Instant::now();
            println!("{idx}: Fetching {}...", url);
            let resp = client.get(url).send().await?;
            let bytes = resp.bytes().await?;
            println!("{idx}: Took {:.2}ms", t.elapsed().as_millis());
            anyhow::Ok((idx, bytes))
        })

Those two exact blocks won’t actually compile, unfortunately. Because the future escapes the closure where it’s defined when the closure returns, you have to prove to the compiler the future captures the right things and owns the right things. If you feel a bit overwhelmed (I did many times), just recall that async Rust loves to annotate futures with a move when spawning tasks, like async move {}. When in doubt, try using the explicit regular closure returning an async block approach, and take clones/immutable references right before the async move {} line. Move things around and let the compiler teach you where you are.

Comment: Again, little did I know when writing this post, I stumbled upon one of the items Rust key contributors highlighted in 2023 and would like to resolve in the near future⁷. Reading up on them and a few of their other posts (see refs) after writing my own truly opened my eyes. I am very grateful for being enlightened from their insights, and remain hopeful these features are stablized for broad adoption.

What about `buffered()`/`buffer_unordered()`?

Look at the signatures:

    fn buffered(self, n: usize) -> Buffered<Self>
       where Self::Item: Future,
             Self: Sized { ... }

    fn buffer_unordered(self, n: usize) -> BufferUnordered<Self>
       where Self::Item: Future,
             Self: Sized { ... }

The Self::Item is just Stream::Item. It is saying, these two methods want to be chained to a stream that produces futures. Well, guess what we have after map()? A stream of futures! This is where async execution is enabled for the futures it consumes from up the chain, bound by CONCURRENT_REQUESTS at a time. By specifying buffered(CONCURRENT_REQUESTS), the results are received in the input futures’ order, and by buffer_unordered(CONCURRENT_REQUESTS), they are received in the order they are completed, potentially less wait from arriving at the input futures’ ordering.

Stacking up

We can create multiple stages of the pipeline by adding more map() and buffered()/buffer_unordered(). Why would we do that? After all, our pipeline is stateless, so we don’t have many resources to manage, so why do we still need different stages? Well, consider a pipeline that downloads from a source, and uploads to a destination. Even if the common “resource” is just a stateless client, the source and destination servers might exhibit different capacities and rate limiting rules. Using the same number of workers might not be compatible with both - one might be under the rate limit, the other might blow over. Our pipeline will then exhibit a modal behavior, swinging from all workers working on one stage, to all workers working on another at a time. In a way, the resources are external, and we need to capture that. By having separate stages, we can configure different numbers of workers to match the external capacities. Different stages will then work continuously to overlap with other stages. Let’s write it out:

const N_TASKS: usize = 8;
const CONCURRENT_REQUESTS_1: usize = 4;
const CONCURRENT_REQUESTS_2: usize = 4;

#[tokio::main]
async fn main() -> Result<()> {
    let client = Client::new();

    let urls = vec!["https://api.ipify.org"; N_TASKS];

    stream::iter(urls)
        .enumerate()
        .map(|(idx, url)| {
            let client = &client;
            async move {
                let t = std::time::Instant::now();
                println!("Stage 1 - {idx}: Fetching {}...", url);
                let resp = client.get(url).send().await?;
                let bytes = resp.bytes().await?;
                println!("Stage 1 - {idx}: Took {:.2}ms", t.elapsed().as_millis());
                anyhow::Ok((idx, bytes))
            }
        })
        .buffered(CONCURRENT_REQUESTS_1)
        .map(|res| async move {
            match res {
                Ok((idx, b)) => {
                    let t = std::time::Instant::now();
                    println!("Stage 2 - {idx}: Starting...");
                    let s = String::from_utf8(b.to_vec())?;
                    tokio::time::sleep(std::time::Duration::from_millis(500)).await;
                    println!("Stage 2 - {idx}: Took {:.2}ms", t.elapsed().as_millis());
                    anyhow::Ok((idx, s))
                }
                Err(e) => Err(e),
            }
        })
        .buffered(CONCURRENT_REQUESTS_2)
        .for_each(|res| async {
            match res {
                Ok((idx, s)) => println!("Finally: {idx}: Got string with length {}", s.len()),
                Err(e) => eprintln!("Finally: Got an error: {}", e),
            }
        })
        .await;

    Ok(())
}

Here is the output. You can see the stages are overlapping - task 0 went to “Finally” early on.

Stage 1 - 0: Fetching https://api.ipify.org...
Stage 1 - 1: Fetching https://api.ipify.org...
Stage 1 - 2: Fetching https://api.ipify.org...
Stage 1 - 3: Fetching https://api.ipify.org...
Stage 1 - 2: Took 425ms
Stage 1 - 0: Took 435ms
Stage 1 - 1: Took 434ms
Stage 1 - 4: Fetching https://api.ipify.org...
Stage 1 - 3: Took 434ms
Stage 2 - 0: Starting...
Stage 2 - 1: Starting...
Stage 2 - 2: Starting...
Stage 2 - 3: Starting...
Stage 2 - 0: Took 502ms
Finally: 0: Got string with length 13
Stage 1 - 5: Fetching https://api.ipify.org...
Stage 1 - 6: Fetching https://api.ipify.org...
Stage 1 - 4: Took 502ms
Stage 2 - 1: Took 502ms
Finally: 1: Got string with length 13
Stage 1 - 7: Fetching https://api.ipify.org...
Stage 2 - 2: Took 502ms
Finally: 2: Got string with length 13
Stage 2 - 3: Took 502ms
Finally: 3: Got string with length 13
Stage 2 - 4: Starting...
Stage 1 - 7: Took 102ms
Stage 1 - 5: Took 102ms
Stage 1 - 6: Took 102ms
Stage 2 - 5: Starting...
Stage 2 - 6: Starting...
Stage 2 - 7: Starting...
Stage 2 - 4: Took 502ms
Finally: 4: Got string with length 13
Stage 2 - 5: Took 502ms
Finally: 5: Got string with length 13
Stage 2 - 6: Took 502ms
Finally: 6: Got string with length 13
Stage 2 - 7: Took 502ms
Finally: 7: Got string with length 13

Extension: Add a Writer Output

There are many ways to extend our pipeline to do more things. Let’s look at the typical input and output - files. We will also think about concurrency and parallelism from the start.

The input could be many files, and concurrency/parallelism can go from there. In that case, put file names into the initial vector, and create a stream by Stream::iter(filenames). Each file will be a unit of work. Alternatively, if you have one large file, you could use the stream! macro from the async-stream crate to define an iterator from file lines.

It’s worth going into the output a bit more. The output writer is usually required to be mut, and it might be tricky to set up in the async context. First, we need a writer, and we modify the for_each closure to use the writer:

use tokio::fs::File;
use tokio::io::{BufWriter, AsyncWriteExt};
// ...
    let mut writer = BufWriter::new(File::create("output.txt").await?);

// ...
        .for_each(|res| {
            async {
                match res {
                    Ok((idx, s)) => {
                        println!("Finally: {idx}: Got string with length {}", s.len());
                        writer.write_all(format!("{}\n", s).as_bytes()).await.unwrap();
                    },
                    Err(e) => eprintln!("Finally: Got an error: {}", e),
                }
            }
        })
        .await;

// ...
    writer.flush().await?;

The compiler is mad:

error: captured variable cannot escape `FnMut` closure body
  --> src/bin/streamext.rs:46:13
   |
13 |       let mut writer = BufWriter::new(File::create("output.txt").await?);
   |           ---------- variable defined here
...
45 |           .for_each(|res| {
   |                         - inferred to be a `FnMut` closure
46 | /             async {
47 | |                 match res {
48 | |                     Ok((idx, s)) => {
49 | |                         println!("Finally: {idx}: Got string with length {}", s.len());
50 | |                         writer.write_all(format!("{}\n", s).as_bytes()).await.unwrap();
   | |                         ------ variable captured here
...  |
53 | |                 }
54 | |             }
   | |_____________^ returns an `async` block that contains a reference to a captured variable, which then escapes the closure body
   |
   = note: `FnMut` closures only have access to their captured variables while they are executing...
   = note: ...therefore, they cannot allow references to captured variables to escape

The message is quite helpful. However, there is one thing left out - if this were a reader, we could take out an immutable reference, and the lifetime rules would be met. The real problem is we need a mutable reference to the writer here. And the fact that async enables concurrency, we are potentially taking out multiple mutable references! And importantly, there is no multi-threading going on here - no spawn(), no executing on other threads, but just plain concurrency. This is not to say multiple for_each() calls are actually executing concurrently, but just that the state machine mechanism in async Rust requires the future to outlive the for_each() closure, thus carrying all captured references with it, immutable or mutable. We just need a proof that we know what we are doing in a safe way. Does it remind you of anything? Yes, internal mutability and Rc<RefCell<T>>! This was one of my first few mental blocks learning Rust. Initially, I could not understand why no more than one mutable references could exist at the same time. If you are like me, think about it in this case and it might further your understanding. The compiler is here to help you with even plain concurrency, to ensure concurrent tasks will not modify something you are not expecting to change!

The following question is, what is changing here? Why do we need a mutable reference to the writer? Well, it’s writing something, and it manages some buffer, or some system file descriptor. The way to get out of it is again, Rc<RefCell<T>>:

use std::cell::RefCell;
use std::rc::Rc;
// ...
    let writer = Rc::new(RefCell::new(BufWriter::new(File::create("output.txt").await?)));

// ...
        .for_each(|res| async {
            match res {
                Ok((idx, s)) => {
                    println!("Finally: {idx}: Got string with length {}", s.len());
                    writer
                        .borrow_mut()
                        .write_all(format!("{}\n", s).as_bytes())
                        .await
                        .unwrap();
                }
                Err(e) => eprintln!("Finally: Got an error: {}", e),
            }
        })
        .await;

// ..
    writer.borrow_mut().flush().await?;

This works! Check out the output.txt! What a journey.

Currently, async Rust does not handle file IO very efficiently (absent io-uring). In fact, the tokio::fs is mostly a wrapper to std::fs with a dedicated big thread pool for blocking IO tasks⁸ ⁹. As a result, there is a lot of overhead moving tasks to the thread pool and results back if IO tasks are small. If your output file IO is slowing you down, consider using async-channel, or Flume to send() tasks in the async land, and receive_blocking() in the sync land with spawn_blocking(). It’s a very useful pattern.

Comparison with Rayon in Sync Land

It’s worthwhile drawing some comparison between async streams to Rayon parallel iterators in the sync land. These are two different beasts. Rayon is designed to achieve multi-threaded/multi-core parallelism for high-compute-density tasks that fully take up CPU cycles. It expects you to feed it tasks that won’t wait on IO for too long, and usually no more than N of (physical or logical with hyperthreading) cores number of tasks are running in parallel. Async Rust solves the complementary problem, where your tasks have heavy IO needs, and you need to spawn many more tasks to take up even one CPU’s cycles (this is concurrency) to saturate computing resources - we are assuming IO limit is high. In Rust, async tasks can also be truly multi-threaded/multi-core parallel, depending on the executor, like Tokio. The downside is you need to ensure IO function calls are not blocking, but yielding, which creates a “paralel universe” to the sync land. This has been a point of frustration to async Rust, and yet is tied to everything that makes modern Rust the way it is - zero-cost abstraction. You are encouraged to read up⁴ on it.

From a user’s perspective, however, the tasks are always a mix of both CPU- and IO- intensive. If you really want it, you could run CPU-intensive tasks with async, and IO-intensive tasks with Rayon. Performance may degrade quite a lot, but it should still work.

With Rayon, there are two approaches to achieve parallelism. You stick a .par_iter()/.into_par_iter() to a collection with a known length (not an iterator), like a Vec, or a HasMap, and Rayon uses divide-and-conquer to cut up the sequence. The advantage of this approach is work division is done on the stack with recursion, and is extremely fast when work scheduling is the bottleneck otherwise. The disadvantage is you can’t convert any iterator with unknown length to a parallel one this way. To do that, Rayon added par_bridge(). However, one notable behavior difference is the collected results from it is not ordered according to the original iterator. You could borrow the logic from here as a remedy.

With async streams, you can express concurrency with or without parallelism, and define stages with an arbitrary number of tasks executed at the same time. There is more flexibility, due to the nature of async, and of course, more complexity.

Pools and Pipeline with Tokio (Part II - Streaming)

2023-12-28T00:00:00+00:00

You can check out the code here.

In Part I of this mini-series, I implemented an actor pool pattern to perform IO-intensive pipeline jobs using the RPC model. The RPC model is relatively easy to program, because the main control flow code is together in one place, which abstracts away the dependencies as services. I don’t know about you, but this paradigm makes me feel like I’m following along in the life cycle of each task, and unless I send a request to some service, nothing moves. There is a psychological sense of safety to it… The biggest downside used to be spawning threads, which is expensive in synchronous programming. However, with async coroutines, it’s less of an issue.

The Streaming Model

Compare this with the streaming model. By that I mean the worker pool exposes send() and recv() methods, behaving like a processing queue. One single caller could keep calling send() sequentially, and tasks get sent off to the worker pool by a bounded input queue/channel and processed in parallel. The caller is not spawning anything but just sending messages. Then, the results would be sent to a bounded output queue within each worker. There is another coroutine/thread that catches the results from the output queue, like what we did in Part I. The queues are bounded, so the task sender feels the back pressure without needing a semaphore. When dealing with multiple worker pools in stages, we spin up connector coroutine/threads that carry tasks from last pool’s output queue to the next pool’s input queue. The input and output messages both contain unique markers, and we handle errors at the connectors and final output receiver.

How does this make you feel? It feels more dynamic and fluid to me, and the composition of the control flow is spread across connectors and receivers. However, this model maps more idiomatically to a pipeline with separate stages represented as boxes on a flow chart, and we imagine tasks flow through these boxes like a stream. It’s a more natural representation than the RPC model for data processing.

Let us revisit the example from Part I, customized to the streaming model. There are more changes than simply converting an ActorHandle to an ActorPool. First the messages and the error. We now have two message types, one for input, and the other for output. We need to encode a unique identifier to each input/output message, such as a sequential number to keep track of them. We also need a custom error type with the unique identifier to encapsulate the errors handling the task. For simplicity’s sake, I converted the original ActorMessage into a struct from an enum. There is no oneshot channel needed because the output is sent through the output channel.

use log::error;
use thiserror::Error;

pub struct MyActorInputMessage {
    pub idx: u64,
}

pub struct MyActorOutputMessage {
    pub idx: u64,
    pub data: u32,
}

#[derive(Error, Debug)]
#[error("Actor error: {idx}")]
pub struct MyActorError {
    pub idx: u64,
}

Then, we implement the private MyActor. It takes the receiver part of the input channel, and the sender part of the output channel. In handle_message(), we evaluate an async task, and map_err whatever errors the underlying business logic code surfaces to our own error type. Over time, this error type will grow more complicated, maybe into an enum with variants that are of interest to the receiving party. The output is a Result type.

struct MyActor {
    in_receiver: async_channel::Receiver<MyActorInputMessage>,
    out_sender: async_channel::Sender<Result<MyActorOutputMessage, MyActorError>>,
    next_id: u32,
}

impl MyActor {
    fn new(
        in_receiver: async_channel::Receiver<MyActorInputMessage>,
        out_sender: async_channel::Sender<Result<MyActorOutputMessage, MyActorError>>,
    ) -> Self {
        MyActor {
            in_receiver,
            out_sender,
            next_id: 0,
        }
    }

    async fn handle_message(&mut self, msg: MyActorInputMessage) {
        let MyActorInputMessage { idx } = msg;
        self.next_id += 1;
        let next_id = self.next_id;
        let res = async move {
            tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
            anyhow::Ok(MyActorOutputMessage { idx, data: next_id })
        }
        .await
        .map_err(|e| {
            error!("MyActor: idx {}: error: {}", idx, e);
            MyActorError { idx }
        });

        // The `let _ =` ignores any errors when sending.
        //
        // This can happen if the `select!` macro is used
        // to cancel waiting for the response.
        let _ = self.out_sender.send(res).await;
    }

    async fn run(&mut self) {
        while let Ok(msg) = self.in_receiver.recv().await {
            self.handle_message(msg).await;
        }
    }
}

Then the streaming actor pool. It takes the sender part of the input channel, and the receiver part of the output channel. There is a send() that creates an input message and sends to the input channel. We propagate the async_channel::SendError to the caller, because the request-response cycle is no longer tightly coupled, and it’s good for the caller to know immediately what’s wrong. There is a recv() that forwards the Result to the receiving party. Note we convert the Result<(), async_channel::RecvError> into an Option<()> with .ok() to create a little more clarity. We don’t have graceful shutdown yet, but when we do, None will indicate there are no more messages and the pool is shut down.

pub struct MyActorStreamerPool {
    in_sender: async_channel::Sender<MyActorInputMessage>,
    out_receiver: async_channel::Receiver<Result<MyActorOutputMessage, MyActorError>>,
}

impl MyActorStreamerPool {
    pub fn new(num_actors: usize) -> Self {
        let (in_sender, in_receiver) = async_channel::bounded(8);
        let (out_sender, out_receiver) = async_channel::bounded(8);
        for _ in 0..num_actors {
            let in_receiver = in_receiver.clone();
            let out_sender = out_sender.clone();
            tokio::spawn(async move { MyActor::new(in_receiver, out_sender).run().await });
        }

        Self {
            in_sender,
            out_receiver,
        }
    }

    pub async fn send(
        &self,
        idx: u64,
    ) -> Result<(), async_channel::SendError<MyActorInputMessage>> {
        let msg = MyActorInputMessage { idx };
        self.in_sender.send(msg).await
    }

    pub async fn recv(&self) -> Option<Result<MyActorOutputMessage, MyActorError>> {
        self.out_receiver.recv().await.ok()
    }
}

Control Flow

Now let’s look at how we wire them together. We use the pool itself, and the Result from the output channel. The input message is not needed directly, because send()’s signature takes care of that for us. We need the pool in a spawned task to send messages sequentially, and in the foreground task, we get into a loop to wait on the messages from the output channel. We handle the Result by a match.

use pools_and_pipeline::my_streamer_pool::{MyActorError, MyActorOutputMessage, MyActorStreamerPool};

const N_TASKS: u64 = 8;

#[tokio::main]
async fn main() {
    // create pool
    let pool = std::sync::Arc::new(MyActorStreamerPool::new(4));

    // submit tasks
    let pool_ = pool.clone();
    tokio::spawn(async move {
        for idx in 0..N_TASKS {
            pool_
                .send(idx)
                .await
                .expect("Failed to send message to pool");
        }
    });

    // wait for all tasks to finish
    loop {
        let t = tokio::time::Instant::now();
        let Some(res) = pool.recv().await else {
            break;
        };
        match res {
            Ok(MyActorOutputMessage { idx, data }) => {
                println!("{idx} ended in {}ms", t.elapsed().as_millis());
                println!("idx {}: data: {}", idx, data);
            }
            Err(MyActorError { idx }) => {
                println!("{idx} ended in {}ms", t.elapsed().as_millis());
                println!("idx {}: error", idx);
            }
        }
    }
}

In Part I, we #[derive(Clone)] earlier on the MyActorPool. Cloning the pool is really cloning the channels the pool owns. To be more clear, here we skip the derive, but instead wrap the pool around an Arc. This makes the relationship more explicit. This is also essential to implement graceful shutdown, which we will visit later.

Let’s run it. We see that ~500ms only showed up twice, and the rest are 0ms. It means tasks are running in parallel by all workers. Again, we haven’t implemented graceful shutdown yet, so we need to CTRL+C out of the loop.

3 ended in 502ms
idx 3: data: 1
2 ended in 0ms
idx 2: data: 1
1 ended in 0ms
idx 1: data: 1
0 ended in 0ms
idx 0: data: 1
7 ended in 501ms
idx 7: data: 2
4 ended in 0ms
idx 4: data: 2
6 ended in 0ms
idx 6: data: 2
5 ended in 0ms
idx 5: data: 2
^C

Graceful Shutdown

On a high level, how does graceful shutdown work? Recall the webserver example from the Book, when things involve channels, usually dropping all clones of one side would prompt the other side to receive a None or Err, exiting out of the loop. How can we do that here?

We hold two clones of the pool object, one for the sending party, and the other for the receiving party. With data processing, the sending party usually knows when the source is finished. So we should drop from the sending party. We need to implement a close() method for the pool, that somehow drops the input channel sender. The private MyActors would break out of their own loops, dropping the output channel sender. Finally the pool’s recv() call will break out of its loop. This is done through wrapping the Senders in Options, and calling take() to swap them out to drop.

The private MyActor:

struct MyActor {
    in_receiver: async_channel::Receiver<MyActorInputMessage>,
    out_sender: Option<async_channel::Sender<Result<MyActorOutputMessage, MyActorError>>>,
    next_id: u32,
}

impl MyActor {
    fn new(
        in_receiver: async_channel::Receiver<MyActorInputMessage>,
        out_sender: async_channel::Sender<Result<MyActorOutputMessage, MyActorError>>,
    ) -> Self {
        MyActor {
            in_receiver,
            out_sender: Some(out_sender),
            next_id: 0,
        }
    }

    async fn handle_message(&mut self, msg: MyActorInputMessage) {
        let MyActorInputMessage { idx } = msg;
        self.next_id += 1;
        let next_id = self.next_id;
        let res = async move {
            tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
            anyhow::Ok(MyActorOutputMessage { idx, data: next_id })
        }
        .await
        .map_err(|e| {
            error!("MyActor: idx {}: error: {}", idx, e);
            MyActorError { idx }
        });

        match &self.out_sender {
            // The `let _ =` ignores any errors when sending.
            //
            // This can happen if the `select!` macro is used
            // to cancel waiting for the response.
            Some(out_sender) => {
                let _ = out_sender.send(res).await;
            }
            None => {
                error!("MyActor: idx {}: error: pool is closed.", idx);
            }
        }
    }

    async fn run(&mut self) {
        while let Ok(msg) = self.in_receiver.recv().await {
            self.handle_message(msg).await;
        }
        self.out_sender.take();
    }
}

The MyActorStreamerPool:

pub struct MyActorStreamerPool {
    in_sender: Option<async_channel::Sender<MyActorInputMessage>>,
    out_receiver: async_channel::Receiver<Result<MyActorOutputMessage, MyActorError>>,
}

impl MyActorStreamerPool {
    pub fn new(num_actors: usize) -> Self {
        let (in_sender, in_receiver) = async_channel::bounded(8);
        let (out_sender, out_receiver) = async_channel::bounded(8);
        for _ in 0..num_actors {
            let in_receiver = in_receiver.clone();
            let out_sender = out_sender.clone();
            tokio::spawn(async move { MyActor::new(in_receiver, out_sender).run().await });
        }

        Self {
            in_sender: Some(in_sender),
            out_receiver,
        }
    }

    pub async fn send(
        &self,
        idx: u64,
    ) -> Result<(), async_channel::SendError<MyActorInputMessage>> {
        let msg = MyActorInputMessage { idx };
        match &self.in_sender {
            Some(in_sender) => in_sender.send(msg).await,
            None => Err(async_channel::SendError(msg)),
        }
    }

    pub async fn recv(&self) -> Option<Result<MyActorOutputMessage, MyActorError>> {
        self.out_receiver.recv().await.ok()
    }

    pub fn close(&mut self) {
        self.in_sender.take();
    }
}

This is where wrapping MyActorStreamerPool with Arc is key to graceful shutdown. When we call close(), we need to make sure it acts on the very same MyActorStreamerPool that is doing the recv(). Otherwise, not all clones of the input channel Sender are dropped, and shutdown wouldn’t happen.

Also notice close() needs a &mut self, so in calling with Arc, we need internal mutability. So we need either a Mutex or RwLock. Would that slow down regular message passing because every action has to be protected by a lock? Luckily for us, in this case both send() and recv() of the async-channel only need a non-mutable &self. I looked into other crates, neither Flume, or the sync world Crossbeam require these methods to take &mut self. There must be a reason for this “luck”, but I haven’t looked deeper into it. Let’s just take advantage of it now! With simply a RwLock, we can incur minimal locking, only when close() is called.

Now the async main becomes:

#[tokio::main]
async fn main() {
    // create pool
    let pool = std::sync::Arc::new(tokio::sync::RwLock::new(MyActorStreamerPool::new(4)));

    // submit tasks
    let pool_ = pool.clone();
    tokio::spawn(async move {
        for idx in 0..N_TASKS {
            pool_
                .read()
                .await
                .send(idx)
                .await
                .expect("Failed to send message to pool");
        }
        println!("All tasks submitted. Closing pool.");
        pool_.write().await.close();
    });

    // wait for all tasks to finish
    let mut res_all = vec![];
    loop {
        let t = tokio::time::Instant::now();
        let Some(res) = pool.read().await.recv().await else {
            break;
        };
        match res {
            Ok(MyActorOutputMessage { idx, data }) => {
                println!("{idx} ended in {}ms", t.elapsed().as_millis());
                res_all.push(data);
            }
            Err(MyActorError { idx }) => {
                println!("{idx} ended in {}ms", t.elapsed().as_millis());
                println!("idx {}: error", idx);
            }
        }
    }
    println!("res_all = {:?}", res_all);
}

And the output looks like:

All tasks submitted. Closing pool.
ended in 502ms
ended in 0ms
ended in 0ms
ended in 0ms
ended in 502ms
ended in 0ms
ended in 0ms
ended in 0ms
res_all = [1, 1, 1, 1, 2, 2, 2, 2]

Potential Customizations

There are a few potential points of adaptation. One is, you might not want to close the pool to break out of the receiver loop. In that case, you might also know how many tasks the sender would send. So just count to that and break. You are able to reuse the pool this way, but also responsible for making sure different occasions of using the pool won’t interfere with each other.

Another consideration: the sender might hold a source that is !Send (like many iterators from Itertools), which means it is not able to be sent between threads across an .await. This could be solved with spawn_local(), or a current_thread runtime. But maybe the simplest approach for this iterator source scenario (it’s mildly blocking anyway) is a sync send_blocking() API (supported by async-channel and Flume), coupled with spawn_blocking(). Note we also need the sync version std::sync::RwLock instead of tokio::sync::RwLock for it.

// impl MyActorStreamerPool {

    pub fn send_blocking(
        &self,
        idx: u64,
    ) -> Result<(), async_channel::SendError<MyActorInputMessage>> {
        let msg = MyActorInputMessage { idx };
        match &self.in_sender {
            Some(in_sender) => in_sender.send_blocking(msg),
            None => Err(async_channel::SendError(msg)),
        }
    }

    pub fn recv_blocking(&self) -> Option<Result<MyActorOutputMessage, MyActorError>> {
        self.out_receiver.recv_blocking().ok()
    }

// }

// async fn main() {

    // submit tasks
    let pool_ = pool.clone();
    tokio::task::spawn_blocking(move || {
        for idx in 0..N_TASKS {
            pool_
                .read()
                .unwrap()
                .send_blocking(idx)
                .expect("Failed to send message to pool");
        }
        println!("All tasks submitted. Closing pool.");
        pool_.write().unwrap().close();
    });

// }

It’s kind of magical that all these can work, right? Rust and Tokio are so flexible and as long as the code compiles, you are almost there with a safe program that will run until the job is finished. Simply amazing!

One other potential adaptation is returning the results in the original order, in a streaming way, without first hoarding up all the results. This is similar to the very useful multiprocessing.Pool.imap. You could borrow the logic from here in the pool’s recv(). Not going more into it, but briefly, you need a next_index counter, another pair of Sender and Receiver, and a HashMap; at pool init, you spawn a task in the background that constantly reorders the results from the raw out_receiver and sends through that extra Sender. If you wish to reuse the pool, you also need a reset_next_index() method to clear the counter state. There are a lot of states, and depending on the order of results, you might end up hoarding a lot of memory anyway. There may be a better way to do it.

Full Program

Finally, with the bells and whistles like in Part I, this is the main program:

use indicatif::{ProgressBar, ProgressStyle};
use pools_and_pipeline::my_actor_streamer_pool::{
    MyActorError, MyActorOutputMessage, MyActorStreamerPool,
};
use pools_and_pipeline::utils::InterruptIndicator;
use std::sync::Arc;
use tokio::sync::RwLock;

const N_TASKS: u64 = 8;

#[tokio::main]
async fn main() {
    // attach CTRL+C handler
    let interrupt_indicator = InterruptIndicator::new();

    // make pretty progress bar
    let pb = ProgressBar::new(N_TASKS);
    let sty = ProgressStyle::with_template(
        "{spinner:.cyan} [{bar:40.cyan/blue}] {pos:>7}/{len:7} [{elapsed_precise}<{eta_precise} {per_sec:.green}] {msg}"
    ).unwrap().progress_chars("#>-");
    pb.set_style(sty);

    // create pool
    let pool = Arc::new(RwLock::new(MyActorStreamerPool::new(4)));

    // submit tasks
    let pool_ = pool.clone();
    tokio::spawn(async move {
        for idx in 0..N_TASKS {
            if interrupt_indicator.is_set() {
                println!("Interrupted! Exiting gracefully...");
                break;
            }

            pool_
                .read()
                .await
                .send(idx)
                .await
                .expect("Failed to send message to pool");
        }
        println!("All tasks submitted. Closing pool.");
        pool_.write().await.close();
    });

    // wait for all tasks to finish
    let mut res_all = vec![];
    loop {
        let t = tokio::time::Instant::now();
        let Some(res) = pool.read().await.recv().await else {
            break;
        };
        match res {
            Ok(MyActorOutputMessage { idx: _idx, data }) => {
                res_all.push(data);
            }
            Err(MyActorError { idx }) => {
                println!("{idx} ended in {}ms", t.elapsed().as_millis());
                println!("idx {}: error", idx);
            }
        }
        pb.inc(1);
    }
    pb.finish();
    println!("res_all = {:?}", res_all);
}

All tasks submitted. Closing pool.
^C⠁ [#####>----------------------------------]       1/8       [00:00:00<00:00:03 1.9906/s]
  [########################################]       8/8       [00:00:01<00:00:00 7.9636/s]
res_all = [1, 1, 1, 1, 2, 2, 2, 2]

Because of the buffering at the input channel, tasks would run further until gracefully exiting.

Again, you probably want to use Anyhow to handle errors and convert the unwrap() into ? propagation for good measure.

To Conclude

The RPC and streaming models in this series are suitable for more elaborate tasks with actors. They require some boilerplate code to set up, but you get great encapsulation, resource management, and separation of responsibilities. On the other hand, what do you do when you don’t need actors, and the pipeline is mostly stateless? Let’s expore that in another post.

Pools and Pipeline with Tokio (Part I - RPC)

2023-12-27T00:00:00+00:00

You can check out the code here.

I am working on a Rust project that goes through a lot of small tasks in parallel. These tasks are primarily IO intensive, like downloading objects from S3, processing them, and storing the results somewhere else. Since this is a fairly representative use case in data engineering, I am keen on working out a generalized design pattern that can be customized to more scenarios. This post lays out such a pattern, and I will explain some of the specific considerations and choices made along the way.

In a distributed setting, there would be a global task queue to feed the inputs from a head node, a first-stage worker pool pulling tasks off the queue, and working through them in parallel. Then those workers would feed the outputs into another node-local queue, which connects to a second-stage worker pool to do some other parallel processing. The stages propagate, and each stage’s worker pool holds some special resources. When started and stabilized, the worker pools would run together and overlap the task execution between stages, benefiting from “pipeline parallelism”.

In the past, I would reach for Python to carry out this type of data pipelining. Due to its slowness and global interpreter lock (GIL), I had to jump through hoops to obtain multi-core usage to get decent throughput, and the resulting design was still riddled with thread contention, high memory usage, and unnecessary (de-)serialization if multiprocessing is used. If you are serious enough, you would probably want to use a distributed framework, like Dask, or Ray. They are all very useful, but I would like to make a case for single-node setups with a fast language - it’s flexible in iterative development and deployment, and it’s really efficient. Often, you don’t need a cluster, but just a node and some code to get similar throughput, especially when there are constraints at the source or the final destination.

The official release of the Rust AWS SDK was a final push for me. Since it’s using Tokio as its async runtime, and Tokio blends concurrency and multi-threaded parallelism together, it’s an ideal substrate for this work. Tokio also has smooth escape hatches in and out of the async world, so you know you are covered with the edge cases. In this post, we use the words “concurrent” and “parallel” in a similar vein, versus “sequential” / “serial”. However, There is an important but different angle to distinguish concurrency and parallelism regarding CPU cores, threads, and task scheduling.

On a high level, given a source of tasks, what do we need to process them in parallel? Of course, we need worker pools to have workers that run in parallel. But that’s not enough. Depending on the worker pool API, we might need a bridge from the serial land to the parallel land. If the worker pool adopts the message streaming model with separate send() and recv(), behaving like a processing queue, this bridge is implied and provided. However, if the worker pool adopts the client-server RPC model, there need to be multiple task submitters hitting the “server” to get the workers working in parallel.

The RPC Model

Let’s zoom in on the RPC model, building off the example provided by the excellent post Actors with Tokio. In that post, an actor design was proposed, with the ability to complete the full request-response cycle of a client call to the actor. Let’s write this out as the starting point. First, the actor itself (modified to add a 500-millisecond sleep in the message handler):

use tokio::sync::{mpsc, oneshot};

enum ActorMessage {
    GetUniqueId { respond_to: oneshot::Sender<u32> },
}

struct MyActor {
    receiver: mpsc::Receiver<ActorMessage>,
    next_id: u32,
}

impl MyActor {
    fn new(receiver: mpsc::Receiver<ActorMessage>) -> Self {
        MyActor {
            receiver,
            next_id: 0,
        }
    }

    async fn handle_message(&mut self, msg: ActorMessage) {
        match msg {
            ActorMessage::GetUniqueId { respond_to } => {
                self.next_id += 1;
                tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;

                // The `let _ =` ignores any errors when sending.
                //
                // This can happen if the `select!` macro is used
                // to cancel waiting for the response.
                let _ = respond_to.send(self.next_id);
            }
        }
    }

    async fn run(&mut self) {
        while let Some(msg) = self.receiver.recv().await {
            self.handle_message(msg).await;
        }
    }
}

Note we store next_id to demonstrate the actor holding a mutable state that changes upon client requests. This is just one reason why you would use the actor model. In practice, you could choose it because you want to hold an immutable resource that is expensive to create and destroy, like a database connection, or an ML model.

Then the actor handle, with which we reference the actor:

#[derive(Clone)]
pub struct MyActorHandle {
    sender: mpsc::Sender<ActorMessage>,
}

impl MyActorHandle {
    pub fn new() -> Self {
        let (sender, receiver) = mpsc::channel(8);
        tokio::spawn(async move { MyActor::new(receiver).run().await });

        Self { sender }
    }

    pub async fn get_unique_id(&self) -> u32 {
        let (send, recv) = oneshot::channel();
        let msg = ActorMessage::GetUniqueId { respond_to: send };

        // Ignore send errors. If this send fails, so does the
        // recv.await below. There's no reason to check the
        // failure twice.
        let _ = self.sender.send(msg).await;
        recv.await.expect("Actor task has been killed")
    }
}

Note that in practice, you could have MyActor().handle_message() send a Result type with a custom Err variant, and MyActorHandle().get_unique_id() can return that to the caller to handle errors. We’ll do that in Part II.

We would create and refer to the actor with:

    let actor = MyActorHandle::new();

Actor Pool

This is just one actor. How can we make a pool of them? We need to consider the ergonomics from the caller’s perspective. We would like to have one single instance to call, instead of having to go through some selection logic. So these actors in the pool should share the same request channel. Tokio’s mpsc channel does not allow multiple actors to share the receiver. So we either need to wrap it with a Arc<Mutex<mpsc::Reciver>>, or use an mpmc channel, such as from async-channel, or Flume. With that, we convert the MyActorHandle into a MyActorPool:

// adapt MyActor to using async_channel
struct MyActor {
    receiver: async_channel::Receiver<ActorMessage>,
    next_id: u32,
}
// ...

#[derive(Clone)]
pub struct MyActorPool {
    sender: async_channel::Sender<ActorMessage>,
}

impl MyActorPool {
    pub fn new(num_actors: usize) -> Self {
        let (sender, receiver) = async_channel::bounded(8);
        for _ in 0..num_actors {
            let receiver = receiver.clone();
            tokio::spawn(async move { MyActor::new(receiver).run().await });
        }

        Self { sender }
    }

    pub async fn get_unique_id(&self) -> u32 {
        let (send, recv) = oneshot::channel();
        let msg = ActorMessage::GetUniqueId { respond_to: send };

        // Ignore send errors. If this send fails, so does the
        // recv.await below. There's no reason to check the
        // failure twice.
        let _ = self.sender.send(msg).await;
        recv.await.expect("Actor task has been killed")
    }
}

Now let’s create 4 actors in a pool and call it 8 times:

use pools_and_pipeline::my_actor_pool::MyActorPool;

const N_TASKS: u64 = 8;

#[tokio::main]
async fn main() {
    let pool = MyActorPool::new(4);

    let mut res_all = vec![];
    for i in 0..N_TASKS {
        let t = tokio::time::Instant::now();
        println!("{i} starting...");
        let res = pool.get_unique_id().await;
        res_all.push(res);
        println!("{i} ended in {}ms", t.elapsed().as_millis());
    }
    println!("res_all = {:?}", res_all);
}

Your output might be different, because different actor workers are taking the calls, incrementing & returning their internal counters.

starting...
ended in 502ms
starting...
ended in 501ms
starting...
ended in 502ms
starting...
ended in 502ms
starting...
ended in 502ms
starting...
ended in 501ms
starting...
ended in 502ms
starting...
ended in 502ms
res_all = [1, 1, 1, 1, 2, 2, 2, 2]

More Callers

The pool is not processing tasks concurrently or in parallel at the moment. Each call has to wait for the response in the RPC model. This is when the “bridge” is needed. Obviously, we need to parallelize the callers:

#[tokio::main]
async fn main() {
    let pool = MyActorPool::new(4);

    let mut join_set = tokio::task::JoinSet::new();

    let mut res_all = vec![];
    for i in 0..N_TASKS {
        let pool = pool.clone();
        join_set.spawn(async move {
            let t = tokio::time::Instant::now();
            println!("{i} starting...");
            let res = pool.get_unique_id().await;
            println!("{i} ended in {}ms", t.elapsed().as_millis());
            res
        });
    }

    while let Some(res) = join_set.join_next().await {
        res_all.push(res.unwrap());
    }
    println!("res_all = {:?}", res_all);
}

This spawns 8 callers in a JoinSet, so that we could harvest the results. This time, the first 4 calls get executed concurrently/in parallel. So do the last 4 calls, but they have to wait until the pool has more capacity.

starting...
starting...
starting...
starting...
starting...
starting...
starting...
starting...
ended in 502ms
ended in 502ms
ended in 502ms
ended in 502ms
ended in 1004ms
ended in 1004ms
ended in 1004ms
ended in 1004ms
res_all = [1, 1, 1, 1, 2, 2, 2, 2]

Big Practical Concerns

There are two main issues with this simplified caller code.

The source task list could be very long. Spawning caller tasks without bounds and letting them wait on the RPC is bad. Async tasks are cheap, but we don’t want the design to not have a bound. It’s also much harder to exit gracefully this way.
The JoinSet will actually hoard the memory of every task it spawns but not consumed. We need to consume from it concurrently/in parallel.

The first one can be resolved by a semaphore. The second one requires a bigger change - we need to spawn a separate task to process the results. As it turns out, we can’t use JoinSet for this, and mpsc channels are the ideal pattern. The task submitter holds the Sender part of the channel, and the task receiver holds the Receiver part. The receiver will break out of its loop when all Senders are dropped by the individual submitted tasks. It is up to you to either put the submitter or the receiver in a background task. I prefer the former, because mentally I am more interested in the results. It also makes graceful shutdown more natural as the receiver loop ends.

#[tokio::main]
async fn main() {
    // create pool
    let pool = MyActorPool::new(4);

    // submitter & receiver comm
    let (sender, mut receiver) = tokio::sync::mpsc::channel(4);

    // submit tasks
    tokio::spawn(async move {
        // concurrency control
        let sem = std::sync::Arc::new(tokio::sync::Semaphore::new(4));

        for _i in 0..N_TASKS {
            let permit = sem.clone().acquire_owned().await.unwrap();

            let pool = pool.clone();
            let sender = sender.clone();
            tokio::spawn(async move {
                let _permit = permit; // own the permit

                let t = tokio::time::Instant::now();
                println!("{_i} starting...");
                let res = pool.get_unique_id().await;
                println!("{_i} ended in {}ms", t.elapsed().as_millis());
                sender.send(res).await.unwrap();
            });
        }
    });

    // wait for all tasks to finish
    let mut res_all = vec![];
    while let Some(res) = receiver.recv().await {
        res_all.push(res);
    }
    println!("res_all = {:?}", res_all);
}

starting...
starting...
starting...
starting...
ended in 502ms
ended in 502ms
ended in 502ms
ended in 502ms
starting...
starting...
starting...
starting...
ended in 502ms
ended in 502ms
ended in 502ms
ended in 502ms
res_all = [1, 1, 1, 1, 2, 2, 2, 2]

If you don’t care about the results, but just would like to process through tasks where internally work is done. You could use a TaskTracker to avoid the extra complexity of spawning a separate task.

Full Program

This pattern will take you quite far in practice. You could wire different pools owning different resources, adjust pool sizes, and the main entry point could blast through millions and billions of tasks. You could intercept the SIGINT, and stop sending tasks, and wait for existing tasks to finish gracefully to resume later. You could use a progress bar crate to track progress.

This is a simple CTRL+C handler:

use std::sync::{
    atomic::{AtomicBool, Ordering},
    Arc,
};

pub struct InterruptIndicator {
    state: Arc<AtomicBool>,
}

impl InterruptIndicator {
    pub fn new() -> Self {
        let state = Arc::new(AtomicBool::new(false));
        let state_ = state.clone();
        tokio::spawn(async move {
            tokio::signal::ctrl_c()
                .await
                .expect("failed to install CTRL+C handler");
            state_.store(true, Ordering::Relaxed);
        });
        Self { state }
    }

    pub fn is_set(&self) -> bool {
        self.state.load(Ordering::Relaxed)
    }
}

This is the full main program with a progress bar from Indicatif:

use indicatif::{ProgressBar, ProgressStyle};
use pools_and_pipeline::my_actor_pool::MyActorPool;
use pools_and_pipeline::utils::InterruptIndicator;
use std::sync::Arc;
use tokio::sync::Semaphore;
use tokio::task::JoinSet;

const N_TASKS: u64 = 8;

#[tokio::main]
async fn main() {
    // attach CTRL+C handler
    let interrupt_indicator = InterruptIndicator::new();

    // create pool
    let pool = MyActorPool::new(4);

    // submitter & receiver comm
    let (sender, mut receiver) = tokio::sync::mpsc::channel(4);

    // submit tasks
    tokio::spawn(async move {
        // concurrency control
        let sem = Arc::new(Semaphore::new(4));

        for _i in 0..N_TASKS {
            if interrupt_indicator.is_set() {
                println!("Interrupted! Exiting gracefully...");
                break;
            }

            let permit = sem.clone().acquire_owned().await.unwrap();

            let pool = pool.clone();
            let sender = sender.clone();
            tokio::spawn(async move {
                let _permit = permit; // own the permit

                // let t = tokio::time::Instant::now();
                // println!("{_i} starting...");
                let res = pool.get_unique_id().await;
                // println!("{_i} ended in {}ms", t.elapsed().as_millis());
                sender.send(res).await.unwrap();
            });
        }
    });

    // wait for all tasks to finish

    // make pretty progress bar
    let pb = ProgressBar::new(N_TASKS);
    let sty = ProgressStyle::with_template(
        "{spinner:.cyan} [{bar:40.cyan/blue}] {pos:>7}/{len:7} [{elapsed_precise}<{eta_precise} {per_sec:.green}] {msg}"
    ).unwrap().progress_chars("#>-");
    pb.set_style(sty);

    let mut res_all = vec![];
    while let Some(res) = receiver.recv().await {
        res_all.push(res);
        pb.inc(1);
    }
    pb.finish();
    println!("res_all = {:?}", res_all);
}

^CInterrupted! Exiting gracefully...
⠚ [####################>-------------------]       4/8       [00:00:00<00:00:02 1.9913/s]
  [########################################]       8/8       [00:00:01<00:00:00 7.9635/s]
res_all = [1, 1, 1, 1, 2]

You probably want to use Anyhow to handle errors and convert the unwrap() into ? propagation for good measure.

To be Continued… The Streaming Model

Earlier we discussed that, if the multi-worker pool behaves like a processing queue instead of a server, the bridge to parallelism was implied and provided. I will spec out that design in a separate post - (Part II).

TB/PB-Scale Data Loading for Deep Learning

2023-09-29T00:00:00+00:00

I’ve been looking into training some big datasets (500GBs ~ several TBs) of small images with PyTorch. Previously, I have been using a really old and inflexible, but quite efficient format, MXNet’s MXIndexedRecordIO. But there are more and more problems with it. As datasets get larger, it takes a long time to pull to disk from wherever the data is stored, and during distributed training, every rank has to have a copy on disk. Indexed random access sucks a lot of IOPs from disks too, so you would need good disks. It takes a lot of conditions to be met to actually start training.

The industry moved fast to find its own solutions. Tensorflow users would use TFRecord. There is WebDataset, which is really just using tar files to store data in a certain way, and other libraries that support these formats. Essentially one would need to be able to stream the data, preferably in shards from the cloud, and train with window-shuffling as new chunks/shards are being downloaded.

Support for WebDataset has been slowly growing.

The original author’s implementation is pretty good, but there are some subtle areas that might trip users. The documentation has outdated parts, interweaved with up-to-date parts. But if you work through them, it’s a good solution.
TorchData, a PyTorch affiliate has full support over it. However recently it announced it had paused development because it needed to reevaluate the long-term vision.
Ray Data has an implementation and it was suggested to me during the recent Ray Summit. Although you would likely need to use the whole platform for it.
NVIDIA DALI supports the basic use of it, but apparently only loading from disk so far. One could however create an external source in Python. The advantage of DALI is doing image transforms/augmentations on GPUs in CUDA streams, taking the load off CPUs. Although usually CPUs are sufficient at simple-ish augs.

I am currently switching between the original author’s impl and the TorchData impl. They work reasonably well for my use case. The challenge has been handling interaction between distributed/multiprocessing and shard splitting. There are many GitHub issues around it, and the documentations are not written for users who just want to shuffle and feed data correctly. Essentially, in the WebDataset scheme, the data shards are treated as coarse indexes, and it requires some finessing to have all ranks receive the same number of batches to synchronize, and not lose some of the data. These two libraries evolved with each other’s influence, and made a lot of genius use of iterators and all kinds of iterator functions. It’s kind of a nod to Rust’s wide support for iterators.

Another solution for small files and images is Parquet, which also has broad support. The challenge is it’s such a flexible format, so one has to read within fine print to see what the library handles, and what you need to handle. Luckily, with data loading during training, the bottleneck is usually the GPU compute part, and even if ad-hoc custom plugins are slowish, it’s still not a big problem.

Just jotting down notes in case others are wondering.

A Week of PyO3 + rust-numpy (How to Speed Up Your Data Pipeline X Times)

2023-06-06T00:00:00+00:00

Updated Oct 7, 2024: use unsafe but correct and faster array content copy to go from 3x to 4x.

You can check out the code here.

If you are like me (and many others), you’d need a strong reason to learn a new programming language (my journey). It’s a big commitment, and requires a lot of days and nights to do it right. Even for languages like Python that boast simplicity, under the hood there is a lot going on. For many developers, Python is the non-negotiable gluing/orchestrating layer that sits closest to them, because it frees them from the distractions that are not part of the main business logic, and has evolved to become the central language for ML/AI.

Rust on the other hand, requires a lot from a developer up front. Beyond a “hello world” toy example, it gets complex quickly, presenting a steep learning curve. It’s an amazing modern systems language, extremely fast and portable. However, my main programming activities have been in ML and data pipelines, operating on the Python (high) level, so I haven’t seriously tried to learn it.

Python’s ML and data story mostly revolves around the Python numeric ecosystem, which really took off when NumPy brought in the array interface and advanced math to become the “open-source researchers’ MATLAB”, that eventually kicked off almost 20 years of ML/AI development. Following that script, there emerged many Python-based workflows and packages that benefited from faster compiled languages as extension modules. They could be exploratory routines in scientific computing (physics, graphics, data analytics) that needed to be flexible yet efficient. They could be deep learning frameworks. They could also be distributed pipelines that ingest & transform large amounts of data, or web servers with heavy computation demands. The interoperating layer was fulfilled by Cython, SWIG, and Boost.Python. pybind11 grew as a successor to Boost.Python (different authors!) to offer C++ integration, and got good traction in late 2010.

On the Rust side, PyO3 has been getting a lot of love. Developers love Rust’s safety guarantees, modern features, and excellent ecosystem, and have been leveraging ndarray, rust-numpy to interoperate with NumPy arrays from Python to speed up performance-critical sections of their code. This has tremendous appeal to me, and has granted me an overwhelming reason to learn Rust with the PyO3 + rust-numpy stack. Let this be my own “command line program” (from the Book) example. It wasn’t easy to get started this way… Took me through confusion, frustration, and finally, enlightenment in a short span of days. I hope this post can help you get started with your own journey.

Before pulling up the sleeves, let’s peek into Rust and PyO3’s ecosystem. PyO3 has great docs, which is much appreciated. I benefited a lot from the Articles section, reading about other developers’ journeys¹ ² ³. (Note: this article also joined the list!)

The GIL and NumPy

It’s worth mentioning that you can already use NumPy for most of the heavy lifting. It is a fast compiled extension module, and already has the ability to release the GIL (global interpreter lock) on many occasions. Releasing the GIL is important for compute-intensive tasks because it enables Python threads to truly utilize multiple CPU cores. It is a big deal with today’s multicore CPU architecture (goodbye Moore’s Law). For example, if you run the code below in enough loops, and check CPU usage, multiple CPUs will be engaged.

import os
# need to set OMP_NUM_THREADS=1 to avoid MKL/OpenBlas using internal multithreading to confound the measurement
# also much faster overall due to less thread contention (only python threads)
os.environ["OMP_NUM_THREADS"] = "1"

from concurrent.futures import ThreadPoolExecutor
import numpy as np
from tqdm import tqdm

def func():
    for _ in range(10):
        return np.dot(np.random.rand(1000, 1000), np.random.rand(1000, 1000))


if __name__ == "__main__":
    N_tasks = 32
    N_workers = 8

    pool = ThreadPoolExecutor(max_workers=N_workers)

    futures = []
    for _ in range(N_tasks):
        futures.append(pool.submit(func))

    for future in tqdm(futures):
        _ = future.result()

So why do you need additional Python extensions? Well, not all operations can release the GIL. And those can be most of the code, even code that interacts with NumPy.

My Data Pipeline Problem

This is what I ran into. I have many msgpack files generated from some ML routine that packs tens of thousands of NumPy arrays into bytes in each, like generated below.

from typing import Iterator
from concurrent.futures import ThreadPoolExecutor
import numpy as np
import msgpack
from tqdm import tqdm
from loguru import logger

SIZE_ARRAY_DIM = 512
COUNT_PER_MSGPACK_UPPERBOUND = 50000


if __name__ == "__main__":
    with open("bytes_vectors.msgpack", "wb") as f:
        for i in range(COUNT_PER_MSGPACK_UPPERBOUND):
            vector = np.random.rand(SIZE_ARRAY_DIM).astype(np.float32)
            msgpack.pack(vector.tobytes(), f)

There are a lot of them sitting somewhere. I need to use this to iterate through them in Python. It’s worth noting in this particular example one could even write this msgpack iterator in Rust, but in reality there was some custom logic that was hard to replicate. So we will stick with Python for the demo’s sake. You could gain a lot of performance by writing this in Rust, especially with many Python threads, because iterating over Python iterators is slow, and bound by the GIL.

def iterate_msgpack(filename):
    with open(filename, "rb") as handle:
        unpacker = msgpack.Unpacker(handle)
        try:
            for item in unpacker:
                yield item

        except Exception as e:
            logger.error(e)

The main logic is to have many thread workers, each processing one msgpack file and returning a large NumPy array with the content to the main thread. The main thread will do something with the arrays from thread workers in sequence. For every iteration of a msgpack file, we copy the vector into the pre-allocated NumPy array. This copy involves individual Python objects, so is bound by the GIL.

def take_iter_py(iterator: Iterator[bytes], np_vectors: np.ndarray) -> int:
    idx = 0
    for bytes_vector in iterator:
        try:
            vector = np.frombuffer(bytes_vector, dtype=np.float32).reshape(SIZE_ARRAY_DIM)
            np_vectors[idx] = vector
        except ValueError:
            print(f"Array size does not match!")
            continue
        idx += 1

    return idx


def process_py(filepath):
    np_vectors = np.empty((COUNT_PER_MSGPACK_UPPERBOUND, SIZE_ARRAY_DIM), dtype=np.float32)
    count = take_iter_py(iterate_msgpack(filepath), np_vectors)
    np_vectors = np_vectors[:count]
    return np_vectors

We use the same ThreadPoolExecutor setup as earlier to get a benchmark. We use 8 thread workers to go through 32 tasks on an M1 Macbook Pro. Each task goes through one msgpack file. Here we simply reuse the same msgpack file.

if __name__ == "__main__":
    filepath = "bytes_vectors.msgpack"

    N_tasks = 32
    N_workers = 8

    pool = ThreadPoolExecutor(max_workers=N_workers)

    logger.info("With Python / thread pool:")
    futures = []
    for _ in range(N_tasks):
        futures.append(pool.submit(process_py, filepath))

    for future in tqdm(futures):
        _ = future.result()

Output:

With Python / thread pool:
32/32 [00:03<00:00,  8.05it/s]

Why Not Process Workers?

A common workaround in Python pipeline processing is to use multiprocessing, or the newer ProcessPoolExecutor. Instead of using threads that are bound by the GIL, one could fork/spawn multiple Python processes to enable true parallelism. Process workers have their own memory space, and only interact with the main process through (de-)serialization of objects. In our case, the input is a sequence of msgpack file paths, and output is a sequence of their NumPy arrays.

from concurrent.futures import ProcessPoolExecutor

# ...

    pool = ProcessPoolExecutor(max_workers=N_workers)
    logger.info("With Python / process pool:")
    futures = []

    for _ in range(N_tasks):
        futures.append(pool.submit(process_py, filepath))

    for future in tqdm(futures):
        _ = future.result()

Output:

With Python / process pool:
32/32 [00:06<00:00,  4.86it/s]

It’s actually slower! Why is that? Usually this trick works quite well, a simple swap to ProcessPoolExecutor to access many CPU cores. Here, it is slower because another source of latency crept in - (de-)serialization of the big NumPy arrays! Each returned array is sized at 50000 * 512 * 4 / 1024**2 = 98MB.

To demonstrate this, let’s construct a simple test, where we don’t do anything but return a newly allocated NumPy array.

def func():
    return np.empty((50000, 512), dtype=np.float32)

# ...

    pool = ProcessPoolExecutor(max_workers=N_workers)

    logger.info("With Python / process pool:")

    futures = []
    for _ in range(N_tasks):
        futures.append(pool.submit(func))

    for future in tqdm(futures):
        _ = future.result()

Output:

With Python / process pool:
32/32 [00:05<00:00,  5.40it/s]

It’s already quite slow, let alone more operations stacked on top. Clearly, going back to thread workers avoids this unnecessary overhead. Try swapping back ThreadPoolExecutor and see it/s go up to the order of hundreds of thousands - we are not doing anything here, so of course it should be fast.

PyO3 + rust-numpy to Rescue

This is the case where compiled extensions can help. PyO3 and rust-numpy provide an excellent interface to deal with anything that might come your way. There are a lot of other great articles I referenced earlier that can get you set up using maturin. Assuming you complete the common steps, let us create a Rust extension package that implements a take_iter() function. What’s really amazing is the interoperability - it can take a Python iterator!

In src/lib.rs, where the main logic lives, we write

use numpy::ndarray::{s, ArrayViewMut1};
use numpy::PyReadwriteArray2;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyIterator};

const SIZE_ARRAY_DIM: usize = 512;
const F32_SIZE: usize = 4;

fn copy_array(src_bytes_vector: &[u8], dst_vector: &mut ArrayViewMut1<f32>) -> Result<(), String> {
    if src_bytes_vector.len() == dst_vector.len() * F32_SIZE {
        // f32 from msgpack

        // copy bytes in f32 le format to dst_vector
        // safe but slower
        // for (dst, src) in dst_vector.iter_mut().zip(
        //     src_bytes_vector
        //         .chunks_exact(F32_SIZE)
        //         .map(|x| f32::from_le_bytes(x.try_into().unwrap())),
        // ) {
        //     *dst = src;
        // }

        // unsafe but correct & faster
        unsafe {
            std::ptr::copy_nonoverlapping(
                src_bytes_vector.as_ptr() as *const f32,
                dst_vector.as_mut_ptr(),
                dst_vector.len(),
            );
        }
        return Ok(());
    } else {
        return Err(format!(
            "Array size is {}, does not match {}!",
            src_bytes_vector.len(),
            dst_vector.len() * F32_SIZE
        ));
    };
}

#[pymodule]
fn rust_ext(m: &Bound<PyModule>) -> PyResult<()> {
    #[pyfn(m)]
    fn take_iter(
        py: Python,
        iter: Bound<PyIterator>,
        mut np_vectors: PyReadwriteArray2<f32>,
    ) -> PyResult<usize> {
        // First collect bytes into a Rust-native vector.
        // We can't release the GIL here because we are dealing with a Python object.
        let mut raw_list: Vec<Vec<u8>> = vec![];
        for item in iter {
            raw_list.push(item?.downcast::<PyBytes>()?.as_bytes().to_vec());
        }

        let mut vectors = np_vectors.as_array_mut();

        // Bytes are read as f32 arrays and copied into the passed in NumPy array.
        // We release the GIL here so other Python threads can run in true parallelism.
        py.allow_threads(|| {
            if raw_list.len() > vectors.dim().0 {
                return Err(PyValueError::new_err("Too many items in iterator!"));
            }
            if vectors.dim().1 != SIZE_ARRAY_DIM {
                return Err(PyValueError::new_err(format!(
                    "2D NumPy array has {} columns, does not match {}!",
                    vectors.dim().1,
                    SIZE_ARRAY_DIM
                )));
            }

            let mut idx = 0;
            for src_bytes_vector in raw_list {
                let mut dst_vector = vectors.slice_mut(s![idx, ..]);
                if let Err(e) = copy_array(&src_bytes_vector, &mut dst_vector) {
                    eprintln!("Error: {}. Skipping...", e);
                    continue;
                }
                idx += 1;
            }
            Ok(idx)
        })
    }

    Ok(())
}

Run maturin develop --release to compile and install as a Python package. And in Python code, write

import rust_ext

# ...

def process_rs(filepath):
    np_vectors = np.empty((COUNT_PER_MSGPACK_UPPERBOUND, SIZE_ARRAY_DIM), dtype=np.float32)
    count = rust_ext.take_iter(iterate_msgpack(filepath), np_vectors)
    np_vectors = np_vectors[:count]
    return np_vectors

# ...

    pool = ThreadPoolExecutor(max_workers=N_workers)

    logger.info("With Rust extension / thread pool:")

    futures = []
    for _ in range(N_tasks):
        futures.append(pool.submit(process_rs, filepath))

    for future in tqdm(futures):
        _ = future.result()

Output:

With Rust extension / thread pool:
32/32 [00:00<00:00, 33.58it/s]

That’s 4x faster than the Python code!

Let’s do a full rerun with a 2x2 combo to take it home:

With Python / process pool:
32/32 [00:05<00:00,  5.45it/s]

With Rust extension / process pool:
32/32 [00:05<00:00,  5.51it/s]

With Python / thread pool:
32/32 [00:03<00:00,  8.21it/s]

With Rust extension / thread pool:
32/32 [00:00<00:00, 33.58it/s]

Closing Thoughts

Before being able to write Python extensions in Rust, I always resorted to workarounds, like using process workers to unlock multicore usage, and paying for their overheads. For this specific task, it used to take me three full days to ingest tens of billions of entries into a database, and I needed to do it many times for experimentation. Now, it can be as short as one single day, a much faster turnaround, leading to more active experimentation loops and bigger findings. I have finally found a powerful alternative, embracing the full potential of performance computing with the mighty combo - Python and Rust. I know for a fact this is going to be the start of something special, and hope you feel the same.

Code Embeddings

2023-04-18T00:00:00+00:00

In this post, I would like to explore the idea of using embedding vectors to represent code snippets, and compute the cosine similarity scores between a few examples. I will compare OpenAI’s text-embedding-ada-002 with two open-source models, SantaCoder and Salesforce CodeGen. The OpenAI model was trained on general text data, but it’s the only embedding model the company currently offers. The other two models were trained on code to generate code, so we would need to do some hacking to get the embedding vectors.

Setup

from itertools import combinations
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

How do we get the embedding vectors from a model? We can use the transformer part of the model to get the hidden states - layer-by-layer outputs of the transformer blocks. Thanks to a conversation with Vlad, I decided to use the last hidden state. The shape of this output is (batch_size, seq_len, hidden_size), so we would need to aggregate over the seq_len dimension to get the “summary”. We could take the mean, or max, or some other way. I tried the max, but the similarity scores looked all close to 1. So here let’s stick with the mean.

It is noteworthy to spell out the variety of choices here. Echoing my previous post, the embeddings from a generative model were not trained to definitively tell things apart.

class EmbeddingModel:
    """
    Wrapper class to get the embedding vector from a model.
    """
    def __init__(self, tokenizer, model):
        self.tokenizer = tokenizer
        self.model = model

    def get_embedding(self, code):
        with torch.no_grad():
            inputs = self.tokenizer(code, return_tensors="pt").to(self.model.device)
            # here we use the transformer part (not the lm head, which is used for generation).
            outputs = self.model.transformer(**inputs)
            # we could use the last_hidden_state, or second last...
            # the shape is (batch_size, seq_len, hidden_size), so we would need to aggregate over the seq_len dimension.
            # we choose to take the mean here.
            return outputs.last_hidden_state.mean(axis=1)[0]


def cosine_similarity(embedding1, embedding2):
    """
    Compute the cosine similarity between two embedding vectors.
    """
    return embedding1.dot(embedding2) / embedding1.norm() / embedding2.norm()

Get some snippets of code to compare.

code_snippets = []

code_snippets.append("""def print_hello_world():
    print("hello world")
""")

code_snippets.append("""def do_something():
    a = 1
    b = a + 1
    return b
""")

code_snippets.append("""def hello_world():
    return "hello world"
""")

OpenAI `text-embedding-ada-002`

Here we just initialize and define one function for later use.

# openai

import openai
openai.api_key = "YOUR_API_KEY"

def get_openai_embedding(text, model="text-embedding-ada-002"):
   text = text.replace("\n", " ")
   return torch.tensor(openai.Embedding.create(input = [text], model=model)['data'][0]['embedding'])

SantaCoder

# bigcode/santacoder

device = "cpu"
tokenizer = AutoTokenizer.from_pretrained("bigcode/santacoder")
model = AutoModelForCausalLM.from_pretrained("bigcode/santacoder", trust_remote_code=True).to(device)
santacoder = EmbeddingModel(tokenizer, model)

Let’s see how the model generates code from a prompt.

code = """def binary_search():
"""
inputs = santacoder.tokenizer(code, return_tensors="pt").to(device)
outputs = santacoder.model.generate(**inputs, max_length=512)
print(santacoder.tokenizer.decode(outputs[0], truncate_before_pattern=[r"\n\n^#", "^'''", "\n\n\n"]))
# don't think `truncate_before_pattern` had any effect here

Output:

Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.


def binary_search():
    # 1. Get the input
    n = int(input())
    # 2. Initialize the list
    arr = [0] * n
    # 3. Get the elements
    for i in range(n):
        arr[i] = int(input())
    # 4. Find the index
    index = binary_search_recursive(arr, 0, n - 1)
    # 5. Print the result
    print(index)

def binary_search_recursive(arr, left, right):
    # 1. Base case
    if left > right:
        return -1
    # 2. Get the middle index
    mid = (left + right) // 2
    # 3. Check if the element is present at the middle index
    if arr[mid] == mid:
        return mid
    # 4. Check if the element is smaller than the middle index
    elif arr[mid] < mid:
        return binary_search_recursive(arr, left, mid - 1)
    # 5. Check if the element is greater than the middle index
    else:
        return binary_search_recursive(arr, mid + 1, right)

if __name__ == '__main__':
    binary_search()<|endoftext|>...

The truncate_before_pattern argument didn’t seem to take effect here and I saw a lot more code after <|endoftext|>. Doesn’t affect this experiment, so I’ll leave it as is.

What does a SantaCoder embedding vector look like?

embedding = santacoder.get_embedding("def hello_world(): return 'hello world'")
print(embedding.shape)
print(embedding)

Output:

torch.Size([2048])
tensor([-0.2162,  0.7902, -1.0303,  ...,  1.9904, -1.1341, -2.7802],
        device='cpu')

Time to compute the cosine similarity scores using the embedding vectors. We denote SIM local as the results from the open-source model, and SIM openai as the results from the OpenAI model.

for code1, code2 in combinations(code_snippets, 2):
    sim_local = cosine_similarity(santacoder.get_embedding(code1), santacoder.get_embedding(code2)).item()
    sim_openai = cosine_similarity(get_openai_embedding(code1), get_openai_embedding(code2)).item()
    print(code1)
    print("----------")
    print(code2)
    print(f"SIM local: {sim_local:.2f} SIM openai: {sim_openai:.2f}")
    print("=================================")

Output:

def print_hello_world():
    print("hello world")

----------
def do_something():
    a = 1
    b = a + 1
    return b

SIM local: 0.79 SIM openai: 0.81
=================================
def print_hello_world():
    print("hello world")

----------
def hello_world():
    return "hello world"

SIM local: 0.90 SIM openai: 0.93
=================================
def do_something():
    a = 1
    b = a + 1
    return b

----------
def hello_world():
    return "hello world"

SIM local: 0.84 SIM openai: 0.82
=================================

You know… They agree quite well. If I don’t want to use the OpenAI model, I can just use SantaCoder.

Salesforce CodeGen

We only use the 2B parameter model for this experiment here.

# salesforce codegen-2B-mono

device = "cpu"
tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-2B-mono")
model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-2B-mono").to(device)
codegen = EmbeddingModel(tokenizer, model)

Same here. Let’s see how the model generates code from a prompt.

code = """def binary_search():
"""
inputs = codegen.tokenizer(code, return_tensors="pt").to(device)
outputs = codegen.model.generate(**inputs, max_length=512)
print(codegen.tokenizer.decode(outputs[0], truncate_before_pattern=[r"\n\n^#", "^'''", "\n\n\n"]))

Output:

Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.


def binary_search():
    print("Binary Search")
    arr = list(map(int, input("Enter the array: ").split()))
    key = int(input("Enter the key: "))
    low = 0
    high = len(arr) - 1
    while low <= high:
        mid = (low + high) // 2
        if arr[mid] == key:
            print("Element found at index: ", mid)
            break
        elif arr[mid] < key:
            low = mid + 1
        else:
            high = mid - 1
    else:
        print("Element not found")

Examine the CodeGen embedding vector.

embedding = codegen.get_embedding("def hello_world(): return 'hello world'")
print(embedding.shape)
print(embedding)

Output:

torch.Size([2560])
tensor([ 1.2350,  1.5245,  1.9772,  ...,  3.1278, -2.5663, -1.3473],
        device='cpu')

Compute the cosine similarity scores. Here SIM local are from the Salesforce model, and SIM openai as the results from the OpenAI model.

for code1, code2 in combinations(code_snippets, 2):
    sim_local = cosine_similarity(codegen.get_embedding(code1), codegen.get_embedding(code2)).item()
    sim_openai = cosine_similarity(get_openai_embedding(code1), get_openai_embedding(code2)).item()
    print(code1)
    print("----------")
    print(code2)
    print(f"SIM local: {sim_local:.2f} SIM openai: {sim_openai:.2f}")
    print("=================================")

Output:

def print_hello_world():
    print("hello world")

----------
def do_something():
    a = 1
    b = a + 1
    return b

SIM local: 0.98 SIM openai: 0.81
=================================
def print_hello_world():
    print("hello world")

----------
def hello_world():
    return "hello world"

SIM local: 0.99 SIM openai: 0.93
=================================
def do_something():
    a = 1
    b = a + 1
    return b

----------
def hello_world():
    return "hello world"

SIM local: 0.98 SIM openai: 0.82
=================================

Here CodeGen is giving scores all close to 1. It doesn’t seem to differentiate between different snippets. I’ve tried to use the 6B model, but the results were similar. CodeGen should be a decent model for code generation, but apparently, the embedding vector obtained from its last hidden state is not suitable for code similarity.

You can read more about the topic in this issue and Q&A.

The World of Embedding Vectors

2023-04-17T00:00:00+00:00

I’ve long been curious but hesitant about using embedding vectors generated from pre-trained neural networks with self-supervision.

Coming from a brief but intense computer vision (CV) background, I can still remember the first bunch of deep learning tutorials were MNIST handwritten digits classification and ImageNet classification. Supervised learning with class labels, softmax, and cross-entropy loss was the standard formulation. Arguably it was facial recognition that pushed this formulation to the extreme. With tens of thousands of identities in the training set, telling the difference between them with certainty became difficult with the standard loss function. So instead of using softmax and cross-entropy loss, practitioners resorted to the contrastive loss formulation, such as triplet loss in the FaceNet paper. The intra-class features from this method were grouped together, while inter-class separated afar, exhibiting strong discriminative power.

However, it was not able to make efficient use of the highly parallel architecture of GPUs. Soon after, a series of modified softmax-based loss formulations proved to be equally as effective, and were friendly to the nature of GPUs. One would apply a margin to the positive softmax term to push the intra-class features more tightly grouped and inter-class ones farther apart, leading to a similar discriminative effect to the triplet loss. This culminated in the ArcFace paper. In addition, to prevent the fully connected (FC) layer for classification from becoming too large (scales with N of classes) to fit into the modicum GPU memory, some model parallelism was worked out to get it evenly chunked onto multiple GPUs with just a small amount of inter-device communication. With the abundance of interconnected GPUs, suddenly millions+ of identities were within line of sight. The facial recognition models trained using this approach were able to tell apart the world’s population, past, present, and future.

You could use the embedding vectors from these neural networks to compute cosine similarity scores to find someone in a vector database, and they would do a fantastic job. The networks were trained with definitive ground-truth labels for a single task: a face either belongs to a certain unique person or not; there is no middle ground. However, not all vectors are born the same. In contrast, for general vision or language tasks, an object could belong to multiple conceptual categories: a “chihuahua” is a “dog”, also an “animal”, oftentimes a “pet”. This gives ambiguity to the concept of similarity. Depending on the training process, be it supervised with explicit class labels, or self-supervised with contrastive pre-training, you would end up with entirely different statistical distributions in the high-dimensional space. Without domain-specific fine-tuning, these embeddings are simply not trained with the explicit and singular purpose of telling things together or apart. Expecting a single number, such as cosine similarity score, to capture the nuances of relations doesn’t seem to make too much sense.

With the progress in large language models (LLMs) that make me feel like a dinosaur every few months, the field of information retrieval and search relevance got a new player in town - embedding vectors. Recently it’s all the hype because it is supposed to be the cure for LLMs’ hallucination problem. By embedding text into vectors and finding top-K similar ones, one limits the scope of LLMs’ mental space, leading to better fidelity. That’s fine. However, what does not get often talked about is in what granularity one should create embeddings, sentences, paragraphs, or articles. Neither too short nor too long is effective. Also, out-of-domain usage is not competitive against the decade-old BM25 baseline. Moreover, is K-nearest neighbors (KNN) based solely on similarity scores the right method? People even find support vector machine (SVM) to deliver more meaningful search results than KNN. You shouldn’t use embedding vectors without going through these thoughts. That’s the point I’m trying to make.

How does this apply to programming code? As a subset of text, it has unique properties and applications. Being able to ask computers to code on their own is a cornerstone for scaling up intelligence. Code is more logical and follows stricter syntactical rules. It’s easier in some ways, and harder in others. We’ve seen the success of GitHub Copilot, so in terms of generative tasks, completing the function, and adding docs, it does a decent job. How well does it embed all of that information into vectors for similarity comparison? I had the chance to look into this the past few days. Will do more and write about it. Please stay tuned!

A Week of Rust

2023-03-21T00:00:00+00:00

The stars are aligned. I finally sat down and went through “the book.”

I tweaked the color combo to “Rust” and went through the first eight chapters slowly and carefully, ending in Common Collections and, in addition, Chapter 15. Smart Pointers. Typically learning to program is best to be hands-on for me, coupled with many random detours. However, I found the structure and flow of this book very helpful, explaining the necessary details while staying in context, and I enjoyed following its lead. I would attribute this partly to the writing style, with focused code demonstrations iteratively simulating the debugging process one would go through themselves. One would also realize it’d be necessary to sit through the “lecture” sections before preemptively getting hands dirty in an editor with Rust because of the philosophy and design decisions that led to the creation of this language - they need more words to be clear.

This book was written for experienced programmers with a few other languages under their belts. As of today, I think most programmers embarking on learning Rust are from the C++ land. Like C++, Rust was designed as a systems language that aimed to be extremely efficient (zero-cost abstraction). So as a user, you must pay attention to resource management, notably RAII (vs garbage collection in other popular languages), copying vs referencing, etc. But it also aims to solve issues practitioners of C++ have struggled with at scale regarding memory safety (undefined behavior, segmentation fault). This imposes much unexpectedness on programmers at various levels.

I wrote about my programming journey in a separate post. Long story short, I had enough exposure to C++-style resource management and found it instrumental to my learning of Rust. If you come from other backgrounds, your mileage will differ. There was a synthesis of properties from other languages into Rust, and I found myself returning to the C++ mental box for comparisons the most. Here I would like to highlight the parts I found note-worthy.

Eager Moves

In C++, when you assign a variable to another variable, the value is copied. With most standard library copiable types that manage resources, this leads to pretty extensive copies. Smart pointers such as shared_ptr are exceptions. Generally, you would need to explicitly use std::move() to get an rvalue reference (see later) to invoke the C++11 move semantics.

vector<int> v1 = {1, 2, 3};
// this copies the entire vector
vector<int> v2 = v1;

string s1 = "hello";
// this copies the entire string
string s2 = s1;

In Rust, when you assign a variable to another variable, the value is moved, if it does not implement the Copy trait. This is a very different behavior from C++.

let v1 = vec![1, 2, 3];
// this moves the entire vector
let v2 = v1;
println!("v2 = {:?}", v2);
// println!("v1 = {:?}", v1); // error: use of moved value: `v1`

let s1 = String::from("hello");
// this moves the entire string
let s2 = s1;
println!("s2 = {}", s2);
// println!("s1 = {}", s1); // error: use of moved value: `s1`

The benefit is that you won’t unnecessarily do expensive copies. The downside is that you need to be careful about the ownership of the value. Although luckily, the compiler will help you out with this. If you want to copy the value, you need to explicitly use the clone() method.

Syntax for Pointers and References

In C++, you have the concepts of pointers and references.

// This is how you get a pointer:
int* p = &a;

// This is how you dereference a pointer:
int a_copy = *p;

The & in &a and * in *p are operators. The int* is a pointer type. You announce it on both ends.

// This is how you assign by reference:
int& r = a;
// This is how you use a reference:
int a_copy = r;

The reference r is functionally no different from the original a. In C++, references are not a first-class type like in Rust, but a mechanism for passing as function arguments, or assigning to variables. You don’t need operators to get a reference, but just make the reference notation on the receiver variable side (int&). There is more nuance between lvalue and rvalue references that goes beyond the scope of this short post.

Whereas in Rust, references look like the below.

// This is how you get a reference:
let r: &i32 = &a;
// This is how you dereference a reference:
let a_copy: i32 = *r;

The & in &a and * in *r are operators. This feels like pointers in C++. However, the &i32 is a first-class reference type in its type system, a rather unique choice among common programming languages. There are a lot more rules on auto-dereferencing that go beyond the superficial resemblance.

Pattern Matching

I like the powerful Enums and pattern matching in Rust. It’s very neat.

enum Color {
    Red(u8), // u8 is the type of the value, indicating the intensity
    Green(u8),
    Blue(u8),
}

let c = Color::Red(255);
match c {
    Color::Red(v) => println!("red {}", v),
    Color::Green(v) => println!("green {}", v),
    Color::Blue(v) => println!("blue {}", v),
}

Borrowing (and Borrow Checker)

This is the most important/novel concept in Rust, but I would just redirect you to read the book. It gets more intuitive as you see it everywhere, and its occurrences logically make sense. Notably, in C++, you don’t have variations of self in the context of an object’s methods, but in Rust, you do because &self/&mut self are explicitly invoked in the method signatures and follow the same rules.

Conclusion

That’s it so far. There might be incorrect wordings, as these are big languages, and I try to write this article as a learner, not an expert. It is a big topic to weigh the pros and cons, comparing languages people love and care about. If you open Stack Overflow or Reddit, you will find heated discussions and debates.

I will continue with the book and write as I learn more. Hope this article helps you in your journey.

My Programming Journey

2023-03-20T00:00:00+00:00

I originally intended to leave this as a short paragraph, in the middle of a thought piece after learning Rust for a week. But the abbreviated turned fully expanded. So here it is.

This is my personal programming journey. I learned a tiny bit of Pascal in high school around 2005, but focused more on playing Counter-Strike and making maps (more productive than being just a gamer?). We were required to learn C in college, but I hated it because exams were intentionally trying to trick us into memorizing syntax quirks. I started seriously learning Python (with A Byte of Python) and Linux in 2012 in grad school for scientific computing needs. It was also around that time when the Python scientific ecosystem really started to flourish, with data science and machine learning on the rise, attracting a lot of cash in the industry. Python is a very intuitive language with plain English-like syntax (progressive disclosure of complexity) and easy-to-expect behaviors. It abstracts much complexity away from novice users, so they won’t have to worry about resource management and automatically get memory-safety guarantees. Its standard library and ecosystem are vast, and you can almost always find something other people created and start using them for your needs after simple installation steps. It accomplishes these tasks at the cost of performance - it is dynamically typed with an interpreter, garbage collected, and comes with a global interpreter lock (GIL for the dominant CPython at least). Many workarounds include interfacing with a lower-level language (C/C++) to get intrinsic speed and multithreading without the GIL. With enough care, you can usually get around the bottlenecks at hand and find new bottlenecks to be something else in a compute-centric system. Python is a T-shaped language.

These properties ushered Python all the way into the era of AI and made it the de-facto language for deep learning. You even see this pattern being played out in different frameworks built within Python. What made PyTorch the dominant framework for AI researchers was that it prioritized intuitiveness and flexibility over performance with eager execution. The framework maintainers then optimized the performance with lower-level graph and compiler techniques. These considerations (abstraction of complexity and division of labor) of Python and PyTorch enabled deep learning researchers who weren’t fluent with less intuitive languages & frameworks, and liberated those who wouldn’t want to be slowed down by them while engaging in domain problems.

Before getting ahead of myself, during grad school years I also learned Java from a book I borrowed from my roommate. It was also the language Algorithm & Data Structure courses typically used to teach, so resources were abundant. The garbage collection and pass-by-assignment were similar enough to Python, so it was just adding static typing and JVM to my mental programming model. However, because I was not actively trying to build commercial software, the Object Oriented verbosity deterred me from moving at full speed.

After graduating from grad school in the summer of 2017, I had a brief learning phase with Swift at a coffee shop. However, I wasn’t much of a fan of mobile programming, so I mostly just read the guide across the span of a few days and focused on comparing it to what I knew best, Python. What stood out to me was that for value types (e.g. structs) in function signatures, you would copy the entire object instead of getting a pointer-like variable for reference types (e.g. classes). In Python or Java, this copying behavior only applies to “primitive” types. The optional unwrapping had some funky syntax, and avoiding reference cycles due to using automatic reference counting (ARC) was worth paying attention to. Other than that, I never had to think about resource management or pointers (of course, they exist in the “unsafe” section).

That summer, I accepted a job offer at Bloomberg. On the job, I took C++ seriously because system programming with OpenSSH and PAM modules demanded it. Of course, you could do it in C, but why not C++ if you could use more modern syntax and libraries? It was a nurturing environment for C++ because NYC finance shops were all about correctness and extracting performance, so I took the time to learn about resource management and caveats about to-dos and not-to-dos. Not having a standard package management system was definitely a deterrent, so I ended up installing a lot of stuff into Docker containers’ system paths and checking out code ad-hoc from GitHub. Reading C++ documentation hurt my head, but I had done more persistent reading into lengthy manuals in grad school, so I persevered. It was relatively easy to go along with an existing project and contribute to it by osmosis, but a lot harder to take full advantage of progressively more C++11/14/17/20 features on top of existing complexities. I could sense they seemed necessary though, because only by using the newer syntaxes could you express more naturally what you would say in Python or Swift. I got sucked into the world of smart pointers, rvalue references, and move semantics. However, it was such a big language with a lot of existing code, and there were just too many ways of doing the same thing with slightly different trade-offs. But when it complied successfully and ran without problems, it felt all the pain was worth it.

Since my job was on backend systems and platforms, something was missing between flexible but slow Python and efficient but complex C++. These network-connected backends and platforms relied heavily on Web API-driven designs and multithreaded execution. A language was practically born for the job - Go. It is intentionally simple, with an amazing package manager and documentation hub. It is very fast and effortlessly supports concurrency/parallelism with goroutines. All those languages I had learned then followed a path towards more complexity, from Python to Java, to Swift, and to C++. But Go does not adhere to the popular Object Oriented norm, preferring interfaces. You also get ease of use with garbage collection, easy-to-return slices and maps without excessive copies. Choosing between Python and Go for a web server boils down to whether you want more features from Python’s vast ecosystem or high concurrent throughput demand and low memory usage.

I kept using Python and C++, occasionally Go for a few more years. And the magic soup started turning - there is something to say about the mental maturing process as a programmer. You start as a novice, surprised and daunted all the time. You keep learning and become mid-level decent at the craft, having seen just enough to say, “everything is pretty similar.” Then your experience deepens, and everything is nuanced and complex again - different languages reflected the times when they were born, the problems they aimed to solve, and they dabbled into an ingenious pool of building blocks, some common, some less so.

Now is the time for some Rust.

Training Deep Learning Models

2022-12-20T00:00:00+00:00

I recently came across an article by Andrej Karpathy, ex-Sr. Director of AI at Tesla. Besides being impressed by the content, I was almost brought to tears. So much of it was what I personally experienced, learned the hard way, came to believe in, and decided to teach others. It felt vindicating to hear him say

…a “fast and furious” approach to training neural networks does not work and only leads to suffering… the qualities that in my experience correlate most strongly to success in deep learning are patience and attention to detail.

because I kept saying this but this point and its implications weren’t going through people’s minds.

Here I offer some of my own rules of training deep learning models. It might end up being a growing list.

There are typically open source code and data to begin with. Replicate their results closely, then multiple times with slightly different configs to explore the param space.
Perfect the training code with each iteration, making logging, charting, resuming extremely easy.
For lack of better tooling for tracking experiments in dynamic ways, create a worklog, form a mental map of the things done, and be sensitive to the file sizes, timestamps, git diff outputs in the commandline to enrich your own working memory. The ability to hold a lot of trial and error outcomes in your head provides chances for breakthroughs. It’s like your brain CPU’s fast-access memory.
Bring new/more data in. Do a lot of experiments at a small scale with big param changes, from data sampling strategy, preprocessing, learning rate schedule, number of epochs to hone in onto a good baseline. Figure out which params are more sensitive to specific datasets, and which are less likely to drift.
Scale up 10x, and change the params that are more sensitive. Observe if experimental results agree with your hypotheses. Also observe if at this scale the trends and rules observed at the smaller scale still hold. Make new hypotheses about scaling to the next level.
In the meantime spend days drilling into relevant parts of the deep learning framework being used. Block out time to go deep as they are the most intimate tooling, and mastering them at a deep level will pay dividends before a big endeavor goes to its finishing line.
Gradually scale up to the largest size of data and model architecture one has & can afford. Apply the scaling rules and all the evidence, working intuition and pick the output models that evaluate best on a large and representative test set.
If the test set starts to get saturated, it’s time to take a break and create a more effective test set.

An openmp Parallel Data Access Pattern in Faiss

2022-10-12T00:00:00+00:00

This is a quick note on how to use openmp or rather, any multithreading library to divide the underlying data. Let’s use some code snippets from Faiss.

Usually you would parallel for, or first parallel then for over a sequence. For example:

        int nt = std::min(omp_get_max_threads(), int(n));

#pragma omp parallel for if (nt > 1)
        for (idx_t slice = 0; slice < nt; slice++) {
            IndexIVFStats local_stats;
            idx_t i0 = n * slice / nt;
            idx_t i1 = n * (slice + 1) / nt;

I came across a different use case that was note-worthy:

    DirectMapAdd dm_adder(direct_map, n, xids);

#pragma omp parallel reduction(+ : nadd)
    {
        int nt = omp_get_num_threads();
        int rank = omp_get_thread_num();

        // each thread takes care of a subset of lists
        for (size_t i = 0; i < n; i++) {
            idx_t list_no = coarse_idx[i];
            if (list_no >= 0 && list_no % nt == rank) {
                idx_t id = xids ? xids[i] : ntotal + i;
                size_t ofs = invlists->add_entry(
                        list_no, id, flat_codes.get() + i * code_size);

                dm_adder.add(i, list_no, ofs);

                nadd++;
            } else if (rank == 0 && list_no == -1) {
                dm_adder.add(i, -1, 0);
            }
        }
    }

Here you just use parallel but not for. Instead, you get the total N of threads nt, and thread rank/number rank. Every single thread enters the entire for loop, iterating over each element. However, there is a conditional, which guarantees each element is processed at most once by one thread alone. Here it is done by the modulo list_no % nt. This allows all elements that are identified by a unique list_no to be exclusively handled by one thread, serially without possible race conditions. The beauty of it is that you get to accumulate states for each element’s turn without ever needing a lock on the accumulator (the one that is identified by list_no in this case). It is very efficient for similar scenarios.

I also happened to encounter this pattern in a write-up. Check out the Using the intrinsic parallelism of the parallel_flat_hash_map to insert values from multiple threads, lock free section. BTW, you NEED to check out parallel-hashmap.

Kids! Use hnswlib for HNSW

2022-09-28T00:00:00+00:00

For people who fool around in the small field of Approximate Nearest Neighbors (ANN) search, Faiss and hnswlib are two big names. Faiss is a much broader library with a lot of in-memory ANN methods, vector compression schemes, GPU counterparts, and utility functions, whereas hnswlib is a lot more specific, with only one graph-based in-memory index construction method called Hierarchical Navigable Small Worlds (HNSW)¹ ². After the open-source implementation of HNSW in hnswlib came out, Faiss also attempted it with its IndexHNSW class.

Which to pick? Being a long-time Faiss user, I had the natural inclination to keep using what it offered. However, issues ensued.

So I whipped up a large HNSW index with scalar quantization (SQ) HNSW32,SQ8, 512 dimensions, ~30M vectors, with efConstruction=100. It amounted to a little over 20GB in size. The general setup was a cloud instance with 32 physical cores and 64 hyperthreads and 120GB ~ 590GB memory, but a smaller or larger config would yield similar outcomes. I also experimented with both Intel and AMD, Ubuntu and RHEL-based OSes, so what I saw should be representative. Note the CPUs all supported AVX-512 and less instructions.

Initially, everything was fine, all cores/threads firing 100% (green in htop meaning all user threads), and memory usage stable. As I kept adding new vectors into it, suddenly, around ~25GB in size, ~30M vectors, searches became twice to three times slower. The CPU utilization was less than total thread count, and there was a lot of red in htop indicating kernel threads firing. The resident memory started fluctuating within a 3GB range. The above screenshot depicts the CPU usage issue, red kernel threads dominating the green user threads, which used to be almost 100% when the index was smaller/fewer vectors.

I checked the SSD, not much IO, and there was no swap set up. All were indicating memory access pattern issues. I found this article about Transparent Hugepage Support, that some databases such as MongoDB and Postgres did not like it to be enabled. However, my situation actually improved if I went the opposite way and enabled always to THP, instead of madvise. Speed went up by about 1.5x, but there was still a lot of kernel threads activity, and memory still fluctuating.

If I had to guess, it was probably tweakable through some OS setting related to memory and caching. But how could I best debug and analyze the issue? I was informed by the Faiss author Matthijs Douze, that I should try the original/reference implementation from hnswlib.

Here is a side-by-side comparison using full precision float32 vectors. Both were benchmarked through bindings using Python 3.8. Faiss 1.7.2 was installed from the Anaconda pytorch channel with AVX2 support, and hnswlib 0.6.2 was compiled with native flag enabling AVX-512 support.

In []: import numpy as np
In []: import time
In []: import hnswlib
In []: import faiss

In []: data = np.load("data.npy", mmap_mode="r")


# hnsw add
In []: num_elements = data.shape[0]
In []: ids = np.arange(num_elements)

In []: p = hnswlib.Index(space = 'ip', dim = 512)
In []: p.init_index(max_elements = num_elements, ef_construction = 200, M = 32)

In []: t = time.time(); p.add_items(data, ids); print(time.time() - t)
11217.803444623947

In []: p.max_elements
Out[]: 37949877

In []: p.set_ef(2048)
In []: %timeit p.knn_query(data[:10000], 5)
22.2 s ± 30.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)


# Faiss add
In []: quantizer = faiss.index_factory(512, "HNSW32", faiss.METRIC_INNER_PRODUCT)
In []: quantizer.hnsw.efConstruction = 200

In []: t = time.time(); quantizer.add(xb); print(time.time() - t)
21627.71038222313

In []: quantizer.ntotal
Out[]: 37949877

In []: quantizer.hnsw.efSearch = 2048
In []: %timeit quantizer.search(data[:10000], 5)
2min 42s ± 3.48 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

This was what CPU usage looked like in htop for hnswlib:

And this for Faiss HNSW32:

Compare hnswlib’s 22s with Faiss IndexHNSW’s 2min42s. That was a huge difference.

Why? I thought first off there were confounding factors - CPU advanced instruction sets. The float32 implementation in hnswlib was at an advantage because it fully utilized AVX-512, whereas in Faiss it used AVX at best (here, here, and here). These powerful vectorization techniques would make a lot of difference in numerical computations.

To highlight this factor, I did some single-core benchmarking at the same param settings with float32 vectors. I did not observe kernel thrashing (red in htop) so the difference must have been mostly due to advanced instruction sets. Faiss’s search queries took ~1.72x (quite consistently) those of hnswlib.

In [58]: p.set_num_threads(1)

In [59]: faiss.omp_set_num_threads(1)

In [60]: %timeit p.knn_query(data[:10], 5)
320 ms ± 1.65 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [61]: %timeit quantizer.search(data[:10], 5)
552 ms ± 9.09 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [62]: %timeit p.knn_query(data[:50], 5)
1.69 s ± 3.37 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [63]: %timeit quantizer.search(data[:50], 5)
2.92 s ± 43.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Assuming this factor would stay the same for fully threaded use, the rest of the difference still needed an explanation. In my earliest experiment with Faiss HNSW,SQ8, which made use of AVX2 instructions, the slowness and kernel thrashing behavior were still present. It was when I encountered this problem in the first place.

My strongest hypothesis was that it was due to the way IndexHNSW variants were implemented, with actual vector storage separate from the HNSW graph structure. Memory was supposed to be fast, but somehow the access pattern did not play well with this type of separation. hnswlib did not adopt this pattern, so it was able to operate at peak capacity.

This seems to be a bummer, because it was such an elegant way to implement the graph structure/storage separation. But it’s probably wise to use the original hnswlib implementation for max performance for the time being. Note it’s not impossible to take Faiss’s compression schemes and apply them to hnswlib, although you need to implement convenient training and metadata dumping/loading to make it more flexible.

This issue was reported by my other account. Looking forward to someone fixing the problem.

Thoughts on ML Research

2022-09-27T00:00:00+00:00

Not every business needs a Research function, certainly not every startup. However if a startup’s bread and butter is advanced technology, a sustained effort has to be put into maintaining it and cutting out new paths. The Research function’s role is to tease and trek into the unknown, to distill the craft into our potential area of expertise. In fulfilling this function, it needs to be comfortable not knowing how the piece of technology exactly fits into the product timeline and coordination - if we are confident in specing out even the big strokes at the beginning, it’s not research but engineering execution. Instead, it needs to do the following things.

It needs to identify and appreciate a big challenge, take comfort and joy in the craft itself, and recognize the fact that the problem is meaningful enough that any findings coming out of it could shape key directions for the business.
It needs to formulate hypotheses, set up environments to prove or disprove them quickly in the most representative yet still efficient way, and change course quickly in a matter of days and weeks. The code written and tools built during the many fail-fast attempts do not get thrown away, but factored into common elements and components in a broad collection of Swiss army knives that can be repurposed easily.
It needs to be fixated on the important details, retain a long working memory, turn every stone and record them, and aim to incrementally become an expert on the subject.
It needs to identify milestones particular to the different branching points, systematically approach them in overlapping time frames, and have the will to drive through to the conclusion.

Perhaps the biggest realization is that it is a constant battle against the unlikely, the seemingly impossible, a race against stasis and inertia. You have to move quickly enough to not let the problem drag you down and by working incredibly resiliently, you start to create a rhythm. Once you have a rhythm you know what you can tackle, and this internal knowledge and comfort gauge can go a very long way. In particular, big computers will be your best friend, and wielding a huge amount of computing resources will put new power and responsibility onto your hands. You’ll plan around off hours trying to fill in jobs that run overnight, and still maintain the pace of the project during the 2/3 of the day not on business hours. I’ve observed this habit up to Big Tech companies’ ML Research function, so it’s not something only frugal startups do. It’s the natural rhythm of work in this field.

In this light, research practitioners need to be able to independently carry out scientific experiments, but also need to be excellent low-level engineers, who understand algorithms, primitives, time complexity, various hardware constraints (compute, memory, networking IO, disk IO), and are handy with a few languages that are close to hardware. In our line of business, they need to be good at parallel programming with CPUs and GPUs to be proficient at wielding large-scale datasets. They need to approach tools from First-Principles “what exactly do I want to do, and how does this tool uniquely help me get there when I can’t do it myself”, instead of adopting some opaque established system that people say is good.

There is a distinction between ML Engineering work and Research work, despite sharing a lot of knowledge and assets in common. These are relevant research topic examples. They all involve dealing with a substantial amount of unknowns, and outcomes that inherently will become premium features to the core product instead of “we also have it” that are 90% good. It’s “best or bust”. If we do a good job, customers will choose no one but us, and we can be ahead of competitors by a good margin, and we get to shape the next step.

Algorithms that demonstrate extremely low error rates and bias
Models that require extensive large-scale parallel cleaning, training and testing.
Model acceleration/quantization with speed and accuracy tradeoffs
Applications that need to be fast at every step for real-time use
Applications that need to be vastly scalable that cannot be easily addressed by duplicating servers

These are problems more suited in Engineering, broadly MLOps.

ML model accessor design, code base organization, CI/CD pipelines, package integration, model artifacts, Docker image management
ML model serving engine with different formats catering to different use cases (latency vs throughput), batching, GPU resource management (slicing/timeshare)
ML-centric pipelines, components for data ingestion, user upload processing and delivery
Orchestration of various ML models, algorithms together for a business function, exposure to Web endpoint for API user or front-end consumption
Custom labeling services, online training & labeling pipelines

The MLOps function as an Engineering unit serves the rest of Engineering/Product and in addition, Research. Cases #1, #2, #3 target the shared infra, #4 for Engineering/Product, and #5 for Research.

It has become clear to me the ML Engineering work and Research work require different types of people. It requires years to train them to be who they are not. When one provisions people to fill the Research function, s/he needs to realize they excel by depth, not “speed to a prototype”. S/he needs to think about the business assuming certain research projects are a reality, and how it could greatly benefit the business, and commit backwards to guide the hiring plan. It needs to be the right people, the right specialists. It is extremely hard to motivate people to do deep research except when it is a topic the person is really into, has great relevant experience and preparation for, and is given an abundant amount of ownership, trust and resources to tackle. But, if you get it right, you get magic in return.

Little-Endian Snafu

2022-09-03T00:00:00+00:00

You learned about this in college classes. You thought working SWE jobs in 2022 you’d never have to deal with this. But it comes back to trick you at your worst.

I was tripped by endian-ness when implementing inverted list listno-offset/LO as a fixed-width byte string key in RocksDB.

// When offsets list id + offset are encoded in an uint64
// we call this LO = list-offset

inline uint64_t lo_build(uint64_t list_id, uint64_t offset) {
    return list_id << 32 | offset;
}

inline uint64_t lo_listno(uint64_t lo) {
    return lo >> 32;
}

inline uint64_t lo_offset(uint64_t lo) {
    return lo & 0xffffffff;
}

The original implementation used rather smart bitwise shifting and masking to accommodate two uint32 integers in uint64 allocations into one uint64 integer. It worked out beautifully on its own. I proceeded to set up my custom comparator in RocksDB reinterpreting the uint64 integer LO as a byte string, and that worked out great. Then I discovered prefix seek and slapped on a NewFixedPrefixTransform(4) (4 bytes for the list_id/listno). This was when the iterator seeks broke down. One should expect all elements prefixed by the same listno to be returned with the fixed prefix, but the results suggested otherwise. It somehow seemed to be treating the offset as the prefix.

I spent hours isolating the problem with trial and error. Then at dinner it dawned on me… Little Endian¹ ²!

When memory bytes are printed sequentially from left to right (e.g. in a hex dump), little-endian representation of integers has the significance increasing from left to right. In other words, it appears backwards when visualized, which can be counter-intuitive.

A (most?) modern x86_64 or aarch64 machine adopts a little-endian system. Here, a naive slice of the first 4 bytes in that memory blob would give you the offset, not the listno!

Realizing that, the fix was to switch the list_id and offset in lo_build, or, build LO with string concatenation or memcpy in the expected order. And the problem was solved.

https://en.wikipedia.org/wiki/Endianness ↩
https://chortle.ccsu.edu/assemblytutorial/Chapter-15/ass15_3.html ↩

Terence Z. Liu

The Wonders of RocksDB in Rust (Part III - MapReduce Group-By)

Map Phase

Shuffle Phase

Reduce Phase

The Wonders of RocksDB in Rust (Part II - Set Operations)

A Tale of Two Pointers

Two RocksDBs, One Pass

Run in Parallel

The Wonders of RocksDB in Rust (Part I - Intro)

RocksDB and its Bindings

Basic Patterns

Parallel Scan

Parallel Bulk Load

Compaction After Bulk Load

Notes on People Management

Case Study: turbopuffer ANN v3

Cost at 1k qps

But Wait…

Re-examining the Original Vectors

The 100 qps Target

The Budget Option for 100B

Conclusion

Individual Drive and Organizational Layout in the Age of AI Agents

What do Streaming Databases Solve Uniquely?

Some AWS EC2 Instance Choices

The t Family

The c/m/r Families

Newer Generations

Variations

The x Family

The i Family

g4dn/g5/g6/g6f

p3 (going away)/p4d/p5/p6

g6e/g7e

EBS News

People are Computers

What Makes a Good Technical Leader?

What's a "Senior Engineer"?

Senior Full Stack Engineer

Senior Backend/Platform Engineer

Senior Data & Systems Engineer

Senior ML Engineer (Computer Vision)

Let's Work Better

Productivity Pressure in Management

Migrating 50TB Data From a Home-Grown Database to ScyllaDB, Fast

The Rust Invasion: New Possibilities in the Python Ecosystem

Big Data Engineering in the 2020s - One Big Machine

Use One Big Machine!

Scaling Out

Comparing Value and Reference Semantics across Languages

Stateless Pipeline with Async Stream

Stateless Pipelines

Type Reasoning

What about map()?

What about buffered()/buffer_unordered()?

Stacking up

Extension: Add a Writer Output

Comparison with Rayon in Sync Land

More on Async

Pools and Pipeline with Tokio (Part II - Streaming)

The Streaming Model

Control Flow

Graceful Shutdown

Potential Customizations

Full Program

To Conclude

Pools and Pipeline with Tokio (Part I - RPC)

The RPC Model

Actor Pool

More Callers

Big Practical Concerns

Full Program

To be Continued… The Streaming Model

TB/PB-Scale Data Loading for Deep Learning

A Week of PyO3 + rust-numpy (How to Speed Up Your Data Pipeline X Times)

The GIL and NumPy

My Data Pipeline Problem

Why Not Process Workers?

PyO3 + rust-numpy to Rescue

The `t` Family

The `c`/`m`/`r` Families

The `x` Family

The `i` Family

`g4dn`/`g5`/`g6`/`g6f`

`p3` (going away)/`p4d`/`p5`/`p6`

`g6e`/`g7e`

What about `map()`?

What about `buffered()`/`buffer_unordered()`?

OpenAI `text-embedding-ada-002`