create metrax/logging package (#116)

* create metrax/logging

* format files

* only change src/metrax/logging

* only modify src/metrax/logging

* add tests

* fix dependency

* fix with statements lint

Author: jshin1394

pyproject.toml

@@ -21,6 +21,8 @@ dependencies = [
   "flax==0.11.1",
   "jax==0.6.2",
   "numpy==2.1.3",
+  "tensorboardX==2.6.4",
+  "wandb==0.22.3",
 ]
 
 [tool.hatch.build]

src/metrax/logging/__init__.py

@@ -0,0 +1,21 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from metrax.logging import protocol
+from metrax.logging import tensorboard_backend
+from metrax.logging import wandb_backend
+
+LoggingBackend = protocol.LoggingBackend
+TensorboardBackend = tensorboard_backend.TensorboardBackend
+WandbBackend = wandb_backend.WandbBackend

src/metrax/logging/protocol.py

@@ -0,0 +1,30 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Metrax LoggingBackend protocol."""
+
+from typing import Protocol
+import numpy as np
+
+
+class LoggingBackend(Protocol):
+  """Defines the interface for a pluggable logging backend."""
+
+  def log_scalar(self, event: str, value: float | np.ndarray, **kwargs):
+    """Logs a scalar value. Must match jax.monitoring listener signature."""
+    ...
+
+  def close(self):
+    """Closes the logger and flushes any pending data."""
+    ...

src/metrax/logging/tensorboard_backend.py

@@ -0,0 +1,56 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Metrax LoggingBackend implemenation for Tensorboard."""
+
+import jax
+import numpy as np
+from tensorboardX import writer
+
+_DEFAULT_STEP = 0
+
+
+def _get_step(kwargs: dict[str, str | int]) -> int:
+  """Returns the step from the kwargs, or 0 if not provided."""
+  step = kwargs.get("step")
+  return _DEFAULT_STEP if step is None else int(step)
+
+
+def _preprocess_event_name(event_name: str) -> str:
+  """Preprocesses the event name before logging."""
+  return event_name.lstrip("/")  # Remove leading slashes
+
+
+class TensorboardBackend:
+  """A logging backend for Tensorboard that conforms to the LoggingBackend protocol."""
+
+  def __init__(self, log_dir: str, flush_every_n_steps: int = 100):
+    self._flush_every_n_steps = flush_every_n_steps
+    if jax.process_index() == 0:
+      self._writer = writer.SummaryWriter(logdir=log_dir)
+    else:
+      self._writer = None
+
+  def log_scalar(self, event: str, value: float | np.ndarray, **kwargs):
+    if self._writer is None:
+      return
+    current_step = _get_step(kwargs)
+    event_name = _preprocess_event_name(event)
+    self._writer.add_scalar(event_name, value, current_step)
+    if current_step % self._flush_every_n_steps == 0:
+      self._writer.flush()
+
+  def close(self):
+    if self._writer:
+      self._writer.close()

src/metrax/logging/tensorboard_backend_test.py

@@ -0,0 +1,64 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest import mock
+
+from absl.testing import absltest
+from metrax import logging as metrax_logging
+
+TensorboardBackend = metrax_logging.TensorboardBackend
+
+
+class TensorboardBackendTest(absltest.TestCase):
+
+  @mock.patch("metrax.logging.tensorboard_backend.writer.SummaryWriter")
+  def test_init_and_log_success_main_process(self, mock_summary_writer):
+    """Tests successful init, logging, and closing on the main process."""
+    mock_writer_instance = mock_summary_writer.return_value
+
+    with mock.patch("jax.process_index", return_value=0):
+      backend = TensorboardBackend(log_dir="/fake/logs", flush_every_n_steps=2)
+
+      mock_summary_writer.assert_called_once_with(logdir="/fake/logs")
+
+      backend.log_scalar("/event1", 1.0, step=1)
+      mock_writer_instance.add_scalar.assert_called_with("event1", 1.0, 1)
+      mock_writer_instance.flush.assert_not_called()
+
+      backend.log_scalar("event2", 2.0, step=2)
+      mock_writer_instance.add_scalar.assert_called_with("event2", 2.0, 2)
+      mock_writer_instance.flush.assert_called_once()
+
+      backend.close()
+      mock_writer_instance.close.assert_called_once()
+
+  @mock.patch("metrax.logging.tensorboard_backend.writer.SummaryWriter")
+  def test_init_non_main_process_is_noop(self, mock_summary_writer):
+    """Tests that the backend does nothing on non-main processes."""
+    mock_writer_instance = mock_summary_writer.return_value
+
+    with mock.patch("jax.process_index", return_value=1):
+      backend = TensorboardBackend(log_dir="/fake/logs")
+
+      mock_summary_writer.assert_not_called()
+
+      backend.log_scalar("myevent", 1.0, step=1)
+      mock_writer_instance.add_scalar.assert_not_called()
+
+      backend.close()
+      mock_writer_instance.close.assert_not_called()
+
+
+if __name__ == "__main__":
+  absltest.main()

src/metrax/logging/wandb_backend.py

@@ -0,0 +1,76 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Metrax LoggingBackend implemenation for Weight & Bias."""
+
+import datetime
+import logging
+
+import jax
+import numpy as np
+
+_DEFAULT_STEP = 0
+
+
+def _get_step(kwargs: dict[str, str | int]) -> int:
+  """Returns the step from the kwargs, or 0 if not provided."""
+  step = kwargs.get("step")
+  return _DEFAULT_STEP if step is None else int(step)
+
+
+def _preprocess_event_name(event_name: str) -> str:
+  """Preprocesses the event name before logging."""
+  return event_name.lstrip("/")  # Remove leading slashes
+
+
+class WandbBackend:
+  """A logging backend for W&B that conforms to the LoggingBackend protocol."""
+
+  def __init__(self, project: str, name: str | None = None, **kwargs):
+    if jax.process_index() != 0:
+      self._is_active = False
+      self.wandb = None  # Ensure the attribute exists
+      return
+
+    try:
+      # pylint: disable=g-import-not-at-top
+      # pytype: disable=import-error
+      import wandb
+    except ImportError as e:
+      raise ImportError(
+          "The 'wandb' library is not installed. Please install it with "
+          "'pip install wandb' to use the WandbBackend."
+      ) from e
+    self.wandb = wandb
+
+    run_name = name or datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    wandb.init(project=project, name=run_name, anonymous="allow", **kwargs)
+    if wandb.run:
+      logging.info("W&B run URL: %s", wandb.run.url)
+      self._is_active = True
+    else:
+      self._is_active = False
+
+  def log_scalar(self, event: str, value: float | np.ndarray, **kwargs):
+    if self.wandb is None or not self._is_active:
+      return
+    current_step = _get_step(kwargs)
+    event_name = _preprocess_event_name(event)
+    self.wandb.log({event_name: value}, step=current_step)
+
+  def close(self):
+    if self.wandb is None or not self._is_active:
+      return
+    if hasattr(self.wandb, "run") and self.wandb.run:
+      self.wandb.finish()

src/metrax/logging/wandb_backend_test.py

@@ -0,0 +1,77 @@
+"""Tests for the W&B backend."""
+
+import builtins
+from unittest import mock
+
+from absl.testing import absltest
+from metrax import logging as metrax_logging
+
+WandbBackend = metrax_logging.WandbBackend
+
+
+_real_import = builtins.__import__
+
+
+class WandbBackendTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.mock_wandb = mock.Mock()
+    self.mock_wandb.run = mock.Mock()
+    self.mock_datetime = mock.Mock()
+    self.mock_datetime.datetime.now.return_value.strftime.return_value = (
+        "fixed-run-name"
+    )
+
+  def _mock_successful_import(self, name, *args, **kwargs):
+    """Mock __import__ to return our mock_wandb for 'wandb'."""
+    if name == "wandb":
+      return self.mock_wandb
+    return _real_import(name, *args, **kwargs)
+
+  def test_init_and_log_success_main_process(self):
+    """Tests successful init, logging, and closing on the main process."""
+    with mock.patch("jax.process_index", return_value=0), mock.patch(
+        "metrax.logging.wandb_backend.datetime", self.mock_datetime
+    ), mock.patch(
+        "builtins.__import__", side_effect=self._mock_successful_import
+    ):
+
+      backend = WandbBackend(project="test-project")
+      self.mock_wandb.init.assert_called_once_with(
+          project="test-project", name="fixed-run-name", anonymous="allow"
+      )
+      self.assertTrue(backend._is_active)
+
+      backend.log_scalar("/myevent", 123.45, step=50)
+      self.mock_wandb.log.assert_called_once_with({"myevent": 123.45}, step=50)
+
+      backend.close()
+      self.mock_wandb.finish.assert_called_once()
+
+  def test_init_non_main_process_is_noop(self):
+    """Tests that the backend does nothing on non-main processes."""
+    with mock.patch("jax.process_index", return_value=1):
+      backend = WandbBackend(project="test-project")
+      self.assertFalse(backend._is_active)
+      self.assertIsNone(backend.wandb)
+
+  def test_init_fails_if_wandb_not_installed(self):
+    """Tests that __init__ raises an ImportError if wandb is missing."""
+
+    def failing_import(name, *args, **kwargs):
+      """Mock __import__ to raise an error for 'wandb'."""
+      if name == "wandb":
+        raise ImportError("Mocked import failure")
+      return _real_import(name, *args, **kwargs)
+
+    with mock.patch(
+        "builtins.__import__", side_effect=failing_import
+    ), mock.patch("jax.process_index", return_value=0):
+      with self.assertRaises(ImportError) as cm:
+        WandbBackend(project="test-project")
+      self.assertIn("pip install wandb", str(cm.exception))
+
+
+if __name__ == "__main__":
+  absltest.main()