MCPcopy Create free account
hub / github.com/Open-Bee/DataStudio / generate

Method generate

datastudio/models/mp_openai_api.py:245–307  ·  view source on GitHub ↗

Generate responses using fork-based multiprocessing. Args: messages: List of input messages. **kwargs: Additional generation arguments. Returns: List of response strings, in same order as input.

(self, messages: List[Any], **kwargs)

Source from the content-addressed store, hash-verified

243 return kwargs
244
245 def generate(self, messages: List[Any], **kwargs) -> List[str]:
246 """Generate responses using fork-based multiprocessing.
247
248 Args:
249 messages: List of input messages.
250 **kwargs: Additional generation arguments.
251
252 Returns:
253 List of response strings, in same order as input.
254 """
255 if not messages:
256 return []
257
258 total = len(messages)
259
260 # Pre-encode PIL images to base64 so they are fork-safe and COW-shared
261 if self.logger:
262 self.logger.info(f"Pre-encoding images for {total} messages...")
263 self._pre_encode_images(messages)
264
265 # Store in module-level global for COW sharing after fork
266 model_kwargs = self._build_model_kwargs()
267 model_kwargs.update(kwargs)
268 _init_shared_data(messages, model_kwargs)
269
270 # Compute index ranges for each worker
271 num_workers = min(self.num_workers, total)
272 chunk_size = math.ceil(total / num_workers)
273
274 worker_args = []
275 for worker_id in range(num_workers):
276 start = worker_id * chunk_size
277 end = min(start + chunk_size, total)
278 worker_args.append((worker_id, start, end, self.worker_concurrency))
279
280 if self.logger:
281 self.logger.info(
282 f"Forking {num_workers} workers for {total} messages "
283 f"(~{chunk_size}/worker, concurrency={self.worker_concurrency}/worker)"
284 )
285
286 # Fork workers (safe: no threads/event loops exist yet in main process)
287 ctx = mp.get_context("fork")
288
289 all_results = []
290 try:
291 with ctx.Pool(processes=num_workers) as pool:
292 with tqdm(total=total, desc="MPOpenAIAPI Progress") as pbar:
293 for chunk_results in pool.imap(_worker_fn, worker_args):
294 all_results.extend(chunk_results)
295 pbar.update(len(chunk_results))
296 finally:
297 # Always clear shared data to release memory
298 _clear_shared_data()
299
300 if self.logger:
301 fail_count = sum(1 for r in all_results if r == self.fail_msg)
302 self.logger.info(

Callers

nothing calls this directly

Calls 6

_build_model_kwargsMethod · 0.95
_init_shared_dataFunction · 0.85
_clear_shared_dataFunction · 0.85
infoMethod · 0.80
_pre_encode_imagesMethod · 0.80
updateMethod · 0.80

Tested by

no test coverage detected