DOCSEARCHbyAlgolia
Submit Docs

Aqueduct Framework

repository·main·Indexed 41 days ago

https://github.com/avito-tech/aqueduct

A framework for building performance-efficient and scalable machine learning prediction services. Aqueduct optimizes throughput using shared memory for data transfer between processes and supports task batching, asyncio compatibility, and web concurrency via SocketFlow and Gunicorn. It allows developers to distribute model steps across subprocesses using BaseTask and BaseTaskHandler abstractions.

Tokens
10K
Snippets
31
Records
43
Agent score
49%

What's inside Aqueduct

  1. Structure of the aiohttp Aqueduct demo project

    main

    The demo project demonstrates how to integrate Aqueduct into an aiohttp web application using the following structure:

    • demoapp/app.py: Responsible for building and running the aiohttp web application.
    • demoapp/pipeline.py: Defines and stores the data processing pipeline.
    • demoapp/flow.py: Contains the logic for wrapping pipeline steps into Aqueduct abstractions and the logic for creating a Flow.
    • data/: Contains sample input data (e.g., images for classification).

  2. How Aqueduct metric names are constructed

    main

    When using the StatsD exporter, metric names are generated by concatenating four parts with dots: aqueduct (constant) + user specified prefix (optional) + metric + metric name.

    Example patterns:

    • aqueduct.transfer_time.from_step2_ModelHandler_to_main_process
    • aqueduct.app_custom_prefix.qsize.from_main_process_to_step1_PreprocessorHandler
    • aqueduct.app_custom_prefix.handle_time.step1_PreprocessorHandler

  3. Best practices for scaling GPU steps in FlowStep

    main

    When using FlowStep with a GPU subscription model, avoid setting nprocs > 1 for a single GPU step. Because GPU models are typically heavy, increasing the number of processes per step significantly increases memory and core consumption on a single instance.

    Recommended Scaling Pattern: Instead of increasing processes per step, scale by increasing the number of service pods (in Kubernetes) or the total number of instances. This keeps pod sizes manageable and simplifies deployment.

  4. How Flow and FlowStep work together

    main

    Aqueduct uses a pipeline model where a Flow manages the lifecycle of multiple FlowStep objects.

    • Flow: The top-level container that manages the pipeline lifecycle, creates queues between steps, monitors process health, and provides the interface to send tasks and receive results.
    • FlowStep: A single stage in the pipeline. Each step can consist of multiple OS processes (nprocs) and uses a Handler to perform the actual work. Steps can be configured with batch_size and batch_timeout to process tasks in groups.

    To build a pipeline, you define one or more FlowStep instances, each with its own BaseTaskHandler, and register them within a Flow.

    from aqueduct import Flow, FlowStep, BaseTaskHandler

# Define your handler
class MyHandler(BaseTaskHandler):
    def on_start(self):
        pass
    def handle(self, *tasks):
        pass

# Create the pipeline
flow = Flow()
flow.add_step(FlowStep(handler=MyHandler(), nprocs=4))
await flow.start()

  5. How Aqueduct differs from ProcessPoolExecutor

    main

    While ProcessPoolExecutor is suitable for running monolithic CPU-bound models in subprocesses, aqueduct is designed for splitting models into multiple steps distributed across subprocesses.

    aqueduct provides several features that ProcessPoolExecutor lacks for model orchestration:

    • Asyncio compatibility: Works seamlessly with asynchronous code.
    • Production metrics: Built-in support for monitoring.
    • Liveness detection: Automatically detects if subprocesses have died.
    • Shared memory: Uses shared memory for efficient data transfer to subprocesses.
    • Automatic queue management: Automatically creates queues for data transfer between steps.
    • Interpreter fixes: Includes fixes for known issues with Python interpreter queues.

  6. Why use Aqueduct instead of external message brokers (Kafka, RabbitMQ, Pulsar)

    main
    Aqueduct uses interprocess queues instead of external messaging brokers like Kafka, RabbitMQ, or Pulsar to prioritize speed and hardware efficiency. External brokers introduce significant latency that makes high-speed inference difficult and are generally considered an anti-pattern when transferring large objects between processes.

  7. Use shared memory to pass large data in Aqueduct

    main

    To avoid high memory consumption when passing large objects (like images or large arrays) between processes, use the share_value_with_data method on a BaseTask instance.

    Instead of attaching large data directly to a task attribute (which would involve expensive serialization), you call await task.share_value_with_data(...). This allows the data to be shared via shared memory.

    Required parameters for share_value_with_data:

    • field_name: The name of the attribute on the BaseTask where the data should be accessible.
    • content: The actual data/bytes to be shared.
    • size: The size of the content in bytes.

    Once the task is processed by a Flow, the results or processed data can be accessed directly from the task object attributes.

    # Example of sharing large data via shared memory
