Update index.py to incrementally add and remove index entries of a repo

Author: sujik18

mlc/index.py

@@ -27,7 +27,6 @@ def __init__(self, repos_path, repos):
         """
         self.repos_path = repos_path
         self.repos = repos
-        #logger.info(repos)
 
         logger.debug(f"Repos path for Index: {self.repos_path}")
         self.index_files = {
@@ -140,7 +139,64 @@ def get_item_mtime(self,file):
         if t > latest:
             latest = t
         return latest
-    
+
+    def _index_single_repo(self, repo, repos_changed=False, current_item_keys=None):
+        repo_path = repo.path
+        if not os.path.isdir(repo_path):
+            return False
+
+        changed = False
+
+        for folder_type in ["script", "cache", "experiment"]:
+            folder_path = os.path.join(repo_path, folder_type)
+            if not os.path.isdir(folder_path):
+                continue
+
+            for automation_dir in os.listdir(folder_path):
+                automation_path = os.path.join(folder_path, automation_dir)
+                if not os.path.isdir(automation_path):
+                    continue
+
+                yaml_path = os.path.join(automation_path, "meta.yaml")
+                json_path = os.path.join(automation_path, "meta.json")
+
+                if os.path.isfile(yaml_path):
+                    config_path = yaml_path
+                elif os.path.isfile(json_path):
+                    config_path = json_path
+                else:
+                    #logger.debug(f"No config file found in {automation_path}, skipping")
+                    delete_flag = False
+                    if automation_dir in self.modified_times:
+                        del self.modified_times[automation_dir]
+                    if any(automation_dir in item["path"] for item in self.indices[folder_type]):
+                        logger.debug(f"Removed index entry (if it exists) for {folder_type} : {automation_dir}")
+                        delete_flag = True
+                        self._remove_index_entry(automation_path)
+                    if delete_flag:
+                        self._save_indices()
+                    continue
+                if current_item_keys is not None:
+                    current_item_keys.add(config_path)
+                mtime = self.get_item_mtime(config_path)
+                old = self.modified_times.get(config_path)
+                old_mtime = old["mtime"] if isinstance(old, dict) else old
+
+                # skip if unchanged
+                if old_mtime == mtime and repos_changed != 1:
+                    continue
+
+                self.modified_times[config_path] = {
+                    "mtime": mtime,
+                    "date_time": datetime.fromtimestamp(mtime).strftime("%Y-%m-%d %H:%M:%S")
+                }
+                
+                # meta file changed, so reindex
+                self._process_config_file(config_path, folder_type, automation_path, repo)
+                changed = True
+        
+        return changed
+
     def build_index(self):
         """
         Build shared indices for script, cache, and experiment folders across all repositories.
@@ -152,125 +208,27 @@ def build_index(self):
         # track all currently detected item paths
         current_item_keys = set()
         changed = False
-        repos_changed = False
-        
-        # load existing modified times
+        force_rebuild = False
+
+        # load modified times
         self.modified_times = self._load_modified_times()
 
+        # if missing index file, then force full rebuild
         index_json_path = os.path.join(self.repos_path, "index_script.json")
-
-        rebuild_index = False
-
-        #file does not exist, rebuild
         if not os.path.exists(index_json_path):
             logger.warning("index_script.json missing. Forcing full index rebuild...")
-            #logger.debug("Resetting modified_times...")
             self.modified_times = {}
-            self._save_modified_times()
-        #else:
-        #    logger.debug("index_script.json exists. Skipping forced rebuild.")
-
-        #check repos.json mtime
-        repos_json_path = os.path.join(self.repos_path, "repos.json")
-        repos_mtime = os.path.getmtime(repos_json_path)
-
-        key = f"{repos_json_path}"
-        old = self.modified_times.get(key)
-        repo_old_mtime = old["mtime"] if isinstance(old, dict) else old
-
-        #logger.debug(f"Current repos.json mtime: {repos_mtime}")
-        #logger.debug(f"Old repos.json mtime: {repo_old_mtime}")
-        current_item_keys.add(key)
-
-        # if changed, reset indexes
-        if repo_old_mtime is None or repo_old_mtime != repos_mtime:
-            logger.debug("repos.json modified. Clearing index ........")
-            # reset indices
-            self.indices = {key: [] for key in self.index_files.keys()}
-            # record repo mtime
-            self.modified_times[key] = {
-                "mtime": repos_mtime,
-                "date_time": datetime.fromtimestamp(repos_mtime).strftime("%Y-%m-%d %H:%M:%S")
-            }
-            # clear modified times except for repos.json
-            self.modified_times = {key: self.modified_times[key]}
-            self._save_indices()
-            self._save_modified_times()
-            repos_changed = True
-        #else:
-        #    logger.debug("Repos.json not modified")
+            self.indices = {k: [] for k in self.index_files.keys()}
+            force_rebuild = True
 
