replaced is_ssd() check with our own version to eliminate a dependency

dbms/postgres/cli.py

@@ -10,11 +10,10 @@
 import subprocess
 from pathlib import Path
 import click
-import ssd_checker
 
 from benchmark.tpch.load_info import TpchLoadInfo
 from dbms.load_info_base_class import LoadInfoBaseClass
-from misc.utils import DBGymConfig, conv_inputpath_to_realabspath, link_result, open_and_save, save_file, get_pgdata_tgz_name, default_pgbin_path, WORKSPACE_PATH_PLACEHOLDER, default_pgdata_parent_dpath
+from misc.utils import DBGymConfig, conv_inputpath_to_realabspath, link_result, open_and_save, save_file, get_pgdata_tgz_name, default_pgbin_path, WORKSPACE_PATH_PLACEHOLDER, default_pgdata_parent_dpath, is_ssd
 from util.shell import subprocess_run
 from sqlalchemy import Connection
 from util.pg import SHARED_PRELOAD_LIBRARIES, conn_execute, sql_file_execute, DBGYM_POSTGRES_DBNAME, create_conn, DEFAULT_POSTGRES_PORT, DBGYM_POSTGRES_USER, DBGYM_POSTGRES_PASS, DEFAULT_POSTGRES_DBNAME
