social-data-attribution-w2/tests/test_draw_samples_modal.py

from __future__ import annotations

import importlib.util
import io
import json
from pathlib import Path
from types import SimpleNamespace

import pandas as pd
import pytest

from dolma.constants import FORMATS, TOPICS


def load_script():
    path = Path("scripts/modal/draw_samples.py")
    spec = importlib.util.spec_from_file_location("draw_samples", path)
    assert spec is not None
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def parquet_bytes(doc_ids: list[str]) -> bytes:
    buffer = io.BytesIO()
    pd.DataFrame({"doc_id": doc_ids}).to_parquet(buffer, index=False)
    return buffer.getvalue()


class FakeS3:
    def __init__(self, objects: dict[str, bytes]) -> None:
        self.objects = objects

    def get_object(self, *, Bucket: str, Key: str) -> dict[str, io.BytesIO]:
        assert Bucket == "soc127-dedup"
        return {"Body": io.BytesIO(self.objects[Key])}


class FakeSampleChunk:
    def __init__(
        self,
        sample_root: Path,
        chunk_results_dir: str,
        rows: list[dict[str, object]],
    ) -> None:
        self.calls: list[list[dict[str, object]]] = []
        self.sample_root = sample_root
        self.chunk_results_dir = chunk_results_dir
        self.rows = rows

    def map(self, chunks: list[str]) -> list[str]:
        payloads = [json.loads(chunk) for chunk in chunks]
        self.calls.append(payloads)
        results = []
        for payload in payloads:
            chunk_id = int(payload["chunk_id"])
            run_id = str(payload["run_id"])
            chunk_dir = self.sample_root / self.chunk_results_dir / run_id
            chunk_dir.mkdir(parents=True, exist_ok=True)
            pd.DataFrame(self.rows).to_parquet(
                chunk_dir / f"chunk_{chunk_id:04d}.parquet",
                index=False,
            )
            results.append(
                json.dumps(
                    {
                        "chunk_id": chunk_id,
                        "status": "ok",
                        "rows": len(self.rows),
                    }
                )
            )
        return results


def test_prepare_exclusion_set_writes_to_manifest_mount(tmp_path, monkeypatch):
    module = load_script()
    commit_calls: list[str] = []
    manifest_root = tmp_path / "manifests"

    monkeypatch.setattr(module, "MANIFEST_MOUNT", str(manifest_root))
    monkeypatch.setattr(
        module,
        "manifest_volume",
        SimpleNamespace(commit=lambda: commit_calls.append("commit")),
    )
    monkeypatch.setattr(
        module,
        "get_s3_client",
        lambda: FakeS3(
            {
                "exclude_a.parquet": parquet_bytes(["doc-a", "doc-b"]),
                "exclude_b.parquet": parquet_bytes(["doc-b", "doc-c"]),
            }
        ),
    )

    result = module._prepare_exclusion_set(
        "exclude_a.parquet,exclude_b.parquet,exclude_a.parquet"
    )

    assert result["exclude_paths_list"] == ["exclude_a.parquet", "exclude_b.parquet"]
    assert result["excluded_doc_count"] == 3
    assert Path(result["exclude_set_path"]).parent == manifest_root / "exclusions"
    assert commit_calls == ["commit"]

    written = pd.read_parquet(result["exclude_set_path"])
    assert set(written["doc_id"]) == {"doc-a", "doc-b", "doc-c"}


def test_run_draw_pipeline_uses_manifest_volume_exclusion_path(tmp_path, monkeypatch):
    module = load_script()
    manifest_mount = module.MANIFEST_MOUNT
    samples_mount = module.SAMPLES_MOUNT
    sample_dir = tmp_path / "samples"
    bin_key = f"{TOPICS[0]}|{FORMATS[0]}"
    fake_sample_chunk = FakeSampleChunk(
        sample_dir,
        module.CHUNK_RESULTS_DIR,
        [
            {
                "priority": 10,
                "doc_id": "doc-1",
                "token_count": 128,
                "shard_path": "s1",
                "bin_key": bin_key,
            }
        ],
    )

    monkeypatch.setattr(module, "SAMPLES_MOUNT", str(sample_dir))
    monkeypatch.setattr(
        module,
        "samples_volume",
        SimpleNamespace(commit=lambda: None, reload=lambda: None),
    )
    monkeypatch.setattr(
        module,
        "_prepare_exclusion_set",
        lambda value: {
            "exclude_paths_list": ["prev_sample.parquet"],
            "exclude_set_path": f"{module.MANIFEST_MOUNT}/exclusions/test.parquet",
            "excluded_doc_count": 7,
        },
    )
    monkeypatch.setattr(
        module,
        "r2_s3_list_keys",
        lambda prefix, suffix="": ["soc95-manifest/data/part-000.parquet"],
    )
    monkeypatch.setattr(module, "sample_chunk", fake_sample_chunk)

    result = json.loads(
        module.run_draw_pipeline.get_raw_f()(
            "1",
            42,
            1,
            512,
            0,
            exclude_manifest_keys="prev_sample.parquet",
        )
    )

    assert manifest_mount in module.run_draw_pipeline.spec.volumes
    assert samples_mount in module.run_draw_pipeline.spec.volumes
    assert result["status"] == "done"
    assert fake_sample_chunk.calls[0][0]["exclude_set_path"] == (
        f"{manifest_mount}/exclusions/test.parquet"
    )

    contract = json.loads(
        (sample_dir / "sample_1_docs" / "sample_contract.json").read_text()
    )
    assert contract["WORKING_SAMPLE_EXCLUDE_MANIFEST_PATHS"] == ["prev_sample.parquet"]
    assert contract["WORKING_SAMPLE_EXCLUDED_DOC_COUNT"] == 7


