Buckets:
| import json | |
| from types import SimpleNamespace | |
| import pytest | |
| import torch | |
| from zebra_sft.evaluate_task_subset import ( | |
| _acquire_output_lock, | |
| _control_collective_device, | |
| _manifest_rows, | |
| _load_prediction_rows, | |
| _limit_manifest_rows, | |
| _merge_rank_shards, | |
| _summarize, | |
| _wait_for_adapter, | |
| _worker_device_index, | |
| ) | |
| from zebra_sft.generated_codes import save_generated_codes | |
| def test_worker_device_index_maps_multiple_local_workers_to_each_gpu(): | |
| assert [_worker_device_index(rank, 1, 3) for rank in range(3)] == [0, 0, 0] | |
| assert [_worker_device_index(rank, 2, 2) for rank in range(4)] == [0, 0, 1, 1] | |
| with pytest.raises(ValueError, match="exceeds capacity"): | |
| _worker_device_index(local_rank=3, visible_devices=1, workers_per_gpu=3) | |
| def test_wait_for_adapter_uses_atomic_complete_marker(tmp_path): | |
| adapter = tmp_path / "adapter" | |
| adapter.mkdir() | |
| args = SimpleNamespace( | |
| wait_for_adapter=True, | |
| mode="interleaved", | |
| adapter_dir=adapter, | |
| ) | |
| sleeps = [] | |
| def publish(seconds): | |
| sleeps.append(seconds) | |
| (adapter / "COMPLETE").write_text("complete\n") | |
| _wait_for_adapter(args, sleep=publish) | |
| assert sleeps == [15.0] | |
| def test_gloo_control_collectives_use_cpu(monkeypatch): | |
| monkeypatch.setattr( | |
| "zebra_sft.evaluate_task_subset.dist.get_backend", lambda: "gloo" | |
| ) | |
| assert _control_collective_device(2, torch.device("cuda", 0)).type == "cpu" | |
| assert _control_collective_device(1, torch.device("cuda", 0)).type == "cuda" | |
| def test_output_lock_rejects_a_concurrent_writer(tmp_path): | |
| first = _acquire_output_lock(tmp_path) | |
| try: | |
| with pytest.raises(RuntimeError, match="already has an active writer"): | |
| _acquire_output_lock(tmp_path) | |
| finally: | |
| first.close() | |
| second = _acquire_output_lock(tmp_path) | |
| second.close() | |
| def test_prediction_resume_truncates_only_torn_tail(tmp_path): | |
| path = tmp_path / "predictions.jsonl" | |
| path.write_bytes(b'{"sample_id":"a"}\n{"sample_id":"tor') | |
| rows = _load_prediction_rows(path, resume=True) | |
| assert rows == [{"sample_id": "a"}] | |
| assert path.read_bytes() == b'{"sample_id":"a"}\n' | |
| def test_prediction_resume_rolls_back_uncommitted_sample_suffix(tmp_path): | |
| path = tmp_path / "predictions.jsonl" | |
| samples = tmp_path / "samples" | |
| complete = samples / "complete" | |
| complete.mkdir(parents=True) | |
| (complete / "target.txt").write_text("1") | |
| (complete / "output.txt").write_text("Final answer: 1") | |
| row = {"sample_id": "complete", "evaluation_index": 0} | |
| (complete / "result.json").write_text(json.dumps(row)) | |
| save_generated_codes(complete / "generated_codes.pt", []) | |
| incomplete = {"sample_id": "incomplete", "evaluation_index": 1} | |
| path.write_text( | |
| json.dumps(row) + "\n" + json.dumps(incomplete) + "\n", | |
| encoding="utf-8", | |
| ) | |
| rows = _load_prediction_rows(path, resume=True, samples_dir=samples) | |
| assert rows == [row] | |
| assert [json.loads(line) for line in path.read_text().splitlines()] == [row] | |
| def test_manifest_selection_supports_indexed_and_complete_splits(tmp_path): | |
| rows = [ | |
| {"sample_id": "train", "split": "train", "evaluation_subset_index": None}, | |
| {"sample_id": "val-b", "split": "validation", "evaluation_subset_index": None}, | |
| {"sample_id": "val-a", "split": "validation", "evaluation_subset_index": None}, | |
| {"sample_id": "test-b", "split": "test", "evaluation_subset_index": 1}, | |
| {"sample_id": "test-a", "split": "test", "evaluation_subset_index": 0}, | |
| {"sample_id": "test-c", "split": "test", "evaluation_subset_index": None}, | |
| ] | |
| (tmp_path / "manifest.jsonl").write_text( | |
| "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8" | |
| ) | |
| indexed = _manifest_rows(tmp_path, "indexed") | |
| validation = _manifest_rows(tmp_path, "validation") | |
| complete = _manifest_rows(tmp_path, "test") | |
| assert [row["sample_id"] for row in indexed] == ["test-a", "test-b"] | |
| assert [row["sample_id"] for row in complete] == [ | |
| "test-a", | |
| "test-b", | |
| "test-c", | |
| ] | |
| assert [row["evaluation_index"] for row in complete] == [0, 1, 2] | |
| assert [row["sample_id"] for row in validation] == ["val-a", "val-b"] | |
| def test_manifest_limit_uses_deterministic_window_and_validates_bounds(): | |
| rows = [{"sample_id": value} for value in ("a", "b", "c")] | |
| assert _limit_manifest_rows(rows, None) is rows | |
| assert _limit_manifest_rows(rows, 2) == rows[:2] | |
| assert _limit_manifest_rows(rows, 2, sample_offset=1) == rows[1:3] | |
| assert _limit_manifest_rows(rows, None, sample_offset=1) == rows[1:] | |
| with pytest.raises(ValueError, match="positive"): | |
| _limit_manifest_rows(rows, 0) | |
| with pytest.raises(ValueError, match="nonnegative"): | |
| _limit_manifest_rows(rows, 1, sample_offset=-1) | |
| def test_rollout_summary_reports_validation_breakdown_and_tokens(): | |
| rows = [ | |
| { | |
| "strict_correct": True, | |
| "score": {"correct": True}, | |
| "stop_reason": "im_end", | |
| "expected_images": 1, | |
| "generated_images": 1, | |
| "expected_grids": [[10, 16]], | |
| "generated_grids": [[10, 16]], | |
| "events": [ | |
| {"type": "text", "token_ids": [1, 2, 3]}, | |
| {"type": "image", "serialized_positions": 160}, | |
| ], | |
| }, | |
| { | |
| "strict_correct": False, | |
| "score": {"correct": False}, | |
| "stop_reason": "cycle", | |
| "expected_images": 4, | |
| "generated_images": 2, | |
| "expected_grids": [[1, 1]] * 4, | |
| "generated_grids": [[1, 1]] * 2, | |
| "events": [{"type": "text", "token_ids": [4, 5]}], | |
| }, | |
| ] | |
| args = SimpleNamespace( | |
| mode="interleaved", | |
| model_label="fixture", | |
| model_path="base", | |
| adapter_dir="adapter", | |
| selection="validation", | |
| workers_per_gpu=2, | |
| scheduling="dynamic", | |
| text_cycle_max_period=128, | |
| text_cycle_min_tokens=64, | |
| decode_media=False, | |
| decoder_steps=28, | |
| decoder_guidance=1.5, | |
| decoder_upsampling=1.0, | |
| ) | |
| summary = _summarize(rows, args, world=2) | |
| assert summary["subset_field"] == "split:validation" | |
| assert summary["accuracy"] == 0.5 | |
| assert summary["by_expected_images"]["1"]["accuracy"] == 1.0 | |
| assert summary["by_expected_images"]["4"]["generated_images"] == 2 | |
| assert summary["complete_grid_sequence_rate"] == 0.5 | |
| assert summary["text_cycle_max_period"] == 128 | |
| assert summary["text_cycle_min_tokens"] == 64 | |
| assert summary["assistant_tokens"] == { | |
| "text": 5, | |
| "visual": 160, | |
| "total": 165, | |
| "mean_text_per_sample": 2.5, | |
| "mean_visual_per_sample": 80.0, | |
| } | |
| def test_merge_rank_shards_enforces_four_rank_boundaries(tmp_path): | |
| manifest_rows = [ | |
| {"sample_id": f"sample-{index}", "evaluation_index": index} | |
| for index in range(9) | |
| ] | |
| for rank in range(4): | |
| rank_dir = tmp_path / "shards" / f"rank-{rank:02d}" | |
| rank_dir.mkdir(parents=True) | |
| rows = manifest_rows[rank::4] | |
| (rank_dir / "predictions.jsonl").write_text( | |
| "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8" | |
| ) | |
| assert _merge_rank_shards(tmp_path, manifest_rows, 4) == manifest_rows | |
| def test_merge_rank_shards_accepts_complete_dynamic_queue_assignments(tmp_path): | |
| manifest_rows = [ | |
| {"sample_id": f"sample-{index}", "evaluation_index": index} | |
| for index in range(7) | |
| ] | |
| assignments = ( | |
| [manifest_rows[index] for index in [0, 3, 4, 6]], | |
| [manifest_rows[index] for index in [1, 2, 5]], | |
| ) | |
| for rank, rows in enumerate(assignments): | |
| rank_dir = tmp_path / "shards" / f"rank-{rank:02d}" | |
| rank_dir.mkdir(parents=True) | |
| (rank_dir / "predictions.jsonl").write_text( | |
| "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8" | |
| ) | |
| assert _merge_rank_shards( | |
| tmp_path, manifest_rows, 2, scheduling="dynamic" | |
| ) == manifest_rows | |
| def test_merge_rank_shards_preserves_nonzero_manifest_indices(tmp_path, scheduling): | |
| manifest_rows = [ | |
| {"sample_id": f"sample-{index}", "evaluation_index": index} | |
| for index in range(100, 104) | |
| ] | |
| for rank in range(2): | |
| rank_dir = tmp_path / "shards" / f"rank-{rank:02d}" | |
| rank_dir.mkdir(parents=True) | |
| rows = manifest_rows[rank::2] | |
| (rank_dir / "predictions.jsonl").write_text( | |
| "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8" | |
| ) | |
| assert _merge_rank_shards( | |
| tmp_path, manifest_rows, 2, scheduling=scheduling | |
| ) == manifest_rows | |
Xet Storage Details
- Size:
- 8.85 kB
- Xet hash:
- 9c257053f6ebd169069366f079bb52b06747ef9afeae0715ca95d6198da0a4cf
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.