arkcompiler_runtime_core/taihe/scripts/build

#!/usr/bin/env python3
# Copyright (c) 2025 Huawei Device Co., Ltd.
# 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
from enum import Enum
from functools import cache
from io import StringIO
import shutil
import subprocess
import sys
from typing import TYPE_CHECKING, Final
from pathlib import Path

# Directory for ".git"
g_repo_dir = Path(__file__).parent.parent
# Directory for "taihe"
g_compiler_dir = g_repo_dir / "compiler"
sys.path.insert(0, str(g_compiler_dir))

# ruff: noqa
from taihe.utils.resources import (
    BUILTIN_RESOURCES,
    DeploymentMode,
    PythonBuild,
    ResourceContext,
)

if TYPE_CHECKING:
    # For lazy import
    from taihe.utils.build_metadata import BuildMetadata


class Platform(Enum):
    LINUX_X86_64 = "linux-x86_64"
    WINDOWS_X86_64 = "windows-x86_64"
    DARWIN_ARM64 = "darwin-arm64"
    DARWIN_X86_64 = "darwin-x86_64"

    def is_windows(self):
        return self.value.startswith("windows-")


def create_tar(in_dir: Path, out_tgz: Path, rename_root: str, timestamp: float):
    out_tgz.unlink(missing_ok=True)
    try:
        tar_cmd = ["tar", "--create", "--gzip", "--format=gnu", "--file", out_tgz]
        # Ensure consistent user name.
        tar_cmd += ["--owner=0", "--group=0", "--numeric-owner"]
        # Ensure consistent ordering and timestamp (2024-01-01)
        tar_cmd += ["--sort=name", f"--mtime=@{timestamp}"]
        # Rename the root directory.
        tar_cmd += [f"--transform=s#^{in_dir.name}#{rename_root}#"]
        # Pass the root dir now.
        tar_cmd += ["-C", in_dir.parent, in_dir.name]
        subprocess.run(tar_cmd, check=True, capture_output=True, text=True)
    except subprocess.CalledProcessError as e:
        print(f"Error creating package: {e}")
        print(f"Command output: {e.stderr}")
        raise
    except Exception as e:
        print(f"Unexpected error: {e}")
        raise


class PythonBundle:
    """A prebuilt Python with bundled packages."""

    def __init__(self, root_dir: Path, platform: Platform):
        self.platform = platform
        self.root_dir = root_dir.absolute()
        self.pyrt_dir = self.root_dir / "lib" / PythonBuild.BUNDLE_DIR_NAME

        shutil.rmtree(self.root_dir, ignore_errors=True)
        self.root_dir.mkdir()

    @property
    @cache
    def sitepkg_dir(self) -> Path:
        return next(self.pyrt_dir.glob("lib/python*/site-packages"))

    def extract(self):
        self.pyrt_dir.parent.mkdir(parents=True, exist_ok=True)
        PythonBuild.resolve().extract_to(self.pyrt_dir, system=self.platform.value)

    def install_package(self, wheel: Path):
        # TODO: rm dummy scripts / fix shebang paths
        uv_args = ["uv", "pip", "install", f"--target={self.sitepkg_dir}"]
        uv_args += ["--", str(wheel)]
        subprocess.run(uv_args, check=True)

    def create_script(self, binary: str, pymod: str):
        if self.platform.is_windows():
            suffix_script = ".bat"
            python_executable = "python3.11.exe"
            path_separator = "\\"
            template = (
                "@echo off\n"
                "set TAIHE_ROOT=%~dp0..\n"
                '"%TAIHE_ROOT%\\_PYTHON_REL_" "%TAIHE_ROOT%\\_PYMOD_REL_" %*\n'
            )
        else:
            suffix_script = ""
            python_executable = "python3"
            path_separator = "/"
            template = (
                "#!/bin/bash -eu\n"
                'export TAIHE_ROOT="$(realpath $(dirname "$0")/..)"\n'
                'exec "$TAIHE_ROOT/_PYTHON_REL_" "$TAIHE_ROOT/_PYMOD_REL_" "$@"\n'
            )

        python_path = self.pyrt_dir / f"bin"/ python_executable
        python_rel = python_path.relative_to(self.root_dir)
        pymod_path = (self.sitepkg_dir / pymod.replace(".", "/")).with_suffix(".py")
        pymod_rel = pymod_path.relative_to(self.root_dir)
        content = (
            template.replace("_PYTHON_REL_", str(python_rel))
            .replace("_PYMOD_REL_", str(pymod_rel))
            .replace("/", path_separator)
        )

        prog = self.root_dir / "bin" / f"{binary}{suffix_script}"
        prog.parent.mkdir(exist_ok=True)
        prog.write_text(content)
        prog.chmod(0o755)


