Fix/async callbacks #102
Fix/async callbacks #102
Conversation
- Split into seperate files - Use list in config to add callbacks - Provide legacy config enabled approach - Fix ruff issues
for more information, see https://pre-commit.ci
- Prefill config with callbacks - Warn on deprecations for old config - Expand config enabled - Add back SWA - Fix logging callback - Add flag to disable checkpointing - Add testing
[feature] Fix trainable attribute callbacks
- Split plots - Rename, lr to optimiser - Refactor plotting callbacks to be more init config
Refactor rollout logic
| self._plot_with_error_catching, | ||
| trainer, | ||
| *args, | ||
| *kwargs.values(), |
There was a problem hiding this comment.
This implies that kwargs will be provided in the correct order?
Why can't this be expanded normally?
There was a problem hiding this comment.
Or maybe something like this ?
def func(s, trainer, args, kwargs):
return s._plot_with_error_catching(trainer, *args, **kwargs)
await loop.run_in_executor(
self._executor,
func,
self,
trainer,
args,
kwargs,
I am not sure we need to pass self as an argument.
There was a problem hiding this comment.
So I had a better look into this. And what's happening is that the async function run_in_executor doesn't support keyword arguments, so using *kwargs.values() bypasses the need to pass keywords directly by transforming the keyword arguments into positional ones. In that way the values are then correctly used by _plot_with_error_catching as positional arguments, which matches the signature that accepts *args.
I have updated the code to reflect that! happy to discuss further!
There was a problem hiding this comment.
Let's discuss, I do like Florian's suggestion, it helps ensure that the kwargs are used correctly.
for more information, see https://pre-commit.ci
…' into fxi/async_callbacks
…g into fxi/async_callbacks
for more information, see https://pre-commit.ci
| self._plot_with_error_catching, | ||
| trainer, | ||
| *args, | ||
| *kwargs.values(), |
There was a problem hiding this comment.
Let's discuss, I do like Florian's suggestion, it helps ensure that the kwargs are used correctly.
| trainer, | ||
| *args, | ||
| **kwargs, | ||
| ) -> None: |
There was a problem hiding this comment.
Missing type hints
| trainer, | |
| *args, | |
| **kwargs, | |
| ) -> None: | |
| trainer: pl.Trainer, | |
| *args: Any, | |
| **kwargs: Any, | |
| ) -> None: |
| async def submit_plot(self, trainer, *args, **kwargs) -> None: | ||
| """Async function or coroutine to schedule the plot function.""" |
There was a problem hiding this comment.
| async def submit_plot(self, trainer, *args, **kwargs) -> None: | |
| """Async function or coroutine to schedule the plot function.""" | |
| async def submit_plot(self, trainer: pl.Trainer, *args: Any, **kwargs: Any) -> None: | |
| """Async function or coroutine to schedule the plot function.""" |
Adding asynchronous callbacks and scatter plots following the former PR by @anaprietonem in aifs-mono (https://github.com/ecmwf-lab/aifs-mono/pull/230/files#diff-817e70e1a151825208d5b55a0b204d08abe52afb9b675b60c48e85376357cc92)