interTwin-eu / itwinai

Advanced AI workflows for digital twins applications in science.
https://itwinai.readthedocs.io
MIT License
15 stars 5 forks source link

Distributed print decorator #222

Closed matbun closed 1 month ago

matbun commented 1 month ago

Update the monior_exec decorator to suppress prints on workers with global rank != 0 only when explicitly asked to do so.

Actually, the logic meant to suppress the print on workers with global rank != 0 has been moved to another decorator, which makes its usage straightforward. Example:

from itwinai.distributed import suppress_workers_print

class SilentComponent(BaseComponent):
    """Fake component class to test the decorator meant to suppress print()."""

    def __init__(self, max_items: int, name: str = 'SilentComponent') -> None:
        super().__init__(name)
        self.save_parameters(max_items=max_items, name=name)
        self.max_items = max_items

    @suppress_workers_print
    def execute(self, foo, bar=123):
        print("Executing SilentComponent")
        return "silent component result"

The two decorators also can be used together. Example:

class FakeComponent(BaseComponent):
    """Fake component class to test the use of two decorators at the same time."""

    def __init__(self, max_items: int, name: str = 'FakeComponent') -> None:
        super().__init__(name)
        self.save_parameters(max_items=max_items, name=name)
        self.max_items = max_items

    @suppress_workers_print
    @monitor_exec
    def execute(self, foo, bar):
        print("Executing FakeComponent")
        return "fake component result"
jarlsondre commented 1 month ago

Everything looks good to me :)