Implement code review suggestions

- Use the class `Sample` to associate xr.DataArray with tensor id coming from the model
- Do the checks only from the client side
- Add checking about axes validity
- Add tests

proto/inference.proto

@@ -120,7 +120,7 @@ message NamedFloat {
 message Tensor {
   bytes buffer = 1;
   string dtype = 2;
-  string specId = 3;
+  string tensorId = 3;
   repeated NamedInt shape = 4;
 }
 

tests/conftest.py

@@ -18,7 +18,9 @@
 TEST_DATA = "data"
 TEST_BIOIMAGEIO_ZIPFOLDER = "unet2d"
 TEST_BIOIMAGEIO_ONNX = "unet2d_onnx"
-TEST_BIOIMAGEIO_DUMMY = "dummy"
+TEST_BIOIMAGEIO_DUMMY_EXPLICIT = "dummy"
+TEST_BIOIMAGEIO_DUMMY_EXPLICIT_RDF = f"{TEST_BIOIMAGEIO_DUMMY_EXPLICIT}/Dummy.model.yaml"
+TEST_BIOIMAGEIO_DUMMY_PARAM_RDF = "dummy_param/Dummy.model_param.yaml"
 TEST_BIOIMAGEIO_TENSORFLOW_DUMMY = "dummy_tensorflow"
 TEST_BIOIMAGEIO_TORCHSCRIPT = "unet2d_torchscript"
 
@@ -98,7 +100,7 @@ def bioimageio_model_zipfile(bioimageio_model_bytes):
 
 @pytest.fixture
 def bioimageio_dummy_model_filepath(data_path, tmpdir):
-    bioimageio_net_dir = Path(data_path) / TEST_BIOIMAGEIO_DUMMY
+    bioimageio_net_dir = Path(data_path) / TEST_BIOIMAGEIO_DUMMY_EXPLICIT
     path = tmpdir / "dummy_model.zip"
 
     with ZipFile(path, mode="w") as zip_model:
@@ -113,17 +115,23 @@ def bioimageio_dummy_model_filepath(data_path, tmpdir):
 
 
 @pytest.fixture
-def bioimageio_dummy_model_bytes(data_path):
-    rdf_source = data_path / TEST_BIOIMAGEIO_DUMMY / "Dummy.model.yaml"
+def bioimageio_dummy_explicit_model_bytes(data_path):
+    rdf_source = data_path / TEST_BIOIMAGEIO_DUMMY_EXPLICIT_RDF
     return _bioimageio_package(rdf_source)
 
 
 @pytest.fixture
 def bioimageio_dummy_param_model_bytes(data_path):
-    rdf_source = data_path / "dummy_param" / "Dummy.model_param.yaml"
+    rdf_source = data_path / TEST_BIOIMAGEIO_DUMMY_PARAM_RDF
     return _bioimageio_package(rdf_source)
 
 
+@pytest.fixture(params=[(TEST_BIOIMAGEIO_DUMMY_PARAM_RDF, "param"), (TEST_BIOIMAGEIO_DUMMY_EXPLICIT_RDF, "input")])
+def bioimageio_dummy_model(request, data_path):
+    path, tensor_id = request.param
+    yield _bioimageio_package(data_path / path), tensor_id
+
+
 def _bioimageio_package(rdf_source):
     data = io.BytesIO()
     export_resource_package(rdf_source, output_path=data)

tests/test_converters.py

@@ -7,12 +7,12 @@
     NamedExplicitOutputShape,
     NamedImplicitOutputShape,
     NamedParametrizedShape,
-    Tensor,
+    Sample,
     input_shape_to_pb_input_shape,
     numpy_to_pb_tensor,
     output_shape_to_pb_output_shape,
     pb_tensor_to_numpy,
-    pb_tensor_to_tensor,
+    pb_tensor_to_xarray,
     xarray_to_pb_tensor,
 )
 from tiktorch.proto import inference_pb2
@@ -28,11 +28,11 @@ def _numpy_to_pb_tensor(arr):
     return parsed
 
 
-def to_pb_tensor(spec_id: str, arr: xr.DataArray):
+def to_pb_tensor(tensor_id: str, arr: xr.DataArray):
     """
     Makes sure that tensor was serialized/deserialized
     """
-    tensor = xarray_to_pb_tensor(spec_id, arr)
+    tensor = xarray_to_pb_tensor(tensor_id, arr)
     parsed = inference_pb2.Tensor()
     parsed.ParseFromString(tensor.SerializeToString())
     return parsed
