def __init__(self):
def get_critical_env(env):
"""
A more strict way to get critical environment variable
"""
try:
var = os.environ[env]
except KeyError as _:
logger.error(f"{env} not defined")
sys.exit(1)
else:
return var
for env in Preprocessor.env_list:
setattr(self, '_' + env.lower(), get_critical_env(env))
self._trtis_conifig_generator = TRTISConfigGenerator()
self.model_root_path = self._model_store
self.model_path = os.path.join(self.model_root_path,
self._serving_name, "1")
def __init__(self):
def get_critical_env(env):
'''
A more strict way to get critical environment variable
'''
try:
var = os.environ[env]
except KeyError:
logger.error(f'{env} not defined')
sys.exit(1)
else:
return var
for env in Preprocessor.env_list:
setattr(self, '_' + env.lower(), get_critical_env(env))
self.using_ormbfile = True if os.getenv(USING_ORMBFILE_ENV, '') else False
self._trtis_conifig_generator = TRTISConfigGenerator()
self.model_root_path = self._model_store
self.model_path = os.path.join(
self.model_root_path, self._serving_name, '1')
class Preprocessor:
'''
Preprocessor formats its directory structure and
generate 'config.pbtxt'
'''
env_list = ['MODEL_STORE', 'SERVING_NAME']
def __init__(self):
def get_critical_env(env):
'''
A more strict way to get critical environment variable
'''
try:
var = os.environ[env]
except KeyError:
logger.error(f'{env} not defined')
sys.exit(1)
else:
return var
for env in Preprocessor.env_list:
setattr(self, '_' + env.lower(), get_critical_env(env))
self.using_ormbfile = True if os.getenv(USING_ORMBFILE_ENV,
'') else False
self._trtis_conifig_generator = TRTISConfigGenerator()
self.model_root_path = self._model_store
self.old_model_path = os.path.join(self.model_root_path,
self._serving_name, 'model')
self.model_path = os.path.join(self.model_root_path,
self._serving_name, '1')
def _extract_yaml(self):
try:
buffer, yaml_data = check_model(self._model_store,
self._serving_name)
except Exception as e:
logger.error('error when checking model: ', e)
sys.exit(1)
else:
logger.info(f'extract yaml_data succeed \n{buffer}')
return yaml_data
def _generate_config_pbtxt(self, yaml_data):
try:
config_pbtext_path = os.path.join(self.model_root_path,
self._serving_name)
self._trtis_conifig_generator.generate_config(
yaml_data, config_pbtext_path, self._serving_name)
except Exception as e:
logger.error('error when generating config.pbtxt: ', e)
sys.exit(1)
def _generate_model_setting(self, format, version):
setting = {}
if format == 'SKLearn':
setting = {
'name': self._serving_name,
'version': version,
'implementation': 'mlserver.models.SKLearnModel',
'parameters': {
'uri': os.path.join(self.model_path, SKLEARN_MODEL)
}
}
elif format == 'XGBoost':
setting = {
'name': self._serving_name,
'version': version,
'implementation': 'mlserver.models.XGBoostModel',
'parameters': {
'uri': os.path.join(self.model_path, XGBOOST_MODEL)
}
}
elif format == 'MLlib':
try:
mllibformat = os.environ['MLLIB_FORMAT']
except KeyError:
logger.error('MLLIB_FORMAT not defined')
sys.exit(1)
setting = {
'name': self._serving_name,
'version': version,
'implementation': 'mlservermllib.models.MLLibModel',
'parameters': {
'uri':
os.path.join(self.model_root_path, self._serving_name,
'1'),
'format':
mllibformat
}
}
json_str = json.dumps(setting)
with open(os.path.join(self.model_root_path, 'model-settings.json'),
'w') as json_file:
json_file.write(json_str)
def _format_model(self, format):
format = format.lower()
try:
formatter = ModelFormatter(format)
formatter.format(self.model_path)
except Exception as e:
logger.error(
f'error when formatting directory {self.model_path}: ', e)
sys.exit(1)
def start(self):
# Phase 1: Extract model_format and yaml
format = ""
yaml_data = {}
if self.using_ormbfile:
yaml_data = self._extract_yaml()
if 'format' not in yaml_data.items():
logger.error('model format missing')
return
format = yaml_data["format"]
else:
inputs = os.getenv(INPUTS_ENV, "[]")
outputs = os.getenv(OUTPUTS_ENV, "[]")
format = os.getenv(FORMAT_ENV, "")
yaml_data["format"] = format
yaml_data["signature"] = {}
yaml_data["signature"]["inputs"] = json.loads(inputs)
yaml_data["signature"]["outputs"] = json.loads(outputs)
# Phase 2: Generate 'config.pbtxt' for triton models
if isTritonModel(format):
self._generate_config_pbtxt(yaml_data)
# Phase 3: Generate 'model setting' for mlserver models
if isMLServerModel(format):
# set env for mlserver
os.putenv('MODEL_FORMAT', format)
# get version from ormbfile
if 'version' in yaml_data.items():
version = yaml_data['version']
else:
version = 'v1.0.0'
self._generate_model_setting(format, version)
# Phase 4: Re-organize directory format
if os.path.isdir(self.model_path):
shutil.rmtree(self.model_path)
os.rename(self.old_model_path, self.model_path)
self._format_model(format)
class Preprocessor:
"""
Preprocessor formats its directory structure and
generate 'config.pbtxt'
"""
env_list = ['MODEL_STORE', 'SERVING_NAME']
def __init__(self):
def get_critical_env(env):
"""
A more strict way to get critical environment variable
"""
try:
var = os.environ[env]
except KeyError as _:
logger.error(f"{env} not defined")
sys.exit(1)
else:
return var
for env in Preprocessor.env_list:
setattr(self, '_' + env.lower(), get_critical_env(env))
self._trtis_conifig_generator = TRTISConfigGenerator()
self.model_root_path = self._model_store
self.model_path = os.path.join(self.model_root_path,
self._serving_name, "1")
def _extract_yaml(self):
try:
buffer, yaml_data = check_model(self._model_store,
self._serving_name)
except Exception as e:
logger.error('error when checking model: ', e)
sys.exit(1)
else:
logger.info(f'extract yaml_data succeed \n{buffer}')
return yaml_data
def _generate_config_pbtxt(self, yaml_data):
try:
config_pbtext_path = os.path.join(self.model_root_path,
self._serving_name)
self._trtis_conifig_generator.generate_config(
yaml_data, config_pbtext_path, self._serving_name)
except Exception as e:
logger.error('error when generating config.pbtxt: ', e)
sys.exit(1)
def _generate_model_setting(self, format):
setting = {}
if format == 'SKLearn':
setting = {
'name': self._serving_name,
'version': 'v0.1.0',
'implementation': 'mlserver.models.SKLearnModel',
'parameters': {
'uri': os.path.join(self.model_path, SKLEARN_MODEL)
}
}
elif format == 'XGBoost':
setting = {
'name': self._serving_name,
'version': 'v0.1.0',
'implementation': 'mlserver.models.XGBoostModel',
'parameters': {
'uri': os.path.join(self.model_path, XGBOOST_MODEL)
}
}
json_str = json.dumps(setting)
with open(os.path.join(self.model_root_path, "model-settings.json"),
'w') as json_file:
json_file.write(json_str)
def _format_model(self, format):
format = format.lower()
try:
formatter = ModelFormatter(format)
formatter.format(self.model_path)
except Exception as e:
logger.error(
f'error when formatting directory {self.model_path}: ', e)
sys.exit(1)
def start(self):
ormb_file_path = os.path.join(self.model_root_path, self._serving_name,
"ormbfile.yaml")
if not os.path.exists(ormb_file_path):
return
# Phase 1: Extract model_format and yaml
yaml_data = self._extract_yaml()
format = yaml_data["format"]
MODEL_NEED_NOT_CONFIG_PBTXT = {'PMML', 'SKLearn', 'XGBoost'}
# Phase 2: Generate 'config.pbtxt' if need
if format not in MODEL_NEED_MODEL_SETTING:
self._generate_config_pbtxt(yaml_data)
MODEL_NEED_MODEL_SETTING = {'SKLearn', 'XGBoost'}
# Phase 3: Generate 'model setting' if need
if format in MODEL_NEED_MODEL_SETTING:
self._generate_model_setting(format)
# Phase 4: Re-organize directory format
self._format_model(format)
os.remove(ormb_file_path)
class Preprocessor:
"""
Preprocessor formats its directory structure and
generate 'config.pbtxt'
"""
env_list = [
'MODEL_STORE', 'SERVING_NAME'
]
def __init__(self):
def get_critical_env(env):
"""
A more strict way to get critical environment variable
"""
try:
var = os.environ[env]
except KeyError:
logger.error(f"{env} not defined")
sys.exit(1)
else:
return var
for env in Preprocessor.env_list:
setattr(self, '_' + env.lower(), get_critical_env(env))
self._trtis_conifig_generator = TRTISConfigGenerator()
self.model_root_path = self._model_store
self.model_path = os.path.join(
self.model_root_path, self._serving_name, "1")
def _extract_yaml(self):
try:
buffer, yaml_data = check_model(
self._model_store, self._serving_name)
except Exception as e:
logger.error('error when checking model: ', e)
sys.exit(1)
else:
logger.info(f'extract yaml_data succeed \n{buffer}')
return yaml_data
def _generate_config_pbtxt(self, yaml_data):
try:
config_pbtext_path = os.path.join(
self.model_root_path, self._serving_name)
self._trtis_conifig_generator.generate_config(
yaml_data, config_pbtext_path, self._serving_name)
except Exception as e:
logger.error('error when generating config.pbtxt: ', e)
sys.exit(1)
def _generate_model_setting(self, format, version):
setting = {}
if format == 'SKLearn':
setting = {
'name': self._serving_name,
'version': version,
'implementation': 'mlserver.models.SKLearnModel',
'parameters': {
'uri': os.path.join(self.model_path, SKLEARN_MODEL)
}
}
elif format == 'XGBoost':
setting = {
'name': self._serving_name,
'version': version,
'implementation': 'mlserver.models.XGBoostModel',
'parameters': {
'uri': os.path.join(self.model_path, XGBOOST_MODEL)
}
}
elif format == 'MLlib':
try:
mllibformat = os.environ["MLLIB_FORMAT"]
except KeyError:
logger.error("MLLIB_FORMAT not defined")
sys.exit(1)
setting = {
'name': self._serving_name,
'version': version,
'implementation': 'mlservermllib.models.MLLibModel',
'parameters': {
'uri': os.path.join(
self.model_root_path, self._serving_name, "1"),
'format': mllibformat
}
}
json_str = json.dumps(setting)
with open(os.path.join(self.model_root_path, "model-settings.json"), 'w') as json_file:
json_file.write(json_str)
def _format_model(self, format):
format = format.lower()
try:
formatter = ModelFormatter(format)
formatter.format(self.model_path)
except Exception as e:
logger.error(
f'error when formatting directory {self.model_path}: ', e)
sys.exit(1)
def start(self):
ormb_file_path = os.path.join(
self.model_root_path, self._serving_name, "ormbfile.yaml")
if not os.path.exists(ormb_file_path):
logger.error(f'{ormb_file_path} does not exist')
return
# Phase 1: Extract model_format and yaml
yaml_data = self._extract_yaml()
if 'format' in yaml_data.items():
logger.error('model format missing')
return
format = yaml_data["format"]
# Phase 2: Generate 'config.pbtxt' for triton models
if isTritonModel(format):
self._generate_config_pbtxt(yaml_data)
# Phase 3: Generate 'model setting' for mlserver models
if isMLServerModel(format):
# set env for mlserver
os.putenv('MODEL_FORMAT', format)
# get version from ormbfile
if 'version' in yaml_data.items():
version = yaml_data["version"]
else:
version = 'v1.0.0'
self._generate_model_setting(format, version)
# Phase 4: Re-organize directory format
self._format_model(format)
os.remove(ormb_file_path)