# Copyright 2021 Google LLC
#
# 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
#
#     https://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.
"""Builds BIG-bench docker image.
"""

import absl.flags
import absl.app
import os
import subprocess

# ----------------------------------------------------------------------------
FLAGS = absl.flags.FLAGS

absl.flags.DEFINE_string("ubuntu", "18.04", "ubuntu version")
absl.flags.DEFINE_string("cuda", "10.1", "cuda version for gpu images")
absl.flags.DEFINE_string("cudnn", "7", "cudnn version for gpu images")


# ----------------------------------------------------------------------------
def main(_):

    host_path = os.path.abspath(os.path.split(os.path.split(__file__)[0])[0]) + "/.."

    tag = "bigbench:latest"
    base_image = (
        f"nvidia/cuda:{FLAGS.cuda}-cudnn{FLAGS.cudnn}-runtime-ubuntu{FLAGS.ubuntu}"
    )

    cmd = [
        "docker",
        "build",
        "-f",
        "scripts/docker/Dockerfile",
        "-t",
        tag,
        "--build-arg",
        f"BASE_IMAGE={base_image}",
        ".",
    ]

    print(" ".join(cmd))

    subprocess.check_call(cmd, cwd=host_path)


# ----------------------------------------------------------------------------
if __name__ == "__main__":
    absl.app.run(main)