@@ -141,40 +141,40 @@ class TestPBTensorToXarray:
     def test_should_raise_on_empty_dtype(self):
         tensor = inference_pb2.Tensor(dtype="", shape=[inference_pb2.NamedInt(size=1), inference_pb2.NamedInt(size=2)])
         with pytest.raises(ValueError):
-            pb_tensor_to_tensor(tensor)
+            pb_tensor_to_xarray(tensor)
 
     def test_should_raise_on_empty_shape(self):
         tensor = inference_pb2.Tensor(dtype="int64", shape=[])
         with pytest.raises(ValueError):
-            pb_tensor_to_tensor(tensor)
+            pb_tensor_to_xarray(tensor)
 
-    def test_should_return_tensor(self):
+    def test_should_return_xarray(self):
         arr = xr.DataArray(np.arange(9))
         parsed = to_pb_tensor("input0", arr)
-        result_tensor = pb_tensor_to_tensor(parsed)
-        assert isinstance(result_tensor, Tensor)
+        result_tensor = pb_tensor_to_xarray(parsed)
+        assert isinstance(result_tensor, xr.DataArray)
 
     @pytest.mark.parametrize("np_dtype,dtype_str", [(np.int64, "int64"), (np.uint8, "uint8"), (np.float32, "float32")])
     def test_should_have_same_dtype(self, np_dtype, dtype_str):
         arr = xr.DataArray(np.arange(9, dtype=np_dtype))
         pb_tensor = to_pb_tensor("input0", arr)
-        tensor = pb_tensor_to_tensor(pb_tensor)
+        result_arr = pb_tensor_to_xarray(pb_tensor)
 
-        assert arr.dtype == tensor.data.dtype
+        assert arr.dtype == result_arr.dtype
 
     @pytest.mark.parametrize("shape", [(3, 3), (1,), (1, 1), (18, 20, 1)])
     def test_should_same_shape(self, shape):
         arr = xr.DataArray(np.zeros(shape))
         pb_tensor = to_pb_tensor("input0", arr)
-        tensor = pb_tensor_to_tensor(pb_tensor)
-        assert arr.shape == tensor.data.shape
+        result_arr = pb_tensor_to_xarray(pb_tensor)
+        assert arr.shape == result_arr.shape
 
     @pytest.mark.parametrize("shape", [(3, 3), (1,), (1, 1), (18, 20, 1)])
     def test_should_same_data(self, shape):
         arr = xr.DataArray(np.random.random(shape))
         pb_tensor = to_pb_tensor("input0", arr)
-        tensor = pb_tensor_to_tensor(pb_tensor)
-        assert_array_equal(arr, tensor.data)
+        result_arr = pb_tensor_to_xarray(pb_tensor)
+        assert_array_equal(arr, result_arr)
 
 
 class TestShapeConversions:
@@ -268,3 +268,64 @@ def test_parametrized_input_shape(self, min_shape, axes, step):
         assert [(d.name, d.size) for d in pb_shape.stepShape.namedInts] == [
             (name, size) for name, size in zip(axes, step)
         ]
