Spaces:
Build error
Build error
| # generates periodic hearbeats for remote expriment monitoring | |
| from pathlib import Path | |
| import json | |
| from inspect import stack | |
| from .ticker import IntervalTicker | |
| _CURRENT_BEAT_STACK = [] | |
| def get_heartbeat(): | |
| """ | |
| Returns: | |
| The :class:`HeartBeat` object that's currently being used. | |
| Throws an error if no :class:`EventStorage` is currently enabled. | |
| """ | |
| assert len( | |
| _CURRENT_BEAT_STACK | |
| ), "get_heartbeat() has to be called inside a 'with EventStorage(...)' context!" | |
| return _CURRENT_BEAT_STACK[-1] | |
| def get_tqdm_meter(pbar, format_dict): | |
| format_dict['bar_format'] = "{r_bar}" | |
| meter_str = pbar.format_meter(**format_dict) | |
| meter_str = meter_str[2:] | |
| return meter_str | |
| def caller_info(n_stack_up): | |
| info = stack()[1 + n_stack_up] # 1 up as base so that it starts from caller | |
| msg = f"{info.filename}:{info.lineno} - {info.function}" | |
| return msg | |
| class HeartBeat(): | |
| def __init__( | |
| self, pbar, write_interval=10, | |
| output_dir="./", fname="heartbeat.json" | |
| ): | |
| self.pbar = pbar | |
| self.fname = Path(output_dir) / fname | |
| self.ticker = IntervalTicker(write_interval) | |
| self.completed = False | |
| # force one write at the beginning | |
| self.beat(force_write=True, n_stack_up=2) | |
| def beat(self, force_write=False, n_stack_up=1): | |
| on_write_period = self.ticker.tick() | |
| if force_write or on_write_period: | |
| stats = self.stats() | |
| stats['caller'] = caller_info(n_stack_up) | |
| with open(self.fname, "w") as f: | |
| json.dump(stats, f) | |
| def done(self): | |
| self.completed = True | |
| self.beat(force_write=True, n_stack_up=2) | |
| def stats(self): | |
| pbar = self.pbar | |
| fdict = pbar.format_dict | |
| stats = { | |
| "beat": self.ticker.tick_str(), | |
| "done": self.completed, | |
| "meter": get_tqdm_meter(pbar, fdict), | |
| "elapsed": int(fdict['elapsed']) | |
| } | |
| return stats | |
| def __enter__(self): | |
| _CURRENT_BEAT_STACK.append(self) | |
| return self | |
| def __exit__(self, exc_type, exc_val, exc_tb): | |
| assert _CURRENT_BEAT_STACK[-1] == self | |
| _CURRENT_BEAT_STACK.pop() | |