Resolve merge conflicts and fix issues with setup.sh

This commit is contained in:
Fred de Gier 2022-10-20 16:09:12 +02:00
commit 2a91018792
9 changed files with 115 additions and 92 deletions

View file

@ -3,12 +3,13 @@ A simple script that sets up the model directory of a given model for Triton.
"""
import argparse
import os
import shutil
from pathlib import Path
from string import Template
SCRIPT_DIR = Path(__file__).parent
CONFIG_TEMPLATE_PATH = SCRIPT_DIR/'config_template.pbtxt'
CONFIG_TEMPLATE_PATH = os.path.join(SCRIPT_DIR, 'config_template.pbtxt')
parser = argparse.ArgumentParser()
parser.add_argument("--model_dir", type=str, required=True)
@ -21,11 +22,11 @@ args = parser.parse_args()
# Step1: Make model directory
model_dir_path = Path(args.model_dir)/f"py-{args.model_name}/py-model/1"
model_dir_path = Path(os.path.join(Path(args.model_dir), f"py-{args.model_name}/py-model/1"))
model_dir_path.mkdir(parents=True, exist_ok=True)
# Step 2: copy model.py
shutil.copy(SCRIPT_DIR/'model.py', model_dir_path/'model.py')
shutil.copy(os.path.join(SCRIPT_DIR, 'model.py'), os.path.join(model_dir_path, 'model.py'))
# Step 3: Generate config.pbtxt
with open(CONFIG_TEMPLATE_PATH, 'r') as f:
@ -38,5 +39,6 @@ config = template.substitute(
use_int8=args.use_int8,
use_auto_device_map=args.use_auto_device_map,
)
with open(model_dir_path/'../config.pbtxt', 'w') as f:
with open(os.path.join(model_dir_path, '../config.pbtxt'), 'w') as f:
f.write(config)
print(f"Config written to")