chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,402 @@
|
||||
"""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.")
|
||||
Reference in New Issue
Block a user