Adding basic preemption code #6161

athitten · 2023-03-09T21:33:07Z

Add preemption functionality in preemption_callback.py under utils
Refactor the code to move NemoModelCheckpoint callback under callbacks

What does this PR do ?

Adding basic functionality for cluster preemption.

Collection: [Note which collection this PR will affect]

Changelog

Add specific line by line info of high level changes in this PR.

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

Related to # (issue)

nemo/collections/common/callbacks/nemomodelcheckpoint.py

nemo/utils/exp_manager.py

nemo/utils/preemption_callback.py

nemo/collections/common/callbacks/nemomodelcheckpoint.py

nemo/utils/preemption_callback.py

nemo/utils/callbacks/nemo_model_checkpoint.py

+                continue
+            index = checkpoint.find(self.monitor) + len(self.monitor) + 1  # Find monitor in str + 1 for '='
+            if index != -1:
+                match = re.search('[A-z]', checkpoint[index:])


nemo/utils/callbacks/nemo_model_checkpoint.py

+            try:
+                self._fs.rm(filepath)
+                logging.info(f"Removed checkpoint: {filepath}")
+            except:


nemo/utils/callbacks/nemo_model_checkpoint.py

ericharper

Everything looks good to me. Once we pass CI and verify it's working on the clusters, let's merge it.

titu1994

Some minor comments

titu1994 · 2023-03-31T23:31:34Z

nemo/utils/callbacks/preemption.py

