Skip to content

Commit 255bc8e

Browse files
DoubleCorePicodet
authored andcommitted
fix(dataset_profiler): bound sparse sample memory and regenerate examples from current serializer
Signed-off-by: DoubleCore <851678303@qq.com> Assisted-by: OpenCode:ark-code-latest
1 parent a03a5ba commit 255bc8e

6 files changed

Lines changed: 255 additions & 93 deletions

File tree

‎tools/dataset_profiler/dataset_profiler/profiler.py‎

Lines changed: 105 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,17 @@
4646
_DEFAULT_DENSE_METRIC = "euclidean"
4747
_DEFAULT_SPARSE_METRIC = "ip"
4848
_MAX_ANISOTROPY_DIM = 50000
49-
_MAX_ANISOTROPY_BYTES = 256 * 1024 * 1024
49+
# One byte budget shared by every sampled-memory guard: the dense float64
50+
# sample, the dense expansion of a sparse sample, and the retained sparse
51+
# sample rows below. Sparse samples must obey the same budget as dense ones
52+
# even though their rows carry a variable number of entries.
53+
_SAMPLE_MEMORY_BUDGET_BYTES = 256 * 1024 * 1024
54+
# A retained sparse sample keeps one (uint32, float32) copy per row and later
55+
# expands into one Python dict entry per finite value; estimate both together
56+
# from nnz with interpreter object overhead so nnz-heavy rows trip the budget
57+
# before any copy or dict expansion happens instead of exhausting memory.
58+
_SPARSE_SAMPLE_ENTRY_BYTES = 32
59+
_SPARSE_SAMPLE_ROW_OVERHEAD_BYTES = 128
5060
_MAX_SPARSE_STATS_DIMENSION = 10_000_000
5161
_PCA_SVD_MAX_DIM = 4096
5262
# Exact norm quantiles/outliers and exact duplicate groups need per-row state.
@@ -279,10 +289,35 @@ def _dense_sample_size(count: int, dim: int, requested: int) -> int:
279289
dataset_limited = min(int(requested), int(count))
280290
if dim <= 0:
281291
return int(dataset_limited)
282-
max_rows = _MAX_ANISOTROPY_BYTES // (int(dim) * 8)
292+
max_rows = _SAMPLE_MEMORY_BUDGET_BYTES // (int(dim) * 8)
283293
return int(min(dataset_limited, max_rows))
284294

285295

