Task Assigner initial implementation

[aleksandr-mokrov]

Example of usage in Jupyter:
Getting registered tasks:

task_keeper = task_interface
tasks = task_keeper.get_registered_tasks()

Creating Random Assigner:

def random_assigner(collaborators, round_number, **kwargs):
    """Assigning task groups randomly while ensuring target distribution"""
    import random
    random.shuffle(collaborators)
    collaborator_task_map = {}
    for idx, col in enumerate(collaborators):
        # select only 30% collaborators for training and validation, 70% for validation
        if (idx+1)/len(collaborators) <= 0.7:
            collaborator_task_map[col] = tasks.values()
        else:
            collaborator_task_map[col] = [tasks['aggregated_model_validate']]
    return collaborator_task_map

or
creating filtering assigner:

def filter_assigner(collaborators, round_number, **kwargs):
    collaborator_task_map = {}
    exclude_collaborators = ['env_two', 'env_three']
    for collaborator in collaborators:
        if collaborator in exclude_collaborators:
            continue
        collaborator_task_map[collaborator] = [
            tasks['train'], 
            tasks['locally_tuned_model_validate'],
            tasks['aggregated_model_validate']
        ]
    return collaborator_task_map

or
creating filtering assigner by gpu info:

shard_registry = federation.get_shard_registry()

def filter_by_shard_registry_assigner(collaborators, round_number, **kwargs):
    collaborator_task_map = {}
    for collaborator in collaborators:
        col_status = shard_registry.get(collaborator)
        if not col_status or not col_status['is_online']:
            continue
        node_info = col_status['shard_info'].node_info
        # Assign train task if collaborator has GPU with total memory more that 8 GB
        if len(node_info.cuda_devices) > 0 and node_info.cuda_devices[0].memory_total > 8 * 1024**3:
            collaborator_task_map[collaborator] = [
                tasks['train'], 
                tasks['locally_tuned_model_validate'],
                tasks['aggregated_model_validate'],
            ]
        else:
            collaborator_task_map[collaborator] = [
                tasks['aggregated_model_validate'],
            ]
    return collaborator_task_map

or
create assigner with additional validation round

rounds_to_train = 3
total_rounds = rounds_to_train + 1 # fl_experiment.start(..., rounds_to_train=total_rounds,...)

def assigner_with_last_round_validation(collaborators, round_number, **kwargs):
    collaborator_task_map = {}
    for collaborator in collaborators:
        if round_number == total_rounds - 1:
            collaborator_task_map[collaborator] = [
                tasks['aggregated_model_validate'],
            ]
        else:
            collaborator_task_map[collaborator] = [
                tasks['train'], 
                tasks['locally_tuned_model_validate'],
                tasks['aggregated_model_validate']
            ]
    return collaborator_task_map

and then pass assigner into start experiment call(task_assigner key):

fl_experiment.start(
    model_provider=model_interface, 
    task_keeper=task_interface,
    data_loader=fed_dataset,
    task_assigner=assigner,
    rounds_to_train=1,
    opt_treatment='CONTINUE_GLOBAL',
    device_assignment_policy='CUDA_PREFERRED'
)

Future opportunities with that assigner function interface:
Filtering by collaborator status (shard_registry)

def filter_by_shard_registry_assigner(collaborators, round_number, **kwargs):
    shard_registry = kwargs.get('shard_registry')
    collaborator_task_map = {}
    for collaborator in collaborators:
        col_status = shard_registry.get(collaborator)
        if not col_status or not col_status['is_online']:
            continue
        collaborator_task_map[collaborator] = [
            tasks['train'], 
            tasks['locally_tuned_model_validate'],
            tasks['aggregated_model_validate']
        ]
    return collaborator_task_map

[psfoley]

Removing the kwargs resolution from the plan will break the task runner API. Maybe add an if statement for this logic:

    if hasattr(self.task_runner, 'TASK_REGISTRY'):
        func_name = task.function_name

        # There no enough information to understand that it's train or validate task
        # Looks like we should provide task type.
        kwargs = {}
        if func_name == 'validate':
            if task.is_local:
                kwargs['apply'] = 'local'
            else:
                kwargs['apply'] = 'global'
    else:
        func_name = self.task_config[task]['function']
        kwargs = self.task_config[task]['kwargs']

[aleksandr-mokrov]

Changed

[psfoley]

Moving the round increment to the else statement will result in N+1 rounds being executed. This should be changed back

[aleksandr-mokrov]

No, in fact difference only in logging. Real checking that it is time-time-quit happens in get_tasks function. I returned it back, but we should rewrite this part of code in the future.