feat: change log time for wandb

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

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.

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.

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", {}))

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)

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

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": {