# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# !/usr/bin/env python3
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#!/usr/bin/env python3
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import configparser
import logging
import pathlib
import sys

import google.protobuf.json_format
import google.protobuf.text_format
import tritonclient.grpc

LOGGER = logging.getLogger(__name__)


def _get_model_parameters(config_ini):
    """
    Read model parameters from FasterTransformer INI format.
    Input:
        config_ini: configparser.ConfigParser configuratio objects
    Returns:
        list of model parameters
    """
    excluded_section_names = ["ft_instance_hyperparameter", "structure"]
    sections_names_with_model_parameters = [s for s in config_ini.sections() if s not in excluded_section_names]

    if not sections_names_with_model_parameters:
        LOGGER.error(
            "Could not find section with model parameters in model config.ini while it is required to fill templates"
        )
        sys.exit(-1)

    def _get_model_name(section_name_):
        model_name = config_ini.get(section_name_, "model_name", fallback=None)
        if model_name is None:
            model_name = config_ini.get(section_name_, "_name_or_path", fallback="unknown")
        return model_name

    params_from_model_config = {
        section_name: {"model_type": config_ini.get(section_name, "model_type", fallback="GPT"),}
        for section_name in sections_names_with_model_parameters
    }

    # ensure that for all models it is obtained same parameters
    parameters_from_all_sections = list(set(map(lambda x: tuple(x.items()), params_from_model_config.values())))[0]
    if len(parameters_from_all_sections) != len(list(params_from_model_config.values())[0]):
        LOGGER.error(
            "Found no consistency between model parameters: %s (%d != %d)",
            params_from_model_config,
            len(parameters_from_all_sections),
            len(list(params_from_model_config.values())[0]),
        )
        sys.exit(-1)

    params_from_model_config = list(params_from_model_config.values())[0]
    return params_from_model_config


def _update_template(
    config, model_name, default_model_filename, max_batch_size, parameters, just_update_parameters: bool = True
):
    """
    Update config.pbtxt decoded from file.
    Input:
        config: Triton config
        model_name: model name
        default_model_filename: model file name
        max_batch_size: maximum batch size for Triton configuration
        just_update_parameters: replace or overwrite selector
    Returns:
        updated config
    """
    config["name"] = model_name
    config["default_model_filename"] = default_model_filename
    config["max_batch_size"] = max_batch_size

    parameters = {k: {"string_value": str(v)} for k, v in parameters.items()}
    replace = not just_update_parameters
    if replace:
        config["parameters"] = parameters
    else:
        # overwrite
        config["parameters"] = {**config["parameters"], **parameters}
    return config


def main():

    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")

    parser = argparse.ArgumentParser(description="Generate Triton model config file")
    parser.add_argument("--model-train-name", help="Name of trained model", required=True)
    parser.add_argument("--template-path", help="Path to template of Triton model config file", required=True)
    parser.add_argument("--config-path", help="Path to output Triton model config file", required=True)
    parser.add_argument("--ft-checkpoint", help="Path to FasterTransformer checkpoint", required=True)
    parser.add_argument("--max-batch-size", type=int, help="Max batch size of Triton batcher", required=True)
    parser.add_argument("--pipeline-model-parallel-size", type=int, help="Pipeline model parallel size", required=True)
    parser.add_argument("--tensor-model-parallel-size", type=int, help="Tensor model parallel size", required=True)
    parser.add_argument(
        "--data-type", choices=["fp32", "fp16", "bf16"], help="Data type of weights in runtime", required=True
    )
    parser.add_argument(
        "--int8-mode", action="store_true", help="Enable int8 mode in FasterTransformer Triton backend"
    )
    parser.add_argument(
        "--enable-custom-all-reduce",
        action="store_true",
        help="Enable custom all reduce ops in FasterTransformer Triton backend",
    )
    args = parser.parse_args()

    ft_checkpoint_path = pathlib.Path(args.ft_checkpoint)
    config_ini_path = ft_checkpoint_path / "config.ini"
    config_ini = configparser.ConfigParser()
    with config_ini_path.open("r") as config_file:
        config_ini.read_file(config_file)

    # parse template
    template_path = pathlib.Path(args.template_path)
    template_payload = template_path.read_text()
    model_config_proto = google.protobuf.text_format.Parse(
        template_payload, tritonclient.grpc.model_config_pb2.ModelConfig()
    )
    triton_model_config_template = google.protobuf.json_format.MessageToDict(
        model_config_proto, preserving_proto_field_name=True
    )

    # update template
    params_from_model_config = _get_model_parameters(config_ini)
    parameters = {
        **{
            "data_type": args.data_type.lower(),
            "pipeline_para_size": args.pipeline_model_parallel_size,
            "tensor_para_size": args.tensor_model_parallel_size,
            "model_checkpoint_path": ft_checkpoint_path.as_posix(),
            "int8_mode": int(args.int8_mode),
            "enable_custom_all_reduce": int(args.enable_custom_all_reduce),
        },
        **params_from_model_config,
    }
    model_name = args.model_train_name
    updated_triton_model_config = _update_template(
        triton_model_config_template, model_name, ft_checkpoint_path.name, args.max_batch_size, parameters
    )

    # store template
    updated_triton_model_config = google.protobuf.json_format.ParseDict(
        updated_triton_model_config, tritonclient.grpc.model_config_pb2.ModelConfig()
    )
    updated_triton_model_config_payload = google.protobuf.text_format.MessageToBytes(updated_triton_model_config)

    config_path = pathlib.Path(args.config_path)
    config_path.parent.mkdir(parents=True, exist_ok=True)

    with config_path.open("wb") as config_file:
        config_file.write(updated_triton_model_config_payload)

    LOGGER.info("Config file successfully generated and written to: %s", config_path)


if __name__ == "__main__":
    main()
