Extract tensorboard callback logic in a separate module (#142)

Author: gsakkis

tests/models/test_pytorch_models.py

@@ -1,10 +1,11 @@
 """Tests for TileDB PyTorch model save and load."""
 
+import glob
 import inspect
 import os
 import pickle
 import platform
-import sys
+import shutil
 
 import pytest
 import torch
@@ -87,23 +88,19 @@ def forward(self, x):
         return x
 
 
-@pytest.mark.parametrize(
-    "optimizer",
-    [
-        getattr(optimizers, name)
-        for name, obj in inspect.getmembers(optimizers)
-        if inspect.isclass(obj) and name != "Optimizer"
-    ],
-)
-@pytest.mark.parametrize(
-    "net",
-    [
-        getattr(sys.modules[__name__], name)
-        for name, obj in inspect.getmembers(sys.modules[__name__])
-        if inspect.isclass(obj) and obj.__module__ == __name__
-    ],
-)
+net = pytest.mark.parametrize("net", [ConvNet, Net, SeqNeuralNetwork])
+
+
 class TestPyTorchModel:
+    @net
+    @pytest.mark.parametrize(
+        "optimizer",
+        [
+            getattr(optimizers, name)
+            for name, obj in inspect.getmembers(optimizers)
+            if inspect.isclass(obj) and name != "Optimizer"
+        ],
+    )
     def test_save(self, tmpdir, net, optimizer):
         EPOCH = 5
         LOSS = 0.4
@@ -139,7 +136,8 @@ def test_save(self, tmpdir, net, optimizer):
         ):
             assert all([a == b for a, b in zip(key_item_1[1], key_item_2[1])])
 
-    def test_preview(self, tmpdir, net, optimizer):
+    @net
+    def test_preview(self, tmpdir, net):
         # With model given as argument
         model = net()
         tiledb_array = os.path.join(tmpdir, "model_array")
@@ -148,7 +146,8 @@ def test_preview(self, tmpdir, net, optimizer):
         tiledb_obj_none = PyTorchTileDBModel(uri=tiledb_array, model=None)
         assert tiledb_obj_none.preview() == ""
 
-    def test_file_properties(self, tmpdir, net, optimizer):
+    @net
+    def test_file_properties(self, tmpdir, net):
         model = net()
         tiledb_array = os.path.join(tmpdir, "model_array")
         tiledb_obj = PyTorchTileDBModel(uri=tiledb_array, model=model)
@@ -165,38 +164,42 @@ def test_file_properties(self, tmpdir, net, optimizer):
         )
         assert tiledb_obj._file_properties["TILEDB_ML_MODEL_PREVIEW"] == str(model)
 
-    def test_tensorboard_callback_meta(self, tmpdir, net, optimizer, mocker):
+    @net
+    def test_tensorboard_callback_meta(self, tmpdir, net):
         model = net()
         tiledb_array = os.path.join(tmpdir, "model_array")
         tiledb_obj = PyTorchTileDBModel(uri=tiledb_array, model=model)
 
-        mocker.patch(
-            "tiledb.ml.models.pytorch.PyTorchTileDBModel._get_tensorboard_files",
-            return_value={
-                f"{tmpdir}/event_file_name_1": b"test_bytes_1",
-                f"{tmpdir}/event_file_name_2": b"test_bytes_2",
-            },
-        )
+        # SummaryWriter creates file(s) under log_dir
+        log_dir = os.path.join(tmpdir, "logs")
+        writer = SummaryWriter(log_dir=log_dir)
+        log_files = read_files(log_dir)
+        assert log_files
 
-        writer = SummaryWriter()
         tiledb_obj.save(update=False, summary_writer=writer)
-
         with tiledb.open(tiledb_array) as A:
-            assert len(pickle.loads(A.meta["__TENSORBOARD__"])) == 2
-            assert pickle.loads(A.meta["__TENSORBOARD__"]) == {
-                f"{tmpdir}/event_file_name_1": b"test_bytes_1",
-                f"{tmpdir}/event_file_name_2": b"test_bytes_2",
-            }
+            assert pickle.loads(A.meta["__TENSORBOARD__"]) == log_files
+        shutil.rmtree(log_dir)
 
         # Loading the event data should create local files
         tiledb_obj.load_tensorboard()
-        assert os.path.exists(f"{tmpdir}/event_file_name_1")
-        assert os.path.exists(f"{tmpdir}/event_file_name_2")
+        new_log_files = read_files(log_dir)
+        assert new_log_files == log_files
 
         custom_dir = os.path.join(tmpdir, "custom_log")
         tiledb_obj.load_tensorboard(target_dir=custom_dir)
-        assert os.path.exists(f"{custom_dir}/event_file_name_1")
-        assert os.path.exists(f"{custom_dir}/event_file_name_2")
+        new_log_files = read_files(custom_dir)
+        assert len(new_log_files) == len(log_files)
+        for new_file, old_file in zip(new_log_files.values(), log_files.values()):
+            assert new_file == old_file
+
+
+def read_files(dirpath):
+    files = {}
+    for path in glob.glob(f"{dirpath}/*"):
+        with open(path, "rb") as f:
+            files[path] = f.read()
+    return files
 
 
 class TestPyTorchModelCloud:

