diff --git a/trinity/cli/launcher.py b/trinity/cli/launcher.py index 2e980632e5..2d72fb7920 100644 --- a/trinity/cli/launcher.py +++ b/trinity/cli/launcher.py @@ -103,6 +103,9 @@ def both(config: Config) -> None: logger.error("Evaluation failed.") raise e + ray.get(explorer.log_finalize.remote(step=explore_iter_num)) + ray.get(trainer.log_finalize.remote(step=train_iter_num)) + def activate_data_module(config_path: str): """Check whether to activate data module and preprocess datasets.""" diff --git a/trinity/explorer/explorer.py b/trinity/explorer/explorer.py index 84c7bfca30..f43d9c753d 100644 --- a/trinity/explorer/explorer.py +++ b/trinity/explorer/explorer.py @@ -247,3 +247,7 @@ def sync_weight(self) -> None: self._offline_weights_update() else: # online weights update self._online_weights_update() + + def log_finalize(self, step: int) -> None: + """Commit the logging results to wandb""" + self.monitor.log({"dummy_log_explorer": step}, step=step, commit=True) diff --git a/trinity/trainer/trainer.py b/trinity/trainer/trainer.py index dd5d184fad..b57f484950 100644 --- a/trinity/trainer/trainer.py +++ b/trinity/trainer/trainer.py @@ -105,6 +105,10 @@ def sync_weight(self) -> None: if self.config.synchronizer.sync_method == "online": self.engine.sync_weight() + def log_finalize(self, step: int) -> None: + """Commit the logging results to wandb""" + self.engine.logger.log({"dummy_log_trainer": step}, step=step, commit=True) + class TrainEngineWrapper(ABC): """A wrapper class to wrap various training engines.""" diff --git a/trinity/utils/monitor.py b/trinity/utils/monitor.py index dc86a9264d..2ff2a61a70 100644 --- a/trinity/utils/monitor.py +++ b/trinity/utils/monitor.py @@ -52,9 +52,9 @@ def calculate_metrics( metrics[key] = val return metrics - def log(self, data: dict, step: int) -> None: + def log(self, data: dict, step: int, commit: bool = False) -> None: """Log metrics.""" - self.logger.log(data, step=step) + self.logger.log(data, step=step, commit=commit) self.console_logger.info(f"Step {step}: {data}") def __del__(self) -> None: