(
self,
mp,
workers_contexts: List[WorkerContext],
result_queue: ShmQueue,
general_task_queue: Optional[ShmQueue],
callback_pickler,
)
| 332 | """ |
| 333 | |
| 334 | def __init__( |
| 335 | self, |
| 336 | mp, |
| 337 | workers_contexts: List[WorkerContext], |
| 338 | result_queue: ShmQueue, |
| 339 | general_task_queue: Optional[ShmQueue], |
| 340 | callback_pickler, |
| 341 | ): |
| 342 | start_method = mp.get_start_method() |
| 343 | if not workers_contexts: |
| 344 | raise RuntimeError("Cannot start a pool with no workers") |
| 345 | if start_method == "fork" and _b.IsDriverInitialized(): |
| 346 | raise RuntimeError( |
| 347 | "Error when starting Python worker threads for DALI parallel External Source. " |
| 348 | "Cannot fork a process when the CUDA has been initialized in the process. " |
| 349 | "CUDA is initialized during ``Pipeline.build()``, or can be initialized by another" |
| 350 | " library that interacts with CUDA, for example a DL framework creating " |
| 351 | "CUDA tensors. If you are trying to build multiple pipelines that use Python " |
| 352 | "workers, you will need to call ``start_py_workers`` method on all of them before " |
| 353 | "calling ``build`` method of any pipeline to start Python workers before CUDA is " |
| 354 | "initialized by ``build`` or other CUDA operation. Alternatively you can change " |
| 355 | "Python workers starting method from ``fork`` to ``spawn`` " |
| 356 | "(see DALI Pipeline's ``py_start_method`` option for details). " |
| 357 | ) |
| 358 | self._workers_contexts = workers_contexts |
| 359 | self._result_queue = result_queue |
| 360 | self._general_task_queue = general_task_queue |
| 361 | self._observer = None |
| 362 | self._processes = [] |
| 363 | write_sockets = [] |
| 364 | try: |
| 365 | for worker_i, worker_context in enumerate(workers_contexts): |
| 366 | if start_method == "fork": |
| 367 | read_socket = None |
| 368 | else: |
| 369 | read_socket, write_socket = socket.socketpair() |
| 370 | write_sockets.append(write_socket) |
| 371 | process_context = WorkerArgs( |
| 372 | worker_id=worker_i, |
| 373 | start_method=start_method, |
| 374 | source_descs=worker_context.source_descs, |
| 375 | shm_chunks=worker_context.shm_chunks, |
| 376 | general_task_queue=general_task_queue, |
| 377 | dedicated_task_queue=worker_context.dedicated_task_queue, |
| 378 | result_queue=result_queue, |
| 379 | setup_socket=read_socket, |
| 380 | callback_pickler=callback_pickler, |
| 381 | ) |
| 382 | process = mp.Process(target=worker, args=(process_context,)) |
| 383 | self._processes.append(process) |
| 384 | self._start_processes(mp, start_method, write_sockets) |
| 385 | finally: |
| 386 | for sock in write_sockets: |
| 387 | sock.shutdown(socket.SHUT_RDWR) |
| 388 | sock.close() |
| 389 | |
| 390 | @classmethod |
| 391 | def from_contexts( |
nothing calls this directly
no test coverage detected