Reporting restructure

[scarere]

PR Type

Feature

Short Description

Clickup Ticket(s): Add Base Reporter Class

This PR restructures and unifies reporting. What was previously the metrics reporter is now JsonReporter. Both the JsonReporter and WandBReporter inherit from the same base class.

Servers and clients now accept any number of reporters and send data to them frequently (During setup, at the end of each round, each epoch and each batch). It is up to the reporter classes internal logic to determine how frequently to report data.

There is no longer any default reporters. Users will not have to pass None to the wandb reporter argument of clients and servers in order to turn wandb reporting off. Additionally if users want to generate the metrics json file frequently used for pytests, they must explicitly pass a JsonReporter to their server or client instance. Creating a JsonReporter instance with all its default arguments exactly mimics the behaviour of what was previously the metrics reporter (Except some of the keys have been changed). There is no longer any need to call dump to write the metrics, this has been moved into the server/client shutdown.

I also upgraded the wandb dependency.

I think there is some work to be done in terms of deciding how to structure what metrics and what to use for the key names. It be great if there was a bit of a more predictable pattern.

Tests Added

Had to modify a few tests to reflect the modified structure of the output metrics json. Also found two issues with the tests.

The MKMMD loss tests only work if your machine does not have a GPU. My temporary fix for this was to set devide to CPU. They still run in a few seconds. If GPU acceleration is needed then someone needs to go through and debug them. Created a ticket here

Tests that made use of both freeze_time and patch could not be run individually in some instances. It seems like somehow their imports were reliant on other tests in the test suite. This makes debugging tests a bit annoying because you have to run the entire test suite just to check whether or not you fixed one specific test that was failing. I created a ticket here