adding context parameter for pool with a warning for not being supported (#6776)

This commit is contained in:
Ameer Haj Ali 2020-01-17 14:57:18 -08:00 committed by Edward Oakes
parent a3a268435f
commit 9f9c3f5026
2 changed files with 14 additions and 2 deletions

View file

@ -40,6 +40,9 @@ instructions to run on a multi-node Ray cluster instead.
The full ``multiprocessing.Pool`` API is currently supported. Please see the
`multiprocessing documentation`_ for details.
.. warning::
The ``context`` argument in the ``Pool`` constructor is ignored when using Ray.
.. _`multiprocessing documentation`: https://docs.python.org/3/library/multiprocessing.html#module-multiprocessing.pool
Run on a Cluster

View file

@ -1,3 +1,4 @@
import logging
from multiprocessing import TimeoutError
import os
import time
@ -9,6 +10,8 @@ import copy
import ray
logger = logging.getLogger(__name__)
RAY_ADDRESS_ENV = "RAY_ADDRESS"
@ -319,6 +322,7 @@ class Pool:
initializer=None,
initargs=None,
maxtasksperchild=None,
context=None,
ray_address=None):
self._closed = False
self._initializer = initializer
@ -326,6 +330,11 @@ class Pool:
self._maxtasksperchild = maxtasksperchild or -1
self._actor_deletion_ids = []
if context:
logger.warning("The 'context' argument is not supported using "
"ray. Please refer to the documentation for how "
"to control ray initialization.")
processes = self._init_ray(processes, ray_address)
self._start_actor_pool(processes)
@ -339,12 +348,12 @@ class Pool:
# Cluster mode.
if ray_address is not None:
print("Connecting to ray cluster at address='{}'".format(
logger.info("Connecting to ray cluster at address='{}'".format(
ray_address))
ray.init(address=ray_address)
# Local mode.
else:
print("Starting local ray cluster")
logger.info("Starting local ray cluster")
ray.init(num_cpus=processes)
ray_cpus = int(ray.state.cluster_resources()["CPU"])