+        # index each repo
         for repo in self.repos:
-            repo_path = repo.path #os.path.join(self.repos_path, repo)
-            if not os.path.isdir(repo_path):
-                continue
-            #logger.debug(f"------------Checking repository: {repo_path}---------------")
-            # Filter for relevant directories in the repo
-            for folder_type in ["script", "cache", "experiment"]:
-                #logger.debug(f"Checking folder type: {folder_type}")
-                folder_path = os.path.join(repo_path, folder_type)
-                if not os.path.isdir(folder_path):
-                    continue
-
-                # Process each automation directory
-                for automation_dir in os.listdir(folder_path):
-                    # logger.debug(f"Checking automation directory: {automation_dir}")
-                    automation_path = os.path.join(folder_path, automation_dir)
-                    if not os.path.isdir(automation_path):
-                        #logger.debug(f"Skipping non-directory automation path: {automation_path}")
-                        continue
-                    
-                    yaml_path = os.path.join(automation_path, "meta.yaml")
-                    json_path = os.path.join(automation_path, "meta.json")
-
-                    if os.path.isfile(yaml_path):
-                        # logger.debug(f"Found YAML config file: {yaml_path}")
-                        config_path = yaml_path
-                    elif os.path.isfile(json_path):
-                        # logger.debug(f"Found JSON config file: {json_path}")
-                        config_path = json_path
-                    else:
-                        #logger.debug(f"No config file found in {automation_path}, skipping")
-                        delete_flag = False
-                        if automation_dir in self.modified_times:
-                            del self.modified_times[automation_dir]
-                        if any(automation_dir in item["path"] for item in self.indices[folder_type]):
-                            logger.debug(f"Removed index entry (if it exists) for {folder_type} : {automation_dir}")
-                            delete_flag = True
-                            self._remove_index_entry(automation_path)
-                        if delete_flag:
-                            self._save_indices()
-                        continue
-                    current_item_keys.add(config_path)
-                    mtime = self.get_item_mtime(config_path)
-
-                    old = self.modified_times.get(config_path)
-                    old_mtime = old["mtime"] if isinstance(old, dict) else old
-
-                    # skip if unchanged
-                    if old_mtime == mtime and repos_changed != 1:
-                        # logger.debug(f"No changes detected for {config_path}, skipping reindexing.")
-                        continue
-                    #if(old_mtime is None):
-                    #    logger.debug(f"New meta.yaml file detected: {config_path}. Adding to index.")
-                    
-                    # update mtime
-                    #logger.debug(f"{config_path} is modified, index getting updated")
-                    #if config_path not in self.modified_times:
-                    #    logger.debug(f"*************{config_path} not found in modified_times; creating new entry***************")
-
-                    self.modified_times[config_path] = {
-                        "mtime": mtime,
-                        "date_time": datetime.fromtimestamp(mtime).strftime("%Y-%m-%d %H:%M:%S")
-                    }
-                    #logger.debug(f"Modified time for {config_path} updated to {mtime}")
-                    changed = True
-                    # meta file changed, so reindex
-                    self._process_config_file(config_path, folder_type, automation_path, repo)
+            repo_changed = self._index_single_repo(repo, force_rebuild, current_item_keys)
+            if repo_changed:
+                changed = True
 
         # remove deleted scripts
-        old_keys = set(self.modified_times.keys())
-        deleted_keys = old_keys - current_item_keys
+        deleted_keys = set(self.modified_times) - current_item_keys
         for key in deleted_keys:
             logger.warning(f"Detected deleted item, removing entry from modified times: {key}")
             del self.modified_times[key]