def test_local_output_uses_remote_exclusion_path(tmp_path, monkeypatch):
    module = load_script()
    bin_key = f"{TOPICS[0]}|{FORMATS[0]}"
    remote_sample_dir = tmp_path / "remote_samples"
    fake_sample_chunk = FakeSampleChunk(
        remote_sample_dir,
        module.CHUNK_RESULTS_DIR,
        [
            {
                "priority": 10,
                "doc_id": "doc-2",
                "token_count": 256,
                "shard_path": "s2",
                "bin_key": bin_key,
            }
        ],
    )
    remote_calls: list[str] = []

    def fake_prepare_remote(exclude_manifest_keys: str) -> dict[str, object]:
        remote_calls.append(exclude_manifest_keys)
        return {
            "exclude_paths_list": ["prev_a.parquet", "prev_b.parquet"],
            "exclude_set_path": f"{module.MANIFEST_MOUNT}/exclusions/helper.parquet",
            "excluded_doc_count": 9,
        }

    monkeypatch.chdir(tmp_path)
    monkeypatch.setattr(
        module,
        "list_keys_remote",
        SimpleNamespace(
            remote=lambda r2_prefix: ["soc95-manifest/data/part-000.parquet"]
        ),
    )
    monkeypatch.setattr(
        module,
        "prepare_exclusion_set_remote",
        SimpleNamespace(remote=fake_prepare_remote),
    )
    monkeypatch.setattr(module, "sample_chunk", fake_sample_chunk)
    monkeypatch.setattr(module, "SAMPLES_MOUNT", str(remote_sample_dir))
    monkeypatch.setattr(
        module,
        "samples_volume",
        SimpleNamespace(commit=lambda: None, reload=lambda: None),
    )

    module.main.info.raw_f(
        configs="1",
        seed=42,
        chunk_count=1,
        min_token_count=512,
        max_token_count=0,
        exclude_manifest="prev_a.parquet,prev_b.parquet",
        local_output=True,
    )

    assert remote_calls == ["prev_a.parquet,prev_b.parquet"]
    assert fake_sample_chunk.calls[0][0]["exclude_set_path"] == (
        f"{module.MANIFEST_MOUNT}/exclusions/helper.parquet"
    )

    contract = json.loads(
        (
            tmp_path / "data" / "samples" / "sample_1_docs" / "sample_contract.json"
        ).read_text()
    )
    assert contract["WORKING_SAMPLE_EXCLUDE_MANIFEST_PATHS"] == [
        "prev_a.parquet",
        "prev_b.parquet",
    ]
    assert contract["WORKING_SAMPLE_EXCLUDED_DOC_COUNT"] == 9


def test_merge_chunk_results_raises_when_expected_chunk_is_missing(
    tmp_path,
    monkeypatch,
):
    module = load_script()
    run_id = "draw_missing_chunk"
    chunk_dir = tmp_path / module.CHUNK_RESULTS_DIR / run_id
    chunk_dir.mkdir(parents=True)
    pd.DataFrame(
        [
            {
                "priority": 10,
                "doc_id": "doc-1",
                "token_count": 128,
                "shard_path": "s1",
                "bin_key": f"{TOPICS[0]}|{FORMATS[0]}",
            }
        ]
    ).to_parquet(chunk_dir / "chunk_0000.parquet", index=False)

    monkeypatch.setattr(module, "SAMPLES_MOUNT", str(tmp_path))
    monkeypatch.setattr(module, "samples_volume", SimpleNamespace(reload=lambda: None))

    with pytest.raises(RuntimeError, match="draw_missing_chunk"):
        module._merge_chunk_results(run_id, max_dpb=1, expected_chunk_count=2)


def test_sample_chunk_raises_when_exclusion_path_is_missing(tmp_path, monkeypatch):
    module = load_script()

    monkeypatch.setattr(module, "manifest_volume", SimpleNamespace(reload=lambda: None))
    monkeypatch.setattr(module, "samples_volume", SimpleNamespace(reload=lambda: None))

    with pytest.raises(FileNotFoundError, match="not available"):
        module.sample_chunk.get_raw_f()(
            json.dumps(
                {
                    "chunk_id": 0,
                    "keys": [],
                    "max_docs_per_bin": 1,
                    "seed": 42,
                    "exclude_set_path": str(tmp_path / "missing.parquet"),
                }
            )
        )