+
+
+class TestSample:
+    def test_create_sample_from_pb_tensors(self):
+        arr_1 = np.arange(32 * 32, dtype=np.int64).reshape(32, 32)
+        tensor_1 = inference_pb2.Tensor(
+            dtype="int64",
+            tensorId="input1",
+            buffer=bytes(arr_1),
+            shape=[inference_pb2.NamedInt(name="x", size=32), inference_pb2.NamedInt(name="y", size=32)],
+        )
+
+        arr_2 = np.arange(64 * 64, dtype=int).reshape(64, 64)
+        tensor_2 = inference_pb2.Tensor(
+            dtype="int64",
+            tensorId="input2",
+            buffer=bytes(arr_2),
+            shape=[inference_pb2.NamedInt(name="x", size=64), inference_pb2.NamedInt(name="y", size=64)],
+        )
+
+        sample = Sample.from_pb_tensors([tensor_1, tensor_2])
+        assert len(sample.tensors) == 2
+        assert sample.tensors["input1"].equals(xr.DataArray(arr_1, dims=["x", "y"]))
+        assert sample.tensors["input2"].equals(xr.DataArray(arr_2, dims=["x", "y"]))
+
+    def test_create_sample_from_raw_data(self):
+        arr_1 = np.arange(32 * 32, dtype=np.int64).reshape(32, 32)
+        tensor_1 = xr.DataArray(arr_1, dims=["x", "y"])
+        arr_2 = np.arange(64 * 64, dtype=np.int64).reshape(64, 64)
+        tensor_2 = xr.DataArray(arr_2, dims=["x", "y"])
+        tensors_ids = ["input1", "input2"]
+        actual_sample = Sample.from_raw_data(tensors_ids, [tensor_1, tensor_2])
+
+        expected_dict = {tensors_ids[0]: tensor_1, tensors_ids[1]: tensor_2}
+        expected_sample = Sample(expected_dict)
+        assert actual_sample == expected_sample
+
+    def test_sample_to_pb_tensors(self):
+        arr_1 = np.arange(32 * 32, dtype=np.int64).reshape(32, 32)
+        tensor_1 = xr.DataArray(arr_1, dims=["x", "y"])
+        arr_2 = np.arange(64 * 64, dtype=np.int64).reshape(64, 64)
+        tensor_2 = xr.DataArray(arr_2, dims=["x", "y"])
+        tensors_ids = ["input1", "input2"]
+        sample = Sample.from_raw_data(tensors_ids, [tensor_1, tensor_2])
+
+        pb_tensor_1 = inference_pb2.Tensor(
+            dtype="int64",
+            tensorId="input1",
+            buffer=bytes(arr_1),
+            shape=[inference_pb2.NamedInt(name="x", size=32), inference_pb2.NamedInt(name="y", size=32)],
+        )
+        pb_tensor_2 = inference_pb2.Tensor(
+            dtype="int64",
+            tensorId="input2",
+            buffer=bytes(arr_2),
+            shape=[inference_pb2.NamedInt(name="x", size=64), inference_pb2.NamedInt(name="y", size=64)],
+        )
+        expected_tensors = [pb_tensor_1, pb_tensor_2]
+
+        actual_tensors = sample.to_pb_tensors()
+        assert expected_tensors == actual_tensors

tests/test_server/test_grpc/test_inference_servicer.py

@@ -49,8 +49,8 @@ def test_model_session_creation(self, grpc_stub, bioimageio_model_bytes):
         assert model.id
         grpc_stub.CloseModelSession(model)
 
-    def test_model_session_creation_using_upload_id(self, grpc_stub, data_store, bioimageio_dummy_model_bytes):
-        id_ = data_store.put(bioimageio_dummy_model_bytes.getvalue())
+    def test_model_session_creation_using_upload_id(self, grpc_stub, data_store, bioimageio_dummy_explicit_model_bytes):
+        id_ = data_store.put(bioimageio_dummy_explicit_model_bytes.getvalue())
 
         rq = inference_pb2.CreateModelSessionRequest(model_uri=f"upload://{id_}", deviceIds=["cpu"])
         model = grpc_stub.CreateModelSession(rq)
@@ -154,30 +154,33 @@ def test_call_fails_with_unknown_model_session_id(self, grpc_stub):
         assert grpc.StatusCode.FAILED_PRECONDITION == e.value.code()
         assert "model-session with id myid1 doesn't exist" in e.value.details()
 
-    def test_call_predict(self, grpc_stub, bioimageio_dummy_model_bytes):
-        model = grpc_stub.CreateModelSession(valid_model_request(bioimageio_dummy_model_bytes))
+    def test_call_predict_valid_explicit(self, grpc_stub, bioimageio_dummy_explicit_model_bytes):
+        model = grpc_stub.CreateModelSession(valid_model_request(bioimageio_dummy_explicit_model_bytes))
         arr = xr.DataArray(np.arange(128 * 128).reshape(1, 1, 128, 128), dims=("b", "c", "x", "y"))
         expected = arr + 1
-        input_spec_id = "input"
-        output_spec_id = "output"
-        input_tensors = [converters.xarray_to_pb_tensor(input_spec_id, arr)]
+        input_tensor_id = "input"
+        output_tensor_id = "output"
+        input_tensors = [converters.xarray_to_pb_tensor(input_tensor_id, arr)]
         res = grpc_stub.Predict(inference_pb2.PredictRequest(modelSessionId=model.id, tensors=input_tensors))
 
         grpc_stub.CloseModelSession(model)
 
         assert len(res.tensors) == 1
-        assert res.tensors[0].specId == output_spec_id
+        assert res.tensors[0].tensorId == output_tensor_id
         assert_array_equal(expected, converters.pb_tensor_to_numpy(res.tensors[0]))
 
