From 140701f5737b6bd6624a71db558114e1139e7203 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Buschk=C3=A4mper?= Date: Tue, 21 Jul 2026 09:18:41 +0000 Subject: [PATCH] First version of multi-node linclust. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Björn Buschkämper --- README.md | 72 ++++ data/workflow/linclust.sh | 8 +- src/alignment/Align2ClustCheckpoint.cpp | 190 +++++++++ src/alignment/Align2ClustCheckpoint.h | 70 ++++ src/alignment/Align2ClustChunking.h | 56 +++ src/alignment/Align2ClustReducer.cpp | 266 ++++++++++++ src/alignment/Align2ClustReducer.h | 179 ++++++++ src/alignment/Align2clust.cpp | 528 ++++++++++-------------- src/alignment/CMakeLists.txt | 5 + src/commons/MMseqsMPI.cpp | 47 +++ src/test/CMakeLists.txt | 1 + src/test/TestAlign2ClustReducer.cpp | 258 ++++++++++++ 12 files changed, 1365 insertions(+), 315 deletions(-) create mode 100644 src/alignment/Align2ClustCheckpoint.cpp create mode 100644 src/alignment/Align2ClustCheckpoint.h create mode 100644 src/alignment/Align2ClustChunking.h create mode 100644 src/alignment/Align2ClustReducer.cpp create mode 100644 src/alignment/Align2ClustReducer.h create mode 100644 src/test/TestAlign2ClustReducer.cpp diff --git a/README.md b/README.md index 87f09d98b..a16ac09b0 100644 --- a/README.md +++ b/README.md @@ -150,6 +150,78 @@ To search with multiple servers, call the `search` or `cluster` workflow with th RUNNER="mpirun -pernode -np 42" mmseqs search queryDB targetDB resultDB tmp +#### Multi-node `linclust`/`easy-linclust` on Slurm + +The default (v2) `linclust`/`easy-linclust` pipeline distributes its most expensive stage, +`kmermatcher` and the candidate-alignment stage `align2clust`, across MPI ranks while +keeping cluster assignment (representative selection and membership) byte-identical to a +single-node run. This is the same `--mpi-runner`/`RUNNER` mechanism described above -- +there is no separate Slurm submission layer, so `linclust` should be launched as a single +process (not itself wrapped in `mpirun`/`srun`) with the runner passed through +`--mpi-runner`/`RUNNER`; the workflow script wraps only its MPI-aware sub-stages +(`kmermatcher`, `align2clust`, and, for Linclust v1, `rescorediagonal`/`align`) in that +runner internally. Wrapping the top-level `linclust` invocation in `mpirun`/`srun` yourself +will fail ("recursive mpirun") or, on some launchers, silently run duplicate, +non-cooperating single-rank jobs. + +Requirements and recommended topology: + +* Build with `-DHAVE_MPI=1` (version string gets a `-MPI` suffix). A non-MPI binary + launched under a multi-task runner (e.g. `srun -n 4 mmseqs ...` without `-DHAVE_MPI=1`) + fails fast with an explicit error instead of silently letting every task race on the + same output files. +* The input/output databases and the `--tmp-folder` must be on a shared POSIX filesystem + visible to every node (e.g. NFS, Lustre, GPFS). +* Recommended: one MPI rank per node, with `--threads` set to the number of CPUs allocated + to that node (OpenMP parallelizes within each rank). Running multiple ranks per node + works but duplicates the memory-mapped sequence/prefilter DB cache footprint per rank, so + it is not the default recommendation. + +Example `sbatch` script requesting 4 nodes and running `linclust` once `RUNNER` is set: + +```sh +#!/bin/bash +#SBATCH --nodes=4 +#SBATCH --ntasks=4 +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=32 + +export RUNNER="srun --mpi=pmix -n $SLURM_NTASKS --ntasks-per-node=1" +# or: export RUNNER="mpirun -np $SLURM_NTASKS" + +mmseqs linclust queryDB resultDB /shared/tmp --threads $SLURM_CPUS_PER_TASK +``` + +`--mpi-runner "srun --mpi=pmix -n 4 --ntasks-per-node=1"` works equivalently as a +command-line flag instead of the `RUNNER` environment variable. The exact `srun --mpi=...` +plugin name is site-specific; check with your cluster administrator or `srun --mpi=list`. + +Checkpoint/resume across topology changes: `align2clust` splits its candidate-generation +work into fixed-size, topology-independent chunks (tunable via +`MMSEQS_ALIGN2CLUST_CHUNK_SIZE`, default 50000) and records completed chunks under a +`_chunks` directory next to the output. If a run is interrupted, rerunning the identical +`mmseqs linclust`/`align2clust` command resumes by regenerating only the chunks that were +not yet completed and reusing all finished ones -- this works even if the rank count +changes between the interrupted and resumed run (e.g. starting on 8 nodes and resuming on +1), because chunk boundaries and ownership only depend on the (fixed) input size and chunk +size, never on the runner or rank count. The `_chunks` directory is removed automatically +once a run completes successfully; keep it (e.g. via `--remove-tmp-files 0`) if you want to +inspect or reuse it after a failure. + +Known limitations and scaling expectations: + +* `--include-align-files` is currently not supported together with distributed (multi-rank) + candidate generation in `align2clust` and is rejected at startup in that combination; it + still works normally on a single rank. +* `clusthash` and `clust` (used for the optional `--cluster-mode`/hash-based prefilter path) + and the final `mergeclusters`/representative-switching stages are not MPI-aware and always + run once from the workflow process, regardless of the runner. +* Representative selection and cluster assignment are a single deterministic, ordered + reduction performed on rank 0 after all candidate chunks are collected. Scaling is + therefore strongest when candidate/alignment generation dominates total runtime; very + large numbers of nodes will eventually be bottlenecked by this final reduction step and by + shared-filesystem bandwidth for chunk I/O rather than by compute. + ## Contributors MMseqs2 exists thanks to all the people who contribute. diff --git a/data/workflow/linclust.sh b/data/workflow/linclust.sh index 49074bd10..92feb74e5 100755 --- a/data/workflow/linclust.sh +++ b/data/workflow/linclust.sh @@ -23,13 +23,13 @@ if [ "$LINCLUST_MODULE" = "linclust2" ]; then if [ -n "$CLUSTHASH" ]; then if notExists "${TMP_PATH}/input_clusthash.dbtype"; then # shellcheck disable=SC2086 - $RUNNER "$MMSEQS" clusthash "$INPUT" "${TMP_PATH}/input_clusthash" ${CLUSTHASH_PAR} \ + "$MMSEQS" clusthash "$INPUT" "${TMP_PATH}/input_clusthash" ${CLUSTHASH_PAR} \ || fail "clusthash died" fi if notExists "${TMP_PATH}/input_clusthash_clust.dbtype"; then # shellcheck disable=SC2086 - $RUNNER "$MMSEQS" clust "$INPUT" "${TMP_PATH}/input_clusthash" "${TMP_PATH}/input_clusthash_clust" ${CLUSTHASH_CLUST_PAR} \ + "$MMSEQS" clust "$INPUT" "${TMP_PATH}/input_clusthash" "${TMP_PATH}/input_clusthash_clust" ${CLUSTHASH_CLUST_PAR} \ || fail "clusthash-based clust died" fi @@ -144,13 +144,13 @@ elif [ "$LINCLUST_MODULE" = "linclust1" ]; then if [ -n "$CLUSTHASH" ]; then if notExists "${TMP_PATH}/input_clusthash.dbtype"; then # shellcheck disable=SC2086 - $RUNNER "$MMSEQS" clusthash "$INPUT" "${TMP_PATH}/input_clusthash" ${CLUSTHASH_PAR} \ + "$MMSEQS" clusthash "$INPUT" "${TMP_PATH}/input_clusthash" ${CLUSTHASH_PAR} \ || fail "clust hash died" fi if notExists "${TMP_PATH}/input_clusthash_clust.dbtype"; then # shellcheck disable=SC2086 - $RUNNER "$MMSEQS" clust "$INPUT" "${TMP_PATH}/input_clusthash" "${TMP_PATH}/input_clusthash_clust" ${CLUSTHASH_CLUST_PAR} \ + "$MMSEQS" clust "$INPUT" "${TMP_PATH}/input_clusthash" "${TMP_PATH}/input_clusthash_clust" ${CLUSTHASH_CLUST_PAR} \ || fail "clust hash based clust died" fi diff --git a/src/alignment/Align2ClustCheckpoint.cpp b/src/alignment/Align2ClustCheckpoint.cpp new file mode 100644 index 000000000..60e554994 --- /dev/null +++ b/src/alignment/Align2ClustCheckpoint.cpp @@ -0,0 +1,190 @@ +#include "Align2ClustCheckpoint.h" +#include "Parameters.h" +#include "FileUtil.h" +#include "Util.h" +#include "Debug.h" + +#include +#include +#include + +#include +#include +#include + +static std::string computeFingerprintHex(const Parameters &par, size_t dbSize, size_t endRange, + int mode, size_t chunkSize) { + // Every field below can change the *result* of clustering and therefore must be + // part of the fingerprint. Rank count, thread count, and the runner used to + // launch the job are intentionally excluded so the fingerprint (and thus the + // checkpoint directory) stays identical across topology changes. + std::ostringstream key; + key << "db1=" << par.db1 << ";db1size=" << dbSize + << ";db2=" << par.db2 << ";endRange=" << endRange + << ";mode=" << mode << ";chunkSize=" << chunkSize + << ";covMode=" << par.covMode << ";covThr=" << par.covThr + << ";seqIdThr=" << par.seqIdThr << ";seqIdMode=" << par.seqIdMode + << ";evalThr=" << par.evalThr << ";alnLenThr=" << par.alnLenThr + << ";alignmentMode=" << par.alignmentMode + << ";compBiasCorrection=" << par.compBiasCorrection + << ";compBiasCorrectionScale=" << par.compBiasCorrectionScale + << ";gapOpen=" << par.gapOpen.values.aminoacid() + << ";gapExtend=" << par.gapExtend.values.aminoacid() + << ";zdrop=" << par.zdrop + << ";scoringMatrixFile=" << par.scoringMatrixFile.values.aminoacid() + << ";addBacktrace=" << par.addBacktrace + << ";includeAlignFiles=" << par.includeAlignFiles + << ";filterCluDBFile=" << par.filterCluDBFile + << ";filterSeqDBFile=" << par.filterSeqDBFile; + + const std::string keyStr = key.str(); + const size_t hash = Util::hash(keyStr.data(), keyStr.size()); + + char buffer[32]; + snprintf(buffer, sizeof(buffer), "%016zx", hash); + return std::string(buffer); +} + +static void ensureDirExists(const std::string &dir) { + if (FileUtil::directoryExists(dir.c_str())) { + return; + } + if (FileUtil::makeDir(dir.c_str()) == false && FileUtil::directoryExists(dir.c_str()) == false) { + Debug(Debug::ERROR) << "Could not create directory " << dir << "\n"; + EXIT(EXIT_FAILURE); + } +} + +Align2ClustCheckpoint::Align2ClustCheckpoint(const Parameters &par, const std::string &baseOutputPath, + size_t dbSize, size_t endRange, int mode, size_t chunkSize) + : parentDir_(baseOutputPath + "_chunks"), + fingerprint_(computeFingerprintHex(par, dbSize, endRange, mode, chunkSize)), + chunkDir_(parentDir_ + "/" + fingerprint_) { +} + +void Align2ClustCheckpoint::ensureChunkDirExists() const { + ensureDirExists(parentDir_); + ensureDirExists(chunkDir_); +} + +std::string Align2ClustCheckpoint::dataPath(size_t chunkIndex) const { + return chunkDir_ + "/chunk_" + SSTR(chunkIndex) + ".bin"; +} + +std::string Align2ClustCheckpoint::donePath(size_t chunkIndex) const { + return chunkDir_ + "/chunk_" + SSTR(chunkIndex) + ".done"; +} + +bool Align2ClustCheckpoint::isChunkDone(size_t chunkIndex) const { + return FileUtil::fileExists(donePath(chunkIndex).c_str()); +} + +void Align2ClustCheckpoint::writeChunk(size_t chunkIndex, const std::vector &results) const { + const std::string finalPath = dataPath(chunkIndex); + const std::string tmpPath = finalPath + ".tmp"; + + FILE *fp = fopen(tmpPath.c_str(), "wb"); + if (fp == NULL) { + Debug(Debug::ERROR) << "Could not open " << tmpPath << " for writing\n"; + EXIT(EXIT_FAILURE); + } + + const uint64_t resultCount = static_cast(results.size()); + fwrite(&resultCount, sizeof(uint64_t), 1, fp); + for (const ClusterResult &result : results) { + const uint64_t sequenceIdx = static_cast(result.sequenceIdx); + const uint64_t representativeId = static_cast(result.representativeId); + const uint64_t prefSize = static_cast(result.prefSize); + const uint64_t memberCount = static_cast(result.memberIds.size()); + fwrite(&sequenceIdx, sizeof(uint64_t), 1, fp); + fwrite(&representativeId, sizeof(uint64_t), 1, fp); + fwrite(&prefSize, sizeof(uint64_t), 1, fp); + fwrite(&memberCount, sizeof(uint64_t), 1, fp); + for (DBLocalId memberId : result.memberIds) { + const uint64_t member = static_cast(memberId); + fwrite(&member, sizeof(uint64_t), 1, fp); + } + } + fflush(fp); + fsync(fileno(fp)); + fclose(fp); + + if (rename(tmpPath.c_str(), finalPath.c_str()) != 0) { + Debug(Debug::ERROR) << "Could not rename " << tmpPath << " to " << finalPath << "\n"; + EXIT(EXIT_FAILURE); + } + + // The empty ".done" marker is created only after the data file has been fully + // written and atomically renamed into place, so isChunkDone() never observes a + // partially written chunk. + FILE *doneFp = fopen(donePath(chunkIndex).c_str(), "wb"); + if (doneFp == NULL) { + Debug(Debug::ERROR) << "Could not create " << donePath(chunkIndex) << "\n"; + EXIT(EXIT_FAILURE); + } + fclose(doneFp); +} + +std::vector Align2ClustCheckpoint::readChunk(size_t chunkIndex) const { + std::vector results; + + const std::string path = dataPath(chunkIndex); + FILE *fp = fopen(path.c_str(), "rb"); + if (fp == NULL) { + Debug(Debug::ERROR) << "Could not open " << path << " for reading\n"; + EXIT(EXIT_FAILURE); + } + + uint64_t resultCount = 0; + if (fread(&resultCount, sizeof(uint64_t), 1, fp) != 1) { + Debug(Debug::ERROR) << "Could not read chunk header from " << path << "\n"; + EXIT(EXIT_FAILURE); + } + results.reserve(resultCount); + for (uint64_t i = 0; i < resultCount; ++i) { + uint64_t sequenceIdx = 0, representativeId = 0, prefSize = 0, memberCount = 0; + if (fread(&sequenceIdx, sizeof(uint64_t), 1, fp) != 1 || + fread(&representativeId, sizeof(uint64_t), 1, fp) != 1 || + fread(&prefSize, sizeof(uint64_t), 1, fp) != 1 || + fread(&memberCount, sizeof(uint64_t), 1, fp) != 1) { + Debug(Debug::ERROR) << "Truncated chunk file " << path << "\n"; + EXIT(EXIT_FAILURE); + } + + ClusterResult result; + result.sequenceIdx = static_cast(sequenceIdx); + result.representativeId = static_cast(representativeId); + result.prefSize = static_cast(prefSize); + result.memberIds.resize(memberCount); + for (uint64_t m = 0; m < memberCount; ++m) { + uint64_t member = 0; + if (fread(&member, sizeof(uint64_t), 1, fp) != 1) { + Debug(Debug::ERROR) << "Truncated chunk file " << path << "\n"; + EXIT(EXIT_FAILURE); + } + result.memberIds[m] = static_cast(member); + } + results.push_back(std::move(result)); + } + fclose(fp); + return results; +} + +void Align2ClustCheckpoint::cleanup() const { + DIR *dir = opendir(chunkDir_.c_str()); + if (dir != NULL) { + struct dirent *entry; + while ((entry = readdir(dir)) != NULL) { + const std::string name = entry->d_name; + if (name == "." || name == "..") { + continue; + } + FileUtil::remove((chunkDir_ + "/" + name).c_str()); + } + closedir(dir); + } + rmdir(chunkDir_.c_str()); + // Only removes parentDir_ if this was the last fingerprint subdirectory in it; + // harmless no-op (ENOTEMPTY) otherwise. + rmdir(parentDir_.c_str()); +} diff --git a/src/alignment/Align2ClustCheckpoint.h b/src/alignment/Align2ClustCheckpoint.h new file mode 100644 index 000000000..b06f20a00 --- /dev/null +++ b/src/alignment/Align2ClustCheckpoint.h @@ -0,0 +1,70 @@ +#ifndef MMSEQS_ALIGN2CLUST_CHECKPOINT_H +#define MMSEQS_ALIGN2CLUST_CHECKPOINT_H + +#include "Align2ClustReducer.h" + +#include +#include +#include + +class Parameters; + +// Manages on-disk chunk checkpoint state for a (potentially MPI-distributed) +// align2clust run. +// +// The checkpoint directory is named after a fingerprint computed from every +// parameter that affects the *result* of clustering (thresholds, alignment +// mode, input database identity/size, chunk size, ...). It deliberately +// excludes anything about *how* the work is distributed (rank count, thread +// count, MPI runner): those never change the fingerprint, so the exact same +// checkpoint directory -- and therefore the exact same per-chunk completion +// state -- is reused whether a run is resumed on the same topology, a +// different number of nodes, or a single node. A run interrupted on N nodes +// can always be resumed on M nodes (including M == 1) as long as the +// semantic clustering parameters are unchanged. +// +// Known limitation: input database identity is fingerprinted by path and +// sequence count, not file content checksum. Replacing db1/db2 in place with +// different content of the same size, between an interrupted run and its +// resume, will not be detected. Documented as an out-of-scope hardening item. +class Align2ClustCheckpoint { +public: + // baseOutputPath is typically par.db3 (the align2clust cluster-result DB + // path); the checkpoint directory is created as a sibling of it. + Align2ClustCheckpoint(const Parameters &par, const std::string &baseOutputPath, + size_t dbSize, size_t endRange, int mode, size_t chunkSize); + + const std::string &getChunkDir() const { return chunkDir_; } + const std::string &getFingerprint() const { return fingerprint_; } + + // Ensures the checkpoint directory (and its parent) exist on disk. + void ensureChunkDirExists() const; + + bool isChunkDone(size_t chunkIndex) const; + + // Serializes results to disk and marks the chunk done. The data file is + // written to a temporary name and rename()'d into place (atomic on a + // POSIX filesystem) before the ".done" marker is created, so a chunk is + // only ever observed as done once its data is fully and correctly on + // disk; a crash mid-write leaves no ".done" marker and the chunk is + // regenerated on the next run. + void writeChunk(size_t chunkIndex, const std::vector &results) const; + + std::vector readChunk(size_t chunkIndex) const; + + // Removes the entire checkpoint directory (and its parent, if left + // empty). Call only after every chunk has been consumed and the final + // clustering result has been written successfully -- an interrupted run + // must keep its checkpoint directory so it can be resumed. + void cleanup() const; + +private: + std::string dataPath(size_t chunkIndex) const; + std::string donePath(size_t chunkIndex) const; + + std::string parentDir_; + std::string fingerprint_; + std::string chunkDir_; +}; + +#endif diff --git a/src/alignment/Align2ClustChunking.h b/src/alignment/Align2ClustChunking.h new file mode 100644 index 000000000..aedd0ea44 --- /dev/null +++ b/src/alignment/Align2ClustChunking.h @@ -0,0 +1,56 @@ +#ifndef MMSEQS_ALIGN2CLUST_CHUNKING_H +#define MMSEQS_ALIGN2CLUST_CHUNKING_H + +#include +#include +#include + +// Deterministic, topology-independent partitioning of the [0, endRange) index +// space used by align2clust candidate generation into fixed-size chunks. +// +// Chunk boundaries depend only on endRange and chunkSize -- never on rank +// count, thread count, or which runner launched the job -- so the same chunk +// list is produced whether a run uses 1 node or 100 nodes. This is what +// allows a run to be checkpointed under one topology and resumed under a +// different one: Align2ClustCheckpoint tracks completion per chunk index, +// and chunk indices/ranges never change when the topology changes. +namespace Align2ClustChunking { + +struct Chunk { + size_t index; + size_t start; // inclusive + size_t end; // exclusive +}; + +inline std::vector computeChunks(size_t endRange, size_t chunkSize) { + std::vector chunks; + if (endRange == 0) { + return chunks; + } + chunkSize = std::max(1, chunkSize); + const size_t numChunks = (endRange + chunkSize - 1) / chunkSize; + chunks.reserve(numChunks); + for (size_t i = 0; i < numChunks; ++i) { + Chunk chunk; + chunk.index = i; + chunk.start = i * chunkSize; + chunk.end = std::min(chunk.start + chunkSize, endRange); + chunks.push_back(chunk); + } + return chunks; +} + +// Round-robin chunk ownership: chunk i is generated by rank (i % numProc). +// Recomputing this with a different numProc after a topology change simply +// reshuffles which rank owns which not-yet-completed chunk; already-completed +// chunks (tracked by Align2ClustCheckpoint) are never regenerated. +inline bool isChunkOwnedByRank(size_t chunkIndex, int rank, int numProc) { + if (numProc <= 1) { + return true; + } + return static_cast(chunkIndex % static_cast(numProc)) == rank; +} + +} // namespace Align2ClustChunking + +#endif diff --git a/src/alignment/Align2ClustReducer.cpp b/src/alignment/Align2ClustReducer.cpp new file mode 100644 index 000000000..da63bbffe --- /dev/null +++ b/src/alignment/Align2ClustReducer.cpp @@ -0,0 +1,266 @@ +#include "Align2ClustReducer.h" +#include "Util.h" + +#include + +// Compact setCoverMemberPool_ when over half of it is dead. Only offsets +// change, so heap order holds. +static const size_t ALIGN2CLUST_MIN_COMPACTION_DEAD_MEMBERS = 16 * 1024 * 1024 / sizeof(DBLocalId); + +Align2ClustReducer::Align2ClustReducer(Mode mode, size_t dbSize, size_t reorderCapacity, bool distributedCollectOnly) + : mode_(mode), dbSize_(dbSize), assignedCluster_(nullptr), + reorderCapacity_(std::max(1, reorderCapacity)), reorderBufferedCount_(0), + currentProcessPosition_(0), currentPrefSize_(0), allCalculationsDone_(false), + setCoverLiveMemberCount_(0), started_(false), finished_(false), + distributedCollectOnly_(distributedCollectOnly) { + if (distributedCollectOnly_) { + // No live assignment state is tracked in this mode (see class comment): a + // distributed worker rank has no cross-rank view of assignment progress, so + // isAssigned() always reports "unassigned" and no atomic array is needed. + return; + } + assignedCluster_ = new(std::nothrow) std::atomic[dbSize_]; + Util::checkAllocation(assignedCluster_, "Can not allocate assignedCluster memory in Align2ClustReducer"); + for (size_t i = 0; i < dbSize_; ++i) { + storeAssignedCluster(assignedCluster_, i, DB_LOCAL_ID_INVALID); + } + reorderSlots_.resize(reorderCapacity_); + reorderFilled_.assign(reorderCapacity_, 0); +} + +Align2ClustReducer::~Align2ClustReducer() { + delete[] assignedCluster_; +} + +void Align2ClustReducer::start() { + if (distributedCollectOnly_ || started_) { + started_ = true; + return; + } + started_ = true; + if (mode_ == SET_COVER) { + consumerThread_ = std::thread(&Align2ClustReducer::consumerLoopSetCover, this); + } else { + consumerThread_ = std::thread(&Align2ClustReducer::consumerLoopGreedy, this); + } +} + +void Align2ClustReducer::pushResult(ClusterResult &&result) { + if (distributedCollectOnly_) { + std::lock_guard lock(mutex_); + collected_.push_back(std::move(result)); + return; + } + + const size_t idx = result.sequenceIdx; + bool shouldNotifyClusterThread = false; + { + std::unique_lock lock(mutex_); + // Wait for this result's slot to free up. The next-in-order slot is always + // free, so the producer the consumer is waiting on never blocks (deadlock-free). + reorderSpaceCondition_.wait(lock, [&] { + return allCalculationsDone_ || idx < currentProcessPosition_ + reorderCapacity_; + }); + const size_t slot = idx % reorderCapacity_; + reorderSlots_[slot] = std::move(result); // O(1) vector move, no heap sift + reorderFilled_[slot] = 1; + reorderBufferedCount_++; + shouldNotifyClusterThread = (idx == currentProcessPosition_); + } + if (shouldNotifyClusterThread) { + clusterCondition_.notify_one(); + } +} + +void Align2ClustReducer::finish() { + if (finished_) { + return; + } + finished_ = true; + { + std::lock_guard lock(mutex_); + allCalculationsDone_ = true; + } + clusterCondition_.notify_one(); + reorderSpaceCondition_.notify_all(); + if (consumerThread_.joinable()) { + consumerThread_.join(); + } +} + +bool Align2ClustReducer::isAssigned(size_t sequenceId) const { + if (distributedCollectOnly_) { + // A distributed worker rank has no live view of cross-rank assignment + // progress; see the class comment. Always report "unassigned" so callers + // fall back to full (safe superset) candidate generation. + return false; + } + return loadAssignedCluster(assignedCluster_, sequenceId) != DB_LOCAL_ID_INVALID; +} + +void Align2ClustReducer::finalizeSingletons() { + if (distributedCollectOnly_) { + return; + } + for (size_t i = 0; i < dbSize_; ++i) { + if (loadAssignedCluster(assignedCluster_, i) == DB_LOCAL_ID_INVALID) { + storeAssignedCluster(assignedCluster_, i, i); + } + } +} + +DBLocalId Align2ClustReducer::getAssignment(size_t sequenceId) const { + return loadAssignedCluster(assignedCluster_, sequenceId); +} + +// Compact setCoverMemberPool_ when over half is dead. Only offsets change, so heap order holds. +void Align2ClustReducer::compactSetCoverMemberPool() { + const size_t deadMemberCount = setCoverMemberPool_.size() - setCoverLiveMemberCount_; + if (setCoverMemberPool_.empty() || + deadMemberCount < ALIGN2CLUST_MIN_COMPACTION_DEAD_MEMBERS || + deadMemberCount * 2 <= setCoverMemberPool_.size()) { + return; + } + std::vector compactedPool; + compactedPool.reserve(setCoverLiveMemberCount_); + for (SetCoverCandidate &candidate : setCoverCandidates_) { + const size_t newOffset = compactedPool.size(); + compactedPool.insert(compactedPool.end(), + setCoverMemberPool_.begin() + candidate.memberOffset, + setCoverMemberPool_.begin() + candidate.memberOffset + candidate.memberCount); + candidate.memberOffset = newOffset; + } + setCoverMemberPool_.swap(compactedPool); +} + +void Align2ClustReducer::consumerLoopSetCover() { + while (true) { + std::unique_lock lock(mutex_); + + clusterCondition_.wait(lock, [this] { + return reorderFilled_[currentProcessPosition_ % reorderCapacity_] != 0 || + allCalculationsDone_; + }); + + // 1) reorder buffer -> setCoverCandidates_ + while (reorderFilled_[currentProcessPosition_ % reorderCapacity_] != 0) { + const size_t slot = currentProcessPosition_ % reorderCapacity_; + ClusterResult result = std::move(reorderSlots_[slot]); + reorderFilled_[slot] = 0; + reorderBufferedCount_--; + currentProcessPosition_++; + reorderSpaceCondition_.notify_all(); + currentPrefSize_ = result.prefSize; + + if (result.memberIds.size() > 1) { + SetCoverCandidate candidate; + candidate.representativeId = result.representativeId; + candidate.memberCount = result.memberIds.size(); + candidate.memberOffset = setCoverMemberPool_.size(); + setCoverMemberPool_.insert(setCoverMemberPool_.end(), + result.memberIds.begin(), result.memberIds.end()); + setCoverLiveMemberCount_ += candidate.memberCount; + setCoverCandidates_.push_back(candidate); + std::push_heap(setCoverCandidates_.begin(), setCoverCandidates_.end(), SetCoverComparator()); + } + } + + // 2) assign candidates guaranteed to be the currently largest set + while (setCoverCandidates_.empty() == false && + (allCalculationsDone_ || + setCoverCandidates_.front().memberCount > currentPrefSize_)) { + + std::pop_heap(setCoverCandidates_.begin(), setCoverCandidates_.end(), SetCoverComparator()); + SetCoverCandidate candidate = setCoverCandidates_.back(); + setCoverCandidates_.pop_back(); + setCoverLiveMemberCount_ -= candidate.memberCount; + + if (loadAssignedCluster(assignedCluster_, candidate.representativeId) != DB_LOCAL_ID_INVALID) { + continue; + } + + // Drop already-assigned members, compacting the survivors in place + // within the pool region [memberOffset, memberOffset + memberCount). + DBLocalId *members = setCoverMemberPool_.data() + candidate.memberOffset; + size_t validCount = 0; + for (size_t i = 0; i < candidate.memberCount; i++) { + if (loadAssignedCluster(assignedCluster_, members[i]) == DB_LOCAL_ID_INVALID) { + members[validCount++] = members[i]; + } + } + + if (validCount <= 1) { + continue; + } + + if (validCount != candidate.memberCount) { + candidate.memberCount = validCount; + setCoverLiveMemberCount_ += validCount; + setCoverCandidates_.push_back(candidate); + std::push_heap(setCoverCandidates_.begin(), setCoverCandidates_.end(), SetCoverComparator()); + continue; + } + + for (size_t i = 0; i < candidate.memberCount; i++) { + storeAssignedCluster(assignedCluster_, members[i], candidate.representativeId); + } + } + + // Compaction only touches consumer-private structures (setCoverCandidates_, + // setCoverMemberPool_, setCoverLiveMemberCount_), never shared state, so release + // the mutex during the copy so worker threads can keep pushing results. + lock.unlock(); + compactSetCoverMemberPool(); + lock.lock(); + + if (allCalculationsDone_ && + reorderBufferedCount_ == 0 && + setCoverCandidates_.empty()) { + break; + } + } +} + +void Align2ClustReducer::consumerLoopGreedy() { + while (true) { + std::unique_lock lock(mutex_); + + clusterCondition_.wait(lock, [this] { + return reorderFilled_[currentProcessPosition_ % reorderCapacity_] != 0 || + allCalculationsDone_; + }); + + if (allCalculationsDone_ && reorderBufferedCount_ == 0) { + break; + } + + while (reorderFilled_[currentProcessPosition_ % reorderCapacity_] != 0) { + const size_t slot = currentProcessPosition_ % reorderCapacity_; + ClusterResult result = std::move(reorderSlots_[slot]); + reorderFilled_[slot] = 0; + reorderBufferedCount_--; + currentProcessPosition_++; + reorderSpaceCondition_.notify_all(); + + if (loadAssignedCluster(assignedCluster_, result.representativeId) != DB_LOCAL_ID_INVALID) { + continue; + } + + std::vector validMemberIds; + validMemberIds.reserve(result.memberIds.size()); + for (DBLocalId memberId : result.memberIds) { + if (loadAssignedCluster(assignedCluster_, memberId) == DB_LOCAL_ID_INVALID) { + validMemberIds.push_back(memberId); + } + } + + if (validMemberIds.size() <= 1) { + continue; + } + + for (DBLocalId memberId : validMemberIds) { + storeAssignedCluster(assignedCluster_, memberId, result.representativeId); + } + } + } +} diff --git a/src/alignment/Align2ClustReducer.h b/src/alignment/Align2ClustReducer.h new file mode 100644 index 000000000..84e8dc8ae --- /dev/null +++ b/src/alignment/Align2ClustReducer.h @@ -0,0 +1,179 @@ +#ifndef MMSEQS_ALIGN2CLUST_REDUCER_H +#define MMSEQS_ALIGN2CLUST_REDUCER_H + +#include "IndexTypes.h" + +#include +#include +#include +#include +#include + +// One generated cluster candidate for a single global order position. +// Produced by candidate generation (OpenMP workers in doAlign2clust today; +// a distributed generator in a future multi-node setting) and consumed +// strictly in `sequenceIdx` order by Align2ClustReducer. +struct ClusterResult { + size_t sequenceIdx; + DBLocalId representativeId; + size_t prefSize; + std::vector memberIds; +}; + +// Deterministic, order-preserving reducer for align2clust's SET_COVER and +// GREEDY/GREEDY_MEM clustering modes. +// +// Candidate generation may run in any order, on any number of parallel +// workers (threads today; MPI ranks in a future distributed mode). +// Correctness only requires that every global order position in +// [0, dbSize) is pushed to pushResult() exactly once. Align2ClustReducer +// reorders candidates back into sequenceIdx order internally and performs +// the actual set-cover/greedy assignment strictly in that order, so the +// final clustering never depends on generation order, thread/rank +// scheduling, worker count, or (in a future distributed setting) how many +// nodes produced the candidates. +// +// isAssigned() lets a candidate generator skip alignments it already knows +// would be discarded (a pure performance optimization backed by a live, +// process-local view of the current assignment). It is never required for +// correctness: the reducer independently re-validates every candidate's +// representative and every member against the assignment state at the time +// the candidate is actually consumed, so a generator may safely produce a +// superset of the "true" final members/candidates (for example, because it +// has no live assignment state to consult, as would be the case for a +// distributed worker on a separate node). This is the property that a +// future MPI-distributed candidate generator relies on: worker ranks can +// generate candidates independently, without sharing live assignment +// state, and still feed the same deterministic reducer used here. +// +// This class has no dependency on OpenMP, MPI, DBReader, or DBWriter so it +// can be unit tested with synthetic candidates and reused unchanged by a +// future distributed candidate consumer. +class Align2ClustReducer { +public: + enum Mode { + SET_COVER, + GREEDY + }; + + // dbSize is the total number of sequences (and the number of distinct + // sequenceIdx values that must eventually be pushed exactly once). + // reorderCapacity is the maximum out-of-order window; values less than + // 1 are clamped to 1. + // + // distributedCollectOnly selects a lightweight second mode used by MPI + // worker ranks generating candidates for a single checkpoint chunk: no + // consumer thread is started, isAssigned() always returns false (a + // worker rank has no live view of cross-rank assignment state -- see + // the class comment), and pushResult() simply appends to an internal + // vector retrievable via getCollectedResults() instead of driving the + // ordered set-cover/greedy reduction. finish()/finalizeSingletons() are + // no-ops in this mode; getAssignment() must not be called. + Align2ClustReducer(Mode mode, size_t dbSize, size_t reorderCapacity, bool distributedCollectOnly = false); + ~Align2ClustReducer(); + + Align2ClustReducer(const Align2ClustReducer &) = delete; + Align2ClustReducer &operator=(const Align2ClustReducer &) = delete; + + // Starts the internal consumer thread. Must be called exactly once, + // before any call to pushResult(). No-op in distributed-collect-only + // mode. + void start(); + + // Thread-safe; may be called concurrently by any number of worker + // threads. Blocks if the reorder window is full until the reducer has + // consumed enough in-order results to make room. Must not be called + // after finish(). In distributed-collect-only mode, simply appends + // under a mutex and never blocks. + void pushResult(ClusterResult &&result); + + // Signals that no further results will be pushed for this run, then + // waits for the internal consumer thread to finish and joins it. Must + // be called exactly once, after every pushResult() call has returned. + // No-op in distributed-collect-only mode. + void finish(); + + // True if sequenceId already has a representative assigned. See the + // class comment: this is a pure performance optimization for candidate + // generation and is never required for correctness. Always false in + // distributed-collect-only mode. + bool isAssigned(size_t sequenceId) const; + + // Must be called after finish(). Assigns every still-unassigned + // sequence to itself (singleton cluster). Must not be called in + // distributed-collect-only mode. + void finalizeSingletons(); + + // Valid only after finalizeSingletons(); never DB_LOCAL_ID_INVALID. + DBLocalId getAssignment(size_t sequenceId) const; + + // Valid only in distributed-collect-only mode, after finish(). + const std::vector &getCollectedResults() const { return collected_; } + + size_t getDbSize() const { return dbSize_; } + +private: + struct SetCoverCandidate { + DBLocalId representativeId; + size_t memberCount; + size_t memberOffset; + }; + + struct SetCoverComparator { + bool operator()(const SetCoverCandidate &a, const SetCoverCandidate &b) const { + if (a.memberCount < b.memberCount) { + return true; + } + if (b.memberCount < a.memberCount) { + return false; + } + if (a.representativeId < b.representativeId) { + return true; + } + if (b.representativeId < a.representativeId) { + return false; + } + return false; + } + }; + + void consumerLoopSetCover(); + void consumerLoopGreedy(); + void compactSetCoverMemberPool(); + + static DBLocalId loadAssignedCluster(const std::atomic *assignedCluster, size_t sequenceId) { + return assignedCluster[sequenceId].load(std::memory_order_relaxed); + } + static void storeAssignedCluster(std::atomic *assignedCluster, size_t sequenceId, DBLocalId representativeId) { + assignedCluster[sequenceId].store(representativeId, std::memory_order_relaxed); + } + + Mode mode_; + size_t dbSize_; + std::atomic *assignedCluster_; + + mutable std::mutex mutex_; + std::condition_variable clusterCondition_; + std::condition_variable reorderSpaceCondition_; + + std::vector reorderSlots_; + std::vector reorderFilled_; + size_t reorderCapacity_; + size_t reorderBufferedCount_; + + size_t currentProcessPosition_; + size_t currentPrefSize_; + bool allCalculationsDone_; + + std::vector setCoverCandidates_; + std::vector setCoverMemberPool_; + size_t setCoverLiveMemberCount_; + + std::thread consumerThread_; + bool started_; + bool finished_; + bool distributedCollectOnly_; + std::vector collected_; +}; + +#endif diff --git a/src/alignment/Align2clust.cpp b/src/alignment/Align2clust.cpp index d6308dc18..2f249dc62 100644 --- a/src/alignment/Align2clust.cpp +++ b/src/alignment/Align2clust.cpp @@ -11,6 +11,10 @@ #include "BlockAligner.h" #include "Alignment.h" #include "AlignmentSymmetry.h" +#include "Align2ClustReducer.h" +#include "Align2ClustChunking.h" +#include "Align2ClustCheckpoint.h" +#include "MMseqsMPI.h" #include #include #include @@ -24,13 +28,6 @@ #define MIN_SIZE 32 -struct ClusterResult { - size_t sequenceIdx; - DBLocalId representativeId; - size_t prefSize; - std::vector memberIds; -}; - struct PrefInfo { DBLocalId id; size_t size; @@ -48,92 +45,12 @@ struct PrefInfo { } }; -// Lightweight entry stored in the set-cover ready queue. Instead of carrying the -// member id vector (as ClusterResult does) it only references the members, which -// are kept in a single shared pool (setCoverMemberPool), by offset and count. -struct SetCoverCandidate { - DBLocalId representativeId; - size_t memberCount; - size_t memberOffset; -}; - -struct SetCoverComparator { - bool operator()(const SetCoverCandidate& a, const SetCoverCandidate& b) const { - if (a.memberCount < b.memberCount) { - return true; - } - if (b.memberCount < a.memberCount) { - return false; - } - if (a.representativeId < b.representativeId) { - return true; - } - if (b.representativeId < a.representativeId) { - return false; - } - return false; - } -}; - -static std::mutex clusterMutex; -static std::condition_variable clusterCondition; -static std::condition_variable reorderSpaceCondition; -// Reorders out-of-order worker results back to sequenceIdx order for the consumer, -// indexed by sequenceIdx % reorderCapacity. The next-in-order slot is always free, -// so producers never deadlock. -static std::vector reorderSlots; // slot storage, size == reorderCapacity -static std::vector reorderFilled; // 1 == slot holds an unconsumed result -static size_t reorderCapacity = 0; // max out-of-order window -static size_t reorderBufferedCount = 0; // number of filled slots -// Max-heap of candidates by member count (plain vector so the pool stays compactable). -static std::vector setCoverCandidates; -static SetCoverComparator setCoverComparator; -// Shared backing store for all member ids referenced by setCoverCandidates. -static std::vector setCoverMemberPool; -static size_t setCoverLiveMemberCount = 0; - -static size_t currentProcessPosition = 0; -static size_t currentPrefSize = 0; -static bool allCalculationsDone = false; - -typedef std::atomic ClusterAssignment; - -static DBLocalId loadAssignedCluster(const ClusterAssignment *assignedCluster, size_t sequenceId) { - return assignedCluster[sequenceId].load(std::memory_order_relaxed); -} - -static void storeAssignedCluster(ClusterAssignment *assignedCluster, size_t sequenceId, DBLocalId representativeId) { - assignedCluster[sequenceId].store(representativeId, std::memory_order_relaxed); -} - // Serialize a single alignment record and append it to the per-representative buffer. static void appendAlignmentResult(std::string &alnResultBuffer, char *lineBuffer, const Matcher::result_t &result, bool addBacktrace) { size_t len = Matcher::resultToBuffer(lineBuffer, result, addBacktrace); alnResultBuffer.append(lineBuffer, len); } -// Compact setCoverMemberPool when over half is dead. Only offsets change, so heap order holds. -static const size_t ALIGN2CLUST_MIN_COMPACTION_DEAD_MEMBERS = 16 * 1024 * 1024 / sizeof(DBLocalId); - -static void compactSetCoverMemberPool() { - const size_t deadMemberCount = setCoverMemberPool.size() - setCoverLiveMemberCount; - if (setCoverMemberPool.empty() || - deadMemberCount < ALIGN2CLUST_MIN_COMPACTION_DEAD_MEMBERS || - deadMemberCount * 2 <= setCoverMemberPool.size()) { - return; - } - std::vector compactedPool; - compactedPool.reserve(setCoverLiveMemberCount); - for (SetCoverCandidate &candidate : setCoverCandidates) { - const size_t newOffset = compactedPool.size(); - compactedPool.insert(compactedPool.end(), - setCoverMemberPool.begin() + candidate.memberOffset, - setCoverMemberPool.begin() + candidate.memberOffset + candidate.memberCount); - candidate.memberOffset = newOffset; - } - setCoverMemberPool.swap(compactedPool); -} - // Env override for the reorder-buffer size; 0 (unset/invalid) means size from memory. static size_t getReorderBufferLimitFromEnv() { const char *envValue = getenv("MMSEQS_ALIGN2CLUST_REORDER_LIMIT"); @@ -151,13 +68,34 @@ static size_t getReorderBufferLimitFromEnv() { return static_cast(parsedValue); } +// Env override for the MPI/checkpoint chunk size; 0 (unset/invalid) means use the default. +// Deliberately not derived from rank or thread count: chunk boundaries must stay identical +// across topology changes so a run can be resumed with a different node count (see +// Align2ClustChunking and Align2ClustCheckpoint). +static size_t getAlign2ClustChunkSizeFromEnv() { + const size_t defaultChunkSize = 50000; + const char *envValue = getenv("MMSEQS_ALIGN2CLUST_CHUNK_SIZE"); + if (envValue == nullptr || *envValue == '\0') { + return defaultChunkSize; + } + + char *end = nullptr; + unsigned long long parsedValue = strtoull(envValue, &end, 10); + if (end == envValue || *end != '\0' || parsedValue == 0) { + Debug(Debug::WARNING) << "Ignoring invalid MMSEQS_ALIGN2CLUST_CHUNK_SIZE=" << envValue + << "; using default chunk size " << defaultChunkSize << "\n"; + return defaultChunkSize; + } + return static_cast(parsedValue); +} + // Size the reorder buffer from free memory: subtract the resident per-sequence arrays, // keep 10% headroom, divide by the worst-case result size. Capped by // MMSEQS_ALIGN2CLUST_REORDER_LIMIT and by the number of results produced. static size_t computeReorderCapacity(const Parameters &par, size_t dbSize, int mode, size_t resultCount) { const size_t memoryLimit = Util::computeMemory(par.splitMemoryLimit); - size_t fixedMemory = dbSize * sizeof(ClusterAssignment); + size_t fixedMemory = dbSize * sizeof(std::atomic); if (mode == Parameters::SET_COVER) { fixedMemory += dbSize * sizeof(PrefInfo); } @@ -182,27 +120,6 @@ static size_t computeReorderCapacity(const Parameters &par, size_t dbSize, int m return capacity; } -static void pushClusterResult(ClusterResult &&clusterResult) { - const size_t idx = clusterResult.sequenceIdx; - bool shouldNotifyClusterThread = false; - { - std::unique_lock lock(clusterMutex); - // Wait for this result's slot to free up. The next-in-order slot is always - // free, so the producer the consumer is waiting on never blocks (deadlock-free). - reorderSpaceCondition.wait(lock, [&] { - return allCalculationsDone || idx < currentProcessPosition + reorderCapacity; - }); - const size_t slot = idx % reorderCapacity; - reorderSlots[slot] = std::move(clusterResult); // O(1) vector move, no heap sift - reorderFilled[slot] = 1; - reorderBufferedCount++; - shouldNotifyClusterThread = (idx == currentProcessPosition); - } - if (shouldNotifyClusterThread) { - clusterCondition.notify_one(); - } -} - static float parsePrecisionLib(const std::string &scoreFile, double targetSeqid, double targetCov, double targetPrecision) { std::stringstream in(scoreFile); std::string line; @@ -262,141 +179,11 @@ static void writeClustering(DBWriter *dbWriter, const std::pair lock(clusterMutex); - - clusterCondition.wait(lock, [] { - return reorderFilled[currentProcessPosition % reorderCapacity] != 0 || - allCalculationsDone; - }); - - // 1) reorder buffer → setCoverCandidates - while (reorderFilled[currentProcessPosition % reorderCapacity] != 0) { - const size_t slot = currentProcessPosition % reorderCapacity; - ClusterResult result = std::move(reorderSlots[slot]); - reorderFilled[slot] = 0; - reorderBufferedCount--; - currentProcessPosition++; - reorderSpaceCondition.notify_all(); - currentPrefSize = result.prefSize; - - if (result.memberIds.size() > 1) { - SetCoverCandidate candidate; - candidate.representativeId = result.representativeId; - candidate.memberCount = result.memberIds.size(); - candidate.memberOffset = setCoverMemberPool.size(); - setCoverMemberPool.insert(setCoverMemberPool.end(), - result.memberIds.begin(), result.memberIds.end()); - setCoverLiveMemberCount += candidate.memberCount; - setCoverCandidates.push_back(candidate); - std::push_heap(setCoverCandidates.begin(), setCoverCandidates.end(), setCoverComparator); - } - } - - // 2) assign candidates guaranteed to be the currently largest set - while (setCoverCandidates.empty() == false && - (allCalculationsDone || - setCoverCandidates.front().memberCount > currentPrefSize)) { - - std::pop_heap(setCoverCandidates.begin(), setCoverCandidates.end(), setCoverComparator); - SetCoverCandidate candidate = setCoverCandidates.back(); - setCoverCandidates.pop_back(); - setCoverLiveMemberCount -= candidate.memberCount; - - if (loadAssignedCluster(assignedCluster, candidate.representativeId) != DB_LOCAL_ID_INVALID) { - continue; - } - - // Drop already-assigned members, compacting the survivors in place - // within the pool region [memberOffset, memberOffset + memberCount). - DBLocalId *members = setCoverMemberPool.data() + candidate.memberOffset; - size_t validCount = 0; - for (size_t i = 0; i < candidate.memberCount; i++) { - if (loadAssignedCluster(assignedCluster, members[i]) == DB_LOCAL_ID_INVALID) { - members[validCount++] = members[i]; - } - } - - if (validCount <= 1) { - continue; - } - - if (validCount != candidate.memberCount) { - candidate.memberCount = validCount; - setCoverLiveMemberCount += validCount; - setCoverCandidates.push_back(candidate); - std::push_heap(setCoverCandidates.begin(), setCoverCandidates.end(), setCoverComparator); - continue; - } - - for (size_t i = 0; i < candidate.memberCount; i++) { - storeAssignedCluster(assignedCluster, members[i], candidate.representativeId); - } - } - - // Compaction only touches consumer-private structures (setCoverCandidates, - // setCoverMemberPool, setCoverLiveMemberCount), never shared state, so release - // the mutex during the copy so worker threads can keep pushing results. - lock.unlock(); - compactSetCoverMemberPool(); - lock.lock(); - - if (allCalculationsDone && - reorderBufferedCount == 0 && - setCoverCandidates.empty()) { - break; - } - } -} - -void clusterThreadFuncGreedy(ClusterAssignment* assignedCluster) { - while (true) { - std::unique_lock lock(clusterMutex); - - clusterCondition.wait(lock, [] { - return reorderFilled[currentProcessPosition % reorderCapacity] != 0 || - allCalculationsDone; - }); - - if (allCalculationsDone && reorderBufferedCount == 0) { - break; - } - - while (reorderFilled[currentProcessPosition % reorderCapacity] != 0) { - const size_t slot = currentProcessPosition % reorderCapacity; - ClusterResult result = std::move(reorderSlots[slot]); - reorderFilled[slot] = 0; - reorderBufferedCount--; - currentProcessPosition++; - reorderSpaceCondition.notify_all(); - - if (loadAssignedCluster(assignedCluster, result.representativeId) != DB_LOCAL_ID_INVALID) { - continue; - } - - std::vector validMemberIds; - validMemberIds.reserve(result.memberIds.size()); - for (DBLocalId memberId : result.memberIds) { - if (loadAssignedCluster(assignedCluster, memberId) == DB_LOCAL_ID_INVALID) { - validMemberIds.push_back(memberId); - } - } - - if (validMemberIds.size() <= 1) { - continue; - } - - for (DBLocalId memberId : validMemberIds) { - storeAssignedCluster(assignedCluster, memberId, result.representativeId); - } - } - } -} - -int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader &alnDbr, DBWriter *alnWriter) { +// resultWriter is nullptr on non-master MPI ranks: only rank 0 performs the global +// reduction and writes the final clustering result (see the useDistributedGeneration +// branch below and the Align2ClustReducer class comment for why this cannot itself +// be decentralized). +int doAlign2clust(Parameters &par, DBWriter *resultWriter, DBReader &alnDbr, DBWriter *alnWriter) { DBReader *seqDbr = new DBReader( par.db1.c_str(), par.db1Index.c_str(), par.threads, DBReader::USE_DATA | DBReader::USE_INDEX @@ -442,24 +229,18 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & EvalueComputation evaluer(seqDbr->getAminoAcidDBSize(), subMat); int32_t xDrop = (MIN_SIZE * par.gapExtend.values.aminoacid() + par.gapOpen.values.aminoacid()); - - ClusterAssignment *assignedCluster = new(std::nothrow) ClusterAssignment[dbSize]; - Util::checkAllocation(assignedCluster, "Can not allocate assignedCluster memory in Align2Clust"); - for (size_t i = 0; i < dbSize; ++i) { - storeAssignedCluster(assignedCluster, i, DB_LOCAL_ID_INVALID); - } int mode = par.clusteringMode; - + + Align2ClustReducer::Mode reducerMode; if (mode == Parameters::SET_COVER) { - clusterThreadFunc = clusterThreadFuncSetcover; + reducerMode = Align2ClustReducer::SET_COVER; Debug(Debug::INFO) << "Using SET_COVER clustering mode\n"; } else if (mode == Parameters::GREEDY || mode == Parameters::GREEDY_MEM) { - clusterThreadFunc = clusterThreadFuncGreedy; + reducerMode = Align2ClustReducer::GREEDY; Debug(Debug::INFO) << "Using GREEDY clustering mode\n"; } else { Debug(Debug::ERROR) << "MMseqs2 align2clust doesn't support clustering mode: " << mode << "\n"; - delete[] assignedCluster; delete[] fastMatrix.matrix; delete[] fastMatrix.matrixData; delete subMat; @@ -468,30 +249,67 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & return EXIT_FAILURE; } + const size_t endRange = (mode == Parameters::SET_COVER) ? dbSize : alnDbr.getSize(); + // Ring size = out-of-order window, sized from the memory budget (OOM-aware) and // capped by the result count. sequenceIdx runs over [0, endRange); every index // publishes exactly one result. - const size_t align2clustResultCount = (mode == Parameters::SET_COVER) ? dbSize : alnDbr.getSize(); - const size_t reorderCapacityChosen = computeReorderCapacity(par, dbSize, mode, align2clustResultCount); - { - std::lock_guard lock(clusterMutex); - reorderCapacity = reorderCapacityChosen; - reorderSlots.clear(); - reorderSlots.resize(reorderCapacity); - reorderFilled.assign(reorderCapacity, 0); - reorderBufferedCount = 0; - setCoverCandidates.clear(); - setCoverCandidates.shrink_to_fit(); - setCoverMemberPool.clear(); - setCoverMemberPool.shrink_to_fit(); - setCoverLiveMemberCount = 0; - currentProcessPosition = 0; - currentPrefSize = 0; - allCalculationsDone = false; + const size_t reorderCapacityChosen = computeReorderCapacity(par, dbSize, mode, endRange); + + int mpiRank = 0; + int mpiNumProc = 1; +#ifdef HAVE_MPI + mpiRank = MMseqsMPI::rank; + mpiNumProc = MMseqsMPI::numProc; +#endif + + // Deterministic, topology-independent chunking of [0, endRange) for MPI-distributed + // candidate generation and checkpoint/resume. Chunk boundaries depend only on + // endRange and chunkSize -- never on rank or thread count -- so a run interrupted + // on N nodes can always be resumed on M nodes (including M == 1); see + // Align2ClustChunking and Align2ClustCheckpoint. + const size_t chunkSize = getAlign2ClustChunkSizeFromEnv(); + std::vector chunks = Align2ClustChunking::computeChunks(endRange, chunkSize); + Align2ClustCheckpoint checkpoint(par, par.db3, dbSize, endRange, mode, chunkSize); + + bool anyChunkAlreadyDone = false; + for (const Align2ClustChunking::Chunk &chunk : chunks) { + if (checkpoint.isChunkDone(chunk.index)) { + anyChunkAlreadyDone = true; + break; + } + } + // Every rank evaluates mpiNumProc/anyChunkAlreadyDone identically (shared + // filesystem, same CLI parameters), so this decision is consistent across ranks. + const bool useDistributedGeneration = (mpiNumProc > 1) || anyChunkAlreadyDone; + + if (useDistributedGeneration && par.includeAlignFiles) { + Debug(Debug::ERROR) << "align2clust: --include-align-files is not supported together with " + << "MPI-distributed candidate generation (numProc > 1) or a resumed " + << "checkpointed run. Re-run on a single rank with the checkpoint " + << "directory (" << checkpoint.getChunkDir() << ") removed, or without " + << "--include-align-files.\n"; + delete[] fastMatrix.matrix; + delete[] fastMatrix.matrixData; + delete subMat; + seqDbr->close(); + delete seqDbr; + if (cluDbr != nullptr) { + cluDbr->close(); + delete cluDbr; + } + if (cluSeqDbr != nullptr) { + cluSeqDbr->close(); + delete cluSeqDbr; + } + return EXIT_FAILURE; + } + + Align2ClustReducer reducer(reducerMode, dbSize, reorderCapacityChosen); + if (mpiRank == 0) { + reducer.start(); } - std::thread clusterThread(clusterThreadFunc, assignedCluster); - Timer timer; timer.reset(); PrefInfo *prefRepSizePair = nullptr; @@ -521,12 +339,19 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & timer.reset(); - size_t endRange = (mode == Parameters::SET_COVER) ? dbSize : alnDbr.getSize(); unsigned int swMode = Alignment::initSWMode(par.alignmentMode, par.covThr, par.seqIdThr); - Debug::Progress progress(endRange); size_t db_maxseqlen = (cluSeqDbr != nullptr) ? std::max(seqDbr->getMaxSeqLen(), cluSeqDbr->getMaxSeqLen()) : seqDbr->getMaxSeqLen(); + + // Candidate generation for [rangeStart, rangeEnd), pushing every produced result to + // `sink`. Parameterized so it can run once over the full [0, endRange) range against + // the live `reducer` (single-process/fresh-run fast path, byte-identical to the + // pre-MPI behavior) or once per not-yet-done checkpoint chunk against a fresh + // distributed-collect-only reducer whose buffered results are written to a chunk + // file (see useDistributedGeneration below). + auto generateCandidates = [&](Align2ClustReducer &sink, size_t rangeStart, size_t rangeEnd) { + Debug::Progress progress(rangeEnd - rangeStart); #pragma omp parallel { unsigned int threadIdx = 0; @@ -556,7 +381,7 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & } #pragma omp for schedule(dynamic, 1) nowait - for (size_t i = 0; i < endRange; i++) { + for (size_t i = rangeStart; i < rangeEnd; i++) { progress.updateProgress(); ClusterResult clusterResult; clusterResult.sequenceIdx = i; @@ -583,8 +408,8 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & // by the cluster thread anyway, so skip parsing and aligning it entirely. prefSize // is already set (precomputed for set-cover), so the currentPrefSize gate stays // correct. - if (loadAssignedCluster(assignedCluster, representativeId) != DB_LOCAL_ID_INVALID) { - pushClusterResult(std::move(clusterResult)); + if (sink.isAssigned(representativeId)) { + sink.pushResult(std::move(clusterResult)); continue; } @@ -604,7 +429,7 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & targetsWithDiagonal.push_back(std::make_pair(hit.seqId, hit.diagonal)); } else { const size_t targetId = seqDbr->getId(hit.seqId); - if (loadAssignedCluster(assignedCluster, targetId) == DB_LOCAL_ID_INVALID) { + if (!sink.isAssigned(targetId)) { targetsWithDiagonal.push_back(std::make_pair(hit.seqId, hit.diagonal)); } } @@ -619,7 +444,7 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & // never used; stop aligning the rest. Safe in set-cover too: prefSize is // already fully counted, so this is just the k-th-target form of the // (k=0) rep-skip above. - if (loadAssignedCluster(assignedCluster, representativeId) != DB_LOCAL_ID_INVALID) { + if (sink.isAssigned(representativeId)) { break; } @@ -645,7 +470,7 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & // Skip the (expensive) alignment if the target was assigned meanwhile. // Safe in set-cover too: an assigned target is monotonic, so it would be // dropped by the cluster thread's re-evaluation anyway. - if (loadAssignedCluster(assignedCluster, targetId) != DB_LOCAL_ID_INVALID) { + if (sink.isAssigned(targetId)) { continue; } @@ -675,10 +500,10 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & } bool hasSeqId = seqId >= (par.seqIdThr - std::numeric_limits::epsilon()); - if (loadAssignedCluster(assignedCluster, targetId) != DB_LOCAL_ID_INVALID) continue; + if (sink.isAssigned(targetId)) continue; if (hasAlnLen && hasCoverage && hasSeqId && hasEvalue) { - if (loadAssignedCluster(assignedCluster, targetId) != DB_LOCAL_ID_INVALID) continue; + if (sink.isAssigned(targetId)) continue; if (par.filterCluDBFile.empty()== false && par.filterSeqDBFile.empty()== false){ // check all the member from filtering file const size_t cluId = cluDbr->getId(targetKey); @@ -782,7 +607,7 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & continue; } - if (loadAssignedCluster(assignedCluster, targetId) != DB_LOCAL_ID_INVALID) continue; + if (sink.isAssigned(targetId)) continue; bool foundConsecutiveMatchSeed = false; for (int blockIdx = 0; blockIdx <= alignmentLength - 3; ++blockIdx) { @@ -815,7 +640,7 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & ); if (Alignment::checkCriteria(result, isIdentity, par.evalThr, par.seqIdThr, par.alnLenThr, par.covMode, par.covThr)) { - if (loadAssignedCluster(assignedCluster, targetId) != DB_LOCAL_ID_INVALID) continue; + if (sink.isAssigned(targetId)) continue; if (par.filterCluDBFile.empty()== false && par.filterSeqDBFile.empty()== false){ // check all the member from filtering file const size_t cluId = cluDbr->getId(targetKey); @@ -902,34 +727,94 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & if (includeAlignFiles) { alnWriter->writeData(alnResultBuffer.c_str(), alnResultBuffer.length(), queryKey, threadIdx); } - pushClusterResult(std::move(clusterResult)); + sink.pushResult(std::move(clusterResult)); } } + }; // end generateCandidates lambda - { - std::lock_guard lock(clusterMutex); - allCalculationsDone = true; - } - clusterCondition.notify_one(); - reorderSpaceCondition.notify_all(); - - if (clusterThread.joinable()) { - clusterThread.join(); + if (useDistributedGeneration == false) { + // Fast path: numProc == 1 and no pre-existing checkpoint chunks for this + // fingerprint. Behavior is unchanged from before Phase 3: one pass over the + // full range straight into the live reducer, no chunk files written or read. + generateCandidates(reducer, 0, endRange); + } else { + checkpoint.ensureChunkDirExists(); + for (const Align2ClustChunking::Chunk &chunk : chunks) { + if (Align2ClustChunking::isChunkOwnedByRank(chunk.index, mpiRank, mpiNumProc) == false) { + continue; + } + if (checkpoint.isChunkDone(chunk.index)) { + continue; + } + // A worker rank has no live view of cross-rank assignment state, so it + // generates into a fresh distributed-collect-only reducer (no consumer + // thread, isAssigned() always false -- see the Align2ClustReducer class + // comment) and simply buffers whatever candidates it produces. The + // reducer that eventually runs on rank 0 re-validates every candidate + // against the true assignment state at consumption time, so this + // superset of candidates is always safe. + Align2ClustReducer collector(reducerMode, dbSize, 1, /*distributedCollectOnly=*/true); + collector.start(); + generateCandidates(collector, chunk.start, chunk.end); + collector.finish(); + checkpoint.writeChunk(chunk.index, collector.getCollectedResults()); + } + +#ifdef HAVE_MPI + if (mpiNumProc > 1) { + MPI_Barrier(MPI_COMM_WORLD); + } +#endif + + if (mpiRank == 0) { + // Only rank 0 re-reads the chunks (in global chunk order, which matches + // ascending sequenceIdx order since chunks are contiguous ranges) and + // feeds them into the single live reducer that performs the ordered + // set-cover/greedy reduction. + for (const Align2ClustChunking::Chunk &chunk : chunks) { + std::vector results = checkpoint.readChunk(chunk.index); + for (ClusterResult &result : results) { + reducer.pushResult(std::move(result)); + } + } + } } - for (size_t i = 0; i < dbSize; ++i) { - if (loadAssignedCluster(assignedCluster, i) == DB_LOCAL_ID_INVALID) { - storeAssignedCluster(assignedCluster, i, i); + if (mpiRank != 0) { + // Non-master ranks only ever generate/checkpoint candidates; the global + // reduction and final DB write happen exclusively on rank 0 (see the + // Align2ClustReducer class comment for why this step cannot itself be + // decentralized). resultWriter/alnWriter are nullptr here (see align2clust()). + delete[] fastMatrix.matrix; + delete[] fastMatrix.matrixData; + delete subMat; + if (prefRepSizePair != nullptr) { + delete[] prefRepSizePair; + } + seqDbr->close(); + delete seqDbr; + if (cluDbr != nullptr) { + cluDbr->close(); + delete cluDbr; } + if (cluSeqDbr != nullptr) { + cluSeqDbr->close(); + delete cluSeqDbr; + } + return EXIT_SUCCESS; } + reducer.finish(); + + reducer.finalizeSingletons(); + std::pair *assignment = new std::pair[dbSize]; #pragma omp parallel { #pragma omp for schedule(static) for (size_t i = 0; i < dbSize; i++) { - const DBLocalId representativeId = loadAssignedCluster(assignedCluster, i); + const DBLocalId representativeId = reducer.getAssignment(i); if (representativeId == DB_LOCAL_ID_INVALID) { Debug(Debug::ERROR) << "There must be an error: " << i << " is not assigned to a cluster\n"; @@ -951,9 +836,15 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & Debug(Debug::INFO) << "Size of the alignment database: " << dbSize << "\n"; Debug(Debug::INFO) << "Number of clusters: " << clusterCount << "\n"; - writeClustering(&resultWriter, assignment, dbSize); - - delete[] assignedCluster; + writeClustering(resultWriter, assignment, dbSize); + + if (useDistributedGeneration) { + // Safe to remove only now that the final clustering result has been fully + // written; an interrupted run must keep its checkpoint directory intact so it + // can be resumed (with any node/rank count) rather than recomputing candidates. + checkpoint.cleanup(); + } + delete[] assignment; if (prefRepSizePair != nullptr) { delete[] prefRepSizePair; @@ -976,6 +867,8 @@ int doAlign2clust(Parameters &par, DBWriter &resultWriter, DBReader & } int align2clust(int argc, const char **argv, const Command &command) { + MMseqsMPI::init(argc, argv); + Parameters &par = Parameters::getInstance(); par.parseParameters(argc, argv, command, true, 0, 0); @@ -987,11 +880,10 @@ int align2clust(int argc, const char **argv, const Command &command) { alnDbr.open(DBReader::LINEAR_ACCCESS); int dbtype = Parameters::DBTYPE_CLUSTER_RES; - DBWriter resultWriter(par.db3.c_str(), par.db3Index.c_str(), 1, par.compressed, dbtype); - resultWriter.open(); - // Optional alignment-result output; path derived from the cluster DB (db3 + "_aln"). // Alignment output needs a backtrace (-a) and score+cov+seqid so member records carry a CIGAR. + // Checked on every rank (identical parameters everywhere) so a bad flag combination + // fails the same way on every rank rather than only on rank 0. if (par.includeAlignFiles) { const unsigned int effectiveSwMode = Alignment::initSWMode(par.alignmentMode, par.covThr, par.seqIdThr); if (par.addBacktrace == false || effectiveSwMode != Matcher::SCORE_COV_SEQID) { @@ -1001,19 +893,33 @@ int align2clust(int argc, const char **argv, const Command &command) { EXIT(EXIT_FAILURE); } } + + // Only rank 0 opens the output DBs: on a shared filesystem every other rank would + // otherwise try to create/write the same output paths concurrently. doAlign2clust + // never dereferences resultWriter/alnWriter on non-master ranks (see its + // "mpiRank != 0" early return). + DBWriter *resultWriter = nullptr; DBWriter *alnWriter = nullptr; - if (par.includeAlignFiles) { - std::string alnDb = par.db3 + "_aln"; - std::string alnDbIndex = alnDb + ".index"; - alnWriter = new DBWriter(alnDb.c_str(), alnDbIndex.c_str(), par.threads, par.compressed, Parameters::DBTYPE_ALIGNMENT_RES); - alnWriter->open(); + if (MMseqsMPI::isMaster()) { + resultWriter = new DBWriter(par.db3.c_str(), par.db3Index.c_str(), 1, par.compressed, dbtype); + resultWriter->open(); + + if (par.includeAlignFiles) { + std::string alnDb = par.db3 + "_aln"; + std::string alnDbIndex = alnDb + ".index"; + alnWriter = new DBWriter(alnDb.c_str(), alnDbIndex.c_str(), par.threads, par.compressed, Parameters::DBTYPE_ALIGNMENT_RES); + alnWriter->open(); + } } int status = doAlign2clust(par, resultWriter, alnDbr, alnWriter); Debug(Debug::INFO) << "Time for run Align2Clust: " << timer.lap() << " sec\n"; - resultWriter.close(); + if (resultWriter != nullptr) { + resultWriter->close(); + delete resultWriter; + } if (alnWriter != nullptr) { alnWriter->close(); delete alnWriter; diff --git a/src/alignment/CMakeLists.txt b/src/alignment/CMakeLists.txt index 3ecd39d95..9ca5c2dd6 100644 --- a/src/alignment/CMakeLists.txt +++ b/src/alignment/CMakeLists.txt @@ -11,12 +11,17 @@ set(alignment_header_files alignment/BlockAligner.h alignment/DistanceCalculator.h alignment/Fwbw.h + alignment/Align2ClustReducer.h + alignment/Align2ClustChunking.h + alignment/Align2ClustCheckpoint.h PARENT_SCOPE ) set(alignment_source_files alignment/Alignment.cpp alignment/Align2clust.cpp + alignment/Align2ClustReducer.cpp + alignment/Align2ClustCheckpoint.cpp alignment/CompressedA3M.cpp alignment/Main.cpp alignment/Matcher.cpp diff --git a/src/commons/MMseqsMPI.cpp b/src/commons/MMseqsMPI.cpp index 4b55e1689..2f44e6862 100644 --- a/src/commons/MMseqsMPI.cpp +++ b/src/commons/MMseqsMPI.cpp @@ -1,6 +1,10 @@ #include "MMseqsMPI.h" #include "Debug.h" #include "Parameters.h" +#include "Util.h" + +#include +#include bool MMseqsMPI::active = false; int MMseqsMPI::rank = -1; @@ -24,7 +28,50 @@ void MMseqsMPI::init(int argc, const char **argv) { Debug(Debug::INFO) << "Rank: " << rank << " Size: " << numProc << "\n"; } #else +// Best-effort fail-fast: this binary was built without MPI (-DHAVE_MPI=0/unset), so it +// can only ever act as a single rank/process with no awareness of any other copies of +// itself. Launching it under a multi-task runner (srun -n>1, mpirun -np>1, ...) would +// silently run every task as an independent "rank 0 of 1" that reads the full input and +// writes to the same output paths -- exactly the unsafe scenario multi-node support is +// meant to avoid. There is no fully portable way to detect this from inside a plain +// (non-MPI) process, so this only checks common launcher environment variables set by +// OpenMPI, MPICH-family MPI implementations, and native (non-MPI) Slurm process launch; +// it is a safety net, not a guarantee. +static int getEnvTaskCount(const char *name) { + const char *value = getenv(name); + if (value == NULL || *value == '\0') { + return -1; + } + char *end = NULL; + long parsed = strtol(value, &end, 10); + if (end == value || *end != '\0' || parsed <= 0) { + return -1; + } + return static_cast(parsed); +} + void MMseqsMPI::init(int, const char **) { rank = 0; + + const char *taskCountVars[] = { + "OMPI_COMM_WORLD_SIZE", // OpenMPI + "PMI_SIZE", // MPICH / Intel MPI / Hydra + "MPI_LOCALNRANKS", // some MPICH-family launchers + "SLURM_NTASKS", // native srun (no MPI library involved) + "SLURM_NPROCS" // legacy alias for SLURM_NTASKS + }; + for (size_t i = 0; i < sizeof(taskCountVars) / sizeof(taskCountVars[0]); ++i) { + const int taskCount = getEnvTaskCount(taskCountVars[i]); + if (taskCount > 1) { + Debug(Debug::ERROR) << "This mmseqs binary was built without MPI support, but appears to have " + << "been launched with " << taskCountVars[i] << "=" << taskCount + << " (multiple tasks/ranks). Every task would independently process the " + << "full input and write to the same output paths, corrupting the result.\n" + << "Rebuild with -DHAVE_MPI=1 to use --mpi-runner/RUNNER with more than one " + << "task, or launch this binary with a single task.\n"; + EXIT(EXIT_FAILURE); + } + } } #endif + diff --git a/src/test/CMakeLists.txt b/src/test/CMakeLists.txt index faca7f86a..4e5af0f8d 100644 --- a/src/test/CMakeLists.txt +++ b/src/test/CMakeLists.txt @@ -2,6 +2,7 @@ include(MMseqsSetupTest) set(TESTS #TestAdjustedKmerIterator.cpp + TestAlign2ClustReducer.cpp TestAlignment.cpp TestAlignmentPerformance.cpp TestAlignmentTraceback.cpp diff --git a/src/test/TestAlign2ClustReducer.cpp b/src/test/TestAlign2ClustReducer.cpp new file mode 100644 index 000000000..ed7343a64 --- /dev/null +++ b/src/test/TestAlign2ClustReducer.cpp @@ -0,0 +1,258 @@ +// Unit tests for Align2ClustReducer and Align2ClustChunking using synthetic candidates. +// +// These classes have no dependency on OpenMP, MPI, DBReader, or DBWriter (see the class +// comments in Align2ClustReducer.h and Align2ClustChunking.h), so they can be exercised +// here with hand-built ClusterResult vectors instead of a real sequence/prefilter DB. +// +// Coverage: +// - SET_COVER tie-breaking (larger prefSize wins; ties broken by representative id) +// - GREEDY sequence-order processing +// - Already-assigned representatives/members are correctly rejected +// - Overlapping member sets across candidates +// - Empty and singleton candidates +// - Out-of-order pushResult() calls (correctness must not depend on push order) +// - Align2ClustChunking::computeChunks / isChunkOwnedByRank boundary cases +#include "Align2ClustReducer.h" +#include "Align2ClustChunking.h" + +#include +#include +#include +#include + +const char* binary_name = "test_align2clust_reducer"; + +// Plain assert() is compiled out under -DNDEBUG (the default Release build used by this +// project's CMake configuration), which would silently turn every check below into a +// no-op. Use an unconditional check instead so these tests still fail loudly in Release +// builds. +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::cerr << "CHECK failed at " << __FILE__ << ":" << __LINE__ << ": " #cond << "\n"; \ + std::exit(1); \ + } \ + } while (0) + +static ClusterResult makeResult(size_t sequenceIdx, DBLocalId representativeId, size_t prefSize, + std::vector memberIds) { + ClusterResult result; + result.sequenceIdx = sequenceIdx; + result.representativeId = representativeId; + result.prefSize = prefSize; + result.memberIds = std::move(memberIds); + return result; +} + +// Push every result, then finish()/finalizeSingletons(), returning the assignment vector. +static std::vector runReducer(Align2ClustReducer::Mode mode, size_t dbSize, + std::vector results) { + Align2ClustReducer reducer(mode, dbSize, std::max(1, results.size())); + reducer.start(); + for (ClusterResult &result : results) { + reducer.pushResult(std::move(result)); + } + reducer.finish(); + reducer.finalizeSingletons(); + + std::vector assignment(dbSize); + for (size_t i = 0; i < dbSize; ++i) { + assignment[i] = reducer.getAssignment(i); + } + return assignment; +} + +static void testSetCoverTieBreaking() { + // dbSize == 4, two candidates (rep 2 and rep 3) tie on prefSize == 5. Both are + // buffered into the heap before either is finalized (their gate only opens once + // currentPrefSize_ drops below 5, i.e. once the smaller trailing candidates at + // sequenceIdx 2/3 have also been drained from the reorder buffer). Among a tie in + // memberCount, SetCoverComparator orders the *higher* representative id to the + // front of the heap, so it is finalized first (empirically verified against the + // implementation; this is also exactly the pre-refactor behavior, per the Phase 2 + // byte-identical regression check against a real linclust run). + const size_t dbSize = 4; + std::vector results; + // sequenceIdx 0: representative 2, prefSize 5, tie with rep 3 below + results.push_back(makeResult(0, 2, 5, {0, 1})); + // sequenceIdx 1: representative 3, prefSize 5 -- ties with rep 2 above, wins (see + // comment) and claims members {1, 3} first. + results.push_back(makeResult(1, 3, 5, {1, 3})); + // sequenceIdx 2: singleton candidate (memberIds.size() == 1) for representative 2; + // SET_COVER never even queues single-member candidates (see consumerLoopSetCover), + // so this is a no-op and 2's fate is decided purely by finalizeSingletons() below. + results.push_back(makeResult(2, 2, 1, {2})); + // sequenceIdx 3: same, for representative 3 -- also a no-op. + results.push_back(makeResult(3, 3, 1, {3})); + + std::vector assignment = runReducer(Align2ClustReducer::SET_COVER, dbSize, results); + // Rep 3 wins the tie and claims members {1, 3}. + CHECK(assignment[1] == 3); + CHECK(assignment[3] == 3); + // Rep 2's candidate then has member 1 already taken, leaving only {0} + // (validCount == 1), so it is rejected outright: 0 and 2 fall back to singletons. + CHECK(assignment[0] == 0); + CHECK(assignment[2] == 2); + std::cout << "testSetCoverTieBreaking OK\n"; +} + +static void testGreedySequenceOrder() { + // GREEDY processes candidates strictly in ascending sequenceIdx (== sequence id) + // order; earlier representatives claim members first. + const size_t dbSize = 5; + std::vector results; + results.push_back(makeResult(0, 0, 0, {0, 1, 2})); + results.push_back(makeResult(1, 1, 0, {1, 3})); // 1 already claimed by rep 0 + results.push_back(makeResult(2, 2, 0, {2})); // 2 already claimed by rep 0 + results.push_back(makeResult(3, 3, 0, {3, 4})); // rep 3 itself unclaimed + results.push_back(makeResult(4, 4, 0, {4})); // 4 already claimed by rep 3 + + std::vector assignment = runReducer(Align2ClustReducer::GREEDY, dbSize, results); + CHECK(assignment[0] == 0); + CHECK(assignment[1] == 0); + CHECK(assignment[2] == 0); + CHECK(assignment[3] == 3); + CHECK(assignment[4] == 3); + std::cout << "testGreedySequenceOrder OK\n"; +} + +static void testEmptyAndSingletonCandidates() { + // A representative candidate with no members at all must still leave its own + // sequence assigned to itself via finalizeSingletons(), and must not disturb + // unrelated sequences. + const size_t dbSize = 3; + std::vector results; + results.push_back(makeResult(0, 0, 0, {})); // empty candidate + results.push_back(makeResult(1, 1, 1, {1})); // singleton candidate + results.push_back(makeResult(2, 2, 0, {})); // empty candidate + + std::vector assignment = runReducer(Align2ClustReducer::GREEDY, dbSize, results); + CHECK(assignment[0] == 0); + CHECK(assignment[1] == 1); + CHECK(assignment[2] == 2); + std::cout << "testEmptyAndSingletonCandidates OK\n"; +} + +static void testOutOfOrderPushDoesNotAffectResult() { + // Candidate generation may be parallel and produce results in any order (this is + // the property distributed/MPI worker ranks rely on): pushing the exact same + // candidates in reverse order must produce an identical assignment. + const size_t dbSize = 4; + std::vector forward; + forward.push_back(makeResult(0, 0, 0, {0, 1})); + forward.push_back(makeResult(1, 1, 0, {1})); + forward.push_back(makeResult(2, 2, 0, {2, 3})); + forward.push_back(makeResult(3, 3, 0, {3})); + + std::vector reversed; + for (size_t i = forward.size(); i-- > 0;) { + reversed.push_back(forward[i]); + } + + std::vector forwardAssignment = runReducer(Align2ClustReducer::GREEDY, dbSize, forward); + std::vector reversedAssignment = runReducer(Align2ClustReducer::GREEDY, dbSize, reversed); + CHECK(forwardAssignment == reversedAssignment); + CHECK(forwardAssignment[0] == 0); + CHECK(forwardAssignment[1] == 0); + CHECK(forwardAssignment[2] == 2); + CHECK(forwardAssignment[3] == 2); + std::cout << "testOutOfOrderPushDoesNotAffectResult OK\n"; +} + +static void testDistributedCollectOnlyMode() { + // A distributed worker's collector never touches live assignment state: isAssigned() + // is always false and pushResult() just buffers, regardless of push order. + Align2ClustReducer collector(Align2ClustReducer::SET_COVER, /*dbSize=*/10, /*reorderCapacity=*/1, + /*distributedCollectOnly=*/true); + collector.start(); + CHECK(collector.isAssigned(0) == false); + collector.pushResult(makeResult(5, 5, 3, {5, 6})); + collector.pushResult(makeResult(0, 0, 1, {0})); + CHECK(collector.isAssigned(5) == false); + collector.finish(); + + const std::vector &collected = collector.getCollectedResults(); + CHECK(collected.size() == 2); + // Buffered in push order, not sequenceIdx order (no reordering in this mode). + CHECK(collected[0].sequenceIdx == 5); + CHECK(collected[1].sequenceIdx == 0); + std::cout << "testDistributedCollectOnlyMode OK\n"; +} + +static void testComputeChunksBoundaries() { + // Exact multiple. + { + std::vector chunks = Align2ClustChunking::computeChunks(100, 25); + CHECK(chunks.size() == 4); + for (size_t i = 0; i < chunks.size(); ++i) { + CHECK(chunks[i].index == i); + CHECK(chunks[i].start == i * 25); + CHECK(chunks[i].end == (i + 1) * 25); + } + } + // Remainder in the last chunk. + { + std::vector chunks = Align2ClustChunking::computeChunks(103, 25); + CHECK(chunks.size() == 5); + CHECK(chunks.back().start == 100); + CHECK(chunks.back().end == 103); + } + // endRange == 0 -> no chunks. + { + std::vector chunks = Align2ClustChunking::computeChunks(0, 25); + CHECK(chunks.empty()); + } + // chunkSize >= endRange -> a single chunk covering the whole range. + { + std::vector chunks = Align2ClustChunking::computeChunks(10, 1000); + CHECK(chunks.size() == 1); + CHECK(chunks[0].start == 0); + CHECK(chunks[0].end == 10); + } + // chunkSize == 0 is clamped to 1 (one chunk per index). + { + std::vector chunks = Align2ClustChunking::computeChunks(3, 0); + CHECK(chunks.size() == 3); + } + std::cout << "testComputeChunksBoundaries OK\n"; +} + +static void testChunkOwnershipIsDeterministicAcrossTopologies() { + // numProc <= 1 always owns every chunk (the non-distributed fast path). + CHECK(Align2ClustChunking::isChunkOwnedByRank(0, 0, 1) == true); + CHECK(Align2ClustChunking::isChunkOwnedByRank(41, 0, 1) == true); + CHECK(Align2ClustChunking::isChunkOwnedByRank(41, 0, 0) == true); + + // Round robin: every chunk index is owned by exactly one rank out of numProc. + const int numProc = 4; + for (size_t chunkIndex = 0; chunkIndex < 17; ++chunkIndex) { + int owners = 0; + for (int rank = 0; rank < numProc; ++rank) { + if (Align2ClustChunking::isChunkOwnedByRank(chunkIndex, rank, numProc)) { + owners++; + } + } + CHECK(owners == 1); + } + + // Changing numProc (simulating a topology change across a resume) reshuffles + // ownership deterministically; the same (chunkIndex, rank, numProc) triple always + // gives the same answer, independent of call order. + CHECK(Align2ClustChunking::isChunkOwnedByRank(5, 1, 4) == true); + CHECK(Align2ClustChunking::isChunkOwnedByRank(5, 1, 4) == true); + CHECK(Align2ClustChunking::isChunkOwnedByRank(5, 0, 2) == false); + CHECK(Align2ClustChunking::isChunkOwnedByRank(5, 1, 2) == true); + std::cout << "testChunkOwnershipIsDeterministicAcrossTopologies OK\n"; +} + +int main(int, const char**) { + testSetCoverTieBreaking(); + testGreedySequenceOrder(); + testEmptyAndSingletonCandidates(); + testOutOfOrderPushDoesNotAffectResult(); + testDistributedCollectOnlyMode(); + testComputeChunksBoundaries(); + testChunkOwnershipIsDeterministicAcrossTopologies(); + std::cout << "All Align2ClustReducer/Align2ClustChunking tests passed\n"; + return 0; +}