def main()

in src/main.py [0:0]


def main():
    parser = argparse.ArgumentParser(description="Program to build docker images")
    parser.add_argument("--buildspec", type=str)
    parser.add_argument("--framework", type=str)
    parser.add_argument("--device_types", type=str, default=constants.ALL)
    parser.add_argument("--image_types", type=str, default=constants.ALL)
    parser.add_argument("--py_versions", type=str, default=constants.ALL)

    args = parser.parse_args()

    image_types = []
    py_versions = []
    device_types = []

    if args.device_types != constants.ALL:
        device_types = args.device_types.split(",")

    if args.image_types != constants.ALL:
        image_types = args.image_types.split(",")

    if args.py_versions != constants.ALL:
        py_versions = args.py_versions.split(",")

    # create the empty json file for images
    build_context = os.getenv("BUILD_CONTEXT")

    # Write empty dict to JSON file, so subsequent buildspec steps do not fail in case we skip
    # this build.
    utils.write_to_json_file(constants.TEST_TYPE_IMAGES_PATH, {})

    # Skip tensorflow-1 PR jobs, as there are no longer patch releases being added for TF1
    # Purposefully not including this in developer config to make this difficult to enable
    # TODO: Remove when we remove these jobs completely
    build_name = get_codebuild_project_name()
    if build_context == "PR" and build_name == "dlc-pr-tensorflow-1":
        return

    # A general build will work if build job and build mode are in non-EI, non-NEURON
    # and non-GRAVITON mode, and its framework and image-type has not been disabled.
    general_builder_enabled = (
        config.is_general_builder_enabled_for_this_pr_build(args.framework)
        and not config.is_deep_canary_mode_enabled()
    )

    # An EI dedicated builder will work if in EI mode and its framework not been disabled
    ei_builder_enabled = (
        config.is_ei_builder_enabled_for_this_pr_build(args.framework)
        and not config.is_deep_canary_mode_enabled()
    )

    # A NEURON dedicated builder will work if in NEURON mode and its framework has not been disabled
    neuron_builder_enabled = (
        config.is_neuron_builder_enabled_for_this_pr_build(args.framework)
        and not config.is_deep_canary_mode_enabled()
    )

    # A NEURONX dedicated builder will work if in NEURONX mode and its framework has not
    # been disabled.
    neuronx_builder_enabled = (
        config.is_neuronx_builder_enabled_for_this_pr_build(args.framework)
        and not config.is_deep_canary_mode_enabled()
    )

    # A GRAVITON dedicated builder will work if in GRAVITON mode and its framework has not
    # been disabled.
    graviton_builder_enabled = (
        config.is_graviton_builder_enabled_for_this_pr_build(args.framework)
        and not config.is_deep_canary_mode_enabled()
    )

    # An ARM64 dedicated builder will work if in ARM64 mode and its framework has not
    # been disabled.
    arm64_builder_enabled = (
        config.is_arm64_builder_enabled_for_this_pr_build(args.framework)
        and not config.is_deep_canary_mode_enabled()
    )

    # A HABANA dedicated builder will work if in HABANA mode and its framework has not been disabled
    habana_builder_enabled = (
        config.is_habana_builder_enabled_for_this_pr_build(args.framework)
        and not config.is_deep_canary_mode_enabled()
    )

    # A HUGGINGFACE TRCOMP dedicated builder will work if in HUGGINGFACE_TRCOMP mode and its
    # framework has not been disabled.
    hf_trcomp_builder_enabled = (
        config.is_hf_trcomp_builder_enabled_for_this_pr_build(args.framework)
        and not config.is_deep_canary_mode_enabled()
    )

    # A TRCOMP dedicated builder will work if in TRCOMP mode and its framework has not been disabled
    trcomp_builder_enabled = (
        config.is_trcomp_builder_enabled_for_this_pr_build(args.framework)
        and not config.is_deep_canary_mode_enabled()
    )

    buildspec_file = config.get_buildspec_override() or args.buildspec

    # Ensure that buildspec_file starts with buildspec and ends with yml
    buildspec_pattern = re.compile(r"buildspec\S*\.yml")
    assert buildspec_pattern.match(
        os.path.basename(buildspec_file)
    ), f"{buildspec_file} must match {buildspec_pattern.pattern}. Please rename file."

    # A builder will always work if it is in non-PR context
    if (
        general_builder_enabled
        or ei_builder_enabled
        or neuron_builder_enabled
        or neuronx_builder_enabled
        or graviton_builder_enabled
        or arm64_builder_enabled
        or habana_builder_enabled
        or hf_trcomp_builder_enabled
        or trcomp_builder_enabled
        or build_context != "PR"
    ):
        utils.build_setup(
            args.framework,
            device_types=device_types,
            image_types=image_types,
            py_versions=py_versions,
        )
        image_builder(buildspec_file, image_types, device_types)