diff --git a/autointent/_callbacks/base.py b/autointent/_callbacks/base.py index b021ad472..5bb04c52e 100644 --- a/autointent/_callbacks/base.py +++ b/autointent/_callbacks/base.py @@ -16,12 +16,13 @@ def __init__(self) -> None: pass @abstractmethod - def start_run(self, run_name: str, dirpath: Path) -> None: + def start_run(self, run_name: str, dirpath: Path, log_interval_time: float) -> None: """Start a new run. Args: run_name: Name of the run. dirpath: Path to the directory where the logs will be saved. + log_interval_time: Sampling interval for the system monitor in seconds. """ @abstractmethod diff --git a/autointent/_callbacks/callback_handler.py b/autointent/_callbacks/callback_handler.py index aff8d1a23..d89429ccc 100644 --- a/autointent/_callbacks/callback_handler.py +++ b/autointent/_callbacks/callback_handler.py @@ -21,14 +21,15 @@ def __init__(self, callbacks: list[type[OptimizerCallback]] | None = None) -> No self.callbacks = [cb() for cb in callbacks] - def start_run(self, run_name: str, dirpath: Path) -> None: + def start_run(self, run_name: str, dirpath: Path, log_interval_time: float) -> None: """Start a new run. Args: run_name: Name of the run. dirpath: Path to the directory where the logs will be saved. + log_interval_time: Sampling interval for the system monitor in seconds. """ - self.call_events("start_run", run_name=run_name, dirpath=dirpath) + self.call_events("start_run", run_name=run_name, dirpath=dirpath, log_interval_time=log_interval_time) def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any]) -> None: """Start a new module. diff --git a/autointent/_callbacks/tensorboard.py b/autointent/_callbacks/tensorboard.py index 1639d6e45..0da46d6cf 100644 --- a/autointent/_callbacks/tensorboard.py +++ b/autointent/_callbacks/tensorboard.py @@ -31,15 +31,17 @@ def __init__(self) -> None: ) raise ImportError(msg) from None - def start_run(self, run_name: str, dirpath: Path) -> None: + def start_run(self, run_name: str, dirpath: Path, log_interval_time: float) -> None: """Starts a new run and sets the directory for storing logs. Args: run_name: Name of the run. dirpath: Path to the directory where logs will be saved. + log_interval_time: Sampling interval for the system monitor in seconds. """ self.run_name = run_name self.dirpath = dirpath + self.log_interval_time = log_interval_time def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any]) -> None: """Starts a new module and initializes a TensorBoard writer for it. diff --git a/autointent/_callbacks/wandb.py b/autointent/_callbacks/wandb.py index 3349e02c9..dd2b12f9d 100644 --- a/autointent/_callbacks/wandb.py +++ b/autointent/_callbacks/wandb.py @@ -33,7 +33,7 @@ def __init__(self) -> None: self.wandb = wandb - def start_run(self, run_name: str, dirpath: Path) -> None: + def start_run(self, run_name: str, dirpath: Path, log_interval_time: float) -> None: """Starts a new W&B run. Initializes the project name and run group. The directory path argument is not @@ -42,10 +42,12 @@ def start_run(self, run_name: str, dirpath: Path) -> None: Args: run_name: Name of the run (used as a W&B group). dirpath: Path to store logs (not utilized in W&B logging). + log_interval_time: Sampling interval for the system monitor in seconds. """ self.project_name = os.getenv("WANDB_PROJECT", "autointent") self.group = run_name self.dirpath = dirpath + self.log_interval_time = log_interval_time def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any]) -> None: """Starts a new module within the W&B logging system. @@ -63,6 +65,7 @@ def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any] group=self.group, name=f"{module_name}_{num}", config=module_kwargs, + settings=self.wandb.Settings(x_stats_sampling_interval=self.log_interval_time), ) def log_value(self, **kwargs: dict[str, Any]) -> None: @@ -96,6 +99,7 @@ def log_final_metrics(self, metrics: dict[str, Any]) -> None: group=self.group, name="final_metrics", config=metrics, + settings=self.wandb.Settings(x_stats_sampling_interval=self.log_interval_time), ) self.wandb.log(metrics.get("pipeline_metrics", {})) diff --git a/autointent/_pipeline/_pipeline.py b/autointent/_pipeline/_pipeline.py index cec6851fc..73a68bb06 100644 --- a/autointent/_pipeline/_pipeline.py +++ b/autointent/_pipeline/_pipeline.py @@ -147,6 +147,7 @@ def _fit(self, context: Context, sampler: SamplerType) -> None: self.context.callback_handler.start_run( run_name=self.context.logging_config.get_run_name(), dirpath=self.context.logging_config.dirpath, + log_interval_time=self.context.logging_config.log_interval_time, ) for node_type in NodeType: node_optimizer = self.nodes.get(node_type, None) diff --git a/autointent/configs/_optimization.py b/autointent/configs/_optimization.py index 990dbf334..c47eb6e91 100644 --- a/autointent/configs/_optimization.py +++ b/autointent/configs/_optimization.py @@ -51,6 +51,9 @@ class LoggingConfig(BaseModel): report_to: list[REPORTERS_NAMES] | None = Field( # type: ignore[valid-type] None, description="List of callbacks to report to. If None, no callbacks will be used" ) + log_interval_time: float = Field( + 0.1, description="Sampling interval for the system monitor in seconds for Wandb logger." + ) """List of callbacks to report to. If None, no callbacks will be used""" @property diff --git a/docs/optimizer_config.schema.json b/docs/optimizer_config.schema.json index aabd685ac..d478ecc41 100644 --- a/docs/optimizer_config.schema.json +++ b/docs/optimizer_config.schema.json @@ -297,6 +297,12 @@ "default": null, "description": "List of callbacks to report to. If None, no callbacks will be used", "title": "Report To" + }, + "log_interval_time": { + "default": 0.1, + "description": "Sampling interval for the system monitor in seconds for Wandb logger.", + "title": "Log Interval Time", + "type": "number" } }, "title": "LoggingConfig", @@ -328,7 +334,8 @@ "run_name": null, "dump_modules": false, "clear_ram": false, - "report_to": null + "report_to": null, + "log_interval_time": 0.1 } }, "embedder_config": {