Custom resources workaround for nested tasks

test_f_g.py

import ray
import time


@ray.remote(num_cpus=0)
class Counter:
    def __init__(self):
        self.num_f = 0
        self.num_g = 0

        self.max_f, self.max_g = 0, 0

    def inc_f(self, val):
        self.num_f += val
        if self.num_f > self.max_f:
            self.max_f = self.num_f

    def inc_g(self, val):
        self.num_g += val
        if self.num_g > self.max_g:
            self.max_g = self.num_g

    def get(self):
        return (self.max_f, self.max_g)

@ray.remote(num_cpus=1)
class G:
    def run(self, counter):
        counter.inc_g.remote(1)
        time.sleep(1)
        ray.get(counter.inc_g.remote(-1))

@ray.remote
def f(counter):
    counter.inc_f.remote(1)
    time.sleep(1)

    g = G.remote()
    ray.get(g.run.remote(counter))

    ray.get(counter.inc_f.remote(-1))


ray.init(num_cpus=2)

counter = Counter.remote()
ray.get(counter.get.remote())
ray.get([f.remote(counter) for _ in range(10)])
print(ray.get(counter.get.remote()))

test_f_g_custom_resources.py

import ray
import time


@ray.remote(num_cpus=0)
class Counter:
    def __init__(self):
        self.num_f = 0
        self.num_g = 0

        self.max_f, self.max_g = 0, 0

    def inc_f(self, val):
        self.num_f += val
        if self.num_f > self.max_f:
            self.max_f = self.num_f

    def inc_g(self, val):
        self.num_g += val
        if self.num_g > self.max_g:
            self.max_g = self.num_g

    def get(self):
        return (self.max_f, self.max_g)

@ray.remote(num_cpus=1)
class G:
    def run(self, counter):
        counter.inc_g.remote(1)
        time.sleep(1)
        ray.get(counter.inc_g.remote(-1))

@ray.remote(resources={"f": 1})
def f(counter):
    counter.inc_f.remote(1)
    time.sleep(1)

    g = G.remote()
    ray.get(g.run.remote(counter))

    ray.get(counter.inc_f.remote(-1))


ray.init(num_cpus=2, resources={"f": 1})

counter = Counter.remote()
ray.get(counter.get.remote())
ray.get([f.remote(counter) for _ in range(10)])
print(ray.get(counter.get.remote()))