2/n Simplify spawn plugins: Spawn immediately

[awaelchli]

What does this PR do?

Fixes #10059 (step 4 & 5)

Moves the spawning of process right at the beginning of Trianer.fit/validate/test/predict.

• Greatly improve readability due to simplified code path
• Aligns the code path of DDPSpawn with other plugins like DDP
• Hooks like setup() and teardown() now consistently run inside the worker processes and process group initialized
• Significantly decreases the chance that the root_device gets accessed in a hook that runs outside the spawned processes in TPU spawn. This today would cause the process to exit with a cryptic message.

The core change of is that the spawning happens a level higher up in the trainer. Here is a simplified illustration using functions for before vs. after:

Before:

def _run(model, ...):
    setup()
    other()
    # <-- processes spawn here
    mp.spawn(train_fn, args=model)
    # <-- processes join here

    teardown()

After:

def _run(model):
    setup()
    other()
    train_fn(model)
    teardown()

# <-- processes spawn here
mp.spawn(_run, args=model)
# <-- processes join here

Possible issues that this PR might fix or help fix (need to check):

• Exception: The wandb backend process has shutdown #10688
• WandBlogger doesn't log when working with TPU cores #9876

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 update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
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?

I made sure I had fun coding 🙃

Part of #1 (it's a lie, this is just here to avoid noisy GitHub bot)

cc @Borda @justusschock @kaushikb11 @awaelchli @ninginthecloud

[awaelchli]

I have addressed all major questions and reviews. Thanks for the input and positive response so far!

[ananthsub]

excellent! super exciting to see this simplification along with all of the followup opportunities this surfaced!

[four4fish]

Looks great to me! I don't have permission to accept, but super happy to see this 💯 A big milestone for Accelerator stable version!!

[kaushikb11]

Woohoo!

[tchaton]

Awesome work !