recipes/contrib/ml_compiler_opt_docker.py

# Copyright 2022 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Recipe for building docker images for ml-compiler-opt training"""

from PB.recipes.fuchsia.contrib.ml_compiler_opt_docker import InputProperties

DEPS = [
    "fuchsia/buildbucket_util",
    "fuchsia/docker",
    "fuchsia/git_checkout",
    "recipe_engine/buildbucket",
    "recipe_engine/context",
    "recipe_engine/properties",
    "recipe_engine/url",
]

PROPERTIES = InputProperties


def RunSteps(api, props):
    gitiles_commit = api.buildbucket.build.input.gitiles_commit
    if (
        gitiles_commit
        and gitiles_commit.host
        and gitiles_commit.project
        and gitiles_commit.id
    ):
        props.remote = f"https://{gitiles_commit.host}/{gitiles_commit.project}"
        props.revision = gitiles_commit.id
    ml_dir, revision = api.git_checkout(
        repo=props.remote,
        revision=props.revision,
    )
    container_tag = api.url.join(props.registry_url, "ml-compiler-opt")
    with api.context(cwd=ml_dir):
        try:
            api.docker.build(
                dockerfile=ml_dir / "Dockerfile",
                tags=[
                    container_tag,
                    container_tag + ":" + revision,
                ],
            )
            api.docker("push", container_tag)
            api.docker("push", container_tag + ":" + revision)
        finally:
            api.docker.cleanup()


def GenTests(api):
    default_properties = api.properties(
        dry_run=False,
        registry_url="us-central1-docker.pkg.dev/fuchsia-artifacts/fuchsia-clang-mlgo",
        revision="main",
        remote="https://fuchsia.googlesource.com/example",
    )
    buildset_properties = api.properties(
        dry_run=False,
        registry_url="us-central1-docker.pkg.dev/fuchsia-artifacts/fuchsia-clang-mlgo",
        remote="https://fuchsia.googlesource.com/fuchsia",
    )

    yield api.test("default") + default_properties
    yield api.buildbucket_util.test(
        "with_buildset", repo="fuchsia", revision="a" * 40
    ) + buildset_properties