Ray: How to run many actors on one GPU?

Question

I have only one gpu, and I want to run many actors on that gpu. Here's what I do using ray, following https://ray.readthedocs.io/en/latest/actors.html

1. first define the network on gpu

class Network():
    def __init__(self, ***some args here***):
        self._graph = tf.Graph()
        os.environ['CUDA_VISIBLE_DIVICES'] = ','.join([str(i) for i in ray.get_gpu_ids()])
        with self._graph.as_default():
            with tf.device('/gpu:0'):
                # network, loss, and optimizer are defined here

        sess_config = tf.ConfigProto(allow_soft_placement=True)
        sess_config.gpu_options.allow_growth=True
        self.sess = tf.Session(graph=self._graph, config=sess_config)
        self.sess.run(tf.global_variables_initializer())
        atexit.register(self.sess.close)

        self.variables = ray.experimental.TensorFlowVariables(self.loss, self.sess)

2. then define the worker class

@ray.remote(num_gpus=1)
class Worker(Network):
    # do something

3. define the learner class

@ray.remote(num_gpus=1)
class Learner(Network):
    # do something

4. train function

def train():
    ray.init(num_gpus=1)
    leaner = Learner.remote(...)
    workers = [Worker.remote(...) for i in range(10)]
    # do something

This process works fine when I don't try to make it work on gpu. That is, it works fine when I remove all with tf.device('/gpu:0') and (num_gpus=1). The trouble arises when I keep them: It seems that only learner is created, but none of the workers is constructed. What should I do to make it work?

Answer 1

When you define an actor class using the decorator @ray.remote(num_gpus=1), you are saying that any actor created from this class must have one GPU reserved for it for the duration of the actor's lifetime. Since you have only one GPU, you will only be able to create one such actor.

If you want to have multiple actors sharing a single GPU, then you need to specify that each actor requires less than 1 GPU, for example, if you wish to share one GPU among 4 actors, then you can have each actor require 1/4th of a GPU. This can be done by declaring the actor class with

@ray.remote(num_gpus=0.25)

In addition, you need to make sure that each actor actually respects the limits that you are placing on it. For example, if you want declare an actor with @ray.remote(num_gpus=0.25), then you should also make sure that TensorFlow uses at most one quarter of the GPU memory. See the answers to How to prevent tensorflow from allocating the totality of a GPU memory? for example.

Answer 2

In case anyone wants to run Ray on multi-GPU system and parallelly run TensorFlow functionality, one can approach the problem as follows. Say you have 2 GPUs and want to run 16 actors.

1. Ray resources are logical and don’t need to have 1-to-1 mapping with physical resources, therefore one can safely init ray as follows:

ray.init(num_gpus=16)

2. Configure memory growth for each each physical GPU:

physical_devices = tf.config.list_physical_devices('GPU')
for device in physical_devices:
    tf.config.experimental.set_memory_growth(device, True)

3. Define the actors and decorate with ray.remote(num_gpus=1):

@ray.remote(num_gpus=1)
    class Simulator:
        def __init__(self, i, *args, **kwargs):
            # do init

        def simulate(
            self,
        ):
            # check the ray logical GPU assigned, it will be in the range [0, 15]
            gpu_id = ray.get_gpu_ids()[0]
            print("gpu_id:", gpu_id)
            # do the tf magic you want
            # ...

4. Create and run the actors:

# Create actors
    simulators = [Simulator.remote(i, args, kwargs) for i in range(16)]
    
    # Run simulations in parallel
    results = ray.get([s.simulate.remote() for s in simulators])

Both GPUs will be utilized and the actors will run parallelly.

This setup works for me with TensorFlow 2.10.1 and Ray 2.6.1.