tests/models/test_tensorflow_keras_models.py

@@ -4,6 +4,7 @@
 import os
 import pickle
 import platform
+import shutil
 
 import numpy as np
 import pytest
@@ -531,36 +532,38 @@ def test_exception_raise_file_property_in_meta_error(self, tmpdir):
             ex.value
         )
 
-    def test_tensorboard_callback_meta(self, tmpdir, mocker):
+    def test_tensorboard_callback_meta(self, tmpdir):
         model = keras.models.Sequential()
         model.add(keras.layers.Flatten(input_shape=(10, 10)))
         tiledb_array = os.path.join(tmpdir, "model_array")
         tiledb_obj = TensorflowKerasTileDBModel(uri=tiledb_array, model=model)
 
         cb = [tf.keras.callbacks.TensorBoard(log_dir=tmpdir)]
 
-        mocker.patch(
-            "tiledb.ml.models.tensorflow_keras.TensorflowKerasTileDBModel._get_tensorboard_files",
-            return_value={
-                f"{tmpdir}/event_file_name_1": b"test_bytes_1",
-                f"{tmpdir}/event_file_name_2": b"test_bytes_2",
-            },
-        )
+        os.mkdir(os.path.join(tmpdir, "train"))
+        with open(os.path.join(tmpdir, "train", "foo_tfevents_1"), "wb") as f:
+            f.write(b"test_bytes_1")
+        with open(os.path.join(tmpdir, "train", "bar_tfevents_2"), "wb") as f:
+            f.write(b"test_bytes_2")
 
         tiledb_obj.save(include_callbacks=cb)
         with tiledb.open(tiledb_array) as A:
-            assert len(pickle.loads(A.meta["__TENSORBOARD__"])) == 2
             assert pickle.loads(A.meta["__TENSORBOARD__"]) == {
-                f"{tmpdir}/event_file_name_1": b"test_bytes_1",
-                f"{tmpdir}/event_file_name_2": b"test_bytes_2",
+                os.path.join(tmpdir, "train", "foo_tfevents_1"): b"test_bytes_1",
+                os.path.join(tmpdir, "train", "bar_tfevents_2"): b"test_bytes_2",
             }
+        shutil.rmtree(os.path.join(tmpdir, "train"))
 
         # Loading the event data should create local files
         tiledb_obj.load_tensorboard()
-        assert os.path.exists(f"{tmpdir}/event_file_name_1")
-        assert os.path.exists(f"{tmpdir}/event_file_name_2")
+        with open(os.path.join(tmpdir, "train", "foo_tfevents_1"), "rb") as f:
+            assert f.read() == b"test_bytes_1"
+        with open(os.path.join(tmpdir, "train", "bar_tfevents_2"), "rb") as f:
+            assert f.read() == b"test_bytes_2"
 
         custom_dir = os.path.join(tmpdir, "custom_log")
         tiledb_obj.load_tensorboard(target_dir=custom_dir)
-        assert os.path.exists(f"{custom_dir}/event_file_name_1")
-        assert os.path.exists(f"{custom_dir}/event_file_name_2")
+        with open(os.path.join(custom_dir, "foo_tfevents_1"), "rb") as f:
+            assert f.read() == b"test_bytes_1"
+        with open(os.path.join(custom_dir, "bar_tfevents_2"), "rb") as f:
+            assert f.read() == b"test_bytes_2"

tiledb/ml/models/_tensorboard.py

@@ -0,0 +1,33 @@
+import glob
+import os
+import pickle
+from typing import Mapping, Optional
+
+import tiledb
+
+from .base import Timestamp
+
+_KEY = "__TENSORBOARD__"
+
+
+def save_tensorboard(log_dir: str) -> Mapping[str, bytes]:
+    event_files = {}
+    for path in glob.glob(f"{log_dir}/*tfevents*"):
+        with open(path, "rb") as f:
+            event_files[path] = f.read()
+    return {_KEY: pickle.dumps(event_files, protocol=4)}
+
+
+def load_tensorboard(
+    uri: str,
+    ctx: Optional[tiledb.Ctx] = None,
+    target_dir: Optional[str] = None,
+    timestamp: Optional[Timestamp] = None,
+) -> None:
+    with tiledb.open(uri, ctx=ctx, timestamp=timestamp) as model_array:
+        for path, file_bytes in pickle.loads(model_array.meta[_KEY]).items():
+            log_dir = target_dir if target_dir else os.path.dirname(path)
+            if not os.path.exists(log_dir):
+                os.mkdir(log_dir)
+            with open(os.path.join(log_dir, os.path.basename(path)), "wb") as f:
+                f.write(file_bytes)