@@ -73,9 +72,9 @@ def postgres_pgdata(dbgym_cfg: DBGymConfig, benchmark_name: str, scale_factor: f
 
  # Check assertions on args
  if intended_pgdata_hardware == "hdd":
- assert not ssd_checker.is_ssd(pgdata_parent_dpath), f"Intended hardware is HDD but pgdata_parent_dpath ({pgdata_parent_dpath}) is an SSD"
+ assert not is_ssd(pgdata_parent_dpath), f"Intended hardware is HDD but pgdata_parent_dpath ({pgdata_parent_dpath}) is an SSD"
  elif intended_pgdata_hardware == "ssd":
- assert ssd_checker.is_ssd(pgdata_parent_dpath), f"Intended hardware is SSD but pgdata_parent_dpath ({pgdata_parent_dpath}) is an HDD"
+ assert is_ssd(pgdata_parent_dpath), f"Intended hardware is SSD but pgdata_parent_dpath ({pgdata_parent_dpath}) is an HDD"
  else:
  assert False
 

dependency/requirements.txt

@@ -122,5 +122,4 @@ virtualenv==20.25.0
 Werkzeug==3.0.1
 wrapt==1.14.1
 zipp==3.17.0
-ssd_checker==1.0.3
 redis==5.0.3

misc/utils.py

@@ -575,3 +575,19 @@ def make_redis_started(port: int):
  # When you start Redis in daemon mode, it won't let you know if it's started, so we ping again to check
  r = redis.Redis(port=port)
  r.ping()
+
+
+def is_ssd(path: Path) -> bool:
+ try:
+ device = subprocess.check_output(['df', path]).decode().split('\n')[1].split()[0]
+ device_basename = os.path.basename(device)
+ lsblk_output = subprocess.check_output(['lsblk', '-d', '-o', 'name,rota']).decode()
+ for line in lsblk_output.split('\n')[1:]:
+ parts = line.split()
+ if parts and parts[0] == device_basename:
+ is_ssd = int(parts[1]) == 0
+ return is_ssd
+ return False
+ except Exception as e:
+ print(f"An error occurred: {e}")
+ return False

task.py

@@ -6,7 +6,7 @@
 
 from benchmark.cli import benchmark_group
 from dbms.cli import dbms_group
-from misc.utils import DBGymConfig
+from misc.utils import DBGymConfig, is_ssd
 from tune.cli import tune_group
 
 # TODO(phw2): save commit, git diff, and run command

tune/protox/agent/hpo.py

@@ -13,7 +13,6 @@
 from typing import Any, Optional, Union
 import random
 import click
-import ssd_checker
 import ray
 from ray.tune import Trainable
 from ray.tune.schedulers import FIFOScheduler
@@ -23,7 +22,7 @@
 from ray.train import SyncConfig
 
 from tune.protox.agent.build_trial import build_trial
-from misc.utils import DEFAULT_BOOT_CONFIG_FPATH, DEFAULT_WORKLOAD_TIMEOUT, DBGymConfig, TuningMode, link_result, open_and_save, restart_ray, conv_inputpath_to_realabspath, default_pristine_pgdata_snapshot_path, default_workload_path, default_embedder_path, default_benchmark_config_path, default_benchbase_config_path, WORKSPACE_PATH_PLACEHOLDER, BENCHMARK_NAME_PLACEHOLDER, WORKLOAD_NAME_PLACEHOLDER, SCALE_FACTOR_PLACEHOLDER, DEFAULT_SYSKNOBS_PATH, default_pgbin_path, workload_name_fn, default_pgdata_parent_dpath, default_hpoed_agent_params_fname
+from misc.utils import DEFAULT_BOOT_CONFIG_FPATH, DEFAULT_WORKLOAD_TIMEOUT, DBGymConfig, TuningMode, link_result, open_and_save, restart_ray, conv_inputpath_to_realabspath, default_pristine_pgdata_snapshot_path, default_workload_path, default_embedder_path, default_benchmark_config_path, default_benchbase_config_path, WORKSPACE_PATH_PLACEHOLDER, BENCHMARK_NAME_PLACEHOLDER, WORKLOAD_NAME_PLACEHOLDER, SCALE_FACTOR_PLACEHOLDER, DEFAULT_SYSKNOBS_PATH, default_pgbin_path, workload_name_fn, default_pgdata_parent_dpath, default_hpoed_agent_params_fname, is_ssd
 
 
 METRIC_NAME = "Best Metric"
@@ -247,9 +246,9 @@ def hpo(
 
  # Check assertions on args
  if intended_pgdata_hardware == "hdd":
- assert not ssd_checker.is_ssd(pgdata_parent_dpath), f"Intended hardware is HDD but pgdata_parent_dpath ({pgdata_parent_dpath}) is an SSD"
+ assert not is_ssd(pgdata_parent_dpath), f"Intended hardware is HDD but pgdata_parent_dpath ({pgdata_parent_dpath}) is an SSD"
  elif intended_pgdata_hardware == "ssd":
- assert ssd_checker.is_ssd(pgdata_parent_dpath), f"Intended hardware is SSD but pgdata_parent_dpath ({pgdata_parent_dpath}) is an HDD"
+ assert is_ssd(pgdata_parent_dpath), f"Intended hardware is SSD but pgdata_parent_dpath ({pgdata_parent_dpath}) is an HDD"
  else:
  assert False
 

tune/protox/embedding/datagen.py

@@ -13,7 +13,6 @@
 import yaml
 from sklearn.preprocessing import quantile_transform
 import shutil
-import ssd_checker
 
 from misc.utils import (
  BENCHMARK_NAME_PLACEHOLDER,
@@ -32,6 +31,7 @@
  save_file,
  workload_name_fn,
  default_pgdata_parent_dpath,
+ is_ssd,
 )
 from tune.protox.embedding.loss import COST_COLUMNS
 from tune.protox.env.space.primitive_space.index_space import IndexSpace
@@ -211,9 +211,9 @@ def datagen(
 
  # Check assertions on args
  if intended_pgdata_hardware == "hdd":
- assert not ssd_checker.is_ssd(pgdata_parent_dpath), f"Intended hardware is HDD but pgdata_parent_dpath ({pgdata_parent_dpath}) is an SSD"
+ assert not is_ssd(pgdata_parent_dpath), f"Intended hardware is HDD but pgdata_parent_dpath ({pgdata_parent_dpath}) is an SSD"
  elif intended_pgdata_hardware == "ssd":
- assert ssd_checker.is_ssd(pgdata_parent_dpath), f"Intended hardware is SSD but pgdata_parent_dpath ({pgdata_parent_dpath}) is an HDD"
+ assert is_ssd(pgdata_parent_dpath), f"Intended hardware is SSD but pgdata_parent_dpath ({pgdata_parent_dpath}) is an HDD"
  else:
  assert False