[MNT] delete_learnware, buggy auto_update

learnware/market/hetergeneous/organizer/__init__.py

@@ -10,6 +10,8 @@ from shutil import copyfile, rmtree
 from typing import List
 
 import pandas as pd
+from torch import nn
+import torch.multiprocessing as mp
 
 from ....learnware import Learnware, get_learnware_from_dirpath
 from ....logger import get_module_logger
@@ -24,7 +26,7 @@ logger = get_module_logger("hetero_market")
 
 
 class HeteroMapTableOrganizer(EasyOrganizer):
-    def reload_market(self, rebuild=False, auto_update_limit=50):
+    def reload_market(self, rebuild=False, auto_update_limit=100):
         self.market_store_path = os.path.join(conf.market_root_path, self.market_id)
         self.market_mapping_path = os.path.join(self.market_store_path, conf.market_model_path)
         self.learnware_pool_path = os.path.join(self.market_store_path, "learnware_pool")
@@ -34,9 +36,21 @@ class HeteroMapTableOrganizer(EasyOrganizer):
         self.learnware_zip_list = {}
         self.learnware_folder_list = {}
         self.count = 0
+        self.last_trained_learnware_num = 0
         self.dbops = DatabaseOperations(conf.database_url, "market_" + self.market_id)
         self.auto_update = False
         self.auto_update_limit = auto_update_limit
+        self.auto_update_lock = mp.Lock()
+        self.is_training_in_progress = mp.Value('i', 0)
+
+        if rebuild:
+            logger.warning("Warning! You are trying to clear current database!")
+            try:
+                self.dbops.clear_learnware_table()
+                rmtree(self.learnware_pool_path)
+            except Exception as err:
+                logger.warning(f"Clear current database failed due to {err}!!")
+                pass
 
         os.makedirs(self.learnware_pool_path, exist_ok=True)
         os.makedirs(self.learnware_zip_pool_path, exist_ok=True)
@@ -49,27 +63,20 @@ class HeteroMapTableOrganizer(EasyOrganizer):
             self.count,
         ) = self.dbops.load_market()
 
-        if rebuild:
-            logger.warning("Warning! You are trying to clear current database!")
-            try:
-                self.dbops.clear_learnware_table()
-                rmtree(self.learnware_pool_path)
-            except:
-                pass
+        if os.path.exists(self.market_mapping_path):
+            logger.info(f"Loading Market Mapping from Default Checkpoint {self.market_mapping_path}")
+            self.market_mapping = HeteroMapping.load(checkpoint=self.market_store_path)
+            # self._update_learnware_list(self.learnware_list)
         else:
-            if os.path.exists(self.market_mapping_path):
-                logger.info(f"Loading Market Mapping from Default Checkpoint {self.market_mapping_path}")
-                self.market_mapping = HeteroMapping.load(checkpoint=self.market_store_path)
-                # self._update_learnware_list(self.learnware_list)
-            else:
-                logger.warning(f"No Existing Market Mapping!!")
-                self.market_mapping = HeteroMapping()
-
-    def reset(self, market_id=None, auto_update=False, **kwargs):
+            logger.warning(f"No Existing Market Mapping!!")
+            self.market_mapping = HeteroMapping()
+
+    def reset(self, market_id=None, auto_update=False, auto_update_limit=None, **kwargs):
         # model training arguments(model architecture + optimization) set via self.reset
         self.auto_update = auto_update
         self.market_id = market_id
         self.training_args = kwargs
