This is the right time to use hooks! It would be very similar to this Memory profiling examples: https://kedro.readthedocs.io/en/latest/07_extend_kedro/02_hooks.html#hooks-examples