Tensorflow Serving with XLA

Question

Is it possible to enable XLA compilation when doing inference with Tensorflow Serving?

(I am hoping it's just a matter of undocumented configs and that I can avoid implementing a custom Servable).

Answer 1

@njs,

It is actually not suggested doing compilations during inference. Compilations at inference time will cause HBM to be out of memory, resulting the chips to be unable to serve requests.

The recommended solution is to:

1. Use batch-function with allowed batch sizes to restrict the number of compilations at run time.

2. Do all compilations for these allowed batch sizes at model load time instead of inference time. This way your model is ready for inference right after load, rather than going through high latency compilations at inference time.