fix: Tokenizer initialization race condition (multiple parallel downloads)

[albertoperdomo2]

Summary

Under high QPS during cold start, multiple concurrent requests triggered duplicate tokenizer downloads due to a check-then-act race condition. It is mentioned in #191 and this caused:

• Parallel downloads for a single model tokenizer
• High request failure rate from file corruption
• High memory usage (for parallel downloads)

The steps I followed to verify the condition with meta-llama/Llama-3.2-1B are:

1. Create venv with requirements

cd services/uds_tokenizer
python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt
pip install psutil  # for memory monitoring

2. Clear Tokenizer cache

rm -rf ~/.cache/huggingface/hub/models--*
rm -rf models/meta-llama

3. Ensure socket directory exists: mkdir -p /tmp/tokenizer

4. Start the gRPC service

python run_grpc_server.py

In a new terminal window, I ran a script to verify the race condition existed.

To solve it, I implemented double-checked locking with per-model locks in, with a fast path (cache hit) with zero locking overhead and a slow path (cache miss) using threading.Lock to ensure only one thread downloads each model.

Test plan

The initial runs of the script (which launches 50 requests at the same time) reported:

============================================================
Memory Summary
============================================================
Initial:  75.7 MB
Peak:     494.5 MB
Final:    298.6 MB
Increase: 223.0 MB

Total test duration: 31.33s

with several errors in the service logs like:

2026-02-04 10:30:15,891 [ERROR] [root] Failed to load tokenizer from /<path>/<to>/llm-d-kv-cache/services/uds_tokenizer/models/meta-llama/Llama-3.2-1B: No such file or directory (os error 2)

2026-02-04 10:30:15,891 [ERROR] [root] Failed to initialize tokenizer for model meta-llama/Llama-3.2-1B: Failed to load tokenizer: No such file or directory (os error 2)

2026-02-04 10:30:20,153 [ERROR] [root] Failed to initialize tokenizer for model meta-llama/Llama-3.2-1B: Failed to load tokenizer: ...

And after the proposed fix:

============================================================
Memory Summary
============================================================
Initial:  75.3 MB
Peak:     173.7 MB
Final:    173.7 MB
Increase: 98.4 MB

Total test duration: 2.16s

Ensuring that the initial memory leak is gone, which directly impacts performance too.

Related issues

• [bug] tokenizer memory-leak on aggressive startup #191

[liu-cong]

does this mean you will never have more than 1 item in acquired_locks given you return immediately after acquiring one?

[albertoperdomo2]

This should be outside of the loop, my bad.

[liu-cong]

don't we need a lock to synchronize access to the _cache_locks dict?

[albertoperdomo2]

You are right, that would be necessary. This is primarily based on my approach to have per model locking, but if we do not need that, we can make it simpler with a single global lock.

[liu-cong]

It looks like we never insert any item to _cache_locks map (relying on defaultdict to get the lock). So this will always be empty?

[albertoperdomo2]

It is a defaultdict() so if you call it and the key does not exist, it will automatically create it on first access.

[vMaroon]

Does this still occur with the new approach of loading tokenizer on module initialization and not after requests come in? Might have missed closing that issue. cc @sagearc

[albertoperdomo2]

I could certainly reproduce the issue (3 days ago).

[github-actions]

Unsigned commits detected! Please sign your commits.

For instructions on how to set up GPG/SSH signing and verify your commits, please see GitHub Documentation.

[sagearc]

@vMaroon Initially, the tokenizer initialization at startup was intended for daulet/tokenizers. I'm unsure of the status for tokenizers initialized within the chat template python file. If I remember correctly, even the python interpreter itself is not safely initialized (it should only happen once) and is currently prone to race conditions within the CGO bindings file.
Soon it won't be an issue.