From cf7b3918cb0f22c8954572404dba9cd5aa4ab31f Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Wed, 8 Apr 2026 18:08:58 -0400 Subject: [PATCH 01/17] added hpc improvements --- bedboss/cli.py | 111 ++++++++++++++++++ bedboss/refgenome_validator/main.py | 18 ++- .../encode_processing/encode_into_parts.py | 23 ++++ 3 files changed, 141 insertions(+), 11 deletions(-) create mode 100644 scripts/encode_processing/encode_into_parts.py diff --git a/bedboss/cli.py b/bedboss/cli.py index 402a538..d4b6d4e 100644 --- a/bedboss/cli.py +++ b/bedboss/cli.py @@ -233,6 +233,117 @@ def run_pep( pm.stop_pipeline() +@app.command( + name="run-pep-hpc", + help="Split a large PEP into N chunks and submit each as a SLURM job. Idempotent: re-run to resume failed/pending chunks.", +) +def run_pep_hpc( + pep: str = typer.Option(..., help="PEP file. Local path or PEPhub registry path."), + workdir: str = typer.Option( + ..., help="Working directory for chunks, sbatch files, manifest, and state." + ), + n_chunks: int = typer.Option(..., help="Number of chunks to split the PEP into."), + # forwarded run-pep options + outfolder: str = typer.Option( + ..., help="Path to the output folder (shared across chunks)." + ), + bedbase_config: str = typer.Option( + ..., + help="Path to the bedbase config file", + exists=True, + file_okay=True, + readable=True, + ), + create_bedset: bool = typer.Option(True, help="Create a new bedset"), + bedset_heavy: bool = typer.Option(False, help="Run heavy bedbuncher"), + rfg_config: str = typer.Option(None, help="Path to the rfg config file"), + check_qc: bool = typer.Option(True, help="Check the quality of the input file?"), + ensdb: str = typer.Option(None, help="Path to the EnsDb database file"), + just_db_commit: bool = typer.Option(False, help="Just commit to the database?"), + force_overwrite: bool = typer.Option( + False, help="Force overwrite the output files" + ), + update: bool = typer.Option(False, help="Update existing records"), + upload_qdrant: bool = typer.Option(True, help="Upload to Qdrant"), + upload_s3: bool = typer.Option(True, help="Upload to S3"), + upload_pephub: bool = typer.Option(True, help="Upload to PEPHub"), + no_fail: bool = typer.Option(False, help="Do not fail on error"), + license_id: str = typer.Option(DEFAULT_LICENSE, help="License ID"), + standardize_pep: bool = typer.Option(False, help="Standardize the PEP using bedMS"), + lite: bool = typer.Option(False, help="Run the pipeline in lite mode."), + rerun: bool = typer.Option(False, help="Rerun already processed samples"), + multi: bool = typer.Option(False, help="Run multiple samples"), + recover: bool = typer.Option(True, help="Recover from previous run"), + dirty: bool = typer.Option(False, help="Run without removing existing files"), + # SLURM options + slurm_template: str = typer.Option( + None, + help="Path to a custom sbatch template. See bedboss/bedboss_hpc.py for placeholders.", + ), + slurm_account: str = typer.Option("shefflab", help="SLURM --account"), + slurm_partition: str = typer.Option("standard", help="SLURM --partition"), + slurm_time: str = typer.Option("72:00:00", help="SLURM --time"), + slurm_mem: str = typer.Option("60000", help="SLURM --mem (MB)"), + slurm_cpus: int = typer.Option(4, help="SLURM --cpus-per-task"), + slurm_ntasks: int = typer.Option(2, help="SLURM --ntasks"), + dry_run: bool = typer.Option( + False, help="Split and write sbatch files but do not submit." + ), +): + from bedboss.bedboss_hpc import RunPepArgs, SlurmConfig + from bedboss.bedboss_hpc import run_pep_hpc as _run_pep_hpc + + run_pep_args = RunPepArgs( + outfolder=outfolder, + bedbase_config=bedbase_config, + create_bedset=create_bedset, + bedset_heavy=bedset_heavy, + rfg_config=rfg_config, + check_qc=check_qc, + ensdb=ensdb, + just_db_commit=just_db_commit, + force_overwrite=force_overwrite, + update=update, + upload_qdrant=upload_qdrant, + upload_s3=upload_s3, + upload_pephub=upload_pephub, + no_fail=no_fail, + license_id=license_id, + standardize_pep=standardize_pep, + lite=lite, + rerun=rerun, + multi=multi, + recover=recover, + dirty=dirty, + ) + slurm_cfg = SlurmConfig( + account=slurm_account, + partition=slurm_partition, + time=slurm_time, + mem=slurm_mem, + cpus_per_task=slurm_cpus, + ntasks=slurm_ntasks, + ) + _run_pep_hpc( + pep=pep, + workdir=workdir, + n_chunks=n_chunks, + run_pep_args=run_pep_args, + slurm_cfg=slurm_cfg, + slurm_template=slurm_template, + dry_run=dry_run, + ) + + +@app.command(name="run-pep-hpc-status", help="Show status of a run-pep-hpc workdir.") +def run_pep_hpc_status( + workdir: str = typer.Option(..., help="Working directory created by run-pep-hpc."), +): + from bedboss.bedboss_hpc import run_pep_hpc_status as _status + + _status(workdir) + + @app.command( help="Run unprocessed files or reprocess them. Currently, only hg38, hg19, and mm10 genomes are supported." ) diff --git a/bedboss/refgenome_validator/main.py b/bedboss/refgenome_validator/main.py index 8b82cfe..71002e6 100644 --- a/bedboss/refgenome_validator/main.py +++ b/bedboss/refgenome_validator/main.py @@ -276,10 +276,8 @@ def determine_compatibility( for genome_model in self.genome_models: # First and Second Layer of Compatibility - model_compat_stats[ - genome_model.genome_digest - ]: CompatibilityStats = self.calculate_chrom_stats( - bed_chrom_info, genome_model.chrom_sizes + model_compat_stats[genome_model.genome_digest]: CompatibilityStats = ( + self.calculate_chrom_stats(bed_chrom_info, genome_model.chrom_sizes) ) # Third layer - IGD, only if layer 1 and layer 2 have passed @@ -291,15 +289,13 @@ def determine_compatibility( genome_model.genome_digest ].chrom_length_stats.beyond_range ): - model_compat_stats[ - genome_model.genome_digest - ].igd_stats = self.get_igd_overlaps(bedfile) + model_compat_stats[genome_model.genome_digest].igd_stats = ( + self.get_igd_overlaps(bedfile) + ) # Calculate compatibility rating - model_compat_stats[ - genome_model.genome_digest - ].compatibility = self.calculate_rating( - model_compat_stats[genome_model.genome_digest] + model_compat_stats[genome_model.genome_digest].compatibility = ( + self.calculate_rating(model_compat_stats[genome_model.genome_digest]) ) if concise: concise_dict = {} diff --git a/scripts/encode_processing/encode_into_parts.py b/scripts/encode_processing/encode_into_parts.py new file mode 100644 index 0000000..32176cb --- /dev/null +++ b/scripts/encode_processing/encode_into_parts.py @@ -0,0 +1,23 @@ +import polars as pl + + +def divide_csv(file_path): + df = pl.read_csv(file_path) + + df = df.with_row_index("index") + + n_parts = 15 + chunk_size = len(df) // n_parts + + for i in range(n_parts): + part = df.filter( + (pl.col("index") >= i * chunk_size) + & (pl.col("index") < (i + 1) * chunk_size) + ) + part.write_csv(f"part_{i}.csv") + + +if __name__ == "__main__": + divide_csv( + "/home/bnt4me/virginia/repos/bedboss/scripts/encode_processing/encode_bed_hg38.csv" + ) From aaa8c5678219ca85585f25e423d5ecf358c5bc2a Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Wed, 8 Apr 2026 18:24:13 -0400 Subject: [PATCH 02/17] added hpc module --- bedboss/bedboss_hpc.py | 715 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 715 insertions(+) create mode 100644 bedboss/bedboss_hpc.py diff --git a/bedboss/bedboss_hpc.py b/bedboss/bedboss_hpc.py new file mode 100644 index 0000000..d6c1385 --- /dev/null +++ b/bedboss/bedboss_hpc.py @@ -0,0 +1,715 @@ +""" +HPC orchestration for `bedboss run-pep`. + +Splits a large PEP into N chunks and submits each as its own SLURM job. +Idempotent: re-running picks up where it left off via per-chunk sentinel files +and `squeue` checks. + +See `hpc_command_plan.md` for the full design. +""" + +from __future__ import annotations + +import logging +import shutil +import subprocess +from datetime import datetime, timezone +from pathlib import Path +from typing import Optional + +import pandas as pd +import yaml +from pephubclient import PEPHubClient +from pephubclient.helpers import is_registry_path +from pydantic import BaseModel, Field + +_LOGGER = logging.getLogger(__name__) + +MANIFEST_NAME = "manifest.json" +SOURCE_PEP_DIR = "source_pep" +CHUNKS_DIR = "chunks" +STATE_DIR = "state" + +DEFAULT_TEMPLATE = """\ +#!/bin/bash + +#SBATCH --account={account} +#SBATCH --ntasks={ntasks} +#SBATCH --cpus-per-task={cpus_per_task} +#SBATCH --mem={mem} +#SBATCH --partition={partition} +#SBATCH --time={time} +#SBATCH --job-name=bedboss-{chunk_id} +#SBATCH -o {logs_dir}/{chunk_id}.out +#SBATCH -e {logs_dir}/{chunk_id}.err + +echo "Hello $USER, this is node $(hostname). Running {chunk_id}." + +bedboss run-pep \\ + --pep {chunk_pep_path} \\ + --outfolder {outfolder} \\ + --bedbase-config {bedbase_config} \\ + {forwarded_flags} +status=$? + +if [ $status -eq 0 ]; then + touch {state_dir}/{chunk_id}.done +else + touch {state_dir}/{chunk_id}.failed +fi +exit $status +""" + + +# --------------------------------------------------------------------------- +# models +# --------------------------------------------------------------------------- + + +class BoolFlagSpec(BaseModel): + """CLI representation of a boolean run-pep option. + + Attributes: + on: Flag emitted when the value is True (e.g. ``--upload-s3``). + off: Flag emitted when False, or None if there is no negative form. + """ + + on: str + off: Optional[str] = None + + +# Maps RunPepArgs field name -> its on/off CLI form. Anything not listed here +# is rendered as ``--key value`` regardless of type. +BOOL_FLAGS: dict[str, BoolFlagSpec] = { + "create_bedset": BoolFlagSpec(on="--create-bedset", off="--no-create-bedset"), + "bedset_heavy": BoolFlagSpec(on="--bedset-heavy"), + "check_qc": BoolFlagSpec(on="--check-qc", off="--no-check-qc"), + "just_db_commit": BoolFlagSpec(on="--just-db-commit"), + "force_overwrite": BoolFlagSpec(on="--force-overwrite"), + "update": BoolFlagSpec(on="--update"), + "upload_qdrant": BoolFlagSpec(on="--upload-qdrant", off="--no-upload-qdrant"), + "upload_s3": BoolFlagSpec(on="--upload-s3", off="--no-upload-s3"), + "upload_pephub": BoolFlagSpec(on="--upload-pephub", off="--no-upload-pephub"), + "no_fail": BoolFlagSpec(on="--no-fail"), + "standardize_pep": BoolFlagSpec(on="--standardize-pep"), + "lite": BoolFlagSpec(on="--lite"), + "rerun": BoolFlagSpec(on="--rerun"), + "multi": BoolFlagSpec(on="--multi"), + "recover": BoolFlagSpec(on="--recover", off="--no-recover"), + "dirty": BoolFlagSpec(on="--dirty"), +} + + +class SlurmConfig(BaseModel): + """SLURM resource settings used to render every chunk's sbatch script.""" + + account: str + partition: str + time: str + mem: str + cpus_per_task: int + ntasks: int + template: Optional[str] = Field( + default=None, + description="Path to a custom sbatch template, or None for the bundled default.", + ) + + +class ChunkMeta(BaseModel): + """Per-chunk metadata persisted in the manifest.""" + + id: str + sample_range: tuple[int, int] + n_samples: int + pep_path: str + sbatch_path: str + logs_dir: str + job_id: Optional[str] = None + submitted_at: Optional[str] = None + + +class Manifest(BaseModel): + """Top-level manifest persisted as ``manifest.json`` in the workdir.""" + + created_at: str + source_pep: str + source_config: str + n_chunks: int + run_pep_args: "RunPepArgs" + slurm: SlurmConfig + chunks: list[ChunkMeta] + + +class RunPepArgs(BaseModel): + """All ``bedboss run-pep`` options forwarded into each chunk job. + + Captured at first invocation and persisted in the manifest so that resume + runs use the exact same arguments and never drift. + """ + + outfolder: str + bedbase_config: str + create_bedset: bool = True + bedset_heavy: bool = False + rfg_config: Optional[str] = None + check_qc: bool = True + ensdb: Optional[str] = None + just_db_commit: bool = False + force_overwrite: bool = False + update: bool = False + upload_qdrant: bool = True + upload_s3: bool = True + upload_pephub: bool = True + no_fail: bool = False + license_id: Optional[str] = None + standardize_pep: bool = False + lite: bool = False + rerun: bool = False + multi: bool = False + recover: bool = True + dirty: bool = False + + +# --------------------------------------------------------------------------- +# manifest helpers +# --------------------------------------------------------------------------- + + +def _manifest_path(workdir: Path) -> Path: + """Return the path to the manifest file inside a workdir.""" + return workdir / MANIFEST_NAME + + +def _load_manifest(workdir: Path) -> Manifest | None: + """Load the manifest from a workdir. + + Args: + workdir: Path to the run-pep-hpc working directory. + + Returns: + Parsed Manifest, or None if no manifest exists yet (first run). + """ + p = _manifest_path(workdir) + if not p.exists(): + return None + return Manifest.model_validate_json(p.read_text()) + + +def _save_manifest(workdir: Path, manifest: Manifest) -> None: + """Persist the manifest to disk as pretty-printed JSON. + + Args: + workdir: Path to the run-pep-hpc working directory. + manifest: Manifest to write. + """ + _manifest_path(workdir).write_text(manifest.model_dump_json(indent=2)) + + +# --------------------------------------------------------------------------- +# source PEP resolution +# --------------------------------------------------------------------------- + + +def _resolve_source_pep(pep: str, workdir: Path) -> Path: + """Materialize the source PEP locally inside the workdir. + + PEPhub registry paths are pulled with `PEPHubClient.pull`. Local paths + (file or directory) are copied so the source is self-contained and + re-runs do not depend on the original location. + + Args: + pep: PEPhub registry path (e.g. ``namespace/name:tag``) or local path. + workdir: Run-pep-hpc working directory. + + Returns: + Path to the project config yaml of the materialized source PEP. + + Raises: + FileNotFoundError: If a local path does not exist. + RuntimeError: If no project config yaml can be located after pulling. + """ + dest = workdir / SOURCE_PEP_DIR + dest.mkdir(parents=True, exist_ok=True) + + if is_registry_path(pep): + # PEPHubClient.pull writes into /// + existing_yaml = list(dest.rglob("*.yaml")) + if not existing_yaml: + _LOGGER.info(f"Pulling PEP {pep} from PEPhub into {dest}") + PEPHubClient().pull(pep, output=str(dest), force=True) + existing_yaml = list(dest.rglob("*_config.yaml")) or list(dest.rglob("*.yaml")) + if not existing_yaml: + raise RuntimeError(f"No project config yaml found after pulling {pep}") + return existing_yaml[0] + + src = Path(pep).expanduser().resolve() + if not src.exists(): + raise FileNotFoundError(pep) + + if src.is_dir(): + src_dir = src + cfgs = list(src_dir.glob("*_config.yaml")) or list(src_dir.glob("*.yaml")) + if not cfgs: + raise RuntimeError(f"No project config yaml found in {src_dir}") + src_cfg = cfgs[0] + else: + src_cfg = src + src_dir = src.parent + + target_dir = dest / src_dir.name + if not target_dir.exists(): + shutil.copytree(src_dir, target_dir) + return target_dir / src_cfg.name + + +# --------------------------------------------------------------------------- +# splitting +# --------------------------------------------------------------------------- + + +def _read_sample_table_path(config_path: Path) -> Path: + """Resolve the sample table path declared in a PEP project config. + + Args: + config_path: Path to the project_config.yaml. + + Returns: + Absolute path to the sample table CSV referenced by the config. + + Raises: + RuntimeError: If the config has no ``sample_table`` entry or declares + multiple sample tables (not yet supported). + FileNotFoundError: If the resolved sample table file does not exist. + """ + with open(config_path) as f: + cfg = yaml.safe_load(f) + st = cfg.get("sample_table") + if st is None: + raise RuntimeError(f"PEP config {config_path} has no `sample_table` entry") + if isinstance(st, list): + if len(st) != 1: + raise RuntimeError("Multiple sample_tables not supported") + st = st[0] + p = (config_path.parent / st).resolve() + if not p.exists(): + raise FileNotFoundError(f"Sample table not found: {p}") + return p + + +def _check_no_subsamples(config_path: Path) -> None: + """Reject PEPs that declare a subsample table. + + Subsample tables are not yet supported by run-pep-hpc because they would + need to be sliced in lock-step with the parent sample table. + + Args: + config_path: Path to the project_config.yaml. + + Raises: + RuntimeError: If the config declares a ``subsample_table``. + """ + with open(config_path) as f: + cfg = yaml.safe_load(f) + if cfg.get("subsample_table"): + raise RuntimeError( + "Subsample tables are not yet supported in run-pep-hpc. " + "Remove `subsample_table` from the PEP config or open an issue." + ) + + +def _split_pep(workdir: Path, source_cfg: Path, n_chunks: int) -> list[ChunkMeta]: + """Slice the source PEP sample table into N chunk PEPs on disk. + + For each chunk, creates ``chunks/chunk_XXXX/{pep,slurm,logs}/`` and writes + a copy of the source project config plus the sliced sample table. Chunk + sizes differ by at most one sample. + + Args: + workdir: Run-pep-hpc working directory. + source_cfg: Path to the source project_config.yaml. + n_chunks: Requested number of chunks. Capped at the total sample count. + + Returns: + List of chunk metadata dicts ready to be stored in the manifest. + + Raises: + RuntimeError: If the source PEP has zero samples. + ValueError: If ``n_chunks`` is less than 1. + """ + _check_no_subsamples(source_cfg) + sample_table = _read_sample_table_path(source_cfg) + df = pd.read_csv(sample_table) + n = len(df) + if n == 0: + raise RuntimeError("Source PEP has 0 samples") + if n_chunks < 1: + raise ValueError("--n-chunks must be >= 1") + n_chunks = min(n_chunks, n) + + with open(source_cfg) as f: + base_config = yaml.safe_load(f) or {} + base_name = base_config.get("name") or source_cfg.stem + config_name = source_cfg.name + sample_table_name = "sample_table.csv" + + # even split, sizes differ by at most 1 + base, extra = divmod(n, n_chunks) + chunks: list[ChunkMeta] = [] + start = 0 + + for i in range(n_chunks): + size = base + (1 if i < extra else 0) + end = start + size + chunk_id = f"chunk_{i:04d}" + chunk_root = workdir / CHUNKS_DIR / chunk_id + pep_dir = chunk_root / "pep" + slurm_dir = chunk_root / "slurm" + logs_dir = chunk_root / "logs" + for d in (pep_dir, slurm_dir, logs_dir): + d.mkdir(parents=True, exist_ok=True) + + chunk_config = dict(base_config) + chunk_config["name"] = f"{base_name}_{chunk_id}" + chunk_config["sample_table"] = sample_table_name + with open(pep_dir / config_name, "w") as f: + yaml.safe_dump(chunk_config, f, sort_keys=False) + df.iloc[start:end].to_csv(pep_dir / sample_table_name, index=False) + + chunks.append( + ChunkMeta( + id=chunk_id, + sample_range=(start, end), + n_samples=size, + pep_path=str(pep_dir / config_name), + sbatch_path=str(slurm_dir / f"{chunk_id}.sbatch"), + logs_dir=str(logs_dir), + ) + ) + start = end + + return chunks + + +# --------------------------------------------------------------------------- +# sbatch rendering +# --------------------------------------------------------------------------- + + +def _forwarded_flags(run_pep_args: RunPepArgs) -> str: + """Render run-pep CLI flags for the generated sbatch script. + + ``--outfolder`` and ``--bedbase-config`` are emitted explicitly by the + template and skipped here. Boolean options are mapped via ``BOOL_FLAGS`` + to their ``--flag`` / ``--no-flag`` forms; everything else is rendered as + ``--key value``. + + Args: + run_pep_args: Run-pep options captured at first invocation. + + Returns: + A backslash-and-newline-joined string ready to be interpolated into + the sbatch template. + """ + skip = {"outfolder", "bedbase_config"} + parts: list[str] = [] + for key, val in run_pep_args.model_dump().items(): + if key in skip or val is None: + continue + spec = BOOL_FLAGS.get(key) + if spec is not None: + if val: + parts.append(spec.on) + elif spec.off is not None: + parts.append(spec.off) + continue + cli_key = "--" + key.replace("_", "-") + parts.append(f"{cli_key} {val}") + return " \\\n ".join(parts) + + +def _render_sbatch( + chunk: ChunkMeta, + slurm_cfg: SlurmConfig, + run_pep_args: RunPepArgs, + state_dir: Path, + template: str, +) -> str: + """Render an sbatch script for a single chunk. + + Args: + chunk: Chunk metadata dict from the manifest. + slurm_cfg: SLURM resource settings (account, partition, time, etc.). + run_pep_args: run-pep options to forward. + state_dir: Directory where the chunk will write its done/failed sentinel. + template: Raw sbatch template string with format placeholders. + + Returns: + The fully-rendered sbatch script as a string. + """ + return template.format( + account=slurm_cfg.account, + ntasks=slurm_cfg.ntasks, + cpus_per_task=slurm_cfg.cpus_per_task, + mem=slurm_cfg.mem, + partition=slurm_cfg.partition, + time=slurm_cfg.time, + chunk_id=chunk.id, + logs_dir=chunk.logs_dir, + state_dir=str(state_dir), + chunk_pep_path=chunk.pep_path, + outfolder=run_pep_args.outfolder, + bedbase_config=run_pep_args.bedbase_config, + forwarded_flags=_forwarded_flags(run_pep_args), + ) + + +def _write_sbatch_files( + chunks: list[ChunkMeta], + slurm_cfg: SlurmConfig, + run_pep_args: RunPepArgs, + state_dir: Path, + template_path: str | None, +) -> None: + """Render and write the sbatch script for every chunk. + + Args: + chunks: All chunk metadata dicts from the manifest. + slurm_cfg: SLURM resource settings. + run_pep_args: run-pep options to forward into each script. + state_dir: Sentinel directory shared across all chunks. + template_path: Optional path to a user-supplied template file. If None, + the bundled ``DEFAULT_TEMPLATE`` is used. + """ + if template_path: + template = Path(template_path).read_text() + else: + template = DEFAULT_TEMPLATE + for chunk in chunks: + content = _render_sbatch(chunk, slurm_cfg, run_pep_args, state_dir, template) + Path(chunk.sbatch_path).write_text(content) + + +# --------------------------------------------------------------------------- +# slurm interaction +# --------------------------------------------------------------------------- + + +def _squeue_alive(job_id: str | None) -> bool: + """Check whether a SLURM job is still queued or running. + + Args: + job_id: SLURM job id returned by a previous ``sbatch``, or None. + + Returns: + True if ``squeue -j `` reports the job as live. False if the + job is missing, finished, or if ``squeue`` is not on PATH. + """ + if not job_id: + return False + try: + out = subprocess.run( + ["squeue", "-j", str(job_id), "-h"], + capture_output=True, + text=True, + check=False, + ) + except FileNotFoundError: + _LOGGER.warning("squeue not found on PATH; cannot check live job state") + return False + return bool(out.stdout.strip()) + + +def _sbatch_submit(sbatch_path: str) -> str: + """Submit an sbatch script and return the assigned SLURM job id. + + Args: + sbatch_path: Path to the sbatch script to submit. + + Returns: + The SLURM job id parsed from sbatch's "Submitted batch job N" output. + + Raises: + subprocess.CalledProcessError: If sbatch exits non-zero. + """ + out = subprocess.run( + ["sbatch", sbatch_path], + capture_output=True, + text=True, + check=True, + ) + # "Submitted batch job 1234567" + line = out.stdout.strip().splitlines()[-1] + return line.split()[-1] + + +def _chunk_status(chunk: ChunkMeta, state_dir: Path) -> str: + """Derive the current status of a chunk from sentinels and squeue. + + Status is derived (never stored) so it cannot go stale across runs. + Precedence: ``done`` > ``running`` > ``failed`` > ``pending``. + + Args: + chunk: Chunk metadata. + state_dir: Directory containing ``.done`` / ``.failed`` sentinels. + + Returns: + One of ``"done"``, ``"running"``, ``"failed"``, ``"pending"``. + """ + cid = chunk.id + if (state_dir / f"{cid}.done").exists(): + return "done" + if _squeue_alive(chunk.job_id): + return "running" + if (state_dir / f"{cid}.failed").exists(): + return "failed" + return "pending" + + +def _submit_pending(manifest: Manifest, workdir: Path) -> None: + """Submit every chunk that is not already done or live in the queue. + + Done chunks are skipped. Live (queued/running) chunks are left alone. + Failed and never-submitted chunks are (re-)submitted; any stale + ``.failed`` sentinel is removed first so the next run starts clean. + The manifest is rewritten with the new job ids. + + Args: + manifest: Manifest dict (mutated in place with new job ids). + workdir: Run-pep-hpc working directory. + """ + state_dir = workdir / STATE_DIR + state_dir.mkdir(parents=True, exist_ok=True) + submitted = 0 + skipped_done = 0 + skipped_running = 0 + for chunk in manifest.chunks: + status = _chunk_status(chunk, state_dir) + if status == "done": + skipped_done += 1 + continue + if status == "running": + skipped_running += 1 + continue + # pending or failed: clear stale failed sentinel and resubmit + failed_sentinel = state_dir / f"{chunk.id}.failed" + if failed_sentinel.exists(): + failed_sentinel.unlink() + job_id = _sbatch_submit(chunk.sbatch_path) + chunk.job_id = job_id + chunk.submitted_at = datetime.now(timezone.utc).isoformat() + submitted += 1 + _LOGGER.info(f"Submitted {chunk.id} as job {job_id}") + _save_manifest(workdir, manifest) + print( + f"Submission summary: submitted={submitted}, " + f"already_done={skipped_done}, still_running={skipped_running}" + ) + + +# --------------------------------------------------------------------------- +# entry points +# --------------------------------------------------------------------------- + + +def run_pep_hpc( + pep: str, + workdir: str, + n_chunks: int, + run_pep_args: RunPepArgs, + slurm_cfg: SlurmConfig, + slurm_template: str | None = None, + dry_run: bool = False, +) -> None: + """Split a PEP into N chunks and submit each as a SLURM job. + + On first invocation: pulls/copies the source PEP, slices its sample table + into N chunk PEPs, renders one sbatch script per chunk, writes the + manifest, and submits all chunks. + + On re-invocation against an existing workdir: skips splitting entirely + and only (re)submits chunks that are not done and not currently live in + the SLURM queue. Failed chunks are resubmitted; run-pep's own per-sample + tracking files in ``outfolder`` cause already-processed samples to be + skipped on retry. + + Args: + pep: Source PEP — PEPhub registry path or local path. + workdir: Working directory for chunks, sbatch files, manifest, state. + Reused across resume invocations. + n_chunks: Number of chunks to split the sample table into. + run_pep_args: Dict of run-pep options to forward into each chunk job. + Must include ``outfolder`` and ``bedbase_config``. + slurm_cfg: SLURM resource settings (account, partition, time, mem, + cpus_per_task, ntasks). + slurm_template: Optional path to a custom sbatch template. If None, + ``DEFAULT_TEMPLATE`` is used. + dry_run: If True, write all chunks and sbatch scripts but do not call + sbatch. + """ + wd = Path(workdir).expanduser().resolve() + wd.mkdir(parents=True, exist_ok=True) + (wd / STATE_DIR).mkdir(exist_ok=True) + + manifest = _load_manifest(wd) + if manifest is None: + _LOGGER.info(f"No manifest in {wd} — splitting PEP") + source_cfg = _resolve_source_pep(pep, wd) + chunks = _split_pep(wd, source_cfg, n_chunks) + slurm_cfg = slurm_cfg.model_copy(update={"template": slurm_template}) + manifest = Manifest( + created_at=datetime.now(timezone.utc).isoformat(), + source_pep=pep, + source_config=str(source_cfg), + n_chunks=len(chunks), + run_pep_args=run_pep_args, + slurm=slurm_cfg, + chunks=chunks, + ) + _write_sbatch_files( + chunks, slurm_cfg, run_pep_args, wd / STATE_DIR, slurm_template + ) + _save_manifest(wd, manifest) + print( + f"Created {len(chunks)} chunks in {wd} " + f"(sizes: {[c.n_samples for c in chunks]})" + ) + else: + _LOGGER.info(f"Resuming from existing manifest at {wd}") + + if dry_run: + print("Dry run: skipping sbatch submission") + return + + _submit_pending(manifest, wd) + + +def run_pep_hpc_status(workdir: str) -> None: + """Print a per-chunk status table and totals for a run-pep-hpc workdir. + + Args: + workdir: Working directory previously created by ``run_pep_hpc``. + + Raises: + RuntimeError: If no manifest is found in the workdir. + """ + wd = Path(workdir).expanduser().resolve() + manifest = _load_manifest(wd) + if manifest is None: + raise RuntimeError(f"No manifest found at {wd}") + state_dir = wd / STATE_DIR + counts = {"done": 0, "failed": 0, "running": 0, "pending": 0} + rows = [] + for chunk in manifest.chunks: + status = _chunk_status(chunk, state_dir) + counts[status] += 1 + rows.append((chunk.id, chunk.n_samples, status, chunk.job_id or "-")) + print(f"{'chunk_id':<14} {'samples':>8} {'status':<10} {'job_id':>12}") + print("-" * 48) + for cid, n, st, jid in rows: + print(f"{cid:<14} {n:>8} {st:<10} {jid:>12}") + print("-" * 48) + total = sum(counts.values()) + print( + f"Totals: done={counts['done']} failed={counts['failed']} " + f"running={counts['running']} pending={counts['pending']} (of {total})" + ) From a8ea8374ca5a02f8aa962ee69861ce3ba437bfb5 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Wed, 8 Apr 2026 18:26:07 -0400 Subject: [PATCH 03/17] updated version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 3b9b232..f303f51 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "bedboss" -version = "0.10.0" +version = "0.10.1" description = "Pipelines for genomic region file to produce bed files, and its statistics" readme = "README.md" license = "BSD-2-Clause" From ff76e9a1adf5057b154b890600c73fbf210e1188 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Wed, 8 Apr 2026 18:29:59 -0400 Subject: [PATCH 04/17] fixed recurstion problem --- bedboss/bedboss_hpc.py | 12 +++++++++++- pyproject.toml | 2 +- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/bedboss/bedboss_hpc.py b/bedboss/bedboss_hpc.py index d6c1385..6cfd279 100644 --- a/bedboss/bedboss_hpc.py +++ b/bedboss/bedboss_hpc.py @@ -258,7 +258,17 @@ def _resolve_source_pep(pep: str, workdir: Path) -> Path: target_dir = dest / src_dir.name if not target_dir.exists(): - shutil.copytree(src_dir, target_dir) + # Avoid infinite recursion when --workdir lives inside src_dir. + workdir_resolved = workdir.resolve() + + def _ignore(dirpath: str, names: list[str]) -> list[str]: + ignored = [] + for name in names: + if (Path(dirpath) / name).resolve() == workdir_resolved: + ignored.append(name) + return ignored + + shutil.copytree(src_dir, target_dir, ignore=_ignore) return target_dir / src_cfg.name diff --git a/pyproject.toml b/pyproject.toml index f303f51..6e0fce9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "bedboss" -version = "0.10.1" +version = "0.10.2" description = "Pipelines for genomic region file to produce bed files, and its statistics" readme = "README.md" license = "BSD-2-Clause" From be12724df6700616b03bf50607032d6305039633 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Wed, 8 Apr 2026 18:39:39 -0400 Subject: [PATCH 05/17] fixed dir problem 2 --- bedboss/bedboss_hpc.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/bedboss/bedboss_hpc.py b/bedboss/bedboss_hpc.py index 6cfd279..0c8901b 100644 --- a/bedboss/bedboss_hpc.py +++ b/bedboss/bedboss_hpc.py @@ -258,15 +258,12 @@ def _resolve_source_pep(pep: str, workdir: Path) -> Path: target_dir = dest / src_dir.name if not target_dir.exists(): - # Avoid infinite recursion when --workdir lives inside src_dir. - workdir_resolved = workdir.resolve() + # Avoid infinite recursion when the workdir (and therefore dest) lives + # inside src_dir: skip both `dest` and `workdir` entries while copying. + skip = {dest.resolve(), workdir.resolve()} def _ignore(dirpath: str, names: list[str]) -> list[str]: - ignored = [] - for name in names: - if (Path(dirpath) / name).resolve() == workdir_resolved: - ignored.append(name) - return ignored + return [n for n in names if (Path(dirpath) / n).resolve() in skip] shutil.copytree(src_dir, target_dir, ignore=_ignore) return target_dir / src_cfg.name From b069031c683c0eef25b781c6aa8a40f3e5479982 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Wed, 8 Apr 2026 18:55:42 -0400 Subject: [PATCH 06/17] small output folder fixes --- bedboss/bedboss_hpc.py | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/bedboss/bedboss_hpc.py b/bedboss/bedboss_hpc.py index 0c8901b..9c5dadc 100644 --- a/bedboss/bedboss_hpc.py +++ b/bedboss/bedboss_hpc.py @@ -464,7 +464,7 @@ def _render_sbatch( logs_dir=chunk.logs_dir, state_dir=str(state_dir), chunk_pep_path=chunk.pep_path, - outfolder=run_pep_args.outfolder, + outfolder=str(Path(run_pep_args.outfolder) / chunk.id), bedbase_config=run_pep_args.bedbase_config, forwarded_flags=_forwarded_flags(run_pep_args), ) diff --git a/pyproject.toml b/pyproject.toml index 6e0fce9..f303f51 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "bedboss" -version = "0.10.2" +version = "0.10.1" description = "Pipelines for genomic region file to produce bed files, and its statistics" readme = "README.md" license = "BSD-2-Clause" From bbe428f42e9b1bbdd3234bcc4e8af7c55e25127f Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Wed, 8 Apr 2026 19:36:58 -0400 Subject: [PATCH 07/17] Fixed bed classifier --- bedboss/bedclassifier/bedclassifier.py | 1 + 1 file changed, 1 insertion(+) diff --git a/bedboss/bedclassifier/bedclassifier.py b/bedboss/bedclassifier/bedclassifier.py index 2ab7208..f2e93c2 100644 --- a/bedboss/bedclassifier/bedclassifier.py +++ b/bedboss/bedclassifier/bedclassifier.py @@ -109,6 +109,7 @@ def _read_bed_file(filepath: str, skiprows: int = 0) -> pd.DataFrame | None: raise BedTypeException(reason="Input is not a string or dataframe.") df = df.dropna(axis=1) + df.columns = range(len(df.columns)) num_cols = len(df.columns) compliant_columns = 0 bed_format_named = DATA_FORMAT.UCSC_BED From de350f5cf5b09e13daea91cd58f743f2b806c2c5 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Wed, 8 Apr 2026 20:08:18 -0400 Subject: [PATCH 08/17] updated bedboss hpc status --- bedboss/bedboss_hpc.py | 70 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 64 insertions(+), 6 deletions(-) diff --git a/bedboss/bedboss_hpc.py b/bedboss/bedboss_hpc.py index 9c5dadc..dd29ffd 100644 --- a/bedboss/bedboss_hpc.py +++ b/bedboss/bedboss_hpc.py @@ -690,6 +690,48 @@ def run_pep_hpc( _submit_pending(manifest, wd) +def _count_log_lines(path: Path) -> int: + """Return the number of non-empty lines in a Skipper log file. + + Args: + path: Path to the log file. Missing files return 0. + + Returns: + Number of non-empty lines, or 0 if the file does not exist. + """ + if not path.exists(): + return 0 + with open(path) as f: + return sum(1 for line in f if line.strip()) + + +def _chunk_sample_counts(chunk: ChunkMeta, base_outfolder: str) -> tuple[int, int]: + """Read run-pep's Skipper logs for a chunk and return processed/failed counts. + + The chunk PEP config's ``name`` field is used to locate + ``//.log`` and ``_fail.log``. + + Args: + chunk: Chunk metadata. + base_outfolder: The shared base outfolder from ``run_pep_args``. + + Returns: + ``(processed, failed)``. Both default to 0 if logs are missing. + """ + try: + with open(chunk.pep_path) as f: + cfg = yaml.safe_load(f) or {} + pep_name = cfg.get("name") + except OSError: + return 0, 0 + if not pep_name: + return 0, 0 + chunk_outfolder = Path(base_outfolder) / chunk.id + processed = _count_log_lines(chunk_outfolder / f"{pep_name}.log") + failed = _count_log_lines(chunk_outfolder / f"{pep_name}_fail.log") + return processed, failed + + def run_pep_hpc_status(workdir: str) -> None: """Print a per-chunk status table and totals for a run-pep-hpc workdir. @@ -704,19 +746,35 @@ def run_pep_hpc_status(workdir: str) -> None: if manifest is None: raise RuntimeError(f"No manifest found at {wd}") state_dir = wd / STATE_DIR + base_outfolder = manifest.run_pep_args.outfolder counts = {"done": 0, "failed": 0, "running": 0, "pending": 0} + total_processed = 0 + total_failed = 0 rows = [] for chunk in manifest.chunks: status = _chunk_status(chunk, state_dir) counts[status] += 1 - rows.append((chunk.id, chunk.n_samples, status, chunk.job_id or "-")) - print(f"{'chunk_id':<14} {'samples':>8} {'status':<10} {'job_id':>12}") - print("-" * 48) - for cid, n, st, jid in rows: - print(f"{cid:<14} {n:>8} {st:<10} {jid:>12}") - print("-" * 48) + processed, failed = _chunk_sample_counts(chunk, base_outfolder) + total_processed += processed + total_failed += failed + rows.append( + (chunk.id, chunk.n_samples, status, chunk.job_id or "-", processed, failed) + ) + header = ( + f"{'chunk_id':<14} {'samples':>8} {'status':<10} {'job_id':>12} " + f"{'processed':>10} {'failed':>8}" + ) + print(header) + print("-" * len(header)) + for cid, n, st, jid, proc, fail in rows: + print(f"{cid:<14} {n:>8} {st:<10} {jid:>12} {proc:>10} {fail:>8}") + print("-" * len(header)) total = sum(counts.values()) print( f"Totals: done={counts['done']} failed={counts['failed']} " f"running={counts['running']} pending={counts['pending']} (of {total})" ) + print( + f"Samples: processed={total_processed} failed={total_failed} " + f"(of {sum(c.n_samples for c in manifest.chunks)})" + ) From 269848c94e188a9d1cf173d8beca764501662411 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Thu, 9 Apr 2026 11:38:26 -0400 Subject: [PATCH 09/17] added some comments --- bedboss/bedboss_hpc.py | 1 + 1 file changed, 1 insertion(+) diff --git a/bedboss/bedboss_hpc.py b/bedboss/bedboss_hpc.py index dd29ffd..3026598 100644 --- a/bedboss/bedboss_hpc.py +++ b/bedboss/bedboss_hpc.py @@ -1,5 +1,6 @@ """ HPC orchestration for `bedboss run-pep`. +! This file is fully generated by AI - Bugs can occure Splits a large PEP into N chunks and submits each as its own SLURM job. Idempotent: re-running picks up where it left off via per-chunk sentinel files From 84e33d0533f29ab1b20e7b13dc99d0964b405fcb Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Tue, 14 Apr 2026 00:57:58 -0400 Subject: [PATCH 10/17] added cli tool that allows to add thousands of files to qdrant --- bedboss/cli.py | 2 + bedboss/qdrant_index/__init__.py | 10 +- bedboss/qdrant_index/qdrant_cli.py | 232 ++++++++++ bedboss/qdrant_index/qdrant_hpc.py | 642 +++++++++++++++++++++++++++ bedboss/qdrant_index/qdrant_index.py | 6 +- bedboss/qdrant_index/upload.py | 238 ++++++++++ bedboss/qdrant_index/vectorize.py | 182 ++++++++ 7 files changed, 1308 insertions(+), 4 deletions(-) create mode 100644 bedboss/qdrant_index/qdrant_cli.py create mode 100644 bedboss/qdrant_index/qdrant_hpc.py create mode 100644 bedboss/qdrant_index/upload.py create mode 100644 bedboss/qdrant_index/vectorize.py diff --git a/bedboss/cli.py b/bedboss/cli.py index d4b6d4e..58956ff 100644 --- a/bedboss/cli.py +++ b/bedboss/cli.py @@ -7,6 +7,7 @@ from pephubclient.helpers import MessageHandler as printm from bedboss.bbuploader.cli import app_bbuploader +from bedboss.qdrant_index.qdrant_cli import qdrant_app # commented and made new const here, because it speeds up help function, # from bbconf.const import DEFAULT_LICENSE @@ -904,3 +905,4 @@ def common( app.add_typer(app_bbuploader, name="geo") +app.add_typer(qdrant_app, name="qdrant") diff --git a/bedboss/qdrant_index/__init__.py b/bedboss/qdrant_index/__init__.py index 5825fc2..2e36202 100644 --- a/bedboss/qdrant_index/__init__.py +++ b/bedboss/qdrant_index/__init__.py @@ -1,3 +1,9 @@ -from bedboss.qdrant_index.qdrant_index import add_to_qdrant - __all__ = ["add_to_qdrant"] + + +def __getattr__(name): + if name == "add_to_qdrant": + from bedboss.qdrant_index.qdrant_index import add_to_qdrant + + return add_to_qdrant + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/bedboss/qdrant_index/qdrant_cli.py b/bedboss/qdrant_index/qdrant_cli.py new file mode 100644 index 0000000..6a3ed76 --- /dev/null +++ b/bedboss/qdrant_index/qdrant_cli.py @@ -0,0 +1,232 @@ +""" +CLI subapp for qdrant reindexing on HPC. + +Registered as ``bedboss qdrant ``. +""" + +import typer + +qdrant_app = typer.Typer( + name="qdrant", + help="Qdrant reindexing commands for HPC.", + pretty_exceptions_short=False, + pretty_exceptions_show_locals=False, +) + +# --------------------------------------------------------------------------- +# region-based commands +# --------------------------------------------------------------------------- + + +@qdrant_app.command( + name="reindex-region-hpc", + help="Fetch hg38 beds from DB, split into chunks, and submit SLURM vectorization jobs. Idempotent: re-run to resume.", +) +def reindex_region_hpc_cmd( + bedbase_config: str = typer.Option( + ..., help="Path to the bedbase config file", exists=True + ), + workdir: str = typer.Option( + ..., help="Working directory for chunks, state, and manifest" + ), + n_chunks: int = typer.Option(..., help="Number of SLURM jobs to create"), + limit: int = typer.Option(None, help="Max number of bed files to process"), + only_unindexed: bool = typer.Option( + False, help="Only fetch beds where file_indexed=False" + ), + # SLURM + slurm_account: str = typer.Option("shefflab", help="SLURM --account"), + slurm_partition: str = typer.Option("standard", help="SLURM --partition"), + slurm_time: str = typer.Option("72:00:00", help="SLURM --time"), + slurm_mem: str = typer.Option("60000", help="SLURM --mem (MB)"), + slurm_cpus: int = typer.Option(4, help="SLURM --cpus-per-task"), + slurm_ntasks: int = typer.Option(2, help="SLURM --ntasks"), + dry_run: bool = typer.Option(False, help="Write chunks but do not submit"), +): + from bedboss.qdrant_index.qdrant_hpc import SlurmConfig + from bedboss.qdrant_index.qdrant_hpc import reindex_region_hpc as _run + + slurm_cfg = SlurmConfig( + account=slurm_account, + partition=slurm_partition, + time=slurm_time, + mem=slurm_mem, + cpus_per_task=slurm_cpus, + ntasks=slurm_ntasks, + ) + _run( + bedbase_config=bedbase_config, + workdir=workdir, + n_chunks=n_chunks, + slurm_cfg=slurm_cfg, + limit=limit, + only_unindexed=only_unindexed, + dry_run=dry_run, + ) + + +@qdrant_app.command( + name="reindex-region-upload", + help="Read parquet vectors from workdir and upload to qdrant (region/file-to-file collection).", +) +def reindex_region_upload_cmd( + bedbase_config: str = typer.Option( + ..., help="Path to the bedbase config file", exists=True + ), + workdir: str = typer.Option( + ..., help="Working directory created by reindex-region-hpc" + ), + batch: int = typer.Option(100, help="Qdrant upsert batch size"), +): + from bedboss.qdrant_index.upload import upload_region_vectors + + upload_region_vectors(config=bedbase_config, workdir=workdir, batch=batch) + + +@qdrant_app.command( + name="reindex-region-status", + help="Show per-chunk status for a region reindex workdir.", +) +def reindex_region_status_cmd( + workdir: str = typer.Option( + ..., help="Working directory created by reindex-region-hpc" + ), + verbose: bool = typer.Option( + False, help="Print full per-chunk table even for large runs" + ), +): + from bedboss.qdrant_index.qdrant_hpc import reindex_hpc_status + + reindex_hpc_status(workdir=workdir, verbose=verbose) + + +# --------------------------------------------------------------------------- +# hybrid commands +# --------------------------------------------------------------------------- + + +@qdrant_app.command( + name="reindex-hybrid-hpc", + help="Fetch all beds from DB, split into chunks, and submit SLURM vectorization jobs. Idempotent: re-run to resume.", +) +def reindex_hybrid_hpc_cmd( + bedbase_config: str = typer.Option( + ..., help="Path to the bedbase config file", exists=True + ), + workdir: str = typer.Option( + ..., help="Working directory for chunks, state, and manifest" + ), + n_chunks: int = typer.Option(..., help="Number of SLURM jobs to create"), + limit: int = typer.Option(None, help="Max number of bed files to process"), + only_unindexed: bool = typer.Option( + False, help="Only fetch beds where indexed=False" + ), + # SLURM + slurm_account: str = typer.Option("shefflab", help="SLURM --account"), + slurm_partition: str = typer.Option("standard", help="SLURM --partition"), + slurm_time: str = typer.Option("72:00:00", help="SLURM --time"), + slurm_mem: str = typer.Option("60000", help="SLURM --mem (MB)"), + slurm_cpus: int = typer.Option(4, help="SLURM --cpus-per-task"), + slurm_ntasks: int = typer.Option(2, help="SLURM --ntasks"), + dry_run: bool = typer.Option(False, help="Write chunks but do not submit"), +): + from bedboss.qdrant_index.qdrant_hpc import SlurmConfig + from bedboss.qdrant_index.qdrant_hpc import reindex_hybrid_hpc as _run + + slurm_cfg = SlurmConfig( + account=slurm_account, + partition=slurm_partition, + time=slurm_time, + mem=slurm_mem, + cpus_per_task=slurm_cpus, + ntasks=slurm_ntasks, + ) + _run( + bedbase_config=bedbase_config, + workdir=workdir, + n_chunks=n_chunks, + slurm_cfg=slurm_cfg, + limit=limit, + only_unindexed=only_unindexed, + dry_run=dry_run, + ) + + +@qdrant_app.command( + name="reindex-hybrid-upload", + help="Read parquet vectors from workdir and upload to qdrant (hybrid/semantic collection).", +) +def reindex_hybrid_upload_cmd( + bedbase_config: str = typer.Option( + ..., help="Path to the bedbase config file", exists=True + ), + workdir: str = typer.Option( + ..., help="Working directory created by reindex-hybrid-hpc" + ), + batch: int = typer.Option(1000, help="Qdrant upsert batch size"), +): + from bedboss.qdrant_index.upload import upload_hybrid_vectors + + upload_hybrid_vectors(config=bedbase_config, workdir=workdir, batch=batch) + + +@qdrant_app.command( + name="reindex-hybrid-status", + help="Show per-chunk status for a hybrid reindex workdir.", +) +def reindex_hybrid_status_cmd( + workdir: str = typer.Option( + ..., help="Working directory created by reindex-hybrid-hpc" + ), + verbose: bool = typer.Option( + False, help="Print full per-chunk table even for large runs" + ), +): + from bedboss.qdrant_index.qdrant_hpc import reindex_hpc_status + + reindex_hpc_status(workdir=workdir, verbose=verbose) + + +# --------------------------------------------------------------------------- +# vectorize commands (called by sbatch scripts, not typically by users) +# --------------------------------------------------------------------------- + + +@qdrant_app.command( + name="vectorize-region", + help="[HPC worker] Vectorize a chunk of bed files using Region2Vec. Writes parquet.", +) +def vectorize_region_cmd( + chunk_pep: str = typer.Option(..., help="Path to chunk CSV"), + output_parquet: str = typer.Option(..., help="Path to write output parquet"), + model_path: str = typer.Option(..., help="Region2Vec model path or HuggingFace ID"), +): + from bedboss.qdrant_index.vectorize import vectorize_region + + vectorize_region( + chunk_pep=chunk_pep, + output_parquet=output_parquet, + model_path=model_path, + ) + + +@qdrant_app.command( + name="vectorize-hybrid", + help="[HPC worker] Vectorize a chunk of bed metadata using dense+sparse text encoders. Writes parquet.", +) +def vectorize_hybrid_cmd( + chunk_pep: str = typer.Option(..., help="Path to chunk CSV"), + output_parquet: str = typer.Option(..., help="Path to write output parquet"), + model_path: str = typer.Option(..., help="Dense text encoder model path"), + sparse_model_path: str = typer.Option( + None, help="Sparse encoder model path (optional)" + ), +): + from bedboss.qdrant_index.vectorize import vectorize_hybrid + + vectorize_hybrid( + chunk_pep=chunk_pep, + output_parquet=output_parquet, + model_path=model_path, + sparse_model_path=sparse_model_path, + ) diff --git a/bedboss/qdrant_index/qdrant_hpc.py b/bedboss/qdrant_index/qdrant_hpc.py new file mode 100644 index 0000000..547b7ba --- /dev/null +++ b/bedboss/qdrant_index/qdrant_hpc.py @@ -0,0 +1,642 @@ +""" +HPC orchestration for qdrant reindexing. + +Fetches bed metadata from the database, splits into chunks, generates +SLURM sbatch scripts for parallel vectorization, and manages job state. + +Modeled after bedboss/bedboss_hpc.py. +""" + +from __future__ import annotations + +import logging +import subprocess +from datetime import datetime, timezone +from pathlib import Path +from typing import Optional + +import pandas as pd +from pydantic import BaseModel +from sqlalchemy import and_, select +from sqlalchemy.orm import Session + +_LOGGER = logging.getLogger(__name__) + +MANIFEST_NAME = "manifest.json" +CHUNKS_DIR = "chunks" +STATE_DIR = "state" + +# --------------------------------------------------------------------------- +# sbatch templates +# --------------------------------------------------------------------------- + +REGION_TEMPLATE = """\ +#!/bin/bash + +#SBATCH --account={account} +#SBATCH --ntasks={ntasks} +#SBATCH --cpus-per-task={cpus_per_task} +#SBATCH --mem={mem} +#SBATCH --partition={partition} +#SBATCH --time={time} +#SBATCH --job-name=qdrant-region-{chunk_id} +#SBATCH -o {logs_dir}/{chunk_id}.out +#SBATCH -e {logs_dir}/{chunk_id}.err + +echo "Hello $USER, node $(hostname). Running {chunk_id} (region vectorization)." + +bedboss qdrant vectorize-region \\ + --chunk-pep {chunk_pep_path} \\ + --output-parquet {output_parquet_path} \\ + --model-path {model_path} +status=$? + +if [ $status -eq 0 ]; then + touch {state_dir}/{chunk_id}.done +else + touch {state_dir}/{chunk_id}.failed +fi +exit $status +""" + +HYBRID_TEMPLATE = """\ +#!/bin/bash + +#SBATCH --account={account} +#SBATCH --ntasks={ntasks} +#SBATCH --cpus-per-task={cpus_per_task} +#SBATCH --mem={mem} +#SBATCH --partition={partition} +#SBATCH --time={time} +#SBATCH --job-name=qdrant-hybrid-{chunk_id} +#SBATCH -o {logs_dir}/{chunk_id}.out +#SBATCH -e {logs_dir}/{chunk_id}.err + +echo "Hello $USER, node $(hostname). Running {chunk_id} (hybrid vectorization)." + +bedboss qdrant vectorize-hybrid \\ + --chunk-pep {chunk_pep_path} \\ + --output-parquet {output_parquet_path} \\ + --model-path {model_path} \\ + {sparse_flag} +status=$? + +if [ $status -eq 0 ]; then + touch {state_dir}/{chunk_id}.done +else + touch {state_dir}/{chunk_id}.failed +fi +exit $status +""" + + +# --------------------------------------------------------------------------- +# models +# --------------------------------------------------------------------------- + + +class SlurmConfig(BaseModel): + account: str + partition: str + time: str + mem: str + cpus_per_task: int + ntasks: int + + +class ChunkMeta(BaseModel): + id: str + sample_range: tuple[int, int] + n_samples: int + pep_path: str + sbatch_path: str + logs_dir: str + output_parquet: str + job_id: Optional[str] = None + submitted_at: Optional[str] = None + + +class QdrantHpcManifest(BaseModel): + created_at: str + search_type: str # "region" or "hybrid" + bedbase_config: str + n_chunks: int + total_samples: int + model_path: str + sparse_model_path: Optional[str] = None + slurm: SlurmConfig + chunks: list[ChunkMeta] + + +# --------------------------------------------------------------------------- +# manifest helpers +# --------------------------------------------------------------------------- + + +def _manifest_path(workdir: Path) -> Path: + return workdir / MANIFEST_NAME + + +def _load_manifest(workdir: Path) -> QdrantHpcManifest | None: + p = _manifest_path(workdir) + if not p.exists(): + return None + return QdrantHpcManifest.model_validate_json(p.read_text()) + + +def _save_manifest(workdir: Path, manifest: QdrantHpcManifest) -> None: + _manifest_path(workdir).write_text(manifest.model_dump_json(indent=2)) + + +# --------------------------------------------------------------------------- +# fetch metadata from database +# --------------------------------------------------------------------------- + + +def _fetch_region_metadata( + config: str, + only_unindexed: bool, + limit: int | None, +) -> tuple[pd.DataFrame, str, str]: + """Fetch hg38 bed metadata for region-based vectorization. + + Returns: + (DataFrame, region2vec_model_path, bedbase_config_path) + """ + from bbconf.bbagent import BedBaseAgent + from bbconf.const import DEFAULT_QDRANT_GENOME_DIGESTS + from bbconf.db_utils import Bed, BedMetadata + + agent = BedBaseAgent(config=config) + model_path = agent.config.config.path.region2vec + + conditions = [Bed.genome_digest.in_(DEFAULT_QDRANT_GENOME_DIGESTS)] + if only_unindexed: + conditions.append(Bed.file_indexed.is_(False)) + + statement = ( + select(Bed) + .join(BedMetadata, Bed.id == BedMetadata.id) + .where(and_(*conditions)) + ) + if limit: + statement = statement.limit(limit) + + with Session(agent.config.db_engine.engine) as session: + results = session.scalars(statement).all() + rows = _records_to_rows(results) + + df = pd.DataFrame(rows) + _LOGGER.info(f"Fetched {len(df)} bed records for region reindexing") + return df, model_path, config + + +def _fetch_hybrid_metadata( + config: str, + only_unindexed: bool, + limit: int | None, +) -> tuple[pd.DataFrame, str, str | None, str]: + """Fetch all bed metadata for hybrid vectorization. + + Returns: + (DataFrame, dense_model_path, sparse_model_path, bedbase_config_path) + """ + from bbconf.bbagent import BedBaseAgent + from bbconf.db_utils import Bed, BedMetadata + + agent = BedBaseAgent(config=config) + model_path = agent.config.config.path.text2vec + sparse_model_path = agent.config.config.path.sparse_model + + conditions = [] + if only_unindexed: + conditions.append(Bed.indexed.is_(False)) + + statement = select(Bed).join(BedMetadata, Bed.id == BedMetadata.id) + if conditions: + statement = statement.where(and_(*conditions)) + if limit: + statement = statement.limit(limit) + + with Session(agent.config.db_engine.engine) as session: + results = session.scalars(statement).all() + rows = _records_to_rows(results) + + df = pd.DataFrame(rows) + _LOGGER.info(f"Fetched {len(df)} bed records for hybrid reindexing") + return df, model_path, sparse_model_path, config + + +def _records_to_rows(results) -> list[dict]: + """Convert SQLAlchemy Bed records to list of dicts for DataFrame.""" + rows = [] + for r in results: + rows.append( + { + "sample_name": r.id, + "name": r.name or "", + "description": r.description or "", + "genome_alias": r.genome_alias or "", + "genome_digest": r.genome_digest or "", + "cell_line": r.annotations.cell_line if r.annotations else "", + "cell_type": r.annotations.cell_type if r.annotations else "", + "tissue": r.annotations.tissue if r.annotations else "", + "target": r.annotations.target if r.annotations else "", + "treatment": r.annotations.treatment if r.annotations else "", + "assay": r.annotations.assay if r.annotations else "", + "species_name": r.annotations.species_name if r.annotations else "", + } + ) + return rows + + +# --------------------------------------------------------------------------- +# splitting +# --------------------------------------------------------------------------- + + +def _split_into_chunks( + workdir: Path, + df: pd.DataFrame, + n_chunks: int, +) -> list[ChunkMeta]: + """Split a DataFrame into N chunk CSVs on disk.""" + n = len(df) + if n == 0: + raise RuntimeError("No bed records to process") + n_chunks = min(n_chunks, n) + + base, extra = divmod(n, n_chunks) + chunks: list[ChunkMeta] = [] + start = 0 + + for i in range(n_chunks): + size = base + (1 if i < extra else 0) + end = start + size + chunk_id = f"chunk_{i:04d}" + chunk_root = workdir / CHUNKS_DIR / chunk_id + pep_dir = chunk_root / "pep" + slurm_dir = chunk_root / "slurm" + logs_dir = chunk_root / "logs" + output_dir = chunk_root / "output" + for d in (pep_dir, slurm_dir, logs_dir, output_dir): + d.mkdir(parents=True, exist_ok=True) + + chunk_csv = pep_dir / "sample_table.csv" + df.iloc[start:end].to_csv(chunk_csv, index=False) + + chunks.append( + ChunkMeta( + id=chunk_id, + sample_range=(start, end), + n_samples=size, + pep_path=str(chunk_csv), + sbatch_path=str(slurm_dir / f"{chunk_id}.sbatch"), + logs_dir=str(logs_dir), + output_parquet=str(output_dir / "vectors.parquet"), + ) + ) + start = end + + return chunks + + +# --------------------------------------------------------------------------- +# sbatch rendering + submission +# --------------------------------------------------------------------------- + + +def _write_region_sbatch_files( + chunks: list[ChunkMeta], + slurm_cfg: SlurmConfig, + model_path: str, + state_dir: Path, +) -> None: + for chunk in chunks: + content = REGION_TEMPLATE.format( + account=slurm_cfg.account, + ntasks=slurm_cfg.ntasks, + cpus_per_task=slurm_cfg.cpus_per_task, + mem=slurm_cfg.mem, + partition=slurm_cfg.partition, + time=slurm_cfg.time, + chunk_id=chunk.id, + logs_dir=chunk.logs_dir, + state_dir=str(state_dir), + chunk_pep_path=chunk.pep_path, + output_parquet_path=chunk.output_parquet, + model_path=model_path, + ) + Path(chunk.sbatch_path).write_text(content) + + +def _write_hybrid_sbatch_files( + chunks: list[ChunkMeta], + slurm_cfg: SlurmConfig, + model_path: str, + sparse_model_path: str | None, + state_dir: Path, +) -> None: + sparse_flag = ( + f"--sparse-model-path {sparse_model_path}" if sparse_model_path else "" + ) + for chunk in chunks: + content = HYBRID_TEMPLATE.format( + account=slurm_cfg.account, + ntasks=slurm_cfg.ntasks, + cpus_per_task=slurm_cfg.cpus_per_task, + mem=slurm_cfg.mem, + partition=slurm_cfg.partition, + time=slurm_cfg.time, + chunk_id=chunk.id, + logs_dir=chunk.logs_dir, + state_dir=str(state_dir), + chunk_pep_path=chunk.pep_path, + output_parquet_path=chunk.output_parquet, + model_path=model_path, + sparse_flag=sparse_flag, + ) + Path(chunk.sbatch_path).write_text(content) + + +def _sbatch_submit(sbatch_path: str) -> str: + out = subprocess.run( + ["sbatch", sbatch_path], + capture_output=True, + text=True, + check=True, + ) + line = out.stdout.strip().splitlines()[-1] + return line.split()[-1] + + +def _get_alive_job_ids(chunks: list[ChunkMeta]) -> set[str]: + """Batch-check which SLURM jobs are still alive with a single squeue call.""" + job_ids = [c.job_id for c in chunks if c.job_id] + if not job_ids: + return set() + try: + out = subprocess.run( + ["squeue", "-j", ",".join(job_ids), "-h", "-o", "%i"], + capture_output=True, + text=True, + check=False, + ) + except FileNotFoundError: + _LOGGER.warning("squeue not found on PATH") + return set() + return {line.strip() for line in out.stdout.strip().splitlines() if line.strip()} + + +def _chunk_status( + chunk: ChunkMeta, state_dir: Path, alive_jobs: set[str] +) -> str: + cid = chunk.id + if (state_dir / f"{cid}.done").exists(): + return "done" + if chunk.job_id and chunk.job_id in alive_jobs: + return "running" + if (state_dir / f"{cid}.failed").exists(): + return "failed" + return "pending" + + +def _submit_pending(manifest: QdrantHpcManifest, workdir: Path) -> None: + state_dir = workdir / STATE_DIR + state_dir.mkdir(parents=True, exist_ok=True) + alive_jobs = _get_alive_job_ids(manifest.chunks) + submitted = 0 + skipped_done = 0 + skipped_running = 0 + for chunk in manifest.chunks: + status = _chunk_status(chunk, state_dir, alive_jobs) + if status == "done": + skipped_done += 1 + continue + if status == "running": + skipped_running += 1 + continue + failed_sentinel = state_dir / f"{chunk.id}.failed" + if failed_sentinel.exists(): + failed_sentinel.unlink() + job_id = _sbatch_submit(chunk.sbatch_path) + chunk.job_id = job_id + chunk.submitted_at = datetime.now(timezone.utc).isoformat() + submitted += 1 + if submitted % 100 == 0: + print(f" submitted {submitted} jobs...") + _save_manifest(workdir, manifest) + print( + f"Submission summary: submitted={submitted}, " + f"already_done={skipped_done}, still_running={skipped_running}" + ) + + +# --------------------------------------------------------------------------- +# entry points +# --------------------------------------------------------------------------- + + +def reindex_region_hpc( + bedbase_config: str, + workdir: str, + n_chunks: int, + slurm_cfg: SlurmConfig, + limit: int | None = None, + only_unindexed: bool = False, + dry_run: bool = False, +) -> None: + """Fetch hg38 bed metadata, split, generate sbatch scripts, and submit.""" + wd = Path(workdir).expanduser().resolve() + wd.mkdir(parents=True, exist_ok=True) + (wd / STATE_DIR).mkdir(exist_ok=True) + + manifest = _load_manifest(wd) + if manifest is None: + _LOGGER.info("Fetching bed metadata from database...") + df, model_path, config_path = _fetch_region_metadata( + bedbase_config, only_unindexed, limit + ) + if df.empty: + print("No bed records found matching criteria. Nothing to do.") + return + + source_dir = wd / "source_pep" + source_dir.mkdir(exist_ok=True) + df.to_csv(source_dir / "sample_table.csv", index=False) + + chunks = _split_into_chunks(wd, df, n_chunks) + state_dir = wd / STATE_DIR + + manifest = QdrantHpcManifest( + created_at=datetime.now(timezone.utc).isoformat(), + search_type="region", + bedbase_config=config_path, + n_chunks=len(chunks), + total_samples=len(df), + model_path=model_path, + slurm=slurm_cfg, + chunks=chunks, + ) + _write_region_sbatch_files(chunks, slurm_cfg, model_path, state_dir) + _save_manifest(wd, manifest) + print( + f"Created {len(chunks)} chunks in {wd} " + f"({len(df)} total samples, sizes: {[c.n_samples for c in chunks]})" + ) + else: + _LOGGER.info(f"Resuming from existing manifest at {wd}") + + if dry_run: + print("Dry run: skipping sbatch submission") + return + + _submit_pending(manifest, wd) + + +def reindex_hybrid_hpc( + bedbase_config: str, + workdir: str, + n_chunks: int, + slurm_cfg: SlurmConfig, + limit: int | None = None, + only_unindexed: bool = False, + dry_run: bool = False, +) -> None: + """Fetch all bed metadata, split, generate sbatch scripts, and submit.""" + wd = Path(workdir).expanduser().resolve() + wd.mkdir(parents=True, exist_ok=True) + (wd / STATE_DIR).mkdir(exist_ok=True) + + manifest = _load_manifest(wd) + if manifest is None: + _LOGGER.info("Fetching bed metadata from database...") + df, model_path, sparse_model_path, config_path = _fetch_hybrid_metadata( + bedbase_config, only_unindexed, limit + ) + if df.empty: + print("No bed records found matching criteria. Nothing to do.") + return + + source_dir = wd / "source_pep" + source_dir.mkdir(exist_ok=True) + df.to_csv(source_dir / "sample_table.csv", index=False) + + chunks = _split_into_chunks(wd, df, n_chunks) + state_dir = wd / STATE_DIR + + manifest = QdrantHpcManifest( + created_at=datetime.now(timezone.utc).isoformat(), + search_type="hybrid", + bedbase_config=config_path, + n_chunks=len(chunks), + total_samples=len(df), + model_path=model_path, + sparse_model_path=sparse_model_path, + slurm=slurm_cfg, + chunks=chunks, + ) + _write_hybrid_sbatch_files( + chunks, slurm_cfg, model_path, sparse_model_path, state_dir + ) + _save_manifest(wd, manifest) + print( + f"Created {len(chunks)} chunks in {wd} " + f"({len(df)} total samples, sizes: {[c.n_samples for c in chunks]})" + ) + else: + _LOGGER.info(f"Resuming from existing manifest at {wd}") + + if dry_run: + print("Dry run: skipping sbatch submission") + return + + _submit_pending(manifest, wd) + + +# --------------------------------------------------------------------------- +# status +# --------------------------------------------------------------------------- + + +def _parquet_row_count(path: Path) -> int: + """Read parquet row count from file metadata without loading data.""" + if not path.exists(): + return 0 + try: + import pyarrow.parquet as pq + + return pq.read_metadata(path).num_rows + except Exception: + return 0 + + +def reindex_hpc_status(workdir: str, verbose: bool = False) -> None: + """Print per-chunk status table for a qdrant reindex workdir. + + Args: + workdir: Working directory created by reindex-*-hpc. + verbose: Print per-chunk table even when chunk count is large. + """ + wd = Path(workdir).expanduser().resolve() + manifest = _load_manifest(wd) + if manifest is None: + raise RuntimeError(f"No manifest found at {wd}") + + state_dir = wd / STATE_DIR + alive_jobs = _get_alive_job_ids(manifest.chunks) + counts = {"done": 0, "failed": 0, "running": 0, "pending": 0} + total_vectors = 0 + failed_chunks: list[str] = [] + rows = [] + + for chunk in manifest.chunks: + status = _chunk_status(chunk, state_dir, alive_jobs) + counts[status] += 1 + if status == "failed": + failed_chunks.append(chunk.id) + + n_vectors = _parquet_row_count(Path(chunk.output_parquet)) + total_vectors += n_vectors + + rows.append( + (chunk.id, chunk.n_samples, status, chunk.job_id or "-", n_vectors) + ) + + print(f"Search type: {manifest.search_type}") + print(f"Model: {manifest.model_path}") + print() + + # Print per-chunk table only when manageable or explicitly requested + show_table = verbose or len(rows) <= 200 + if show_table: + header = ( + f"{'chunk_id':<14} {'samples':>8} {'status':<10} " + f"{'job_id':>12} {'vectors':>10}" + ) + print(header) + print("-" * len(header)) + for cid, n, st, jid, nv in rows: + print(f"{cid:<14} {n:>8} {st:<10} {jid:>12} {nv:>10}") + print("-" * len(header)) + else: + print(f"({len(rows)} chunks — use --verbose to print full table)") + print() + + total = sum(counts.values()) + print( + f"Chunks: done={counts['done']} failed={counts['failed']} " + f"running={counts['running']} pending={counts['pending']} (of {total})" + ) + print( + f"Samples: {manifest.total_samples} total, " + f"{total_vectors} vectors produced" + ) + + if failed_chunks: + show_n = min(20, len(failed_chunks)) + print(f"\nFailed chunks ({len(failed_chunks)} total):") + for cid in failed_chunks[:show_n]: + print(f" {cid}") + if len(failed_chunks) > show_n: + print(f" ... and {len(failed_chunks) - show_n} more") diff --git a/bedboss/qdrant_index/qdrant_index.py b/bedboss/qdrant_index/qdrant_index.py index a0e58c1..6e4c106 100644 --- a/bedboss/qdrant_index/qdrant_index.py +++ b/bedboss/qdrant_index/qdrant_index.py @@ -1,7 +1,5 @@ import logging -from bbconf.bbagent import BedBaseAgent - _LOGGER = logging.getLogger("bedboss") @@ -14,6 +12,8 @@ def add_to_qdrant(config: str, batch: int = 100, purge: bool = False) -> None: batch: Number of items to upload in one batch. purge: Whether to purge the existing index before reindexing. """ + from bbconf.bbagent import BedBaseAgent + agent = BedBaseAgent(config=config) agent.bed.reindex_qdrant(batch=batch, purge=purge) @@ -29,6 +29,8 @@ def reindex_semantic_search( purge: Whether to purge the existing index before reindexing. batch: Number of items that will be uploaded to qdrant in one batch. """ + from bbconf.bbagent import BedBaseAgent + agent = BedBaseAgent(config=config) agent.bed.reindex_hybrid_search(batch=batch, purge=purge) _LOGGER.info("Semantic search reindexing completed.") diff --git a/bedboss/qdrant_index/upload.py b/bedboss/qdrant_index/upload.py new file mode 100644 index 0000000..0a5fc18 --- /dev/null +++ b/bedboss/qdrant_index/upload.py @@ -0,0 +1,238 @@ +""" +Upload pre-computed vectors from parquet files to Qdrant and update DB flags. + +This phase requires a bbagent connection (database + qdrant). +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import pandas as pd +from bbconf.bbagent import BedBaseAgent +from bbconf.db_utils import Bed +from bbconf.models.bed_models import VectorMetadata +from qdrant_client.http.models import PointStruct +from qdrant_client.models import SparseVector +from sqlalchemy.orm import Session + +_LOGGER = logging.getLogger(__name__) + + +def _load_parquet_files(workdir: Path) -> pd.DataFrame: + """Glob and concatenate all vectors.parquet files from chunk output dirs.""" + parquet_files = sorted(workdir.glob("chunks/*/output/vectors.parquet")) + if not parquet_files: + raise FileNotFoundError(f"No vectors.parquet files found in {workdir}/chunks/") + + dfs = [] + for pf in parquet_files: + df = pd.read_parquet(pf) + if len(df) > 0: + dfs.append(df) + _LOGGER.info(f"Loaded {len(df)} vectors from {pf}") + + if not dfs: + raise RuntimeError("All parquet files are empty — nothing to upload") + + combined = pd.concat(dfs, ignore_index=True) + _LOGGER.info(f"Total vectors to upload: {len(combined)}") + return combined + + +def upload_region_vectors( + config: str, + workdir: str, + batch: int = 100, +) -> None: + """Upload region-based vectors to qdrant and mark file_indexed=True. + + Args: + config: Path to the bedbase config file. + workdir: Working directory containing chunk output parquet files. + batch: Number of points to upload in one qdrant upsert call. + """ + wd = Path(workdir).expanduser().resolve() + df = _load_parquet_files(wd) + + agent = BedBaseAgent(config=config) + qd_client = agent.config.qdrant_file_backend.qd_client + collection = agent.config.config.qdrant.file_collection + + uploaded = 0 + points_batch: list[PointStruct] = [] + ids_batch: list[str] = [] + + with Session(agent.config.db_engine) as session: + for _, row in df.iterrows(): + bed_id = row["sample_name"] + vector = row["vector"] + if isinstance(vector, str): + import ast + + vector = ast.literal_eval(vector) + + metadata = VectorMetadata( + id=bed_id, + name=row.get("name") or "", + description=row.get("description") or "", + genome_alias=row.get("genome_alias") or "", + genome_digest=row.get("genome_digest"), + cell_line=row.get("cell_line") or "", + cell_type=row.get("cell_type") or "", + tissue=row.get("tissue") or "", + target=row.get("target") or "", + treatment=row.get("treatment") or "", + assay=row.get("assay") or "", + species_name=row.get("species_name") or "", + ) + + points_batch.append( + PointStruct( + id=bed_id, + vector=list(vector), + payload=metadata.model_dump(), + ) + ) + ids_batch.append(bed_id) + + if len(points_batch) >= batch: + _upsert_and_mark( + session, qd_client, collection, points_batch, ids_batch, "file_indexed" + ) + uploaded += len(points_batch) + _LOGGER.info(f"Uploaded {uploaded} points") + points_batch = [] + ids_batch = [] + + # Upload remaining + if points_batch: + _upsert_and_mark( + session, qd_client, collection, points_batch, ids_batch, "file_indexed" + ) + uploaded += len(points_batch) + + _LOGGER.info(f"Region upload complete: {uploaded} points uploaded to '{collection}'") + print(f"Upload complete: {uploaded} points to collection '{collection}'") + + +def upload_hybrid_vectors( + config: str, + workdir: str, + batch: int = 1000, +) -> None: + """Upload hybrid (dense+sparse) vectors to qdrant and mark indexed=True. + + Args: + config: Path to the bedbase config file. + workdir: Working directory containing chunk output parquet files. + batch: Number of points to upload in one qdrant upsert call. + """ + wd = Path(workdir).expanduser().resolve() + df = _load_parquet_files(wd) + + agent = BedBaseAgent(config=config) + qd_client = agent.config.qdrant_client + collection = agent.config.config.qdrant.hybrid_collection + + uploaded = 0 + points_batch: list[PointStruct] = [] + ids_batch: list[str] = [] + + with Session(agent.config.db_engine) as session: + for _, row in df.iterrows(): + bed_id = row["sample_name"] + dense_vec = row["dense_vector"] + if isinstance(dense_vec, str): + import ast + + dense_vec = ast.literal_eval(dense_vec) + + point_vectors = {"dense": list(dense_vec)} + + sparse_indices = row.get("sparse_indices") + sparse_values = row.get("sparse_values") + if sparse_indices is not None and sparse_values is not None: + if isinstance(sparse_indices, str): + import ast + + sparse_indices = ast.literal_eval(sparse_indices) + sparse_values = ast.literal_eval(sparse_values) + point_vectors["sparse"] = SparseVector( + indices=list(sparse_indices), + values=list(sparse_values), + ) + + metadata = VectorMetadata( + id=bed_id, + name=row.get("name") or "", + description=row.get("description") or "", + genome_alias=row.get("genome_alias") or "", + genome_digest=row.get("genome_digest"), + cell_line=row.get("cell_line") or "", + cell_type=row.get("cell_type") or "", + tissue=row.get("tissue") or "", + target=row.get("target") or "", + treatment=row.get("treatment") or "", + assay=row.get("assay") or "", + species_name=row.get("species_name") or "", + ) + + points_batch.append( + PointStruct( + id=bed_id, + vector=point_vectors, + payload=metadata.model_dump(), + ) + ) + ids_batch.append(bed_id) + + if len(points_batch) >= batch: + _upsert_and_mark( + session, qd_client, collection, points_batch, ids_batch, "indexed" + ) + uploaded += len(points_batch) + _LOGGER.info(f"Uploaded {uploaded} points") + points_batch = [] + ids_batch = [] + + if points_batch: + _upsert_and_mark( + session, qd_client, collection, points_batch, ids_batch, "indexed" + ) + uploaded += len(points_batch) + + _LOGGER.info(f"Hybrid upload complete: {uploaded} points uploaded to '{collection}'") + print(f"Upload complete: {uploaded} points to collection '{collection}'") + + +def _upsert_and_mark( + session: Session, + qd_client, + collection: str, + points: list[PointStruct], + bed_ids: list[str], + flag_column: str, +) -> None: + """Upsert a batch of points to qdrant and update the DB indexed flag. + + Args: + session: Active SQLAlchemy session. + qd_client: Qdrant client instance. + collection: Qdrant collection name. + points: List of PointStruct to upsert. + bed_ids: Corresponding bed IDs to mark in the DB. + flag_column: DB column to set True ('file_indexed' or 'indexed'). + """ + operation_info = qd_client.upsert( + collection_name=collection, + points=points, + ) + assert operation_info.status in ("completed", "acknowledged") + + session.query(Bed).filter(Bed.id.in_(bed_ids)).update( + {getattr(Bed, flag_column): True}, + synchronize_session=False, + ) + session.commit() \ No newline at end of file diff --git a/bedboss/qdrant_index/vectorize.py b/bedboss/qdrant_index/vectorize.py new file mode 100644 index 0000000..545378c --- /dev/null +++ b/bedboss/qdrant_index/vectorize.py @@ -0,0 +1,182 @@ +""" +Connection-free vectorization for HPC jobs. + +Each function loads a model by path, processes a chunk PEP CSV, +and writes vectors + metadata to a parquet file. No bbagent or +database connection is needed. +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import numpy as np +import pandas as pd + +_LOGGER = logging.getLogger(__name__) + +# Columns carried through from PEP CSV to parquet output (metadata payload) +METADATA_COLUMNS = [ + "sample_name", + "name", + "description", + "genome_alias", + "genome_digest", + "cell_line", + "cell_type", + "tissue", + "target", + "treatment", + "assay", + "species_name", +] + + +def vectorize_region( + chunk_pep: str, + output_parquet: str, + model_path: str, +) -> None: + """Vectorize bed files using Region2Vec. Writes results to parquet. + + Args: + chunk_pep: Path to chunk CSV with sample_name + metadata columns. + output_parquet: Path to write output parquet file. + model_path: Path or HuggingFace ID for the Region2VecExModel. + """ + from geniml.bbclient import BBClient + from geniml.region2vec.main import Region2VecExModel + from gtars.models import RegionSet as GRegionSet + + _LOGGER.info(f"Loading Region2Vec model from {model_path}") + encoder = Region2VecExModel(model_path) + + df = pd.read_csv(chunk_pep) + _LOGGER.info(f"Processing {len(df)} bed files") + + bb_client = BBClient() + rows = [] + failed_ids = [] + + for _, row in df.iterrows(): + bed_id = row["sample_name"] + try: + try: + bed_region_set = GRegionSet(bb_client.seek(bed_id)) + except FileNotFoundError: + bed_region_set = bb_client.load_bed(bed_id) + + embedding = np.mean(encoder.encode(bed_region_set), axis=0) + vector = embedding.tolist() + + record = {"vector": vector} + for col in METADATA_COLUMNS: + record[col] = row.get(col) + rows.append(record) + + except Exception as e: + _LOGGER.warning(f"Failed to vectorize {bed_id}: {e}") + failed_ids.append(bed_id) + continue + + if not rows: + _LOGGER.warning("No files were successfully vectorized") + # Write empty parquet so the upload step doesn't break + out = pd.DataFrame(columns=["vector"] + METADATA_COLUMNS) + out.to_parquet(output_parquet, index=False) + else: + out = pd.DataFrame(rows) + out.to_parquet(output_parquet, index=False) + _LOGGER.info( + f"Wrote {len(rows)} vectors to {output_parquet} " + f"({len(failed_ids)} failures)" + ) + + # Write failed IDs for debugging + if failed_ids: + failed_path = Path(output_parquet).parent / "failed_ids.txt" + failed_path.write_text("\n".join(failed_ids) + "\n") + + +def vectorize_hybrid( + chunk_pep: str, + output_parquet: str, + model_path: str, + sparse_model_path: str | None = None, +) -> None: + """Vectorize bed metadata text using dense + sparse encoders. Writes to parquet. + + Args: + chunk_pep: Path to chunk CSV with sample_name + metadata columns. + output_parquet: Path to write output parquet file. + model_path: Path or HuggingFace ID for the dense text encoder. + sparse_model_path: Path or HuggingFace ID for the sparse encoder. Optional. + """ + from fastembed import TextEmbedding + + _LOGGER.info(f"Loading dense encoder from {model_path}") + dense_encoder = TextEmbedding(model_path) + + sparse_encoder = None + if sparse_model_path: + from sentence_transformers import SparseEncoder as STSparseEncoder + + _LOGGER.info(f"Loading sparse encoder from {sparse_model_path}") + sparse_encoder = STSparseEncoder(sparse_model_path) + + df = pd.read_csv(chunk_pep) + _LOGGER.info(f"Processing {len(df)} bed files for hybrid search") + + rows = [] + failed_ids = [] + + for _, row in df.iterrows(): + bed_id = row["sample_name"] + try: + text = ( + f"biosample is {row.get('cell_line')} / {row.get('cell_type')} / " + f"{row.get('tissue')} with target {row.get('target')} " + f"assay {row.get('assay')}." + f"File name {row.get('name')} with summary {row.get('description')}" + ) + + dense_vec = list(list(dense_encoder.embed(text))[0]) + + record = {"dense_vector": dense_vec} + + if sparse_encoder: + sparse_result = sparse_encoder.encode(text).coalesce() + record["sparse_indices"] = sparse_result.indices().tolist()[0] + record["sparse_values"] = sparse_result.values().tolist() + else: + record["sparse_indices"] = None + record["sparse_values"] = None + + for col in METADATA_COLUMNS: + record[col] = row.get(col) + rows.append(record) + + except Exception as e: + _LOGGER.warning(f"Failed to vectorize {bed_id}: {e}") + failed_ids.append(bed_id) + continue + + if not rows: + _LOGGER.warning("No files were successfully vectorized") + out = pd.DataFrame( + columns=["dense_vector", "sparse_indices", "sparse_values"] + + METADATA_COLUMNS + ) + out.to_parquet(output_parquet, index=False) + else: + out = pd.DataFrame(rows) + out.to_parquet(output_parquet, index=False) + _LOGGER.info( + f"Wrote {len(rows)} vectors to {output_parquet} " + f"({len(failed_ids)} failures)" + ) + + if failed_ids: + failed_path = Path(output_parquet).parent / "failed_ids.txt" + failed_path.write_text("\n".join(failed_ids) + "\n") \ No newline at end of file From 62e573a72a106aec5f099db3e46093e6ebc2cde8 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Tue, 14 Apr 2026 14:26:25 -0400 Subject: [PATCH 11/17] Fixed db connection --- bedboss/qdrant_index/upload.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bedboss/qdrant_index/upload.py b/bedboss/qdrant_index/upload.py index 0a5fc18..68331c5 100644 --- a/bedboss/qdrant_index/upload.py +++ b/bedboss/qdrant_index/upload.py @@ -64,7 +64,7 @@ def upload_region_vectors( points_batch: list[PointStruct] = [] ids_batch: list[str] = [] - with Session(agent.config.db_engine) as session: + with Session(agent.config.db_engine.engine) as session: for _, row in df.iterrows(): bed_id = row["sample_name"] vector = row["vector"] @@ -140,7 +140,7 @@ def upload_hybrid_vectors( points_batch: list[PointStruct] = [] ids_batch: list[str] = [] - with Session(agent.config.db_engine) as session: + with Session(agent.config.db_engine.engine) as session: for _, row in df.iterrows(): bed_id = row["sample_name"] dense_vec = row["dense_vector"] From c15655ad75a81ea5fbf6da934be6e3739dc743d2 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Tue, 14 Apr 2026 14:37:49 -0400 Subject: [PATCH 12/17] Fixed string line parsing --- bedboss/qdrant_index/upload.py | 47 +++++++++++++++++++--------------- 1 file changed, 27 insertions(+), 20 deletions(-) diff --git a/bedboss/qdrant_index/upload.py b/bedboss/qdrant_index/upload.py index 68331c5..a0b921e 100644 --- a/bedboss/qdrant_index/upload.py +++ b/bedboss/qdrant_index/upload.py @@ -20,6 +20,13 @@ _LOGGER = logging.getLogger(__name__) +def _str(val) -> str: + """Coerce a value to str, treating pandas NaN/None as empty string.""" + if val is None or (isinstance(val, float) and pd.isna(val)): + return "" + return str(val) + + def _load_parquet_files(workdir: Path) -> pd.DataFrame: """Glob and concatenate all vectors.parquet files from chunk output dirs.""" parquet_files = sorted(workdir.glob("chunks/*/output/vectors.parquet")) @@ -75,17 +82,17 @@ def upload_region_vectors( metadata = VectorMetadata( id=bed_id, - name=row.get("name") or "", - description=row.get("description") or "", - genome_alias=row.get("genome_alias") or "", + name=_str(row.get("name")), + description=_str(row.get("description")), + genome_alias=_str(row.get("genome_alias")), genome_digest=row.get("genome_digest"), - cell_line=row.get("cell_line") or "", - cell_type=row.get("cell_type") or "", - tissue=row.get("tissue") or "", - target=row.get("target") or "", - treatment=row.get("treatment") or "", - assay=row.get("assay") or "", - species_name=row.get("species_name") or "", + cell_line=_str(row.get("cell_line")), + cell_type=_str(row.get("cell_type")), + tissue=_str(row.get("tissue")), + target=_str(row.get("target")), + treatment=_str(row.get("treatment")), + assay=_str(row.get("assay")), + species_name=_str(row.get("species_name")), ) points_batch.append( @@ -166,17 +173,17 @@ def upload_hybrid_vectors( metadata = VectorMetadata( id=bed_id, - name=row.get("name") or "", - description=row.get("description") or "", - genome_alias=row.get("genome_alias") or "", + name=_str(row.get("name")), + description=_str(row.get("description")), + genome_alias=_str(row.get("genome_alias")), genome_digest=row.get("genome_digest"), - cell_line=row.get("cell_line") or "", - cell_type=row.get("cell_type") or "", - tissue=row.get("tissue") or "", - target=row.get("target") or "", - treatment=row.get("treatment") or "", - assay=row.get("assay") or "", - species_name=row.get("species_name") or "", + cell_line=_str(row.get("cell_line")), + cell_type=_str(row.get("cell_type")), + tissue=_str(row.get("tissue")), + target=_str(row.get("target")), + treatment=_str(row.get("treatment")), + assay=_str(row.get("assay")), + species_name=_str(row.get("species_name")), ) points_batch.append( From 9bbf9e836acf8c91fe1ca574da414a8829815628 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Fri, 17 Apr 2026 12:13:07 -0400 Subject: [PATCH 13/17] Updated function saving umap --- bedboss/cli.py | 16 +- bedboss/const.py | 26 +- .../refgenie_chrom_sizes.py | 22 +- bedboss/scripts/make_umap.py | 245 +++++------------- 4 files changed, 94 insertions(+), 215 deletions(-) diff --git a/bedboss/cli.py b/bedboss/cli.py index 58956ff..f3a8c4b 100644 --- a/bedboss/cli.py +++ b/bedboss/cli.py @@ -799,9 +799,9 @@ def download_umap( "umap", help="Dimensionality reduction method to use. Options: 'umap', 'pca', or 'tsne'. To use UMAP, 'umap-learn' package must be installed.", ), - save_parquet: bool = typer.Option( - False, - help="Whether to save Parquet tier files alongside JSON", + output_format: str = typer.Option( + "parquet", + help="Output format: 'json', 'parquet', or 'both'", ), ): from bedboss.scripts.make_umap import get_embeddings @@ -815,11 +815,11 @@ def download_umap( top_assays=top_assays, top_cell_lines=top_cell_lines, method=method, - save_parquet=save_parquet, + output_format=output_format, ) -@app.command(help="Update UMAP metadata Parquet tiers without regenerating geometry") +@app.command(help="Update UMAP parquet metadata without regenerating geometry") def update_umap_metadata( config: str = typer.Option( ..., @@ -828,9 +828,9 @@ def update_umap_metadata( file_okay=True, readable=True, ), - output_dir: str = typer.Option( + output_path: str = typer.Option( ..., - help="Directory to write Parquet tier files", + help="Path to write parquet file (without extension)", ), geometry: str = typer.Option( None, @@ -839,7 +839,7 @@ def update_umap_metadata( ): from bedboss.scripts.make_umap import update_umap_metadata as _update - _update(bbconf=config, output_dir=output_dir, geometry=geometry) + _update(bbconf=config, output_path=output_path, geometry=geometry) @app.command(help="Check installed R packages") diff --git a/bedboss/const.py b/bedboss/const.py index bd0e420..f47e2c6 100644 --- a/bedboss/const.py +++ b/bedboss/const.py @@ -37,34 +37,18 @@ BED_PEP_REGISTRY: str = "databio/allbeds:bedbase" -# UMAP Parquet tier column definitions +# UMAP constants DB_QUERY_BATCH_SIZE: int = 5000 +UMAP_GENOME: str = "hg38" -TIER1_COLUMNS: list[str] = [ +UMAP_PARQUET_COLUMNS: list[str] = [ "id", + "x", + "y", "name", "description", "assay", - "target", "cell_line", "cell_type", "tissue", - "number_of_regions", - "mean_region_width", - "gc_content", -] - -TIER2_COLUMNS: list[str] = [ - "id", - "treatment", - "antibody", - "species_name", - "genome_alias", - "bed_compliance", - "data_format", - "median_tss_dist", - "library_source", - "global_sample_id", - "global_experiment_id", - "original_file_name", ] diff --git a/bedboss/refgenome_validator/refgenie_chrom_sizes.py b/bedboss/refgenome_validator/refgenie_chrom_sizes.py index d801215..7237dea 100644 --- a/bedboss/refgenome_validator/refgenie_chrom_sizes.py +++ b/bedboss/refgenome_validator/refgenie_chrom_sizes.py @@ -21,6 +21,7 @@ SEQ_COL_URL = os.path.join( BASE_URL, "seqcol/collection/{digest}?collated=true&attribute=name_length_pairs" ) +SEQ_COL_URL_JSON_URL = "https://huggingface.co/databio/bedbase-umap/resolve/main/genome_seqcol.json" _LOGGER = logging.getLogger(PKG_NAME) @@ -157,6 +158,18 @@ def read_seq_col_from_json(input_path: str = "genome_seqcol.json") -> Genomes: data = json.load(f) return Genomes(**data) +def read_seq_col_from_url(input_path: str = "genome_seqcol.json") -> Genomes: + """ + Read sequence collections from a JSON URL. + + :param input_path: URL to the JSON file. + :return: Genomes object containing the sequence collections. + """ + data = run_requests(input_path) + if not data: + raise BedBossException(f"Failed to fetch sequence collections from URL: {input_path}") + return Genomes(**data) + def modify_for_analysis(genomes: Genomes) -> list[GenomeModel]: """ @@ -197,9 +210,12 @@ def get_chrom_sizes() -> list[GenomeModel]: try: ret = read_seq_col_from_json(input_path=cached_file_path) except FileNotFoundError: - _LOGGER.info("No genome_seqcol.json found, downloading from refgenie...") - ret = get_seq_col() - + try: + _LOGGER.info("No genome_seqcol.json found, downloading from refgenie...") + ret = get_seq_col() + except Exception as e: + _LOGGER.info(f"Failed to fetch genome data from Refgenie: {e}") + ret = read_seq_col_from_url(input_path=SEQ_COL_URL_JSON_URL) save_seq_col_to_json(ret, output_path=cached_file_path) return modify_for_analysis(ret) diff --git a/bedboss/scripts/make_umap.py b/bedboss/scripts/make_umap.py index b5d31d5..8302f2e 100644 --- a/bedboss/scripts/make_umap.py +++ b/bedboss/scripts/make_umap.py @@ -10,15 +10,13 @@ import pandas as pd import seaborn as sns from bbconf import BedBaseAgent -from bbconf.db_utils import Bed, BedStats from pydantic import BaseModel, ConfigDict from qdrant_client import QdrantClient from sklearn.decomposition import PCA from sklearn.manifold import TSNE -from sqlalchemy.orm import Session, joinedload from umap import UMAP -from bedboss.const import DB_QUERY_BATCH_SIZE, PKG_NAME, TIER1_COLUMNS, TIER2_COLUMNS +from bedboss.const import PKG_NAME, UMAP_PARQUET_COLUMNS _LOGGER = logging.getLogger(PKG_NAME) @@ -32,24 +30,6 @@ class umapReturn(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) -class BedDbMetadata(BaseModel): - """Schema for per-file metadata fetched from PostgreSQL.""" - - id: str - number_of_regions: float | None = None - mean_region_width: float | None = None - gc_content: float | None = None - median_tss_dist: float | None = None - antibody: str | None = None - library_source: str | None = None - original_file_name: str | None = None - global_sample_id: str | None = None - global_experiment_id: str | None = None - bed_compliance: str | None = None - data_format: str | None = None - - model_config = ConfigDict(from_attributes=True) - def save_umap_model(umap_model: UMAP | PCA | TSNE, model_path: str) -> None: """ @@ -113,79 +93,6 @@ def fetch_data(agent: BedBaseAgent) -> pd.DataFrame: return merged -def fetch_db_metadata(agent: BedBaseAgent, bed_ids: list[str]) -> pd.DataFrame: - """ - Fetch bed_stats and annotation metadata from PostgreSQL for the given bed IDs. - - Returns a DataFrame indexed by bed ID with stats and annotation columns. - """ - _LOGGER.info(f"Fetching DB metadata for {len(bed_ids)} beds...") - - rows = [] - with Session(agent.config.db_engine.engine) as session: - for i in range(0, len(bed_ids), DB_QUERY_BATCH_SIZE): - batch = bed_ids[i : i + DB_QUERY_BATCH_SIZE] - for bed_obj in ( - session.query(Bed) - .options( - joinedload(Bed.stats).load_only( - BedStats.number_of_regions, - BedStats.mean_region_width, - BedStats.gc_content, - BedStats.median_tss_dist, - ), - joinedload(Bed.annotations), - ) - .filter(Bed.id.in_(batch)) - .all() - ): - meta = BedDbMetadata( - id=bed_obj.id, - number_of_regions=( - bed_obj.stats.number_of_regions if bed_obj.stats else None - ), - mean_region_width=( - bed_obj.stats.mean_region_width if bed_obj.stats else None - ), - gc_content=(bed_obj.stats.gc_content if bed_obj.stats else None), - median_tss_dist=( - bed_obj.stats.median_tss_dist if bed_obj.stats else None - ), - antibody=( - bed_obj.annotations.antibody if bed_obj.annotations else None - ), - library_source=( - bed_obj.annotations.library_source - if bed_obj.annotations - else None - ), - original_file_name=( - bed_obj.annotations.original_file_name - if bed_obj.annotations - else None - ), - global_sample_id=( - ";".join(bed_obj.annotations.global_sample_id) - if bed_obj.annotations and bed_obj.annotations.global_sample_id - else None - ), - global_experiment_id=( - ";".join(bed_obj.annotations.global_experiment_id) - if bed_obj.annotations - and bed_obj.annotations.global_experiment_id - else None - ), - bed_compliance=bed_obj.bed_compliance, - data_format=bed_obj.data_format, - ) - rows.append(meta.model_dump()) - - df = pd.DataFrame(rows) - if not df.empty: - df = df.set_index("id") - _LOGGER.info(f"Fetched DB metadata for {len(df)} beds.") - return df - def save_df_as_json(df: pd.DataFrame, output_path: str) -> None: """ @@ -208,6 +115,8 @@ def save_df_as_json(df: pd.DataFrame, output_path: str) -> None: # "bed_compliance", "assay", "cell_line", + "cell_type", + "tissue", ] if "z" in df.columns: @@ -226,7 +135,7 @@ def save_df_as_json(df: pd.DataFrame, output_path: str) -> None: for node in nodes: for col in coord_cols: if col in node: - node[col] = round(float(node[col]), 2) + node[col] = round(float(node[col]), 3) # Create the final JSON structure json_data = {"nodes": nodes, "links": []} @@ -239,64 +148,44 @@ def save_df_as_json(df: pd.DataFrame, output_path: str) -> None: _LOGGER.info(f"Data saved to {output_path} successfully.") -def save_parquet_tiers( +def save_parquet( df: pd.DataFrame, - db_meta: pd.DataFrame, - output_dir: str, + output_path: str, ) -> None: """ - Save UMAP data as tiered Parquet files. + Save UMAP data as a single Parquet file. - Produces: - - hg38_geometry.parquet (x, y, id) - - hg38_meta_t1.parquet (core biological annotation + region stats) - - hg38_meta_t2.parquet (extended annotation) - """ - os.makedirs(output_dir, exist_ok=True) + Columns: id, x, y, name, description, assay, cell_line (+ z if 3D). - # Ensure id is a column (not just index) + :param df: DataFrame with UMAP coordinates and metadata. + :param output_path: Path to save the parquet file (without extension). + """ if "id" not in df.columns: df = df.copy() df["id"] = df.index.astype(str) - # Join DB metadata - if not db_meta.empty: - combined = df.join(db_meta, how="left", rsuffix="_db") - else: - combined = df - - # --- Geometry (only when coordinates are present) --- - if "x" in combined.columns and "y" in combined.columns: - coord_cols = ["id", "x", "y"] - if "z" in combined.columns: - coord_cols.append("z") - geometry = combined[coord_cols].copy() - for col in ["x", "y", "z"]: - if col in geometry.columns: - geometry[col] = geometry[col].round(2).astype("float32") - geometry.to_parquet( - os.path.join(output_dir, "hg38_geometry.parquet"), index=False - ) - _LOGGER.info(f"Geometry: {len(geometry)} rows") - else: - _LOGGER.info("No coordinates found, skipping geometry file.") + cols = list(UMAP_PARQUET_COLUMNS) + if "z" in df.columns: + cols.insert(3, "z") + + out = df[[c for c in cols if c in df.columns]].copy() - # --- Tier 1: Core metadata --- - t1 = combined[[c for c in TIER1_COLUMNS if c in combined.columns]].copy() - str_cols = t1.select_dtypes(include="object").columns - t1[str_cols] = t1[str_cols].fillna("") - t1.to_parquet(os.path.join(output_dir, "hg38_meta_t1.parquet"), index=False) - _LOGGER.info(f"Tier 1: {len(t1)} rows, {len(t1.columns)} columns") + # float32 for coordinates, dictionary encoding for categorical strings + for col in ["x", "y", "z"]: + if col in out.columns: + out[col] = out[col].round(3).astype("float32") - # --- Tier 2: Extended annotation --- - t2 = combined[[c for c in TIER2_COLUMNS if c in combined.columns]].copy() - str_cols = t2.select_dtypes(include="object").columns - t2[str_cols] = t2[str_cols].fillna("") - t2.to_parquet(os.path.join(output_dir, "hg38_meta_t2.parquet"), index=False) - _LOGGER.info(f"Tier 2: {len(t2)} rows, {len(t2.columns)} columns") + str_cols = out.select_dtypes(include="object").columns + out[str_cols] = out[str_cols].fillna("") - # Tier 3 reserved for future analysis results (gtars genomic distributions, - # enrichment profiles, embedding quality scores). Not generated in this version. + parquet_path = f"{output_path}_{python_version}.parquet" + out.to_parquet( + parquet_path, + index=False, + engine="pyarrow", + compression="snappy", + ) + _LOGGER.info(f"Parquet saved to {parquet_path}: {len(out)} rows, {len(out.columns)} columns") def create_umap( @@ -423,17 +312,16 @@ def plot_umap(value, label, name="default") -> None: def update_umap_metadata( bbconf: str, - output_dir: str, + output_path: str, geometry: str = None, ) -> None: """ - Update UMAP metadata Parquet tiers without regenerating geometry. + Update UMAP parquet without regenerating geometry. - Args: - bbconf: Path to bedbase configuration file. - output_dir: Directory to write Parquet tier files. - geometry: Path to existing geometry Parquet to read bed IDs from. - If not provided, fetches IDs from Qdrant. + :param bbconf: Path to bedbase configuration file. + :param output_path: Path to write parquet file (without extension). + :param geometry: Path to existing geometry Parquet to read bed IDs from. + If not provided, fetches IDs from Qdrant. """ if isinstance(bbconf, str): agent = BedBaseAgent(config=bbconf) @@ -451,10 +339,8 @@ def update_umap_metadata( qdrant_df = qdrant_df.loc[qdrant_df.index.isin(bed_ids)] else: qdrant_df = fetch_data(agent=agent) - bed_ids = list(qdrant_df.index) - db_meta = fetch_db_metadata(agent, bed_ids) - save_parquet_tiers(qdrant_df, db_meta, output_dir) + save_parquet(qdrant_df, output_path) def get_embeddings( @@ -467,22 +353,21 @@ def get_embeddings( top_cell_lines: int | None = 15, save_model: bool = True, method: str = "umap", - save_parquet: bool = False, + output_format: str = "parquet", ) -> None: """ - Get embeddings from Qdrant, create UMAP/PCA/t-SNE, and save results to a JSON file. - - Args: - bbconf: Path to bedbase configuration file. - output_file: Path to save the output JSON file. - n_components: Number of dimensions for UMAP/PCA/t-SNE. Default: 2. - plot_name: Name for the output plot file. If None, no plot will be saved. - plot_label: Column name to use for labeling plot points (e.g. "cell_line" or "assay"). - top_assays: Number of top assays to consider. If None, all assays are used. Default: 15. - top_cell_lines: Number of top cell lines to consider. If None, all. Default: 15. - save_model: Whether to save the fitted model. Default: True. - method: Dimensionality reduction method. Options: "umap", "pca", or "tsne". Default: "umap". - save_parquet: Whether to save Parquet tier files alongside JSON. Default: False. + Get embeddings from Qdrant, create UMAP/PCA/t-SNE, and save results. + + :param bbconf: Path to bedbase configuration file. + :param output_file: Path to save the output file (without extension). + :param n_components: Number of dimensions for UMAP/PCA/t-SNE. + :param plot_name: Name for the output plot file. If None, no plot will be saved. + :param plot_label: Column name to use for labeling plot points. + :param top_assays: Number of top assays to consider. If None, all assays are used. + :param top_cell_lines: Number of top cell lines to consider. If None, all. + :param save_model: Whether to save the fitted model. + :param method: Dimensionality reduction method: "umap", "pca", or "tsne". + :param output_format: Output format: "json", "parquet", or "both". """ if isinstance(bbconf, str): @@ -496,10 +381,13 @@ def get_embeddings( merged = fetch_data(agent=agent) - if output_file.endswith(".json"): - output_file = output_file[:-5] - # output_file += ".json" - os.makedirs(os.path.dirname(output_file), exist_ok=True) + for ext in (".json", ".parquet"): + if output_file.endswith(ext): + output_file = output_file[: -len(ext)] + break + output_dir = os.path.dirname(output_file) + if output_dir: + os.makedirs(output_dir, exist_ok=True) CELL_LINE = "cell_line" ASSAY = "assay" @@ -565,20 +453,11 @@ def get_embeddings( method=method, ) - # Legacy JSON output - save_df_as_json(umap_return.dataframe, output_file) - - # Parquet tiered output - if save_parquet: - parquet_dir = os.path.dirname(os.path.abspath(output_file)) - try: - db_meta = fetch_db_metadata(agent, list(umap_return.dataframe.index)) - except Exception as e: - _LOGGER.warning( - f"Failed to fetch DB metadata, Parquet tiers will lack stats/annotation: {e}" - ) - db_meta = pd.DataFrame() - save_parquet_tiers(umap_return.dataframe, db_meta, parquet_dir) + if output_format in ("json", "both"): + save_df_as_json(umap_return.dataframe, output_file) + + if output_format in ("parquet", "both"): + save_parquet(umap_return.dataframe, output_file) if save_model: # controls the random initialization and stochastic optimization during UMAP fitting. But removing, because it causes issues during saving/loading From d09e24f486b6b15e22ba07b975fb13f1acdd4c23 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Fri, 17 Apr 2026 12:22:46 -0400 Subject: [PATCH 14/17] fmt and new version --- bedboss/qdrant_index/qdrant_hpc.py | 17 ++++------------- bedboss/qdrant_index/upload.py | 17 +++++++++++++---- bedboss/qdrant_index/vectorize.py | 2 +- bedboss/refgenome_validator/main.py | 18 +++++++++++------- .../refgenie_chrom_sizes.py | 9 +++++++-- pyproject.toml | 2 +- 6 files changed, 37 insertions(+), 28 deletions(-) diff --git a/bedboss/qdrant_index/qdrant_hpc.py b/bedboss/qdrant_index/qdrant_hpc.py index 547b7ba..44e50da 100644 --- a/bedboss/qdrant_index/qdrant_hpc.py +++ b/bedboss/qdrant_index/qdrant_hpc.py @@ -175,9 +175,7 @@ def _fetch_region_metadata( conditions.append(Bed.file_indexed.is_(False)) statement = ( - select(Bed) - .join(BedMetadata, Bed.id == BedMetadata.id) - .where(and_(*conditions)) + select(Bed).join(BedMetadata, Bed.id == BedMetadata.id).where(and_(*conditions)) ) if limit: statement = statement.limit(limit) @@ -388,9 +386,7 @@ def _get_alive_job_ids(chunks: list[ChunkMeta]) -> set[str]: return {line.strip() for line in out.stdout.strip().splitlines() if line.strip()} -def _chunk_status( - chunk: ChunkMeta, state_dir: Path, alive_jobs: set[str] -) -> str: +def _chunk_status(chunk: ChunkMeta, state_dir: Path, alive_jobs: set[str]) -> str: cid = chunk.id if (state_dir / f"{cid}.done").exists(): return "done" @@ -599,9 +595,7 @@ def reindex_hpc_status(workdir: str, verbose: bool = False) -> None: n_vectors = _parquet_row_count(Path(chunk.output_parquet)) total_vectors += n_vectors - rows.append( - (chunk.id, chunk.n_samples, status, chunk.job_id or "-", n_vectors) - ) + rows.append((chunk.id, chunk.n_samples, status, chunk.job_id or "-", n_vectors)) print(f"Search type: {manifest.search_type}") print(f"Model: {manifest.model_path}") @@ -628,10 +622,7 @@ def reindex_hpc_status(workdir: str, verbose: bool = False) -> None: f"Chunks: done={counts['done']} failed={counts['failed']} " f"running={counts['running']} pending={counts['pending']} (of {total})" ) - print( - f"Samples: {manifest.total_samples} total, " - f"{total_vectors} vectors produced" - ) + print(f"Samples: {manifest.total_samples} total, {total_vectors} vectors produced") if failed_chunks: show_n = min(20, len(failed_chunks)) diff --git a/bedboss/qdrant_index/upload.py b/bedboss/qdrant_index/upload.py index a0b921e..892f595 100644 --- a/bedboss/qdrant_index/upload.py +++ b/bedboss/qdrant_index/upload.py @@ -106,7 +106,12 @@ def upload_region_vectors( if len(points_batch) >= batch: _upsert_and_mark( - session, qd_client, collection, points_batch, ids_batch, "file_indexed" + session, + qd_client, + collection, + points_batch, + ids_batch, + "file_indexed", ) uploaded += len(points_batch) _LOGGER.info(f"Uploaded {uploaded} points") @@ -120,7 +125,9 @@ def upload_region_vectors( ) uploaded += len(points_batch) - _LOGGER.info(f"Region upload complete: {uploaded} points uploaded to '{collection}'") + _LOGGER.info( + f"Region upload complete: {uploaded} points uploaded to '{collection}'" + ) print(f"Upload complete: {uploaded} points to collection '{collection}'") @@ -210,7 +217,9 @@ def upload_hybrid_vectors( ) uploaded += len(points_batch) - _LOGGER.info(f"Hybrid upload complete: {uploaded} points uploaded to '{collection}'") + _LOGGER.info( + f"Hybrid upload complete: {uploaded} points uploaded to '{collection}'" + ) print(f"Upload complete: {uploaded} points to collection '{collection}'") @@ -242,4 +251,4 @@ def _upsert_and_mark( {getattr(Bed, flag_column): True}, synchronize_session=False, ) - session.commit() \ No newline at end of file + session.commit() diff --git a/bedboss/qdrant_index/vectorize.py b/bedboss/qdrant_index/vectorize.py index 545378c..239a66b 100644 --- a/bedboss/qdrant_index/vectorize.py +++ b/bedboss/qdrant_index/vectorize.py @@ -179,4 +179,4 @@ def vectorize_hybrid( if failed_ids: failed_path = Path(output_parquet).parent / "failed_ids.txt" - failed_path.write_text("\n".join(failed_ids) + "\n") \ No newline at end of file + failed_path.write_text("\n".join(failed_ids) + "\n") diff --git a/bedboss/refgenome_validator/main.py b/bedboss/refgenome_validator/main.py index 71002e6..8b82cfe 100644 --- a/bedboss/refgenome_validator/main.py +++ b/bedboss/refgenome_validator/main.py @@ -276,8 +276,10 @@ def determine_compatibility( for genome_model in self.genome_models: # First and Second Layer of Compatibility - model_compat_stats[genome_model.genome_digest]: CompatibilityStats = ( - self.calculate_chrom_stats(bed_chrom_info, genome_model.chrom_sizes) + model_compat_stats[ + genome_model.genome_digest + ]: CompatibilityStats = self.calculate_chrom_stats( + bed_chrom_info, genome_model.chrom_sizes ) # Third layer - IGD, only if layer 1 and layer 2 have passed @@ -289,13 +291,15 @@ def determine_compatibility( genome_model.genome_digest ].chrom_length_stats.beyond_range ): - model_compat_stats[genome_model.genome_digest].igd_stats = ( - self.get_igd_overlaps(bedfile) - ) + model_compat_stats[ + genome_model.genome_digest + ].igd_stats = self.get_igd_overlaps(bedfile) # Calculate compatibility rating - model_compat_stats[genome_model.genome_digest].compatibility = ( - self.calculate_rating(model_compat_stats[genome_model.genome_digest]) + model_compat_stats[ + genome_model.genome_digest + ].compatibility = self.calculate_rating( + model_compat_stats[genome_model.genome_digest] ) if concise: concise_dict = {} diff --git a/bedboss/refgenome_validator/refgenie_chrom_sizes.py b/bedboss/refgenome_validator/refgenie_chrom_sizes.py index 7237dea..3519b57 100644 --- a/bedboss/refgenome_validator/refgenie_chrom_sizes.py +++ b/bedboss/refgenome_validator/refgenie_chrom_sizes.py @@ -21,7 +21,9 @@ SEQ_COL_URL = os.path.join( BASE_URL, "seqcol/collection/{digest}?collated=true&attribute=name_length_pairs" ) -SEQ_COL_URL_JSON_URL = "https://huggingface.co/databio/bedbase-umap/resolve/main/genome_seqcol.json" +SEQ_COL_URL_JSON_URL = ( + "https://huggingface.co/databio/bedbase-umap/resolve/main/genome_seqcol.json" +) _LOGGER = logging.getLogger(PKG_NAME) @@ -158,6 +160,7 @@ def read_seq_col_from_json(input_path: str = "genome_seqcol.json") -> Genomes: data = json.load(f) return Genomes(**data) + def read_seq_col_from_url(input_path: str = "genome_seqcol.json") -> Genomes: """ Read sequence collections from a JSON URL. @@ -167,7 +170,9 @@ def read_seq_col_from_url(input_path: str = "genome_seqcol.json") -> Genomes: """ data = run_requests(input_path) if not data: - raise BedBossException(f"Failed to fetch sequence collections from URL: {input_path}") + raise BedBossException( + f"Failed to fetch sequence collections from URL: {input_path}" + ) return Genomes(**data) diff --git a/pyproject.toml b/pyproject.toml index f303f51..9ee2fb6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "bedboss" -version = "0.10.1" +version = "0.11.0" description = "Pipelines for genomic region file to produce bed files, and its statistics" readme = "README.md" license = "BSD-2-Clause" From 9546170b12efaab2db85e17a9a47f0e10ab577e2 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Fri, 17 Apr 2026 12:25:42 -0400 Subject: [PATCH 15/17] Added changelog --- docs/changelog.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/changelog.md b/docs/changelog.md index b2f6eac..7f7f965 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -2,6 +2,16 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html) and [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) format. +# [0.11.0] - 2026-04-17 +## Added: +- HPC module for running bedboss on HPC +- CLI tool for bulk uploading files to Qdrant + +## Fixed: +- Bed classifier +- DB connection + + # [0.10.0] - 2026-04-05 ## Added: - Creation of parquet file for umap with more metadata inside From 4bbe586510305a7f1e2005ed2e27c50d991e7e71 Mon Sep 17 00:00:00 2001 From: Oleksandr <41573628+khoroshevskyi@users.noreply.github.com> Date: Fri, 17 Apr 2026 12:43:33 -0400 Subject: [PATCH 16/17] Update bedboss/qdrant_index/upload.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- bedboss/qdrant_index/upload.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/bedboss/qdrant_index/upload.py b/bedboss/qdrant_index/upload.py index 892f595..2a9fd99 100644 --- a/bedboss/qdrant_index/upload.py +++ b/bedboss/qdrant_index/upload.py @@ -245,7 +245,13 @@ def _upsert_and_mark( collection_name=collection, points=points, ) - assert operation_info.status in ("completed", "acknowledged") + status = operation_info.status + if status not in ("completed", "acknowledged"): + raise RuntimeError( + f"Qdrant upsert failed for collection '{collection}': " + f"unexpected status {status!r} for {len(points)} point(s). " + f"operation_info={operation_info!r}" + ) session.query(Bed).filter(Bed.id.in_(bed_ids)).update( {getattr(Bed, flag_column): True}, From 43a5e13614409e8d9e4b1733e5245c3800b2fdbb Mon Sep 17 00:00:00 2001 From: Oleksandr <41573628+khoroshevskyi@users.noreply.github.com> Date: Fri, 17 Apr 2026 12:44:12 -0400 Subject: [PATCH 17/17] Update bedboss/bedboss_hpc.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- bedboss/bedboss_hpc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bedboss/bedboss_hpc.py b/bedboss/bedboss_hpc.py index 3026598..b305b6f 100644 --- a/bedboss/bedboss_hpc.py +++ b/bedboss/bedboss_hpc.py @@ -1,6 +1,6 @@ """ HPC orchestration for `bedboss run-pep`. -! This file is fully generated by AI - Bugs can occure +! This file is fully generated by AI - Bugs can occur Splits a large PEP into N chunks and submits each as its own SLURM job. Idempotent: re-running picks up where it left off via per-chunk sentinel files