def main()

in scripts/get_tperf_args.py [0:0]

• 30 lines of code
• 5 McCabe index (conditional complexity)

def main() -> int:
    """
    Script entry point.
    """
    init_logger()
    args = parse_args()

    # Load settings.
    path = Path(args.path)
    meta_path = path / "train" / "meta.json"
    assert meta_path.is_file(), "meta.json not found!"
    with meta_path.open() as f:
        meta = json.load(f)
    flags = meta["flags"]

    # Check that the config is valid.
    assert not flags[
        "use_job_id_in_actor"
    ], "providing the job ID is not currently supported"

    # Trace model if needed.
    traced_path = path / "traced_model.pt"
    if traced_path.exists():
        logger.info("Traced model already exists")
    else:
        trace(path, flags)

    # Extract relevant tperf settings.
    tperf_args = {f: flags[f] for f in TPERF_FLAGS}

    # Post-process / add arguments.
    tperf_args.update(
        {
            "congestion": "rl",
            "cc_env_job_id": -1,
            "cc_env_mode": "local",
            "cc_env_model_file": str(traced_path),
            "cc_env_actions": get_cc_env_actions(flags["cc_env_actions"]),
        }
    )

    # Output tperf args.
    tperf_cmd = " ".join(f"-{k}='{to_str(v)}'" for k, v in tperf_args.items())
    logger.info("tperf command line arguments:\n%s", tperf_cmd)

    return 0