Add task_expires_seconds for async endpoint task expiration

[lukasewecker]

Pull Request Summary

Add configurable task expiration for async endpoints, allowing users to specify how long a task can wait in the queue before being discarded by Celery workers. Defaults to 86400 seconds (24 hours) if not specified.

Keep in mind that this change does not automatically discard a task from the queue once it is stale. It only marks a task as stale, which then waits in the queue to be picked up by the worker. Once picked up by a worker, the worker then notices that the task is stale and will not process it. This means that when a stale task is never picked up by the worker (e.g. no worker is available), it stays in the queue.

Changes:

• Add task_expires_seconds field to DTOs, entities, and ORM model
• Pass expires parameter to Celery task queue gateway
• Add K8s annotations for celery autoscaler configuration
• Include database migration for new column
• Add comprehensive unit tests for the feature

Note

• Before deploying, the db migration must be run (alembic update head).
• There was a paremeter task_timeout_seconds before, which was not used in the downstream tasks though. It was all renamed to task_expires_seconds, and is now passed properly to the celery tasks.

Example API usage

When creating a new model endpoint via the api, you can now pass the parameter:

curl -X POST "https://api.example.com/v1/model-endpoints" \                                                                                                                                                                                                                                             
  -H "Content-Type: application/json" \                                                                                                                                                                                                                                                                 
  -d '{                                                                                                                                                                                                                                                                                                 
    "name": "my-endpoint",                                                                                                                                                                                                                                                                              
    "model_bundle_id": "bundle-123",                                                                                                                                                                                                                                                                    
    "endpoint_type": "async",                                                                                                                                                                                                                                                                           
    "min_workers": 1,                                                                                                                                                                                                                                                                                   
    "max_workers": 3,                                                                                                                                                                                                                                                                                   
    "per_worker": 1,                                                                                                                                                                                                                                                                                    
    "cpus": 1,                                                                                                                                                                                                                                                                                          
    "memory": "4Gi",                                                                                                                                                                                                                                                                                    
    "storage": "10Gi",                                                                                                                                                                                                                                                                                  
    "gpus": 0,                                                                                                                                                                                                                                                                                          
    "labels": {},                                                                                                                                                                                                                                                                                       
    "metadata": {},                                                                                                                                                                                                                                                                                     
    "task_expires_seconds": 3600                                                                                                                                                                                                                                                                        
  }'                                                                                                                                                                                                                                                                                                    

Test Plan and Usage Guide

Ran unit tests and added new tests for the new feature.

