ray Actor Pool

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

import ray
from ray.util.actor_pool import ActorPool
from ray.util.accelerators import NVIDIA_TESLA_V100

ray.init(ignore_reinit_error=True, include_dashboard=False)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

NUM_GPU = 4

@ray.remote(num_cpus=1, num_gpus=1, accelerator_type=NVIDIA_TESLA_V100)
class MyWorker:
    def hello(self, payload):
	"""use GPU here, for example, load a model into GPU
	"""
	return 'predicted'

MODEL_POOL = ActorPool([WayspotEmbedding.remote() for _ in range(NUM_GPU)])


for payload in payloads:
    MODEL_POOL.submit(lambda a, v: a.predict.remote(v), row)

results = []
for idx in range(len(payloads)):
    try:
        results.append(MODEL_POOL.get_next())
    except Exception as e:
        pass

print(results)