Skip to content

Commit 9efeef5

Browse files
apstenku123apstenku123
committed
fix: keep corpus imports portable without MLX
1 parent 7be95ab commit 9efeef5

3 files changed

Lines changed: 126 additions & 135 deletions

File tree

‎.github/workflows/ci-self-hosted.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ jobs:
2828
tests/test_ast_fim.py \
2929
tests/test_audit_sidecar_parquet.py \
3030
tests/test_cpp_jsonl_generation_compile_eval.py \
31+
tests/test_data_package_imports.py \
3132
tests/test_domain_graph_routes.py \
3233
tests/test_eval_domain_routed_codegen.py \
3334
tests/test_inference_generation.py \
@@ -58,6 +59,7 @@ jobs:
5859
# the macOS-only root conftest/package initialization.
5960
python -m pytest -q --noconftest \
6061
tests/test_audit_sidecar_parquet.py \
62+
tests/test_data_package_imports.py \
6163
tests/test_streaming_conveyor_progress.py \
6264
tests/test_process_commits_fail_loud.py \
6365
tests/test_repair_packed_document_boundaries.py

‎cppmega_mlx/data/__init__.py‎

Lines changed: 88 additions & 135 deletions
Original file line numberDiff line numberDiff line change
@@ -1,137 +1,90 @@
1-
"""Data readers and batch collation helpers."""
1+
"""Data readers and batch collation helpers.
22
3-
from cppmega_mlx.data.batch import LMTokenBatch, ensure_lm_batch, synthetic_token_batch
4-
from cppmega_mlx.data.dataloader_bridge import (
5-
LocalTokenBatchDataset,
6-
TorchDataLoaderBridgeConfig,
7-
TorchDataLoaderBridgeError,
8-
build_spawn_dataloader,
9-
is_torch_dataloader_available,
10-
iter_mlx_batches,
11-
)
12-
from cppmega_mlx.data.fim import (
13-
EOT_ID,
14-
FIMSpecialTokenIds,
15-
FIM_INSTRUCTION_ID,
16-
FIMMode,
17-
FIM_MIDDLE_ID,
18-
FIM_PREFIX_ID,
19-
FIM_SPECIAL_TOKEN_IDS,
20-
FIM_SUFFIX_ID,
21-
apply_fim_permutation,
22-
apply_fim_transform,
23-
apply_ifim_permutation,
24-
apply_ifim_transform,
25-
extract_ifim_instruction_text,
26-
sample_middle_span,
27-
)
28-
from cppmega_mlx.data.megatron_indexed import (
29-
MegatronIndexedDataset,
30-
MegatronIndexedMetadata,
31-
MegatronIndexedMultiShardDataset,
32-
MegatronIndexedMultiShardMetadata,
33-
megatron_indexed_side_channel_schema,
34-
open_megatron_indexed_dataset,
35-
)
36-
from cppmega_mlx.data.packing import (
37-
OversizedSamplePolicy,
38-
PackedSequences,
39-
PackingStrategy,
40-
cumulative_doc_ids_from_eos,
41-
document_boundary_mask,
42-
mlx_cumulative_doc_ids_from_eos,
43-
mlx_document_boundary_mask,
44-
mlx_sequence_packing_attention_mask,
45-
pack_bos_aligned_best_fit,
46-
pack_documents_with_eos,
47-
)
48-
from cppmega_mlx.data.parquet_dataset import (
49-
MultiShardTokenParquetDataset,
50-
ParquetColumns,
51-
TokenParquetDataset,
52-
)
53-
from cppmega_mlx.data.platform_context import (
54-
MAX_PLATFORM_IDS,
55-
PLATFORM_VOCAB,
56-
PLATFORM_VOCAB_SIZE,
57-
PlatformContext,
58-
encode_platform_context,
59-
parse_platform_context,
60-
platform_ids_array,
61-
render_platform_context,
62-
)
63-
from cppmega_mlx.data.tokenizer_contract import (
64-
REQUIRED_SPECIAL_TOKEN_IDS,
65-
SpecialTokenMapping,
66-
TOOL_USE_SPECIAL_TOKEN_IDS,
67-
validate_required_special_token_ids,
68-
)
69-
from cppmega_mlx.data.token_dataset import (
70-
BatchCursor,
71-
TokenDatasetMetadata,
72-
TokenNpzDataset,
73-
iterate_token_batches,
74-
open_token_dataset,
75-
)
3+
The package also contains portable corpus/indexer modules used on Linux hosts
4+
where Apple's MLX runtime is unavailable. Public MLX-backed exports are loaded
5+
on first attribute access so importing a portable submodule does not eagerly
6+
import the entire training runtime.
7+
"""
768

