axon icon indicating copy to clipboard operation
axon copied to clipboard

Add Axon Serving API

Open seanmor5 opened this issue 1 year ago • 0 comments

WIP: Need to add tests

From our original implementation I added a few features, and there's still more I think we need to determine before merging:

  1. preprocess/postprocess - I found when trying to deploy a pipeline, it makes sense to have something like this. For example, imagine an image classification task, you probably want the logic of featurizing and processing config to be batched/handled in batches rather than worry about doing it yourself. Both preprocess and postprocess take an input and state (discussed later) and return tensors. The issue right now with is that if we focus on pre-processing raw inputs, it's not possible to determine batch size unless we allow predict to have a signature which accepts either tensors, raw, or list of raw inputs.

  2. state - Again imagine the image classification app. You have a featurizer which is loaded on startup. it makes sense to include this as a part of the serving state, then you can use it in preprocess and postprocess. The state is arbitrary, so it can also be something like a tokenizer or other.

  3. We still do not correctly handle container outputs. I think we can safely just do the padding of map outputs by just padding. I don't think it's possible to have an unbatched output.

  4. We still do not correctly handle unbatched inputs. I handled optional inputs by discarding those not specified in :shape, but for inputs like head_mask where they are not batched, the API will have a fit and raise if their leading dim is not the same as batch_size

  5. We also still have the question of starting a task to pad/split batches and do inference. I guess the only benefit here is we do not block the queue from filling up? Not sure I have enough info to make a decision ATM.

Otherwise, I will add tests tomorrow and then once we address these concerns I think this is good to merge :)

seanmor5 avatar Oct 03 '22 01:10 seanmor5