Function tuning_configs

lightllm/utils/tuning_utils.py:95–153

(device_id, device_count, **configs)

def tuning_configs(device_id, device_count, **configs):
    test_func = configs.pop("test_func")
    test_kwargs = configs.pop("test_func_args")
    get_test_configs_func = configs.pop("get_test_configs_func")

    os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
    best_config, best_cost_time = None, float("inf")
    queue = mp.Queue()
    test_configs_list = []
    for t_config in get_test_configs_func(device_id, device_count, **test_kwargs):
        test_configs_list.append(t_config)
        if len(test_configs_list) < 64:
            continue

        p = mp.Process(
            target=workder,
            args=(test_func, test_configs_list, test_kwargs, queue),
        )
        p.start()
        p.join()

        while len(test_configs_list) > 0:
            try:
                cost_time = queue.get_nowait()
                logger.info(f"get {test_configs_list[0]} cost_time: {cost_time}")
                if cost_time < best_cost_time:
                    best_config = test_configs_list[0]
                    best_cost_time = cost_time
                    logger.info(f"current best: {best_config}, cost_time: {best_cost_time}")
                del test_configs_list[0]
            except:
                logger.info(f"current best: {best_config}, cost_time: {best_cost_time}")
                del test_configs_list[0]
                break

    while len(test_configs_list) > 0:
        p = mp.Process(
            target=workder,
            args=(test_func, test_configs_list, test_kwargs, queue),
        )
        p.start()
        p.join()

        while len(test_configs_list) > 0:
            try:
                cost_time = queue.get_nowait()
                logger.info(f"get {test_configs_list[0]} cost_time: {cost_time}")
                if cost_time < best_cost_time:
                    best_config = test_configs_list[0]
                    best_cost_time = cost_time
                    logger.info(f"current best: {best_config}, cost_time: {best_cost_time}")
                del test_configs_list[0]
            except:
                logger.info(f"current best: {best_config}, cost_time: {best_cost_time}")
                del test_configs_list[0]
                break

    logger.info(f"Final best config: {best_config}, cost_time: {best_cost_time}")