fix drop bug

Author: CRZbulabula

iotdb-core/ainode/ainode/core/manager/model_manager.py

@@ -56,19 +56,16 @@ def register_model(self, req: TRegisterModelReq) -> TRegisterModelResp:
             )
         except InvalidUriError as e:
             logger.warning(e)
-            self.model_storage.delete_model(req.modelId)
             return TRegisterModelResp(
                 get_status(TSStatusCode.INVALID_URI_ERROR, e.message)
             )
         except BadConfigValueError as e:
             logger.warning(e)
-            self.model_storage.delete_model(req.modelId)
             return TRegisterModelResp(
                 get_status(TSStatusCode.INVALID_INFERENCE_CONFIG, e.message)
             )
         except YAMLError as e:
             logger.warning(e)
-            self.model_storage.delete_model(req.modelId)
             if hasattr(e, "problem_mark"):
                 mark = e.problem_mark
                 return TRegisterModelResp(
@@ -86,7 +83,6 @@ def register_model(self, req: TRegisterModelReq) -> TRegisterModelResp:
             )
         except Exception as e:
             logger.warning(e)
-            self.model_storage.delete_model(req.modelId)
             return TRegisterModelResp(get_status(TSStatusCode.AINODE_INTERNAL_ERROR))
 
     def delete_model(self, req: TDeleteModelReq) -> TSStatus:

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

@@ -211,22 +211,30 @@ def register_model(self, model_id: str, uri: str):
             configs: TConfigs
             attributes: str
         """
-        storage_path = os.path.join(self._model_dir, f"{model_id}")
-        # create storage dir if not exist
-        if not os.path.exists(storage_path):
-            os.makedirs(storage_path)
-        model_storage_path = os.path.join(storage_path, DEFAULT_MODEL_FILE_NAME)
-        config_storage_path = os.path.join(storage_path, DEFAULT_CONFIG_FILE_NAME)
-        configs, attributes = fetch_model_by_uri(
-            uri, model_storage_path, config_storage_path
-        )
-        self._model_info_map[model_id] = ModelInfo(
-            model_id=model_id,
-            model_type="",
-            category=ModelCategory.USER_DEFINED,
-            state=ModelStates.ACTIVE,
-        )
-        return configs, attributes
+        with self._lock_pool.get_lock(model_id).write_lock():
+            storage_path = os.path.join(self._model_dir, f"{model_id}")
+            # create storage dir if not exist
+            if not os.path.exists(storage_path):
+                os.makedirs(storage_path)
+            model_storage_path = os.path.join(storage_path, DEFAULT_MODEL_FILE_NAME)
+            config_storage_path = os.path.join(storage_path, DEFAULT_CONFIG_FILE_NAME)
+            self._model_info_map[model_id] = ModelInfo(
+                model_id=model_id,
+                model_type="",
+                category=ModelCategory.USER_DEFINED,
+                state=ModelStates.LOADING,
+            )
+            try:
+                # TODO: The uri should be fetched asynchronously
+                configs, attributes = fetch_model_by_uri(
+                    uri, model_storage_path, config_storage_path
+                )
+                self._model_info_map[model_id].state = ModelStates.ACTIVE
+                return configs, attributes
+            except Exception as e:
+                logger.error(f"Failed to register model {model_id}: {e}")
+                self._model_info_map[model_id].state = ModelStates.INACTIVE
+                raise e
 
     def delete_model(self, model_id: str) -> None:
         """
@@ -235,12 +243,11 @@ def delete_model(self, model_id: str) -> None:
         Returns:
             None
         """
-        storage_path = os.path.join(self._model_dir, f"{model_id}")
         with self._lock_pool.get_lock(model_id).write_lock():
+            storage_path = os.path.join(self._model_dir, f"{model_id}")
             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():
+            storage_path = os.path.join(self._builtin_model_dir, f"{model_id}")
             if os.path.exists(storage_path):
                 shutil.rmtree(storage_path)
             if model_id in self._model_info_map:

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/model/DropModelProcedure.java

@@ -19,6 +19,7 @@
 
 package org.apache.iotdb.confignode.procedure.impl.model;
 
+import org.apache.iotdb.common.rpc.thrift.TAINodeConfiguration;
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.commons.client.ainode.AINodeClient;
 import org.apache.iotdb.commons.client.ainode.AINodeClientManager;
@@ -101,33 +102,35 @@ protected Flow executeFromState(ConfigNodeProcedureEnv env, DropModelState state
   private void dropModelOnAINode(ConfigNodeProcedureEnv env) {
     LOGGER.info("Start to drop model file [{}] on AI Node", modelName);
 
-    List<Integer> nodeIds =
-        env.getConfigManager().getModelManager().getModelDistributions(modelName);
-    for (Integer nodeId : nodeIds) {
-      try (AINodeClient client =
-          AINodeClientManager.getInstance()
-              .borrowClient(
-                  env.getConfigManager()
-                      .getNodeManager()
-                      .getRegisteredAINode(nodeId)
-                      .getLocation()
-                      .getInternalEndPoint())) {
-        TSStatus status = client.deleteModel(modelName);
-        if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
-          LOGGER.warn(
-              "Failed to drop model [{}] on AINode [{}], status: {}",
-              modelName,
-              nodeId,
-              status.getMessage());
-        }
-      } catch (Exception e) {
-        LOGGER.warn(
-            "Failed to drop model [{}] on AINode [{}], status: {}",
-            modelName,
-            nodeId,
-            e.getMessage());
-      }
-    }
+    List<TAINodeConfiguration> aiNodes =
+        env.getConfigManager().getNodeManager().getRegisteredAINodes();
+    aiNodes.forEach(
+        aiNode -> {
+          int nodeId = aiNode.getLocation().getAiNodeId();
+          try (AINodeClient client =
+              AINodeClientManager.getInstance()
+                  .borrowClient(
+                      env.getConfigManager()
+                          .getNodeManager()
+                          .getRegisteredAINode(nodeId)
+                          .getLocation()
+                          .getInternalEndPoint())) {
+            TSStatus status = client.deleteModel(modelName);
+            if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+              LOGGER.warn(
+                  "Failed to drop model [{}] on AINode [{}], status: {}",
+                  modelName,
+                  nodeId,
+                  status.getMessage());
+            }
+          } catch (Exception e) {
+            LOGGER.warn(
+                "Failed to drop model [{}] on AINode [{}], status: {}",
+                modelName,
+                nodeId,
+                e.getMessage());
+          }
+        });
   }
 
   private void dropModelOnConfigNode(ConfigNodeProcedureEnv env) {