Check CUDA out of memory

environment.yml

@@ -49,6 +49,7 @@ dependencies:
  - bump2version
  - mypy
  - pre_commit
+ - grpcio-tools
 
  - pip
 

proto/inference.proto

@@ -2,6 +2,7 @@ syntax = "proto3";
 
 service Inference {
  rpc CreateModelSession(CreateModelSessionRequest) returns (ModelSession) {}
+
  rpc CloseModelSession(ModelSession) returns (Empty) {}
 
  rpc CreateDatasetDescription(CreateDatasetDescriptionRequest) returns (DatasetDescription) {}
@@ -11,6 +12,10 @@ service Inference {
  rpc ListDevices(Empty) returns (Devices) {}
 
  rpc Predict(PredictRequest) returns (PredictResponse) {}
+
+ rpc IsCudaOutOfMemory(IsCudaOutOfMemoryRequest) returns (IsCudaOutOfMemoryResponse) {}
+
+ rpc MaxCudaMemoryShape(MaxCudaMemoryShapeRequest) returns (MaxCudaMemoryShapeResponse) {}
 }
 
 message Device {
@@ -55,35 +60,6 @@ message NamedFloats {
  repeated NamedFloat namedFloats = 1;
 }
 
-
-/* InputShape will always be expected to have `shape` set.
- * For `ShapeType` PARAMETRIZED, also a `stepShape` has to be given.
- * ref: https://github.com/bioimage-io/spec-bioimage-io/blob/gh-pages/model_spec_latest.md */
-message InputShape {
- enum ShapeType {
- EXPLICIT = 0;
- PARAMETRIZED = 1;
- }
-
- ShapeType shapeType = 1;
- // shape is min, when PARAMETRIZED
- NamedInts shape = 2;
- NamedInts stepShape = 4;
-}
-
-message OutputShape {
- enum ShapeType {
- EXPLICIT = 0;
- IMPLICIT = 1;
- }
- ShapeType shapeType = 1;
- NamedInts shape = 2;
- NamedInts halo = 3;
- string referenceTensor = 4;
- NamedFloats scale = 5;
- NamedFloats offset = 6;
-}
-
 message ModelSession {
  string id = 1;
 }
@@ -124,6 +100,30 @@ message Tensor {
  repeated NamedInt shape = 4;
 }
 
+message IsCudaOutOfMemoryRequest {
+ string modelSessionId = 1;
+ string tensorId = 2;
+ NamedInts shape = 3;
+ string deviceId = 4;
+}
+
+message IsCudaOutOfMemoryResponse {
+ bool isCudaOutOfMemory = 1;
+}
+
+message MaxCudaMemoryShapeRequest {
+ string modelSessionId = 1;
+ string tensorId = 2;
+ NamedInts stepShape = 3;
+ NamedInts minShape = 4;
+ NamedInts maxShape = 5;
+ string deviceId = 6;
+}
+
+message MaxCudaMemoryShapeResponse {
+ NamedInts maxShape = 1;
+}
+
 message PredictRequest {
  string modelSessionId = 1;
  string datasetId = 2;
@@ -142,13 +142,3 @@ service FlightControl {
  rpc Shutdown(Empty) returns (Empty) {}
 }
 
-message ModelInfo {
- repeated string deviceIds = 1;
-}
-
-message CreateModelSessionChunkedRequest {
- oneof data {
- ModelInfo info = 1;
- Blob chunk = 2;
- }
-}

tests/conftest.py

@@ -21,6 +21,9 @@
 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_DUMMY = "dummy"
+TEST_BIOIMAGEIO_DUMMY_CUDA_OUT_OF_MEMORY = "dummy_cuda_out_of_memory"
+
 TEST_BIOIMAGEIO_TENSORFLOW_DUMMY = "dummy_tensorflow"
 TEST_BIOIMAGEIO_TORCHSCRIPT = "unet2d_torchscript"
 
@@ -138,6 +141,14 @@ def _bioimageio_package(rdf_source):
  return data
 
 
+@pytest.fixture
+def bioimageio_dummy_cuda_out_of_memory_model_bytes(data_path):
+ rdf_source = data_path / TEST_BIOIMAGEIO_DUMMY_CUDA_OUT_OF_MEMORY / "rdf.yaml"
+ data = io.BytesIO()
+ export_resource_package(rdf_source, output_path=data)
+ return data
+
+
 def archive(directory):
  result = io.BytesIO()
 

tests/data/dummy_cuda_out_of_memory/dummy_cuda_out_of_memory.py

@@ -0,0 +1,13 @@
+import numpy as np
+from torch import Tensor, nn
+
+MAX_SHAPE = (1, 1, 10, 10)
+
+
+class Dummy(nn.Module):
+ def forward(self, input: Tensor):
+ input_size = np.prod(input.shape)
+ max_size = np.prod(MAX_SHAPE)
+ if input_size > max_size:
+ raise RuntimeError("out of memory")
+ return input + 1

tests/data/dummy_cuda_out_of_memory/rdf.yaml

@@ -0,0 +1,57 @@
+format_version: 0.3.3
+language: python
+framework: pytorch
+
+name: UNet2DNucleiBroad
+description: A 2d U-Net pretrained on broad nucleus dataset.
+cite:
+ - text: "Ronneberger, Olaf et al. U-net: Convolutional networks for biomedical image segmentation. MICCAI 2015."
+ doi: https://doi.org/10.1007/978-3-319-24574-4_28
+authors:
+ - name: "ilastik-team"
+ affiliation: "EMBL Heidelberg"
+
+documentation: dummy.md
+tags: [pytorch, nucleus-segmentation]
+license: MIT
+git_repo: https://github.com/ilastik/tiktorch
+covers: []
+
+source: dummy_cuda_out_of_memory.py::Dummy
+sha256: 00ffb1647cf7ec524892206dce6258d9da498fe040c62838f31b501a09bfd573
+timestamp: 2019-12-11T12:22:32Z # ISO 8601
+
+test_inputs: [dummy_in.npy]
+test_outputs: [dummy_out.npy]
+
+weights:
+ pytorch_state_dict:
+ source: ./weights
+ sha256: 518cb80bad2eb3ec3dfbe6bab74920951391ce8fb24e15cf59b9b9f052a575a6
+ authors:
+ - name: "ilastik-team"
+ affiliation: "EMBL Heidelberg"
+
+
+# TODO double check inputs/outputs
+inputs:
+ - name: input
+ axes: bcyx
+ data_type: float32
+ data_range: [-inf, inf]
+ shape:
+ min: [1, 1, 5, 5]
+ step: [0, 0, 1, 1]
+
+outputs:
+ - name: output
+ axes: bcyx
+ data_type: float32
+ data_range: [0, 1]
+ shape:
+ reference_tensor: input # FIXME(m-novikov) ignoring for now
+ scale: [1, 1, 1, 1]
+ offset: [0, 0, 0, 0]
+ halo: [0, 0, 2, 2] # Should be moved to outputs
+
+type: model

tests/test_converters.py

@@ -3,18 +3,7 @@
 import xarray as xr
 from numpy.testing import assert_array_equal
 
-from tiktorch.converters import (
- NamedExplicitOutputShape,
- NamedImplicitOutputShape,
- NamedParametrizedShape,
- Sample,
- input_shape_to_pb_input_shape,
- numpy_to_pb_tensor,
- output_shape_to_pb_output_shape,
- pb_tensor_to_numpy,
- pb_tensor_to_xarray,
- xarray_to_pb_tensor,
-)
+from tiktorch.converters import Sample, numpy_to_pb_tensor, pb_tensor_to_numpy, pb_tensor_to_xarray, xarray_to_pb_tensor
 from tiktorch.proto import inference_pb2
 
 
@@ -177,99 +166,6 @@ def test_should_same_data(self, shape):
  assert_array_equal(arr, result_arr)
 
 
-class TestShapeConversions:
- def to_named_explicit_shape(self, shape, axes, halo):
- return NamedExplicitOutputShape(
- halo=[(name, dim) for name, dim in zip(axes, halo)], shape=[(name, dim) for name, dim in zip(axes, shape)]
- )
-
- def to_named_implicit_shape(self, axes, halo, offset, scales, reference_tensor):
- return NamedImplicitOutputShape(
- halo=[(name, dim) for name, dim in zip(axes, halo)],
- offset=[(name, dim) for name, dim in zip(axes, offset)],
- scale=[(name, scale) for name, scale in zip(axes, scales)],
- reference_tensor=reference_tensor,
- )
-
- def to_named_paramtrized_shape(self, min_shape, axes, step):
- return NamedParametrizedShape(
- min_shape=[(name, dim) for name, dim in zip(axes, min_shape)],
- step_shape=[(name, dim) for name, dim in zip(axes, step)],
- )
-
- @pytest.mark.parametrize(
- "shape,axes,halo",
- [((42,), "x", (0,)), ((42, 128, 5), "abc", (1, 1, 1)), ((5, 4, 3, 2, 1, 42), "btzyxc", (1, 2, 3, 4, 5, 24))],
- )
- def test_explicit_output_shape(self, shape, axes, halo):
- named_shape = self.to_named_explicit_shape(shape, axes, halo)
- pb_shape = output_shape_to_pb_output_shape(named_shape)
-
- assert pb_shape.shapeType == 0
- assert pb_shape.referenceTensor == ""
- assert len(pb_shape.scale.namedFloats) == 0
- assert len(pb_shape.offset.namedFloats) == 0
-
- assert [(d.name, d.size) for d in pb_shape.halo.namedInts] == [(name, size) for name, size in zip(axes, halo)]
- assert [(d.name, d.size) for d in pb_shape.shape.namedInts] == [(name, size) for name, size in zip(axes, shape)]
-
- @pytest.mark.parametrize(
- "axes,halo,offset,scales,reference_tensor",
- [("x", (0,), (10,), (1.0,), "forty-two"), ("abc", (1, 1, 1), (1, 2, 3), (1.0, 2.0, 3.0), "helloworld")],
- )
- def test_implicit_output_shape(self, axes, halo, offset, scales, reference_tensor):
- named_shape = self.to_named_implicit_shape(axes, halo, offset, scales, reference_tensor)
- pb_shape = output_shape_to_pb_output_shape(named_shape)
-
- assert pb_shape.shapeType == 1
- assert pb_shape.referenceTensor == reference_tensor
- assert [(d.name, d.size) for d in pb_shape.scale.namedFloats] == [
- (name, size) for name, size in zip(axes, scales)
- ]
- assert [(d.name, d.size) for d in pb_shape.offset.namedFloats] == [
- (name, size) for name, size in zip(axes, offset)
- ]
-
- assert [(d.name, d.size) for d in pb_shape.halo.namedInts] == [(name, size) for name, size in zip(axes, halo)]
- assert len(pb_shape.shape.namedInts) == 0
-
- def test_output_shape_raises(self):
- shape = [("a", 1)]
- with pytest.raises(TypeError):
- _ = output_shape_to_pb_output_shape(shape)
-
- @pytest.mark.parametrize(
- "shape,axes",
- [((42,), "x"), ((42, 128, 5), "abc"), ((5, 4, 3, 2, 1, 42), "btzyxc")],
- )
- def test_explicit_input_shape(self, shape, axes):
- named_shape = [(name, dim) for name, dim in zip(axes, shape)]
- pb_shape = input_shape_to_pb_input_shape(named_shape)
-
- assert pb_shape.shapeType == 0
- assert [(d.name, d.size) for d in pb_shape.shape.namedInts] == [(name, size) for name, size in zip(axes, shape)]
-
- @pytest.mark.parametrize(
- "min_shape,axes,step",
- [
- ((42,), "x", (5,)),
- ((42, 128, 5), "abc", (1, 2, 3)),
- ((5, 4, 3, 2, 1, 42), "btzyxc", (15, 24, 33, 42, 51, 642)),
- ],
- )
- def test_parametrized_input_shape(self, min_shape, axes, step):
- named_shape = self.to_named_paramtrized_shape(min_shape, axes, step)
- pb_shape = input_shape_to_pb_input_shape(named_shape)
-
- assert pb_shape.shapeType == 1
- assert [(d.name, d.size) for d in pb_shape.shape.namedInts] == [
- (name, size) for name, size in zip(axes, min_shape)
- ]
- 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)
@@ -293,7 +189,7 @@ def test_create_sample_from_pb_tensors(self):
  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):
+ def test_create_sample_from_xr_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)