Increase unit test coverage (#201)

* cloud utils unit tests

* pytorch model unit tests

* tensorflow keras model unit tests

* _use_legacy_schema became a static method

* sklearn model unit tests

* cloud utils unit tests

* pytorch model unit tests

* tensorflow keras model unit tests

* _use_legacy_schema became a static method

* sklearn model unit tests

* removed tensorflow custom layer commented code

* PR changes

Author: georgeSkoumas

tests/models/test_cloud_utils.py

@@ -1,3 +1,5 @@
+import pytest
+
 from tiledb.ml.models._cloud_utils import (
     get_cloud_uri,
     get_s3_prefix,
@@ -8,6 +10,8 @@
 class TestCloudUtils:
     def test_get_s3_prefix(self, mocker):
 
+        assert get_s3_prefix(None) is None
+
         profile = mocker.patch(
             "tiledb.cloud.client.user_profile",
             return_value=mocker.Mock(username="foo", default_s3_path="bar"),
@@ -47,6 +51,12 @@ def test_get_cloud_uri(self, mocker):
             uri="tiledb_array", namespace="test_namespace"
         )
 
+        mocker.patch("tiledb.ml.models._cloud_utils.get_s3_prefix", return_value=None)
+        with pytest.raises(ValueError) as ex:
+            get_cloud_uri(uri="tiledb_array", namespace="test_namespace")
+
+        assert "You must set the default s3 prefix path for ML models" in str(ex.value)
+
     def test_update_file_properties(self, mocker):
         mock_tiledb_cloud_update_file_properties = mocker.patch(
             "tiledb.cloud.array.update_file_properties"

tests/models/test_pytorch_models.py

@@ -127,6 +127,12 @@ def test_save(self, tmpdir, net, optimizer):
         ):
             assert all([a == b for a, b in zip(key_item_1[1], key_item_2[1])])
 
+        with pytest.raises(RuntimeError) as ex:
+            tiledb_obj = PyTorchTileDBModel(uri="")
+            tiledb_obj.save()
+
+        assert "Model is not initialized" in str(ex.value)
+
     @net
     def test_preview(self, tmpdir, net):
         # With model given as argument

tests/models/test_sklearn_models.py

@@ -36,6 +36,12 @@ def test_save_load(self, tmpdir, net):
             ]
         )
 
+        with pytest.raises(RuntimeError) as ex:
+            tiledb_obj = SklearnTileDBModel(uri="")
+            tiledb_obj.save()
+
+        assert "Model is not initialized" in str(ex.value)
+
     def test_preview(self, tmpdir, net):
         # With model as argument
         tiledb_array = os.path.join(tmpdir, "test_array")

tests/models/test_tensorflow_keras_models.py

@@ -111,6 +111,12 @@ def test_save_model_to_tiledb_array(self, tmpdir, api, loss, optimizer, metrics)
         tiledb_model_obj.save(include_optimizer=True if optimizer else False)
         assert tiledb.array_exists(tiledb_uri)
 
