490 lines
21 KiB
Python
490 lines
21 KiB
Python
"""The rotating training pool: slices chosen, written, read, retired, and the
|
|
producer and trainer that share them.
|
|
|
|
What must hold: a slice lies inside one pristine segment with the model's span of
|
|
margin; a reader never sees a half-built slice and keeps reading one retired under
|
|
it; unreadable slots split a slice instead of entering it; the pool stays inside its
|
|
budget, oldest out first; and the manifest keeps evaluation and holdout data out.
|
|
"""
|
|
|
|
import datetime
|
|
import os
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import pool_producer
|
|
from conftest import solar_disc, write_fits
|
|
from suvi import availability, dataset, db, detectors, paths, pool
|
|
|
|
BASE = 1715400000 // paths.CADENCE * paths.CADENCE
|
|
SPAN = 4
|
|
SATELLITES = (16, 18)
|
|
|
|
|
|
def block(value=1000):
|
|
return np.full((6, dataset.SHARD_SIZE, dataset.SHARD_SIZE), value, dtype=np.uint16)
|
|
|
|
|
|
def all_good(satellite, times):
|
|
for t in times:
|
|
yield t, block()
|
|
|
|
|
|
def slots(count, start=0):
|
|
return (BASE + start * paths.CADENCE, BASE + (start + count) * paths.CADENCE)
|
|
|
|
|
|
# --------------------------------------------------------------------- choosing
|
|
|
|
|
|
def test_a_slice_lies_inside_one_segment_with_its_margin():
|
|
rng = np.random.default_rng(0)
|
|
segments = [slots(40), slots(30, start=100)]
|
|
for _ in range(200):
|
|
t0, t1 = pool.choose_slice(segments, SPAN, rng, targets=10)
|
|
assert any(a <= t0 and t1 <= b for a, b in segments)
|
|
assert (t1 - t0) // paths.CADENCE == 10 + 2 * SPAN
|
|
|
|
|
|
def test_a_short_segment_gives_its_whole_length():
|
|
t0, t1 = pool.choose_slice([slots(12)], SPAN, np.random.default_rng(1), targets=50)
|
|
assert (t0, t1) == slots(12)
|
|
|
|
|
|
def test_slices_are_drawn_by_target_count_and_never_from_spanless_segments():
|
|
rng = np.random.default_rng(2)
|
|
short = slots(2 * SPAN, start=500) # holds no target at this span
|
|
segments = [slots(20), slots(200, start=50), short]
|
|
picks = [pool.choose_slice(segments, SPAN, rng, targets=4) for _ in range(600)]
|
|
inside = [sum(a <= t0 < b for t0, _ in picks) for a, b in segments]
|
|
assert inside[2] == 0
|
|
assert inside[1] > 5 * inside[0] # 192 targets against 12
|
|
|
|
|
|
def test_choosing_from_nothing_trainable_is_an_error():
|
|
with pytest.raises(ValueError):
|
|
pool.choose_slice([slots(2 * SPAN)], SPAN, np.random.default_rng(0))
|
|
|
|
|
|
# ---------------------------------------------------------------------- slices
|
|
|
|
|
|
def test_a_written_slice_reads_back_whole(tmp_path):
|
|
piece = slots(20)
|
|
path = pool.write_slice(str(tmp_path), "0000001_x", piece, SATELLITES, all_good, SPAN)
|
|
shards, segments = pool.open_slice(path)
|
|
assert segments == [piece]
|
|
assert {satellite for _, satellite in shards} == set(SATELLITES)
|
|
for shard in shards.values():
|
|
assert len(shard.times()) == 20
|
|
np.testing.assert_array_equal(shard.raw(shard.times()[0]), block())
|
|
shard.close()
|
|
|
|
|
|
def test_a_slice_across_midnight_has_a_shard_per_day(tmp_path):
|
|
midnight = (BASE // 86400 + 1) * 86400
|
|
piece = (midnight - 5 * paths.CADENCE, midnight + 15 * paths.CADENCE)
|
|
path = pool.write_slice(str(tmp_path), "s", piece, SATELLITES, all_good, SPAN)
|
|
shards, _ = pool.open_slice(path)
|
|
assert len(shards) == 4
|
|
assert sum(len(s.times()) for s in shards.values()) == 2 * 20
|
|
for shard in shards.values():
|
|
shard.close()
|
|
|
|
|
|
def test_an_unreadable_slot_splits_the_slice(tmp_path):
|
|
piece = slots(30)
|
|
hole = BASE + 15 * paths.CADENCE
|
|
|
|
def holed(satellite, times):
|
|
for t in times:
|
|
yield t, None if (satellite == 18 and t == hole) else block()
|
|
|
|
path = pool.write_slice(str(tmp_path), "s", piece, SATELLITES, holed, SPAN)
|
|
_, segments = pool.open_slice(path)
|
|
assert segments == [(BASE, hole), (hole + paths.CADENCE, piece[1])]
|
|
|
|
|
|
def test_a_slice_with_nothing_left_to_train_on_is_dropped(tmp_path):
|
|
def nothing(satellite, times):
|
|
for t in times:
|
|
yield t, None
|
|
|
|
assert pool.write_slice(str(tmp_path), "s", slots(20), SATELLITES, nothing, SPAN) is None
|
|
assert os.listdir(tmp_path) == []
|
|
|
|
|
|
def test_a_failed_write_leaves_nothing_behind(tmp_path):
|
|
def broken(satellite, times):
|
|
yield times[0], block()
|
|
raise OSError("share went away")
|
|
|
|
with pytest.raises(OSError):
|
|
pool.write_slice(str(tmp_path), "s", slots(20), SATELLITES, broken, SPAN)
|
|
assert os.listdir(tmp_path) == []
|
|
|
|
|
|
def test_readers_see_only_complete_slices_oldest_first(tmp_path):
|
|
first = pool.write_slice(str(tmp_path), "b", slots(20), SATELLITES, all_good, SPAN)
|
|
second = pool.write_slice(str(tmp_path), "a", slots(20, 40), SATELLITES, all_good, SPAN)
|
|
os.makedirs(tmp_path / ".c.partial")
|
|
assert [path for path, _ in pool.ready_slices(str(tmp_path))] == [first, second]
|
|
assert pool.ready_slices(str(tmp_path / "absent")) == []
|
|
|
|
|
|
def test_retiring_drops_the_oldest_until_the_pool_fits(tmp_path):
|
|
made = [pool.write_slice(str(tmp_path), f"{i}", slots(20, 40 * i), SATELLITES,
|
|
all_good, SPAN) for i in range(4)]
|
|
size = pool.ready_slices(str(tmp_path))[0][1]
|
|
gone = pool.retire(str(tmp_path), budget_bytes=2.5 * size)
|
|
assert gone == made[:2]
|
|
assert [p for p, _ in pool.ready_slices(str(tmp_path))] == made[2:]
|
|
assert pool.retire(str(tmp_path), budget_bytes=10 * size) == []
|
|
|
|
|
|
def test_partials_from_a_stopped_producer_are_cleared(tmp_path):
|
|
os.makedirs(tmp_path / ".x.partial")
|
|
kept = pool.write_slice(str(tmp_path), "k", slots(20), SATELLITES, all_good, SPAN)
|
|
pool.clear_partials(str(tmp_path))
|
|
assert os.listdir(tmp_path) == [os.path.basename(kept)]
|
|
|
|
|
|
def test_the_newest_slice_is_never_retired(tmp_path):
|
|
only = pool.write_slice(str(tmp_path), "s", slots(20), SATELLITES, all_good, SPAN)
|
|
assert pool.retire(str(tmp_path), budget_bytes=0) == []
|
|
assert [p for p, _ in pool.ready_slices(str(tmp_path))] == [only]
|
|
|
|
|
|
def test_a_slice_retired_mid_visit_stays_readable(tmp_path):
|
|
path = pool.write_slice(str(tmp_path), "s", slots(20), SATELLITES, all_good, SPAN)
|
|
pool.write_slice(str(tmp_path), "t", slots(20, 40), SATELLITES, all_good, SPAN)
|
|
shards, _ = pool.open_slice(path)
|
|
pool.retire(str(tmp_path), budget_bytes=0)
|
|
assert not os.path.exists(path)
|
|
for shard in shards.values():
|
|
np.testing.assert_array_equal(shard.raw(shard.times()[-1]), block())
|
|
shard.close()
|
|
with pytest.raises(OSError):
|
|
pool.open_slice(path)
|
|
|
|
|
|
def test_slice_names_sort_by_counter():
|
|
names = [pool.slice_name(i, slots(20)) for i in (9, 10, 100)]
|
|
assert names == sorted(names)
|
|
|
|
|
|
# -------------------------------------------------------------------- manifest
|
|
|
|
|
|
def test_a_manifest_round_trips_and_a_foreign_version_is_refused(tmp_path):
|
|
manifest = {"version": pool.MANIFEST_VERSION, "segments": [[1, 2]], "frames": {}}
|
|
pool.save_manifest(manifest, str(tmp_path / "m.zst"))
|
|
assert pool.load_manifest(str(tmp_path / "m.zst")) == manifest
|
|
pool.save_manifest(dict(manifest, version=0), str(tmp_path / "old.zst"))
|
|
with pytest.raises(ValueError, match="pool_manifest.py"):
|
|
pool.load_manifest(str(tmp_path / "old.zst"))
|
|
|
|
|
|
def index(tmp_path, count, bad=(), missing=()):
|
|
"""`count` slots per satellite from BASE, every band judged good unless named."""
|
|
conn = db.connect(str(tmp_path / "index.sqlite"))
|
|
run = db.create_detector_run(conn, detectors.JUDGMENT_RUN, detectors.judgment_config())
|
|
rows = []
|
|
for satellite in SATELLITES:
|
|
for slot in range(count):
|
|
for wavelength in paths.WAVELENGTHS:
|
|
key = (satellite, slot, wavelength)
|
|
if key in missing:
|
|
continue
|
|
when = BASE + slot * paths.CADENCE
|
|
name = paths.FrameName(satellite=satellite, wavelength=wavelength,
|
|
t_start=when, t_end=when + paths.CADENCE,
|
|
version="1-0-2")
|
|
frame_id = db.upsert_frame(conn, name, name.relpath(), size_bytes=1_500_000)
|
|
rows.append((frame_id, "bad" if key in bad else "good", None, None, 0))
|
|
db.record_detections(conn, run, rows)
|
|
conn.commit()
|
|
return conn, run
|
|
|
|
|
|
def pristine(count, cuts=()):
|
|
"""Every band of both satellites GOOD, except one band BAD at each slot in `cuts`
|
|
-- which splits the era into segments there."""
|
|
states = np.full((2, len(availability.BANDS), count), availability.GOOD, np.uint8)
|
|
for cut in cuts:
|
|
states[0, 0, cut] = availability.BAD
|
|
return availability.Availability(states=states, t0=BASE, cadence=paths.CADENCE,
|
|
satellites=SATELLITES)
|
|
|
|
|
|
def manifest_for(conn, run, model, old_val=(), test=(), holdout=(), **kwargs):
|
|
kwargs.setdefault("val_new", 0)
|
|
return pool.build_manifest(conn, model, run, SPAN, list(old_val), list(test),
|
|
list(holdout), SATELLITES, **kwargs)
|
|
|
|
|
|
def test_the_manifest_lists_every_frame_of_its_segments(tmp_path):
|
|
conn, run = index(tmp_path, 30)
|
|
manifest = manifest_for(conn, run, pristine(30))
|
|
assert manifest["segments"] == [list(slots(30))]
|
|
assert manifest["val"] == [] and manifest["val_segments"] == []
|
|
assert len(manifest["frames"]["16"]) == len(manifest["frames"]["18"]) == 30
|
|
first = manifest["frames"]["18"][str(BASE)]
|
|
assert len(first) == 6 and all("g18" in path for path in first)
|
|
|
|
|
|
def test_a_slot_without_a_good_frame_of_every_band_splits_its_segment(tmp_path):
|
|
conn, run = index(tmp_path, 30, bad={(16, 12, 171)}, missing={(18, 20, 304)})
|
|
manifest = manifest_for(conn, run, pristine(30))
|
|
# 13..20 is 7 slots, below the 2*span+1 a target needs, so it goes too.
|
|
assert manifest["segments"] == [list(slots(12)),
|
|
[BASE + 21 * paths.CADENCE, BASE + 30 * paths.CADENCE]]
|
|
assert str(BASE + 12 * paths.CADENCE) not in manifest["frames"]["16"]
|
|
|
|
|
|
def test_test_segments_and_holdout_days_stay_out_of_the_pool(tmp_path):
|
|
conn, run = index(tmp_path, 30)
|
|
model = pristine(30)
|
|
near = (BASE + 200 * paths.CADENCE, BASE + 300 * paths.CADENCE) # 12 h later
|
|
assert manifest_for(conn, run, model, test=[near])["segments"] == []
|
|
far = (BASE + 10 * 86400, BASE + 10 * 86400 + 100 * paths.CADENCE)
|
|
assert manifest_for(conn, run, model, test=[far])["segments"] == [list(slots(30))]
|
|
day = datetime.datetime.fromtimestamp(BASE, datetime.timezone.utc).strftime("%Y-%m-%d")
|
|
assert manifest_for(conn, run, model, holdout=[day])["segments"] == []
|
|
|
|
|
|
#: Four segments: slots 0..49, 50..99 and 100..149 (49 each), and 150..200 (50).
|
|
CUTS = (49, 99, 149)
|
|
|
|
|
|
def segment(index):
|
|
end = 50 * index + 49 if index < 3 else 200
|
|
return [BASE + 50 * index * paths.CADENCE, BASE + end * paths.CADENCE]
|
|
|
|
|
|
def test_validation_spreads_over_the_era_and_leaves_the_pool(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(pool, "GUARD_SECONDS", 0)
|
|
conn, run = index(tmp_path, 200)
|
|
manifest = manifest_for(conn, run, pristine(200, CUTS), val_new=2, val_targets=10)
|
|
assert manifest["val_segments"] == [segment(1), segment(3)]
|
|
assert segment(0) in manifest["segments"] and segment(2) in manifest["segments"]
|
|
for a, b in manifest["segments"]: # the rest of 1 and 3 stays trainable
|
|
assert not any(a < high and b > low for low, high in manifest["val"])
|
|
assert len(manifest["val"]) == 2
|
|
for (a, b), (low, high) in zip(manifest["val"], manifest["val_segments"]):
|
|
assert low <= a and b <= high and (b - a) // paths.CADENCE == 10 + 2 * SPAN
|
|
for a, b in manifest["val"]:
|
|
assert all(str(t) in manifest["frames"]["16"] for t in range(a, b, paths.CADENCE))
|
|
|
|
|
|
def test_the_old_validation_segments_stay_in_validation(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(pool, "GUARD_SECONDS", 0)
|
|
conn, run = index(tmp_path, 200)
|
|
manifest = manifest_for(conn, run, pristine(200, CUTS), old_val=[segment(0)],
|
|
val_new=1, val_targets=10)
|
|
assert manifest["val_segments"] == [segment(0), segment(2)]
|
|
assert segment(1) in manifest["segments"] and segment(3) in manifest["segments"]
|
|
|
|
|
|
def test_new_validation_avoids_the_previous_models_training_data(tmp_path, monkeypatch):
|
|
"""So the previous model can be scored on the new set without having seen it."""
|
|
monkeypatch.setattr(pool, "GUARD_SECONDS", 0)
|
|
conn, run = index(tmp_path, 200)
|
|
manifest = manifest_for(conn, run, pristine(200, CUTS), val_new=4, val_targets=10,
|
|
old_train=[segment(2)])
|
|
assert manifest["val_segments"] == [segment(0), segment(1), segment(3)]
|
|
assert segment(2) in manifest["segments"]
|
|
|
|
|
|
def test_subtract_cuts_out_the_guarded_spans_and_drops_crumbs(monkeypatch):
|
|
monkeypatch.setattr(pool, "GUARD_SECONDS", 2 * paths.CADENCE)
|
|
a, b = slots(100)
|
|
cut = (a + 40 * paths.CADENCE, a + 50 * paths.CADENCE)
|
|
assert pool.subtract([(a, b)], [cut], minimum=9) == [
|
|
(a, a + 38 * paths.CADENCE), (a + 52 * paths.CADENCE, b)]
|
|
near_end = (b - 5 * paths.CADENCE, b)
|
|
assert pool.subtract([(a, b)], [near_end], minimum=9) == [(a, b - 7 * paths.CADENCE)]
|
|
assert pool.subtract([(a, a + 10 * paths.CADENCE)], [(a + 4 * paths.CADENCE,
|
|
a + 5 * paths.CADENCE)], 9) == []
|
|
|
|
|
|
def test_a_validation_window_never_comes_from_a_test_segment(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(pool, "GUARD_SECONDS", 0)
|
|
conn, run = index(tmp_path, 200)
|
|
manifest = manifest_for(conn, run, pristine(200, CUTS), test=[segment(1)],
|
|
val_new=3, val_targets=10)
|
|
assert segment(1) not in manifest["val_segments"]
|
|
assert segment(1) not in manifest["segments"]
|
|
|
|
|
|
def test_test_segments_are_kept_whole_with_their_frames(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(pool, "GUARD_SECONDS", 0)
|
|
conn, run = index(tmp_path, 200)
|
|
manifest = manifest_for(conn, run, pristine(200, CUTS), test=[segment(1)])
|
|
assert manifest["test"] == [segment(1)]
|
|
a, b = segment(1)
|
|
assert all(str(t) in manifest["frames"]["18"] for t in range(a, b, paths.CADENCE))
|
|
|
|
|
|
def test_a_short_validation_segment_gives_its_whole_length(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(pool, "GUARD_SECONDS", 0)
|
|
conn, run = index(tmp_path, 200)
|
|
old = [BASE + 3 * paths.CADENCE, BASE + 15 * paths.CADENCE] # 12 slots
|
|
manifest = manifest_for(conn, run, pristine(200, CUTS), old_val=[old],
|
|
val_targets=40)
|
|
assert manifest["val"] == [old]
|
|
|
|
|
|
def test_the_pool_keeps_a_day_clear_of_validation(tmp_path):
|
|
conn, run = index(tmp_path, 200)
|
|
manifest = manifest_for(conn, run, pristine(200, CUTS), val_new=1, val_targets=10)
|
|
assert manifest["val_segments"] and manifest["segments"] == []
|
|
|
|
|
|
# -------------------------------------------------------------------- producer
|
|
|
|
|
|
def archive(tmp_path, count):
|
|
"""A FITS archive of `count` slots per satellite plus its pool manifest."""
|
|
root = tmp_path / "Data"
|
|
frames = {str(s): {} for s in SATELLITES}
|
|
image = solar_disc(size=dataset.SHARD_SIZE, radius=200)
|
|
for satellite in SATELLITES:
|
|
for slot in range(count):
|
|
when = BASE + slot * paths.CADENCE
|
|
files = []
|
|
for wavelength in paths.WAVELENGTHS:
|
|
rel = f"g{satellite}/{wavelength}/{when}.fits"
|
|
os.makedirs(root / os.path.dirname(rel), exist_ok=True)
|
|
write_fits(str(root / rel), image)
|
|
files.append(rel)
|
|
frames[str(satellite)][str(when)] = files
|
|
manifest = {"version": pool.MANIFEST_VERSION, "satellites": list(SATELLITES),
|
|
"span": SPAN, "judgment_run": 1, "segments": [list(slots(count))],
|
|
"val": [list(slots(count))], "val_segments": [list(slots(count))],
|
|
"test": [list(slots(count // 2 + SPAN))], "frames": frames}
|
|
pool.save_manifest(manifest, str(tmp_path / "manifest.zst"))
|
|
return root
|
|
|
|
|
|
def test_the_producer_fills_the_pool_and_keeps_it_in_budget(tmp_path):
|
|
root = archive(tmp_path, 12)
|
|
directory = tmp_path / "pool"
|
|
pool_producer.main(["--manifest", str(tmp_path / "manifest.zst"), "--archive",
|
|
str(root), "--pool", str(directory), "--workers", "2",
|
|
"--slices", "3", "--seed", "0", "--budget-gb", "0"])
|
|
resident = pool.ready_slices(str(directory))
|
|
assert len(resident) == 1 # a zero budget keeps only the newest
|
|
assert os.path.basename(resident[0][0]).startswith("0000003")
|
|
shards, segments = pool.open_slice(resident[0][0])
|
|
assert segments == [slots(12)]
|
|
frame = next(iter(shards.values())).frames(BASE)
|
|
assert frame.shape == (6, dataset.SHARD_SIZE, dataset.SHARD_SIZE)
|
|
for shard in shards.values():
|
|
shard.close()
|
|
|
|
|
|
def test_an_unreadable_archive_file_becomes_a_hole_not_a_crash(tmp_path):
|
|
root = archive(tmp_path, 12)
|
|
os.remove(root / f"g18/171/{BASE + 6 * paths.CADENCE}.fits")
|
|
directory = tmp_path / "pool"
|
|
pool_producer.main(["--manifest", str(tmp_path / "manifest.zst"), "--archive",
|
|
str(root), "--pool", str(directory), "--workers", "1",
|
|
"--slices", "1", "--seed", "0"])
|
|
assert pool.ready_slices(str(directory)) == [] # 6 + 5 slots: below 2*span+1 each
|
|
|
|
|
|
def test_fixed_mode_writes_each_validation_window_once(tmp_path):
|
|
root = archive(tmp_path, 12)
|
|
directory = tmp_path / "val"
|
|
argv = ["--manifest", str(tmp_path / "manifest.zst"), "--archive", str(root),
|
|
"--pool", str(directory), "--workers", "1", "--fixed", "val"]
|
|
pool_producer.main(argv)
|
|
assert sorted(os.listdir(directory)) == ["val000"]
|
|
made = os.path.getmtime(directory / "val000" / pool.SLICE_FILE)
|
|
pool_producer.main(argv) # a rerun skips what is complete
|
|
assert os.path.getmtime(directory / "val000" / pool.SLICE_FILE) == made
|
|
test_dir = tmp_path / "test"
|
|
pool_producer.main(["--manifest", str(tmp_path / "manifest.zst"), "--archive",
|
|
str(root), "--pool", str(test_dir), "--workers", "1",
|
|
"--fixed", "test"])
|
|
_, segments = pool.open_slice(str(test_dir / "test000"))
|
|
assert segments == [slots(12 // 2 + SPAN)]
|
|
|
|
|
|
def test_encode_slot_declines_a_band_it_cannot_read(tmp_path):
|
|
assert pool_producer.encode_slot([str(tmp_path / "missing.fits")] * 6) is None
|
|
|
|
|
|
# ------------------------------------------------------------------ the trainer
|
|
|
|
|
|
torch = pytest.importorskip("torch")
|
|
|
|
|
|
def test_the_pool_set_yields_model_ready_samples(tmp_path):
|
|
import train
|
|
|
|
root = archive(tmp_path, 14)
|
|
directory = tmp_path / "pool"
|
|
pool_producer.main(["--manifest", str(tmp_path / "manifest.zst"), "--archive",
|
|
str(root), "--pool", str(directory), "--workers", "1",
|
|
"--slices", "1", "--seed", "0"])
|
|
data = train.PoolSet(str(directory), SATELLITES, offsets=(1, 4), per_visit=3)
|
|
stream = iter(data)
|
|
items = [next(stream) for _ in range(5)] # crosses a visit boundary
|
|
for item in items:
|
|
assert item["stack"].shape[1:] == (6, dataset.SHARD_SIZE, dataset.SHARD_SIZE)
|
|
assert item["lost"].sum() >= 1
|
|
batch = train.collate(items[:2])
|
|
assert batch["stack"].shape[0] == 2
|
|
|
|
|
|
def test_the_slice_set_is_a_fixed_validation_set(tmp_path):
|
|
import train
|
|
|
|
root = archive(tmp_path, 14)
|
|
directory = tmp_path / "val"
|
|
pool_producer.main(["--manifest", str(tmp_path / "manifest.zst"), "--archive",
|
|
str(root), "--pool", str(directory), "--workers", "1",
|
|
"--fixed", "val"])
|
|
data = train.SliceSet(str(directory), SATELLITES, offsets=(1, 4))
|
|
assert len(data) == 2 * (14 - 2 * SPAN)
|
|
first, again = data[3], data[3]
|
|
assert torch.equal(first["stack"], again["stack"]) # damage pinned to the index
|
|
assert torch.equal(first["lost"], again["lost"])
|
|
|
|
|
|
def test_the_slice_set_refuses_an_empty_directory(tmp_path):
|
|
import train
|
|
|
|
with pytest.raises(SystemExit):
|
|
train.SliceSet(str(tmp_path), SATELLITES)
|
|
|
|
|
|
def test_the_pool_set_gives_up_when_no_producer_appears(tmp_path, monkeypatch):
|
|
import train
|
|
|
|
monkeypatch.setattr(train.time, "sleep", lambda seconds: None)
|
|
data = train.PoolSet(str(tmp_path / "empty"), SATELLITES, wait=10.0)
|
|
with pytest.raises(RuntimeError, match="producer"):
|
|
next(iter(data))
|
|
|
|
|
|
def test_gpu_train_runs_the_producer_beside_a_pool_run():
|
|
import gpu_train
|
|
|
|
args = gpu_train.argparse.Namespace(
|
|
shards_dir="shards", name="v10", epochs=6, batch=4, lr=2e-4, base=32,
|
|
depth=3, workers=6, val_samples=96, eval_every=500, steps=8923,
|
|
lr_range=None, span=64, warmup=None, relative=False, corr_flow=False,
|
|
attention=True, max_steps=57000, overfit=None, resume=False, pool=True,
|
|
pool_manifest="pool_manifest.json.zst", pool_budget_gb=48.0, pool_workers=16,
|
|
val_dir="val10", evaluate=None)
|
|
assert f"--pool {gpu_train.POOL_DIR}" in gpu_train.command(args)
|
|
assert "--val-slices /work/val10" in gpu_train.command(args)
|
|
producer = gpu_train.producer_command(args)
|
|
assert "--archive /archive/Data" in producer and gpu_train.POOL_DIR in producer
|
|
args.pool = False
|
|
assert "--pool" not in gpu_train.command(args)
|