From 66a120a584dc2617feeb0f70c69a2889ceafc1c7 Mon Sep 17 00:00:00 2001 From: Liu Zhengyun Date: Tue, 8 Jul 2025 16:43:10 +0800 Subject: [PATCH] fix_drop_built_in_model --- iotdb-core/ainode/ainode/core/exception.py | 5 +++++ .../ainode/core/model/built_in_model_factory.py | 2 +- .../ainode/ainode/core/model/model_storage.py | 15 +++++++++------ 3 files changed, 15 insertions(+), 7 deletions(-) diff --git a/iotdb-core/ainode/ainode/core/exception.py b/iotdb-core/ainode/ainode/core/exception.py index 977b10cfa04f..e703fe1aaf78 100644 --- a/iotdb-core/ainode/ainode/core/exception.py +++ b/iotdb-core/ainode/ainode/core/exception.py @@ -96,6 +96,11 @@ def __init__(self, msg: str): self.message = "Built-in model not support: {0}".format(msg) +class BuiltInModelDeletionError(_BaseError): + def __init__(self, model_id: str): + self.message = "Cannot delete built-in model: {0}".format(model_id) + + class WrongAttributeTypeError(_BaseError): def __init__(self, attribute_name: str, expected_type: str): self.message = "Wrong type for attribute: {0}, expected: {1}".format( diff --git a/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py b/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py index b822357158b3..dd71aa4b6259 100644 --- a/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py +++ b/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py @@ -103,7 +103,7 @@ def get_model_attributes(model_type: BuiltInModelType): attribute_map = naive_forecaster_attribute_map elif ( model_type == BuiltInModelType.EXPONENTIAL_SMOOTHING - or model_type == BuiltInModelType.HOLTWINTERS.value + or model_type == BuiltInModelType.HOLTWINTERS ): attribute_map = exponential_smoothing_attribute_map elif model_type == BuiltInModelType.STL_FORECASTER: diff --git a/iotdb-core/ainode/ainode/core/model/model_storage.py b/iotdb-core/ainode/ainode/core/model/model_storage.py index 3d745e7ce3b2..15727392e0be 100644 --- a/iotdb-core/ainode/ainode/core/model/model_storage.py +++ b/iotdb-core/ainode/ainode/core/model/model_storage.py @@ -33,7 +33,7 @@ MODEL_CONFIG_FILE_IN_JSON, TSStatusCode, ) -from ainode.core.exception import ModelNotExistError +from ainode.core.exception import BuiltInModelDeletionError, ModelNotExistError from ainode.core.log import Logger from ainode.core.model.built_in_model_factory import ( download_ltsm_if_necessary, @@ -220,12 +220,13 @@ def register_model(self, model_id: str, uri: str): configs, attributes = fetch_model_by_uri( uri, model_storage_path, config_storage_path ) - self._model_info_map[model_id] = ModelInfo( + model_info = ModelInfo( model_id=model_id, model_type="", category=ModelCategory.USER_DEFINED, state=ModelStates.ACTIVE, ) + self.register_built_in_model(model_info) return configs, attributes def delete_model(self, model_id: str) -> None: @@ -235,11 +236,13 @@ def delete_model(self, model_id: str) -> None: Returns: None """ + # check if the model is built-in + with self._lock_pool.get_lock(model_id).read_lock(): + if self._is_built_in(model_id): + raise BuiltInModelDeletionError(model_id) + + # delete the user-defined model storage_path = os.path.join(self._model_dir, f"{model_id}") - with self._lock_pool.get_lock(model_id).write_lock(): - if os.path.exists(storage_path): - shutil.rmtree(storage_path) - storage_path = os.path.join(self._builtin_model_dir, f"{model_id}") with self._lock_pool.get_lock(model_id).write_lock(): if os.path.exists(storage_path): shutil.rmtree(storage_path)