-    def test_call_predict_invalid_shape_explicit(self, grpc_stub, bioimageio_dummy_model_bytes):
-        model = grpc_stub.CreateModelSession(valid_model_request(bioimageio_dummy_model_bytes))
+    def test_call_predict_invalid_shape_explicit(self, grpc_stub, bioimageio_dummy_explicit_model_bytes):
+        model = grpc_stub.CreateModelSession(valid_model_request(bioimageio_dummy_explicit_model_bytes))
         arr = xr.DataArray(np.arange(32 * 32).reshape(1, 1, 32, 32), dims=("b", "c", "x", "y"))
         input_tensors = [converters.xarray_to_pb_tensor("input", arr)]
         with pytest.raises(grpc.RpcError):
             grpc_stub.Predict(inference_pb2.PredictRequest(modelSessionId=model.id, tensors=input_tensors))
         grpc_stub.CloseModelSession(model)
 
-    @pytest.mark.parametrize("shape", [(1, 1, 64, 32), (1, 1, 32, 64), (1, 1, 64, 32), (0, 1, 64, 64), (1, 0, 64, 64)])
+    @pytest.mark.parametrize(
+        "shape",
+        [(1, 1, 64, 32), (1, 1, 32, 64), (1, 1, 64, 32), (0, 1, 64, 64), (1, 0, 64, 64)],
+    )
     def test_call_predict_invalid_shape_parameterized(self, grpc_stub, shape, bioimageio_dummy_param_model_bytes):
         model = grpc_stub.CreateModelSession(valid_model_request(bioimageio_dummy_param_model_bytes))
         arr = xr.DataArray(np.arange(np.prod(shape)).reshape(*shape), dims=("b", "c", "x", "y"))
