def main()

in scripts/inf2_env.py [0:0]

• 36 lines of code
• 6 McCabe index (conditional complexity)

def main():
    """
    This script determines proper default TGI env variables for the neuron precompiled models to
    work properly
    :return:
    """
    args = parse_cmdline_and_set_env()

    for env_var in env_vars:
        if not os.getenv(env_var):
            break
    else:
        logger.info(
            "All env vars %s already set, skipping, user know what they are doing",
            env_vars,
        )
        sys.exit(0)

    cache_dir = constants.HF_HUB_CACHE

    logger.info("Cache dir %s, model %s", cache_dir, args.model_id)

    config = AutoConfig.from_pretrained(args.model_id, revision=args.revision)
    neuron_config = getattr(config, "neuron", None)
    if neuron_config is not None:
        compatible = check_env_and_neuron_config_compatibility(
            neuron_config, check_compiler_version=False
        )
        if not compatible:
            env_dict = get_env_dict()
            msg = (
                "Invalid neuron config and env. Config {}, env {}, neuronxcc version {}"
            ).format(neuron_config, env_dict, neuronxcc_version)
            logger.error(msg)
            raise Exception(msg)
    else:
        neuron_config = lookup_compatible_cached_model(args.model_id, args.revision)

    if not neuron_config:
        neuron_config = {'static_batch_size': 1, 'static_sequence_length': 128}
        msg = (
            "No compatible neuron config found. Provided env {}, neuronxcc version {}. Falling back to default"
        ).format(get_env_dict(), neuronxcc_version, neuron_config)
        logger.info(msg)

    logger.info("Final neuron config %s", neuron_config)

    neuron_config_to_env(neuron_config)