How to use Ray with closures?

[mil-ad]

Is it possible to define a remote function through closures?

In the example below I want get_func and my_func to share the same resources (e.g. one GPU):

def get_func(a):

   a = a.to(torch.device("cuda")

    def my_func(e, f):
        return (e*f) + a

    return my_func

worker1 = get_func(10)
worker2 = get_func(20)

worker1(100)
worker2(100)

I think just decorating get_func won't be enough and I guess below won't work either?

@ray.remote(num_cpus=4, num_gpus=1)
def get_func(a):

   a = a.to(torch.device("cuda")

    @ray.remote(num_cpus=4, num_gpus=1)
    def my_func(e, f):
        return (e*f) + a

    return my_func

worker1 = get_func(10)
worker2 = get_func(20)

worker1(100)
worker2(100)

[richardliaw]

Thanks for making this issue! Can you provide some more context about what you're trying to do (i.e., elaborating more about "making your functions share the same resources (e.g. one GPU)")?

[mil-ad]

So for example if you use something like Ignite with PyTorch, what you often do is construct trainer functions that are later called to actually do the training. It would be nice if the resource allocation could propagate from the trainer factory function to the actual trainer function.