celery_worker.py

'''
    Sets up the appropriate Celery consumer(s):
    - AIController
    - AIWorker
    - DataManagement

    depending on the "AIDE_MODULES" environment
    variable.

    2020-21 Benjamin Kellenberger
'''

import os
import celery
from celery import Celery
from celery.signals import celeryd_after_setup
from kombu import Queue
from kombu.common import Broadcast
from util.configDef import Config

# force enable passive mode
os.environ['PASSIVE_MODE'] = '1'

# parse system config
if not 'AIDE_CONFIG_PATH' in os.environ:
    raise ValueError('Missing system environment variable "AIDE_CONFIG_PATH".')
if not 'AIDE_MODULES' in os.environ:
    raise ValueError('Missing system environment variable "AIDE_MODULES".')
config = Config()

# get Celery worker hostname
hostname='aide'
@celeryd_after_setup.connect
def capture_worker_name(sender, instance, **kwargs):
    hostname = '{0}'.format(sender)

# parse AIDE modules and set up queues
queues = []
aideModules = os.environ['AIDE_MODULES'].split(',')
aideModules = set([a.strip().lower() for a in aideModules])
for m in aideModules:
    module = m.strip().lower()
    if module == 'aicontroller':
        queues.append(Queue('AIController'))
    elif module == 'fileserver':
        queues.append(Queue('FileServer'))
    elif module == 'aiworker':
        queues.append(Queue('AIWorker'))
    elif module == 'labelui':
        queues.append(Queue('ModelMarketplace'))        #TODO: allow outsourcing? Move to FileServer?

queues.extend([
    Broadcast('aide_broadcast'),
    Queue(hostname+'@'+celery.utils.nodenames.gethostname())
])



app = Celery('AIDE',
            broker=config.getProperty('AIController', 'broker_URL'),        #TODO
            backend=config.getProperty('AIController', 'result_backend'))   #TODO
app.conf.update(
    result_backend=config.getProperty('AIController', 'result_backend'),    #TODO
    task_ignore_result=False,
    result_persistent=True,
    accept_content = ['json'],
    task_serializer = 'json',
    result_serializer = 'json',
    task_track_started = True,
    broker_pool_limit=None,                 # required to avoid peer connection resets
    broker_heartbeat = 0,                   # required to avoid peer connection resets
    worker_max_tasks_per_child = 1,         # required to free memory (also CUDA) after each process
    # task_default_rate_limit = 3,            #TODO
    worker_prefetch_multiplier = 1,         #TODO
    task_acks_late = True,
    task_create_missing_queues = True,
    task_queues = tuple(queues),
    task_routes = {
        'general.get_worker_details': {
            'queue': 'aide@'+celery.utils.nodenames.gethostname(),
            'routing_key': 'worker_details'
        },
        'AIController.get_training_images': {
            'queue': 'AIController',
            'routing_key': 'get_training_images'
        },
        'AIController.get_inference_images': {
            'queue': 'AIController',
            'routing_key': 'get_inference_images'
        },
        'AIController.delete_model_states': {
            'queue': 'AIController',
            'routing_key': 'delete_model_states'
        },
        'AIController.get_model_training_statistics': {
            'queue': 'AIController',
            'routing_key': 'get_model_training_statistics'
        },
        'AIController.duplicate_model_state': {
            'queue': 'AIController',
            'routing_key': 'duplicate_model_state'
        },
        'AIWorker.call_update_model': {
            'queue': 'AIWorker',
            'routing_key': 'call_update_model'
        },
        'AIWorker.call_train': {
            'queue': 'AIWorker',
            'routing_key': 'call_train'
        },
        'AIWorker.call_average_model_states': {
            'queue': 'AIWorker',
            'routing_key': 'call_average_model_states'
        },
        'AIWorker.call_inference': {
            'queue': 'AIWorker',
            'routing_key': 'call_inference'
        },
        'DataAdministration.list_images': {
            'queue': 'FileServer',
            'routing_key': 'list_images'
        },
        'DataAdministration.scan_for_images': {
            'queue': 'FileServer',
            'routing_key': 'scan_for_images'
        },
        'DataAdministration.add_existing_images': {
            'queue': 'FileServer',
            'routing_key': 'add_existing_images'
        },
        'DataAdministration.remove_images': {
            'queue': 'FileServer',
            'routing_key': 'remove_images'
        },
        'DataAdministration.prepare_data_download': {
            'queue': 'FileServer',
            'routing_key': 'prepare_data_download'
        },
        'DataAdministration.watch_image_folders': {
            'queue': 'FileServer',
            'routing_key': 'watch_image_folders'
        },
        'DataAdministration.delete_project': {
            'queue': 'FileServer',
            'routing_key': 'delete_project'
        },
        'ModelMarketplace.shareModel': {
            'queue': 'ModelMarketplace',
            'routing_key': 'shareModel'
        },
        'ModelMarketplace.importModelDatabase': {
            'queue': 'ModelMarketplace',
            'routing_key': 'importModelDatabase'
        },
        'ModelMarketplace.importModelURI': {
            'queue': 'ModelMarketplace',
            'routing_key': 'importModelURI'
        },
        'ModelMarketplace.requestModelDownload': {
            'queue': 'ModelMarketplace',
            'routing_key': 'requestModelDownload'
        }
    }
    #task_default_queue = Broadcast('aide_admin')
)


# initialize appropriate consumer functionalities
from util import celeryWorkerCommons

num_modules = 0
if 'aicontroller' in aideModules:
    from modules.AIController.backend import celery_interface as aic_int
    num_modules += 1
if 'aiworker' in aideModules:
    from modules.AIWorker.backend import celery_interface as aiw_int
    num_modules += 1
if 'labelui' in aideModules:
    from modules.ModelMarketplace.backend import celery_interface as mm_int
    num_modules += 1
if 'fileserver' in aideModules:
    from modules.DataAdministration.backend import celery_interface as da_int
    num_modules += 1

    # scanning project folders for new images: set up periodic task
    scanInterval = config.getProperty('FileServer', 'watch_folder_interval', type=float, fallback=60)
    if scanInterval > 0:
        @app.on_after_configure.connect
        def setup_periodic_tasks(sender, **kwargs):
            sender.add_periodic_task(scanInterval, da_int.watchImageFolders.s())


if __name__ == '__main__':
    # launch Celery consumer
    if num_modules:
        app.start()