diff --git a/.github/workflows/ci-self-hosted.yml b/.github/workflows/ci-self-hosted.yml new file mode 100644 index 00000000..e7ccc7bb --- /dev/null +++ b/.github/workflows/ci-self-hosted.yml @@ -0,0 +1,67 @@ +name: CppMega MLX self-hosted CI + +on: + pull_request: + branches: [main] + push: + branches: [main] + workflow_dispatch: + +concurrency: + group: cppmega-mlx-ci-${{ github.ref }} + cancel-in-progress: true + +jobs: + mac-contracts: + name: macOS MLX data, model, and eval contracts + runs-on: [self-hosted, macOS, ARM64, cppmega-mlx-macos] + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + - name: Run focused MLX contract suite + shell: bash + run: | + set -euo pipefail + python_bin=/Volumes/external/sources/cppmega.mlx/.venv/bin/python + test -x "${python_bin}" + "${python_bin}" -m pytest -q \ + tests/test_ast_fim.py \ + tests/test_audit_sidecar_parquet.py \ + tests/test_cpp_jsonl_generation_compile_eval.py \ + tests/test_data_package_imports.py \ + tests/test_domain_graph_routes.py \ + tests/test_eval_domain_routed_codegen.py \ + tests/test_inference_generation.py \ + tests/test_megatron_indexed.py \ + tests/test_pack_enriched_rows.py \ + tests/test_streaming_conveyor_progress.py \ + tests/test_process_commits_fail_loud.py \ + tests/test_repair_packed_document_boundaries.py + git diff --check + + linux-portable: + name: Linux portable syntax and conveyor contracts + runs-on: [self-hosted, Linux, X64, cppmega-mlx] + timeout-minutes: 20 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.13" + - name: Install portable test dependencies + run: python -m pip install --disable-pip-version-check pytest numpy pyarrow tokenizers zstandard + - name: Run portable script contracts + shell: bash + run: | + set -euo pipefail + # cppmega.mlx imports Apple's MLX runtime at package import time. The + # Linux runner exercises only portable scripts and deliberately skips + # the macOS-only root conftest/package initialization. + python -m pytest -q --noconftest \ + tests/test_audit_sidecar_parquet.py \ + tests/test_data_package_imports.py \ + tests/test_streaming_conveyor_progress.py \ + tests/test_process_commits_fail_loud.py \ + tests/test_repair_packed_document_boundaries.py + python -m compileall -q scripts tools/clang_indexer + git diff --check diff --git a/.github/workflows/cuda-lane.yml b/.github/workflows/cuda-lane.yml index 85da6b29..f17d7dd4 100644 --- a/.github/workflows/cuda-lane.yml +++ b/.github/workflows/cuda-lane.yml @@ -1,10 +1,9 @@ name: CUDA Lane on: - push: - branches: [main] - pull_request: - branches: [main] + # CUDA capacity is provisioned explicitly (for example on Nebius H200). + # Do not leave ordinary PRs queued forever when no self-hosted CUDA runner + # is attached to this repository. workflow_dispatch: jobs: diff --git a/AGENTS.md b/AGENTS.md index ae23b499..7e653f29 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -28,3 +28,12 @@ See CLAUDE.md. This overrides every rule below. format directly; - only use casts/reshapes as explicit graph decisions, not hidden adapter staging. + +## Native subagents + +- Use Codex native subagents for independent, bounded parallel work when it + materially improves throughput. +- Every native subagent must use `gpt-5.6-sol` with `ultra` or `max` + reasoning. Use `ultra` by default. +- Do not use another model or a lower reasoning level unless the user + explicitly overrides this rule. diff --git a/cppmega_mlx/data/__init__.py b/cppmega_mlx/data/__init__.py index e10bb1af..66402dcd 100644 --- a/cppmega_mlx/data/__init__.py +++ b/cppmega_mlx/data/__init__.py @@ -1,137 +1,90 @@ -"""Data readers and batch collation helpers.""" +"""Data readers and batch collation helpers. -from cppmega_mlx.data.batch import LMTokenBatch, ensure_lm_batch, synthetic_token_batch -from cppmega_mlx.data.dataloader_bridge import ( - LocalTokenBatchDataset, - TorchDataLoaderBridgeConfig, - TorchDataLoaderBridgeError, - build_spawn_dataloader, - is_torch_dataloader_available, - iter_mlx_batches, -) -from cppmega_mlx.data.fim import ( - EOT_ID, - FIMSpecialTokenIds, - FIM_INSTRUCTION_ID, - FIMMode, - FIM_MIDDLE_ID, - FIM_PREFIX_ID, - FIM_SPECIAL_TOKEN_IDS, - FIM_SUFFIX_ID, - apply_fim_permutation, - apply_fim_transform, - apply_ifim_permutation, - apply_ifim_transform, - extract_ifim_instruction_text, - sample_middle_span, -) -from cppmega_mlx.data.megatron_indexed import ( - MegatronIndexedDataset, - MegatronIndexedMetadata, - MegatronIndexedMultiShardDataset, - MegatronIndexedMultiShardMetadata, - megatron_indexed_side_channel_schema, - open_megatron_indexed_dataset, -) -from cppmega_mlx.data.packing import ( - OversizedSamplePolicy, - PackedSequences, - PackingStrategy, - cumulative_doc_ids_from_eos, - document_boundary_mask, - mlx_cumulative_doc_ids_from_eos, - mlx_document_boundary_mask, - mlx_sequence_packing_attention_mask, - pack_bos_aligned_best_fit, - pack_documents_with_eos, -) -from cppmega_mlx.data.parquet_dataset import ( - MultiShardTokenParquetDataset, - ParquetColumns, - TokenParquetDataset, -) -from cppmega_mlx.data.platform_context import ( - MAX_PLATFORM_IDS, - PLATFORM_VOCAB, - PLATFORM_VOCAB_SIZE, - PlatformContext, - encode_platform_context, - parse_platform_context, - platform_ids_array, - render_platform_context, -) -from cppmega_mlx.data.tokenizer_contract import ( - REQUIRED_SPECIAL_TOKEN_IDS, - SpecialTokenMapping, - TOOL_USE_SPECIAL_TOKEN_IDS, - validate_required_special_token_ids, -) -from cppmega_mlx.data.token_dataset import ( - BatchCursor, - TokenDatasetMetadata, - TokenNpzDataset, - iterate_token_batches, - open_token_dataset, -) +The package also contains portable corpus/indexer modules used on Linux hosts +where Apple's MLX runtime is unavailable. Public MLX-backed exports are loaded +on first attribute access so importing a portable submodule does not eagerly +import the entire training runtime. +""" -__all__ = [ - "BatchCursor", - "EOT_ID", - "FIMSpecialTokenIds", - "FIM_INSTRUCTION_ID", - "FIMMode", - "FIM_MIDDLE_ID", - "FIM_PREFIX_ID", - "FIM_SPECIAL_TOKEN_IDS", - "FIM_SUFFIX_ID", - "LMTokenBatch", - "LocalTokenBatchDataset", - "MAX_PLATFORM_IDS", - "MegatronIndexedDataset", - "MegatronIndexedMetadata", - "MegatronIndexedMultiShardDataset", - "MegatronIndexedMultiShardMetadata", - "MultiShardTokenParquetDataset", - "OversizedSamplePolicy", - "PLATFORM_VOCAB", - "PLATFORM_VOCAB_SIZE", - "ParquetColumns", - "PackedSequences", - "PackingStrategy", - "PlatformContext", - "REQUIRED_SPECIAL_TOKEN_IDS", - "SpecialTokenMapping", - "TokenDatasetMetadata", - "TorchDataLoaderBridgeConfig", - "TorchDataLoaderBridgeError", - "TOOL_USE_SPECIAL_TOKEN_IDS", - "TokenNpzDataset", - "TokenParquetDataset", - "apply_fim_permutation", - "apply_fim_transform", - "apply_ifim_permutation", - "apply_ifim_transform", - "build_spawn_dataloader", - "cumulative_doc_ids_from_eos", - "document_boundary_mask", - "encode_platform_context", - "ensure_lm_batch", - "extract_ifim_instruction_text", - "is_torch_dataloader_available", - "iter_mlx_batches", - "iterate_token_batches", - "megatron_indexed_side_channel_schema", - "mlx_cumulative_doc_ids_from_eos", - "mlx_document_boundary_mask", - "mlx_sequence_packing_attention_mask", - "open_megatron_indexed_dataset", - "open_token_dataset", - "pack_bos_aligned_best_fit", - "pack_documents_with_eos", - "parse_platform_context", - "platform_ids_array", - "render_platform_context", - "sample_middle_span", - "synthetic_token_batch", - "validate_required_special_token_ids", -] +from __future__ import annotations + +from importlib import import_module +from typing import Any + + +_EXPORT_MODULES = { + "LMTokenBatch": "cppmega_mlx.data.batch", + "ensure_lm_batch": "cppmega_mlx.data.batch", + "synthetic_token_batch": "cppmega_mlx.data.batch", + "LocalTokenBatchDataset": "cppmega_mlx.data.dataloader_bridge", + "TorchDataLoaderBridgeConfig": "cppmega_mlx.data.dataloader_bridge", + "TorchDataLoaderBridgeError": "cppmega_mlx.data.dataloader_bridge", + "build_spawn_dataloader": "cppmega_mlx.data.dataloader_bridge", + "is_torch_dataloader_available": "cppmega_mlx.data.dataloader_bridge", + "iter_mlx_batches": "cppmega_mlx.data.dataloader_bridge", + "EOT_ID": "cppmega_mlx.data.fim", + "FIMSpecialTokenIds": "cppmega_mlx.data.fim", + "FIM_INSTRUCTION_ID": "cppmega_mlx.data.fim", + "FIMMode": "cppmega_mlx.data.fim", + "FIM_MIDDLE_ID": "cppmega_mlx.data.fim", + "FIM_PREFIX_ID": "cppmega_mlx.data.fim", + "FIM_SPECIAL_TOKEN_IDS": "cppmega_mlx.data.fim", + "FIM_SUFFIX_ID": "cppmega_mlx.data.fim", + "apply_fim_permutation": "cppmega_mlx.data.fim", + "apply_fim_transform": "cppmega_mlx.data.fim", + "apply_ifim_permutation": "cppmega_mlx.data.fim", + "apply_ifim_transform": "cppmega_mlx.data.fim", + "extract_ifim_instruction_text": "cppmega_mlx.data.fim", + "sample_middle_span": "cppmega_mlx.data.fim", + "MegatronIndexedDataset": "cppmega_mlx.data.megatron_indexed", + "MegatronIndexedMetadata": "cppmega_mlx.data.megatron_indexed", + "MegatronIndexedMultiShardDataset": "cppmega_mlx.data.megatron_indexed", + "MegatronIndexedMultiShardMetadata": "cppmega_mlx.data.megatron_indexed", + "megatron_indexed_side_channel_schema": "cppmega_mlx.data.megatron_indexed", + "open_megatron_indexed_dataset": "cppmega_mlx.data.megatron_indexed", + "OversizedSamplePolicy": "cppmega_mlx.data.packing", + "PackedSequences": "cppmega_mlx.data.packing", + "PackingStrategy": "cppmega_mlx.data.packing", + "cumulative_doc_ids_from_eos": "cppmega_mlx.data.packing", + "document_boundary_mask": "cppmega_mlx.data.packing", + "mlx_cumulative_doc_ids_from_eos": "cppmega_mlx.data.packing", + "mlx_document_boundary_mask": "cppmega_mlx.data.packing", + "mlx_sequence_packing_attention_mask": "cppmega_mlx.data.packing", + "pack_bos_aligned_best_fit": "cppmega_mlx.data.packing", + "pack_documents_with_eos": "cppmega_mlx.data.packing", + "MultiShardTokenParquetDataset": "cppmega_mlx.data.parquet_dataset", + "ParquetColumns": "cppmega_mlx.data.parquet_dataset", + "TokenParquetDataset": "cppmega_mlx.data.parquet_dataset", + "MAX_PLATFORM_IDS": "cppmega_mlx.data.platform_context", + "PLATFORM_VOCAB": "cppmega_mlx.data.platform_context", + "PLATFORM_VOCAB_SIZE": "cppmega_mlx.data.platform_context", + "PlatformContext": "cppmega_mlx.data.platform_context", + "encode_platform_context": "cppmega_mlx.data.platform_context", + "parse_platform_context": "cppmega_mlx.data.platform_context", + "platform_ids_array": "cppmega_mlx.data.platform_context", + "render_platform_context": "cppmega_mlx.data.platform_context", + "REQUIRED_SPECIAL_TOKEN_IDS": "cppmega_mlx.data.tokenizer_contract", + "SpecialTokenMapping": "cppmega_mlx.data.tokenizer_contract", + "TOOL_USE_SPECIAL_TOKEN_IDS": "cppmega_mlx.data.tokenizer_contract", + "validate_required_special_token_ids": "cppmega_mlx.data.tokenizer_contract", + "BatchCursor": "cppmega_mlx.data.token_dataset", + "TokenDatasetMetadata": "cppmega_mlx.data.token_dataset", + "TokenNpzDataset": "cppmega_mlx.data.token_dataset", + "iterate_token_batches": "cppmega_mlx.data.token_dataset", + "open_token_dataset": "cppmega_mlx.data.token_dataset", +} + +__all__ = list(_EXPORT_MODULES) + + +def __getattr__(name: str) -> Any: + module_name = _EXPORT_MODULES.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(import_module(module_name), name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(__all__)) diff --git a/cppmega_mlx/data/ast_fim.py b/cppmega_mlx/data/ast_fim.py index 3afedbd6..b8a41ae9 100644 --- a/cppmega_mlx/data/ast_fim.py +++ b/cppmega_mlx/data/ast_fim.py @@ -22,8 +22,8 @@ * 10% of the time we emit a random-CHAR-FIM example whose middle is a random token span (``fim.sample_middle_span``) that may start/end mid-statement — this teaches robustness to arbitrary cursor positions. The fall-back to the - char path is RECORDED in the returned ``AstFimResult.kind`` ("char_fim"), - never silent. + char path is RECORDED in the returned ``AstFimResult.kind`` ("char_fim" + or "char_ifim" for instruction-aware examples), never silent. iFIM variant: when the document carries a leading comment / docstring we extract it via ``fim.extract_ifim_instruction_text`` and emit through @@ -58,7 +58,7 @@ # Fraction of AST-FIM examples; the remaining slice is random-char-FIM. DEFAULT_AST_FIM_RATE = 0.9 -AstFimKind = Literal["ast_fim", "char_fim", "ast_ifim"] +AstFimKind = Literal["ast_fim", "char_fim", "ast_ifim", "char_ifim"] class NoEligibleChunkError(ValueError): @@ -306,11 +306,11 @@ def apply_ast_ifim( except NoEligibleChunkError: start, end = sample_middle_span(length, rng=rand) chunk_index = None - kind = "char_fim" + kind = "char_ifim" else: start, end = sample_middle_span(length, rng=rand) chunk_index = None - kind = "char_fim" + kind = "char_ifim" mode: FIMMode = "spm" if rand.random() < spm_rate else "psm" permuted = apply_ifim_permutation( diff --git a/cppmega_mlx/data/megatron_indexed.py b/cppmega_mlx/data/megatron_indexed.py index 42d8a576..a85bc1b1 100644 --- a/cppmega_mlx/data/megatron_indexed.py +++ b/cppmega_mlx/data/megatron_indexed.py @@ -4,7 +4,7 @@ import json import struct -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Any, Iterator @@ -22,6 +22,7 @@ _INDEX_HEADER = b"MMIDIDX\x00\x00" _INDEX_VERSION = 1 _ATTENTION_SIDE_CHANNEL_KEY = "attention_mask" +_LOSS_MASK_SIDE_CHANNEL_KEY = "loss_mask" _STRUCTURE_SIDE_CHANNEL_KEYS = ( "structure_ids", "dep_levels", @@ -39,22 +40,36 @@ _TEMPORAL_SIDE_CHANNEL_KEYS = ( "change_mask_pre", "change_mask_post", + "hunk_ids", + "edit_ops", +) +_DOMAIN_SIDE_CHANNEL_KEYS = ( + "domain_ids", + "role_ids", + "entity_ids", + "scope_ids", + "source_doc_ids", + "confidence_ids", ) _SIDE_CHANNEL_KEYS = ( _ATTENTION_SIDE_CHANNEL_KEY, + _LOSS_MASK_SIDE_CHANNEL_KEY, *_STRUCTURE_SIDE_CHANNEL_KEYS, *_PLATFORM_SIDE_CHANNEL_KEYS, *_SEMANTIC_SIDE_CHANNEL_KEYS, *_TEMPORAL_SIDE_CHANNEL_KEYS, + *_DOMAIN_SIDE_CHANNEL_KEYS, ) _LM_BATCH_DIRECT_SIDE_CHANNEL_KEYS = ( _ATTENTION_SIDE_CHANNEL_KEY, + _LOSS_MASK_SIDE_CHANNEL_KEY, *_STRUCTURE_SIDE_CHANNEL_KEYS, *_PLATFORM_SIDE_CHANNEL_KEYS, ) _FAMILY_SIDE_CHANNEL_KEYS: dict[str, tuple[str, ...]] = { "semantic_graph": _SEMANTIC_SIDE_CHANNEL_KEYS, "temporal_diff": _TEMPORAL_SIDE_CHANNEL_KEYS, + "domain_routes": _DOMAIN_SIDE_CHANNEL_KEYS, } _SIDE_CHANNEL_SCHEMA_VERSION = 1 _SIDE_CHANNEL_KEY_ALIASES: dict[str, tuple[str, ...]] = { @@ -71,6 +86,14 @@ "def_use": ("token_def_use",), "change_mask_pre": ("token_change_mask_pre",), "change_mask_post": ("token_change_mask_post",), + "hunk_ids": ("hunk_id_per_token",), + "edit_ops": ("edit_op_per_token",), + "domain_ids": ("token_domain_ids",), + "role_ids": ("token_role_ids",), + "entity_ids": ("token_entity_ids",), + "scope_ids": ("token_scope_ids",), + "source_doc_ids": ("token_source_doc_ids",), + "confidence_ids": ("token_confidence_ids",), } _DOCUMENT_ID_KEYS = ("document_ids", "doc_ids", "packing_document_ids") _SIDE_CHANNEL_ALIAS_TO_KEY = { @@ -92,8 +115,16 @@ "ngram_sidecar", "ngrams", ) -_GRAPH_SIDECAR_SCHEMA = "cppmega_graph_routes_v1" -_GRAPH_EDGE_KEYS = ("token_call_edges", "token_type_edges") +_GRAPH_SIDECAR_SCHEMAS = ("cppmega_graph_routes_v1", "cppmega_graph_routes_v2") +_GRAPH_PAIR_EDGE_KEYS = ("token_call_edges", "token_type_edges") +_GRAPH_TRIPLE_EDGE_KEYS = ( + "token_domain_edges", + "token_build_edges", + "token_shell_edges", + "token_diagnostic_edges", + "token_cross_domain_edges", +) +_GRAPH_EDGE_KEYS = (*_GRAPH_PAIR_EDGE_KEYS, *_GRAPH_TRIPLE_EDGE_KEYS) _GRAPH_CHUNK_KEYS = ( "token_chunk_starts", "token_chunk_ends", @@ -104,6 +135,11 @@ _GRAPH_RELATION_BY_KEY = { "token_call_edges": "call", "token_type_edges": "type", + "token_domain_edges": "domain", + "token_build_edges": "build", + "token_shell_edges": "shell", + "token_diagnostic_edges": "diagnostic", + "token_cross_domain_edges": "cross_domain", } _INDEX_DTYPES: dict[int, np.dtype] = { @@ -168,6 +204,7 @@ class MegatronGraphRoutePacket: chunk_ends: np.ndarray chunk_kinds: np.ndarray chunk_dep_levels: np.ndarray + edge_kinds: dict[str, np.ndarray] = field(default_factory=dict) @property def num_chunks(self) -> int: @@ -248,12 +285,22 @@ def __init__( offsets=sequence_offsets, lengths=sequence_lengths, ) - windows = _build_windows( - sequence_offsets, - sequence_lengths, - self.seq_len, - itemsize=token_dtype.itemsize, + self._restores_compact_fixed_rows = _is_compact_fixed_row_shard( + sidecar, + sequence_lengths=sequence_lengths, + seq_len=self.seq_len, ) + if self._restores_compact_fixed_rows: + windows = sequence_offsets.astype(np.int64, copy=True) + window_lengths = sequence_lengths.astype(np.int64, copy=True) + else: + windows = _build_windows( + sequence_offsets, + sequence_lengths, + self.seq_len, + itemsize=token_dtype.itemsize, + ) + window_lengths = np.full(windows.shape, self.seq_len, dtype=np.int64) if not len(windows): raise ValueError("Megatron token data does not contain a full fixed-shape sample") @@ -262,6 +309,7 @@ def __init__( self._sequence_lengths = sequence_lengths.astype(np.int64, copy=True) self._bin_mmap = np.memmap(self.bin_path, mode="r", dtype=np.uint8) self._windows = windows + self._window_lengths = window_lengths self._side_channels = _load_side_channels( sidecar, prefix=resolved.prefix, @@ -419,15 +467,35 @@ def graph_route_packet_for_sample(self, sample_index: int) -> MegatronGraphRoute if len(set(lengths.values())) != 1: raise ValueError(f"graph chunk sidecar length mismatch: {lengths}") num_nodes = int(starts.shape[0]) + if ( + np.any(starts < 0) + or np.any(ends <= starts) + or np.any(ends > self.seq_len) + ): + raise ValueError( + "graph chunk spans must satisfy 0 <= start < end <= sequence length" + ) edges: dict[str, EdgeIndex] = {} + edge_kinds: dict[str, np.ndarray] = {} for key, relation in _GRAPH_RELATION_BY_KEY.items(): if key not in self._graph_sidecars: continue - pairs = _read_graph_pairs(self._graph_sidecars[key], sequence_index) + if key in _GRAPH_PAIR_EDGE_KEYS: + pairs = _read_graph_pairs(self._graph_sidecars[key], sequence_index) + relation_num_nodes = num_nodes + else: + triples = _read_graph_triples(self._graph_sidecars[key], sequence_index) + pairs = triples[:, :2] + edge_kinds[relation] = triples[:, 2].astype(np.int32, copy=False) + relation_num_nodes = self.seq_len + if pairs.size and (np.any(pairs < 0) or np.any(pairs >= relation_num_nodes)): + raise ValueError( + f"{key} endpoints exceed {relation_num_nodes} nodes for sample {sample_index}" + ) edges[relation] = EdgeIndex.from_pairs( pairs, relation=relation, - num_nodes=num_nodes, + num_nodes=relation_num_nodes, ) return MegatronGraphRoutePacket( graph=GraphPacket(edges=edges, num_nodes=num_nodes), @@ -435,6 +503,7 @@ def graph_route_packet_for_sample(self, sample_index: int) -> MegatronGraphRoute chunk_ends=ends.astype(np.int32, copy=False), chunk_kinds=kinds.astype(np.int32, copy=False), chunk_dep_levels=dep_levels.astype(np.int32, copy=False), + edge_kinds=edge_kinds, ) def _sequence_index_for_window(self, window_index: int) -> int: @@ -447,7 +516,13 @@ def _sequence_index_for_window(self, window_index: int) -> int: "graph route sidecars are document-aligned; token-window slicing " "requires a sequence-aligned sample start" ) - if int(self._sequence_lengths[sequence_index]) != self.seq_len: + sequence_length = int(self._sequence_lengths[sequence_index]) + if self._restores_compact_fixed_rows: + if sequence_length > self.seq_len: + raise ValueError( + "compact fixed-row sequence length exceeds the configured seq_len" + ) + elif sequence_length != self.seq_len: raise NotImplementedError( "graph route sidecars require sequence length to equal seq_len; " "regenerate packed graph sidecars at the training block size" @@ -458,44 +533,50 @@ def _make_batch(self, sample_idx: np.ndarray) -> LMTokenBatch: return _lm_batch_from_numpy(self._make_numpy_batch(sample_idx)) def _make_numpy_batch(self, sample_idx: np.ndarray) -> dict[str, np.ndarray]: - tokens = np.empty((sample_idx.shape[0], self.seq_len), dtype=np.int32) + tokens = np.zeros((sample_idx.shape[0], self.seq_len), dtype=np.int32) side_channels = { - key: np.empty( + key: np.zeros( (sample_idx.shape[0], self.seq_len), dtype=_target_side_channel_dtype(key), ) for key in self._side_channels } + if "hunk_ids" in side_channels: + side_channels["hunk_ids"].fill(-1) document_ids = ( - np.empty((sample_idx.shape[0], self.seq_len), dtype=np.int32) + np.zeros((sample_idx.shape[0], self.seq_len), dtype=np.int32) if self._document_ids is not None else None ) for row, window_index in enumerate(sample_idx): - byte_offset = int(self._windows[int(window_index)]) + resolved_window = int(window_index) + byte_offset = int(self._windows[resolved_window]) + window_length = int(self._window_lengths[resolved_window]) token_view = np.frombuffer( self._bin_mmap, dtype=self._dtype, - count=self.seq_len, + count=window_length, offset=byte_offset, ) - tokens[row] = _to_int32_tokens(token_view) + tokens[row, :window_length] = _to_int32_tokens(token_view) for key, storage in self._side_channels.items(): side_view = np.frombuffer( storage.mmap, dtype=storage.dtype, - count=self.seq_len, - offset=int(storage.windows[int(window_index)]), + count=window_length, + offset=int(storage.windows[resolved_window]), + ) + side_channels[key][row, :window_length] = _to_side_channel_values( + key, side_view ) - side_channels[key][row] = _to_side_channel_values(key, side_view) if self._document_ids is not None and document_ids is not None: doc_view = np.frombuffer( self._document_ids.mmap, dtype=self._document_ids.dtype, - count=self.seq_len, - offset=int(self._document_ids.windows[int(window_index)]), + count=window_length, + offset=int(self._document_ids.windows[resolved_window]), ) - document_ids[row] = _to_document_id_values(doc_view) + document_ids[row, :window_length] = _to_document_id_values(doc_view) batch: dict[str, np.ndarray] = {"tokens": tokens} for key in _SIDE_CHANNEL_KEYS: @@ -789,14 +870,22 @@ def megatron_indexed_graph_sidecar_schema() -> dict[str, dict[str, object]]: return { key: { - "kind": "edge_pairs" if key in _GRAPH_EDGE_KEYS else "ragged_1d", + "kind": "edge_pairs" + if key in _GRAPH_PAIR_EDGE_KEYS + else "edge_triples" + if key in _GRAPH_TRIPLE_EDGE_KEYS + else "ragged_1d", "dtype": "int32" if key in _GRAPH_EDGE_KEYS else "uint32" if key in ("token_chunk_starts", "token_chunk_ends") else "uint16", - "shape_tail": [2] if key in _GRAPH_EDGE_KEYS else [], - "schema": _GRAPH_SIDECAR_SCHEMA, + "shape_tail": [2] + if key in _GRAPH_PAIR_EDGE_KEYS + else [3] + if key in _GRAPH_TRIPLE_EDGE_KEYS + else [], + "schema": "cppmega_graph_routes_v2", } for key in _GRAPH_SIDECAR_KEYS } @@ -863,6 +952,7 @@ class _GraphSidecarStorage: data: np.ndarray item_count: int shape_tail: tuple[int, ...] + coordinate_space: str | None @dataclass(frozen=True) @@ -1094,9 +1184,10 @@ def _load_graph_sidecars( if not isinstance(entries, dict): raise ValueError("graph_sidecar_paths must be a mapping of key to CSR metadata") schema = sidecar.get("graph_sidecar_schema") - if schema != _GRAPH_SIDECAR_SCHEMA: + if schema not in _GRAPH_SIDECAR_SCHEMAS: raise ValueError( - f"unsupported graph_sidecar_schema {schema!r}; expected {_GRAPH_SIDECAR_SCHEMA!r}" + f"unsupported graph_sidecar_schema {schema!r}; expected one of " + f"{_GRAPH_SIDECAR_SCHEMAS!r}" ) base_dir = metadata_path.parent if metadata_path is not None else prefix.parent storages: dict[str, _GraphSidecarStorage] = {} @@ -1104,6 +1195,8 @@ def _load_graph_sidecars( key = str(raw_key) if key not in _GRAPH_SIDECAR_KEYS: raise NotImplementedError(f"unsupported graph sidecar key {key!r}") + if schema == "cppmega_graph_routes_v1" and key in _GRAPH_TRIPLE_EDGE_KEYS: + raise ValueError(f"{key} requires cppmega_graph_routes_v2") if key in storages: raise ValueError(f"{key} graph sidecar declared more than once") storages[key] = _parse_graph_sidecar_entry( @@ -1111,6 +1204,7 @@ def _load_graph_sidecars( entry, base_dir=base_dir, sequence_count=sequence_count, + schema=str(schema), ) missing = [key for key in _GRAPH_CHUNK_KEYS if key not in storages] if missing: @@ -1125,6 +1219,7 @@ def _parse_graph_sidecar_entry( *, base_dir: Path, sequence_count: int, + schema: str, ) -> _GraphSidecarStorage: if not isinstance(entry, dict): raise ValueError(f"{key} graph sidecar entry must be an object") @@ -1136,18 +1231,37 @@ def _parse_graph_sidecar_entry( dtype = _coerce_graph_sidecar_dtype(key, entry.get("dtype")) except KeyError as error: raise ValueError(f"{key} graph sidecar entry missing required field {error.args[0]!r}") from error - if key in _GRAPH_EDGE_KEYS: + if key in _GRAPH_PAIR_EDGE_KEYS: if kind != "edge_pairs": raise ValueError(f"{key} graph sidecar kind must be edge_pairs, got {kind!r}") shape_tail = tuple(int(x) for x in entry.get("shape_tail", [2])) if shape_tail != (2,): raise ValueError(f"{key} edge sidecar shape_tail must be [2], got {shape_tail}") + elif key in _GRAPH_TRIPLE_EDGE_KEYS: + if kind != "edge_triples": + raise ValueError(f"{key} graph sidecar kind must be edge_triples, got {kind!r}") + shape_tail = tuple(int(x) for x in entry.get("shape_tail", [3])) + if shape_tail != (3,): + raise ValueError(f"{key} edge sidecar shape_tail must be [3], got {shape_tail}") elif key in _GRAPH_CHUNK_KEYS: if kind != "ragged_1d": raise ValueError(f"{key} graph sidecar kind must be ragged_1d, got {kind!r}") shape_tail = () else: raise NotImplementedError(f"unsupported graph sidecar key {key!r}") + coordinate_space = entry.get("coordinate_space") + if schema == "cppmega_graph_routes_v2": + expected_coordinate = ( + "chunk_index" + if key in _GRAPH_PAIR_EDGE_KEYS + or key in {"token_chunk_kinds", "token_chunk_dep_levels"} + else "token_index" + ) + if coordinate_space != expected_coordinate: + raise ValueError( + f"{key} coordinate_space must be {expected_coordinate!r}, " + f"got {coordinate_space!r}" + ) if offset_dtype != "int64": raise ValueError(f"{key} graph sidecar offsets must be int64, got {offset_dtype!r}") if not offsets_path.is_absolute(): @@ -1204,6 +1318,7 @@ def _parse_graph_sidecar_entry( data=data, item_count=item_count, shape_tail=shape_tail, + coordinate_space=None if coordinate_space is None else str(coordinate_space), ) @@ -1243,6 +1358,19 @@ def _read_graph_pairs(storage: _GraphSidecarStorage, sequence_index: int) -> np. return pairs.astype(np.int32, copy=False) +def _read_graph_triples(storage: _GraphSidecarStorage, sequence_index: int) -> np.ndarray: + if storage.kind != "edge_triples" or storage.shape_tail != (3,): + raise ValueError(f"{storage.key} graph storage is not edge_triples") + start, end = _graph_range(storage, sequence_index) + values = np.asarray(storage.data[start * 3 : end * 3], dtype=np.int64) + if values.size == 0: + return np.zeros((0, 3), dtype=np.int32) + triples = values.reshape(-1, 3) + if np.any(triples < 0) or np.any(triples > np.iinfo(np.int32).max): + raise ValueError(f"{storage.key} graph edge values must fit non-negative int32") + return triples.astype(np.int32, copy=False) + + def _read_graph_vector(storage: _GraphSidecarStorage, sequence_index: int) -> np.ndarray: if storage.kind != "ragged_1d" or storage.shape_tail: raise ValueError(f"{storage.key} graph storage is not ragged_1d") @@ -1397,7 +1525,7 @@ def _allowed_side_channel_dtype_names(key: str) -> list[str]: def _side_channel_family(key: str) -> str: - if key == _ATTENTION_SIDE_CHANNEL_KEY: + if key in {_ATTENTION_SIDE_CHANNEL_KEY, _LOSS_MASK_SIDE_CHANNEL_KEY}: return "attention" if key in _STRUCTURE_SIDE_CHANNEL_KEYS: return "structure" @@ -1407,14 +1535,18 @@ def _side_channel_family(key: str) -> str: return "semantic_graph" if key in _TEMPORAL_SIDE_CHANNEL_KEYS: return "temporal_diff" + if key in _DOMAIN_SIDE_CHANNEL_KEYS: + return "domain_routes" raise ValueError(f"unknown side-channel key {key!r}") def _to_side_channel_values(key: str, values: np.ndarray) -> np.ndarray: if key == _ATTENTION_SIDE_CHANNEL_KEY: return values.astype(np.float32, copy=False) - if np.any(values < 0): + if key != "hunk_ids" and np.any(values < 0): raise ValueError(f"{key} side-channel IDs must be non-negative") + if key == "hunk_ids" and np.any(values < np.iinfo(np.int32).min): + raise ValueError("hunk_ids side-channel values are below int32 range") if np.any(values > np.iinfo(np.int32).max): raise ValueError(f"{key} side-channel IDs exceed int32 range") return values.astype(np.int32, copy=False) @@ -1584,6 +1716,40 @@ def _validate_bin_references( raise ValueError("Megatron .idx sequence pointers must align to dtype size") +def _is_compact_fixed_row_shard( + sidecar: dict[str, Any], + *, + sequence_lengths: np.ndarray, + seq_len: int, +) -> bool: + """Recognize converter output that omits only fixed-row padding.""" + + declared_capacity = sidecar.get("source_capacity_token_count") + if declared_capacity is None: + return False + + sequence_count = int(sequence_lengths.shape[0]) + expected_capacity = sequence_count * seq_len + if int(declared_capacity) != expected_capacity: + raise ValueError( + "source_capacity_token_count does not match sequence_count * seq_len" + ) + if int(sidecar.get("document_count", -1)) != sequence_count: + raise ValueError( + "compact fixed-row document_count does not match MMIDIDX sequence count" + ) + token_count = int(sequence_lengths.sum(dtype=np.int64)) + if int(sidecar.get("token_count", -1)) != token_count: + raise ValueError( + "compact fixed-row token_count does not match MMIDIDX sequence lengths" + ) + if np.any(sequence_lengths <= 0) or np.any(sequence_lengths > seq_len): + raise ValueError( + "compact fixed-row sequence lengths must satisfy 0 < length <= seq_len" + ) + return True + + def _build_windows( sequence_offsets: np.ndarray, sequence_lengths: np.ndarray, diff --git a/cppmega_mlx/inference/generation.py b/cppmega_mlx/inference/generation.py index bc76d21f..5338a1c2 100644 --- a/cppmega_mlx/inference/generation.py +++ b/cppmega_mlx/inference/generation.py @@ -764,10 +764,18 @@ def _resolve_kv_cache( def _standard_generation_logits(model_output: Any, tokens: mx.array) -> mx.array: if isinstance(model_output, tuple | list): - raise ValueError( - "MTP/draft tuple outputs are not supported by standard next-token " - "inference" - ) + if ( + len(model_output) == 2 + and isinstance(model_output[0], mx.array) + and model_output[1] is None + ): + model_output = model_output[0] + else: + raise ValueError( + "MTP/draft tuple outputs are not supported by standard next-token " + "inference; only the DenseCppLM (logits, None) inference contract " + "is accepted" + ) if isinstance(model_output, dict): raise ValueError( "structured model outputs are not supported by standard next-token " diff --git a/cppmega_mlx/nn/domain_graph_routes.py b/cppmega_mlx/nn/domain_graph_routes.py index 4aae4d74..d624db71 100644 --- a/cppmega_mlx/nn/domain_graph_routes.py +++ b/cppmega_mlx/nn/domain_graph_routes.py @@ -16,7 +16,9 @@ class DomainGraphRouteConfig: num_blocks: int = 64 normalize: str = "binary" - edge_weights: Mapping[int, float] = field(default_factory=dict) + edge_weights: Mapping[int, float] = field( + default_factory=lambda: DEFAULT_EDGE_WEIGHTS.copy() + ) default_weight: float = 1.0 max_candidates_per_query: int | None = None diff --git a/cppmega_mlx/nn/mamba3.py b/cppmega_mlx/nn/mamba3.py index 3df95bea..44803122 100644 --- a/cppmega_mlx/nn/mamba3.py +++ b/cppmega_mlx/nn/mamba3.py @@ -523,7 +523,6 @@ def _dispatch_mamba3_scan( lowered TileLang DSL fwd+bwd kernel and fails closed if unavailable. """ - from cppmega_mlx.nn._tilelang.mamba3 import mamba3_mimo_metal_status from cppmega_mlx.runtime.kernel_policy import ( KernelPath, record_dispatch, @@ -576,6 +575,8 @@ def _dispatch_mamba3_scan( ) # AUTO + PATH_B share the Metal-availability check. + from cppmega_mlx.nn._tilelang.mamba3 import mamba3_mimo_metal_status + status = mamba3_mimo_metal_status(x) if path is KernelPath.PATH_B and not status.available: # Device-aware Path B: on a CUDA host (Metal unavailable, default diff --git a/evals/domain_routed_prompts.jsonl b/evals/domain_routed_prompts.jsonl index 2134c088..c8309fa8 100644 --- a/evals/domain_routed_prompts.jsonl +++ b/evals/domain_routed_prompts.jsonl @@ -1,4 +1,4 @@ -{"id":"cpp_docstring_add","task_type":"cpp_docstring_to_code","required_domains":["CPP"],"prompt":"// Return the sum of two integers.\nint add(int a, int b) {\n","compile_prefix":"","compile_suffix":"\nint main(){return add(2,3)==5 ? 0 : 1;}\n","expected_sidecars":["token_domain_ids","token_role_ids","token_call_edges"]} +{"id":"cpp_docstring_add","task_type":"cpp_docstring_to_code","required_domains":["CPP"],"prompt":"// Return the sum of two integers.\nint add(int a, int b) {\n","compile_prefix":"","compile_suffix":"\nint main(){return add(2,3)==5 ? 0 : 1;}\n","run_binary":true,"expected_sidecars":["token_domain_ids","token_role_ids","token_call_edges"]} {"id":"fim_compile_fix_missing_header","task_type":"fim_compile_fix","required_domains":["CPP","COMPILER_ERROR"],"prompt":" src/main.cpp:1:10: fatal error: 'vector' file not found \n fix the compile error by adding the missing standard header\nint main(){ std::vector xs; return xs.size(); }\n\n","compile_prefix":"","compile_suffix":"","expected_sidecars":["token_diagnostic_edges","token_cross_domain_edges"]} {"id":"cmake_link_library_fix","task_type":"cmake_fix","required_domains":["CPP","CMAKE","BUILD_ERROR","LINKER_ERROR"],"prompt":" add_executable(app main.cpp) \n undefined reference to `pthread_create' \n update CMake to link the required thread library\nadd_executable(app main.cpp)\n\n","compile_prefix":"","compile_suffix":"","expected_sidecars":["token_build_edges","token_diagnostic_edges","token_cross_domain_edges"]} {"id":"shell_build_command","task_type":"shell_build_understanding","required_domains":["SH","BUILD_DIAGNOSTIC"],"prompt":" cmake --build build -j8 \n ninja: build stopped: subcommand failed \nExplain the failing build command target in one line.","expected_sidecars":["token_shell_edges","token_diagnostic_edges"]} diff --git a/scripts/audit_sidecar_parquet.py b/scripts/audit_sidecar_parquet.py index 21c420d0..d7d6cc8f 100644 --- a/scripts/audit_sidecar_parquet.py +++ b/scripts/audit_sidecar_parquet.py @@ -15,12 +15,17 @@ import json import os from pathlib import Path +import sys from typing import Any import numpy as np import pyarrow.compute as pc import pyarrow.parquet as pq +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + from cppmega_mlx.data.domain_schema import DomainKind, DomainRoleKind, ParseConfidence from cppmega_mlx.data.tokenizer_contract import DOMAIN_DELIMITER_TOKEN_IDS @@ -378,70 +383,74 @@ def _count_bad_flat_values( row_bad |= rows -def _validate_loss_mask_against_doc_ids( +def _validate_packed_document_boundaries( *, stats: AuditStats, loss_mask_col: Any, doc_ids_col: Any, + source_doc_token_lengths_col: Any, input_lengths: np.ndarray, valid: np.ndarray, row_bad: np.ndarray, ) -> None: - """Value-level check that `loss_mask` equals the producer's doc-boundary rule. - - This is the check that catches C1-style corruption (an all-ones loss_mask - rebuilt over a multi-document packed row). Lengths are preserved by that - corruption, so the length checks alone cannot see it; only a value-level - comparison against `doc_ids` can. - - Canonical rule (see `pack_enriched_rows._loss_mask_for_packed_docs`): + """Validate boundaries against independent logical-document lengths. - loss_mask[pos] == 1 iff pos + 1 < valid AND doc_ids[pos] == doc_ids[pos + 1] - - i.e. 1 only when the next token belongs to the SAME document and is itself a - real (non-pad) token; 0 at every inter-document boundary, at the last valid - token (no next token to predict), and across the entire pad region. + ``doc_ids`` cannot be its own oracle: an old producer assigned the same ID + to distinct functions sharing file provenance, and a mask derived from those + IDs looked internally consistent while still training across documents. + ``source_doc_token_lengths`` is the independent packed-document contract. """ - rows = len(input_lengths) - if rows == 0: - return - lengths = input_lengths.astype(np.int64) - n = int(lengths.sum()) - doc_flat = _flat_numpy(doc_ids_col) - lm_flat = _flat_numpy(loss_mask_col) - # The vectorized derivation below is only meaningful when both columns share - # the canonical per-row layout (length == input_ids length). Any row whose - # doc_ids/loss_mask length differs is already flagged as a bad row by the - # length checks (and blocks the upload under the fail-closed default), so a - # layout-inconsistent file is ALREADY failing the gate -- skip deriving the - # mask against wrong offsets rather than fabricate a value comparison. - if len(doc_flat) != n or len(lm_flat) != n: - return - doc_flat = doc_flat.astype(np.int64) - lm_flat = lm_flat.astype(np.int64) + doc_rows = doc_ids_col.to_pylist() + mask_rows = loss_mask_col.to_pylist() + source_length_rows = source_doc_token_lengths_col.to_pylist() + bad_doc_ids = 0 + bad_masks = 0 + + for row_index, (capacity, valid_count) in enumerate(zip(input_lengths, valid)): + capacity = int(capacity) + valid_count = int(valid_count) + raw_lengths = source_length_rows[row_index] + if raw_lengths is None: + row_bad[row_index] = True + stats.errors.append(f"row {row_index}: missing source_doc_token_lengths") + continue + logical_lengths = [int(value) for value in raw_lengths] + if ( + (valid_count > 0 and not logical_lengths) + or any(length <= 0 for length in logical_lengths) + or sum(logical_lengths) != valid_count + ): + row_bad[row_index] = True + stats.field_stats["source_doc_token_lengths"].bad_value_rows += 1 + stats.errors.append( + f"row {row_index}: source_doc_token_lengths={logical_lengths!r} " + f"do not partition valid_token_count={valid_count}" + ) + continue - row_idx = np.repeat(np.arange(rows, dtype=np.int64), lengths) - starts = np.zeros(rows, dtype=np.int64) - if rows > 1: - np.cumsum(lengths[:-1], out=starts[1:]) - pos = np.arange(n, dtype=np.int64) - starts[row_idx] - - same_next = np.zeros(n, dtype=bool) - if n > 1: - same_next[:-1] = doc_flat[:-1] == doc_flat[1:] - # Never let the neighbour comparison cross a row boundary. - is_row_last = pos == (lengths[row_idx] - 1) - same_next &= ~is_row_last - # pos + 1 < valid (excludes the last valid token and the whole pad region). - within_valid = pos < (valid[row_idx].astype(np.int64) - 1) - - expected = (same_next & within_valid).astype(np.int64) - mismatch = lm_flat != expected - bad_rows = _rows_with_flat_mask(lengths, mismatch) - count = int(np.count_nonzero(bad_rows)) - if count: - stats.field_stats["loss_mask"].bad_value_rows += count - row_bad |= bad_rows + expected_doc_ids: list[int] = [] + expected_loss_mask: list[int] = [] + for local_doc_id, length in enumerate(logical_lengths, start=1): + expected_doc_ids.extend([local_doc_id] * length) + expected_loss_mask.extend([1] * (length - 1)) + expected_loss_mask.append(0) + pad_doc_id = len(logical_lengths) if logical_lengths else 0 + expected_doc_ids.extend([pad_doc_id] * (capacity - valid_count)) + expected_loss_mask.extend([0] * (capacity - valid_count)) + + actual_doc_ids = [int(value) for value in (doc_rows[row_index] or [])] + actual_loss_mask = [int(value) for value in (mask_rows[row_index] or [])] + if actual_doc_ids != expected_doc_ids: + row_bad[row_index] = True + bad_doc_ids += 1 + if actual_loss_mask != expected_loss_mask: + row_bad[row_index] = True + bad_masks += 1 + + if bad_doc_ids: + stats.field_stats["doc_ids"].bad_value_rows += bad_doc_ids + if bad_masks: + stats.field_stats["loss_mask"].bad_value_rows += bad_masks def _validate_target_ids_shift( @@ -689,47 +698,25 @@ def _validate_diagnostic_rows_have_edges_or_raw( ) -def _audit_file(path_str: str, kind: str, bucket: str, vocab_size: int | None) -> dict[str, Any]: - path = Path(path_str) - stats = AuditStats(files=1) - try: - pf = pq.ParquetFile(path) - top_level_names = set(pf.schema_arrow.names) - cols = [name for name in ALL_FIELDS if name in top_level_names] - cols += [ - name - for name in ("valid_token_count", "trained_token_count", "slack_tokens") - if name in top_level_names - ] - table = pf.read(columns=cols) - except Exception as exc: - # Fail loud: a shard we cannot even read cannot be certified for upload. - # Re-raise with where+what so the gate crashes instead of recording a - # degraded "bad file" entry that a caller might overlook. - raise RuntimeError( - f"{path}: parquet read failed: {type(exc).__name__}: {exc}" - ) from exc - - names = set(cols) - for name in ALL_FIELDS: - if name not in names: - stats.field_stats[name].missing_files += 1 +def _audit_table( + *, + path: Path, + table: object, + names: set[str], + expected_len: int | None, + vocab_size: int | None, +) -> AuditStats: + """Audit one bounded parquet row group and return additive statistics.""" - expected_len = int(bucket) if bucket.isdigit() else None + stats = AuditStats() stats.rows = table.num_rows row_bad = np.zeros(stats.rows, dtype=bool) try: if "input_ids" not in names: - stats.bad_files += 1 stats.bad_rows = stats.rows stats.errors.append(f"{path}: missing input_ids") - return { - "kind": kind, - "bucket": bucket, - "path": path_str, - "stats": stats.as_dict(), - } + return stats input_ids = table.column("input_ids") input_lengths = _add_numeric_list_stats( @@ -804,14 +791,14 @@ def _audit_file(path_str: str, kind: str, bucket: str, vocab_size: int | None) - row_bad=row_bad, ) - # Value-level loss-target correctness: loss_mask MUST match the - # doc_ids-derived boundary rule. Length checks alone cannot catch a - # corrupted (e.g. all-ones) loss_mask because lengths are preserved. - if "loss_mask" in names and "doc_ids" in names: - _validate_loss_mask_against_doc_ids( + # Derive both doc_ids and loss_mask from logical source-document lengths. + # Never trust doc_ids as the oracle for its own boundary correctness. + if all(name in names for name in ("loss_mask", "doc_ids", "source_doc_token_lengths")): + _validate_packed_document_boundaries( stats=stats, loss_mask_col=table.column("loss_mask"), doc_ids_col=table.column("doc_ids"), + source_doc_token_lengths_col=table.column("source_doc_token_lengths"), input_lengths=input_lengths, valid=valid, row_bad=row_bad, @@ -894,6 +881,20 @@ def _audit_file(path_str: str, kind: str, bucket: str, vocab_size: int | None) - bad_mask=bad_ends, row_bad=row_bad, ) + if "token_chunk_starts" in names and "token_chunk_ends" in names: + start_rows = table.column("token_chunk_starts").to_pylist() + end_rows = table.column("token_chunk_ends").to_pylist() + for row_idx, (row_starts, row_ends) in enumerate( + zip(start_rows, end_rows, strict=True) + ): + if len(row_starts or []) != len(row_ends or []): + continue + if any( + int(start) >= int(end) + for start, end in zip(row_starts or [], row_ends or [], strict=True) + ): + row_bad[row_idx] = True + stats.field_stats["token_chunk_ends"].bad_value_rows += 1 if "platform_ids" in names: _add_numeric_list_stats( @@ -969,6 +970,56 @@ def _audit_file(path_str: str, kind: str, bucket: str, vocab_size: int | None) - ) from exc stats.bad_rows = int(np.count_nonzero(row_bad)) + return stats + + +def _audit_file(path_str: str, kind: str, bucket: str, vocab_size: int | None) -> dict[str, Any]: + path = Path(path_str) + stats = AuditStats(files=1) + try: + pf = pq.ParquetFile(path) + top_level_names = set(pf.schema_arrow.names) + cols = [name for name in ALL_FIELDS if name in top_level_names] + cols += [ + name + for name in ("valid_token_count", "trained_token_count", "slack_tokens") + if name in top_level_names + ] + except Exception as exc: + raise RuntimeError( + f"{path}: parquet open failed: {type(exc).__name__}: {exc}" + ) from exc + + names = set(cols) + for name in ALL_FIELDS: + if name not in names: + stats.field_stats[name].missing_files += 1 + + if "input_ids" not in names: + stats.rows = int(pf.metadata.num_rows) + stats.bad_files = 1 + stats.bad_rows = stats.rows + stats.errors.append(f"{path}: missing input_ids") + else: + expected_len = int(bucket) if bucket.isdigit() else None + for row_group_idx in range(pf.metadata.num_row_groups): + try: + table = pf.read_row_group(row_group_idx, columns=cols) + except Exception as exc: + raise RuntimeError( + f"{path}#row_group{row_group_idx}: parquet read failed: " + f"{type(exc).__name__}: {exc}" + ) from exc + stats.add( + _audit_table( + path=path, + table=table, + names=names, + expected_len=expected_len, + vocab_size=vocab_size, + ) + ) + return { "kind": kind, "bucket": bucket, diff --git a/scripts/cpp_jsonl_generation_compile_eval.py b/scripts/cpp_jsonl_generation_compile_eval.py index 68e5ab73..c84ef992 100644 --- a/scripts/cpp_jsonl_generation_compile_eval.py +++ b/scripts/cpp_jsonl_generation_compile_eval.py @@ -194,6 +194,14 @@ def __init__( tokenizer.tool_result_id, tokenizer.code_start_id, } + token_for_id = getattr(tokenizer, "token_for_id", None) + vocab_size = getattr(tokenizer, "vocab_size", None) + if callable(token_for_id) and isinstance(vocab_size, int): + self._always_banned.update( + token_id + for token_id in range(vocab_size) + if str(token_for_id(token_id) or "").startswith(" mx.array: if logits.ndim != 2 or tokens.ndim != 2: @@ -255,20 +263,88 @@ def trim_body_completion(text: str) -> str: "", "", "= 0: stripped = stripped[:pos] - kept: list[str] = [] - for line in stripped.splitlines(): - if line.startswith("}"): - break - kept.append(line) - body = "\n".join(kept).strip() + stripped = _trim_at_function_closing_brace(stripped) + body = stripped.strip() return body + ("\n" if body else "") +def _trim_at_function_closing_brace(text: str) -> str: + """Drop only the brace that closes the prompt's already-open function.""" + depth = 1 + index = 0 + state = "code" + while index < len(text): + char = text[index] + following = text[index + 1] if index + 1 < len(text) else "" + + if state == "line-comment": + if char == "\n": + state = "code" + index += 1 + continue + if state == "block-comment": + if char == "*" and following == "/": + state = "code" + index += 2 + else: + index += 1 + continue + if state in {"string", "character"}: + quote = '"' if state == "string" else "'" + if char == "\\": + index += 2 + elif char == quote: + state = "code" + index += 1 + else: + index += 1 + continue + + if char == "/" and following == "/": + state = "line-comment" + index += 2 + continue + if char == "/" and following == "*": + state = "block-comment" + index += 2 + continue + if char == "R" and following == '"': + delimiter_end = text.find("(", index + 2, min(len(text), index + 20)) + if delimiter_end >= 0: + delimiter = text[index + 2 : delimiter_end] + if len(delimiter) <= 16 and not any( + item.isspace() or item in "()\\" for item in delimiter + ): + terminator = ")" + delimiter + '"' + raw_end = text.find(terminator, delimiter_end + 1) + if raw_end < 0: + return text + index = raw_end + len(terminator) + continue + if char == '"': + state = "string" + index += 1 + continue + if char == "'": + state = "character" + index += 1 + continue + if char == "{": + depth += 1 + elif char == "}": + depth -= 1 + if depth == 0: + return text[:index] + index += 1 + return text + + def reject_unsupported_checkpoint(path: Path) -> None: """Fail closed for Megatron distributed checkpoints. diff --git a/scripts/eval_domain_routed_codegen.py b/scripts/eval_domain_routed_codegen.py index a51444ea..0dd4f54d 100644 --- a/scripts/eval_domain_routed_codegen.py +++ b/scripts/eval_domain_routed_codegen.py @@ -27,6 +27,7 @@ class DomainEvalPrompt: expected_sidecars: tuple[str, ...] compile_prefix: str = "" compile_suffix: str = "" + run_binary: bool = False @classmethod def from_row(cls, row: dict[str, Any], *, path: Path, line_no: int) -> "DomainEvalPrompt": @@ -44,6 +45,7 @@ def from_row(cls, row: dict[str, Any], *, path: Path, line_no: int) -> "DomainEv expected_sidecars=tuple(str(x) for x in row.get("expected_sidecars", ())), compile_prefix=str(row.get("compile_prefix", "")), compile_suffix=str(row.get("compile_suffix", "")), + run_binary=bool(row.get("run_binary", False)), ) @@ -90,7 +92,14 @@ def _find_cpp_compiler() -> str | None: return None -def compile_cpp_completion(prompt: DomainEvalPrompt, completion: str, *, compiler: str | None = None) -> dict[str, Any]: +def compile_cpp_completion( + prompt: DomainEvalPrompt, + completion: str, + *, + compiler: str | None = None, + compile_timeout_s: float = 120.0, + run_timeout_s: float = 10.0, +) -> dict[str, Any]: compiler = compiler or _find_cpp_compiler() if compiler is None: return {"status": "compile_skipped", "reason": "no local C++ compiler found"} @@ -99,14 +108,52 @@ def compile_cpp_completion(prompt: DomainEvalPrompt, completion: str, *, compile src = Path(tmp) / "main.cpp" exe = Path(tmp) / "a.out" src.write_text(source, encoding="utf-8") - proc = subprocess.run( - [compiler, "-std=c++20", "-Wall", "-Wextra", str(src), "-o", str(exe)], - capture_output=True, - text=True, - timeout=30, - ) + try: + proc = subprocess.run( + [compiler, "-std=c++20", "-Wall", "-Wextra", str(src), "-o", str(exe)], + capture_output=True, + text=True, + timeout=compile_timeout_s, + ) + except subprocess.TimeoutExpired as exc: + return { + "status": "compile_timeout", + "timeout_s": compile_timeout_s, + "stderr": str(exc.stderr or "")[-4000:], + } + if proc.returncode != 0: + return { + "status": "compile_failed", + "returncode": proc.returncode, + "stderr": proc.stderr[-4000:], + } + if prompt.run_binary: + try: + run = subprocess.run( + [str(exe)], + capture_output=True, + text=True, + timeout=run_timeout_s, + ) + except subprocess.TimeoutExpired as exc: + return { + "status": "runtime_timeout", + "returncode": proc.returncode, + "timeout_s": run_timeout_s, + "stderr": proc.stderr[-4000:], + "runtime_stdout": str(exc.stdout or "")[-4000:], + "runtime_stderr": str(exc.stderr or "")[-4000:], + } + return { + "status": "compile_passed" if run.returncode == 0 else "runtime_failed", + "returncode": proc.returncode, + "runtime_returncode": run.returncode, + "stderr": proc.stderr[-4000:], + "runtime_stdout": run.stdout[-4000:], + "runtime_stderr": run.stderr[-4000:], + } return { - "status": "compile_passed" if proc.returncode == 0 else "compile_failed", + "status": "compile_passed", "returncode": proc.returncode, "stderr": proc.stderr[-4000:], } @@ -139,12 +186,14 @@ def evaluate( rows.append(row) passed = sum(1 for row in rows if row.get("status") == "compile_passed") failed = sum(1 for row in rows if row.get("status") == "compile_failed") + runtime_failed = sum(1 for row in rows if row.get("status") == "runtime_failed") missing = sum(1 for row in rows if row.get("status") == "missing_completion") return { "prompts": len(prompts), "completion_rows": len(completions), "compile_passed": passed, "compile_failed": failed, + "runtime_failed": runtime_failed, "missing_completion": missing, "rows": rows, } @@ -163,8 +212,8 @@ def main() -> int: report = evaluate(prompts, completions, compile=not args.no_compile) args.out.parent.mkdir(parents=True, exist_ok=True) args.out.write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8") - print(json.dumps({k: report[k] for k in ("prompts", "completion_rows", "compile_passed", "compile_failed", "missing_completion")}, sort_keys=True)) - return 2 if report["compile_failed"] else 0 + print(json.dumps({k: report[k] for k in ("prompts", "completion_rows", "compile_passed", "compile_failed", "runtime_failed", "missing_completion")}, sort_keys=True)) + return 2 if report["compile_failed"] or report["runtime_failed"] else 0 if __name__ == "__main__": diff --git a/scripts/fix_packed_parquet_boundaries.py b/scripts/fix_packed_parquet_boundaries.py index 7d0d4055..41372a6d 100644 --- a/scripts/fix_packed_parquet_boundaries.py +++ b/scripts/fix_packed_parquet_boundaries.py @@ -282,27 +282,30 @@ def _shift_span(span: Any, positions: list[int]) -> Any: def _update_source_doc_lengths(row: dict[str, Any], positions: list[int], valid: int) -> None: lengths = row.get("source_doc_token_lengths") - doc_ids = row.get("doc_ids") - source_doc_ids = row.get("source_doc_ids") - if not isinstance(lengths, list) or not isinstance(doc_ids, list): + if not isinstance(lengths, list): return if not lengths: return if len(lengths) == 1: row["source_doc_token_lengths"] = [int(lengths[0]) + len(positions)] return - if not isinstance(source_doc_ids, list): - return - increments = {int(doc_id): 0 for doc_id in source_doc_ids} + normalized_lengths = [int(length) for length in lengths] + increments = [0] * len(normalized_lengths) + cumulative: list[int] = [] + running = 0 + for length in normalized_lengths: + running += length + cumulative.append(running) for pos in positions: source_idx = pos - 1 if pos > 0 else pos - if 0 <= source_idx < min(valid, len(doc_ids)): - doc_id = int(doc_ids[source_idx]) - if doc_id in increments: - increments[doc_id] += 1 + if 0 <= source_idx < valid: + for doc_index, end in enumerate(cumulative): + if source_idx < end: + increments[doc_index] += 1 + break row["source_doc_token_lengths"] = [ - int(length) + increments.get(int(doc_id), 0) - for length, doc_id in zip(lengths, source_doc_ids) + length + increments[index] + for index, length in enumerate(normalized_lengths) ] diff --git a/scripts/nanochat_data/pack_enriched_rows.py b/scripts/nanochat_data/pack_enriched_rows.py index 44e39e6e..c4849e5a 100644 --- a/scripts/nanochat_data/pack_enriched_rows.py +++ b/scripts/nanochat_data/pack_enriched_rows.py @@ -104,7 +104,6 @@ TOKEN_TYPE_EDGES_COLUMN, TOKEN_TYPE_REFS_COLUMN, ) -from cppmega_mlx.data.nanochat_pipeline.platform_vocab import MAX_PLATFORM_IDS from cppmega_mlx.data.packing import PackingStrategy TRAINED_TOKEN_COUNT_COLUMN = "trained_token_count" @@ -540,7 +539,12 @@ def _shared_chronology_for_docs(docs: list[NormalizedDoc]) -> dict[str, Any]: def _merged_platform_ids_for_docs(docs: list[NormalizedDoc]) -> list[int]: - ids = sorted( + # This row-level value is provenance, not the fixed-width model input. + # Mixed packed rows may legitimately describe more platforms than one + # source document. Exact per-document bags remain in + # ``source_platform_ids`` and are expanded through row-local ``doc_ids`` by + # the data loader, so capping the union here only makes valid packs fail. + return sorted( { int(platform_id) for doc in docs @@ -548,12 +552,6 @@ def _merged_platform_ids_for_docs(docs: list[NormalizedDoc]) -> list[int]: if int(platform_id) > 0 } ) - if len(ids) > MAX_PLATFORM_IDS: - raise ValueError( - f"packed row platform_ids has {len(ids)} unique IDs; " - f"MAX_PLATFORM_IDS={MAX_PLATFORM_IDS}" - ) - return ids def _coerce_optional_int(value: object) -> int | None: @@ -1274,10 +1272,14 @@ def _materialize_packed_row( token_offset = 0 chunk_offset = 0 - for doc in ordered_docs: - doc_id = max(0, int(doc.stable_doc_id)) + # ``doc_ids`` is a row-local attention/loss boundary channel. It must + # identify every logical packed document independently, even when two + # documents share file-level provenance (for example two functions from the + # same header). ``stable_doc_id`` remains provenance metadata; using it here + # collapsed those functions into one segment and trained across the boundary. + for row_doc_id, doc in enumerate(ordered_docs, start=1): concatenated_tokens.extend(doc.token_ids) - doc_ids.extend([doc_id] * doc.token_count) + doc_ids.extend([row_doc_id] * doc.token_count) source_doc_indices.append(doc.source_doc_index) source_doc_token_lengths.append(doc.token_count) source_platform_ids.append(list(doc.platform_ids)) diff --git a/scripts/pr_ingest/pr_store.py b/scripts/pr_ingest/pr_store.py index 88be6e89..e279fbab 100755 --- a/scripts/pr_ingest/pr_store.py +++ b/scripts/pr_ingest/pr_store.py @@ -49,6 +49,7 @@ import glob import json import os +from pathlib import Path import sqlite3 import sys from typing import Optional @@ -100,13 +101,24 @@ """ -def connect(store_path: str, create: bool = True) -> sqlite3.Connection: +def connect( + store_path: str, + create: bool = True, + *, + readonly: bool = False, +) -> sqlite3.Connection: + if readonly and create: + raise ValueError("readonly PR store connections cannot create a database") if not create and not os.path.exists(store_path): raise SystemExit(f"[pr_store] store does not exist: {store_path}") if create: d = os.path.dirname(os.path.abspath(store_path)) os.makedirs(d, exist_ok=True) - conn = sqlite3.connect(store_path, timeout=60.0) + if readonly: + uri = Path(store_path).resolve().as_uri() + "?mode=ro" + conn = sqlite3.connect(uri, uri=True, timeout=60.0) + else: + conn = sqlite3.connect(store_path, timeout=60.0) conn.row_factory = sqlite3.Row # Set busy_timeout FIRST. The parallel graphql_pr_stream --workers path opens # many writer connections concurrently; with the timeout in place the WAL @@ -114,6 +126,8 @@ def connect(store_path: str, create: bool = True) -> sqlite3.Connection: # instead of erroring out with "database is locked" (fail-loud preserved: # SQLite still raises if the lock is not released within the timeout). conn.execute("PRAGMA busy_timeout=60000") + if readonly: + return conn conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA synchronous=NORMAL") conn.executescript(SCHEMA) diff --git a/scripts/repair_packed_document_boundaries.py b/scripts/repair_packed_document_boundaries.py new file mode 100644 index 00000000..99fb82de --- /dev/null +++ b/scripts/repair_packed_document_boundaries.py @@ -0,0 +1,348 @@ +#!/usr/bin/env python3 +"""Repair packed-row document boundaries from logical source lengths. + +Older packed shards reused a file-level stable ID for multiple logical +documents from the same source file. Their ``doc_ids`` and derived +``loss_mask`` were internally consistent but trained across function/document +boundaries. This tool treats ``source_doc_token_lengths`` as the independent +truth, rebuilds row-local ``doc_ids`` and ``loss_mask``, and updates +``trained_token_count``. + +Files are scanned before any write. Changed parquet files are rewritten one +row group at a time to a sibling temporary file and published with +``os.replace``. This makes it safe to run against a hard-linked snapshot: the +live source inode is never modified in place. +""" + +from __future__ import annotations + +import argparse +from concurrent.futures import ProcessPoolExecutor, as_completed +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +import json +import os +from pathlib import Path +import tempfile +from typing import Any, Iterable + +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq + + +BOUNDARY_COLUMNS = ( + "input_ids", + "doc_ids", + "loss_mask", + "source_doc_ids", + "source_doc_token_lengths", + "valid_token_count", + "trained_token_count", + "num_docs", +) + + +@dataclass(frozen=True) +class FileScan: + path: str + rows: int + changed_rows: int + restored_boundaries: int + old_trained_tokens: int + new_trained_tokens: int + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _write_json_atomic(path: Path, payload: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(f".{path.name}.tmp-{os.getpid()}") + tmp.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + os.replace(tmp, path) + + +def _canonical_boundary_values(row: dict[str, Any], *, where: str) -> tuple[list[int], list[int], int]: + input_ids = row.get("input_ids") + doc_ids = row.get("doc_ids") + loss_mask = row.get("loss_mask") + source_lengths = row.get("source_doc_token_lengths") + if not isinstance(input_ids, list) or not isinstance(doc_ids, list) or not isinstance(loss_mask, list): + raise ValueError(f"{where}: missing list input_ids/doc_ids/loss_mask") + capacity = len(input_ids) + if len(doc_ids) != capacity or len(loss_mask) != capacity: + raise ValueError( + f"{where}: token-aligned length mismatch: capacity={capacity} " + f"doc_ids={len(doc_ids)} loss_mask={len(loss_mask)}" + ) + if not isinstance(source_lengths, list): + raise ValueError(f"{where}: missing source_doc_token_lengths") + valid = int(row.get("valid_token_count", capacity)) + source_doc_ids = row.get("source_doc_ids") + return _canonical_from_lengths( + capacity=capacity, + valid=valid, + logical_lengths=[int(value) for value in source_lengths], + source_doc_count=len(source_doc_ids) if isinstance(source_doc_ids, list) else None, + num_docs=int(row["num_docs"]) if row.get("num_docs") is not None else None, + where=where, + ) + + +def _canonical_from_lengths( + *, + capacity: int, + valid: int, + logical_lengths: list[int], + source_doc_count: int | None, + num_docs: int | None, + where: str, +) -> tuple[list[int], list[int], int]: + if valid < 0 or valid > capacity: + raise ValueError(f"{where}: invalid valid_token_count={valid} capacity={capacity}") + if (valid > 0 and not logical_lengths) or any(length <= 0 for length in logical_lengths): + raise ValueError(f"{where}: invalid source_doc_token_lengths={logical_lengths!r}") + if sum(logical_lengths) != valid: + raise ValueError( + f"{where}: source_doc_token_lengths sum={sum(logical_lengths)} " + f"!= valid_token_count={valid}" + ) + if source_doc_count is not None and source_doc_count != len(logical_lengths): + raise ValueError( + f"{where}: source_doc_ids count={source_doc_count} " + f"!= source_doc_token_lengths count={len(logical_lengths)}" + ) + if num_docs is not None and num_docs != len(logical_lengths): + raise ValueError( + f"{where}: num_docs={num_docs} != source_doc_token_lengths count={len(logical_lengths)}" + ) + + expected_doc_ids: list[int] = [] + expected_loss_mask: list[int] = [] + for local_doc_id, length in enumerate(logical_lengths, start=1): + expected_doc_ids.extend([local_doc_id] * length) + expected_loss_mask.extend([1] * (length - 1)) + expected_loss_mask.append(0) + pad_doc_id = len(logical_lengths) if logical_lengths else 0 + expected_doc_ids.extend([pad_doc_id] * (capacity - valid)) + expected_loss_mask.extend([0] * (capacity - valid)) + return expected_doc_ids, expected_loss_mask, sum(expected_loss_mask) + + +def repair_row(row: dict[str, Any], *, where: str = "row") -> tuple[dict[str, Any], bool, int]: + expected_doc_ids, expected_loss_mask, trained = _canonical_boundary_values(row, where=where) + changed = ( + [int(value) for value in row["doc_ids"]] != expected_doc_ids + or [int(value) for value in row["loss_mask"]] != expected_loss_mask + or int(row.get("trained_token_count", -1)) != trained + ) + old_mask = [int(value) for value in row["loss_mask"]] + restored = sum( + 1 + for old, new in zip(old_mask, expected_loss_mask) + if old != 0 and new == 0 + ) + if not changed: + return row, False, 0 + repaired = dict(row) + repaired["doc_ids"] = expected_doc_ids + repaired["loss_mask"] = expected_loss_mask + repaired["trained_token_count"] = trained + return repaired, True, restored + + +def _scan_file(path_str: str) -> FileScan: + path = Path(path_str) + parquet = pq.ParquetFile(path) + missing = [name for name in BOUNDARY_COLUMNS[:7] if name not in parquet.schema_arrow.names] + if missing: + raise ValueError(f"{path}: missing required columns: {', '.join(missing)}") + selected = [name for name in BOUNDARY_COLUMNS if name in parquet.schema_arrow.names] + rows = changed_rows = restored_boundaries = 0 + old_trained = new_trained = 0 + bucket = int(path.parent.name) if path.parent.name.isdigit() else 16384 + batch_size = max(1, 131_072 // bucket) + for batch in parquet.iter_batches(columns=selected, batch_size=batch_size): + for row in batch.to_pylist(): + where = f"{path}:row={rows}" + repaired, changed, restored = repair_row(row, where=where) + rows += 1 + changed_rows += int(changed) + restored_boundaries += restored + old_trained += int(row.get("trained_token_count", sum(row["loss_mask"]))) + new_trained += int(repaired.get("trained_token_count", sum(repaired["loss_mask"]))) + return FileScan( + path=str(path), + rows=rows, + changed_rows=changed_rows, + restored_boundaries=restored_boundaries, + old_trained_tokens=old_trained, + new_trained_tokens=new_trained, + ) + + +def _rewrite_file(path_str: str, compression_level: int) -> None: + path = Path(path_str) + parquet = pq.ParquetFile(path) + schema = parquet.schema_arrow + fd, tmp_name = tempfile.mkstemp(prefix=f".{path.name}.repair-", suffix=".parquet", dir=path.parent) + os.close(fd) + tmp = Path(tmp_name) + writer: pq.ParquetWriter | None = None + try: + writer = pq.ParquetWriter( + tmp, + schema, + compression="zstd", + compression_level=compression_level, + use_dictionary=True, + ) + bucket = int(path.parent.name) if path.parent.name.isdigit() else 16384 + batch_size = max(1, 131_072 // bucket) + absolute_row = 0 + indices = {name: schema.get_field_index(name) for name in BOUNDARY_COLUMNS} + for batch in parquet.iter_batches(batch_size=batch_size): + table = pa.Table.from_batches([batch], schema=schema) + capacities = pc.list_value_length(batch.column(indices["input_ids"])).to_pylist() + logical_lengths_rows = batch.column( + indices["source_doc_token_lengths"] + ).to_pylist() + source_doc_ids_rows = batch.column(indices["source_doc_ids"]).to_pylist() + valids = batch.column(indices["valid_token_count"]).to_pylist() + num_docs_rows = ( + batch.column(indices["num_docs"]).to_pylist() + if indices["num_docs"] >= 0 + else [None] * len(batch) + ) + + doc_id_rows: list[list[int]] = [] + loss_mask_rows: list[list[int]] = [] + trained_rows: list[int] = [] + for offset, (capacity, logical_lengths, source_doc_ids, valid, num_docs) in enumerate( + zip( + capacities, + logical_lengths_rows, + source_doc_ids_rows, + valids, + num_docs_rows, + strict=True, + ) + ): + logical_lengths = [int(value) for value in logical_lengths or []] + expected_doc_ids, expected_loss_mask, trained = _canonical_from_lengths( + capacity=int(capacity), + valid=int(valid), + logical_lengths=logical_lengths, + source_doc_count=( + len(source_doc_ids) if isinstance(source_doc_ids, list) else None + ), + num_docs=int(num_docs) if num_docs is not None else None, + where=f"{path}:row={absolute_row + offset}", + ) + doc_id_rows.append(expected_doc_ids) + loss_mask_rows.append(expected_loss_mask) + trained_rows.append(trained) + + table = table.set_column( + indices["doc_ids"], + schema.field(indices["doc_ids"]), + pa.array(doc_id_rows, type=schema.field(indices["doc_ids"]).type), + ) + table = table.set_column( + indices["loss_mask"], + schema.field(indices["loss_mask"]), + pa.array(loss_mask_rows, type=schema.field(indices["loss_mask"]).type), + ) + table = table.set_column( + indices["trained_token_count"], + schema.field(indices["trained_token_count"]), + pa.array( + trained_rows, + type=schema.field(indices["trained_token_count"]).type, + ), + ) + writer.write_table(table) + absolute_row += len(batch) + writer.close() + writer = None + check = pq.ParquetFile(tmp) + if check.metadata.num_rows != parquet.metadata.num_rows or check.schema_arrow != schema: + raise RuntimeError(f"{path}: rewritten parquet validation failed") + os.replace(tmp, path) + finally: + if writer is not None: + writer.close() + tmp.unlink(missing_ok=True) + + +def _discover(roots: list[Path], buckets: tuple[int, ...]) -> list[Path]: + paths: list[Path] = [] + for root in roots: + for bucket in buckets: + bucket_root = root / str(bucket) + if not bucket_root.is_dir(): + raise FileNotFoundError(bucket_root) + paths.extend(sorted(bucket_root.glob("*.parquet"))) + if not paths: + raise RuntimeError("no parquet files selected") + return paths + + +def _parse_buckets(value: str) -> tuple[int, ...]: + buckets = tuple(int(part) for part in value.split(",") if part.strip()) + if not buckets or any(bucket <= 0 for bucket in buckets): + raise argparse.ArgumentTypeError("buckets must be a non-empty comma-separated list") + return buckets + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", action="append", type=Path, required=True) + parser.add_argument("--buckets", type=_parse_buckets, default=(1024, 2048, 4096, 8192, 16384)) + parser.add_argument("--workers", type=int, default=max(1, min(4, os.cpu_count() or 1))) + parser.add_argument("--compression-level", type=int, default=6) + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--receipt", type=Path) + return parser + + +def main(argv: Iterable[str] | None = None) -> int: + args = build_parser().parse_args(argv) + paths = _discover([root.resolve() for root in args.root], args.buckets) + scans: list[FileScan] = [] + with ProcessPoolExecutor(max_workers=max(1, args.workers)) as pool: + futures = {pool.submit(_scan_file, str(path)): path for path in paths} + for future in as_completed(futures): + scans.append(future.result()) + scans.sort(key=lambda item: item.path) + + changed = [scan for scan in scans if scan.changed_rows] + if changed and not args.dry_run: + with ProcessPoolExecutor(max_workers=max(1, args.workers)) as pool: + futures = [pool.submit(_rewrite_file, scan.path, args.compression_level) for scan in changed] + for future in as_completed(futures): + future.result() + + payload = { + "schema": "cppmega_packed_document_boundary_repair_v1", + "created_at": _utc_now(), + "dry_run": bool(args.dry_run), + "files": len(scans), + "rows": sum(scan.rows for scan in scans), + "changed_files": len(changed), + "changed_rows": sum(scan.changed_rows for scan in scans), + "restored_boundaries": sum(scan.restored_boundaries for scan in scans), + "old_trained_tokens": sum(scan.old_trained_tokens for scan in scans), + "new_trained_tokens": sum(scan.new_trained_tokens for scan in scans), + "file_scans": [asdict(scan) for scan in changed], + } + if args.receipt: + _write_json_atomic(args.receipt, payload) + print(json.dumps({key: value for key, value in payload.items() if key != "file_scans"}, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/streaming_conveyor.py b/scripts/streaming_conveyor.py index 5b5f8aec..7beed2b6 100644 --- a/scripts/streaming_conveyor.py +++ b/scripts/streaming_conveyor.py @@ -64,6 +64,7 @@ from concurrent.futures import ( FIRST_COMPLETED, CancelledError, + Future, ThreadPoolExecutor, wait, ) @@ -221,18 +222,41 @@ class BackgroundRecompressor: def __init__(self, max_workers: int) -> None: self._pool = ThreadPoolExecutor(max_workers=max(1, int(max_workers))) self._lock = threading.Lock() - self._futures = [] + self._futures: list[tuple[Path, Future]] = [] + self._handled: set[int] = set() - def submit(self, path: Path) -> None: + def submit(self, path: Path) -> Future: fut = self._pool.submit(src.recompress_zstd_max, path) with self._lock: self._futures.append((path, fut)) + return fut + + def wait(self, jobs: Sequence[tuple[Path, Future]]) -> None: + failures: list[str] = [] + try: + for path, fut in jobs: + try: + fut.result() + except Exception as exc: + failures.append(f"{path}: {type(exc).__name__}: {exc}") + finally: + with self._lock: + self._handled.update(id(fut) for _path, fut in jobs) + if failures: + raise RuntimeError( + "background code parquet recompress failed:\n" + + "\n".join(failures[:20]) + ) def shutdown(self) -> None: self._pool.shutdown(wait=True) failures: list[str] = [] with self._lock: - futures = list(self._futures) + futures = [ + (path, fut) + for path, fut in self._futures + if id(fut) not in self._handled + ] for path, fut in futures: try: fut.result() @@ -567,25 +591,18 @@ def repo_fully_done( return code_ok and commits_ok +def commit_plan_key(repo: str) -> str: + return f"{repo}::commit_plan" + + def manifest_complete_commit_ranges( repo: str, manifest: Manifest, range_size: int ) -> tuple[tuple[int, int], ...] | None: - """Return complete manifest ranges for ``repo`` when completion is provable. - - This is deliberately conservative. A finished repo may have lost its stable - extract cache after cleanup; on resume we must not spend hours re-running - extract_git_history only to skip already-done ranges. The manifest already - records each processed range as ``::r`` with ``range=[start,end]``. - - We can prove completion without the jsonl only when those done ranges: - * have no failed commit-range entry for this repo, - * start at 0 and are contiguous, and - * end with a short final range (``end - start < range_size``). - - If total records were exactly divisible by ``range_size`` there is no final - short range, so the manifest alone cannot distinguish complete coverage from - "one more full range exists". In that case return None and let the old - extract/count path decide. + """Return done ranges only with an authoritative extracted record count. + + A range shorter than ``range_size`` is not EOF evidence because adaptive + planning also cuts ranges at a byte target. Completion is proven only by a + persisted commit-plan record and exact coverage of ``[0, n_records)``. """ if range_size <= 0: raise ValueError(f"range_size must be positive, got {range_size}") @@ -593,6 +610,16 @@ def manifest_complete_commit_ranges( for k in manifest.failed): return None + plan = manifest.done.get(commit_plan_key(repo)) + if not isinstance(plan, dict) or plan.get("source") != "commit_plan": + return None + try: + n_records = int(plan["n_records"]) + except (KeyError, TypeError, ValueError): + return None + if n_records <= 0: + return None + ranges: list[tuple[int, int]] = [] prefix = f"{repo}::r" for key, info in manifest.done.items(): @@ -628,7 +655,7 @@ def manifest_complete_commit_ranges( if end - start > range_size: return None expected_start = end - if ranges[-1][1] - ranges[-1][0] >= range_size: + if expected_start != n_records: return None return tuple(ranges) @@ -1370,13 +1397,13 @@ def run_code_half( promote_dedup_on_success=False, ) except RepoFailure as exc: - if is_no_trainable_source_failure(exc): - promoted = True + skip_reason = code_skip_reason(exc) + if skip_reason is not None: return { "source": "code", "repo": repo, "skipped": True, - "skip_reason": "no_trainable_source", + "skip_reason": skip_reason, "lengths": {}, "stage_timings_s": {}, "detail": exc.detail, @@ -1385,6 +1412,26 @@ def run_code_half( if info.get("skipped"): return info timings = dict(info.get("stage_timings_s", {})) + # Recompression is part of publication, not deferred cleanup. Multiple + # bucket files can still run concurrently, but this unit cannot promote + # dedup or become manifest-done until every recompress future succeeds. + started = time.monotonic() + jobs: list[tuple[Path, Future]] = [] + try: + for L in info.get("lengths", {}): + dest = sr.OUTPUT_ROOT / str(L) / f"{repo}.parquet" + if dest.exists(): + if recompressor is None: + src.recompress_zstd_max(dest) + else: + jobs.append((dest, recompressor.submit(dest))) + if recompressor is not None: + recompressor.wait(jobs) + except Exception as exc: + remove_code_outputs(repo, info.get("lengths", {}).keys()) + raise RepoFailure(repo, "recompress", f"{type(exc).__name__}: {exc}") from exc + timings["recompress_s"] = round(time.monotonic() - started, 6) + try: timings.update(sr.promote_dedup_stage(dedup_db, stage_id, stage_db)) except Exception as exc: @@ -1395,20 +1442,6 @@ def run_code_half( f"{type(exc).__name__}: {exc}", ) from exc promoted = True - # zstd-max the per-length code parquet files this repo just wrote. - started = time.monotonic() - deferred = 0 - for L in info.get("lengths", {}): - dest = sr.OUTPUT_ROOT / str(L) / f"{repo}.parquet" - if dest.exists(): - if recompressor is None: - src.recompress_zstd_max(dest) - else: - recompressor.submit(dest) - deferred += 1 - timings["recompress_s"] = round(time.monotonic() - started, 6) - if deferred: - timings["recompress_deferred"] = float(deferred) info["stage_timings_s"] = timings return info finally: @@ -1449,6 +1482,17 @@ def is_no_trainable_source_failure(exc: RepoFailure) -> bool: ) +def code_skip_reason(exc: RepoFailure) -> str | None: + if is_no_trainable_source_failure(exc): + return "no_trainable_source" + if ( + exc.stage == "index_project" + and "no training docs (dedup_exhausted)" in exc.detail.lower() + ): + return "dedup_exhausted" + return None + + def is_retryable_index_project_failure(exc: RepoFailure) -> bool: detail = exc.detail.lower() return ( @@ -1520,7 +1564,6 @@ def run_code_half_adaptive( index_stall_timeout_s: int | None = None, *, runner: CodeRunner | None = None, - isolate_dedup_on_retry: bool = True, recompressor: BackgroundRecompressor | None = None, ) -> dict: """Run the code half, retrying index_project peaks/stalls with one parser. @@ -1551,18 +1594,16 @@ def run_code_half_adaptive( raise _log( f"RETRY {code_key(repo)}: {exc.stage} retryable failure with " - f"parse_workers={parse_workers}; retrying parse_workers=1" - + (" with isolated per-repo dedup" if isolate_dedup_on_retry else "") + f"parse_workers={parse_workers}; retrying parse_workers=1 with " + "global exact/chunk dedup and near-dedup disabled" ) - retry_dedup_db = None if isolate_dedup_on_retry else dedup_db - retry_dedup_near = False if isolate_dedup_on_retry else dedup_near return active_runner( repo, repo_dir, lengths_code, work_root, - retry_dedup_db, - retry_dedup_near, + dedup_db, + False, global_symbol_index, memory_limit_gb, 1, @@ -1863,6 +1904,20 @@ def run_commits_half( records_jsonl=str(records_jsonl), ) + # This is the only safe EOF proof. Old manifests without it are deliberately + # re-staged and counted from the extraction cache before ranges are skipped. + records_stat = records_jsonl.stat() + with manifest_lock: + manifest.mark_done( + commit_plan_key(repo), + { + "source": "commit_plan", + "repo": repo, + "n_records": int(n_records), + "records_size_bytes": int(records_stat.st_size), + }, + ) + # .git no longer needed -> free disk now (records already captured/cached). git_dir = repo_dir / ".git" if git_dir.exists(): @@ -2313,12 +2368,12 @@ def process_one_repo( ) ): code_parse_workers = 1 - code_dedup_db = None + code_dedup_db = dedup_db code_dedup_near = False _log( f"RETRY {ck}: prior manifest failure was " "retryable index_project failure; starting parse_workers=1 " - "with isolated per-repo dedup" + "with global exact/chunk dedup and near-dedup disabled" ) if progress is not None: progress.emit( @@ -3240,7 +3295,7 @@ def _empty_totals(lengths): "persistent extract cache were retained. RESUME WITH THE SAME ARGS " "to continue exactly where this left off (zero work lost).") return 130 - return 0 if (not manifest.failed or processed_repos > 0) else 1 + return 0 if not manifest.failed else 1 if __name__ == "__main__": diff --git a/scripts/streaming_reindex.py b/scripts/streaming_reindex.py index b63fc3a1..dd625528 100644 --- a/scripts/streaming_reindex.py +++ b/scripts/streaming_reindex.py @@ -1429,7 +1429,7 @@ def should_process(repo: str) -> bool: "manifest": str(MANIFEST_PATH), } print(json.dumps(summary, indent=2)) - return 0 if not manifest.failed or processed > 0 else 1 + return 0 if not manifest.failed else 1 if __name__ == "__main__": diff --git a/scripts/streaming_reindex_commits.py b/scripts/streaming_reindex_commits.py index 6fbb0fc3..fa284038 100644 --- a/scripts/streaming_reindex_commits.py +++ b/scripts/streaming_reindex_commits.py @@ -944,7 +944,7 @@ def should_process(repo: str) -> bool: "manifest": str(COMMIT_MANIFEST), } print(json.dumps(summary, indent=2)) - return 0 if not manifest.failed or ranges_done > 0 else 1 + return 0 if not manifest.failed else 1 if __name__ == "__main__": diff --git a/tests/test_ast_fim.py b/tests/test_ast_fim.py index 57c7f44e..20dbd9f2 100644 --- a/tests/test_ast_fim.py +++ b/tests/test_ast_fim.py @@ -119,6 +119,26 @@ def encoder(text: str) -> list[int]: assert FIM_PREFIX_ID in result.token_ids assert FIM_SUFFIX_ID in result.token_ids assert result.token_ids[-1] == EOT_ID + assert result.kind == "ast_ifim" + + +def test_ast_ifim_char_fallback_remains_instruction_aware() -> None: + packet = _packet( + list(range(10)), + [(0, 10)], + source_text="// Fill the function body\nint f() {\n", + ) + + result = apply_ast_ifim( + packet, + instruction_encoder=lambda _text: [101, 102], + seed=1, + spm_rate=0.0, + ast_fim_rate=1.0, + ) + + assert result.kind == "char_ifim" + assert result.token_ids[0] == FIM_INSTRUCTION_ID def test_ast_ifim_missing_source_text_raises() -> None: diff --git a/tests/test_audit_sidecar_parquet.py b/tests/test_audit_sidecar_parquet.py index 864a9b1a..0574fa2d 100644 --- a/tests/test_audit_sidecar_parquet.py +++ b/tests/test_audit_sidecar_parquet.py @@ -1,6 +1,9 @@ from __future__ import annotations +import importlib.util import json +import os +from pathlib import Path import subprocess import sys @@ -8,6 +11,16 @@ import pyarrow.parquet as pq +def _load_audit_module(): + path = Path(__file__).parents[1] / "scripts" / "audit_sidecar_parquet.py" + spec = importlib.util.spec_from_file_location("cppmega_test_audit_sidecars", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + def _write_tiny_parquet( path, *, @@ -18,7 +31,7 @@ def _write_tiny_parquet( # single document (doc_ids all equal) and valid=4 the rule yields # [1,1,1,0, 0,0,0,0] and trained_token_count = sum = 3 = valid - num_docs. loss_mask=(1, 1, 1, 0, 0, 0, 0, 0), - doc_ids=(0, 0, 0, 0, 0, 0, 0, 0), + doc_ids=(1, 1, 1, 1, 1, 1, 1, 1), valid_token_count=4, trained_token_count=3, extra=None, @@ -69,6 +82,8 @@ def _write_tiny_parquet( def _run_audit(tmp_path, code_root, commit_root, pr_root, *, extra_args=()): out_dir = tmp_path / "audit" + env = os.environ.copy() + env.pop("PYTHONPATH", None) proc = subprocess.run( [ sys.executable, @@ -88,6 +103,7 @@ def _run_audit(tmp_path, code_root, commit_root, pr_root, *, extra_args=()): *extra_args, ], capture_output=True, + env=env, text=True, ) report = None @@ -120,6 +136,55 @@ def test_sidecar_audit_accepts_valid_chunk_indexed_edges(tmp_path): assert report["total"]["edge_count"]["token_domain_edges"] == 0 +def test_sidecar_audit_reads_large_files_by_row_group(tmp_path, monkeypatch): + path = tmp_path / "code" / "8" / "code.parquet" + _write_tiny_parquet(path) + one_row = pq.read_table(path) + pq.write_table(pa.concat_tables([one_row, one_row]), path, row_group_size=1) + + audit = _load_audit_module() + real_factory = audit.pq.ParquetFile + calls: list[int] = [] + + class RowGroupOnlyParquetFile: + def __init__(self, parquet_path): + self._inner = real_factory(parquet_path) + self.schema_arrow = self._inner.schema_arrow + self.metadata = self._inner.metadata + + def read(self, *args, **kwargs): + raise AssertionError("whole-file parquet reads are forbidden") + + def read_row_group(self, index, *, columns): + calls.append(index) + return self._inner.read_row_group(index, columns=columns) + + monkeypatch.setattr(audit.pq, "ParquetFile", RowGroupOnlyParquetFile) + result = audit._audit_file(str(path), "code", "8", 65536) + + assert calls == [0, 1] + assert result["stats"]["rows"] == 2 + assert result["stats"]["valid_tokens"] == 8 + + +def test_sidecar_audit_rejects_empty_or_reversed_chunk_span(tmp_path): + code_root = tmp_path / "code" + commit_root = tmp_path / "commits" + pr_root = tmp_path / "pr" + _write_tiny_parquet( + code_root / "8" / "code.parquet", + extra={"token_chunk_starts": [0, 2], "token_chunk_ends": [0, 4]}, + ) + _write_tiny_parquet(commit_root / "8" / "commit.parquet") + _write_tiny_parquet(pr_root / "8" / "pr.parquet") + + proc, report = _run_audit(tmp_path, code_root, commit_root, pr_root) + + assert proc.returncode == 2 + assert report["total"]["bad_rows"] == 1 + assert report["total"]["fields"]["token_chunk_ends"]["bad_value_rows"] >= 1 + + def test_sidecar_audit_rejects_allones_loss_mask_on_multidoc_row(tmp_path): """C1 regression: an all-ones loss_mask over a MULTI-document packed row. @@ -145,10 +210,11 @@ def test_sidecar_audit_rejects_allones_loss_mask_on_multidoc_row(tmp_path): commit_root / "8" / "commit.parquet", input_ids=(10, 11, 12, 13, 14, 0, 0, 0), target_ids=(11, 12, 13, 14, 0, 0, 0, 0), - doc_ids=(7, 7, 7, 9, 9, 0, 0, 0), + doc_ids=(1, 1, 1, 2, 2, 2, 2, 2), loss_mask=(1, 1, 1, 1, 0, 0, 0, 0), # WRONG: boundary at pos 2 is masked 1 valid_token_count=5, trained_token_count=4, + extra={"source_doc_ids": [7, 9], "source_doc_token_lengths": [3, 2]}, ) proc, report = _run_audit(tmp_path, code_root, commit_root, pr_root) @@ -180,10 +246,11 @@ def test_sidecar_audit_accepts_correct_multidoc_loss_mask(tmp_path): commit_root / "8" / "commit.parquet", input_ids=(10, 11, 12, 13, 14, 0, 0, 0), target_ids=(11, 12, 13, 14, 0, 0, 0, 0), - doc_ids=(7, 7, 7, 9, 9, 0, 0, 0), + doc_ids=(1, 1, 1, 2, 2, 2, 2, 2), loss_mask=(1, 1, 0, 1, 0, 0, 0, 0), # canonical: 0 at the inter-doc boundary valid_token_count=5, trained_token_count=3, + extra={"source_doc_ids": [7, 9], "source_doc_token_lengths": [3, 2]}, ) proc, report = _run_audit(tmp_path, code_root, commit_root, pr_root) @@ -193,6 +260,31 @@ def test_sidecar_audit_accepts_correct_multidoc_loss_mask(tmp_path): assert report["total"]["fields"]["loss_mask"]["bad_value_rows"] == 0 +def test_sidecar_audit_rejects_collapsed_doc_ids_even_when_mask_matches_them(tmp_path): + code_root = tmp_path / "code" + commit_root = tmp_path / "commits" + pr_root = tmp_path / "pr" + + _write_tiny_parquet(code_root / "8" / "code.parquet") + _write_tiny_parquet(pr_root / "8" / "pr.parquet") + _write_tiny_parquet( + commit_root / "8" / "commit.parquet", + input_ids=(10, 11, 12, 13, 14, 0, 0, 0), + target_ids=(11, 12, 13, 14, 0, 0, 0, 0), + doc_ids=(11, 11, 11, 11, 11, 11, 11, 11), + loss_mask=(1, 1, 1, 1, 0, 0, 0, 0), + valid_token_count=5, + trained_token_count=4, + extra={"source_doc_ids": [38, 41], "source_doc_token_lengths": [3, 2]}, + ) + + proc, report = _run_audit(tmp_path, code_root, commit_root, pr_root) + + assert proc.returncode == 2, proc.stdout + proc.stderr + assert report["total"]["fields"]["doc_ids"]["bad_value_rows"] >= 1 + assert report["total"]["fields"]["loss_mask"]["bad_value_rows"] >= 1 + + def test_sidecar_audit_rejects_bad_target_shift(tmp_path): code_root = tmp_path / "code" commit_root = tmp_path / "commits" diff --git a/tests/test_cpp_jsonl_generation_compile_eval.py b/tests/test_cpp_jsonl_generation_compile_eval.py index 56f39759..ae441d29 100644 --- a/tests/test_cpp_jsonl_generation_compile_eval.py +++ b/tests/test_cpp_jsonl_generation_compile_eval.py @@ -104,6 +104,7 @@ def test_body_decode_constraints_ban_specials_and_degenerate_token_run(): mod = _load_module() class Tok: + vocab_size = 64 bos_token_id = 2 eos_token_id = 3 fim_prefix_id = 4 @@ -116,6 +117,10 @@ class Tok: query_tool_id = 11 tool_result_id = 19 + @staticmethod + def token_for_id(token_id): + return "" if token_id == 48 else None + constraints = mod.BodyDecodeConstraints(Tok(), prompt_len=2, max_token_run=4) logits = mx.zeros((1, 64), dtype=mx.float32) tokens = mx.array([[100, 101, 42, 42, 42, 42]], dtype=mx.int32) @@ -125,6 +130,37 @@ class Tok: assert float(masked[0, 42].item()) == float("-inf") assert float(masked[0, Tok.code_start_id].item()) == float("-inf") assert float(masked[0, Tok.fim_prefix_id].item()) == float("-inf") + assert float(masked[0, 48].item()) == float("-inf") + + +def test_trim_body_completion_preserves_nested_blocks_and_trailing_statements(): + mod = _load_module() + completion = """if (value < lo) { + value = lo; +} +// A brace in a comment does not close the function: } +const char* marker = "}"; +return value; +} +int main() { return 0; } +""" + + assert mod.trim_body_completion(completion) == """if (value < lo) { + value = lo; +} +// A brace in a comment does not close the function: } +const char* marker = "}"; +return value; +""" + + +def test_trim_body_completion_ignores_braces_in_raw_strings(): + mod = _load_module() + completion = 'auto text = R"tag(})tag";\nreturn text.size();\n}\n' + + assert mod.trim_body_completion(completion) == ( + 'auto text = R"tag(})tag";\nreturn text.size();\n' + ) def test_script_help_bootstraps_repo_root_from_sibling_cwd(): diff --git a/tests/test_data_package_imports.py b/tests/test_data_package_imports.py new file mode 100644 index 00000000..0002c083 --- /dev/null +++ b/tests/test_data_package_imports.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] + + +def test_portable_data_submodule_does_not_import_mlx() -> None: + code = """ +import builtins + +real_import = builtins.__import__ + +def reject_mlx(name, globals=None, locals=None, fromlist=(), level=0): + if name == "mlx" or name.startswith("mlx."): + raise AssertionError(f"portable data import reached MLX: {name}") + return real_import(name, globals, locals, fromlist, level) + +builtins.__import__ = reject_mlx +from cppmega_mlx.data.nanochat_pipeline.language_info import detect_language_info + +info = detect_language_info("int main() { return 0; }", filepath="src/example.cpp") +assert info["primary_language"] == "c++" +""" + result = subprocess.run( + [sys.executable, "-c", code], + cwd=ROOT, + text=True, + capture_output=True, + timeout=30, + check=False, + ) + assert result.returncode == 0, result.stderr diff --git a/tests/test_domain_graph_routes.py b/tests/test_domain_graph_routes.py index d4210187..ca92efab 100644 --- a/tests/test_domain_graph_routes.py +++ b/tests/test_domain_graph_routes.py @@ -11,6 +11,12 @@ ) +def test_default_config_prioritizes_diagnostic_locations(): + assert DomainGraphRouteConfig().edge_weights[ + int(DomainEdgeKind.DIAG_PRIMARY_LOCATION) + ] == 2.0 + + def test_domain_attention_bias_routes_edge_triples_to_blocks() -> None: packet = DomainPacket.filled( list(range(8)), diff --git a/tests/test_eval_domain_routed_codegen.py b/tests/test_eval_domain_routed_codegen.py index 86d38961..61c123ff 100644 --- a/tests/test_eval_domain_routed_codegen.py +++ b/tests/test_eval_domain_routed_codegen.py @@ -3,9 +3,16 @@ import importlib.util import json from pathlib import Path +import shutil +import subprocess import sys +def _test_cpp_compiler() -> str | None: + apple_clang = Path("/usr/bin/clang++") + return str(apple_clang) if apple_clang.exists() else shutil.which("clang++") + + def _load_eval_module(): module_path = ( Path(__file__).resolve().parents[1] @@ -65,8 +72,63 @@ def test_domain_eval_compile_gate_accepts_simple_cpp_completion(): expected_sidecars=("token_domain_ids",), compile_prefix="", compile_suffix="\nint main(){return add(2,3)==5 ? 0 : 1;}\n", + run_binary=True, ) - result = mod.compile_cpp_completion(prompt, "int add(int a, int b){ return a + b; }") + result = mod.compile_cpp_completion( + prompt, + "int add(int a, int b){ return a + b; }", + compiler=_test_cpp_compiler(), + ) assert result["status"] in {"compile_passed", "compile_skipped"} + + +def test_domain_eval_runtime_oracle_rejects_wrong_compilable_completion(): + mod = _load_eval_module() + prompt = mod.DomainEvalPrompt( + id="add", + task_type="cpp_docstring_to_code", + prompt="// add\n", + required_domains=("CPP",), + expected_sidecars=("token_domain_ids",), + compile_suffix="\nint main(){return add(2,3)==5 ? 0 : 1;}\n", + run_binary=True, + ) + + result = mod.compile_cpp_completion( + prompt, + "int add(int, int){ return 0; }", + compiler=_test_cpp_compiler(), + ) + + assert result["status"] in {"runtime_failed", "compile_skipped"} + + +def test_domain_eval_reports_compile_timeout_instead_of_raising(monkeypatch): + mod = _load_eval_module() + prompt = mod.DomainEvalPrompt( + id="timeout", + task_type="cpp_docstring_to_code", + prompt="// timeout\n", + required_domains=("CPP",), + expected_sidecars=(), + ) + + def timeout(*_args, **_kwargs): + raise subprocess.TimeoutExpired(["clang++"], timeout=0.01, stderr="busy") + + monkeypatch.setattr(mod.subprocess, "run", timeout) + + result = mod.compile_cpp_completion( + prompt, + "int f() { return 0; }", + compiler="clang++", + compile_timeout_s=0.01, + ) + + assert result == { + "status": "compile_timeout", + "timeout_s": 0.01, + "stderr": "busy", + } diff --git a/tests/test_inference_generation.py b/tests/test_inference_generation.py index 142136bb..c30a2ef7 100644 --- a/tests/test_inference_generation.py +++ b/tests/test_inference_generation.py @@ -487,6 +487,36 @@ def __call__(self, tokens: mx.array) -> tuple[mx.array, mx.array]: ) +def test_generate_tokens_accepts_dense_cpp_lm_inference_tuple() -> None: + class DenseInferenceModel: + def __call__(self, tokens: mx.array) -> tuple[mx.array, None]: + logits = mx.concatenate( + [ + mx.full( + (tokens.shape[0], tokens.shape[1], 3), + -100.0, + dtype=mx.float32, + ), + mx.full( + (tokens.shape[0], tokens.shape[1], 1), + 100.0, + dtype=mx.float32, + ), + ], + axis=-1, + ) + return logits, None + + generated = generate_tokens( + DenseInferenceModel(), + mx.array([[1, 2]], dtype=mx.int32), + max_new_tokens=1, + temperature=0.0, + ) + + np.testing.assert_array_equal(_as_numpy(generated), np.array([[1, 2, 3]])) + + def test_generate_tokens_rejects_structured_mtp_output() -> None: class StructuredMTPModel: def __call__(self, tokens: mx.array) -> dict[str, mx.array]: diff --git a/tests/test_mamba3_dispatch.py b/tests/test_mamba3_dispatch.py index 738110c1..59a9120e 100644 --- a/tests/test_mamba3_dispatch.py +++ b/tests/test_mamba3_dispatch.py @@ -9,6 +9,8 @@ from __future__ import annotations +import sys + import numpy as np import pytest @@ -85,6 +87,20 @@ def test_reference_policy_forces_pure_mlx(monkeypatch: pytest.MonkeyPatch) -> No assert matches[-1]["kernel_used"] == "reference_pure_mlx" +def test_reference_policy_does_not_import_tilelang_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("CPPMEGA_KERNEL_PATH", "ref") + monkeypatch.setitem(sys.modules, "cppmega_mlx.nn._tilelang.mamba3", None) + + block, hidden = _make_block() + out, _ = block(hidden) + mx.eval(out) + + matches = [e for e in get_dispatch_log() if e["op_name"] == "mamba3_mimo"] + assert matches[-1]["kernel_used"] == "reference_pure_mlx" + + def test_path_b_policy_forces_metal(monkeypatch: pytest.MonkeyPatch) -> None: if not _METAL_AVAILABLE: pytest.skip("Metal not available") diff --git a/tests/test_megatron_indexed.py b/tests/test_megatron_indexed.py index 5de3372f..ef49712e 100644 --- a/tests/test_megatron_indexed.py +++ b/tests/test_megatron_indexed.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import os import subprocess import sys import struct @@ -255,9 +256,15 @@ def _write_structured_multishard_fixture( def _run_train_hybrid_tiny(*args: str) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + # This suite validates indexed-dataset ingress, not optional kernel discovery. + # An explicit reference route keeps the subprocess independent of any shared + # TileLang/TVM development checkout configured by the parent pytest process. + env["CPPMEGA_KERNEL_PATH"] = "ref" return subprocess.run( [sys.executable, str(TRAIN_HYBRID_TINY), *args], cwd=ROOT, + env=env, text=True, capture_output=True, timeout=45, @@ -1249,6 +1256,185 @@ def test_mmididx_graph_sidecars_are_sequence_aligned_routes(tmp_path) -> None: assert packet1.graph.edge("type").num_edges == 0 +def test_compact_fixed_rows_restore_padding_and_document_graphs(tmp_path) -> None: + prefix = tmp_path / "compact_fixed_rows" + docs = [ + np.array([1, 2, 3], dtype=np.int32), + np.array([10, 11], dtype=np.int32), + ] + _write_mmididx(prefix, docs, dtype=np.int32) + sidecar = _write_graph_sidecars( + prefix, + call_edges=[ + np.array([[0, 1]], dtype=np.int32), + np.zeros((0, 2), dtype=np.int32), + ], + type_edges=[ + np.zeros((0, 2), dtype=np.int32), + np.array([[0, 0]], dtype=np.int32), + ], + chunk_starts=[np.array([0, 1]), np.array([0])], + chunk_ends=[np.array([1, 3]), np.array([2])], + chunk_kinds=[np.array([4, 5]), np.array([6])], + chunk_dep_levels=[np.array([0, 1]), np.array([0])], + ) + loss_mask = np.array([1, 1, 0, 1, 0], dtype=np.uint8) + document_ids = np.array([1, 1, 1, 2, 2], dtype=np.uint16) + hunk_ids = np.array([0, 0, -1, 2, -1], dtype=np.int32) + loss_mask.tofile(tmp_path / "loss_mask.bin") + document_ids.tofile(tmp_path / "doc_ids.bin") + hunk_ids.tofile(tmp_path / "hunk_ids.bin") + sidecar.update( + { + "token_count": 5, + "source_capacity_token_count": 8, + "document_count": 2, + "side_channel_paths": { + "loss_mask": {"path": "loss_mask.bin", "dtype": "uint8"}, + "doc_ids": {"path": "doc_ids.bin", "dtype": "uint16"}, + "hunk_id_per_token": {"path": "hunk_ids.bin", "dtype": "int32"}, + }, + } + ) + prefix.with_suffix(".idx.json").write_text(json.dumps(sidecar), encoding="utf-8") + + dataset = MegatronIndexedDataset(prefix, seq_len=4, batch_size=2) + batch = next(dataset.iter_batches()) + + assert dataset.num_samples == 2 + np.testing.assert_array_equal( + np.array(batch.tokens), + np.array([[1, 2, 3, 0], [10, 11, 0, 0]], dtype=np.int32), + ) + np.testing.assert_array_equal( + np.array(batch.loss_mask), + np.array([[1, 1, 0, 0], [1, 0, 0, 0]], dtype=np.float32), + ) + np.testing.assert_array_equal( + np.array(batch.document_ids), + np.array([[1, 1, 1, 0], [2, 2, 0, 0]], dtype=np.int32), + ) + np.testing.assert_array_equal( + np.array(batch.side_channel_map()["temporal_diff"]["hunk_ids"]), + np.array([[0, 0, -1, -1], [2, -1, -1, -1]], dtype=np.int32), + ) + assert dataset.graph_route_packet_for_sample(0).graph.edge("call").to_pairs() == [ + (0, 1) + ] + assert dataset.graph_route_packet_for_sample(1).graph.edge("type").to_pairs() == [ + (0, 0) + ] + + +def test_compact_fixed_rows_reject_inconsistent_capacity_metadata(tmp_path) -> None: + prefix = tmp_path / "bad_compact_fixed_rows" + _write_mmididx(prefix, [np.arange(3, dtype=np.int32)], dtype=np.int32) + prefix.with_suffix(".idx.json").write_text( + json.dumps( + { + "token_count": 3, + "source_capacity_token_count": 5, + "document_count": 1, + } + ), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match=r"sequence_count \* seq_len"): + MegatronIndexedDataset(prefix, seq_len=4, batch_size=1) + + +def test_mmididx_v2_domain_routes_and_token_sidecars_round_trip(tmp_path) -> None: + prefix = tmp_path / "domain_routes_v2" + docs = [np.arange(4, dtype=np.int32), np.arange(10, 14, dtype=np.int32)] + _write_mmididx(prefix, docs, dtype=np.int32) + sidecar = _write_graph_sidecars( + prefix, + call_edges=[np.array([[0, 1]], dtype=np.int32), np.zeros((0, 2), dtype=np.int32)], + type_edges=[np.zeros((0, 2), dtype=np.int32), np.zeros((0, 2), dtype=np.int32)], + chunk_starts=[np.array([0, 2]), np.array([0, 2])], + chunk_ends=[np.array([2, 4]), np.array([2, 4])], + chunk_kinds=[np.array([1, 2]), np.array([1, 2])], + chunk_dep_levels=[np.array([0, 1]), np.array([0, 1])], + ) + sidecar["graph_sidecar_schema"] = "cppmega_graph_routes_v2" + graph_paths = sidecar["graph_sidecar_paths"] + for key, entry in graph_paths.items(): + entry["coordinate_space"] = ( + "chunk_index" + if key in {"token_call_edges", "token_type_edges", "token_chunk_kinds", "token_chunk_dep_levels"} + else "token_index" + ) + + offsets_path = tmp_path / "domain_offsets.bin" + data_path = tmp_path / "domain_data.bin" + np.array([0, 1, 2], dtype=np.int64).tofile(offsets_path) + np.array([[0, 3, 20], [1, 2, 21]], dtype=np.int32).tofile(data_path) + graph_paths["token_domain_edges"] = { + "kind": "edge_triples", + "offsets_path": offsets_path.name, + "data_path": data_path.name, + "offset_dtype": "int64", + "dtype": "int32", + "item_count": 2, + "shape_tail": [3], + "coordinate_space": "token_index", + } + for name in ( + "token_build_edges", + "token_shell_edges", + "token_diagnostic_edges", + "token_cross_domain_edges", + ): + empty_offsets = tmp_path / f"{name}_offsets.bin" + empty_data = tmp_path / f"{name}_data.bin" + np.zeros(3, dtype=np.int64).tofile(empty_offsets) + np.zeros((0, 3), dtype=np.int32).tofile(empty_data) + graph_paths[name] = { + "kind": "edge_triples", + "offsets_path": empty_offsets.name, + "data_path": empty_data.name, + "offset_dtype": "int64", + "dtype": "int32", + "item_count": 0, + "shape_tail": [3], + "coordinate_space": "token_index", + } + + side_channel_paths = {} + token_sidecars = { + "loss_mask": np.array([1, 1, 1, 0, 1, 1, 1, 0], dtype=np.uint8), + "doc_ids": np.array([1] * 8, dtype=np.int32), + "token_domain_ids": np.array([20, 20, 20, 20, 21, 21, 21, 21], dtype=np.uint16), + "token_role_ids": np.arange(8, dtype=np.uint16), + "token_entity_ids": np.arange(8, dtype=np.uint32), + "token_scope_ids": np.arange(8, dtype=np.uint32), + "token_source_doc_ids": np.array([1, 1, 1, 1, 2, 2, 2, 2], dtype=np.uint32), + "token_confidence_ids": np.ones(8, dtype=np.uint8), + "hunk_id_per_token": np.array([-1, 0, 0, -1, -1, 1, 1, -1], dtype=np.int32), + "edit_op_per_token": np.arange(8, dtype=np.uint8), + } + for name, values in token_sidecars.items(): + path = tmp_path / f"{name}.bin" + values.tofile(path) + side_channel_paths[name] = {"path": path.name, "dtype": values.dtype.name} + sidecar["side_channel_paths"] = side_channel_paths + prefix.with_suffix(".idx.json").write_text(json.dumps(sidecar), encoding="utf-8") + + dataset = MegatronIndexedDataset(prefix, seq_len=4, batch_size=1) + batch = next(dataset.iter_batches()) + packet = dataset.graph_route_packet_for_sample(0) + + np.testing.assert_array_equal(np.array(batch.loss_mask), [np.array([1, 1, 1, 0])]) + domain = batch.side_channel_map()["domain_routes"] + np.testing.assert_array_equal(np.array(domain["domain_ids"]), [[20, 20, 20, 20]]) + temporal = batch.side_channel_map()["temporal_diff"] + np.testing.assert_array_equal(np.array(temporal["hunk_ids"]), [[-1, 0, 0, -1]]) + assert packet.graph.edge("call").to_pairs() == [(0, 1)] + assert packet.graph.edge("domain").to_pairs() == [(0, 3)] + np.testing.assert_array_equal(packet.edge_kinds["domain"], np.array([20], dtype=np.int32)) + + def test_mmididx_graph_sidecars_fail_closed_for_window_slicing(tmp_path) -> None: prefix = tmp_path / "graph_routes_long_sequence" _write_mmididx(prefix, [np.arange(8, dtype=np.int32)], dtype=np.int32) diff --git a/tests/test_pack_enriched_rows.py b/tests/test_pack_enriched_rows.py index 548d8793..6f95faa9 100644 --- a/tests/test_pack_enriched_rows.py +++ b/tests/test_pack_enriched_rows.py @@ -169,6 +169,17 @@ def test_sequential_strategy_preserves_document_order() -> None: assert rows[2][INPUT_IDS_COLUMN] == [20, 21, 22, 30, 0, 0] +def test_row_platform_provenance_union_is_not_limited_by_model_bag_width() -> None: + docs = _normalize( + [_doc([100 + index], platform_ids=[index + 1]) for index in range(22)] + ) + + rows, overflow = pack_documents(docs, target_length=32, pad_token_id=0) + + assert overflow == [] + assert rows[0][PLATFORM_IDS_COLUMN] == list(range(1, 23)) + + def test_loss_mask_excludes_padding_and_cross_document_targets() -> None: docs = _normalize([_doc([1, 2]), _doc([10, 11])]) @@ -182,6 +193,35 @@ def test_loss_mask_excludes_padding_and_cross_document_targets() -> None: assert row[LOSS_MASK_COLUMN] == [1, 0, 1, 0, 0] +def test_packed_doc_ids_do_not_collapse_same_file_documents() -> None: + first = normalize_document_record( + { + **_doc([1, 2]), + REPO_STABLE_ID_COLUMN: "same-repo", + FILEPATH_STABLE_ID_COLUMN: "same-file", + }, + source_doc_index=38, + stable_doc_id=11, + ) + second = normalize_document_record( + { + **_doc([10, 11]), + REPO_STABLE_ID_COLUMN: "same-repo", + FILEPATH_STABLE_ID_COLUMN: "same-file", + }, + source_doc_index=41, + stable_doc_id=11, + ) + + rows, overflow = pack_documents([first, second], target_length=5, pad_token_id=0) + + assert overflow == [] + row = rows[0] + assert row[SOURCE_DOC_INDICES_COLUMN] == [38, 41] + assert row[DOC_IDS_COLUMN] == [1, 1, 2, 2, 2] + assert row[LOSS_MASK_COLUMN] == [1, 0, 1, 0, 0] + + def test_pack_documents_preserves_pr_discussion_source_metadata() -> None: first = _doc([1, 2]) first.update( diff --git a/tests/test_process_commits_fail_loud.py b/tests/test_process_commits_fail_loud.py new file mode 100644 index 00000000..414afc60 --- /dev/null +++ b/tests/test_process_commits_fail_loud.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import json +import sqlite3 +import sys +from pathlib import Path +from types import SimpleNamespace + + +CLANG_INDEXER = Path(__file__).resolve().parents[1] / "tools" / "clang_indexer" +if str(CLANG_INDEXER) not in sys.path: + sys.path.insert(0, str(CLANG_INDEXER)) + +import process_commits # noqa: E402 +import pytest # noqa: E402 + + +def _stats(*, parse_errors: int) -> dict[str, int]: + return { + "records_read": 2, + "documents_written": 1, + "records_skipped": 0, + "records_empty": 0, + "parse_errors": parse_errors, + "commit_chunks_claimed": 0, + "commit_chunks_skipped": 0, + "analysis_cache_hits": 0, + "analysis_cache_misses": 0, + "analysis_cache_evictions": 0, + } + + +def _run_main(tmp_path: Path, monkeypatch, *, allow: bool) -> int: + source = tmp_path / "records.jsonl" + source.write_text("{}\n", encoding="utf-8") + output = tmp_path / "out.jsonl" + argv = [ + "process_commits.py", + "--inputs", + str(source), + "--output", + str(output), + ] + if allow: + argv.append("--allow-parse-errors") + monkeypatch.setattr(sys, "argv", argv) + monkeypatch.setattr(process_commits, "start_memory_guard", lambda *_a, **_k: None) + monkeypatch.setattr(process_commits, "_configure_libclang", lambda _path: None) + monkeypatch.setattr(process_commits, "_load_tokenizer", lambda _path: None) + monkeypatch.setattr(process_commits, "Index", SimpleNamespace(create=lambda: object())) + monkeypatch.setattr( + process_commits, + "process_jsonl_file", + lambda *_args, **_kwargs: _stats(parse_errors=1), + ) + return process_commits.main() + + +def test_process_commits_rejects_partial_range_by_default(tmp_path, monkeypatch) -> None: + assert _run_main(tmp_path, monkeypatch, allow=False) == 1 + + +def test_process_commits_allows_explicit_lossy_mode(tmp_path, monkeypatch) -> None: + assert _run_main(tmp_path, monkeypatch, allow=True) == 0 + + +def test_process_commits_rejects_missing_input_before_clang_init(tmp_path, monkeypatch) -> None: + output = tmp_path / "out.jsonl" + monkeypatch.setattr( + sys, + "argv", + [ + "process_commits.py", + "--inputs", + str(tmp_path / "missing.jsonl"), + "--output", + str(output), + ], + ) + monkeypatch.setattr( + process_commits, + "_configure_libclang", + lambda _path: (_ for _ in ()).throw(AssertionError("must not initialize clang")), + ) + + assert process_commits.main() == 1 + assert not output.exists() + + +def test_sha_pr_lookup_restores_number_title_and_uses_readonly_store( + tmp_path: Path, +) -> None: + store = tmp_path / "prs.sqlite" + conn = process_commits._pr_store_mod.connect(str(store), create=True) + process_commits._pr_store_mod.upsert_record( + conn, + { + "repo": "owner/repo", + "pr_number": 1467, + "merge_commit_sha": "abc123", + "pr_title": "Fix the build", + "pr_body": "Body", + "comments": [], + "reviews": [], + "linked_issues": [], + }, + ) + conn.close() + repo_list = tmp_path / "repos.json" + repo_list.write_text(json.dumps({"repos": []}), encoding="utf-8") + + lookup = process_commits.PRDiscussionLookup(str(store), str(repo_list)) + record = {"repo": "owner/repo", "commit_hash": "abc123"} + try: + assert lookup.attach(record) is True + assert record["pr_number"] == 1467 + assert record["pr_title"] == "Fix the build" + assert "Body" in record["pr_discussion"] + with pytest.raises(sqlite3.OperationalError, match="readonly"): + lookup._conn.execute("CREATE TABLE forbidden(value INTEGER)") + finally: + lookup.close() diff --git a/tests/test_repair_packed_document_boundaries.py b/tests/test_repair_packed_document_boundaries.py new file mode 100644 index 00000000..4b18e1bf --- /dev/null +++ b/tests/test_repair_packed_document_boundaries.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from scripts.repair_packed_document_boundaries import _rewrite_file, repair_row + + +def _collapsed_row() -> dict[str, object]: + return { + "input_ids": [10, 11, 12, 13, 14, 0, 0, 0], + "target_ids": [11, 12, 13, 14, 0, 0, 0, 0], + "doc_ids": [11] * 8, + "loss_mask": [1, 1, 1, 1, 0, 0, 0, 0], + "source_doc_ids": [38, 41], + "source_doc_token_lengths": [3, 2], + "valid_token_count": 5, + "trained_token_count": 4, + "num_docs": 2, + } + + +def test_repair_row_restores_independent_same_file_boundary() -> None: + row = _collapsed_row() + + repaired, changed, restored = repair_row(row) + + assert changed is True + assert restored == 1 + assert repaired["doc_ids"] == [1, 1, 1, 2, 2, 2, 2, 2] + assert repaired["loss_mask"] == [1, 1, 0, 1, 0, 0, 0, 0] + assert repaired["trained_token_count"] == 3 + assert repaired["target_ids"] == row["target_ids"] + + +def test_repair_row_rejects_lengths_that_do_not_partition_valid_tokens() -> None: + row = _collapsed_row() + row["source_doc_token_lengths"] = [3, 1] + + with pytest.raises(ValueError, match="source_doc_token_lengths sum=4"): + repair_row(row, where="fixture") + + +def test_rewrite_file_replaces_hardlink_without_mutating_source(tmp_path: Path) -> None: + source = tmp_path / "source.parquet" + snapshot = tmp_path / "snapshot.parquet" + pq.write_table(pa.Table.from_pylist([_collapsed_row()]), source, compression="zstd") + snapshot.hardlink_to(source) + source_inode = source.stat().st_ino + + _rewrite_file(str(snapshot), 3) + + assert source.stat().st_ino == source_inode + assert snapshot.stat().st_ino != source_inode + assert pq.read_table(source).column("doc_ids").to_pylist()[0] == [11] * 8 + repaired = pq.read_table(snapshot).to_pylist()[0] + assert repaired["doc_ids"] == [1, 1, 1, 2, 2, 2, 2, 2] + assert repaired["loss_mask"] == [1, 1, 0, 1, 0, 0, 0, 0] diff --git a/tests/test_streaming_conveyor_progress.py b/tests/test_streaming_conveyor_progress.py index 67c6abb7..0f3a31c5 100644 --- a/tests/test_streaming_conveyor_progress.py +++ b/tests/test_streaming_conveyor_progress.py @@ -5,6 +5,8 @@ import sys from pathlib import Path +import pytest + MLX_ROOT = Path(__file__).resolve().parents[1] SCRIPTS = MLX_ROOT / "scripts" @@ -144,12 +146,13 @@ def fake_recompress(path): assert sorted(calls) == ["a.parquet", "b.parquet"] -def test_manifest_complete_commit_ranges_detects_short_final_range(tmp_path): +def test_manifest_complete_commit_ranges_uses_authoritative_record_count(tmp_path): import streaming_conveyor manifest = streaming_conveyor.Manifest( path=tmp_path / "_done.json", done={ + "repo::commit_plan": {"source": "commit_plan", "n_records": 823}, "repo::r0": {"range": [0, 500], "source": "commits"}, "repo::r500": {"range": [500, 823], "source": "commits"}, }, @@ -178,6 +181,23 @@ def test_manifest_complete_commit_ranges_refuses_exact_multiple(tmp_path): ) is None +def test_manifest_complete_commit_ranges_refuses_adaptive_short_nonfinal_range(tmp_path): + import streaming_conveyor + + manifest = streaming_conveyor.Manifest( + path=tmp_path / "_done.json", + done={ + "repo::commit_plan": {"source": "commit_plan", "n_records": 1200}, + "repo::r0": {"range": [0, 468], "source": "commits"}, + }, + failed={}, + ) + + assert streaming_conveyor.manifest_complete_commit_ranges( + "repo", manifest, 500 + ) is None + + def test_missing_commit_subranges_respects_split_done_intervals(tmp_path): import streaming_conveyor @@ -406,7 +426,7 @@ def runner( assert calls == [ (2, tmp_path / "global.sqlite", True, 7200, 900, None), - (1, None, False, 7200, 900, None), + (1, tmp_path / "global.sqlite", False, 7200, 900, None), ] assert info["lengths"]["1024"]["valid_tokens"] == 1024 @@ -460,7 +480,7 @@ def runner( assert calls == [ (2, tmp_path / "global.sqlite", True, 7200, 900), - (1, None, False, 7200, 900), + (1, tmp_path / "global.sqlite", False, 7200, 900), ] assert info["lengths"]["1024"]["valid_tokens"] == 1024 @@ -486,6 +506,7 @@ def fail_promote(*_args, **_kwargs): raise sqlite3.OperationalError("unable to open database file") monkeypatch.setattr(streaming_conveyor.sr, "process_one_repo", fake_process_one_repo) + monkeypatch.setattr(streaming_conveyor.src, "recompress_zstd_max", lambda _path: None) monkeypatch.setattr(streaming_conveyor.sr, "promote_dedup_stage", fail_promote) try: @@ -506,6 +527,99 @@ def fail_promote(*_args, **_kwargs): assert not dest.exists() +def test_run_code_half_recompresses_before_dedup_promote(tmp_path, monkeypatch): + import streaming_conveyor + + repo = "repo" + output_root = tmp_path / "out" + dest = output_root / "1024" / f"{repo}.parquet" + dest.parent.mkdir(parents=True) + dest.write_bytes(b"parquet") + monkeypatch.setattr(streaming_conveyor.sr, "OUTPUT_ROOT", output_root) + events: list[str] = [] + + monkeypatch.setattr( + streaming_conveyor.sr, + "process_one_repo", + lambda *_args, **_kwargs: { + "source": "code", + "lengths": {"1024": _stat(rows=1, valid=1024, target_length=1024)}, + "stage_timings_s": {}, + }, + ) + monkeypatch.setattr( + streaming_conveyor.src, + "recompress_zstd_max", + lambda _path: events.append("recompress"), + ) + + def promote(*_args, **_kwargs): + events.append("promote") + return {} + + monkeypatch.setattr(streaming_conveyor.sr, "promote_dedup_stage", promote) + recompressor = streaming_conveyor.BackgroundRecompressor(max_workers=1) + try: + streaming_conveyor.run_code_half( + repo, + tmp_path / "src", + (1024,), + tmp_path / "work", + tmp_path / "dedup.sqlite", + True, + recompressor=recompressor, + ) + finally: + recompressor.shutdown() + + assert events == ["recompress", "promote"] + + +def test_run_code_half_recompress_failure_rolls_back_before_promote(tmp_path, monkeypatch): + import streaming_conveyor + + repo = "repo" + output_root = tmp_path / "out" + dest = output_root / "1024" / f"{repo}.parquet" + dest.parent.mkdir(parents=True) + dest.write_bytes(b"parquet") + monkeypatch.setattr(streaming_conveyor.sr, "OUTPUT_ROOT", output_root) + monkeypatch.setattr( + streaming_conveyor.sr, + "process_one_repo", + lambda *_args, **_kwargs: { + "source": "code", + "lengths": {"1024": _stat(rows=1, valid=1024, target_length=1024)}, + "stage_timings_s": {}, + }, + ) + monkeypatch.setattr( + streaming_conveyor.src, + "recompress_zstd_max", + lambda _path: (_ for _ in ()).throw(RuntimeError("zstd failed")), + ) + monkeypatch.setattr( + streaming_conveyor.sr, + "promote_dedup_stage", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("dedup must not promote before recompress") + ), + ) + + with pytest.raises(streaming_conveyor.RepoFailure, match="zstd failed") as exc_info: + streaming_conveyor.run_code_half( + repo, + tmp_path / "src", + (1024,), + tmp_path / "work", + tmp_path / "dedup.sqlite", + True, + ) + + assert exc_info.value.stage == "recompress" + assert not dest.exists() + + def test_run_code_half_skips_no_trainable_source(tmp_path, monkeypatch): import streaming_conveyor @@ -532,6 +646,40 @@ def no_trainable_source(*_args, **_kwargs): assert info["lengths"] == {} +def test_run_code_half_skips_dedup_exhausted_and_discards_stage( + tmp_path, monkeypatch +): + import streaming_conveyor + + def dedup_exhausted(*_args, **_kwargs): + raise streaming_conveyor.RepoFailure( + "repo", + "index_project", + "no training docs (dedup_exhausted): all functions already claimed", + ) + + discarded = [] + monkeypatch.setattr(streaming_conveyor.sr, "process_one_repo", dedup_exhausted) + monkeypatch.setattr( + streaming_conveyor.sr, + "discard_dedup_stage", + lambda *args: discarded.append(args), + ) + + info = streaming_conveyor.run_code_half( + "repo", + tmp_path / "src", + (1024,), + tmp_path / "work", + tmp_path / "dedup.sqlite", + True, + ) + + assert info["skipped"] is True + assert info["skip_reason"] == "dedup_exhausted" + assert len(discarded) == 1 + + def test_failed_code_unit_was_index_memory_reads_manifest_receipt(tmp_path): import streaming_conveyor @@ -588,6 +736,7 @@ def test_run_commits_half_skips_extract_when_manifest_proves_complete(tmp_path): manifest = streaming_conveyor.Manifest( path=tmp_path / "_done.json", done={ + "repo::commit_plan": {"source": "commit_plan", "n_records": 612}, "repo::r0": {"range": [0, 500], "source": "commits"}, "repo::r500": {"range": [500, 612], "source": "commits"}, }, @@ -908,7 +1057,9 @@ def runner( ) assert (done, failed, all_done) == (3, 0, True) - assert sorted(manifest.done) == ["repo::r0", "repo::r1", "repo::r2"] + range_keys = sorted(key for key in manifest.done if key.startswith("repo::r")) + assert range_keys == ["repo::r0", "repo::r1", "repo::r2"] + assert manifest.done["repo::commit_plan"]["n_records"] == 3 assert len(observed_kwargs) == 3 reader = DedupStore(str(db), near=False, commit_every=1) try: @@ -1147,6 +1298,7 @@ def test_should_stage_repo_skips_manifest_proven_complete_both_streams(tmp_path) path=tmp_path / "_done.json", done={ "repo::code": {"source": "code"}, + "repo::commit_plan": {"source": "commit_plan", "n_records": 612}, "repo::r0": {"range": [0, 500], "source": "commits"}, "repo::r500": {"range": [500, 612], "source": "commits"}, }, diff --git a/tools/clang_indexer/process_commits.py b/tools/clang_indexer/process_commits.py index 72f4b298..3d711ce7 100644 --- a/tools/clang_indexer/process_commits.py +++ b/tools/clang_indexer/process_commits.py @@ -101,7 +101,11 @@ def __init__(self, store_path: str, repo_list_path: str | None) -> None: if not os.path.exists(store_path): raise FileNotFoundError(f"--pr-store does not exist: {store_path}") # Read-only connection; create=False RAISES on a missing store (fail-loud). - self._conn = _pr_store_mod.connect(store_path, create=False) + self._conn = _pr_store_mod.connect( + store_path, + create=False, + readonly=True, + ) self._name_to_owner_repo: dict[str, str] = {} if repo_list_path: if not os.path.exists(repo_list_path): @@ -158,6 +162,13 @@ def attach(self, record: dict) -> bool: if not discussion: self.misses += 1 return False + canonical_pr_number = rec.get("pr_number") + if canonical_pr_number is None: + raise ValueError( + f"PR store hit for {owner_repo} has no canonical pr_number" + ) + record["pr_number"] = int(canonical_pr_number) + record["pr_title"] = str(rec.get("pr_title") or "") record["pr_discussion"] = discussion self.hits += 1 return True @@ -2396,6 +2407,14 @@ def main() -> int: help='Bounded per-range LRU for expensive libclang file analyses. ' 'Default 128; use 0 to disable.', ) + parser.add_argument( + '--allow-parse-errors', + action='store_true', + help=( + 'Explicitly allow malformed/failed records to be skipped. Default ' + 'fails the whole range so staged dedup and parquet publication roll back.' + ), + ) parser.add_argument( '--dedup-db', default=None, help='Path to the SHARED global dedup SQLite store. When set, whole commit ' @@ -2438,6 +2457,13 @@ def main() -> int: "'/' are used as-is.", ) args = parser.parse_args() + missing_inputs = [path for path in args.inputs if not os.path.exists(path)] + if missing_inputs: + print( + "ERROR: missing input file(s): " + ", ".join(missing_inputs), + file=sys.stderr, + ) + return 1 start_memory_guard(args.memory_limit_gb, label="process_commits") print("Clang commit processor starting", file=sys.stderr) @@ -2546,10 +2572,6 @@ def main() -> int: try: with open(args.output, 'w') as out_f: for input_path in args.inputs: - if not os.path.exists(input_path): - print(f" WARN: {input_path} not found, skipping", file=sys.stderr) - continue - print(f"\n Processing {input_path}...", file=sys.stderr) stats = process_jsonl_file( input_path, out_f, clang_index, actual_tmpdir, @@ -2586,6 +2608,14 @@ def main() -> int: if total_stats['records_read'] > 0: rate = total_stats['records_read'] / elapsed print(f"Rate: {rate:.1f} records/sec", file=sys.stderr) + if total_stats['parse_errors'] and not args.allow_parse_errors: + print( + "ERROR: refusing partial commit range: " + f"{total_stats['parse_errors']} record parse error(s); " + "use --allow-parse-errors only for an explicitly lossy run", + file=sys.stderr, + ) + return 1 return 0