Coverage for sacred/sacred/observers/telegram_obs.py: 17%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

105 statements  

1#!/usr/bin/env python 

2# coding=utf-8 

3 

4from sacred.observers.base import RunObserver, td_format 

5from sacred.config.config_files import load_config_file 

6import logging 

7 

8 

9DEFAULT_TELEGRAM_PRIORITY = 10 

10 

11 

12class TelegramObserver(RunObserver): 

13 """Sends a message to Telegram upon completion/failing of an experiment.""" 

14 

15 @staticmethod 

16 def get_proxy_request(telegram_config): 

17 from telegram.utils.request import Request 

18 

19 if telegram_config["proxy_url"].startswith("socks5"): 

20 urllib3_proxy_kwargs = dict() 

21 for key in ["username", "password"]: 

22 if key in telegram_config: 

23 urllib3_proxy_kwargs[key] = telegram_config[key] 

24 return Request( 

25 proxy_url=telegram_config["proxy_url"], 

26 urllib3_proxy_kwargs=urllib3_proxy_kwargs, 

27 ) 

28 elif telegram_config["proxy_url"].startswith("http"): 

29 cred_string = "" 

30 if "username" in telegram_config: 

31 cred_string += telegram_config["username"] 

32 if "password" in telegram_config: 

33 cred_string += ":" + telegram_config["password"] 

34 if len(cred_string) > 0: 

35 domain = telegram_config["proxy_url"].split("/")[-1].split("@")[-1] 

36 cred_string += "@" 

37 proxy_url = "http://{}{}".format(cred_string, domain) 

38 return Request(proxy_url=proxy_url) 

39 else: 

40 return Request(proxy_url=telegram_config["proxy_url"]) 

41 else: 

42 raise Exception( 

43 "Proxy URL should be in format " 

44 "PROTOCOL://PROXY_HOST[:PROXY_PORT].\n" 

45 "HTTP and Socks5 are supported." 

46 ) 

47 

48 @classmethod 

49 def from_config(cls, filename): 

50 """ 

51 Create a TelegramObserver from a given configuration file. 

52 

53 The file can be in any format supported by Sacred 

54 (.json, .pickle, [.yaml]). 

55 It has to specify a ``token`` and a ``chat_id`` and can optionally set 

56 ``silent_completion``, ``started_text``, ``completed_text``, 

57 ``interrupted_text``, and ``failed_text``. 

58 """ 

59 import telegram 

60 

61 d = load_config_file(filename) 

62 request = cls.get_proxy_request(d) if "proxy_url" in d else None 

63 

64 if "token" in d and "chat_id" in d: 

65 bot = telegram.Bot(d["token"], request=request) 

66 obs = cls(bot, **d) 

67 else: 

68 raise ValueError( 

69 "Telegram configuration file must contain " 

70 "entries for 'token' and 'chat_id'!" 

71 ) 

72 for k in ["started_text", "completed_text", "interrupted_text", "failed_text"]: 

73 if k in d: 

74 setattr(obs, k, d[k]) 

75 return obs 

76 

77 def __init__( 

78 self, 

79 bot, 

80 chat_id, 

81 silent_completion=False, 

82 priority=DEFAULT_TELEGRAM_PRIORITY, 

83 **kwargs, 

84 ): 

85 self.silent_completion = silent_completion 

86 self.chat_id = chat_id 

87 self.bot = bot 

88 

89 self.started_text = ( 

90 "♻ *{experiment[name]}* " 

91 "started at _{start_time}_ " 

92 "on host `{host_info[hostname]}`" 

93 ) 

94 self.completed_text = ( 

95 "✅ *{experiment[name]}* " 

96 "completed after _{elapsed_time}_ " 

97 "with result=`{result}`" 

98 ) 

99 self.interrupted_text = ( 

100 "⚠ *{experiment[name]}* " "interrupted after _{elapsed_time}_" 

101 ) 

