Automated workspace pruning/cleaning (#35)

**Summary**: A command to find and remove unneeded files from the
workspace directory.

**Demo**:
Passing 25 unit tests with many edge cases which create
files/directories, call `clean_workspace()`, and then verify its
contents.
![Screenshot 2024-07-07 at 17 05
07](https://github.com/cmu-db/dbgym/assets/20631215/ed2edeae-062a-40ae-b38e-e2ad3718d6b3)

**Details**
* "Aggressive" mode removes all task_runs/\*/ directories that are not
directly pointed to by a symlink in symlinks/.
* "Safe" mode also keeps task_runs/\*/ directories which are indirectly
pointed to by a symlink. This can happen if a symlink points to a
task_runs/\*/ directory which has a symlink in it that points to another
task_runs/*/ directory.
* I chose to write so many unit tests because this operation must be
bug-free.

.gitignore

@@ -1,6 +1,7 @@
 __pycache__/
 .conda/
 .idea/
+test_clean_scratchspace/
 
 workspace/
 default_*_benchbase_config_*.xml

README.md

@@ -9,7 +9,7 @@ These steps were tested on a fresh repository clone, Ubuntu 22.04.
 
 ```
 # Setup dependencies.
-# You may want to create a Python virtual environment (e.g. with conda) before doing this.
+# You may want to create a Python 3.10 virtual environment (e.g. with conda) before doing this.
 ./dependency/install_dependencies.sh
 
 # Compile a custom fork of PostgreSQL, load TPC-H (SF 0.01), train the Proto-X agent, and tune.

dependencies/rust.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

manage/cli.py

@@ -0,0 +1,217 @@
+import shutil
+from typing import List, Set
+import click
+import yaml
+import logging
+from pathlib import Path
+from misc.utils import DBGymConfig, is_child_path, parent_dpath_of_path
+from itertools import chain
+import os
+
+
+task_logger = logging.getLogger("task")
+task_logger.setLevel(logging.INFO)
+
+
+@click.group(name="manage")
+def manage_group():
+    pass
+
+
+@click.command(name="show")
+@click.argument("keys", nargs=-1)
+@click.pass_obj
+def manage_show(dbgym_cfg, keys):
+    config_path = dbgym_cfg.path
+    config_yaml = dbgym_cfg.yaml
+
+    # Traverse the YAML.
+    for key in keys:
+        config_yaml = config_yaml[key]
+
+    # Pretty-print the requested YAML value.
+    output_str = None
+    if type(config_yaml) != dict:
+        output_str = config_yaml
+    else:
+        output_str = yaml.dump(config_yaml, default_flow_style=False)
+        if len(keys) > 0:
+            output_str = "  " + output_str.replace("\n", "\n  ")
+        output_str = output_str.rstrip()
+    print(output_str)
+
+    task_logger.info(f"Read: {Path(config_path)}")
+
+
+@click.command(name="write")
+@click.argument("keys", nargs=-1)
+@click.argument("value_type")
+@click.argument("value")
+@click.pass_obj
+def manage_write(dbgym_cfg, keys, value_type, value):
+    config_path = dbgym_cfg.path
+    config_yaml = dbgym_cfg.yaml
+
+    # Traverse the YAML.
+    root_yaml = config_yaml
+    for key in keys[:-1]:
+        config_yaml = config_yaml[key]
+
+    # Modify the requested YAML value and write the YAML file.
+    assert type(config_yaml[keys[-1]]) != dict
+    config_yaml[keys[-1]] = getattr(__builtins__, value_type)(value)
+    new_yaml = yaml.dump(root_yaml, default_flow_style=False).rstrip()
+    Path(config_path).write_text(new_yaml)
+
+    task_logger.info(f"Updated: {Path(config_path)}")
+
+
+@click.command(name="standardize")
+@click.pass_obj
+def manage_standardize(dbgym_cfg):
+    config_path = dbgym_cfg.path
+    config_yaml = dbgym_cfg.yaml
+
+    # Write the YAML file.
+    new_yaml = yaml.dump(config_yaml, default_flow_style=False).rstrip()
+    Path(config_path).write_text(new_yaml)
+
+    task_logger.info(f"Updated: {Path(config_path)}")
+
+
+@click.command("clean")
+@click.pass_obj
+@click.option(
+    "--mode",
+    type=click.Choice(["safe", "aggressive"]),
+    default="safe",
+    help="The mode to clean the workspace (default=\"safe\"). \"aggressive\" means \"only keep run_*/ folders referenced by a file in symlinks/\". \"safe\" means \"in addition to that, recursively keep any run_*/ folders referenced by any symlinks in run_*/ folders we are keeping.\""
+)
+def manage_clean(dbgym_cfg: DBGymConfig, mode: str):
+    clean_workspace(dbgym_cfg, mode=mode, verbose=True)
+
+
+@click.command("count")
+@click.pass_obj
+def manage_count(dbgym_cfg: DBGymConfig):
+    num_files = _count_files_in_workspace(dbgym_cfg)
+    print(f"The workspace ({dbgym_cfg.dbgym_workspace_path}) has {num_files} total files/dirs/symlinks.")
+
+
+def add_symlinks_in_dpath(symlinks_stack: List[Path], root_dpath: Path, processed_symlinks: Set[Path]) -> None:
+    """
+    Will modify symlinks_stack and processed_symlinks.
+    """
+    for root_pathstr, dir_names, file_names in os.walk(root_dpath):
+        root_path = Path(root_pathstr)
+        # symlinks can either be files or directories, so we go through both dir_names and file_names
+        for file_name in chain(dir_names, file_names):
+            file_path = root_path / file_name
+            if file_path.is_symlink() and file_path not in processed_symlinks:
+                symlinks_stack.append(file_path)
+                processed_symlinks.add(file_path)
+
+
+def _count_files_in_workspace(dbgym_cfg: DBGymConfig) -> int:
+    """
+    Counts the number of files (regular file or dir or symlink) in the workspace.
+    """
+    total_count = 0
+    for dirpath, dirnames, filenames in os.walk(dbgym_cfg.dbgym_workspace_path, followlinks=False):
+        # Check if any of the directories are symbolic links and remove them from dirnames
+        dirnames[:] = [d for d in dirnames if not os.path.islink(os.path.join(dirpath, d))]
+
+        # Count files and directories (non-symlink directories already filtered)
+        total_count += len(filenames) + len(dirnames)
+
+    return total_count
+
+
+def clean_workspace(dbgym_cfg: DBGymConfig, mode: str="safe", verbose=False) -> None:
+    """
+    Clean all [workspace]/task_runs/run_*/ directories that are not referenced by any "active symlinks".
+    If mode is "aggressive", "active symlinks" means *only* the symlinks directly in [workspace]/symlinks/.
+    If mode is "safe", "active symlinks" means the symlinks directly in [workspace]/symlinks/ as well as
+      any symlinks referenced in task_runs/run_*/ directories we have already decided to keep.
+    """
+    # This stack holds the symlinks that are left to be processed
+    symlink_fpaths_to_process = []
+    # This set holds the symlinks that have already been processed to avoid infinite loops
+    processed_symlinks = set()
+
+    # 1. Initialize paths to process
+    if dbgym_cfg.dbgym_symlinks_path.exists():
+        add_symlinks_in_dpath(symlink_fpaths_to_process, dbgym_cfg.dbgym_symlinks_path, processed_symlinks)
+
+    # 2. Go through symlinks, figuring out which "children of task runs" to keep
+    # Based on the rules of the framework, "children of task runs" should be run_*/ directories.
+    # However, the user's workspace might happen to break these rules by putting directories not
+    #   named "run_*/" or files directly in task_runs/. Thus, I use the term "task_run_child_fordpaths"
+    #   instead of "run_dpaths".
+    task_run_child_fordpaths_to_keep = set()
+
+    if dbgym_cfg.dbgym_runs_path.exists():
+        while symlink_fpaths_to_process:
+            symlink_fpath: Path = symlink_fpaths_to_process.pop()
+            assert symlink_fpath.is_symlink()
+            # Path.resolve() resolves all layers of symlinks while os.readlink() only resolves one layer.
+            # However, os.readlink() literally reads the string contents of the link. We need to do some
+            #   processing on the result of os.readlink() to convert it to an absolute path
+            real_fordpath = symlink_fpath.resolve()
+            one_layer_resolved_fordpath = os.readlink(symlink_fpath)
+            assert str(real_fordpath) == str(os.readlink(symlink_fpath)), f"symlink_fpath ({symlink_fpath}) seems to point to *another* symlink. This is difficult to handle, so it is currently disallowed. Please resolve this situation manually."
+
+            # If the file doesn't exist, we'll just ignore it.
+            if not real_fordpath.exists():
+                continue
+            # We're only trying to figure out which direct children of task_runs/ to save. If the file isn't
+            #   even a descendant, we don't care about it.
+            if not is_child_path(real_fordpath, dbgym_cfg.dbgym_runs_path):
+                continue
+
+            assert not os.path.samefile(real_fordpath, dbgym_cfg.dbgym_runs_path)
+
+            # Figure out the task_run_child_fordpath to put into task_run_child_fordpaths_to_keep
+            task_run_child_fordpath = None
+            if os.path.samefile(parent_dpath_of_path(real_fordpath), dbgym_cfg.dbgym_runs_path):
+                # While it's true that it shouldn't be possible to symlink to a directory directly in task_runs/,
+                #   we'll just not delete it if the user happens to have one like this. Even if the user messed up
+                #   the structure somehow, it's just a good idea not to delete it.
+                task_run_child_fordpath = real_fordpath
+            else:
+                # Technically, it's not allowed to symlink to any files not in task_runs/run_*/[codebase]/[organization]/.
+                #   However, as with above, we won't just nuke files if the workspace doesn't follow this rule for
+                #   some reason.
+                task_run_child_fordpath = real_fordpath
+                while not os.path.samefile(parent_dpath_of_path(task_run_child_fordpath), dbgym_cfg.dbgym_runs_path):
+                    task_run_child_fordpath = parent_dpath_of_path(task_run_child_fordpath)
+            assert task_run_child_fordpath != None
+            assert os.path.samefile(parent_dpath_of_path(task_run_child_fordpath), dbgym_cfg.dbgym_runs_path), f"task_run_child_fordpath ({task_run_child_fordpath}) is not a direct child of dbgym_cfg.dbgym_runs_path"
+            task_run_child_fordpaths_to_keep.add(task_run_child_fordpath)
+
+            # If on safe mode, add symlinks inside the task_run_child_fordpath to be processed
+            if mode == "safe":
+                add_symlinks_in_dpath(symlink_fpaths_to_process, task_run_child_fordpath, processed_symlinks)
+
+    # 3. Go through all children of task_runs/*, deleting any that we weren't told to keep
+    # It's true that symlinks might link outside of task_runs/*. We'll just not care about those
+    starting_num_files = _count_files_in_workspace(dbgym_cfg)
+    if dbgym_cfg.dbgym_runs_path.exists():
+        for child_fordpath in dbgym_cfg.dbgym_runs_path.iterdir():
+            if child_fordpath not in task_run_child_fordpaths_to_keep:
+                if child_fordpath.is_dir():
+                    shutil.rmtree(child_fordpath)
+                else:
+                    os.remove(child_fordpath)
+    ending_num_files = _count_files_in_workspace(dbgym_cfg)
+
+    if verbose:
+        task_logger.info(f"Removed {starting_num_files - ending_num_files} out of {starting_num_files} files")
+        task_logger.info(f"Workspace went from {starting_num_files - ending_num_files} to {starting_num_files}")
+
+
+manage_group.add_command(manage_show)
+manage_group.add_command(manage_write)
+manage_group.add_command(manage_standardize)
+manage_group.add_command(manage_clean)
+manage_group.add_command(manage_count)