iap_messenger

The IA Parc inference plugin allows developers to easily integrate their inference pipeline into IA Parc's production module.

Installation

pip install iaparc-inference

Usage

• If your inference pipeline support batching:

  from iap_messenger import MsgListener, Message
  
  # Define a callback to query your inference pipeline
  # To load your model only once it is recommended to use a class:
  class MyModel:
      def __init__(self, model_path: str):
          ## Load your model in pytorch, tensorflow or any other backend
      
      def batch_query(msgs: list[Message]) -> list[Message]:
          ''' execute your pipeline on a batch input
              Note:   "parameters" is an optional argument.
                      It can be used to handle URL's query parameters
                      It's a list of key(string)/value(string) dictionaries
          '''
  
  if __name__ == '__main__':
      # Initiate your model class
      my_model = MyModel("path/to/my/model")
  
      # Initiate IAParc listener
      listener = MsgListener(my_model.batch_query)
      # Start the listener
      listener.run()
  

• If your inference pipeline do not support batching:

  from iap_messenger import MsgListener, Message
  
  # Define a callback to query your inference pipeline
  # To load your model only once it is recommended to use a class:
  class MyModel:
      def __init__(self, model_path: str):
          ## Load your model in pytorch, tensorflow or any other backend
      
      def single_query(msg: Message) -> Message:
          ''' execute your pipeline on a single input
              Note:   "parameters" is an optional argument.
                      It can be used to handle URL's query parameters
                      It's a key(string)/value(string) dictionary
          '''
  
  if __name__ == '__main__':
      # Initiate your model class
      my_model = MyModel("path/to/my/model")
  
      # Initiate IAParc listener
      listener = MsgListener(my_model.single_query, batch=1)  # Note that batch size is forced to 1 here
      # Start the listener
      listener.run()
  

Message consumption and replicas

Each input link is served by one durable JetStream consumer named {queue}-{link}, shared by every replica of the worker. Since the IAP_CONSUMER_MODE release the consumer is a pull consumer: a replica only fetches a message (or a batch) when it is free, so nothing ever waits in the buffer of a busy or dead replica while another one is idle, and a crash only leaves the batch in flight to be redelivered after ack_wait. Older versions used a push consumer with a queue group, where the server spreads messages across replicas at random and up to max_ack_pending messages could sit in a dead pod until ack_wait expired.

Concurrent callbacks without losing messages

A callback must return only once its work is done: the message is acked when it returns. Returning at once and finishing in a background task (to overlap requests on an engine that batches internally) acks every message before its result exists, and a crash then loses them all for good. Make the callback wait for its result instead (an async def awaiting the work suits this), and set IAP_MAX_IN_FLIGHT to the number of requests the backend should overlap. A sync callback also works: each one runs in its own thread, so it must be safe to call concurrently.

Upgrading a pipeline from push to pull

The stream is a WorkQueue, which refuses two consumers on the same subjects, so a pipeline is either all push or all pull for a given link. auto makes a rolling update safe: new pods join the push consumer while old pods are bound to it, nothing is delivered twice and no old pod breaks. Once every pod of the queue runs a pull-capable version, the migration happens on its own the next time no pod is bound (scale to zero, NATS restart), or immediately with a rollout restart after setting IAP_CONSUMER_MODE=pull. Note that after a NATS restart, a pod of the old version reconnecting to a consumer that has just been migrated fails to subscribe and exits; the rollout replaces it.

Graceful shutdown

On SIGTERM (or SIGINT) the listener stops fetching, lets the task or batch in flight (every one of them with IAP_MAX_IN_FLIGHT above 1) complete and acks it, hands back with a NAK any message already pushed to the process but not started (push mode), waits up to IAP_SHUTDOWN_DDL_GRACE_SECONDS for large replies to be fetched downstream, then exits with code 0. Set the pod's terminationGracePeriodSeconds above the longest inference plus that grace, otherwise Kubernetes kills the pod mid-task and the message is redelivered after ack_wait. A second signal exits at once.

Large payloads (DDL)

Payloads above 8 MB never travel through NATS: the sender keeps them in its local DDL store and the message only carries a uuid@host:port pointer that the recipient resolves over ZMQ. Environment variables:

The reply format is negotiated per request: a client that can read a two-frame reply ([status], [payload]) asks with the MP: prefix and the server then hands the stored bytes to ZMQ without copying them; any older peer keeps getting the historical single frame. No configuration and no deployment order are required — mixed versions of workers and inferences-api interoperate.

Features

• Dynamic batching
• Autoscalling
• Support both synchronous and asynchronous queries
• Data agnostic

License

This project is licensed under the Apache License Version 2.0 - see the Apache LICENSE file for details.