Skip to content

Commit

Permalink
[local-worker-mgr] Use _get_context() for cpu count, too
Browse files Browse the repository at this point in the history
Not strictly necessary, but for consistency.
  • Loading branch information
mtrofin committed Oct 18, 2023
1 parent 25c5254 commit a3edb98
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion compiler_opt/distributed/local/local_worker_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,7 @@ def create_local_worker_pool(worker_cls: 'type[worker.Worker]',
**kwargs) -> worker.FixedWorkerPool:
"""Create a local worker pool for worker_cls."""
if not count:
count = multiprocessing.get_context().cpu_count()
count = _get_context().cpu_count()
final_kwargs = worker.get_full_worker_args(worker_cls, **kwargs)
stubs = [_make_stub(worker_cls, *args, **final_kwargs) for _ in range(count)]
return worker.FixedWorkerPool(workers=stubs, worker_concurrency=16)
Expand Down

0 comments on commit a3edb98

Please sign in to comment.