| """Real Accelerator/Gloo entry-point tests with empty ranks and windows.""" |
| from copy import deepcopy |
| from datetime import timedelta |
| import json |
| import os |
| from pathlib import Path |
| import sys |
|
|
| import pytest |
| import torch |
|
|
| from fastfill.v2.tests.test_active_objective_training import _row, _yaw_config |
|
|
|
|
| def _worker(rank, rendezvous, directory, mode): |
| from torch import distributed as dist |
| from fastfill.v2 import train |
| os.environ.update(RANK=str(rank), WORLD_SIZE="2", LOCAL_RANK=str(rank), LOCAL_WORLD_SIZE="2", |
| OMP_NUM_THREADS="1", MASTER_ADDR="127.0.0.1", MASTER_PORT="29513") |
| dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=2, |
| timeout=timedelta(seconds=20)) |
| patches = pytest.MonkeyPatch() |
| try: |
| if sys.platform == "darwin": |
| |
| |
| |
| patches.setattr(torch._C, "_get_accelerator", lambda: torch.device("cpu")) |
| root = Path(directory) |
| rows_count = {"peer_empty": 2, "last_empty": 4, "empty_then_valid": 8, "all_empty": 4}[mode] |
| data = root / "train.jsonl" |
| if rank == 0: |
| data.write_text("".join(json.dumps(_row(str(i))) + "\n" for i in range(rows_count))) |
| dist.barrier() |
| original_collate, original_preflight, original_loader = train.collate_samples, train._preflight, train.DataLoader |
| state = {"after_preflight": False} |
| def collate(rows, *args, **kwargs): |
| batch = original_collate(rows, *args, **kwargs) |
| if state["after_preflight"]: |
| yaw = batch["validity"]["yaw"].clone() |
| for i, row in enumerate(rows): |
| identity = int(row["provenance"]["scene_id"]) |
| empty = (mode == "all_empty" or (mode == "peer_empty" and rank == 0) |
| or (mode == "last_empty" and identity >= 2) |
| or (mode == "empty_then_valid" and identity < 4)) |
| if empty: |
| yaw[i] = False |
| batch = {**batch, "validity": {**batch["validity"], "yaw": yaw}} |
| return batch |
| def preflight(*args, **kwargs): |
| state["after_preflight"] = False |
| result = original_preflight(*args, **kwargs) |
| state["after_preflight"] = True |
| return result |
| def loader(rows, **kwargs): |
| return original_loader(rows, **{**kwargs, "shuffle": False}) |
| models, before = [], [] |
| original_build = train.build_model |
| def build(config): |
| model = original_build(config) |
| models.append(model) |
| before.append(deepcopy(model.state_dict())) |
| return model |
| patches.setattr(train, "collate_samples", collate) |
| patches.setattr(train, "_preflight", preflight) |
| patches.setattr(train, "DataLoader", loader) |
| patches.setattr(train, "build_model", build) |
| config = _yaw_config() |
| config["training"]["gradient_accumulation_steps"] = 1 if mode == "peer_empty" else 2 |
| if mode == "all_empty": |
| with pytest.raises(RuntimeError, match="no active objective"): |
| train.run_training(config, data, root / "run") |
| changed = sum(not torch.equal(before[0][name], value) for name, value in models[0].state_dict().items()) |
| assert changed == 0 |
| observation = {"optimizer_updates": 0, "changed_parameter_tensors": changed} |
| else: |
| logs = train.run_training(config, data, root / "run") |
| assert len(logs) == 1 and logs[0]["step"] == 1 |
| assert logs[0]["active_objective_count"] == {"peer_empty": 1, "last_empty": 2, "empty_then_valid": 4}[mode] |
| assert logs[0]["gradient_norms"]["yaw_logits_head"] > 0 |
| changed = sum(not torch.equal(before[0][name], value) for name, value in models[0].state_dict().items()) |
| assert changed > 0 |
| observation = {"optimizer_updates": 1, "active_objective_count": logs[0]["active_objective_count"], |
| "changed_parameter_tensors": changed} |
| if rank == 0: |
| manifest = json.loads((root / "run/run_manifest.json").read_text()) |
| assert manifest["world_size"] == 2 |
| assert manifest["skipped_no_objective_windows"] == int(mode == "empty_then_valid") |
| (root / f"rank-{rank}.json").write_text(json.dumps(observation, indent=2) + "\n") |
| finally: |
| patches.undo() |
| if dist.is_initialized(): |
| dist.destroy_process_group() |
|
|
|
|
| @pytest.mark.skipif(not torch.distributed.is_available() or not torch.distributed.is_gloo_available(), reason="requires Gloo") |
| @pytest.mark.parametrize("mode", ["peer_empty", "last_empty", "empty_then_valid", "all_empty"]) |
| def test_actual_two_rank_objective_windows(tmp_path, monkeypatch, mode): |
| monkeypatch.setenv("GLOO_SOCKET_IFNAME", "lo0" if sys.platform == "darwin" else "lo") |
| monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") |
| torch.multiprocessing.spawn(_worker, args=("file://" + str(tmp_path / "rendezvous"), str(tmp_path), mode), |
| nprocs=2, join=True) |
| assert all((tmp_path / f"rank-{rank}.json").exists() for rank in range(2)) |
|
|
