[feat] Repurpose aim reindex command for index db recreation

CHANGELOG.md

@@ -1,6 +1,11 @@
 # Changelog
 
-## 3.19.4
+## 3.20.0 
+
+### Enhancements:
+- Repurpose aim reindex command for index db recreation (mihran113)
+
+### Fixes
 - Handle index db corruption and warn in UI (mihran113) 
 
 ## 3.19.3 Apr 17, 2024

aim/cli/cli.py

@@ -6,7 +6,6 @@
 from aim.cli.version import commands as version_commands
 from aim.cli.up import commands as up_commands
 from aim.cli.server import commands as server_commands
-from aim.cli.reindex import commands as reindex_commands
 from aim.cli.runs import commands as runs_commands
 from aim.cli.convert import commands as convert_commands
 from aim.cli.storage import commands as storage_commands
@@ -25,7 +24,6 @@ def cli_entry_point(verbose):
 cli_entry_point.add_command(version_commands.version, VERSION_NAME)
 cli_entry_point.add_command(up_commands.up, UP_NAME)
 cli_entry_point.add_command(server_commands.server, SERVER_NAME)
-cli_entry_point.add_command(reindex_commands.reindex, REINDEX_NAME)
 cli_entry_point.add_command(runs_commands.runs, RUNS_NAME)
 cli_entry_point.add_command(convert_commands.convert, CONVERT)
 cli_entry_point.add_command(storage_commands.storage, STORAGE)

aim/cli/configs.py

@@ -1,7 +1,6 @@
 INIT_NAME = 'init'
 VERSION_NAME = 'version'
 UP_NAME = 'up'
-REINDEX_NAME = 'reindex'
 SERVER_NAME = 'server'
 RUNS_NAME = 'runs'
 CONVERT = 'convert'

aim/cli/storage/commands.py

@@ -122,12 +122,22 @@ def prune(ctx):
 
 
 @storage.command('reindex')
-@click.option('--finalize-only', required=False, is_flag=True, default=False)
+@click.option('-y', '--yes', is_flag=True, help='Automatically confirm prompt')
 @click.pass_context
-def reindex(ctx, finalize_only):
-    """ Process runs left in 'in progress' state. """
-    from aim.utils.deprecation import deprecation_warning
+def reindex(ctx, yes):
+    """ Recreate index database from scratch. """
+    repo_path = ctx.obj['repo']
+    repo = Repo.from_path(repo_path)
+    click.secho(f'This command will forcefully recreate index db for Aim Repo \'{repo_path}\'.\n'
+                f'Please stop the Aim UI to not interfere with the procedure.')
+
+    if yes:
+        confirmed = True
+    else:
+        confirmed = click.confirm('Do you want to proceed?')
+    if not confirmed:
+        return
+
+    repo._recreate_index()
 
-    deprecation_warning(remove_version='3.16', msg='`aim storage reindex` is deprecated! '
-                                                   'Use `aim runs close` command instead.')
     return

aim/sdk/remote_repo_proxy.py

@@ -46,3 +46,6 @@ def _restore_run(self, hash_):
 
     def _close_run(self, hash_):
         return self._rpc_client.run_instruction(-1, self._handler, '_close_run', [hash_])
+
+    def _recreate_index(self):
+        return self._rpc_client.run_instruction(-1, self._handler, '_recreate_index', [])

aim/sdk/repo.py

@@ -974,3 +974,26 @@ def optimize_container(path, extra_options):
             optimize_container(seqs_db_path, extra_options={})
         if index_manager.run_needs_indexing(run_hash):
             index_manager.index(run_hash)
+
+    def _recreate_index(self):
+        from tqdm import tqdm
+        if self.is_remote_repo:
+            self._remote_repo_proxy._recreate_index()
+            return
+
+        from aim.sdk.index_manager import RepoIndexManager
+        index_manager = RepoIndexManager.get_index_manager(self)
+
+        # force delete the index db and the locks
+
+        index_lock_path = os.path.join(self.path, 'locks', 'index')
+        if os.path.exists(index_lock_path):
+            os.remove(index_lock_path)
+
+        index_db_path = os.path.join(self.path, 'meta', 'index')
+        shutil.rmtree(index_db_path, ignore_errors=True)
+
+        # recreate the index db
+        run_hashes = self._all_run_hashes()
+        for run_hash in tqdm(run_hashes, desc='Indexing runs', total=len(run_hashes)):
+            index_manager.index(run_hash)

aim/storage/rockscontainer.pyx

@@ -165,7 +165,8 @@ class RocksContainer(Container):
         for k, v in self.items():
             index[k] = v
 
-        self._progress_path.unlink()
+        if self._progress_path.exists():
+            self._progress_path.unlink()
         self._progress_path = None
 
     def close(self):