Also, tested locally with this script (spun up Redis locally):

 """Test task_expires_seconds is passed to Celery"""                                                                                                                                                                                                                                                    
 import os                                                                                                                                                                                                                                                                                              
 os.environ["USE_REDIS_LOCALHOST"] = "1"                                                                                                                                                                                                                                                                
 os.environ["GIT_TAG"] = "test"                                                                                                                                                                                                                                                                         
                                                                                                                                                                                                                                                                                                        
 from model_engine_server.common.dtos.model_endpoints import BrokerType                                                                                                                                                                                                                                 
 from model_engine_server.core.tracing.live_tracing_gateway import LiveTracingGateway                                                                                                                                                                                                                   
 from model_engine_server.infra.gateways.celery_task_queue_gateway import CeleryTaskQueueGateway                                                                                                                                                                                                        
                                                                                                                                                                                                                                                                                                        
 # Create gateway                                                                                                                                                                                                                                                                                       
 gateway = CeleryTaskQueueGateway(                                                                                                                                                                                                                                                                      
     broker_type=BrokerType.REDIS_24H,                                                                                                                                                                                                                                                                  
     tracing_gateway=LiveTracingGateway()                                                                                                                                                                                                                                                               
 )                                                                                                                                                                                                                                                                                                      
                                                                                                                                                                                                                                                                                                        
 print("Testing task submission with expires parameter...\n")                                                                                                                                                                                                                                           
                                                                                                                                                                                                                                                                                                        
 from unittest.mock import patch, MagicMock                                                                                                                                                                                                                                                             
                                                                                                                                                                                                                                                                                                        
 with patch.object(gateway, '_get_celery_dest') as mock_get_dest:                                                                                                                                                                                                                                       
     mock_celery = MagicMock()                                                                                                                                                                                                                                                                          
     mock_celery.send_task.return_value = MagicMock(id="test-task-id")                                                                                                                                                                                                                                  
     mock_get_dest.return_value = mock_celery                                                                                                                                                                                                                                                           
                                                                                                                                                                                                                                                                                                        
     # Test 1: Custom expiration (1 hour)                                                                                                                                                                                                                                                               
     gateway.send_task(                                                                                                                                                                                                                                                                                 
         task_name="test.task",                                                                                                                                                                                                                                                                         
         queue_name="test-queue",                                                                                                                                                                                                                                                                       
         args=[{"test": "data"}],                                                                                                                                                                                                                                                                       
         expires=3600,                                                                                                                                                                                                                                                                                  
     )                                                                                                                                                                                                                                                                                                  
                                                                                                                                                                                                                                                                                                        
     call_kwargs = mock_celery.send_task.call_args.kwargs                                                                                                                                                                                                                                               
     print(f"Test 1 - Custom expiration:")                                                                                                                                                                                                                                                              
     print(f"  expires={call_kwargs.get('expires')} seconds")                                                                                                                                                                                                                                           
     assert call_kwargs['expires'] == 3600, "expires should be 3600"                                                                                                                                                                                                                                    
     print("  ✅ PASSED\n")                                                                                                                                                                                                                                                                             
                                                                                                                                                                                                                                                                                                        
     # Test 2: Default expiration (86400 = 24 hours)                                                                                                                                                                                                                                                    
     mock_celery.reset_mock()                                                                                                                                                                                                                                                                           
     gateway.send_task(                                                                                                                                                                                                                                                                                 
         task_name="test.task",                                                                                                                                                                                                                                                                         
         queue_name="test-queue",                                                                                                                                                                                                                                                                       
         args=[{"test": "data"}],                                                                                                                                                                                                                                                                       
         expires=86400,                                                                                                                                                                                                                                                                                 
     )                                                                                                                                                                                                                                                                                                  
                                                                                                                                                                                                                                                                                                        
     call_kwargs = mock_celery.send_task.call_args.kwargs                                                                                                                                                                                                                                               
     print(f"Test 2 - Default expiration (24 hours):")                                                                                                                                                                                                                                                  
     print(f"  expires={call_kwargs.get('expires')} seconds")                                                                                                                                                                                                                                           
     assert call_kwargs['expires'] == 86400, "expires should be 86400"                                                                                                                                                                                                                                  
     print("  ✅ PASSED\n")                                                                                                                                                                                                                                                                             
                                                                                                                                                                                                                                                                                                        

[dmchoiboi]

I believe this is being used in async_inference_use_cases.py

[lukasewecker]

I just renamed the parameter to task_expires_seconds here. Can you explain more what you mean?

[dmchoiboi]

oh i was mentioning that this is currently used in https://github.com/scaleapi/llm-engine/blob/main/model-engine/model_engine_server/domain/use_cases/async_inference_use_cases.py#L72

[lukasewecker]

The issue was that it wasnt properly passed down to when celery creates the task:

llm-engine/model-engine/model_engine_server/infra/gateways/celery_task_queue_gateway.py

Lines 192 to 197 in 429cf01

	res = celery_dest.send_task(
	name=task_name,
	args=args,
	kwargs=kwargs,
	queue=queue_name,
	)

Also it wasnt possible to configure this parameter. So basically what I did was

• rename it (you could argue if thats necessary)
• pass it down properly to downstream code
• make it configurable from outside