@@ -186,6 +189,26 @@ def test_call_predict_invalid_shape_parameterized(self, grpc_stub, shape, bioima
             grpc_stub.Predict(inference_pb2.PredictRequest(modelSessionId=model.id, tensors=input_tensors))
         grpc_stub.CloseModelSession(model)
 
+    def test_call_predict_invalid_tensor_ids(self, grpc_stub, bioimageio_dummy_model):
+        model_bytes, _ = bioimageio_dummy_model
+        model = grpc_stub.CreateModelSession(valid_model_request(model_bytes))
+        arr = xr.DataArray(np.arange(32 * 32).reshape(32, 32), dims=("x", "y"))
+        input_tensors = [converters.xarray_to_pb_tensor("invalidTensorName", arr)]
+        with pytest.raises(grpc.RpcError) as error:
+            grpc_stub.Predict(inference_pb2.PredictRequest(modelSessionId=model.id, tensors=input_tensors))
+        assert error.value.details().startswith("Exception calling application: Spec invalidTensorName doesn't exist")
+        grpc_stub.CloseModelSession(model)
+
+    def test_call_predict_invalid_axes(self, grpc_stub, bioimageio_dummy_model):
+        model_bytes, tensor_id = bioimageio_dummy_model
+        model = grpc_stub.CreateModelSession(valid_model_request(model_bytes))
+        arr = xr.DataArray(np.arange(32 * 32).reshape(32, 32), dims=("invalidAxis", "y"))
+        input_tensors = [converters.xarray_to_pb_tensor(tensor_id, arr)]
+        with pytest.raises(grpc.RpcError) as error:
+            grpc_stub.Predict(inference_pb2.PredictRequest(modelSessionId=model.id, tensors=input_tensors))
+        assert error.value.details().startswith("Exception calling application: Incompatible axes")
+        grpc_stub.CloseModelSession(model)
+
     @pytest.mark.parametrize("shape", [(1, 1, 64, 64), (1, 1, 66, 65), (1, 1, 68, 66), (1, 1, 70, 67)])
     def test_call_predict_valid_shape_parameterized(self, grpc_stub, shape, bioimageio_dummy_param_model_bytes):
         model = grpc_stub.CreateModelSession(valid_model_request(bioimageio_dummy_param_model_bytes))
@@ -199,13 +222,13 @@ def test_call_predict_tf(self, grpc_stub, bioimageio_dummy_tensorflow_model_byte
         model = grpc_stub.CreateModelSession(valid_model_request(bioimageio_dummy_tensorflow_model_bytes))
         arr = xr.DataArray(np.arange(32 * 32).reshape(1, 1, 32, 32), dims=("b", "c", "x", "y"))
         expected = arr * -1
-        input_spec_id = "input"
-        output_spec_id = "output"
-        input_tensors = [converters.xarray_to_pb_tensor(input_spec_id, arr)]
+        input_tensor_id = "input"
+        output_tensor_id = "output"
+        input_tensors = [converters.xarray_to_pb_tensor(input_tensor_id, arr)]
         res = grpc_stub.Predict(inference_pb2.PredictRequest(modelSessionId=model.id, tensors=input_tensors))
 
         grpc_stub.CloseModelSession(model)
 
         assert len(res.tensors) == 1
-        assert res.tensors[0].specId == output_spec_id
+        assert res.tensors[0].tensorId == output_tensor_id
         assert_array_equal(expected, converters.pb_tensor_to_numpy(res.tensors[0]))

tiktorch/converters.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import dataclasses
-from typing import List, Tuple, Union
+from typing import Dict, List, Tuple, Union
 
 import numpy as np
 import xarray as xr
@@ -34,22 +36,20 @@ class NamedImplicitOutputShape:
 
 
 @dataclasses.dataclass(frozen=True)
-class Tensor:
-    spec_id: str
-    data: xr.DataArray
+class Sample:
+    tensors: Dict[str, xr.DataArray]
 
-    def __hash__(self):
-        return hash(self.spec_id)
+    @classmethod
+    def from_pb_tensors(cls, pb_tensors: List[inference_pb2.Tensor]) -> Sample:
+        return Sample({tensor.tensorId: pb_tensor_to_xarray(tensor) for tensor in pb_tensors})
 
-    def __eq__(self, other):
-        if isinstance(other, Tensor):
-            return self.spec_id == other.spec_id
-        return False
+    @classmethod
+    def from_raw_data(cls, tensor_ids: List[str], tensors_data: List[xr.DataArray]):
+        assert len(tensor_ids) == len(tensors_data)
+        return Sample({tensor_id: tensor_data for tensor_id, tensor_data in zip(tensor_ids, tensors_data)})
 
-    def equals(self, other):
-        if not isinstance(other, Tensor):
-            return False
-        return self.__dict__.items() == other.__dict__.items()
+    def to_pb_tensors(self) -> List[inference_pb2.Tensor]:
+        return [xarray_to_pb_tensor(tensor_id, res_tensor) for tensor_id, res_tensor in self.tensors.items()]
 
 
 def numpy_to_pb_tensor(array: np.ndarray, axistags=None) -> inference_pb2.Tensor:
@@ -60,9 +60,9 @@ def numpy_to_pb_tensor(array: np.ndarray, axistags=None) -> inference_pb2.Tensor
     return inference_pb2.Tensor(dtype=str(array.dtype), shape=shape, buffer=bytes(array))
 
 
-def xarray_to_pb_tensor(spec_id: str, array: xr.DataArray) -> inference_pb2.Tensor:
+def xarray_to_pb_tensor(tensor_id: str, array: xr.DataArray) -> inference_pb2.Tensor:
     shape = [inference_pb2.NamedInt(size=dim, name=name) for dim, name in zip(array.shape, array.dims)]
-    return inference_pb2.Tensor(specId=spec_id, dtype=str(array.dtype), shape=shape, buffer=bytes(array.data))
+    return inference_pb2.Tensor(tensorId=tensor_id, dtype=str(array.dtype), shape=shape, buffer=bytes(array.data))
 
 
 def name_int_tuples_to_pb_NamedInts(name_int_tuples) -> inference_pb2.NamedInts:
@@ -112,7 +112,7 @@ def output_shape_to_pb_output_shape(
         raise TypeError(f"Conversion not supported for type {type(output_shape)}")
 
 
-def pb_tensor_to_tensor(tensor: inference_pb2.Tensor) -> inference_pb2.Tensor:
+def pb_tensor_to_xarray(tensor: inference_pb2.Tensor) -> inference_pb2.Tensor:
     if not tensor.dtype:
         raise ValueError("Tensor dtype is not specified")
 
@@ -121,7 +121,7 @@ def pb_tensor_to_tensor(tensor: inference_pb2.Tensor) -> inference_pb2.Tensor:
 
     data = np.frombuffer(tensor.buffer, dtype=tensor.dtype).reshape(*[dim.size for dim in tensor.shape])
 
-    return Tensor(spec_id=tensor.specId, data=xr.DataArray(data, dims=[d.name for d in tensor.shape]))
+    return xr.DataArray(data, dims=[d.name for d in tensor.shape])
 
 
 def pb_tensor_to_numpy(tensor: inference_pb2.Tensor) -> np.ndarray: