wehub-resource-sync
403 lines
16 KiB
Python
403 lines
16 KiB
Python
"""Python entrypoint of package."""
|
|
|
|
import dataclasses
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Literal # noqa: UP035
|
|
|
|
from mlc_llm.interface import jit
|
|
from mlc_llm.support import download_cache, logging, style
|
|
|
|
logging.enable_logging()
|
|
logger = logging.getLogger(__name__)
|
|
|
|
SUPPORTED_DEVICES = ["iphone", "macabi", "android"]
|
|
|
|
|
|
def build_model_library(
|
|
package_config: Dict[str, Any], # noqa: UP006
|
|
device: str,
|
|
bundle_dir: Path,
|
|
app_config_path: Path,
|
|
) -> Dict[str, str]: # noqa: UP006
|
|
"""Build model libraries. Return the dictionary of "library prefix to lib path"."""
|
|
# - Create the bundle directory.
|
|
os.makedirs(bundle_dir, exist_ok=True)
|
|
# Clean up all the directories in `output/bundle`.
|
|
logger.info('Clean up all directories under "%s"', str(bundle_dir))
|
|
for content_path in bundle_dir.iterdir():
|
|
if content_path.is_dir():
|
|
shutil.rmtree(content_path)
|
|
|
|
# - Process each model, and prepare the app config.
|
|
app_config_model_list = []
|
|
|
|
model_entries = package_config.get("model_list", [])
|
|
if not isinstance(model_entries, list):
|
|
raise ValueError('The "model_list" in "mlc-package-config.json" is expected to be a list.')
|
|
model_lib_path_for_prepare_libs = package_config.get("model_lib_path_for_prepare_libs", {})
|
|
if not isinstance(model_lib_path_for_prepare_libs, dict):
|
|
raise ValueError(
|
|
'The "model_lib_path_for_prepare_libs" in "mlc-package-config.json" is expected to be '
|
|
"a dict."
|
|
)
|
|
|
|
jit.log_jit_policy()
|
|
|
|
for model_entry in package_config.get("model_list", []):
|
|
# - Parse model entry.
|
|
if not isinstance(model_entry, dict):
|
|
raise ValueError('The element of "model_list" is expected to be a dict.')
|
|
model = model_entry["model"]
|
|
model_id = model_entry["model_id"]
|
|
bundle_weight = model_entry.get("bundle_weight", False)
|
|
overrides = model_entry.get("overrides", {})
|
|
model_lib = model_entry.get("model_lib", None)
|
|
|
|
estimated_vram_bytes = model_entry["estimated_vram_bytes"]
|
|
if not isinstance(model, str):
|
|
raise ValueError('The value of "model" in "model_list" is expected to be a string.')
|
|
if not isinstance(model_id, str):
|
|
raise ValueError('The value of "model_id" in "model_list" is expected to be a string.')
|
|
if not isinstance(bundle_weight, bool):
|
|
raise ValueError(
|
|
'The value of "bundle_weight" in "model_list" is expected to be a boolean.'
|
|
)
|
|
if not isinstance(overrides, dict):
|
|
raise ValueError('The value of "overrides" in "model_list" is expected to be a dict.')
|
|
if model_lib is not None and not isinstance(model_lib, str):
|
|
raise ValueError('The value of "model_lib" in "model_list" is expected to be string.')
|
|
|
|
# - Load model config. Download happens when needed.
|
|
model_path = download_cache.get_or_download_model(model)
|
|
|
|
# - Jit compile if the model lib path is not specified.
|
|
model_lib_path = (
|
|
model_lib_path_for_prepare_libs.get(model_lib, None) if model_lib is not None else None
|
|
)
|
|
if model_lib_path is None:
|
|
if model_lib is None:
|
|
logger.info(
|
|
'Model lib is not specified for model "%s". Now jit compile the model library.',
|
|
model_id,
|
|
)
|
|
else:
|
|
logger.info(
|
|
'Model lib path for "%s" is not specified in "model_lib_path_for_prepare_libs".'
|
|
"Now jit compile the model library.",
|
|
model_lib,
|
|
)
|
|
model_lib_path, model_lib = dataclasses.astuple(
|
|
jit.jit(
|
|
model_path=model_path,
|
|
overrides=overrides,
|
|
device=device,
|
|
system_lib_prefix=model_lib,
|
|
skip_log_jit_policy=True,
|
|
)
|
|
)
|
|
assert model_lib is not None
|
|
model_lib_path_for_prepare_libs[model_lib] = model_lib_path
|
|
|
|
# - Set "model_url"/"model_path" and "model_id"
|
|
app_config_model_entry = {}
|
|
is_local_model = not model.startswith("HF://") and not model.startswith("https://")
|
|
app_config_model_entry["model_id"] = model_id
|
|
app_config_model_entry["model_lib"] = model_lib
|
|
|
|
# - Bundle weight
|
|
if is_local_model and not bundle_weight:
|
|
raise ValueError(
|
|
f'Model "{model}" in "model_list" is a local path.'
|
|
f'Please set \'"bundle_weight": true\' in the entry of model "{model}".'
|
|
)
|
|
if bundle_weight:
|
|
if not os.path.isfile(model_path / "tensor-cache.json"):
|
|
raise ValueError(
|
|
f'Bundle weight is set for model "{model}". However, model weights are not'
|
|
f'found under the directory "{model}". '
|
|
+ (
|
|
"Please follow https://llm.mlc.ai/docs/compilation/convert_weights.html to "
|
|
"convert model weights."
|
|
if is_local_model
|
|
else "Please report this issue to https://github.com/mlc-ai/mlc-llm/issues."
|
|
)
|
|
)
|
|
# Overwrite the model weight directory in bundle.
|
|
bundle_model_weight_path = bundle_dir / model_id
|
|
logger.info(
|
|
"Bundle weight for %s, copy into %s",
|
|
style.bold(model_id),
|
|
style.bold(str(bundle_model_weight_path)),
|
|
)
|
|
if bundle_model_weight_path.exists():
|
|
shutil.rmtree(bundle_model_weight_path)
|
|
shutil.copytree(model_path, bundle_model_weight_path)
|
|
if bundle_weight and device in ["iphone", "macabi"]:
|
|
app_config_model_entry["model_path"] = model_id
|
|
else:
|
|
app_config_model_entry["model_url"] = model.replace("HF://", "https://huggingface.co/")
|
|
|
|
# - estimated_vram_bytes
|
|
app_config_model_entry["estimated_vram_bytes"] = estimated_vram_bytes
|
|
|
|
app_config_model_list.append(app_config_model_entry)
|
|
|
|
# - Dump "mlc-app-config.json".
|
|
app_config_json_str = json.dumps(
|
|
{"model_list": app_config_model_list},
|
|
indent=2,
|
|
)
|
|
with open(app_config_path, "w", encoding="utf-8") as file:
|
|
print(app_config_json_str, file=file)
|
|
logger.info(
|
|
'Dump the app config below to "%s":\n%s',
|
|
str(app_config_path),
|
|
style.green(app_config_json_str),
|
|
)
|
|
return model_lib_path_for_prepare_libs
|
|
|
|
|
|
def validate_model_lib(
|
|
app_config_path: Path,
|
|
package_config_path: Path,
|
|
model_lib_path_for_prepare_libs: dict,
|
|
device: Literal["iphone", "macabi", "android"],
|
|
output: Path,
|
|
) -> None:
|
|
"""Validate the model lib prefixes of model libraries."""
|
|
if device == "android":
|
|
from tvm.support import ndk as cc
|
|
else:
|
|
from tvm.support import cc
|
|
|
|
with open(app_config_path, encoding="utf-8") as file:
|
|
app_config = json.load(file)
|
|
|
|
tar_list = []
|
|
model_set = set()
|
|
|
|
for model, model_lib_path in model_lib_path_for_prepare_libs.items():
|
|
model_lib_path = os.path.join(model_lib_path)
|
|
lib_path_valid = os.path.isfile(model_lib_path)
|
|
if not lib_path_valid:
|
|
raise RuntimeError(f"Cannot find file {model_lib_path} as an {device} model library")
|
|
tar_list.append(model_lib_path)
|
|
model_set.add(model)
|
|
|
|
os.makedirs(output / "lib", exist_ok=True)
|
|
if device in ["iphone", "macabi"]:
|
|
lib_name = "libmodel_iphone.a"
|
|
else:
|
|
lib_name = "libmodel_android.a"
|
|
lib_path = output / "lib" / lib_name
|
|
|
|
def _get_model_libs(lib_path: Path) -> List[str]: # noqa: UP006
|
|
"""Get the model lib prefixes in the given static lib path."""
|
|
global_symbol_map = cc.get_global_symbol_section_map(lib_path)
|
|
libs = []
|
|
suffix = "___tvm_ffi__library_bin"
|
|
for name, _ in global_symbol_map.items():
|
|
if name.endswith(suffix):
|
|
model_lib = name[: -len(suffix)]
|
|
if model_lib.startswith("_"):
|
|
model_lib = model_lib[1:]
|
|
libs.append(model_lib)
|
|
return libs
|
|
|
|
cc.create_staticlib(lib_path, tar_list)
|
|
available_model_libs = _get_model_libs(lib_path)
|
|
logger.info("Creating lib from %s", str(tar_list))
|
|
logger.info("Validating the library %s", str(lib_path))
|
|
logger.info(
|
|
"List of available model libs packaged: %s,"
|
|
" if we have '-' in the model_lib string, it will be turned into '_'",
|
|
str(available_model_libs),
|
|
)
|
|
global_symbol_map = cc.get_global_symbol_section_map(lib_path)
|
|
error_happened = False
|
|
|
|
for item in app_config["model_list"]:
|
|
model_lib = item["model_lib"]
|
|
model_id = item["model_id"]
|
|
if model_lib not in model_set:
|
|
# NOTE: this cannot happen under new setting
|
|
# since if model_lib is not included, it will be jitted
|
|
raise RuntimeError(
|
|
f"ValidationError: model_lib={model_lib} specified for model_id={model_id} "
|
|
"is not included in model_lib_path_for_prepare_libs argument, "
|
|
"This will cause the specific model not being able to load, "
|
|
f"model_lib_path_for_prepare_libs={model_lib_path_for_prepare_libs}"
|
|
)
|
|
|
|
model_prefix_pattern = model_lib.replace("-", "_") + "___tvm_ffi__library_bin"
|
|
if (
|
|
model_prefix_pattern not in global_symbol_map
|
|
and "_" + model_prefix_pattern not in global_symbol_map
|
|
):
|
|
# NOTE: no lazy format is ok since this is a slow pass
|
|
model_lib_path = model_lib_path_for_prepare_libs[model_lib]
|
|
log_msg = (
|
|
"ValidationError:\n"
|
|
f"\tmodel_lib {model_lib} requested in {str(app_config_path)}"
|
|
f" is not found in {str(lib_path)}\n"
|
|
f"\tspecifically the model_lib for {model_lib_path}.\n"
|
|
f"\tcurrent available model_libs in {str(lib_path)}: {available_model_libs}\n"
|
|
f"\tThis can happen when we manually specified model_lib_path_for_prepare_libs"
|
|
f" in {str(package_config_path)}\n"
|
|
f"\tConsider remove model_lib_path_for_prepare_libs (so library can be jitted)"
|
|
"or check the compile command"
|
|
)
|
|
logger.info(log_msg)
|
|
error_happened = True
|
|
|
|
if not error_happened:
|
|
logger.info(style.green("Validation pass"))
|
|
else:
|
|
logger.info(style.red("Validation failed"))
|
|
sys.exit(255)
|
|
|
|
|
|
def build_android_binding(mlc_llm_source_dir: Path, output: Path) -> None:
|
|
"""Build android binding in MLC LLM"""
|
|
mlc4j_path = mlc_llm_source_dir / "android" / "mlc4j"
|
|
|
|
# Move the model libraries to "build/lib/" for linking
|
|
os.makedirs(Path("build") / "lib", exist_ok=True)
|
|
src_path = str(output / "lib" / "libmodel_android.a")
|
|
dst_path = str(Path("build") / "lib" / "libmodel_android.a")
|
|
logger.info('Moving "%s" to "%s"', src_path, dst_path)
|
|
shutil.move(src_path, dst_path)
|
|
|
|
# Build mlc4j
|
|
logger.info("Building mlc4j")
|
|
subprocess.run([sys.executable, mlc4j_path / "prepare_libs.py"], check=True, env=os.environ)
|
|
# Copy built files back to output directory.
|
|
lib_path = output / "lib" / "mlc4j"
|
|
os.makedirs(lib_path, exist_ok=True)
|
|
logger.info('Clean up all directories under "%s"', str(lib_path))
|
|
for content_path in lib_path.iterdir():
|
|
if content_path.is_dir():
|
|
shutil.rmtree(content_path)
|
|
|
|
src_path = str(mlc4j_path / "src")
|
|
dst_path = str(lib_path / "src")
|
|
logger.info('Copying "%s" to "%s"', src_path, dst_path)
|
|
shutil.copytree(src_path, dst_path)
|
|
|
|
src_path = str(mlc4j_path / "build.gradle")
|
|
dst_path = str(lib_path / "build.gradle")
|
|
logger.info('Copying "%s" to "%s"', src_path, dst_path)
|
|
shutil.copy(src_path, dst_path)
|
|
|
|
src_path = str(Path("build") / "output")
|
|
dst_path = str(lib_path / "output")
|
|
logger.info('Copying "%s" to "%s"', src_path, dst_path)
|
|
shutil.copytree(src_path, dst_path)
|
|
|
|
os.makedirs(lib_path / "src" / "main" / "assets")
|
|
src_path = str(output / "bundle" / "mlc-app-config.json")
|
|
dst_path = str(lib_path / "src" / "main" / "assets" / "mlc-app-config.json")
|
|
logger.info('Moving "%s" to "%s"', src_path, dst_path)
|
|
shutil.move(src_path, dst_path)
|
|
|
|
|
|
def build_iphone_binding(mlc_llm_source_dir: Path, output: Path) -> None:
|
|
"""Build iOS binding in MLC LLM"""
|
|
# Build iphone binding
|
|
logger.info("Build iphone binding")
|
|
subprocess.run(
|
|
["bash", mlc_llm_source_dir / "ios" / "prepare_libs.sh"],
|
|
check=True,
|
|
env=os.environ,
|
|
)
|
|
|
|
# Copy built libraries back to output directory.
|
|
for static_library in (Path("build") / "lib").iterdir():
|
|
dst_path = str(output / "lib" / static_library.name)
|
|
logger.info('Copying "%s" to "%s"', static_library, dst_path)
|
|
shutil.copy(static_library, dst_path)
|
|
|
|
|
|
def build_macabi_binding(mlc_llm_source_dir: Path, output: Path) -> None:
|
|
"""Build Mac Catalyst binding in MLC LLM"""
|
|
deployment_target = os.environ.get("MLC_MACABI_DEPLOYMENT_TARGET", "18.0")
|
|
macabi_arch = os.environ.get("MLC_MACABI_ARCH", "").strip() or "arm64"
|
|
logger.info("Build macabi binding (deployment target %s)", deployment_target)
|
|
cmd = [
|
|
"bash",
|
|
str(mlc_llm_source_dir / "ios" / "prepare_libs.sh"),
|
|
"--catalyst",
|
|
"--deployment-target",
|
|
deployment_target,
|
|
]
|
|
if macabi_arch:
|
|
cmd += ["--arch", macabi_arch]
|
|
subprocess.run(cmd, check=True, env=os.environ)
|
|
|
|
# Copy built libraries back to output directory.
|
|
build_dir = Path(f"build-maccatalyst-{macabi_arch}")
|
|
for static_library in (build_dir / "lib").iterdir():
|
|
dst_path = str(output / "lib" / static_library.name)
|
|
logger.info('Copying "%s" to "%s"', static_library, dst_path)
|
|
shutil.copy(static_library, dst_path)
|
|
|
|
|
|
def package(
|
|
package_config_path: Path,
|
|
mlc_llm_source_dir: Path,
|
|
output: Path,
|
|
) -> None:
|
|
"""Python entrypoint of package."""
|
|
logger.info('MLC LLM HOME: "%s"', mlc_llm_source_dir)
|
|
|
|
# - Read package config.
|
|
with open(package_config_path, encoding="utf-8") as file:
|
|
package_config = json.load(file)
|
|
if not isinstance(package_config, dict):
|
|
raise ValueError(
|
|
"The content of MLC package config is expected to be a dict with "
|
|
f'field "model_list". However, the content of "{package_config_path}" is not a dict.'
|
|
)
|
|
|
|
# - Read device.
|
|
if "device" not in package_config:
|
|
raise ValueError(f'JSON file "{package_config_path}" is required to have field "device".')
|
|
device = package_config["device"]
|
|
if device not in SUPPORTED_DEVICES:
|
|
raise ValueError(
|
|
f'The "device" field of JSON file {package_config_path} is expected to be one of '
|
|
f'{SUPPORTED_DEVICES}, while "{device}" is given in the JSON.'
|
|
)
|
|
|
|
bundle_dir = output / "bundle"
|
|
app_config_path = bundle_dir / "mlc-app-config.json"
|
|
# - Build model libraries.
|
|
model_lib_path_for_prepare_libs = build_model_library(
|
|
package_config, device, bundle_dir, app_config_path
|
|
)
|
|
# - Validate model libraries.
|
|
validate_model_lib(
|
|
app_config_path,
|
|
package_config_path,
|
|
model_lib_path_for_prepare_libs,
|
|
device,
|
|
output,
|
|
)
|
|
|
|
# - Copy model libraries
|
|
if device == "android":
|
|
build_android_binding(mlc_llm_source_dir, output)
|
|
elif device == "iphone":
|
|
build_iphone_binding(mlc_llm_source_dir, output)
|
|
elif device == "macabi":
|
|
build_macabi_binding(mlc_llm_source_dir, output)
|
|
else:
|
|
assert False, "Cannot reach here"
|
|
|
|
logger.info("All finished.")
|