+        with pytest.raises(RuntimeError) as ex:
+            tiledb_obj = TensorflowKerasTileDBModel(uri="")
+            tiledb_obj.save()
+
+        assert "Model is not initialized" in str(ex.value)
+
     @api
     @loss_optimizer_metrics
     def test_save_model_to_tiledb_array_predictions(

tiledb/ml/models/_base.py

@@ -217,6 +217,7 @@ def _load_tensorboard(self, model_array: tiledb.Array) -> None:
             with open(path, "wb") as f:
                 f.write(file_bytes)
 
-    def _use_legacy_schema(self, model_array: tiledb.Array) -> bool:
+    @staticmethod
+    def _use_legacy_schema(model_array: tiledb.Array) -> bool:
         # TODO: Decide based on tiledb-ml version and not on schema characteristics, like "offset".
         return str(model_array.schema.domain.dim(0).name) != "offset"

tiledb/ml/models/sklearn.py

@@ -63,7 +63,8 @@ def load(self, *, timestamp: Optional[Timestamp] = None) -> BaseEstimator:
             else:
                 return self.__load(model_array)
 
-    def __load_legacy(self, model_array: tiledb.Array) -> BaseEstimator:
+    @staticmethod
+    def __load_legacy(model_array: tiledb.Array) -> BaseEstimator:
         return pickle.loads(model_array[:]["model_params"].item(0))
 
     def __load(self, model_array: tiledb.Array) -> BaseEstimator:

tiledb/ml/models/tensorflow_keras.py

@@ -6,10 +6,9 @@
 import os
 import pickle
 from collections import ChainMap
-from typing import Any, List, Mapping, Optional, Tuple
+from typing import Any, Mapping, Optional, Tuple
 
 import keras
-import numpy as np
 import tensorflow as tf
 
 import tiledb
@@ -128,7 +127,7 @@ def load(
         with tiledb.open(self.uri, ctx=self.ctx, timestamp=timestamp) as model_array:
             if self._use_legacy_schema(model_array):
                 return self.__load_legacy(
-                    model_array, compile_model, callback, custom_objects, input_shape
+                    model_array, compile_model, callback, custom_objects
                 )
             else:
                 return self.__load(model_array, compile_model, callback)
@@ -139,30 +138,15 @@ def __load_legacy(
         compile_model: bool,
         callback: bool,
         custom_objects: Optional[Mapping[str, Any]],
-        input_shape: Optional[Tuple[int, ...]],
     ) -> tf.keras.Model:
         model_array_results = model_array[:]
         model_config = json.loads(model_array.meta["model_config"])
         model_class = model_config["class_name"]
 
-        if model_class not in ("Functional", "Sequential"):
-            with SharedObjectLoadingScope():
-                with tf.keras.utils.CustomObjectScope(custom_objects or {}):
-                    if hasattr(model_config, "decode"):
-                        model_config = model_config.decode("utf-8")
-                    model = tf.keras.models.model_from_config(
-                        model_config, custom_objects=custom_objects
-                    )
-                    if not model.built:
-                        model.build(input_shape)
-
-                    # Load weights for layers
-                    self._load_custom_subclassed_model(model, model_array)
-        else:
-            cls = tf.keras.Sequential if model_class == "Sequential" else tf.keras.Model
-            model = cls.from_config(model_config["config"])
-            model_weights = pickle.loads(model_array_results["model_weights"].item(0))
-            model.set_weights(model_weights)
+        cls = tf.keras.Sequential if model_class == "Sequential" else tf.keras.Model
+        model = cls.from_config(model_config["config"])
+        model_weights = pickle.loads(model_array_results["model_weights"].item(0))
+        model.set_weights(model_weights)
 
         if compile_model:
             optimizer_weights = pickle.loads(
@@ -198,6 +182,7 @@ def __load_legacy(
                         "starting with a freshly initialized "
                         "optimizer."
                     )
+
         if callback:
             try:
                 with tiledb.open(f"{self.uri}-tensorboard") as tb_array:
@@ -239,6 +224,7 @@ def __load(
             saving_utils.try_build_compiled_arguments(model)
 
             optimizer_weights = self._get_model_param(model_array, "optimizer")
+
             # Set optimizer weights.
             if optimizer_weights:
                 try:
@@ -284,70 +270,3 @@ def _serialize_optimizer_weights(
             optimizer_weights = tf.keras.backend.batch_get_value(optimizer.weights)
             return pickle.dumps(optimizer_weights, protocol=4)
         return b""
-
-    def _load_custom_subclassed_model(
-        self, model: tf.keras.Model, model_array: tiledb.Array
-    ) -> None:
-        if "keras_version" in model_array.meta:
-            original_keras_version = model_array.meta["keras_version"]
-            if hasattr(original_keras_version, "decode"):
-                original_keras_version = original_keras_version.decode("utf8")
-        else:
-            original_keras_version = "1"
-        if "backend" in model_array.meta:
-            original_backend = model_array.meta["backend"]
-            if hasattr(original_backend, "decode"):
-                original_backend = original_backend.decode("utf8")
-        else:
-            original_backend = None
-
-        # Load weights for layers
-        self._load_weights_from_tiledb(
-            model_array[:], model, original_keras_version, original_backend
-        )
-
-    @staticmethod
-    def _load_weights_from_tiledb(
-        model_array_results: Mapping[str, Any],
-        model: tf.keras.Model,
-        original_keras_version: Optional[str],
-        original_backend: Optional[str],
-    ) -> None:
-        num_layers = 0
-        for layer in model.layers:
-            weights = layer.trainable_weights + layer.non_trainable_weights
-            if weights:
-                num_layers += 1
-
-        read_layer_names = []
-        for k, name in enumerate(model_array_results["layer_name"]):
-            layer_weight_names = pickle.loads(
-                model_array_results["weight_names"].item(k)
-            )
-            if layer_weight_names:
-                read_layer_names.append(name)
-
-        if len(read_layer_names) != num_layers:
-            raise ValueError(
-                f"You are trying to load a weight file with {len(read_layer_names)} "
-                f"layers into a model with {num_layers} layers"
-            )
-
-        var_value_tuples: List[Tuple[tf.Variable, np.ndarray]] = []
-        for k, layer in enumerate(model.layers):
-            weight_vars = layer.trainable_weights + layer.non_trainable_weights
-            read_weight_values = pickle.loads(
-                model_array_results["weight_values"].item(k)
-            )
-            read_weight_values = preprocess_weights_for_loading(
-                layer, read_weight_values, original_keras_version, original_backend
-            )
-            if len(read_weight_values) != len(weight_vars):
-                raise ValueError(
-                    f'Layer #{k}  (named "{layer.name}" in the current model) was found '
-                    f"to correspond to layer {layer} in the save file. However the new "
-                    f"layer {layer.name} expects {len(weight_vars)} weights, "
-                    f"but the saved weights have {len(read_weight_values)} elements"
-                )
-            var_value_tuples.extend(zip(weight_vars, read_weight_values))
-        tf.keras.backend.batch_set_value(var_value_tuples)