3333 MODEL_CONFIG_FILE_IN_JSON ,
3434 TSStatusCode ,
3535)
36- from ainode .core .exception import ModelNotExistError
36+ from ainode .core .exception import BuiltInModelDeletionError , ModelNotExistError
3737from ainode .core .log import Logger
3838from ainode .core .model .built_in_model_factory import (
3939 download_ltsm_if_necessary ,
@@ -220,12 +220,13 @@ def register_model(self, model_id: str, uri: str):
220220 configs , attributes = fetch_model_by_uri (
221221 uri , model_storage_path , config_storage_path
222222 )
223- self . _model_info_map [ model_id ] = ModelInfo (
223+ model_info = ModelInfo (
224224 model_id = model_id ,
225225 model_type = "" ,
226226 category = ModelCategory .USER_DEFINED ,
227227 state = ModelStates .ACTIVE ,
228228 )
229+ self .register_built_in_model (model_info )
229230 return configs , attributes
230231
231232 def delete_model (self , model_id : str ) -> None :
@@ -235,11 +236,13 @@ def delete_model(self, model_id: str) -> None:
235236 Returns:
236237 None
237238 """
239+ # check if the model is built-in
240+ with self ._lock_pool .get_lock (model_id ).read_lock ():
241+ if self ._is_built_in (model_id ):
242+ raise BuiltInModelDeletionError (model_id )
243+
244+ # delete the user-defined model
238245 storage_path = os .path .join (self ._model_dir , f"{ model_id } " )
239- with self ._lock_pool .get_lock (model_id ).write_lock ():
240- if os .path .exists (storage_path ):
241- shutil .rmtree (storage_path )
242- storage_path = os .path .join (self ._builtin_model_dir , f"{ model_id } " )
243246 with self ._lock_pool .get_lock (model_id ).write_lock ():
244247 if os .path .exists (storage_path ):
245248 shutil .rmtree (storage_path )
0 commit comments