Skip to content

[RNE Rewrite] Clean up the tokenizer host object: locking, error unwrapping and native default arg #1316

Description

@msluszniak

While reviewing cpp/extensions/nlp/tokenizer.cpp alongside the std::span refactor (#1308 / #1315), three related deviations surfaced. They all sit in the same file and overlap enough that they are probably best resolved together rather than piecemeal. Splitting them out of #1315 kept that PR scoped to the span change.

1. TokenizerHostObject uses std::mutex, while TensorHostObject uses shared_mutex

The three host objects don't split the way the mutex choice implies:

  • TensorHostObject uses std::shared_mutex because it has genuine read-only concurrency: many ops read the same src tensor at once via tryLockShared, and only writers need exclusivity.
  • ModelHostObject uses std::mutex for a real reason: Module::execute() mutates internal ExecuTorch state and is not const, so every method must be exclusive.
  • TokenizerHostObject uses std::mutex — but every one of its operations is const upstream: encode, decode, id_to_piece, piece_to_id and vocab_size are all declared const in pytorch/tokenizers/tokenizer.h and hf_tokenizer.h. Only load() mutates, and that runs in the constructor before the object reaches JS.

Structurally the tokenizer therefore resembles Tensor (read-mostly, with a single exclusive dispose) rather than Model, and looks to have inherited std::mutex from the ModelHostObject skeleton.

This has a user-visible consequence. Because the lock is taken with std::try_to_lock, two concurrent encode calls from different threads do not queue — one throws "encode: Tokenizer is currently in use". A pure read operation fails spuriously. Under a shared_mutex both would succeed, with dispose taking the unique lock.

Caveat that must be settled before acting: const does not imply thread-safety. Grepping hf_tokenizer.h, bpe_tokenizer_base.h, tokenizer.h and the regex headers turns up no mutable state, but third-party/include ships headers only, so the implementation could not be inspected for internal scratch state (regex match contexts are a classic offender — re2 is thread-safe, pcre2 match data is not). Switching to shared_mutex on the strength of const alone risks a data race, so this needs the upstream tokenizers sources checked for reentrancy first.

2. The try-lock + disposed-check boilerplate is duplicated 8×

This block is repeated verbatim (modulo the context string) in 5 places in tokenizer.cppencode, decode, getVocabSize, idToToken, tokenToId:

std::unique_lock<std::mutex> lock(self->mutex_, std::try_to_lock);
if (!lock.owns_lock()) {
    throw jsi::JSError(rt, "encode: Tokenizer is currently in use");
}
if (!self->tokenizer_) {
    throw jsi::JSError(rt, "encode: Tokenizer has been disposed");
}

It is not only a tokenizer problem: ModelHostObject has the same duplication at core/model.cpp:231, :263 and :332. So a tokenizer-local helper would fix 5 of 8 sites and leave the pattern duplicated in core.

The existing tensor::tryLockShared / tryLockUnique (core/tensor_helpers.cpp:12-34) cannot be reused as-is: they are non-template functions fixed on the parameter type (TensorHostObject), the mutex type (std::shared_mutex, written literally in the signature and body), the liveness check (data_) and the error wording ("tensor"). TokenizerHostObject differs on every axis, and its members are private where TensorHostObject's are public.

Resolving this properly means a generic helper in core (templated over the host object and a liveness accessor) shared by tensor, model and tokenizer. That edits core, which core-guidelines says to leave alone — hence a deliberate decision rather than a drive-by. Note this interacts with (1): the right helper signature depends on whether the tokenizer keeps std::mutex or moves to shared_mutex.

Suggested order

  1. Settle the thread-safety question upstream, which decides (1).
  2. Do (2) with the helper shape that (1) implies, covering model.cpp as well.

Metadata

Metadata

Assignees

Type

Projects

No projects

Relationships

None yet

Development

No branches or pull requests

Issue actions