1#!/usr/bin/env python3 2# Copyright (c) 2025 Huawei Device Co., Ltd. 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15import shutil 16import subprocess 17import tarfile 18from datetime import datetime 19from pathlib import Path 20import struct 21import argparse 22import zlib 23import sys 24 25sys.path.insert(0, str(Path(__file__).parent / "../compiler")) 26from taihe.utils.resources import ( 27 BUNDLE_PYTHON_RUNTIME_DIR_NAME, 28 DeploymentMode, 29 ResourceLocator, 30 ResourceType, 31) 32 33 34SYSTEM_TYPES = ["linux-x86_64", "windows-x86_64", "darwin-arm64", "darwin-x86_64"] 35 36 37class TaiheBuilder: 38 def __init__(self, project_dir: Path, system: str, version: str): 39 self.system = system 40 self.version = version 41 42 self.project_dir = project_dir 43 self.compiler_dir = self.project_dir / "compiler" 44 self.resources = ResourceLocator.detect() 45 46 self.dist = self.project_dir / "dist" / f"taihe-{system}" / "taihe" 47 if self.dist.exists(): 48 shutil.rmtree(self.dist) 49 self.dist_pyrt_dir = self.dist / "lib" / BUNDLE_PYTHON_RUNTIME_DIR_NAME 50 self.dist_bin_dir = self.dist / "bin" 51 self.dist_resources = ResourceLocator(DeploymentMode.BUNDLE, self.dist) 52 53 def copy_resources(self): 54 for ty in ResourceType: 55 if ty.is_packagable(): 56 shutil.copytree( 57 self.resources.get(ty), 58 self.dist_resources.get(ty), 59 dirs_exist_ok=True, 60 ignore=shutil.ignore_patterns("build", "generated"), 61 ) 62 63 def install_compiler(self): 64 print(f"Copying the compiler for {self.system}") 65 site_packages_dir = next(self.dist_pyrt_dir.glob("lib/python*/site-packages")) 66 67 # TODO rm dummy scripts / fix shebang paths 68 uv_args = [ 69 "uv", 70 "pip", 71 "install", 72 f"--target={site_packages_dir}", 73 "--", 74 str(self.compiler_dir), 75 ] 76 subprocess.run(uv_args, check=True, cwd=self.compiler_dir) 77 78 def extract_python(self): 79 python_tar = ( 80 self.resources.get(ResourceType.DEV_PYTHON_BUILD) 81 / f"{self.system}-python.tar.gz" 82 ) 83 # First, extract to "dist/lib". 84 # It creates a sole directory named "python" 85 parent_dir = self.dist_pyrt_dir.parent 86 with tarfile.open(python_tar, "r:gz") as tar: 87 tar.extractall(parent_dir, filter="tar") 88 # Next, rename "dist/lib/python" to "dist/lib/pyrt" 89 old_dir = parent_dir / "python" 90 if self.dist_pyrt_dir.exists(): 91 self.dist_pyrt_dir.rmdir() 92 old_dir.rename(self.dist_pyrt_dir) 93 94 def create_script(self, binary: str, pymod: str): 95 # TODO use absolute script path 96 is_windows = self.system.startswith("windows") 97 prog = self.dist_bin_dir / (f"{binary}.bat" if is_windows else binary) 98 if is_windows: 99 content = ( 100 "@echo off\n" 101 "set TAIHE_ROOT=%~dp0..\n" 102 f'%TAIHE_ROOT%\\lib\\pyrt\\bin\\python3.11.exe -m "{pymod}" %*\n' 103 ) 104 else: 105 content = ( 106 "#!/bin/bash -eu\n" 107 'export TAIHE_ROOT="$(realpath $(dirname "$0")/..)"\n' 108 f'exec "$TAIHE_ROOT/lib/pyrt/bin/python3" -m \'{pymod}\' "$@"\n' 109 ) 110 prog.write_text(content) 111 prog.chmod(0o755) 112 113 def create_taihec_script(self): 114 self.dist_bin_dir.mkdir() 115 print(f"Creating taihec script for {self.system}") 116 if self.system.startswith("linux"): 117 self.create_script("taihe-tryit", "taihe.cli.tryit") 118 self.create_script("taihec", "taihe.cli.compiler") 119 120 def write_version(self): 121 print(f"Writing version information for {self.system}") 122 version_path = self.dist / "version.txt" 123 124 git_commit = subprocess.run( 125 ["git", "rev-parse", "HEAD"], 126 capture_output=True, 127 text=True, 128 check=True, 129 cwd=self.project_dir, 130 ).stdout.strip() 131 132 git_message = ( 133 subprocess.run( 134 ["git", "log", "-1", "--pretty=%B"], 135 capture_output=True, 136 text=True, 137 check=True, 138 cwd=self.project_dir, 139 ) 140 .stdout.splitlines()[0] 141 .strip() 142 ) 143 144 with open(version_path, "w") as version_file: 145 version_file.write(f"System : {self.system}\n") 146 version_file.write(f"Version : {self.version}\n") 147 version_file.write(f"Last Commit : {git_commit}\n") 148 version_file.write(f"Commit Message: {git_message}\n") 149 150 def create_package(self): 151 print(f"Creating package for {self.system}") 152 package_name = ( 153 f"taihe-{self.system}-{self.version}-{datetime.now().strftime('%Y%m%d')}" 154 ) 155 tar_path = self.project_dir / f"{package_name}.tar" 156 zip_path = self.project_dir / f"{package_name}.tar.gz" 157 158 def reset_tarinfo(tarinfo: tarfile.TarInfo): 159 tarinfo.uid = 0 160 tarinfo.gid = 0 161 tarinfo.uname = "" 162 tarinfo.gname = "" 163 tarinfo.mtime = 1704067200 164 return tarinfo 165 166 try: 167 with tarfile.open(tar_path, "w:", format=tarfile.GNU_FORMAT) as tar: 168 for file_path in sorted(self.dist.rglob("*")): 169 if file_path.is_file(): 170 arcname = f"{self.dist.name}/{file_path.relative_to(self.dist)}" 171 tar.add(file_path, arcname=arcname, filter=reset_tarinfo) 172 with open(tar_path, "rb") as tar_file, open(zip_path, "wb") as zip_file: 173 zip_file.write(b"\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\x03") 174 compressor = zlib.compressobj( 175 9, 176 zlib.DEFLATED, 177 -zlib.MAX_WBITS, 178 zlib.DEF_MEM_LEVEL, 179 0, 180 ) 181 data = tar_file.read() 182 compressed = compressor.compress(data) + compressor.flush() 183 zip_file.write(compressed) 184 crc = zlib.crc32(data) 185 size = len(data) 186 zip_file.write(struct.pack("<II", crc & 0xFFFFFFFF, size & 0xFFFFFFFF)) 187 finally: 188 if tar_path.exists(): 189 tar_path.unlink() 190 191 print(f"Created package for {self.system}: {zip_path}") 192 193 def build(self): 194 self.extract_python() 195 self.install_compiler() 196 self.copy_resources() 197 self.create_taihec_script() 198 self.write_version() 199 self.create_package() 200 201 202def main(): 203 parser = argparse.ArgumentParser(description="Build the Taihe project.") 204 parser.add_argument( 205 "--system", 206 choices=SYSTEM_TYPES, 207 nargs="*", 208 default=None, 209 help="Specify the system type to build for. Can be used multiple times.", 210 ) 211 args = parser.parse_args() 212 213 project_dir = Path(__file__).parent.parent 214 215 # TODO detect version use uv / hatch 216 try: 217 version = subprocess.run( 218 ["git", "describe", "--tags", "--abbrev=0"], 219 capture_output=True, 220 text=True, 221 check=True, 222 cwd=project_dir, 223 ).stdout.strip() 224 except subprocess.CalledProcessError: 225 short_hash = subprocess.run( 226 ["git", "rev-parse", "--short=8", "HEAD"], 227 capture_output=True, 228 text=True, 229 check=True, 230 cwd=project_dir, 231 ).stdout.strip() 232 version = f"TRUNK-{short_hash}" 233 234 if args.system is None: 235 systems = SYSTEM_TYPES 236 else: 237 systems = args.system 238 239 print(f"Version is {version}") 240 for system in systems: 241 print(f"Building for {system}") 242 builder = TaiheBuilder(project_dir, system, version) 243 builder.build() 244 245 246if __name__ == "__main__": 247 main() 248