pytorch-lightning

1
import logging
2
import sys
3
import time
4

5
from lightning.app import LightningApp, LightningFlow, LightningWork
6

7
logger = logging.getLogger(__name__)
8
logger.setLevel(logging.INFO)
9
logger.addHandler(logging.StreamHandler(sys.stdout))
10

11

12
class SimpleWork(LightningWork):
13
    def __init__(self):
14
        super().__init__(cache_calls=False, parallel=True, raise_exception=False)
15
        self.is_running_now = False
16

17
    def run(self):
18
        self.is_running_now = True
19
        print("work_is_running")
20
        for i in range(1, 10):
21
            time.sleep(1)
22
            if i % 5 == 0:
23
                raise Exception(f"invalid_value_of_i_{i}")
24
            print(f"good_value_of_i_{i}")
25

26

27
class RootFlow(LightningFlow):
28
    def __init__(self):
29
        super().__init__()
30
        self.simple_work = SimpleWork()
31

32
    def run(self):
33
        print("useless_garbage_log_that_is_always_there_to_overload_logs")
34
        self.simple_work.run()
35
        if not self.simple_work.is_running_now:
36
            pass
37
            # work is not ready yet
38
            print("waiting_for_work_to_be_ready")
39
        else:
40
            print("flow_and_work_are_running")
41
            logger.info("logger_flow_work")
42
            time.sleep(0.1)
43

44

45
if __name__ == "__main__":
46
    app = LightningApp(RootFlow(), log_level="debug")
47