296+
def _sparse_sample_row_bytes(ids: np.ndarray) -> int:
297+
"""Estimate the bytes one sampled sparse row keeps after retention and
298+
after the Python dict expansion performed by :func:`_pairwise_sparse`."""
299+
return (
300+
_SPARSE_SAMPLE_ROW_OVERHEAD_BYTES + int(ids.size) * _SPARSE_SAMPLE_ENTRY_BYTES
301+
)
302+
303+
304+
def _sparse_sample_budget_reason(
305+
next_row_bytes: int, retained_bytes: int, retained_rows: int
306+
) -> str:
307+
"""Return why retention of the sampled sparse rows stopped."""
308+
return (
309+
"sparse sample limited to %d retained rows: continuing would exceed the "
310+
"%d-byte sample memory budget (estimated %d total bytes so far plus the "
311+
"pending row); the sampled sparse pairwise and distribution metrics are "
312+
"unavailable, while streaming and full-split metrics are still reported"
313+
% (
314+
retained_rows,
315+
_SAMPLE_MEMORY_BUDGET_BYTES,
316+
int(retained_bytes) + int(next_row_bytes),
317+
)
318+
)
319+
320+
286321
def _full_split_state_reason(
287322
metric_name: str, count: int, bytes_per_vector: int
288323
) -> Optional[str]:
@@ -531,6 +566,11 @@ def _profile_sparse(
531566
vectors_with_invalid = 0
532567
all_zero_vectors = 0
533568
sample_records: List[Tuple[int, np.ndarray, np.ndarray]] = []
569+
# Sample retention is bounded by estimated bytes, not by the row count:
570+
# valid sparse files can carry so many entries per vector that the full
571+
# row count would exhaust memory before any report is produced.
572+
retained_sample_bytes = 0
573+
sample_budget_reason: Optional[str] = None
534574

535575
for index, raw_ids, raw_values in dataset.iter_sparse_records(
536576
max_chunk_bytes=config.max_chunk_bytes
@@ -582,10 +622,21 @@ def _profile_sparse(
582622
"<f4", copy=False
583623
).tobytes()
584624
hashes[index] = np.frombuffer(_digest(payload), dtype=np.uint8)
585-
if index in sample_set:
586-
sample_records.append((index, ids.copy(), values.copy()))
625+
if sample_budget_reason is None and index in sample_set:
626+
row_bytes = _sparse_sample_row_bytes(ids)
627+
if retained_sample_bytes + row_bytes > _SAMPLE_MEMORY_BUDGET_BYTES:
628+
sample_budget_reason = _sparse_sample_budget_reason(
629+
row_bytes, retained_sample_bytes, len(sample_records)
630+
)
631+
warnings.append(sample_budget_reason)
632+
else:
633+
retained_sample_bytes += row_bytes
634+
sample_records.append((index, ids.copy(), values.copy()))
587635

588-
if len(sample_records) != len(sample_set):
636+
if (
637+
sample_budget_reason is None
638+
and len(sample_records) != len(sample_set)
639+
):
589640
raise ValueError("sparse sampling failed: not all sampled records were found")
590641
sample_records.sort(key=lambda item: item[0])
591642
pair_records = [
@@ -621,20 +672,27 @@ def _profile_sparse(
621672
reason=dimension_stats_reason,
622673
dim=int(dim),
623674
)
624-
if any(
625-
not bool(np.isfinite(values).all()) for _, _, values in sample_records
626-
):
627-
warnings.append(_INVALID_SAMPLE_WARNING)
628-
pair_matrix = _pairwise_sparse(pair_records, metric)
629-
pair_values = _upper_triangle(pair_matrix)
630-
duplicate_pair_values = pair_values
631-
duplicate_metric = metric
632-
if metric == "ip":
633-
duplicate_pair_values = _upper_triangle(
634-
_pairwise_sparse(pair_records, "euclidean")
635-
)
636-
duplicate_metric = "euclidean"
637-
sample_matrix = _sparse_sample_to_dense(sample_records, dim)
675+
if sample_budget_reason is None:
676+
if any(
677+
not bool(np.isfinite(values).all()) for _, _, values in sample_records
678+
):
679+
warnings.append(_INVALID_SAMPLE_WARNING)
680+
pair_matrix = _pairwise_sparse(pair_records, metric)
681+
pair_values = _upper_triangle(pair_matrix)
682+
duplicate_pair_values = pair_values
683+
duplicate_metric = metric
684+
if metric == "ip":
685+
duplicate_pair_values = _upper_triangle(
686+
_pairwise_sparse(pair_records, "euclidean")
687+
)
688+
duplicate_metric = "euclidean"
689+
sample_matrix = _sparse_sample_to_dense(sample_records, dim)
690+
else:
691+
pair_matrix = np.zeros((0, 0), dtype=np.float64)
692+
pair_values = np.empty(0, dtype=np.float64)
693+
duplicate_pair_values = pair_values
694+
duplicate_metric = metric
695+
sample_matrix = None
638696

639697
metrics: Dict[str, Any] = {
640698
"invalid_values": _invalid_metric(
@@ -662,6 +720,7 @@ def _profile_sparse(
662720
duplicate_metric,
663721
int(pair_matrix.shape[0]),
664722
duplicates_reason,
723+
near_unavailable_reason=sample_budget_reason,
665724
),
666725
}
667726
metrics.update(
@@ -674,6 +733,7 @@ def _profile_sparse(
674733
rng=rng,
675734
warnings=warnings,
676735
sparse_pairwise=True,
736+
unavailable_reason=sample_budget_reason,
677737
)
678738
)
679739
return metrics
@@ -1058,6 +1118,7 @@ def _duplicates_metric(
10581118
near_metric: str,
10591119
near_sample_size: int,
10601120
unavailable_reason: Optional[str] = None,
1121+
near_unavailable_reason: Optional[str] = None,
10611122
) -> Dict[str, Any]:
10621123
if hashes is None:
10631124
exact = {
@@ -1107,7 +1168,7 @@ def _duplicates_metric(
11071168
"indicator covers the bounded pairwise sample"
11081169
)
11091170
near = _near_duplicate_summary(
1110-
pair_values, epsilon, near_metric, near_sample_size
1171+
pair_values, epsilon, near_metric, near_sample_size, near_unavailable_reason
11111172
)
11121173
return metric_entry(
11131174
DEFINITIONS["duplicates"],
@@ -1120,12 +1181,16 @@ def _duplicates_metric(
11201181

11211182

11221183
def _near_duplicate_summary(
1123-
pair_values: np.ndarray, epsilon: float, metric: str, sample_size: int
1184+
pair_values: np.ndarray,
1185+
epsilon: float,
1186+
metric: str,
1187+
sample_size: int,
1188+
unavailable_reason: Optional[str] = None,
11241189
) -> Dict[str, Any]:
11251190
if pair_values.size == 0:
11261191
return {
11271192
"available": False,
1128-
"reason": "fewer than two sampled vectors",
1193+
"reason": unavailable_reason or "fewer than two sampled vectors",
11291194
"scope": "sample",
11301195
"sample_size": int(sample_size),
11311196
"pair_count": 0,
@@ -1170,9 +1235,26 @@ def _sampled_metrics(
11701235
rng: np.random.Generator,
11711236
warnings: List[str],
11721237
sparse_pairwise: bool = False,
1238+
unavailable_reason: Optional[str] = None,
11731239
) -> Dict[str, Any]:
11741240
sample_size = int(pair_matrix.shape[0])
11751241
pair_count = int(pair_values.size)
1242+
if unavailable_reason is not None:
1243+
return {
1244+
name: metric_entry(
1245+
DEFINITIONS[name],
1246+
"sample",
1247+
sample_size,
1248+
available=False,
1249+
reason=unavailable_reason,
1250+
)
1251+
for name in (
1252+
"pairwise_distances",
1253+
"nearest_neighbor_distances",
1254+
"anisotropy",
1255+
"clustering",
1256+
)
1257+
}
11761258
pairwise = metric_entry(
11771259
DEFINITIONS["pairwise_distances"],
11781260
"sample",
@@ -1482,7 +1564,7 @@ def _sparse_sample_to_dense(
14821564
if dim <= 0 or not records:
14831565
return None
14841566
rows = len(records)
1485-
if rows * dim * 8 > _MAX_ANISOTROPY_BYTES or dim > _MAX_ANISOTROPY_DIM:
1567+
if rows * dim * 8 > _SAMPLE_MEMORY_BUDGET_BYTES or dim > _MAX_ANISOTROPY_DIM:
14861568
return None
14871569
matrix = np.zeros((rows, dim), dtype=np.float64)
14881570
for position, (_, ids, values) in enumerate(records):

0 commit comments

Comments
 (0)