Wait for launch pad program to finish in a subprocess

[ethanluoyc]

Hi,

I am currently using LaunchPad for a local distributed RL program. I am using the local multiprocessing backend but would like to have multiple runs of the program in a single process to perform a sweep.

I have the following script (ignoring the details of the agent implementation):

def run():
    program = impala_agent.DistributedIMPALA(
        environment_factory=environment_factory,
        network_factory=network_factory,
        num_actors=8,
        max_actor_steps=int(1e8),
        log_every=10.0,
    ).build()
    local_resources = {
        "learner": PythonProcess(
            env={"JAX_PLATFORM_NAME": "gpu", "CUDA_VISIBLE_DEVICES": "0"}
        ),
        "evaluator": PythonProcess(env={"CUDA_VISIBLE_DEVICES": ""}),
        "actor": PythonProcess(env={"CUDA_VISIBLE_DEVICES": ""}),
        "counter": PythonProcess(env={"CUDA_VISIBLE_DEVICES": ""}),
        "replay": PythonProcess(env={"CUDA_VISIBLE_DEVICES": ""}),
        "coordinator": PythonProcess(env={"CUDA_VISIBLE_DEVICES": ""}),
    }
    lp.launch(
        program,
        lp.LaunchType.LOCAL_MULTI_PROCESSING,
        terminal="current_terminal",
        local_resources=local_resources,
    )

def main(_):
    p = multiprocessing.Process(target=run)
    p.start()
    p.join() # This returns immediately

This currently does not work as the LaunchPad program would kill the entire current process tree at exit (at least in my experience when I have a reverb node that never terminates properly and has to be killed with os.kill). This prevents me from performing a sweep in the same process with something like

for params in SWEEPS:
  p = multiprocessing.Process(target=run, args=(params,))
  p.start()
  p.join() # This returns immediately

With the script, the process would return immediately and exit. However, the LP processes will not be shut down and I end up with the agent processes running in the background. In the ideal case, I would love to be able to wait on the LP program to finish and then execute another LP program in the same process with out having to write a shell script to wrap the sweeps.

Thanks!

[qstanczyk]

Hi Yicheng, please give this change a try. Now lp.lauch() returns back a worker manager, you can call wait() on it to sleep until program is over. Let me know how that works for you.

[qstanczyk]

Please reopen in case of any problems.

[ethanluoyc]

Hi Yicheng, please give this change a try. Now lp.lauch() returns back a worker manager, you can call wait() on it to sleep until program is over. Let me know how that works for you.

I can confirm it works for me :)