From 8f114a9c1ae4493353821799bd1470016b056080 Mon Sep 17 00:00:00 2001 From: yrchen-tw Date: Wed, 9 Sep 2026 07:11:19 +0000 Subject: [PATCH 1/2] Remove redundant slice event trigger in preemption recovery --- axlearn/common/elastic_utils.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/axlearn/common/elastic_utils.py b/axlearn/common/elastic_utils.py index b8ae576ea..354092ba0 100644 --- a/axlearn/common/elastic_utils.py +++ b/axlearn/common/elastic_utils.py @@ -223,9 +223,6 @@ def handle_preemption_recovery( ) wait_for_slices(active_count, timeout_seconds=pause_timeout_seconds) - if elastic_manager: - elastic_manager.new_slice_event.set() - return active_count From d29e526b96f267de2d941df353f55ad03286a70e Mon Sep 17 00:00:00 2001 From: yrchen-tw Date: Wed, 9 Sep 2026 07:30:10 +0000 Subject: [PATCH 2/2] Increase timeout for waiting on all devices during elastic snapshotting --- axlearn/common/launch_trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/axlearn/common/launch_trainer.py b/axlearn/common/launch_trainer.py index 3014fe989..85bd1b761 100644 --- a/axlearn/common/launch_trainer.py +++ b/axlearn/common/launch_trainer.py @@ -251,7 +251,7 @@ def sigterm_handler(signum, frame): ) if elastic_snapshotting_enabled: - wait_for_all_devices() + wait_for_all_devices(timeout_seconds=1800) elastic_manager = None elastic_manager_initialized = False