def main()

in tools/auto_fill_neuron_cache.py [0:0]
82 lines of code
13 McCabe index (conditional complexity)

def main():
    args = parse_args()

    now = datetime.now()
    now_str = now.strftime("%d%m%Y_%H_%M")
    success_filename = f"success_{now_str}.txt"
    failure_filename = f"failure_{now_str}.txt"

    if args.models == "all":
        models = list(ARCHITECTURES_TO_COMMON_PRETRAINED_WEIGHTS.keys())
    else:
        models = args.models

    # Test examples are slow, this allows us to run them.
    os.environ["RUN_SLOW"] = "1"

    total_num_cores = get_num_neuron_cores()
    num_cores_configurations = [1, 2, 8, total_num_cores]
    num_cores_configurations = sorted(set(num_cores_configurations))

    for model_type in models:
        testers = get_testers_for_model_type(model_type)
        for task, tester, method_name in testers:
            if args.tasks != "all" and task not in args.tasks:
                continue
            for model_name, shape_values in ARCHITECTURES_TO_COMMON_PRETRAINED_WEIGHTS[model_type].items():
                print(f"Running precompilation for {model_name} on {task}...")
                for num_cores in num_cores_configurations:
                    print(f"For {num_cores} neuron cores")
                    if num_cores == 1:
                        tester.MULTI_PROC = "false"
                    else:
                        tester.MULTI_PROC = "true"
                    start = time.time()
                    shape_values_for_task = shape_values.get(task)
                    if shape_values_for_task is None:
                        shape_values_for_task = shape_values["default"]
                    example = ["*" * 20]
                    example.append(f"Model:\t{model_name}")
                    example.append("Shapes:")
                    for name, value in shape_values_for_task.items():
                        example.append(f"\t{name} = {value}")
                    example.append(f"Num cores:\t{num_cores}")
                    example_str = ""
                    try:
                        start = time.time()
                        example.append("Precision:\tBF16")
                        example.append("*" * 20)
                        example_str = "\n".join(example)
                        run_auto_fill_cache_for_model_name(
                            model_type,
                            model_name,
                            shape_values_for_task,
                            tester,
                            method_name,
                            args.cache_path,
                            num_cores,
                            True,
                        )
                    except Exception as e:
                        print(e)
                        open_and_append_to_file(failure_filename, example_str)
                    else:
                        open_and_append_to_file(success_filename, example_str)

                    end = time.time()
                    print(f"Done! Duration: {end - start:.3f}.")
                    try:
                        start = time.time()
                        example.append("Precision:\tFull-precision")
                        example.append("*" * 20)
                        example_str = "\n".join(example)
                        run_auto_fill_cache_for_model_name(
                            model_type,
                            model_name,
                            shape_values_for_task,
                            tester,
                            method_name,
                            args.cache_path,
                            num_cores,
                            False,
                        )
                    except Exception:
                        open_and_append_to_file(failure_filename, example_str)
                    else:
                        open_and_append_to_file(success_filename, example_str)
                    end = time.time()
                    print(f"Done! Duration: {end - start:.3f}.")

    os.environ["RUN_SLOW"] = "0"