+        if auto_update_limit is not None: self.auto_update_limit = auto_update_limit
 
     def add_learnware(
         self, zip_path: str, semantic_spec: dict, check_status: int, learnware_id: str = None
@@ -103,22 +110,34 @@ class HeteroMapTableOrganizer(EasyOrganizer):
 
         learnwere_status = check_status if check_status is not None else BaseChecker.NONUSABLE_LEARNWARE
 
+        self.dbops.add_learnware(
+            id=learnware_id,
+            semantic_spec=semantic_spec,
+            zip_path=target_zip_dir,
+            folder_path=target_folder_dir,
+            use_flag=learnwere_status,
+        )
+
         self._update_learnware_list([new_learnware])
         self.learnware_list[learnware_id] = new_learnware
         self.learnware_zip_list[learnware_id] = target_zip_dir
         self.learnware_folder_list[learnware_id] = target_folder_dir
         self.use_flags[learnware_id] = learnwere_status
-        self.count += 1
-
-        if self.auto_update and self.count >= self.auto_update_limit:
-            train_process = multiprocessing.Process(target=self.train, args=(self.learnware_list.values(),))
-            train_process.start()
-            # train_process.join()
+        self.count += 1  
+
+        with self.auto_update_lock:
+            if self.auto_update and not self.is_training_in_progress.value and self.count - self.last_trained_learnware_num >= self.auto_update_limit:
+                self.is_training_in_progress.value = 1
+                curr_learnware_list = copy.deepcopy(self.learnware_list)
+                train_process = mp.Process(target=self.train, args=(curr_learnware_list.values(),))
+                train_process.start()
+                # train_process.join()
         
         return learnware_id, learnwere_status
 
     def train(self, learnware_list: List[Learnware] = None):
         learnware_list = learnware_list or self.learnware_list.values()
+        logger.warning(f"Leanwares for training: {[learnware.id for learnware in learnware_list]}")
         allset = self._learnwares_to_dataframes(learnware_list)
         self.market_mapping = HeteroMapping(**self.training_args)
         market_mapping_trainer = Trainer(
@@ -134,14 +153,23 @@ class HeteroMapTableOrganizer(EasyOrganizer):
 
         # essential hetero-mapping update for each market learnware when market model retrained
         self._update_learnware_list(learnware_list)
+        self.last_trained_learnware_num = self.count
+
+        logger.warning(f"Updataed Specification For: {[learnware.id for learnware in learnware_list]}")
+
+        with self.auto_update_lock:
+            self.is_training_in_progress.value = 0
 
     def _update_learnware_list(self, learnware_list: List[Learnware]):
-        hetero_mappings_save_path = os.path.join(self.market_store_path, "hetero_mappings")
-        os.makedirs(hetero_mappings_save_path, exist_ok=True)
-        for learnware in learnware_list:
-            learnware.id = learnware.id.replace(",", "_")
-            hetero_spec_path = os.path.join(hetero_mappings_save_path, f"{learnware.id}.npy")
-            self._update_learnware_specification(learnware, save_path=hetero_spec_path)
+        try:
+            hetero_mappings_save_path = os.path.join(self.market_store_path, "hetero_mappings")
+            os.makedirs(hetero_mappings_save_path, exist_ok=True)
+            for learnware in learnware_list:
+                hetero_spec_path = os.path.join(hetero_mappings_save_path, f"{learnware.id}.npy")
+                self._update_learnware_specification(learnware, save_path=hetero_spec_path)
+                logger.info(f"Learnware {learnware.id} HeteroSpecification Successfully Saved")
+        except Exception as err:
+            logger.warning(f"Update learnware HeteroSpecification failed! Due to {err}")
 
     def _update_learnware_specification(self, learnware: Learnware, save_path: str) -> Learnware:
         specification = learnware.specification

learnware/market/hetergeneous/organizer/hetero_mapping/__init__.py

@@ -60,6 +60,7 @@ class HeteroMapping(nn.Module):
             device=device,
         )
 
+        ##todo: BUG!!!!!!
         self.encoder = TransformerMultiLayer(
             hidden_dim=hidden_dim,
             num_layer=num_layer,
@@ -69,9 +70,6 @@ class HeteroMapping(nn.Module):
             activation=activation,
         )
         self.cls_token = CLSToken(hidden_dim=hidden_dim)
-        self.device = device
-        self.to(device)
-
         self.collate_fn = TransTabCollatorForCL(
             feature_tokenizer=feature_tokenizer, overlap_ratio=overlap_ratio, num_partition=num_partition
         )
@@ -103,7 +101,7 @@ class HeteroMapping(nn.Module):
         # load model weight state dict
         market_model_path = os.path.join(checkpoint, conf.market_model_path)
         model_info = torch.load(market_model_path, map_location="cpu")
-        model = HeteroMapping(feature_tokenizer=model_info["feature_tokenizer"], **model_info["model_args"])
+        model = HeteroMapping(**model_info["model_args"])
         model.load_state_dict(model_info["model_state_dict"], strict=False)
         return model
         # self.feature_tokenizer.load(checkpoint)
@@ -126,7 +124,7 @@ class HeteroMapping(nn.Module):
         model_info = {
             "model_state_dict": self.state_dict(),
             "model_args": self.model_args,
-            "feature_tokenizer": self.feature_tokenizer,
+            # "feature_tokenizer": self.feature_tokenizer,
         }
         torch.save(model_info, os.path.join(ckpt_dir, conf.market_model_path))
 
@@ -407,6 +405,7 @@ class TransformerMultiLayer(nn.Module):
                 use_layer_norm=True,
                 activation=activation,
             )
+            ##todo: BUG!!!!!!
             stacked_transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layer - 1)
             self.transformer_encoder.append(stacked_transformer)
 

learnware/market/hetergeneous/organizer/hetero_mapping/feature_extractor.py

@@ -68,7 +68,7 @@ class FeatureTokenizer:
 
     def __init__(
         self,
-        disable_tokenizer_parallel=False,
+        disable_tokenizer_parallel=True,
         **kwargs,
     ) -> None:
         """args:

learnware/market/hetergeneous/organizer/hetero_mapping/trainer.py

@@ -56,7 +56,7 @@ class Trainer:
         self.args["steps_per_epoch"] = int(self.args["num_training_steps"] / (num_epoch * len(self.train_set_list)))
         self.optimizer = None
 
-    def train(self):
+    def train(self, verbose: bool = True):
         self._create_optimizer()
         start_time = time.time()
         final_train_loss = 0
@@ -72,18 +72,17 @@ class Trainer:
                     train_loss_all += loss.item()
                     ite += 1
 
-            logger.info(
-                "epoch: {}, train loss: {:.4f}, lr: {:.6f}, spent: {:.1f} secs".format(
-                    epoch,
-                    train_loss_all,
-                    self.optimizer.param_groups[0]["lr"],
-                    time.time() - start_time,
+            if verbose:
+                logger.info(
+                    "epoch: {}, train loss: {:.4f}, lr: {:.6f}, spent: {:.1f} secs".format(
+                        epoch,
+                        train_loss_all,
+                        self.optimizer.param_groups[0]["lr"],
+                        time.time() - start_time,
+                    )
                 )
-            )
             final_train_loss = train_loss_all
 
-        # self.save_model(self.output_dir)
-
         logger.info("training complete, cost {:.1f} secs.".format(time.time() - start_time))
         return final_train_loss