@@ -281,13 +239,10 @@ def build_index(self):
         if deleted_keys:
             logger.debug(f"Deleted keys removed from modified times and indices: {deleted_keys}")
 
-        if changed:
+        if force_rebuild or changed:
             logger.debug("Changes detected, saving updated index and modified times.")
             self._save_modified_times()
             self._save_indices()
-            #logger.debug("**************Index updated (changes detected).*************************")
-        #else:
-            #logger.debug("**************Index unchanged (no changes detected).********************")
 
     def _remove_index_entry(self, key):
         logger.debug(f"Removing index entry for {key}")
@@ -379,3 +334,48 @@ def _save_indices(self):
                 #logger.debug(f"Shared index for {folder_type} saved to {output_file}.")
             except Exception as e:
                 logger.error(f"Error saving shared index for {folder_type}: {e}")
+
+
+    def add_repo(self, repo):
+        """
+        Incrementally index a newly registered repository.
+        """
+        changed = self._index_single_repo(repo, repos_changed=True)
+
+        if changed:
+            self._save_indices()
+            self._save_modified_times()
+
+
+    def remove_repo_from_index(self, repo_path):
+        """
+        Remove all index entries and modified times belonging to a repo.
+        Called when a repo is unregistered from repos.json.
+        """
+
+        logger.info(f"Removing repo from index: {repo_path}")
+        changed = False
+
+        # remove index entries
+        for folder_type in self.indices:
+            before = len(self.indices[folder_type])
+            self.indices[folder_type] = [
+                item for item in self.indices[folder_type]
+                if not item["path"].startswith(repo_path)
+            ]
+            if len(self.indices[folder_type]) != before:
+                changed = True
+
+        # remove modified times
+        keys_to_delete = [
+            k for k in self.modified_times
+            if k.startswith(repo_path)
+        ]
+
+        for k in keys_to_delete:
+            del self.modified_times[k]
+            changed = True
+
+        if changed:
+            self._save_indices()
+            self._save_modified_times()

mlc/repo_action.py

@@ -166,31 +166,19 @@ def register_repo(self, repo_path, repo_meta, ignore_on_conflict=False):
 
         with open(repos_file_path, 'w') as f:
             json.dump(repos_list, f, indent=2)
-
-        logger.debug("Forcing Index rebuild ...")
-        # reload repos list 
-        with open(repos_file_path, 'r') as f:
-            repos_list = json.load(f)
-
-        repos = []
-        for p in repos_list:
-            meta = {}
-            yaml_file = os.path.join(p, "meta.yaml")
-            json_file = os.path.join(p, "meta.json")
-
-            if os.path.exists(yaml_file):
-                meta = utils.read_yaml(yaml_file)
-            elif os.path.exists(json_file):
-                meta = utils.read_json(json_file)
-            else:
-                logger.info(f"No meta file found in {self.repos_path}")
-                continue
-
-            repos.append(Repo(path=p, meta=meta))
-
-        # rebuild index via constructor
-        Index(self.repos_path, repos)
-        logger.debug("repos.json and index file has been updated")
+            logger.info(f"Updated repos.json at {repos_file_path}")
+        
+        self.repos = self.load_repos_and_meta()
+        repo_obj = next(
+            (r for r in self.repos if r.path == repo_path),
+            None
+        )
+
+        if repo_obj:
+            index = Index(self.repos_path, self.repos)
+            index.add_repo(repo_obj)
+            logger.debug("Index file has been updated")
+        
         return {'return': 0}
 
     def unregister_repo(self, repo_path):
@@ -531,7 +519,8 @@ def rm(self, run_args):
     Action: rm
     ####################################################################################################################
 
-    The `rm` action removes a specified repository from MLCFlow, deleting both the repo folder and its registration.  
+    The `rm` action removes a specified repository from MLCFlow, deleting the repository folder, its index entries, 
+    and its registration.
     If there are any modified local changes, the user will be prompted for confirmation unless the `-f` flag is used  
     for force removal.
  
@@ -578,6 +567,8 @@ def rm(self, run_args):
         repos_file_path = os.path.join(self.repos_path, 'repos.json')
 
         force_remove = True if run_args.get('f') else False
+        index = Index(self.repos_path, self.repos)
+        index.remove_repo_from_index(repo_path)
 
         return rm_repo(repo_path, repos_file_path, force_remove)