class Builder:
    def __init__(self, pyproject_dir: Path, dist_dir: Path):
        self.pyproject_dir = pyproject_dir
        self.dist_dir = dist_dir
        self.brand: Final = "taihe"
        self._has_built_compiler = False

    @property
    def metadata(self) -> "BuildMetadata":
        # Perhaps we should use importlib to conform with self.pyproject_dir?
        # Lazy import, after "rebuild_compiler"
        from taihe.utils.build_metadata import BuildMetadata

        return BuildMetadata.get()

    @property
    @cache
    def release_note(self) -> str:
        metadata = self.metadata
        version_txt_buf = StringIO()
        metadata.print_info(output=version_txt_buf, auto_exit=False)
        return version_txt_buf.getvalue()

    @property
    @cache
    def wheel(self) -> Path:
        return next(self.dist_dir.glob(f"{self.brand}-{self.metadata.version}-*.whl"))

    def rebuild_compiler(self):
        if self._has_built_compiler:
            return  # Only build once.
        self._has_built_compiler = True
        subprocess.run(["uv", "build"], cwd=self.pyproject_dir, check=True)

    def build(self, platform: Platform):
        ### Preparation ###
        self.rebuild_compiler()
        dist_prefix = f"{self.brand}-{platform.value}"  # taihe-linux-x86_64
        bundle_root = self.dist_dir / dist_prefix  # repo/dist/taihe-linux-x86_64
        # bundle_out = repo/dist/taihe-linux-x86_64-0.35.0-20250701.tar.gz # adapt to pre built download scripts
        bundle_out = (
            self.dist_dir
            / f"{dist_prefix}-{self.metadata.version}-{self.metadata.build_date}.tar.gz"
        )

        ### Create the bundle ###
        bundle = PythonBundle(bundle_root, platform)
        bundle.extract()
        bundle.install_package(self.wheel)
        bundle.create_script("taihec", "taihe.cli.compiler")
        if platform == Platform.LINUX_X86_64:
            bundle.create_script("taihe-tryit", "taihe.cli.tryit")

        ### Move resources ###
        src_ctx = ResourceContext(DeploymentMode.PKG, bundle.sitepkg_dir / "taihe/data")
        dst_ctx = ResourceContext(DeploymentMode.BUNDLE, bundle.root_dir)
        for r in BUILTIN_RESOURCES:
            src = r.resolve_path(src_ctx)
            dst = r.resolve_path(dst_ctx)
            dst.parent.mkdir(parents=True, exist_ok=True)
            src.rename(dst)

        ### Finalize ###
        (bundle.root_dir / "version.txt").write_text(self.release_note)
        create_tar(
            bundle_root,
            bundle_out,
            rename_root=self.brand,
            timestamp=self.metadata.build_time_utc,
        )


def main():
    parser = argparse.ArgumentParser(description="Build the Taihe project.")
    parser.add_argument(
        "--system",
        choices=[p.value for p in Platform],
        nargs="*",
        default=None,
        help="Specify the system type to build for. Can be used multiple times.",
    )
    args = parser.parse_args()

    ResourceContext.initialize()
    builder = Builder(g_repo_dir, g_repo_dir / "dist")
    for p in args.system or Platform:
        builder.build(Platform(p))


if __name__ == "__main__":
    main()