Fix client max workers

ProbablyFaiz · Sep 3, 2024 · 80e41b7 · 80e41b7
1 parent 1ca8ba6
commit 80e41b7
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/rl/llm/engines/client.py b/rl/llm/engines/client.py
@@ -24,9 +24,6 @@
     from transformers import PreTrainedTokenizer
 
 
-_CLIENT_ENGINE_MAX_WORKERS = int(rl.utils.io.getenv("RL_MAX_WORKERS", 4))
-
-
 class ClientEngine(InferenceEngine, ABC):
     BASE_URL: str
     API_KEY_NAME: str
@@ -37,7 +34,9 @@ def generate(self, prompt: ChatInput) -> InferenceOutput:
 
     def batch_generate(self, prompts: list[ChatInput]) -> InferenceOutput:
         return thread_map(
-            self.generate, prompts, max_workers=_CLIENT_ENGINE_MAX_WORKERS
+            self.generate,
+            prompts,
+            max_workers=int(rl.utils.io.getenv("RL_MAX_WORKERS", 4)),
         )