tiledb/ml/models/pytorch.py

@@ -1,7 +1,5 @@
 """Functionality for saving and loading PytTorch models as TileDB arrays"""
 
-import glob
-import os
 import pickle
 from typing import Any, Mapping, Optional
 
@@ -13,6 +11,7 @@
 import tiledb
 
 from ._cloud_utils import update_file_properties
+from ._tensorboard import load_tensorboard, save_tensorboard
 from .base import Meta, TileDBModel, Timestamp, current_milli_time
 
 
@@ -81,11 +80,8 @@ def save(
 
         # Summary writer
         if summary_writer:
-            event_files = self._get_tensorboard_files(summary_writer.log_dir)
-            meta = {
-                "__TENSORBOARD__": pickle.dumps(event_files, protocol=4),
-                **(meta or {}),
-            }
+            cb_meta = save_tensorboard(summary_writer.log_dir)
+            meta = {**meta, **cb_meta} if meta else cb_meta
 
         # Create TileDB model array
         if not update:
@@ -150,14 +146,7 @@ def load_tensorboard(
         target_dir: Optional[str] = None,
         timestamp: Optional[Timestamp] = None,
     ) -> None:
-        with tiledb.open(self.uri, ctx=self.ctx, timestamp=timestamp) as model_array:
-            tb_data = pickle.loads(model_array.meta["__TENSORBOARD__"])
-            for path in tb_data.keys():
-                log_dir = target_dir if target_dir else os.path.dirname(path)
-                if not os.path.exists(log_dir):
-                    os.mkdir(log_dir)
-                with open(os.path.join(log_dir, os.path.basename(path)), "wb") as f:
-                    f.write(tb_data[path])
+        return load_tensorboard(self.uri, self.ctx, target_dir, timestamp)
 
     def preview(self) -> str:
         """
@@ -250,11 +239,3 @@ def _write_array(
                 key: np.array([value]) for key, value in serialized_model_dict.items()
             }
             self.update_model_metadata(array=tf_model_tiledb, meta=meta)
-
-    @staticmethod
-    def _get_tensorboard_files(log_dir: str) -> Mapping[str, bytes]:
-        event_files = {}
-        for path in glob.glob(f"{log_dir}/*tfevents*"):
-            with open(path, "rb") as f:
-                event_files[path] = f.read()
-        return event_files

tiledb/ml/models/tensorflow_keras.py

@@ -1,11 +1,10 @@
 """Functionality for saving and loading Tensorflow Keras models as TileDB arrays"""
 
 import contextlib
-import glob
 import io
 import json
 import logging
-import os.path
+import os
 import pickle
 from operator import attrgetter
 from typing import Any, List, Mapping, Optional, Tuple
@@ -25,6 +24,7 @@
 import tiledb
 
 from ._cloud_utils import update_file_properties
+from ._tensorboard import load_tensorboard, save_tensorboard
 from .base import Meta, TileDBModel, Timestamp, current_milli_time
 
 # SharedObjectLoadingScope was introduced in TensorFlow 2.5
@@ -73,11 +73,8 @@ def save(
         if include_callbacks:
             for cb in include_callbacks:
                 if isinstance(cb, TensorBoard):
-                    event_files = self._get_tensorboard_files(cb.log_dir)
-                    meta = {
-                        "__TENSORBOARD__": pickle.dumps(event_files, protocol=4),
-                        **(meta or {}),
-                    }
+                    cb_meta = save_tensorboard(os.path.join(cb.log_dir, "train"))
+                    meta = {**meta, **cb_meta} if meta else cb_meta
 
         # Create TileDB model array
         if not update:
@@ -182,14 +179,7 @@ def load_tensorboard(
         target_dir: Optional[str] = None,
         timestamp: Optional[Timestamp] = None,
     ) -> None:
-        with tiledb.open(self.uri, ctx=self.ctx, timestamp=timestamp) as model_array:
-            tb_data = pickle.loads(model_array.meta["__TENSORBOARD__"])
-            for path in tb_data.keys():
-                log_dir = target_dir if target_dir else os.path.dirname(path)
-                if not os.path.exists(log_dir):
-                    os.mkdir(log_dir)
-                with open(os.path.join(log_dir, os.path.basename(path)), "wb") as f:
-                    f.write(tb_data[path])
+        return load_tensorboard(self.uri, self.ctx, target_dir, timestamp)
 
     def preview(self) -> str:
         """Create a string representation of the model."""
@@ -419,11 +409,3 @@ def _load_weights_from_tiledb(
                 )
             var_value_tuples.extend(zip(weight_vars, read_weight_values))
         backend.batch_set_value(var_value_tuples)
-
-    @staticmethod
-    def _get_tensorboard_files(log_dir: str) -> Mapping[str, bytes]:
-        event_files = {}
-        for path in glob.glob(f"{log_dir}/train/*tfevents*"):
-            with open(path, "rb") as f:
-                event_files[path] = f.read()
-        return event_files