task = Task()
await task.share_value_with_data(
    field_name='image', 
    content=self.request.content,
    size=self.request.content_length,
)

# After processing, the result is available on the task object
await self.request.app['flow'].process(task)
return web.json_response(data={'result': task.image_processed})

  9. Integrate Sentry for error logging

    main

    To capture logger events from Aqueduct workers and the main process in Sentry, you must manually attach a SentryHandler to the Aqueduct log object. This is typically done by checking for a SENTRY_ENABLED environment variable and using the SENTRY_DSN to configure the raven.Client.

    import logging
import os

from raven import Client
from raven.handlers.logging import SentryHandler
from raven.transport.http import HTTPTransport

from aqueduct.logger import log


if os.getenv('SENTRY_ENABLED') is True:
    dsn = os.getenv('SENTRY_DSN')
    sentry_handler = SentryHandler(client=Client(dsn=dsn, transport=HTTPTransport), level=logging.ERROR)
    log.addHandler(sentry_handler)

  10. Export metrics to StatsD using ToStatsDMetricsExporter

    main

    To export metrics to a StatsD server, use the ToStatsDMetricsExporter class. This requires a target object that implements the StatsDBuffer protocol.

    To implement StatsDBuffer, your target class must provide the following methods:

    • count(name: str, value: Union[float, int])
    • timing(name: str, value: Union[float, int])
    • gauge(name: str, value: Union[float, int])

    You can pass an optional prefix string to the exporter to customize the metric name hierarchy.

    import statsd
from typing import Union
from aqueduct import Flow, FlowStep
from aqueduct.metrics import ToStatsDMetricsExporter
from aqueduct.metrics.export import StatsDBuffer

# 1. Implement the StatsDBuffer protocol
class AqueductStatsd(StatsDBuffer):
    def __init__(self):
        self.statsd_client = statsd.StatsClient('localhost', 8125)

    def count(self, name: str, value: Union[float, int]):
        self.statsd_client.incr(name, value)

    def timing(self, name: str, value: Union[float, int]):
        self.statsd_client.timing(name, value)

    def gauge(self, name: str, value: Union[float, int]):
        self.statsd_client.gauge(name, value)

def create_flow() -> Flow:
    aqueduct_statsd = AqueductStatsd()
    return Flow(
        FlowStep(
            SumHandler()
        ),
        # 2. Pass the exporter to the Flow
        metrics_exporter=ToStatsDMetricsExporter(
            target=aqueduct_statsd,
            prefix='app_custom_prefix'
        ),
    )

  11. Improve performance using task batching

    main

    Aqueduct supports processing tasks in batches to improve model performance. To enable batching, specify the batch_size parameter when defining a FlowStep. This parameter indicates the maximum number of tasks that can be included in a single batch.

    Batch Size Behavior:

    • By default, Aqueduct does not guarantee that a handler will receive a batch of the exact batch_size.
    • The actual batch size may be less than the specified batch_size to avoid latency in low-load scenarios.
    • The batch size will never exceed the specified batch_size.
    • If your handler requires a specific, exact batch size for optimal performance, use the batch_timeout parameter to limit how long the system waits to assemble a full batch.
    flow_with_batch_handler = Flow(FlowStep(CatDetectorHandler(), batch_size=20))

  12. Install Aqueduct via pip

    main

    Install the core Aqueduct library using pip:

    pip install aqueduct

    Aqueduct is a framework designed to make machine learning prediction services performance-efficient and scalable by increasing throughput and using shared memory for instantaneous data transfer between processes.