Datatoy Logo

Workers avec état

Un modèle ML par worker, chargé une seule fois

.class_task(MyClass, MyClass.method, args) instancie MyClass une seule fois par worker. L'instance persiste entre les items — parfait pour les modèles ML, les connexions BDD, les accumulateurs.

Signature

Pipeline.class_task(class_constructor, class_method, class_args=None,
                    close_method=None, class_kwargs=None, count=1) -> Pipeline[OutPacket]
python
from olympipe import Pipeline
from transformers import pipeline as hf_pipeline

class EmbeddingWorker:
    def __init__(self, model_name: str):
        # Chargé une fois au démarrage du worker
        self.model = hf_pipeline("feature-extraction", model=model_name)

    def embed(self, text: str) -> list:
        return self.model(text)[0][0]  # CLS token

# 2 workers GPU, chacun charge le modèle une fois
embeddings = (
    Pipeline(texts)          # 50 000 textes
    .batch(32)
    .class_task(
        EmbeddingWorker,
        EmbeddingWorker.embed,
        ["sentence-transformers/all-MiniLM-L6-v2"],
        count=2,
    )
    .wait_for_result()
)

⚡ Performance

50 000 textes, embeddings HuggingFace

Séquentiel 2.0min
1 worker GPU 40.0s
2 workers GPU 21.0s
4 workers GPU 11.0s
🚀 10.9× plus rapide

Comment ça fonctionne

Au démarrage du processus worker, le constructeur est appelé avec args. Toutes les invocations suivantes réutilisent la même instance — les poids du modèle restent en RAM (ou VRAM) sans rechargement.

class_kwargs et close_method

Signature complète : .class_task(class_constructor, class_method, class_args=None, close_method=None, class_kwargs=None, count=1). class_args et class_kwargs sont passés au constructeur dans chaque worker. close_method est appelée une fois par worker à l'arrêt, pour fermer une connexion ou libérer un modèle.

python
import os
import time
from olympipe import Pipeline

class Scaler:
    def __init__(self, factor: int, offset: int = 0):
        self.factor, self.offset = factor, offset
        self.seen = 0

    def apply(self, x: int) -> int:
        time.sleep(0.05)
        self.seen += 1
        return x * self.factor + self.offset

    def close(self) -> None:          # appelée une fois par worker à l'arrêt
        print(f"worker {os.getpid()} closed after {self.seen} items")

res = (
    Pipeline(range(6))
    .class_task(
        Scaler,
        Scaler.apply,
        class_args=[10],              # arguments positionnels du constructeur
        class_kwargs={"offset": 1},   # arguments nommés du constructeur
        close_method=Scaler.close,
        count=2,
    )
    .wait_for_result()
)
print(sorted(res))
Sortie
worker 62409 closed after 3 items
worker 62411 closed after 3 items
[1, 11, 21, 31, 41, 51]

Exemples liés