Use a static wandb run_id if none specified (#1175)

Author: cbalioglu

requirements-devel.txt

@@ -13,3 +13,4 @@ types-tqdm~=4.67.0
 types-editdistance~=0.6
 types-psutil~=5.9
 types-setuptools~=75.8
+wandb~=0.19

src/fairseq2/metrics/recorders/__init__.py

@@ -38,7 +38,6 @@
 from fairseq2.metrics.recorders._recorder import (
     NoopMetricRecorder as NoopMetricRecorder,
 )
-from fairseq2.metrics.recorders._recorder import record_metrics as record_metrics
 from fairseq2.metrics.recorders._tensorboard import (
     TENSORBOARD_RECORDER as TENSORBOARD_RECORDER,
 )

src/fairseq2/metrics/recorders/_composite.py

@@ -21,16 +21,11 @@ def __init__(self, recorders: Sequence[MetricRecorder]) -> None:
         self._inner_recorders = recorders
 
     @override
-    def record_metrics(
-        self,
-        section: str,
-        values: Mapping[str, object],
-        step_nr: int | None = None,
-        *,
-        flush: bool = True,
+    def record_metric_values(
+        self, section: str, values: Mapping[str, object], step_nr: int | None = None
     ) -> None:
         for recorder in self._inner_recorders:
-            recorder.record_metrics(section, values, step_nr, flush=flush)
+            recorder.record_metric_values(section, values, step_nr)
 
     @override
     def close(self) -> None:

src/fairseq2/metrics/recorders/_jsonl.py

@@ -61,13 +61,8 @@ def __init__(
         self._streams = {}
 
     @override
-    def record_metrics(
-        self,
-        section: str,
-        values: Mapping[str, object],
-        step_nr: int | None = None,
-        *,
-        flush: bool = True,
+    def record_metric_values(
+        self, section: str, values: Mapping[str, object], step_nr: int | None = None
     ) -> None:
         section = section.strip()
 
@@ -123,8 +118,7 @@ def sanitize(value: object, descriptor: MetricDescriptor) -> object:
 
             stream.write("\n")
 
-            if flush:
-                stream.flush()
+            stream.flush()
         except OSError as ex:
             raise MetricRecordError(
                 f"The metric values of the '{section}' cannot be saved to the JSON file. See the nested exception for details."

src/fairseq2/metrics/recorders/_log.py

@@ -42,13 +42,8 @@ def __init__(
         self._display_names = {"valid": "Validation", "eval": "Evaluation"}
 
     @override
-    def record_metrics(
-        self,
-        section: str,
-        values: Mapping[str, object],
-        step_nr: int | None = None,
-        *,
-        flush: bool = True,
+    def record_metric_values(
+        self, section: str, values: Mapping[str, object], step_nr: int | None = None
     ) -> None:
         if not self._log.is_enabled_for_info():
             return

src/fairseq2/metrics/recorders/_recorder.py

@@ -7,7 +7,7 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from collections.abc import Mapping, Sequence
+from collections.abc import Mapping
 from typing import final
 
 from typing_extensions import override
@@ -19,20 +19,14 @@ class MetricRecorder(Closable):
     """Records metric values."""
 
     @abstractmethod
-    def record_metrics(
-        self,
-        section: str,
-        values: Mapping[str, object],
-        step_nr: int | None = None,
-        *,
-        flush: bool = True,
+    def record_metric_values(
+        self, section: str, values: Mapping[str, object], step_nr: int | None = None
     ) -> None:
         """Record ``values``.
 
         :param section: The run section (e.g. 'train', 'eval').
         :param values: The metric values.
         :param step_nr: The step number of the run.
-        :param flush: If ``True``, flushes any buffers after recording.
         """
 
 
@@ -43,36 +37,11 @@ class MetricRecordError(Exception):
 @final
 class NoopMetricRecorder(MetricRecorder):
     @override
-    def record_metrics(
-        self,
-        section: str,
-        values: Mapping[str, object],
-        step_nr: int | None = None,
-        *,
-        flush: bool = True,
+    def record_metric_values(
+        self, section: str, values: Mapping[str, object], step_nr: int | None = None
     ) -> None:
         pass
 
     @override
     def close(self) -> None:
         pass
-
-
-def record_metrics(
-    recorders: Sequence[MetricRecorder],
-    section: str,
-    values: Mapping[str, object],
-    step_nr: int | None = None,
-    *,
-    flush: bool = True,
-) -> None:
-    """Record ``values`` to ``recorders``.
-
-    :param recorders: The recorders to record to.
-    :param section: The run section (e.g. 'train', 'eval').
-    :param values: The metric values.
-    :param step_nr: The step number of the run.
-    :param flush: If ``True``, flushes any buffers after recording.
-    """
-    for recorder in recorders:
-        recorder.record_metrics(section, values, step_nr, flush=flush)

src/fairseq2/metrics/recorders/_tensorboard.py

@@ -60,13 +60,8 @@ def __init__(
         self._writers = {}
 
     @override
-    def record_metrics(
-        self,
-        section: str,
-        values: Mapping[str, object],
-        step_nr: int | None = None,
-        *,
-        flush: bool = True,
+    def record_metric_values(
+        self, section: str, values: Mapping[str, object], step_nr: int | None = None
     ) -> None:
         writer = self._get_writer(section)
         if writer is None:
@@ -86,8 +81,7 @@ def record_metrics(
 
                 writer.add_scalar(display_name, value, step_nr)
 
-            if flush:
-                writer.flush()
+            writer.flush()
         except RuntimeError as ex:
             raise MetricRecordError(
                 f"The metric values of the '{section}' section cannot be saved to TensorBoard. See the nested exception for details."

src/fairseq2/metrics/recorders/_wandb.py

@@ -20,6 +20,7 @@
 else:
     _has_wandb = True
 
+from fairseq2.file_system import FileMode, FileSystem
 from fairseq2.logging import log
 from fairseq2.metrics import MetricDescriptor
 from fairseq2.registry import Provider
@@ -35,8 +36,6 @@
     NoopMetricRecorder,
 )
 
-WandbResume: TypeAlias = Literal["allow", "never", "auto"]
-
 
 @final
 class WandbRecorder(MetricRecorder):
@@ -57,13 +56,8 @@ def __init__(
         self._metric_descriptors = metric_descriptors
 
     @override
-    def record_metrics(
-        self,
-        section: str,
-        values: Mapping[str, object],
-        step_nr: int | None = None,
-        *,
-        flush: bool = True,
+    def record_metric_values(
+        self, section: str, values: Mapping[str, object], step_nr: int | None = None
     ) -> None:
         for name, value in values.items():
             try:
@@ -91,28 +85,37 @@ def close(self) -> None:
 WANDB_RECORDER: Final = "wandb"
 
 
+WandbResumeMode: TypeAlias = Literal["allow", "never", "must", "auto"]
+
+
 @dataclass(kw_only=True)
 class WandbRecorderConfig:
     enabled: bool = False
 
+    entity: str | None = None
+
     project: str | None = None
 
-    run_id: str | None = None
+    run_id: str | None = "auto"
 
     run_name: str | None = None
 
     group: str | None = None
 
     job_type: str | None = None
 
-    resume: WandbResume = "allow"
+    resume_mode: WandbResumeMode = "allow"
 
 
 @final
 class WandbRecorderHandler(MetricRecorderHandler):
+    _file_system: FileSystem
     _metric_descriptors: Provider[MetricDescriptor]
 
-    def __init__(self, metric_descriptors: Provider[MetricDescriptor]) -> None:
+    def __init__(
+        self, file_system: FileSystem, metric_descriptors: Provider[MetricDescriptor]
+    ) -> None:
+        self._file_system = file_system
         self._metric_descriptors = metric_descriptors
 
     @override
@@ -131,28 +134,32 @@ def create(
 
             return NoopMetricRecorder()
 
-        try:
-            hyper_params = unstructure(hyper_params)
-        except StructureError as ex:
-            raise ValueError(
-                "`hyper_params` cannot be unstructured. See the nested exception for details."
-            ) from ex
+        if hyper_params is not None:
+            try:
+                hyper_params = unstructure(hyper_params)
+            except StructureError as ex:
+                raise ValueError(
+                    "`hyper_params` cannot be unstructured. See the nested exception for details."
+                ) from ex
 
-        if not isinstance(hyper_params, dict):
-            raise TypeError(
-                f"The unstructured form of `hyper_params` must be of type `dict`, but is of type `{type(hyper_params)}` instead."
-            )
+            if not isinstance(hyper_params, dict):
+                raise TypeError(
+                    f"The unstructured form of `hyper_params` must be of type `dict`, but is of type `{type(hyper_params)}` instead."
+                )
+
+        run_id = self._get_run_id(output_dir, config)
 
         try:
             run = wandb.init(
+                entity=config.entity,
                 project=config.project,
                 dir=output_dir,
-                id=config.run_id,
+                id=run_id,
                 name=config.run_name,
                 config=hyper_params,
                 group=config.group,
                 job_type=config.job_type,
-                resume=config.resume,
+                resume=config.resume_mode,
             )
         except (RuntimeError, ValueError) as ex:
             raise MetricRecordError(
@@ -161,6 +168,43 @@ def create(
 
         return WandbRecorder(run, self._metric_descriptors)
 
+    def _get_run_id(self, output_dir: Path, config: WandbRecorderConfig) -> str:
+        run_id = config.run_id
+
+        if run_id is None:
+            return wandb.util.generate_id()
+
+        if run_id != "auto":
+            return run_id
+
+        wandb_file = output_dir.joinpath("wandb_run_id")
+
+        try:
+            fp = self._file_system.open_text(wandb_file)
+
+            with fp:
+                return fp.read()
+        except FileNotFoundError:
+            pass
+        except OSError as ex:
+            raise MetricRecordError(
+                "The Weights & Biases run ID cannot be loaded. See the nested exception for details."
+            ) from ex
+
+        run_id = wandb.util.generate_id()
+
+        try:
+            fp = self._file_system.open_text(wandb_file, mode=FileMode.WRITE)
+
+            with fp:
+                fp.write(run_id)
+        except OSError as ex:
+            raise MetricRecordError(
+                "The Weights & Biases run ID cannot be saved. See the nested exception for details."
+            ) from ex
+
+        return run_id
+
     @property
     @override
     def name(self) -> str:

src/fairseq2/recipes/_evaluator.py

@@ -352,7 +352,7 @@ def _publish_metrics(self, unit: EvalUnit[BatchT]) -> None:
                 section = f"{section}/{unit.name}"
 
             try:
-                self._metric_recorder.record_metrics(section, values)
+                self._metric_recorder.record_metric_values(section, values)
             except MetricRecordError as ex:
                 s = "evaluation"
 

src/fairseq2/recipes/_generator.py

@@ -284,7 +284,7 @@ def _publish_metrics(self) -> None:
             values["wall_time"] = self._wall_watch.get_elapsed_time()
 
             try:
-                self._metric_recorder.record_metrics("generation", values)
+                self._metric_recorder.record_metric_values("generation", values)
             except MetricRecordError as ex:
                 raise RecipeError(
                     "The generation metric values cannot recorded. See the nested exception for details."