77-
__all__ = [
78-
"BatchCursor",
79-
"EOT_ID",
80-
"FIMSpecialTokenIds",
81-
"FIM_INSTRUCTION_ID",
82-
"FIMMode",
83-
"FIM_MIDDLE_ID",
84-
"FIM_PREFIX_ID",
85-
"FIM_SPECIAL_TOKEN_IDS",
86-
"FIM_SUFFIX_ID",
87-
"LMTokenBatch",
88-
"LocalTokenBatchDataset",
89-
"MAX_PLATFORM_IDS",
90-
"MegatronIndexedDataset",
91-
"MegatronIndexedMetadata",
92-
"MegatronIndexedMultiShardDataset",
93-
"MegatronIndexedMultiShardMetadata",
94-
"MultiShardTokenParquetDataset",
95-
"OversizedSamplePolicy",
96-
"PLATFORM_VOCAB",
97-
"PLATFORM_VOCAB_SIZE",
98-
"ParquetColumns",
99-
"PackedSequences",
100-
"PackingStrategy",
101-
"PlatformContext",
102-
"REQUIRED_SPECIAL_TOKEN_IDS",
103-
"SpecialTokenMapping",
104-
"TokenDatasetMetadata",
105-
"TorchDataLoaderBridgeConfig",
106-
"TorchDataLoaderBridgeError",
107-
"TOOL_USE_SPECIAL_TOKEN_IDS",
108-
"TokenNpzDataset",
109-
"TokenParquetDataset",
110-
"apply_fim_permutation",
111-
"apply_fim_transform",
112-
"apply_ifim_permutation",
113-
"apply_ifim_transform",
114-
"build_spawn_dataloader",
115-
"cumulative_doc_ids_from_eos",
116-
"document_boundary_mask",
117-
"encode_platform_context",
118-
"ensure_lm_batch",
119-
"extract_ifim_instruction_text",
120-
"is_torch_dataloader_available",
121-
"iter_mlx_batches",
122-
"iterate_token_batches",
123-
"megatron_indexed_side_channel_schema",
124-
"mlx_cumulative_doc_ids_from_eos",
125-
"mlx_document_boundary_mask",
126-
"mlx_sequence_packing_attention_mask",
127-
"open_megatron_indexed_dataset",
128-
"open_token_dataset",
129-
"pack_bos_aligned_best_fit",
130-
"pack_documents_with_eos",
131-
"parse_platform_context",
132-
"platform_ids_array",
133-
"render_platform_context",
134-
"sample_middle_span",
135-
"synthetic_token_batch",
136-
"validate_required_special_token_ids",
137-
]
9+
from __future__ import annotations
10+
11+
from importlib import import_module
12+
from typing import Any
13+
14+
15+
_EXPORT_MODULES = {
16+
"LMTokenBatch": "cppmega_mlx.data.batch",
17+
"ensure_lm_batch": "cppmega_mlx.data.batch",
18+
"synthetic_token_batch": "cppmega_mlx.data.batch",
19+
"LocalTokenBatchDataset": "cppmega_mlx.data.dataloader_bridge",
20+
"TorchDataLoaderBridgeConfig": "cppmega_mlx.data.dataloader_bridge",
21+
"TorchDataLoaderBridgeError": "cppmega_mlx.data.dataloader_bridge",
22+
"build_spawn_dataloader": "cppmega_mlx.data.dataloader_bridge",
23+
"is_torch_dataloader_available": "cppmega_mlx.data.dataloader_bridge",
24+
"iter_mlx_batches": "cppmega_mlx.data.dataloader_bridge",
25+
"EOT_ID": "cppmega_mlx.data.fim",
26+
"FIMSpecialTokenIds": "cppmega_mlx.data.fim",
27+
"FIM_INSTRUCTION_ID": "cppmega_mlx.data.fim",
28+
"FIMMode": "cppmega_mlx.data.fim",
29+
"FIM_MIDDLE_ID": "cppmega_mlx.data.fim",
30+
"FIM_PREFIX_ID": "cppmega_mlx.data.fim",
31+
"FIM_SPECIAL_TOKEN_IDS": "cppmega_mlx.data.fim",
32+
"FIM_SUFFIX_ID": "cppmega_mlx.data.fim",
33+
"apply_fim_permutation": "cppmega_mlx.data.fim",
34+
"apply_fim_transform": "cppmega_mlx.data.fim",
35+
"apply_ifim_permutation": "cppmega_mlx.data.fim",
36+
"apply_ifim_transform": "cppmega_mlx.data.fim",
37+
"extract_ifim_instruction_text": "cppmega_mlx.data.fim",
38+
"sample_middle_span": "cppmega_mlx.data.fim",
39+
"MegatronIndexedDataset": "cppmega_mlx.data.megatron_indexed",
40+
"MegatronIndexedMetadata": "cppmega_mlx.data.megatron_indexed",
41+
"MegatronIndexedMultiShardDataset": "cppmega_mlx.data.megatron_indexed",
42+
"MegatronIndexedMultiShardMetadata": "cppmega_mlx.data.megatron_indexed",
43+
"megatron_indexed_side_channel_schema": "cppmega_mlx.data.megatron_indexed",
44+
"open_megatron_indexed_dataset": "cppmega_mlx.data.megatron_indexed",
45+
"OversizedSamplePolicy": "cppmega_mlx.data.packing",
46+
"PackedSequences": "cppmega_mlx.data.packing",
47+
"PackingStrategy": "cppmega_mlx.data.packing",
48+
"cumulative_doc_ids_from_eos": "cppmega_mlx.data.packing",
49+
"document_boundary_mask": "cppmega_mlx.data.packing",
50+
"mlx_cumulative_doc_ids_from_eos": "cppmega_mlx.data.packing",
51+
"mlx_document_boundary_mask": "cppmega_mlx.data.packing",
52+
"mlx_sequence_packing_attention_mask": "cppmega_mlx.data.packing",
53+
"pack_bos_aligned_best_fit": "cppmega_mlx.data.packing",
54+
"pack_documents_with_eos": "cppmega_mlx.data.packing",
55+
"MultiShardTokenParquetDataset": "cppmega_mlx.data.parquet_dataset",
56+
"ParquetColumns": "cppmega_mlx.data.parquet_dataset",
57+
"TokenParquetDataset": "cppmega_mlx.data.parquet_dataset",
58+
"MAX_PLATFORM_IDS": "cppmega_mlx.data.platform_context",
59+
"PLATFORM_VOCAB": "cppmega_mlx.data.platform_context",
60+
"PLATFORM_VOCAB_SIZE": "cppmega_mlx.data.platform_context",
61+
"PlatformContext": "cppmega_mlx.data.platform_context",
62+
"encode_platform_context": "cppmega_mlx.data.platform_context",
63+
"parse_platform_context": "cppmega_mlx.data.platform_context",
64+
"platform_ids_array": "cppmega_mlx.data.platform_context",
65+
"render_platform_context": "cppmega_mlx.data.platform_context",
66+
"REQUIRED_SPECIAL_TOKEN_IDS": "cppmega_mlx.data.tokenizer_contract",
67+
"SpecialTokenMapping": "cppmega_mlx.data.tokenizer_contract",
68+
"TOOL_USE_SPECIAL_TOKEN_IDS": "cppmega_mlx.data.tokenizer_contract",
69+
"validate_required_special_token_ids": "cppmega_mlx.data.tokenizer_contract",
70+
"BatchCursor": "cppmega_mlx.data.token_dataset",
71+
"TokenDatasetMetadata": "cppmega_mlx.data.token_dataset",
72+
"TokenNpzDataset": "cppmega_mlx.data.token_dataset",
73+
"iterate_token_batches": "cppmega_mlx.data.token_dataset",
74+
"open_token_dataset": "cppmega_mlx.data.token_dataset",
75+
}
76+
77+
__all__ = list(_EXPORT_MODULES)
78+
79+
80+
def __getattr__(name: str) -> Any:
81+
module_name = _EXPORT_MODULES.get(name)
82+
if module_name is None:
83+
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
84+
value = getattr(import_module(module_name), name)
85+
globals()[name] = value
86+
return value
87+
88+
89+
def __dir__() -> list[str]:
90+
return sorted(set(globals()) | set(__all__))

