Skip to content

Commit 57f22e7

Browse files
authored
[AINode] Fix the bug that the built-in model be deleted (#15888)
1 parent acaa72d commit 57f22e7

File tree

3 files changed

+15
-7
lines changed

3 files changed

+15
-7
lines changed

iotdb-core/ainode/ainode/core/exception.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,11 @@ def __init__(self, msg: str):
9696
self.message = "Built-in model not support: {0}".format(msg)
9797

9898

99+
class BuiltInModelDeletionError(_BaseError):
100+
def __init__(self, model_id: str):
101+
self.message = "Cannot delete built-in model: {0}".format(model_id)
102+
103+
99104
class WrongAttributeTypeError(_BaseError):
100105
def __init__(self, attribute_name: str, expected_type: str):
101106
self.message = "Wrong type for attribute: {0}, expected: {1}".format(

iotdb-core/ainode/ainode/core/model/built_in_model_factory.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ def get_model_attributes(model_type: BuiltInModelType):
103103
attribute_map = naive_forecaster_attribute_map
104104
elif (
105105
model_type == BuiltInModelType.EXPONENTIAL_SMOOTHING
106-
or model_type == BuiltInModelType.HOLTWINTERS.value
106+
or model_type == BuiltInModelType.HOLTWINTERS
107107
):
108108
attribute_map = exponential_smoothing_attribute_map
109109
elif model_type == BuiltInModelType.STL_FORECASTER:

iotdb-core/ainode/ainode/core/model/model_storage.py

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
MODEL_CONFIG_FILE_IN_JSON,
3434
TSStatusCode,
3535
)
36-
from ainode.core.exception import ModelNotExistError
36+
from ainode.core.exception import BuiltInModelDeletionError, ModelNotExistError
3737
from ainode.core.log import Logger
3838
from 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

Comments
 (0)