def pb_session():
return inference_pb2.ModelSession(
halo=[
inference_pb2.TensorDim(name="x", size=256),
inference_pb2.TensorDim(name="y", size=128),
inference_pb2.TensorDim(name="c", size=1),
]
)
def test_should_raise_on_empty_dtype(self):
tensor = inference_pb2.Tensor(dtype="",
shape=[
inference_pb2.TensorDim(size=1),
inference_pb2.TensorDim(size=2)
])
with pytest.raises(ValueError):
result_arr = pb_tensor_to_numpy(tensor)
def numpy_to_pb_tensor(array: np.ndarray,
axistags=None) -> inference_pb2.Tensor:
if axistags:
shape = [
inference_pb2.TensorDim(size=dim, name=name)
for dim, name in zip(array.shape, axistags)
]
else:
shape = [inference_pb2.TensorDim(size=dim) for dim in array.shape]
return inference_pb2.Tensor(dtype=str(array.dtype),
shape=shape,
buffer=bytes(array))
def CreateModelSession(self,
request: inference_pb2.CreateModelSessionRequest,
context) -> inference_pb2.ModelSession:
if request.HasField("model_uri"):
if not request.model_uri.startswith("upload://"):
raise NotImplementedError("Only upload:// URI supported")
upload_id = request.model_uri.replace("upload://", "")
content = self.__data_store.get(upload_id)
else:
content = request.model_blob.content
lease = self.__device_pool.lease(request.deviceIds)
try:
_, client = start_model_session_process(
model_zip=content, devices=[d.id for d in lease.devices])
except Exception:
lease.terminate()
raise
session = self.__session_manager.create_session()
session.on_close(lease.terminate)
session.on_close(client.shutdown)
session.client = client
try:
model_info = session.client.get_model_info()
except Exception:
lease.terminate()
raise
pb_valid_shapes = []
for shape in model_info.valid_shapes:
pb_shape = []
for tag, size in shape:
pb_shape.append(inference_pb2.TensorDim(size=size, name=tag))
pb_valid_shapes.append(inference_pb2.Shape(dims=pb_shape))
return inference_pb2.ModelSession(
id=session.id,
name=model_info.name,
inputAxes=model_info.input_axes,
outputAxes=model_info.output_axes,
validShapes=pb_valid_shapes,
hasTraining=False,
halo=[
inference_pb2.TensorDim(size=size, name=tag)
for tag, size in model_info.halo
],
)
def xarray_to_pb_tensor(array: xr.DataArray) -> inference_pb2.Tensor:
shape = [
inference_pb2.TensorDim(size=dim, name=name)
for dim, name in zip(array.shape, array.dims)
]
return inference_pb2.Tensor(dtype=str(array.dtype),
shape=shape,
buffer=bytes(array.data))
def numpy_to_pb_tensor(array: np.ndarray) -> inference_pb2.Tensor:
shape = [inference_pb2.TensorDim(size=dim) for dim in array.shape]
return inference_pb2.Tensor(dtype=str(array.dtype), shape=shape, buffer=bytes(array))