Converting component creation to be async

[vishwanath1306]

One of the bottlenecks in scaling MLTrace occurs when there is a large stream of requests. One of the ways to achieve this is by asynchronously writing the entries to database. Currently, this process is synchronous, thus it becomes difficult to scale.

Implementation

The Celery Task Queue will be deployed with Redis as a broker, this should ensure ordering of writes, and multiple workers can be deployed to speed up the tasks.

Creating a central celery task manager to coordinate multiple tasks that need to be performed. Codeblock will be similar to

app = Flask(__name__)
   
   def make_celery(app):
       celery = Celery(
           app.import_name,
           broker=<BROKER_URL>, backend=<BACKEND_URL>,
       )
       celery.conf.update(app.config)
   
       class ContextTask(celery.Task):
           def __call__(self, *args, **kwargs):
               with app.app_context():
                   return self.run(*args, **kwargs)
   
       celery.Task = ContextTask
       return celery
   
   celery = make_celery(app)

This should also help developers to offload any future tasks to the task queue.

In client.py, create a new function asyn_create_component with following code

@celery.task
 def async_create_component(
     name: str, description: str, owner: str, tags: typing.List[str] = []
 ):
    """ Same content as create_component function """

Todos to enable integration

Storage / DB Layer

• Deploying Redis (KV Store) as Broker for Celery
• Writing redis configuration for celery
• Adding an async_create_component function with @celery.task operator to enable async writes.

API Layer

No changes since this would go under the hood. Optional → adding a tag to keep a run synchronously.

Query Layer

No changes in UI layer

Benchmarking / Testing

• Locust as workload generator for testing the difference in time between

Deployment

• Creating containers for Redis and Celery to run and linking with existing docker-compose.

[shreyashankar]

The bottleneck is not in create_component -- it is largely in commit_component_run. I think you can make async functions for both, and I'll add async functions similarly for other functions if we need more speedups.

Otherwise LGTM! Thanks for taking this on!