[wandb] Use wandb Run as a context manager

[timoffex]

Why are these changes needed?

This fixes two problems:

• Using the global wandb.finish() function is discouraged as it relies on there being exactly one run in the current Python interpreter. Instead, the finish() method of the desired Run object should be called directly.
• We recommend using the Run as a context manager to correctly finish it (and mark it as Crashed) in case of exceptions.

I found this implementation while helping debug an issue for a user who's using this integration together with PyTorch Lightning's W&B integration. I don't think this PR fixes a bug, but the original code is risky and doesn't follow our recommendations.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[matthewdeng]

Thanks for contributing this! Is there any W&B documentation that describes the best practices for this? Would love to brush up on this information and make sure we're applying the best practices everywhere we mention W&B (e.g. the documentation).

[timoffex]

Slightly embarrassingly, I think our docs don't do a great job of mentioning this particular recommendation. I'm working on fixing that with our documentation team, after which https://docs.wandb.ai/guides/track/launch/ and especially https://docs.wandb.ai/guides/integrations/add-wandb-to-any-library would mention proper run.finish() or run-as-a-context-manager usage.

[timoffex]

Took a while, but it's in the works :) wandb/docs#1146

[stale]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 14 days if no further activity occurs. Thank you for your contributions.

• If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[timoffex]

Friendly ping! Is merging this PR my responsibility or the repo owners'? I do not appear to have the ability to initiate a merge.

[pcmoritz]

After the tests pass, we can merge it :)

[timoffex]

Aha, I didn't realize I could see the tests. I see some failures that look relevant; I'll try to investigate soon.

[timoffex]

Is there a good reason to use the _MockWandbAPI class instead of using wandb directly in these tests? I think the test failures are simply a result of _MockWandbAPI.init() returning a value that doesn't implement __enter__ and __exit__.

I could update the mock, but I think the tests would be better if they used wandb directly. There's a way to run wandb in offline mode and to make sure that it doesn't touch files in the current directory.

Also, I'm having trouble getting the tests to run locally: I rebased to master, installed test requirements as in https://docs.ray.io/en/latest/ray-contribute/getting-involved.html#testing and followed the instructions at https://docs.ray.io/en/latest/ray-contribute/development.html#building-ray-python-only, but I get the following error

❯ python -m pytest -s python/ray/air/tests/test_integration_wandb.py
ImportError while loading conftest '/Users/timoffex/Documents/workspace/ray/python/ray/air/tests/conftest.py'.
python/ray/__init__.py:85: in <module>
    import ray._raylet  # noqa: E402
E   ModuleNotFoundError: No module named 'ray._raylet'

Any advice?