From 4d4091281a43e58ae5364603e5c822911ee4d92d Mon Sep 17 00:00:00 2001 From: allegroai Date: Sat, 8 Oct 2022 02:10:34 +0300 Subject: [PATCH] add more model frameworks support --- clearml_serving/__main__.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/clearml_serving/__main__.py b/clearml_serving/__main__.py index f590619..8023b38 100644 --- a/clearml_serving/__main__.py +++ b/clearml_serving/__main__.py @@ -4,6 +4,8 @@ import sys from argparse import ArgumentParser from pathlib import Path +from clearml.model import Framework + from clearml_serving.serving.model_request_processor import ModelRequestProcessor, CanaryEP from clearml_serving.serving.endpoints import ModelMonitoring, ModelEndpoint, EndpointMetricLogging @@ -420,7 +422,9 @@ def cli(): '--project', type=str, required=True, help='Specifying the project for the model tp be registered in') parser_model_upload.add_argument( - '--framework', type=str, choices=("scikit-learn", "xgboost", "lightgbm", "tensorflow", "pytorch"), + '--framework', type=str, + choices=[p for p in Framework.__dict__.keys() + if not p.startswith("_") and not callable(getattr(Framework, p))]+["custom"], help='[Optional] Specify the model framework: "scikit-learn", "xgboost", "lightgbm", "tensorflow", "pytorch"') parser_model_upload.add_argument( '--publish', action='store_true', @@ -460,7 +464,7 @@ def cli(): help='Base Model endpoint (must be unique)') parser_model_monitor.add_argument( '--engine', type=str, required=True, - help='Model endpoint serving engine (triton, sklearn, xgboost, lightgbm)') + help='Model endpoint serving engine (triton, sklearn, xgboost, lightgbm, custom, custom_async)') parser_model_monitor.add_argument( '--max-versions', type=int, default=1, help='max versions to store (and create endpoints) for the model. highest number is the latest version')