colossalai

1
from abc import ABC
2

3
from coati.experience_maker import Experience
4

5

6
class TrainerCallback(ABC):
7
    """
8
    Base callback class. It defines the interface for callbacks.
9
    """
10

11
    def on_fit_start(self) -> None:
12
        pass
13

14
    def on_fit_end(self) -> None:
15
        pass
16

17
    def on_episode_start(self, episode: int) -> None:
18
        pass
19

20
    def on_episode_end(self, episode: int) -> None:
21
        pass
22

23
    def on_epoch_start(self, epoch: int) -> None:
24
        pass
25

26
    def on_epoch_end(self, epoch: int) -> None:
27
        pass
28

29
    def on_batch_start(self) -> None:
30
        pass
31

32
    def on_batch_end(self, metrics: dict, experience: Experience) -> None:
33
        pass
34

35
    def on_update_start(self) -> None:
36
        pass
37

38
    def on_update_end(self) -> None:
39
        pass
40

41

42
class MakerCallback(ABC):
43
    def on_loop_start(self) -> None:
44
        pass
45

46
    def on_loop_end(self) -> None:
47
        pass
48

49
    def on_make_experience_start(self) -> None:
50
        pass
51

52
    def on_make_experience_end(self, experience: Experience) -> None:
53
        pass
54

55
    def on_send_start(self) -> None:
56
        pass
57

58
    def on_send_end(self) -> None:
59
        pass
60

61
    def on_batch_start(self) -> None:
62
        pass
63

64
    def on_batch_end(self) -> None:
65
        pass
66