Coverage for trnbl\loggers\wandb.py: 0%

1import datetime 

2import json 

3import logging 

4from typing import Any 

5from pathlib import Path 

6import sys 

7 

8import wandb 

9from wandb.sdk.wandb_run import Run, Artifact 

10 

11from trnbl.loggers.base import TrainingLoggerBase 

12 

13 

14class WandbLogger(TrainingLoggerBase): 

15 """wrapper around wandb logging for `TrainingLoggerBase`. create using `WandbLogger.create(config, project, job_type)`""" 

16 

17 def __init__(self, run: Run): 

18 self._run: Run = run 

19 

20 @classmethod 

21 def create( 

22 cls, 

23 config: dict, 

24 project: str | None = None, 

25 job_type: str | None = None, 

26 logging_fmt: str = "%(asctime)s %(levelname)s %(message)s", 

27 logging_datefmt: str = "%Y-%m-%d %H:%M:%S", 

28 wandb_kwargs: dict | None = None, 

29 ) -> "WandbLogger": 

30 logging.basicConfig( 

31 stream=sys.stdout, 

32 level=logging.INFO, 

33 format=logging_fmt, 

34 datefmt=logging_datefmt, 

35 ) 

36 

37 run: Run # type: ignore[return-value] 

38 run = wandb.init( 

39 config=config, 

40 project=project, 

41 job_type=job_type, 

42 **(wandb_kwargs if wandb_kwargs else {}), 

43 ) 

44 

45 assert run is not None, f"wandb.init returned None: {wandb_kwargs}" 

46 

47 logger: WandbLogger = WandbLogger(run) 

48 # TODO: why are we ignoring type checking here? 

49 logger.progress(f"{config =}") # type: ignore[attr-defined] 

50 return logger 

51 

52 def debug(self, message: str, **kwargs) -> None: 

53 if kwargs: 

54 message += f" {kwargs =}" 

55 logging.debug(message) 

56 

57 def message(self, message: str, **kwargs) -> None: 

58 if kwargs: 

59 message += f" {kwargs =}" 

60 logging.info(message) 

61 

62 def metrics(self, data: dict[str, Any]) -> None: 

63 self._run.log(data) 

64 

65 def artifact( 

66 self, 

67 path: Path, 

68 type: str, 

69 aliases: list[str] | None = None, 

70 metadata: dict | None = None, 

71 ) -> None: 

72 artifact: Artifact = wandb.Artifact(name=self._run.id, type=type) 

73 artifact.add_file(str(path)) 

74 if metadata: 

75 artifact.description = json.dumps( 

76 dict( 

77 timestamp=datetime.datetime.now().isoformat(), 

78 path=path.as_posix(), 

79 type=type, 

80 aliases=aliases, 

81 metadata=metadata if metadata else {}, 

82 ) 

83 ) 

84 self._run.log_artifact(artifact, aliases=aliases) 

85 

86 @property 

87 def url(self) -> str: 

88 # TODO: get_url returns `None` for offline runs. need to adjust allowed return types in superclass 

89 return str(self._run.get_url()) 

90 

91 @property 

92 def run_path(self) -> Path: 

93 return Path(self._run._get_path()) 

94 

95 def flush(self) -> None: 

96 self._run.save() 

97 

98 def finish(self) -> None: 

99 """Finish logging""" 

100 self._run.finish()