def configure_epoch()

in ssiog/training.py [0:0]

• 43 lines of code
• 8 McCabe index (conditional complexity)

def configure_epoch(sources: dict[str, Source], args: argparse.Namespace):
    prefix = [random.choice(args.prefix)]
    
    if td.get_world_size() > 1:
        broadcast_time_start = time.monotonic_ns()
        td.broadcast_object_list(prefix, src=0)
        td.barrier()
        broadcast_time_end = time.monotonic_ns()
        logger.info(f"Prefix broadcast took {(broadcast_time_end - broadcast_time_start) / 1000000} ms.")
    else:
        logger.info("Broadcasting[prefix] is not required as world size is 1.")

    p = prefix[0]
    name = sources[p].name
    filesystem = sources[p].filesystem
    epoch_objects = sources[p].objects.copy()
    random.shuffle(epoch_objects)
    if len(epoch_objects) > args.object_count_limit:
        epoch_objects = epoch_objects[0 : args.object_count_limit]

    if td.get_world_size() > 1:
        broadcast_time_start = time.monotonic_ns()
        td.broadcast_object_list(epoch_objects, src=0)
        td.barrier()
        broadcast_time_end = time.monotonic_ns()
        logger.info(f"Epoch-objects broadcast took {(broadcast_time_end - broadcast_time_start) / 1000000} ms.")
    else:
        logger.info("Broadcasting[epoch-objects] is not required as world size is 1.")
        
    read_order = [random.choice(args.read_order)]
    if td.get_world_size() > 1:
        broadcast_time_start = time.monotonic_ns()
        td.broadcast_object_list(read_order, src=0)
        td.barrier()
        broadcast_time_end = time.monotonic_ns()
        logger.info(f"Read-order broadcast took {(broadcast_time_end - broadcast_time_start) / 1000000} ms.")
    else:
        logger.info("Broadcasting[read-order] is not required as world size is 1.")

    if read_order[0] == "Sequential":
        reader = sequential_reader
    elif read_order[0] == "FileRandom":
        reader = file_random_reader
    elif read_order[0] == "FullRandom":
        reader = full_random_reader
    else:
        raise Exception(f"Unknown reading order {read_order[0]}")

    return (reader, read_order[0], name, filesystem, epoch_objects)