From 3e22b1fd14b470b45fbcc718cccb6923db3995dc Mon Sep 17 00:00:00 2001 From: natinew77-creator Date: Wed, 10 Dec 2025 17:10:55 -0500 Subject: [PATCH] fix: use launch_with_termination_handler for clean experiment termination This PR addresses issue #312 by updating all example scripts to use lp_utils.launch_with_termination_handler() instead of lp.launch() directly. The launch_with_termination_handler() function wraps lp.launch() with a signal handler that ensures all Launchpad processes are properly terminated when the parent process receives a SIGTERM or SIGINT signal. This prevents the messy teardown messages like: 'Worker groups that did not terminate in time: ['actor'] Killing entire runtime.' This is especially important when using external loggers like Weights and Biases, which previously reported experiments as 'crashed' even though they completed successfully. Files updated: - examples/baselines/imitation/*.py - examples/baselines/rl_continuous/*.py - examples/baselines/rl_discrete/*.py - examples/multiagent/multigrid/run_multigrid.py Fixes #312 --- examples/baselines/imitation/run_bc.py | 3 ++- examples/baselines/imitation/run_gail.py | 3 ++- examples/baselines/imitation/run_iqlearn.py | 3 ++- examples/baselines/imitation/run_pwil.py | 3 ++- examples/baselines/imitation/run_sqil.py | 3 ++- examples/baselines/rl_continuous/run_d4pg.py | 3 ++- examples/baselines/rl_continuous/run_dmpo.py | 3 ++- examples/baselines/rl_continuous/run_mogmpo.py | 3 ++- examples/baselines/rl_continuous/run_mpo.py | 3 ++- examples/baselines/rl_continuous/run_ppo.py | 3 ++- examples/baselines/rl_continuous/run_sac.py | 3 ++- examples/baselines/rl_continuous/run_td3.py | 3 ++- examples/baselines/rl_continuous/run_wpo.py | 3 ++- examples/baselines/rl_discrete/run_dqn.py | 3 ++- examples/baselines/rl_discrete/run_impala.py | 3 ++- examples/baselines/rl_discrete/run_mdqn.py | 3 ++- examples/baselines/rl_discrete/run_muzero.py | 3 ++- examples/baselines/rl_discrete/run_qr_dqn.py | 3 ++- examples/baselines/rl_discrete/run_r2d2.py | 3 ++- examples/multiagent/multigrid/run_multigrid.py | 3 ++- 20 files changed, 40 insertions(+), 20 deletions(-) diff --git a/examples/baselines/imitation/run_bc.py b/examples/baselines/imitation/run_bc.py index ccbb2b29c6..be2402d017 100644 --- a/examples/baselines/imitation/run_bc.py +++ b/examples/baselines/imitation/run_bc.py @@ -183,7 +183,8 @@ def main(_): config = build_experiment_config() if FLAGS.run_distributed: program = experiments.make_distributed_offline_experiment(experiment=config) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_offline_experiment( experiment=config, diff --git a/examples/baselines/imitation/run_gail.py b/examples/baselines/imitation/run_gail.py index a7c43a6c04..bc42ae96d3 100644 --- a/examples/baselines/imitation/run_gail.py +++ b/examples/baselines/imitation/run_gail.py @@ -142,7 +142,8 @@ def main(_): if FLAGS.run_distributed: program = experiments.make_distributed_experiment( experiment=config, num_actors=4) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, diff --git a/examples/baselines/imitation/run_iqlearn.py b/examples/baselines/imitation/run_iqlearn.py index 24e118059b..cbc73b5871 100644 --- a/examples/baselines/imitation/run_iqlearn.py +++ b/examples/baselines/imitation/run_iqlearn.py @@ -133,7 +133,8 @@ def main(_): program = experiments.make_distributed_experiment( experiment=config, num_actors=4 ) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, diff --git a/examples/baselines/imitation/run_pwil.py b/examples/baselines/imitation/run_pwil.py index 179b77945a..449480627c 100644 --- a/examples/baselines/imitation/run_pwil.py +++ b/examples/baselines/imitation/run_pwil.py @@ -148,7 +148,8 @@ def main(_): if FLAGS.run_distributed: program = experiments.make_distributed_experiment( experiment=config, num_actors=4) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, diff --git a/examples/baselines/imitation/run_sqil.py b/examples/baselines/imitation/run_sqil.py index aeb77c645c..418749bb47 100644 --- a/examples/baselines/imitation/run_sqil.py +++ b/examples/baselines/imitation/run_sqil.py @@ -92,7 +92,8 @@ def main(_): if FLAGS.run_distributed: program = experiments.make_distributed_experiment( experiment=config, num_actors=4) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, diff --git a/examples/baselines/rl_continuous/run_d4pg.py b/examples/baselines/rl_continuous/run_d4pg.py index e6e9326b38..0fd5baa466 100644 --- a/examples/baselines/rl_continuous/run_d4pg.py +++ b/examples/baselines/rl_continuous/run_d4pg.py @@ -73,7 +73,8 @@ def main(_): if FLAGS.run_distributed: program = experiments.make_distributed_experiment( experiment=config, num_actors=4) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, diff --git a/examples/baselines/rl_continuous/run_dmpo.py b/examples/baselines/rl_continuous/run_dmpo.py index a477ce86f4..db9bee1b62 100644 --- a/examples/baselines/rl_continuous/run_dmpo.py +++ b/examples/baselines/rl_continuous/run_dmpo.py @@ -87,7 +87,8 @@ def main(_): if RUN_DISTRIBUTED.value: program = experiments.make_distributed_experiment( experiment=config, num_actors=4) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, diff --git a/examples/baselines/rl_continuous/run_mogmpo.py b/examples/baselines/rl_continuous/run_mogmpo.py index d4da09d2d3..519a1c7033 100644 --- a/examples/baselines/rl_continuous/run_mogmpo.py +++ b/examples/baselines/rl_continuous/run_mogmpo.py @@ -79,7 +79,8 @@ def main(_): if RUN_DISTRIBUTED.value: program = experiments.make_distributed_experiment( experiment=config, num_actors=4) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, diff --git a/examples/baselines/rl_continuous/run_mpo.py b/examples/baselines/rl_continuous/run_mpo.py index 51c928a62a..6370cc6424 100644 --- a/examples/baselines/rl_continuous/run_mpo.py +++ b/examples/baselines/rl_continuous/run_mpo.py @@ -79,7 +79,8 @@ def main(_): if RUN_DISTRIBUTED.value: program = experiments.make_distributed_experiment( experiment=config, num_actors=4) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, diff --git a/examples/baselines/rl_continuous/run_ppo.py b/examples/baselines/rl_continuous/run_ppo.py index ae2fb12831..c55f3873fa 100644 --- a/examples/baselines/rl_continuous/run_ppo.py +++ b/examples/baselines/rl_continuous/run_ppo.py @@ -61,7 +61,8 @@ def main(_): if FLAGS.run_distributed: program = experiments.make_distributed_experiment( experiment=config, num_actors=FLAGS.num_distributed_actors) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, diff --git a/examples/baselines/rl_continuous/run_sac.py b/examples/baselines/rl_continuous/run_sac.py index c80e84b8fe..407f909324 100644 --- a/examples/baselines/rl_continuous/run_sac.py +++ b/examples/baselines/rl_continuous/run_sac.py @@ -69,7 +69,8 @@ def main(_): if FLAGS.run_distributed: program = experiments.make_distributed_experiment( experiment=config, num_actors=4) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, diff --git a/examples/baselines/rl_continuous/run_td3.py b/examples/baselines/rl_continuous/run_td3.py index 9bb446a900..88885896a8 100644 --- a/examples/baselines/rl_continuous/run_td3.py +++ b/examples/baselines/rl_continuous/run_td3.py @@ -64,7 +64,8 @@ def main(_): if FLAGS.run_distributed: program = experiments.make_distributed_experiment( experiment=config, num_actors=4) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, diff --git a/examples/baselines/rl_continuous/run_wpo.py b/examples/baselines/rl_continuous/run_wpo.py index 633b75dd2a..9119c02756 100644 --- a/examples/baselines/rl_continuous/run_wpo.py +++ b/examples/baselines/rl_continuous/run_wpo.py @@ -77,7 +77,8 @@ def main(_): if RUN_DISTRIBUTED.value: program = experiments.make_distributed_experiment( experiment=config, num_actors=4) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, diff --git a/examples/baselines/rl_discrete/run_dqn.py b/examples/baselines/rl_discrete/run_dqn.py index 760e84ea52..dea9a7a8eb 100644 --- a/examples/baselines/rl_discrete/run_dqn.py +++ b/examples/baselines/rl_discrete/run_dqn.py @@ -74,7 +74,8 @@ def main(_): program = experiments.make_distributed_experiment( experiment=experiment_config, num_actors=4 if lp_utils.is_local_run() else 128) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment(experiment_config) diff --git a/examples/baselines/rl_discrete/run_impala.py b/examples/baselines/rl_discrete/run_impala.py index c566a7558f..b9e16023c2 100644 --- a/examples/baselines/rl_discrete/run_impala.py +++ b/examples/baselines/rl_discrete/run_impala.py @@ -79,7 +79,8 @@ def main(_): program = experiments.make_distributed_experiment( experiment=experiment_config, num_actors=4 if lp_utils.is_local_run() else 256) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment(experiment_config) diff --git a/examples/baselines/rl_discrete/run_mdqn.py b/examples/baselines/rl_discrete/run_mdqn.py index 3b4c940097..11d7ff2933 100644 --- a/examples/baselines/rl_discrete/run_mdqn.py +++ b/examples/baselines/rl_discrete/run_mdqn.py @@ -75,7 +75,8 @@ def main(_): program = experiments.make_distributed_experiment( experiment=experiment_config, num_actors=4 if lp_utils.is_local_run() else 128) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment(experiment_config) diff --git a/examples/baselines/rl_discrete/run_muzero.py b/examples/baselines/rl_discrete/run_muzero.py index ca8fc83eec..9144b9900a 100644 --- a/examples/baselines/rl_discrete/run_muzero.py +++ b/examples/baselines/rl_discrete/run_muzero.py @@ -122,7 +122,8 @@ def main(_): num_inference_servers=num_inference_servers, inference_server_config=inference_server_config, ) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program,),) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) if __name__ == '__main__': diff --git a/examples/baselines/rl_discrete/run_qr_dqn.py b/examples/baselines/rl_discrete/run_qr_dqn.py index ce3ba0852a..9a4536a712 100644 --- a/examples/baselines/rl_discrete/run_qr_dqn.py +++ b/examples/baselines/rl_discrete/run_qr_dqn.py @@ -79,7 +79,8 @@ def main(_): program = experiments.make_distributed_experiment( experiment=experiment_config, num_actors=4 if lp_utils.is_local_run() else 16) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment(experiment_config) diff --git a/examples/baselines/rl_discrete/run_r2d2.py b/examples/baselines/rl_discrete/run_r2d2.py index c22166cbb5..9d0905040b 100644 --- a/examples/baselines/rl_discrete/run_r2d2.py +++ b/examples/baselines/rl_discrete/run_r2d2.py @@ -84,7 +84,8 @@ def main(_): if FLAGS.run_distributed: program = experiments.make_distributed_experiment( experiment=config, num_actors=4 if lp_utils.is_local_run() else 80) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment(experiment=config) diff --git a/examples/multiagent/multigrid/run_multigrid.py b/examples/multiagent/multigrid/run_multigrid.py index 0b5ef0a06d..52ecf0159c 100644 --- a/examples/multiagent/multigrid/run_multigrid.py +++ b/examples/multiagent/multigrid/run_multigrid.py @@ -101,7 +101,8 @@ def main(_): if _RUN_DISTRIBUTED.value: program = experiments.make_distributed_experiment( experiment=config, num_actors=4) - lp.launch(program, xm_resources=lp_utils.make_xm_docker_resources(program)) + lp_utils.launch_with_termination_handler( + program, xm_resources=lp_utils.make_xm_docker_resources(program)) else: experiments.run_experiment( experiment=config, eval_every=_EVAL_EVERY.value, num_eval_episodes=5)