102 self.failed_text = ( 

103 "❌ *{experiment[name]}* failed after " 

104 "_{elapsed_time}_ with `{error}`\n\n" 

105 "Backtrace:\n```{backtrace}```" 

106 ) 

107 self.run = None 

108 self.priority = priority 

109 

110 def started_event( 

111 self, ex_info, command, host_info, start_time, config, meta_info, _id 

112 ): 

113 import telegram 

114 

115 self.run = { 

116 "_id": _id, 

117 "config": config, 

118 "start_time": start_time, 

119 "experiment": ex_info, 

120 "command": command, 

121 "host_info": host_info, 

122 } 

123 if self.started_text is None: 

124 return 

125 try: 

126 self.bot.send_message( 

127 chat_id=self.chat_id, 

128 text=self.get_started_text(), 

129 disable_notification=True, 

130 parse_mode=telegram.ParseMode.MARKDOWN, 

131 ) 

132 except Exception as e: 

133 log = logging.getLogger("telegram-observer") 

134 log.warning("failed to send start_event message via telegram.", exc_info=e) 

135 

136 def get_started_text(self): 

137 return self.started_text.format(**self.run) 

138 

139 def get_completed_text(self): 

140 return self.completed_text.format(**self.run) 

141 

142 def get_interrupted_text(self): 

143 return self.interrupted_text.format(**self.run) 

144 

145 def get_failed_text(self): 

146 return self.failed_text.format( 

147 backtrace="".join(self.run["fail_trace"]), **self.run 

148 ) 

149 

150 def completed_event(self, stop_time, result): 

151 import telegram 

152 

153 if self.completed_text is None: 

154 return 

155 

156 self.run["result"] = result 

157 self.run["stop_time"] = stop_time 

158 self.run["elapsed_time"] = td_format(stop_time - self.run["start_time"]) 

159 

160 try: 

161 self.bot.send_message( 

162 chat_id=self.chat_id, 

163 text=self.get_completed_text(), 

164 disable_notification=self.silent_completion, 

165 parse_mode=telegram.ParseMode.MARKDOWN, 

166 ) 

167 except Exception as e: 

168 log = logging.getLogger("telegram-observer") 

169 log.warning( 

170 "failed to send completed_event message via telegram.", exc_info=e 

171 ) 

172 

173 def interrupted_event(self, interrupt_time, status): 

174 import telegram 

175 

176 if self.interrupted_text is None: 

177 return 

178 

179 self.run["status"] = status 

180 self.run["interrupt_time"] = interrupt_time 

181 self.run["elapsed_time"] = td_format(interrupt_time - self.run["start_time"]) 

182 

183 try: 

184 self.bot.send_message( 

185 chat_id=self.chat_id, 

186 text=self.get_interrupted_text(), 

187 disable_notification=False, 

188 parse_mode=telegram.ParseMode.MARKDOWN, 

189 ) 

190 except Exception as e: 

191 log = logging.getLogger("telegram-observer") 

192 log.warning( 

193 "failed to send interrupted_event message " "via telegram.", exc_info=e 

194 ) 

195 

196 def failed_event(self, fail_time, fail_trace): 

197 import telegram 

198 

199 if self.failed_text is None: 

200 return 

201 

202 self.run["fail_trace"] = fail_trace 

203 self.run["error"] = fail_trace[-1].strip() 

204 self.run["fail_time"] = fail_time 

205 self.run["elapsed_time"] = td_format(fail_time - self.run["start_time"]) 

206 

207 try: 

208 self.bot.send_message( 

209 chat_id=self.chat_id, 

210 text=self.get_failed_text(), 

211 disable_notification=False, 

212 parse_mode=telegram.ParseMode.MARKDOWN, 

213 ) 

214 except Exception as e: 

215 log = logging.getLogger("telegram-observer") 

216 log.warning("failed to send failed_event message via telegram.", exc_info=e)