feat(wandb): support distributed modes

[borisdayma]

What does this PR do?

Fixes #11602

Improves support of distributed mode.

For now it also requires calling wandb.require("service") in training script.

Does your PR introduce any breaking changes? If yes, please list them.

The run is created earlier than it used to be as it happens when we instantiate WandbLogger while it used to be created when first logging metric, hyper-parameters or logging any other supported method.

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

cc @awaelchli @kptkin @raubitsj who are also working on this PR
cc @morganmcg1 @scottire @manangoel99

[awaelchli]

Amazing! I tried and it works great so far. Couldn't push to your branch for some reason. I suggest to update the tests:

diff --git a/tests/loggers/test_wandb.py b/tests/loggers/test_wandb.py
index 85b20c562..d347d32d1 100644
--- a/tests/loggers/test_wandb.py
+++ b/tests/loggers/test_wandb.py
@@ -22,6 +22,7 @@ from pytorch_lightning import Trainer
 from pytorch_lightning.loggers import WandbLogger
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers import BoringModel
+from tests.helpers.utils import no_warning_call
 
 
 @mock.patch("pytorch_lightning.loggers.wandb.wandb")
@@ -57,7 +58,7 @@ def test_wandb_logger_init(wandb):
     # verify default resume value
     assert logger._wandb_init["resume"] == "allow"
 
-    with pytest.warns(UserWarning, match="There is a wandb run already in progress"):
+    with no_warning_call(UserWarning, match="There is a wandb run already in progress"):
         _ = logger.experiment
 
     logger.log_metrics({"acc": 1.0}, step=3)
@@ -124,8 +125,10 @@ def test_wandb_pickle(wandb, tmpdir):
 def test_wandb_logger_dirs_creation(wandb, tmpdir):
     """Test that the logger creates the folders and files in the right place."""
     logger = WandbLogger(save_dir=str(tmpdir), offline=True)
-    assert logger.version is None
-    assert logger.name is None
+
+    # the wandb Run gets initialized on instantiation, version and name should be available
+    assert logger.version is not None
+    assert logger.name is not None
 
     # mock return values of experiment
     wandb.run = None

[justusschock]

LGTM in general, some very minor nit :)

[borisdayma]

Confused about the pre-commit error.
I have the hook installed locally and same isort version.

[rohitgr7]

Confused about the pre-commit error. I have the hook installed locally and same isort version.

pre-commit errors seem correct since there is an import that isn't used.

[borisdayma]

I was not able to understand the problem with isort.
All imports seem to be used…

Also  I tried using monkeypatch for setting wandb version or _WANDB_GREATER_EQUAL_XXX but ran into a bunch of errors…

[rohitgr7]

hey @borisdayma

here is the updated code with monkeypatch you can use: https://pastebin.com/hvkyGxqc

[borisdayma]

Many thanks @rohitgr7
I'll make sure to use a non-protected fork in the future!