+    PreemptionCallback is always enabled.
+    """
+
+    def __init__(self, checkpoint_callback, sig=signal.SIGTERM):


Should be a None by default, and if None then self.sig = [signal.SIGTERM]. That way BCP can pass other signals.

@titu1994 thanks, made this change

titu1994 · 2023-03-31T23:32:56Z

nemo/utils/callbacks/preemption.py

+
+    @property
+    def interrupted(self):
+        interrupted = torch.tensor(self._interrupted).int().to(torch.device('cuda'))


Instead of .to(), use device= inside of the torch.tensor(..., device=torch.cuda.current_device(), dtype=torch.int32).

Added this too

nemo/utils/callbacks/preemption.py

titu1994 · 2023-03-31T23:37:46Z

nemo/utils/exp_manager.py

@@ -1122,6 +888,8 @@ def configure_checkpointing(
    if 'mp_rank' in checkpoint_callback.last_model_path or 'tp_rank' in checkpoint_callback.last_model_path:
        checkpoint_callback.last_model_path = uninject_model_parallel_rank(checkpoint_callback.last_model_path)
    trainer.callbacks.append(checkpoint_callback)
+    preemption_callback = PreemptionCallback(checkpoint_callback)


Will this be a default ? I would suggest to add a bool flag that enables by default but can be disabled if wanted.

Makes sense, changed this in the latest commit

titu1994

Preemption signal is currentlt not configurable from config

nemo/utils/exp_manager.py

nemo/utils/callbacks/preemption.py

titu1994

I think this current version doesn't work as it's just overriding all self.xyz with last signal.

Let's revert the commit and merge earlier version.

titu1994 · 2023-04-05T08:58:06Z

nemo/utils/callbacks/preemption.py

+            # 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(sig)


This just gets overriden by whatever is the last signal now

We don't need to do special logic for each signal. Listen for all, if any of those are received, ignore the rest and call preemption code c

If this logic is getting too complicated revert this commit and go with previous tested version for now.

Yes thanks for the catch, you are right self.original_handler = signal.getsignal(sig) will have the last signal. When I tested with 2 signals, it worked though bcoz the actual handlers are functioning fine and self.original_handler = signal.getsignal(sig) just seems to be used to reset the handlers when the signal is received. Probably, that's why it dint show up as an error in my testing.

For now, I'll go ahead and revert the commit and have the previous code merged first and then look into this.

nemo/utils/callbacks/preemption.py

titu1994

Ready to merge

nemo/utils/callbacks/preemption.py

titu1994

Lgtm

Add preemption functionality in preemption_callback.py under utils Refactor the code to move NemoModelCheckpoint callback under callbacks Signed-off-by: Abhishree <[email protected]>

1) Rename nemo/collections/common/callbacks/nemomodelcheckpoint.py to nemo/utils/callbacks/nemo_model_checkpoint.py 2) Rename nemo/utils/preemption_callback.py to nemo/utils/callbacks/preemption.py 3) Add docstrings, headers, logging and check for torch distributed Signed-off-by: Abhishree <[email protected]>

Signed-off-by: Abhishree <[email protected]>

1) Add boolean flag for createing preemption callback 2) Make sig arg in PreemptionCallback as None 3) Other minor modifications and code comments Signed-off-by: Abhishree <[email protected]>

… if unavailable Signed-off-by: Abhishree <[email protected]>

Signed-off-by: Abhishree <[email protected]>

…lass from exp_manager.py Signed-off-by: Abhishree <[email protected]>

for more information, see https://pre-commit.ci

titu1994 · 2023-04-06T23:30:38Z

nemo/utils/callbacks/nemo_model_checkpoint.py

+import re
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union


In a later PR lets clean this part up.

nemo/utils/callbacks/nemo_model_checkpoint.py

+        self.best_model_path = best_k_models[0]
+        self.best_model_score = self.best_k_models[self.best_model_path]
+
+    def on_save_checkpoint(self, trainer, pl_module, checkpoint):


titu1994

Cool ! Looks good, thanks !

titu1994 · 2023-04-06T23:30:38Z

nemo/utils/callbacks/nemo_model_checkpoint.py

+import re
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union


In a later PR lets clean this part up.

* Adding basic preemption code Add preemption functionality in preemption_callback.py under utils Refactor the code to move NemoModelCheckpoint callback under callbacks Signed-off-by: Abhishree <[email protected]> * Adding the following modifications 1) Rename nemo/collections/common/callbacks/nemomodelcheckpoint.py to nemo/utils/callbacks/nemo_model_checkpoint.py 2) Rename nemo/utils/preemption_callback.py to nemo/utils/callbacks/preemption.py 3) Add docstrings, headers, logging and check for torch distributed Signed-off-by: Abhishree <[email protected]> * Minor edit in preemption.py Signed-off-by: Abhishree <[email protected]> * Removing unused imports Signed-off-by: Abhishree <[email protected]> * Remove device arg from PreemptionCallback class Signed-off-by: Abhishree <[email protected]> * Add more details in the NeMoModelCheckpointdocstring Signed-off-by: Abhishree <[email protected]> * Add the following modifications: 1) Add boolean flag for createing preemption callback 2) Make sig arg in PreemptionCallback as None 3) Other minor modifications and code comments Signed-off-by: Abhishree <[email protected]> * Modify torch cuda and distributed available checks to skip preemption if unavailable Signed-off-by: Abhishree <[email protected]> * Add preemption_enabled flag to preemption.py Signed-off-by: Abhishree <[email protected]> * Update nemo_model_checkpoint.py with the latest NemoModelCheckpoint class from exp_manager.py Signed-off-by: Abhishree <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Abhishree <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: hsiehjackson <[email protected]>

github-actions bot added the common label Mar 9, 2023

ericharper reviewed Mar 9, 2023

View reviewed changes

nemo/collections/common/callbacks/nemomodelcheckpoint.py Outdated Show resolved Hide resolved

ericharper reviewed Mar 9, 2023

View reviewed changes

nemo/utils/exp_manager.py Outdated Show resolved Hide resolved

ericharper reviewed Mar 9, 2023

View reviewed changes

nemo/utils/preemption_callback.py Outdated Show resolved Hide resolved

github-advanced-security bot found potential problems Mar 9, 2023

View reviewed changes

ericharper reviewed Mar 9, 2023

View reviewed changes

nemo/utils/preemption_callback.py Outdated Show resolved Hide resolved

ericharper reviewed Mar 9, 2023

View reviewed changes

nemo/utils/preemption_callback.py Outdated Show resolved Hide resolved

ericharper reviewed Mar 9, 2023

View reviewed changes

nemo/utils/preemption_callback.py Outdated Show resolved Hide resolved

athitten force-pushed the athitten/cluster_preemption branch 2 times, most recently from b50e092 to 483f932 Compare March 16, 2023 04:35

github-actions bot removed the common label Mar 16, 2023

github-advanced-security bot found potential problems Mar 16, 2023

View reviewed changes

athitten force-pushed the athitten/cluster_preemption branch 3 times, most recently from 27891dc to 34a5964 Compare March 30, 2023 17:11

athitten marked this pull request as ready for review March 30, 2023 17:13

athitten requested a review from titu1994 March 30, 2023 17:14

ericharper reviewed Mar 30, 2023

View reviewed changes

nemo/utils/callbacks/nemo_model_checkpoint.py Outdated Show resolved Hide resolved

ericharper reviewed Mar 30, 2023

View reviewed changes

athitten force-pushed the athitten/cluster_preemption branch from 99553e5 to 65bc34d Compare March 30, 2023 21:56

titu1994 reviewed Mar 31, 2023

View reviewed changes

athitten force-pushed the athitten/cluster_preemption branch from fc56c8a to 34bcc1f Compare April 3, 2023 19:19

titu1994 reviewed Apr 3, 2023

View reviewed changes

nemo/utils/exp_manager.py Outdated Show resolved Hide resolved

titu1994 reviewed Apr 4, 2023

View reviewed changes

nemo/utils/exp_manager.py Show resolved Hide resolved

nemo/utils/exp_manager.py Outdated Show resolved Hide resolved

athitten force-pushed the athitten/cluster_preemption branch from 8a5eb1d to 4f3a280 Compare April 5, 2023 05:47

github-advanced-security bot found potential problems Apr 5, 2023

View reviewed changes

nemo/utils/callbacks/preemption.py Fixed Show resolved Hide resolved

nemo/utils/callbacks/preemption.py Fixed Show fixed Hide fixed

titu1994 reviewed Apr 5, 2023

View reviewed changes

athitten force-pushed the athitten/cluster_preemption branch from 74b17c4 to 2f1e826 Compare April 5, 2023 15:25

titu1994 previously approved these changes Apr 5, 2023

View reviewed changes

athitten dismissed titu1994’s stale review via c04f57a April 5, 2023 18:27

athitten force-pushed the athitten/cluster_preemption branch from fd50914 to c04f57a Compare April 5, 2023 18:27

github-advanced-security bot found potential problems Apr 5, 2023

View reviewed changes

nemo/utils/callbacks/preemption.py Fixed Show fixed Hide fixed

titu1994 previously approved these changes Apr 5, 2023

View reviewed changes

athitten dismissed titu1994’s stale review via e42618e April 6, 2023 06:24

athitten force-pushed the athitten/cluster_preemption branch from dbb20f4 to e42618e Compare April 6, 2023 06:24

athitten added 10 commits April 6, 2023 19:11

Adding basic preemption code

6ebcfaa

Add preemption functionality in preemption_callback.py under utils Refactor the code to move NemoModelCheckpoint callback under callbacks Signed-off-by: Abhishree <[email protected]>

Minor edit in preemption.py

604b8c3

Signed-off-by: Abhishree <[email protected]>

Removing unused imports

e418e16

Signed-off-by: Abhishree <[email protected]>

Remove device arg from PreemptionCallback class

e632930

Signed-off-by: Abhishree <[email protected]>

Add more details in the NeMoModelCheckpointdocstring

69a0c97

Signed-off-by: Abhishree <[email protected]>

Add the following modifications:

50cd904

1) Add boolean flag for createing preemption callback 2) Make sig arg in PreemptionCallback as None 3) Other minor modifications and code comments Signed-off-by: Abhishree <[email protected]>

Modify torch cuda and distributed available checks to skip preemption…

3b56d55

… if unavailable Signed-off-by: Abhishree <[email protected]>

Add preemption_enabled flag to preemption.py

772239d

Signed-off-by: Abhishree <[email protected]>

Update nemo_model_checkpoint.py with the latest NemoModelCheckpoint c…

0a329da

…lass from exp_manager.py Signed-off-by: Abhishree <[email protected]>

athitten force-pushed the athitten/cluster_preemption branch from c77f90f to 0a329da Compare April 6, 2023 19:12

[pre-commit.ci] auto fixes from pre-commit.com hooks

ae9217c

for more information, see https://pre-commit.ci

github-advanced-security bot found potential problems Apr 6, 2023

View reviewed changes

titu1994 approved these changes Apr 6, 2023

View reviewed changes

titu1994 merged commit 9ba4abd into main Apr 6, 2023

titu1994 deleted the athitten/cluster_preemption branch April 6, 2023 23:32

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Adding basic preemption code #6161

Adding basic preemption code #6161

athitten commented Mar 9, 2023

ericharper left a comment

titu1994 left a comment

titu1994 Mar 31, 2023

athitten Apr 3, 2023

titu1994 Mar 31, 2023

athitten Apr 3, 2023

titu1994 Mar 31, 2023

athitten Apr 3, 2023

titu1994 left a comment

titu1994 left a comment

titu1994 Apr 5, 2023

titu1994 Apr 5, 2023

titu1994 Apr 5, 2023

athitten Apr 5, 2023

athitten Apr 5, 2023

titu1994 left a comment

titu1994 left a comment

titu1994 Apr 6, 2023

titu1994 left a comment

titu1994 Apr 6, 2023

Adding basic preemption code #6161

Adding basic preemption code #6161

Conversation

athitten commented Mar 9, 2023

What does this PR do ?

Changelog

Usage

Before your PR is "Ready for review"

Who can review?

Additional Information

ericharper left a comment

Choose a reason for hiding this comment

titu1994 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

titu1994 left a comment

Choose a reason for hiding this comment

titu1994 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

titu1994 left a comment

Choose a reason for hiding this comment

titu1994 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

titu1994 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment