Batch embed workflow across actors

The actor pool is useless if you cannot split work across it correctly. We shard the batch round-robin across actors, fire every shard in parallel, then re-stitch results so the output order matches the input order. Callers never know parallelism happened.

pipeline/embed.py
python
@ray.remote
class _EmbedActor:
    def __init__(self, name: str):
        from sentence_transformers import SentenceTransformer
        self.model = SentenceTransformer(name)

    def encode(self, texts: List[str]) -> List[List[float]]:
        if not texts:
            return []
        vecs = self.model.encode(texts, convert_to_tensor=False, show_progress_bar=False)
        if hasattr(vecs, 'tolist'):
            vecs = vecs.tolist()
        return [list(v) for v in vecs]

self._actors = [_EmbedActor.remote(model_name) for _ in range(self.num_workers)]

Each actor loads SentenceTransformer once in its own process. We spin up RAY_NUM_WORKERS actors at startup and reuse them for every batch.

pipeline/embed.py
python
def embed_batch(self, texts: List[str]) -> List[List[float]]:
    if not texts:
        return []

    if self._ray is None or not self._actors:
        return self._fallback.encode(texts)

    ray = self._ray
    n = len(texts)
    shards: List[List[str]] = [[] for _ in range(self.num_workers)]
    for i, t in enumerate(texts):
        shards[i % self.num_workers].append(t)

    futures = [
        actor.encode.remote(shard)
        for actor, shard in zip(self._actors, shards)
        if shard
    ]
    results = ray.get(futures)

    stitched: List[List[float]] = [None] * n
    cursors = [0] * self.num_workers
    for i in range(n):
        worker = i % self.num_workers
        stitched[i] = results[worker][cursors[worker]]
        cursors[worker] += 1
    return stitched

Round-robin sharding keeps every actor roughly equally loaded. The stitching loop is the inverse: it walks the original index and pulls from the right actor result. Output order matches input order exactly.

Ray actor pool sharding

Round-robin sharding fans a batch across the actor pool, then stitches results back in input order.

Ordering exercise: Order the parallel embed workflow

Loading practice…