425 fix progress bars on jupyter (#426)

* fix(events): Add sync handlers in event dispatcher

- show_progress as sync to fix broken progress bars in jupyter environments
- Introduce HandlerMetadata with sync flag in
- registry Allow decorator usage with
- @..._event_handler(sync=True|False) Execute sync
- handlers on calling thread; queue only async
- ones Submit only async handlers to thread pool;
- adjust shutdown Add unregister_handler API and
- update tests accordingly Mark progress
- test verifying sync vs async execution threads

* feat(events): Emit COMMAND_FINISHED and reset

 progress

- Introduce EventType.COMMAND_FINISHED
- Emit command.finished when engine signals exit Stop and recreate Progress on command finish
- clear bars Reconstruct renderable only in textual output to avoid control chars
- Force monochrome output in VSCode terminals
- Rename new_task_with_reconstructed_renderable to new_task_and_reconstruct_renderable_maybe

* Fix union type in handler;

use sum generic in test_events

Author: srnnkls

src/python-api/getml/communication.py

@@ -43,10 +43,12 @@
     PoolEventDispatcher,
 )
 from getml.events.types import (
+    Event,
     EventContext,
     EventEmitter,
     EventParser,
     EventSource,
+    EventType,
 )
 from getml.exceptions import handle_engine_exception
 from getml.helpers import _is_iterable_not_str
@@ -95,6 +97,12 @@ def listen(self, socket: socket.socket, exit_on: Callable[[str], bool]) -> str:
                 self.emitter.emit(events)
 
             if exit_on(exit_status := message):
+                finished_event = Event(
+                    source=self.parser.context.source,
+                    type=EventType.COMMAND_FINISHED,
+                    attributes={},
+                )
+                self.emitter.emit([finished_event])
                 return exit_status
 
 

src/python-api/getml/events/dispatchers.py

@@ -16,9 +16,9 @@
 import warnings
 from concurrent.futures import Executor, Future, ThreadPoolExecutor
 from types import TracebackType
-from typing import Callable, Dict, Optional, Sequence, Type, Union
+from typing import Dict, Optional, Sequence, Type, Union
 
-from getml.events.handlers import EventHandlerRegistry
+from getml.events.handlers import EventHandlerRegistry, HandlerMetadata
 from getml.events.types import Event, EventSource, Shutdown
 
 
@@ -40,8 +40,8 @@ def __init__(
     ):
         self.queues = {
             source: {
-                handler: queue.Queue()
-                for handler in EventHandlerRegistry.handlers[source]
+                metadata: queue.Queue()
+                for metadata in EventHandlerRegistry.handlers[source]
             }
             for source in EventSource
         }
