diff --git a/python/ray/_private/ray_logging.py b/python/ray/_private/ray_logging.py index 5db3805a8a58..5b3017c5ae5c 100644 --- a/python/ray/_private/ray_logging.py +++ b/python/ray/_private/ray_logging.py @@ -91,6 +91,28 @@ def setup_component_logger( return logger +def run_callback_on_events_in_ipython(event: str, cb: Callable): + """ + Register a callback to be run after each cell completes in IPython. + E.g.: + This is used to flush the logs after each cell completes. + + If IPython is not installed, this function does nothing. + + Args: + cb: The callback to run. + """ + try: + from IPython import get_ipython + + ipython = get_ipython() + # Register a callback on cell completion. + if ipython is not None: + ipython.events.register(event, cb) + except ImportError: + pass + + """ All components underneath here is used specifically for the default_worker.py. """ @@ -226,6 +248,8 @@ def __init__( self.recent: Dict[str, DedupState] = {} self.timesource = _timesource or (lambda: time.time()) + run_callback_on_events_in_ipython("post_execute", self.flush) + def deduplicate(self, batch: LogBatch) -> List[LogBatch]: """Rewrite a batch of lines to reduce duplicate log messages.