‎tests/test_data_package_imports.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
from __future__ import annotations
2+
3+
import subprocess
4+
import sys
5+
from pathlib import Path
6+
7+
8+
ROOT = Path(__file__).resolve().parents[1]
9+
10+
11+
def test_portable_data_submodule_does_not_import_mlx() -> None:
12+
code = """
13+
import builtins
14+
15+
real_import = builtins.__import__
16+
17+
def reject_mlx(name, globals=None, locals=None, fromlist=(), level=0):
18+
if name == "mlx" or name.startswith("mlx."):
19+
raise AssertionError(f"portable data import reached MLX: {name}")
20+
return real_import(name, globals, locals, fromlist, level)
21+
22+
builtins.__import__ = reject_mlx
23+
from cppmega_mlx.data.nanochat_pipeline.language_info import detect_language_info
24+
25+
info = detect_language_info("int main() { return 0; }", filepath="src/example.cpp")
26+
assert info["primary_language"] == "c++"
27+
"""
28+
result = subprocess.run(
29+
[sys.executable, "-c", code],
30+
cwd=ROOT,
31+
text=True,
32+
capture_output=True,
33+
timeout=30,
34+
check=False,
35+
)
36+
assert result.returncode == 0, result.stderr

0 commit comments

Comments
 (0)