@@ -50,7 +50,7 @@ def __init__(
             executor = ThreadPoolExecutor()
         self.executor = executor
 
-        self.futures: Dict[Callable[[Event], None], Future] = {}
+        self.futures: Dict[HandlerMetadata, Future] = {}
 
     def __enter__(self):
         self.start()
@@ -69,12 +69,23 @@ def __exit__(
 
     def dispatch(self, events: Sequence[Event]):
         for event in events:
-            for queue in self.queues[event.source].values():
-                queue.put(event)
+            for metadata, handler_queue in self.queues[event.source].items():
+                if metadata.sync:
+                    # Execute synchronously in the calling thread
+                    try:
+                        metadata.handler(event)  # type: ignore
+                    except Exception as e:
+                        warnings.warn(
+                            f"An exception occurred while dispatching event {event} to handler {metadata.handler}: {e}",
+                            RuntimeWarning,
+                        )
+                else:
+                    # Dispatch to worker thread via queue
+                    handler_queue.put(event)
 
     def process(
         self,
-        handler: Callable[[Event], None],
+        metadata: HandlerMetadata,
         handler_queue: queue.Queue[Union[Event, Shutdown]],
     ):
         while True:
@@ -88,29 +99,33 @@ def process(
                 return
 
             try:
-                handler(event)  # type: ignore
+                metadata.handler(event)  # type: ignore
             except Exception as e:
                 warnings.warn(
-                    f"An exception occurred while dispatching event {event} to handler {handler}: {e}",
+                    f"An exception occurred while dispatching event {event} to handler {metadata.handler}: {e}",
                     RuntimeWarning,
                 )
             finally:
                 handler_queue.task_done()
 
     def start(self):
         for source in self.queues:
-            for handler, queue in self.queues[source].items():
-                self.futures[handler] = self.executor.submit(
-                    self.process, handler, queue
-                )
+            for metadata, handler_queue in self.queues[source].items():
+                if not metadata.sync:
+                    # Only submit async handlers to the thread pool
+                    self.futures[metadata] = self.executor.submit(
+                        self.process, metadata, handler_queue
+                    )
 
     def stop(self, wait: bool = True):
         if not wait:
             self.executor.shutdown(wait=False)
             return
 
         for source in self.queues:
-            for queue in self.queues[source].values():
-                queue.join()
-                queue.put(Shutdown)
+            for metadata, handler_queue in self.queues[source].items():
+                if not metadata.sync:
+                    # Only wait for and shutdown async handlers
+                    handler_queue.join()
+                    handler_queue.put(Shutdown)
         self.executor.shutdown(wait=True)

src/python-api/getml/events/handlers.py

@@ -10,33 +10,93 @@
 Event system.
 """
 
-from typing import Callable, Dict, List
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional, TypeVar, Union, overload
 
 from getml.events.types import Event, EventSource
 
 EventHandler = Callable[[Event], None]
+H = TypeVar("H", bound=EventHandler)
+
+
+@dataclass(frozen=True)
+class HandlerMetadata:
+    """Metadata for an event handler."""
+
+    handler: EventHandler
+    sync: bool = False
 
 
 class EventHandlerRegistry:
-    handlers: Dict[EventSource, List[EventHandler]] = {}
+    handlers: Dict[EventSource, List[HandlerMetadata]] = {}
 
     def __init__(self, source: EventSource):
         self.source = source
         type(self).handlers[source] = []
 
-    def __call__(self, handler: EventHandler) -> EventHandler:
+    @overload
+    def __call__(self, handler: H) -> H: ...
+
+    @overload
+    def __call__(self, *, sync: bool = False) -> Callable[[H], H]: ...
+
+    def __call__(
+        self, handler: Optional[H] = None, *, sync: bool = False
+    ) -> Union[Callable[[H], H], H]:
         """
         Register an event handler for the source associated with the instance.
+
+        Can be used as a decorator with or without arguments:
+        - @event_handler
+        - @event_handler(sync=True)
+
+        Args:
+            handler: The event handler function
+            sync: If True, handler executes synchronously in the calling thread
+                  (no dispatching). If False, handler executes in a worker thread pool.
+
+        Returns:
+            The handler function (for decorator usage)
         """
-        self.register_handler(self.source, handler)
+        if handler is None:
+            # Called with arguments: @event_handler(sync=...)
+            def decorator(func: H) -> H:
+                self.register_handler(self.source, func, sync=sync)
+                return func
+
+            return decorator
+
+        # Called without arguments: @event_handler
+        self.register_handler(self.source, handler, sync=sync)
         return handler
 
     @classmethod
-    def register_handler(cls, source: EventSource, handler: EventHandler):
+    def register_handler(
+        cls, source: EventSource, handler: EventHandler, sync: bool = False
+    ):
         """
         Register an event handler for a specific source.
+
+        Args:
+            source: The event source
+            handler: The event handler function
+            sync: If True, handler executes synchronously in the calling thread
+        """
+        metadata = HandlerMetadata(handler=handler, sync=sync)
+        cls.handlers[source].append(metadata)
+
+    @classmethod
+    def unregister_handler(cls, source: EventSource, handler: EventHandler):
+        """
+        Unregister an event handler for a specific source.
+
+        Args:
+            source: The event source
+            handler: The event handler function to remove
         """
-        cls.handlers[source].append(handler)
+        cls.handlers[source] = [
+            metadata for metadata in cls.handlers[source] if metadata.handler != handler
+        ]
 
 
 engine_event_handler = EventHandlerRegistry(EventSource.ENGINE)

src/python-api/getml/events/types.py

@@ -30,6 +30,7 @@ class EventTypeState(str, Enum):
 
 
 class EventType(str, Enum):
+    COMMAND_FINISHED = "command.finished"
     PIPELINE_CHECK_STAGE_START = "pipeline.check.stage.start"
     PIPELINE_CHECK_STAGE_PROGRESS = "pipeline.check.stage.progress"
     PIPELINE_CHECK_STAGE_END = "pipeline.check.stage.end"

src/python-api/getml/utilities/progress.py

@@ -33,7 +33,7 @@
 from rich.style import Style
 
 from getml.events.handlers import engine_event_handler, monitor_event_handler
-from getml.events.types import Event, EventTypeState
+from getml.events.types import Event, EventType, EventTypeState
 
 
 def _is_emacs_kernel() -> bool:
@@ -89,6 +89,8 @@ def _should_enforce_textual_output() -> bool:
 def _should_enforce_monochrome_output() -> bool:
     if env_value := os.getenv("GETML_PROGRESS_FORCE_MONOCHROME_OUTPUT"):
         return env_value in _ENV_VAR_TRUTHY_VALUES
+    if _is_vscode():
+        return True
     return False
 
 
@@ -314,7 +316,7 @@ def new_task(
         )
         return task_id
 
-    def new_task_with_reconstructed_renderable(
+    def new_task_and_reconstruct_renderable_maybe(
         self,
         description: str,
         *,
@@ -327,7 +329,9 @@ def new_task_with_reconstructed_renderable(
         renderable to ensure that the progress bar is displayed correctly
         with interleaved `print` statements.
         """
-        self._reconstruct_renderable()
+        if _should_enforce_textual_output():
+            self._reconstruct_renderable()
+
         task_id = self.new_task(
             description, total=total, completed=completed, visible=visible
         )
@@ -426,17 +430,22 @@ def __init__(self, progress: Progress):
     @contextmanager
     def live_renderable(self, event: Event) -> Iterator[Live]:
         yield self.progress.live
-        if event.type.state is EventTypeState.END:
+        if event.type is EventType.COMMAND_FINISHED:
             self.progress.stop()
+            self.progress = Progress()
 
     def create_or_update_progress(self, event: Event):
         """
         Create a new progress task or update an existing one.
+
+        A new Progress is created lazily if all previous tasks have finished. Progresses
+        cannot be created eagerly, because construction a live renderable in a
+        textual environment results in control characters being printed to the output.
         """
 
         with self.live_renderable(event):
             if event.type.state is EventTypeState.START:
-                task_id = self.progress.new_task_with_reconstructed_renderable(
+                task_id = self.progress.new_task_and_reconstruct_renderable_maybe(
                     event.attributes["description"]
                 )
             elif event.type.state is EventTypeState.PROGRESS and self.progress.tasks:
@@ -451,8 +460,8 @@ def create_or_update_progress(self, event: Event):
 progress_event_handler = ProgressEventHandler(progress=Progress())
 
 
-@monitor_event_handler
-@engine_event_handler
+@monitor_event_handler(sync=True)
+@engine_event_handler(sync=True)
 def show_progress(event: Event):
     if DISABLE:
         return