Files
OpenLLM/make.py
bojiang 193d480b09 fix
2024-05-23 02:00:45 +08:00

67 lines
2.0 KiB
Python

import yaml
import sys
import subprocess
import os
import shutil
import tempfile
import pathlib
with open("recipe.yaml") as f:
RECIPE = yaml.safe_load(f)
BENTOML_HOME = pathlib.Path(os.environ["BENTOML_HOME"])
CONSTANT_YAML_TMPL = r"""
CONSTANT_YAML = '''
{}
'''
"""
if __name__ == "__main__":
if len(sys.argv) == 2:
specified_model = sys.argv[1]
if specified_model not in RECIPE:
raise ValueError(f"Model {specified_model} not found in recipe")
else:
specified_model = None
for model_name, config in RECIPE.items():
if specified_model and model_name != specified_model:
continue
project = config["project"]
model_repo, model_version = model_name.split(":")
with tempfile.TemporaryDirectory() as tempdir:
tempdir = pathlib.Path(tempdir)
shutil.copytree(project, tempdir, dirs_exist_ok=True)
with open(tempdir / "bento_constants.py", "w") as f:
f.write(CONSTANT_YAML_TMPL.format(yaml.dump(config)))
subprocess.run(
["bentoml", "build", str(tempdir), "--version", model_version],
check=True,
cwd=tempdir,
env=os.environ,
)
# delete latest
(BENTOML_HOME / "bentos" / model_repo / "latest").unlink(missing_ok=True)
# link alias
for alias in config.get("alias", []):
if alias == "latest":
ALIAS_PATH = BENTOML_HOME / "bentos" / model_repo / alias
if ALIAS_PATH.exists():
continue
with open(ALIAS_PATH, "w") as f:
f.write(model_version)
else: # bentoml currently only support latest alias, copy to other alias
shutil.copytree(
BENTOML_HOME / "bentos" / model_repo / model_version,
BENTOML_HOME / "bentos" / model_repo / alias,
)