@@ -27,6 +27,12 @@ def grpc_servicer(data_store):
2727 return inference_servicer .InferenceServicer (TorchDevicePool (), SessionManager (), data_store )
2828
2929
30+ @pytest .fixture (autouse = True )
31+ def clean (grpc_servicer ):
32+ yield
33+ grpc_servicer .close_all_sessions ()
34+
35+
3036@pytest .fixture (scope = "module" )
3137def grpc_stub_cls (grpc_channel ):
3238 return inference_pb2_grpc .InferenceStub
@@ -47,15 +53,13 @@ def method_requiring_session(self, request, grpc_stub):
4753 def test_model_session_creation (self , grpc_stub , bioimageio_model_bytes ):
4854 model = grpc_stub .CreateModelSession (valid_model_request (bioimageio_model_bytes ))
4955 assert model .id
50- grpc_stub .CloseModelSession (model )
5156
5257 def test_model_session_creation_using_upload_id (self , grpc_stub , data_store , bioimageio_dummy_explicit_model_bytes ):
5358 id_ = data_store .put (bioimageio_dummy_explicit_model_bytes .getvalue ())
5459
5560 rq = inference_pb2 .CreateModelSessionRequest (model_uri = f"upload://{ id_ } " , deviceIds = ["cpu" ])
5661 model = grpc_stub .CreateModelSession (rq )
5762 assert model .id
58- grpc_stub .CloseModelSession (model )
5963
6064 def test_model_session_creation_using_random_uri (self , grpc_stub ):
6165 rq = inference_pb2 .CreateModelSessionRequest (model_uri = "randomSchema://" , deviceIds = ["cpu" ])
@@ -92,36 +96,28 @@ def test_if_model_create_fails_devices_are_released(self, grpc_stub):
9296 model_blob = inference_pb2 .Blob (content = b"" ), deviceIds = ["cpu" ]
9397 )
9498
95- model = None
9699 with pytest .raises (Exception ):
97- model = grpc_stub .CreateModelSession (model_req )
100+ grpc_stub .CreateModelSession (model_req )
98101
99102 device_by_id = self ._query_devices (grpc_stub )
100103 assert "cpu" in device_by_id
101104 assert inference_pb2 .Device .Status .AVAILABLE == device_by_id ["cpu" ].status
102105
103- if model :
104- grpc_stub .CloseModelSession (model )
105-
106106 def test_use_device (self , grpc_stub , bioimageio_model_bytes ):
107107 device_by_id = self ._query_devices (grpc_stub )
108108 assert "cpu" in device_by_id
109109 assert inference_pb2 .Device .Status .AVAILABLE == device_by_id ["cpu" ].status
110110
111- model = grpc_stub .CreateModelSession (valid_model_request (bioimageio_model_bytes , device_ids = ["cpu" ]))
111+ grpc_stub .CreateModelSession (valid_model_request (bioimageio_model_bytes , device_ids = ["cpu" ]))
112112
113113 device_by_id = self ._query_devices (grpc_stub )
114114 assert "cpu" in device_by_id
115115 assert inference_pb2 .Device .Status .IN_USE == device_by_id ["cpu" ].status
116116
117- grpc_stub .CloseModelSession (model )
118-
119117 def test_using_same_device_fails (self , grpc_stub , bioimageio_model_bytes ):
120- model = grpc_stub .CreateModelSession (valid_model_request (bioimageio_model_bytes , device_ids = ["cpu" ]))
118+ grpc_stub .CreateModelSession (valid_model_request (bioimageio_model_bytes , device_ids = ["cpu" ]))
121119 with pytest .raises (grpc .RpcError ):
122- model = grpc_stub .CreateModelSession (valid_model_request (bioimageio_model_bytes , device_ids = ["cpu" ]))
123-
124- grpc_stub .CloseModelSession (model )
120+ grpc_stub .CreateModelSession (valid_model_request (bioimageio_model_bytes , device_ids = ["cpu" ]))
125121
126122 def test_closing_session_releases_devices (self , grpc_stub , bioimageio_model_bytes ):
127123 model = grpc_stub .CreateModelSession (valid_model_request (bioimageio_model_bytes , device_ids = ["cpu" ]))
@@ -163,8 +159,6 @@ def test_call_predict_valid_explicit(self, grpc_stub, bioimageio_dummy_explicit_
163159 input_tensors = [converters .xarray_to_pb_tensor (input_tensor_id , arr )]
164160 res = grpc_stub .Predict (inference_pb2 .PredictRequest (modelSessionId = model .id , tensors = input_tensors ))
165161
166- grpc_stub .CloseModelSession (model )
167-
168162 assert len (res .tensors ) == 1
169163 assert res .tensors [0 ].tensorId == output_tensor_id
170164 assert_array_equal (expected , converters .pb_tensor_to_numpy (res .tensors [0 ]))
@@ -175,7 +169,6 @@ def test_call_predict_invalid_shape_explicit(self, grpc_stub, bioimageio_dummy_e
175169 input_tensors = [converters .xarray_to_pb_tensor ("input" , arr )]
176170 with pytest .raises (grpc .RpcError ):
177171 grpc_stub .Predict (inference_pb2 .PredictRequest (modelSessionId = model .id , tensors = input_tensors ))
178- grpc_stub .CloseModelSession (model )
179172
180173 @pytest .mark .parametrize (
181174 "shape" ,
@@ -187,7 +180,6 @@ def test_call_predict_invalid_shape_parameterized(self, grpc_stub, shape, bioima
187180 input_tensors = [converters .xarray_to_pb_tensor ("param" , arr )]
188181 with pytest .raises (grpc .RpcError ):
189182 grpc_stub .Predict (inference_pb2 .PredictRequest (modelSessionId = model .id , tensors = input_tensors ))
190- grpc_stub .CloseModelSession (model )
191183
192184 def test_call_predict_invalid_tensor_ids (self , grpc_stub , bioimageio_dummy_model ):
193185 model_bytes , _ = bioimageio_dummy_model
@@ -197,7 +189,6 @@ def test_call_predict_invalid_tensor_ids(self, grpc_stub, bioimageio_dummy_model
197189 with pytest .raises (grpc .RpcError ) as error :
198190 grpc_stub .Predict (inference_pb2 .PredictRequest (modelSessionId = model .id , tensors = input_tensors ))
199191 assert error .value .details ().startswith ("Exception calling application: Spec invalidTensorName doesn't exist" )
200- grpc_stub .CloseModelSession (model )
201192
202193 def test_call_predict_invalid_axes (self , grpc_stub , bioimageio_dummy_model ):
203194 model_bytes , tensor_id = bioimageio_dummy_model
@@ -207,15 +198,13 @@ def test_call_predict_invalid_axes(self, grpc_stub, bioimageio_dummy_model):
207198 with pytest .raises (grpc .RpcError ) as error :
208199 grpc_stub .Predict (inference_pb2 .PredictRequest (modelSessionId = model .id , tensors = input_tensors ))
209200 assert error .value .details ().startswith ("Exception calling application: Incompatible axes" )
210- grpc_stub .CloseModelSession (model )
211201
212202 @pytest .mark .parametrize ("shape" , [(1 , 1 , 64 , 64 ), (1 , 1 , 66 , 65 ), (1 , 1 , 68 , 66 ), (1 , 1 , 70 , 67 )])
213203 def test_call_predict_valid_shape_parameterized (self , grpc_stub , shape , bioimageio_dummy_param_model_bytes ):
214204 model = grpc_stub .CreateModelSession (valid_model_request (bioimageio_dummy_param_model_bytes ))
215205 arr = xr .DataArray (np .arange (np .prod (shape )).reshape (* shape ), dims = ("b" , "c" , "x" , "y" ))
216206 input_tensors = [converters .xarray_to_pb_tensor ("param" , arr )]
217207 grpc_stub .Predict (inference_pb2 .PredictRequest (modelSessionId = model .id , tensors = input_tensors ))
218- grpc_stub .CloseModelSession (model )
219208
220209 @pytest .mark .skip
221210 def test_call_predict_tf (self , grpc_stub , bioimageio_dummy_tensorflow_model_bytes ):
@@ -227,8 +216,6 @@ def test_call_predict_tf(self, grpc_stub, bioimageio_dummy_tensorflow_model_byte
227216 input_tensors = [converters .xarray_to_pb_tensor (input_tensor_id , arr )]
228217 res = grpc_stub .Predict (inference_pb2 .PredictRequest (modelSessionId = model .id , tensors = input_tensors ))
229218
230- grpc_stub .CloseModelSession (model )
231-
232219 assert len (res .tensors ) == 1
233220 assert res .tensors [0 ].tensorId == output_tensor_id
234221 assert_array_equal (expected , converters .pb_tensor_to_numpy (res .tensors [0 ]))
0 commit comments