Spaces:
Runtime error
Runtime error
File size: 4,469 Bytes
0558aa4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import signal
import sys
import torch
from lightning.pytorch.callbacks import Callback
from nemo.utils import logging
class PreemptionCallback(Callback):
"""
PreemptionCallback class creates a callback that checks for preemption during training at the end of every step.
Upon preemption the callback provides a function to gracefully exit the training immediately and also saves the current state in a checkpoint as *last.ckpt.
(to be able to start from the same step without wasting any compute while resuming the next time).
PreemptionCallback is always enabled by default via the arg create_preemption_callback under ExpManagerConfig. To disable please pass
create_preemption_callback: False in your config file.
"""
def __init__(self, checkpoint_callback, sig=None):
self.sig = sig
if self.sig is None:
self.sig = signal.SIGTERM
self.checkpoint_callback = checkpoint_callback
self.preemption_enabled = False
@property
def interrupted(self):
interrupted = torch.tensor(self._interrupted, device=torch.cuda.current_device(), dtype=torch.int32)
torch.distributed.broadcast(interrupted, 0)
interrupted = bool(interrupted.item())
return interrupted
def on_train_start(self, trainer, pl_module):
"""
Defines custom handlers at the beginning of training to be executed when the
preemption signal is received.
"""
# Check if torch distributed is initialised, as its needed for broadcasting the preemption signal to all the ranks
if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
logging.info("Preemption requires torch distributed to be initialized, disabling preemption")
else:
self.preemption_enabled = True
# Bool var that's initialized to false and made True upon receving the preemption signal
self._interrupted = False
self.released = False
self.original_handler = signal.getsignal(self.sig)
# Master handler executed only by rank 0 when the preemption siganal is received, to avoid deadlock conditions
def master_handler(signum, frame):
self.release()
self._interrupted = True
# Handler executed by the non zero ranks
def ignoring_handler(signum, frame):
self.release()
self.private_rank = torch.distributed.get_rank()
if self.private_rank == 0:
signal.signal(self.sig, master_handler)
else:
signal.signal(self.sig, ignoring_handler)
return self
def on_train_end(self, trainer, pl_module):
if self.preemption_enabled:
self.release()
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx: int):
if self.preemption_enabled:
# check if the job was preempted at the end of every training step/iteration
# NOTE: "self.interrupted" is a property which triggers a
# distributed broadcast of "_interrupted" flag from rank 0 to all other
# ranks, to avoid performance overheads it's best to store the result in
# a regular local variable
interrupted = self.interrupted
if interrupted:
logging.info("Received SIGTERM, saving checkpoint and exiting")
monitor_candidates = self.checkpoint_callback._monitor_candidates(trainer)
self.checkpoint_callback._save_last_checkpoint(trainer, monitor_candidates)
sys.exit(0)
def release(self):
if self.released:
return False
signal.signal(self.sig, self.original_handler)
self.released = True
return True
|