[train][tune] fix WandbLoggerCallback to reuse actors upon restore

[matthewdeng]

Why are these changes needed?

This fixes an issue in the restoration flow of Trials. The _WandbLoggingActor actor is cleaned up as part of the trial completion flow, but not appropriately handled during restoration, leading to leaked actors.

This fix reuses the same actor across restoration.

Before	Header

Repo:

import ray.train
from ray.train import RunConfig, FailureConfig
from ray.train.torch import TorchTrainer
from ray.air.integrations.wandb import WandbLoggerCallback

def train_func():
    import time
    time.sleep(30)
    ray.train.report({"a":1})
    1/0

trainer = TorchTrainer(train_func, run_config=RunConfig(failure_config=FailureConfig(max_failures=100), callbacks=[WandbLoggerCallback()]))
trainer.fit()

_logging_future_to_trial would continue to grow and contain references to the old futures.

self._trial_logging_actors={TorchTrainer_300e0_00000: Actor(_WandbLoggingActor, d4612996084220d2e81a0a1c05000000)}

self._trial_logging_futures={TorchTrainer_300e0_00000: ObjectRef(5fca57afecc76a7bd4612996084220d2e81a0a1c0500000001000000)}

self._logging_future_to_trial={ObjectRef(45efae1e94df359f2c025b00bc0903ded351a7690500000001000000): TorchTrainer_300e0_00000, ObjectRef(81a3509f4f3351b28f43f41babde7280bbf087930500000001000000): TorchTrainer_300e0_00000, ObjectRef(f72a3a2cf9d68a40687e9e7fe24faa8aac63ddf70500000001000000): TorchTrainer_300e0_00000, ObjectRef(cd38da7afea09ef59d7fd91d7e72f4b1e011ab920500000001000000): TorchTrainer_300e0_00000, ObjectRef(5fca57afecc76a7bd4612996084220d2e81a0a1c0500000001000000): TorchTrainer_300e0_00000}

self._trial_queues={TorchTrainer_300e0_00000: <ray.util.queue.Queue object at 0x7718d8171070>}

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 :(

[justinvyu]

Thanks for fixing this!