Add Callback Support

[mattdeitke]

Background

Adds initial support for Callbacks, inspired by PyTorch Lightning.

The immediate use case is to enable logging during training with Weights and Biases.

Motivation

The motivation is to make it easier to log, debug, and inspect the training setup without having to manually modify runner.py.

Down the line, I suspect callbacks will also be the best place to write tests, where the tests may be in callback functions like on_checkpoint_load(model).

Example

An example usage might be to define a Callback class under the file training/callbacks/wandb_logging.py:

from typing import Any, Dict, Optional

import wandb
from allenact.base_abstractions.callbacks import Callback


class WandbLogging(Callback):
    def setup(self, name: str, **kwargs) -> None:
        wandb.init(
            project="test-project",
            entity="prior-ai2",
            name=name,
            config=kwargs,
        )

    def on_train_log(self, metric_means: Dict[str, float], step: int, **kwargs) -> None:
        wandb.log({**metric_means, "step": step})

    def on_valid_log(
        self,
        metrics: Optional[Dict[str, Any]],
        metric_means: Dict[str, float],
        step: int,
        **kwargs
    ) -> None:
        wandb.log({**metric_means, "step": step})

    def on_test_log(
        self,
        checkpoint: str,
        metrics: Dict[str, Any],
        metric_means: Dict[str, float],
        step: int,
        **kwargs
    ) -> None:
        wandb.log({**metric_means, "step": step})

and to use it, one would add the file to the --callbacks flag in the allenact command:

allenact <...> --callbacks training/callbacks/wandb_logging.py

Note that this doesn't require modifying the experiment configs at all, and hence is fully opt in functionality.

Notes

I'm still thinking about what callbacks would be best, and what should be passed into each of them.

Right now, I think the best approach I have for logging videos, images, or other more complex information, is to save that information to disk, and then process, log, and delete it inside of on_train_log(), but perhaps there's a cleaner solution.

[lgtm-com]

This pull request introduces 2 alerts when merging bc49d47 into cc0d123 - view on LGTM.com

new alerts:

• 2 for Variable defined multiple times

[jordis-ai2]

I like a lot the idea. Just wondering if we should actually include it in the experiment config, but other than that, and some picky details, LGTM.

[jordis-ai2]

Throughout the code base we try to keep third party imports before allenact ones, so I would move this a few lines below.

[Lucaweihs]

If you're using PyCharm I think this would be handled by running Code -> Optimize Imports.

[mattdeitke]

Ah, I will revert these import updates.

I am using isort, which works well with Black and is pretty popular for sorting and organizing imports. It autoformats on save for VSCode so it ended up changing the import order automatically.

[jordis-ai2]

I see many formatting changes - are you also using black to ensure consistency?

[Lucaweihs]

Yes, note that we're using version 19.10b0 of black.

[mattdeitke]

Yes, using Black, but must be using some different preferences with it (such as the number of characters in a line). I will try with 19.10b0 :)

[Lucaweihs]

Really like this but do have a few suggestions here and there that I think could improve usability. I think the main question for us is if we see people wanting to mix and match callbacks across repeated runs of the same experiment. If not then maybe its sufficient to define callbacks at the experiment config level. One important thing to remember for this discussion is that it need not be the same person running the experiment and so if we "hard code" certain callbacks (e.g. wandb logging) then this will make it a bit more annoying for people who don't have all the appropriate permissions set up.

[Lucaweihs]

If you're using PyCharm I think this would be handled by running Code -> Optimize Imports.

[Lucaweihs]

Yes, note that we're using version 19.10b0 of black.

[Lucaweihs]

I would prefer if tasks didn't know about callbacks as this creates a bidirectional dependency between the tasks and callbacks. Could we instead pass a function (or callable object) to the VectorSampledTask and then do something like task_callback_data = task_callback_data_fn(current_task). This task_callback_data_fn could be returned by the experiment config similarly as to how make_sampler_fn is.

[Lucaweihs]

Actually even better would be if each Callback was required to have this function defined on itself as a static method and we just passed in a list of these methods (that way the experiment config wouldn't need to know which callbacks were going to be used with it making it easier to mix and match from the command line).

[mattdeitke]

That’s an interesting idea! Like an on_task_end(task: Task) callback.

[Lucaweihs]

If doing my above suggestion then this would be removed.

[jordis-ai2]

Just a minor comment and a short question about usage, bot other than that LGTM.

[jordis-ai2]

This is nice - not sure why it was never added as an argument, but good you did. I guess it would also make sense to make it an arg in main?

[mattdeitke]

Sorry, have just been periodically adding to this branch for any issues I’ve come across. I think I’m just about ready for it, other than addressing the existing comments above. I’ve been waiting a bit since it’s nice to sometimes add arguments to the callbacks, but if it’s merged into main, it’d be hard to add anything without breaking backwards compatibility.

[jordis-ai2]

I meant main.py :)

[jordis-ai2]

I guess it's also possible to make the callback start a thread and return immediately to allow logger to keep showing stats for train in "real time". Is that right?

[mattdeitke]

This would be possible, but one of the issues is that if wandb is initialized in one process, you cannot log from a thread, and hence why logging from a thread and then destroying it wouldn’t work as one would expect. If you try to log from all threads and processes, logging becomes prohibitively expensive.

See their notes on distributed training: https://docs.wandb.ai/guides/track/advanced/distributed-training

[jordis-ai2]

I think I didn't express correctly what I meant. Currently, if the callback function takes a long time to process, packages in the runner queue for logging wait for a long time. When they're read, we quickly flush that queue, resulting in e.g. spikes of FPS (note that there's no timestamp in the logging packages sent from the trainers). I thought that passing the data for the callback to a thread (in the same process) could do the job, but if using an open wandb session from a thread in the same process is a no go, then there's no discussion. 👍

[Lucaweihs]

Having merged the callback sensor API, LGTM.