579 lines
23 KiB
Python
579 lines
23 KiB
Python
"""Generating the corpus straight into the transfer.
|
|
|
|
The properties under test are the ones the old two-stage build got for free and
|
|
this one has to earn: that a shard is never staged locally, that nothing larger
|
|
than one record is ever held, that an interrupted run resumes without forgetting
|
|
what it learned, and that the vetting contract did not move when the pipeline
|
|
around it did.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import dataset_build
|
|
from conftest import solar_disc, write_fits
|
|
from suvi import corpus, dataset, paths
|
|
|
|
CAD = paths.CADENCE
|
|
BASE = 1715385600 # 2024-05-11 00:00 UTC
|
|
DAY = "2024-05-11"
|
|
|
|
#: Scale putting a peak-1.0 synthetic disc's IMG_MEAN inside each band's
|
|
#: physical bounds (header_v1 rightly rejects a 304A frame as dim as a 94A one).
|
|
BAND_SCALE = {94: 0.5, 131: 0.5, 171: 1.0, 195: 2.0, 284: 2.0, 304: 3.0}
|
|
|
|
|
|
# ------------------------------------------------------------------ the archive
|
|
|
|
|
|
class FakeRow(dict):
|
|
"""sqlite3.Row-alike: the builder reads columns by name."""
|
|
|
|
|
|
class FakeConn:
|
|
"""Stands in for the index, counting reads so 'was this re-read' is testable."""
|
|
|
|
def __init__(self, rows):
|
|
self.rows = rows # [(satellite, wavelength, t_start, relpath)]
|
|
self.queries = 0
|
|
|
|
def execute(self, sql, params):
|
|
self.queries += 1
|
|
satellite, times = params[0], set(params[1:])
|
|
return [FakeRow(wavelength=w, t_start=t, path=p)
|
|
for (s, w, t, p) in self.rows if s == satellite and t in times]
|
|
|
|
|
|
def make_archive(root, times, satellites=(16, 18), broken=()):
|
|
"""Write a synthetic archive and return the rows an index would hold."""
|
|
rows = []
|
|
for satellite in satellites:
|
|
for t in times:
|
|
for wavelength in paths.WAVELENGTHS:
|
|
name = f"g{satellite}_{wavelength}_{t}.fits"
|
|
image = (np.zeros((1280, 1280), dtype=np.float32)
|
|
if (satellite, t) in broken
|
|
else solar_disc(size=1280, radius=386, peak=1.0))
|
|
write_fits(os.path.join(root, name),
|
|
image * BAND_SCALE[wavelength],
|
|
headers={"WAVELNTH": wavelength})
|
|
rows.append((satellite, wavelength, t, name))
|
|
return rows
|
|
|
|
|
|
# ------------------------------------------------------------------ destinations
|
|
|
|
|
|
def test_a_local_destination_writes_only_the_shard(tmp_path):
|
|
"""Nothing is staged: the destination is the only thing that appears."""
|
|
root = tmp_path / "archive"
|
|
out = tmp_path / "dest"
|
|
root.mkdir()
|
|
times = [BASE + i * CAD for i in range(3)]
|
|
rows = make_archive(str(root), times, satellites=(16,))
|
|
conn = FakeConn(rows)
|
|
destination = dataset_build.Destination(str(out))
|
|
|
|
size, kept, dropped, unusable = dataset_build.build_shard(
|
|
conn, str(root), DAY, 16, times, destination, workers=1)
|
|
|
|
assert kept == 3 and not dropped and not unusable and size > 0
|
|
assert os.listdir(out) == [dataset.shard_name(DAY, 16)]
|
|
with dataset.Shard(str(out / dataset.shard_name(DAY, 16))) as shard:
|
|
assert shard.times() == times
|
|
assert all(shard.usable(t) for t in times)
|
|
|
|
|
|
def test_a_failed_stream_leaves_nothing_at_the_real_name(tmp_path):
|
|
"""A shard is committed by its rename. A run that dies mid-stream must leave
|
|
no file the resume scan would mistake for finished work."""
|
|
out = tmp_path / "dest"
|
|
destination = dataset_build.Destination(str(out))
|
|
with pytest.raises(RuntimeError):
|
|
with destination.open("x.shard") as handle:
|
|
handle.write(b"partial")
|
|
raise RuntimeError("stream died")
|
|
assert not os.path.exists(out / "x.shard")
|
|
assert not os.path.exists(out / "x.shard.partial")
|
|
|
|
|
|
def test_existing_names_are_what_resume_is_keyed_on(tmp_path):
|
|
out = tmp_path / "dest"
|
|
out.mkdir()
|
|
(out / "a.shard").write_bytes(b"")
|
|
destination = dataset_build.Destination(str(out))
|
|
assert destination.existing() == {"a.shard"}
|
|
assert dataset_build.Destination(str(tmp_path / "absent")).existing() == set()
|
|
|
|
|
|
def test_put_and_get_round_trip_small_files(tmp_path):
|
|
destination = dataset_build.Destination(str(tmp_path / "dest"))
|
|
assert destination.get("vetting.json") is None
|
|
destination.put("vetting.json", b'[["2024-05-11", 16, 1000]]')
|
|
assert json.loads(destination.get("vetting.json")) == [["2024-05-11", 16, 1000]]
|
|
|
|
|
|
def test_a_remote_destination_streams_through_ssh_without_staging():
|
|
"""The remote path builds a command that pipes into cat and commits by rename;
|
|
no local file appears in it at any point."""
|
|
destination = dataset_build.Destination("host:/var/shards")
|
|
assert destination.host == "host" and destination.root == "/var/shards"
|
|
writer = dataset_build._RemoteWrite(destination, "a.shard")
|
|
assert "cat > " in writer.command
|
|
assert "a.shard.partial" in writer.command
|
|
assert "mv " in writer.command
|
|
|
|
|
|
# ------------------------------------------------------------------- generation
|
|
|
|
|
|
def test_the_shard_is_never_held_whole(tmp_path):
|
|
"""The memory property. Records are pulled one at a time from the pool, so a
|
|
600-slot shard is written with one ~4 MB payload live, not 1 GB."""
|
|
root = tmp_path / "archive"
|
|
root.mkdir()
|
|
times = [BASE + i * CAD for i in range(4)]
|
|
rows = make_archive(str(root), times, satellites=(16,))
|
|
conn = FakeConn(rows)
|
|
|
|
live = []
|
|
real_write = dataset.write_shard_stream
|
|
|
|
def watching(handle, day, satellite, wavelengths, records, level=3):
|
|
def counted():
|
|
for item in records:
|
|
live.append(1)
|
|
yield item
|
|
live.pop()
|
|
return real_write(handle, day, satellite, wavelengths, counted(), level)
|
|
|
|
dataset.write_shard_stream = watching
|
|
try:
|
|
dataset_build.build_shard(conn, str(root), DAY, 16, times,
|
|
dataset_build.Destination(str(tmp_path / "d")),
|
|
workers=1)
|
|
finally:
|
|
dataset.write_shard_stream = real_write
|
|
assert live == [] # every record released before the next was taken
|
|
|
|
|
|
def test_a_vetoed_slot_is_stored_as_a_sentinel_and_reported(tmp_path):
|
|
"""An all-zero frame parses fine and would enter the shard as 'truth' unvetted.
|
|
It must land as a sentinel *and* be named, because the segment it sits in has
|
|
to be split afterwards."""
|
|
root = tmp_path / "archive"
|
|
out = tmp_path / "dest"
|
|
root.mkdir()
|
|
times = [BASE + i * CAD for i in range(3)]
|
|
rows = make_archive(str(root), times, satellites=(16,),
|
|
broken={(16, times[1])})
|
|
conn = FakeConn(rows)
|
|
|
|
size, kept, dropped, unusable = dataset_build.build_shard(
|
|
conn, str(root), DAY, 16, times,
|
|
dataset_build.Destination(str(out)), workers=1)
|
|
|
|
assert kept == 2
|
|
assert unusable == {times[1]}
|
|
assert sum(dropped.values()) == 1
|
|
with dataset.Shard(str(out / dataset.shard_name(DAY, 16))) as shard:
|
|
assert times[1] in shard # we looked
|
|
assert not shard.usable(times[1]) # and there was nothing usable
|
|
assert shard.raw(times[1]) is None
|
|
|
|
|
|
def test_a_slot_missing_from_the_index_becomes_a_sentinel(tmp_path):
|
|
"""A wanted slot the archive does not have must still occupy its place, or the
|
|
shard's slots stop lining up with the segment that asked for them."""
|
|
root = tmp_path / "archive"
|
|
out = tmp_path / "dest"
|
|
root.mkdir()
|
|
times = [BASE + i * CAD for i in range(3)]
|
|
rows = [r for r in make_archive(str(root), times, satellites=(16,))
|
|
if r[2] != times[2]]
|
|
conn = FakeConn(rows)
|
|
|
|
_, kept, _, unusable = dataset_build.build_shard(
|
|
conn, str(root), DAY, 16, times,
|
|
dataset_build.Destination(str(out)), workers=1)
|
|
assert kept == 2 and unusable == {times[2]}
|
|
with dataset.Shard(str(out / dataset.shard_name(DAY, 16))) as shard:
|
|
assert shard.times() == times
|
|
|
|
|
|
def test_generation_defaults_to_a_single_reader():
|
|
"""One outstanding read against a shared mount is a promise to the other
|
|
consumers of it, not an implementation detail free to drift.
|
|
|
|
Read concurrency is the one lever that would make this faster -- 48 readers
|
|
reach ~207 MB/s against one reader's ~62 -- so the temptation to raise it is
|
|
permanent and the default is worth pinning.
|
|
"""
|
|
args = dataset_build.build_parser().parse_args(
|
|
["--availability", "a.npz", "--dest", "/tmp/x", "--budget-gb", "1"])
|
|
assert args.workers == 1
|
|
|
|
|
|
# ------------------------------------------------------------------ verification
|
|
|
|
|
|
def test_a_vetoed_slot_splits_the_segment_that_contained_it():
|
|
"""The profile is header-only; disc_v1 rejects inside a nominally pristine run.
|
|
Verification runs off what extraction observed, so nothing is read back."""
|
|
segment = (BASE, BASE + 10 * CAD)
|
|
vetoed = {(DAY, 16, BASE + 4 * CAD)}
|
|
verified = dataset_build.verified_segments(
|
|
{"train": [segment]}, vetoed, satellites=(16, 18), minimum_slots=3)
|
|
assert [pair for _, pair in verified] == [
|
|
(BASE, BASE + 4 * CAD), (BASE + 5 * CAD, BASE + 10 * CAD)]
|
|
|
|
|
|
def test_fragments_below_the_span_minimum_are_dropped():
|
|
"""A stub too short to hold one model input is not a segment."""
|
|
segment = (BASE, BASE + 10 * CAD)
|
|
vetoed = {(DAY, 16, BASE + 2 * CAD)}
|
|
verified = dataset_build.verified_segments(
|
|
{"train": [segment]}, vetoed, satellites=(16,), minimum_slots=5)
|
|
assert [pair for _, pair in verified] == [(BASE + 3 * CAD, BASE + 10 * CAD)]
|
|
|
|
|
|
def test_a_veto_on_either_satellite_splits_the_segment():
|
|
"""Segments are jointly pristine: a slot is usable only if every satellite is."""
|
|
segment = (BASE, BASE + 6 * CAD)
|
|
verified = dataset_build.verified_segments(
|
|
{"train": [segment]}, {(DAY, 18, BASE + 3 * CAD)},
|
|
satellites=(16, 18), minimum_slots=2)
|
|
assert [pair for _, pair in verified] == [
|
|
(BASE, BASE + 3 * CAD), (BASE + 4 * CAD, BASE + 6 * CAD)]
|
|
|
|
|
|
# --------------------------------------------------------------------- end to end
|
|
|
|
|
|
def test_a_whole_corpus_is_generated_into_the_destination(tmp_path, monkeypatch):
|
|
"""The pieces composed: profile -> slices -> shards -> manifest, with the
|
|
destination the only thing written and no library anywhere."""
|
|
from suvi import availability as availability_module, db
|
|
|
|
root = tmp_path / "archive"
|
|
out = tmp_path / "dest"
|
|
root.mkdir()
|
|
|
|
span_min = 2 # segments need 2*2+1 = 5 slots
|
|
slots = 24
|
|
times = [BASE + i * CAD for i in range(slots)]
|
|
rows = make_archive(str(root), times, satellites=(16, 18))
|
|
conn = FakeConn(rows)
|
|
|
|
states = np.full((2, len(availability_module.BANDS), slots), availability_module.GOOD,
|
|
dtype=np.uint8)
|
|
profile = availability_module.Availability(
|
|
states=states, t0=BASE, cadence=CAD, satellites=(16, 18))
|
|
profile.save(str(tmp_path / "availability.npz"))
|
|
|
|
monkeypatch.setattr(db, "connect", lambda *a, **k: conn)
|
|
monkeypatch.setattr(dataset_build.db, "connect", lambda *a, **k: conn)
|
|
monkeypatch.setattr(dataset_build.paths, "data_root", lambda: str(root))
|
|
monkeypatch.setattr(dataset_build, "holdout_days", lambda *a, **k: [])
|
|
monkeypatch.chdir(tmp_path / "archive") # cwd must stay untouched
|
|
|
|
before = set(os.listdir(root))
|
|
assert dataset_build.main([
|
|
"--availability", str(tmp_path / "availability.npz"),
|
|
"--dest", str(out), "--budget-gb", "1",
|
|
"--span-min", str(span_min), "--val-segments", "0",
|
|
"--test-segments", "0", "--workers", "1",
|
|
]) == 0
|
|
|
|
# Nothing was staged: the archive directory is exactly as it was.
|
|
assert set(os.listdir(root)) == before
|
|
|
|
produced = sorted(os.listdir(out))
|
|
assert "manifest.json" in produced and "segments.json" in produced
|
|
manifest = json.loads((out / "manifest.json").read_text())
|
|
assert manifest["satellites"] == [16, 18]
|
|
assert manifest["span_min"] == span_min
|
|
assert manifest["digest"]
|
|
assert manifest["segments"]["train"]
|
|
|
|
for satellite in (16, 18):
|
|
name = dataset.shard_name(DAY, satellite)
|
|
assert name in produced
|
|
with dataset.Shard(str(out / name)) as shard:
|
|
assert shard.times() # real slots, readable
|
|
assert all(shard.usable(t) for t in shard.times())
|
|
assert shard.raw(shard.times()[0]).shape == (6, 640, 640)
|
|
|
|
|
|
def test_a_resumed_build_skips_finished_shards_and_keeps_their_vetoes(tmp_path):
|
|
"""Nothing is local, so the destination is the whole memory of a run. A
|
|
resume that forgot the vetoes would emit a manifest claiming segments longer
|
|
than the data behind them."""
|
|
out = tmp_path / "dest"
|
|
destination = dataset_build.Destination(str(out))
|
|
destination.put("vetting.json", json.dumps([[DAY, 16, BASE + 4 * CAD]]).encode())
|
|
|
|
carried = {tuple(entry) for entry in json.loads(destination.get("vetting.json"))}
|
|
assert carried == {(DAY, 16, BASE + 4 * CAD)}
|
|
verified = dataset_build.verified_segments(
|
|
{"train": [(BASE, BASE + 10 * CAD)]}, carried,
|
|
satellites=(16,), minimum_slots=3)
|
|
assert [pair for _, pair in verified] == [
|
|
(BASE, BASE + 4 * CAD), (BASE + 5 * CAD, BASE + 10 * CAD)]
|
|
|
|
|
|
def test_roles_survive_verification():
|
|
"""A split segment stays in the split it was chosen for -- a val fragment that
|
|
drifted into train would unblind the model."""
|
|
verified = dataset_build.verified_segments(
|
|
{"train": [(BASE, BASE + 4 * CAD)],
|
|
"val": [(BASE + 100 * CAD, BASE + 104 * CAD)]},
|
|
vetoed=set(), satellites=(16,), minimum_slots=2)
|
|
assert {role for role, _ in verified} == {"train", "val"}
|
|
for role, pair in verified:
|
|
assert (role == "val") == (pair[0] >= BASE + 100 * CAD)
|
|
|
|
|
|
def test_the_corpus_carries_its_availability_profile(tmp_path, monkeypatch):
|
|
"""train.py looks for availability.npz beside the shards and exits without it,
|
|
so a corpus that does not ship it cannot be trained on at all."""
|
|
import dataset_build
|
|
|
|
profile = tmp_path / "availability.npz"
|
|
profile.write_bytes(b"not really an npz, but bytes are bytes")
|
|
destination = dataset_build.Destination(str(tmp_path / "dest"))
|
|
destination.put("availability.npz", profile.read_bytes())
|
|
assert (tmp_path / "dest" / "availability.npz").read_bytes() == profile.read_bytes()
|
|
|
|
|
|
def test_an_unreachable_destination_raises_rather_than_reading_as_empty(monkeypatch):
|
|
"""A failed listing once came back empty, which reads as 'nothing has been
|
|
built yet' -- that silently restarts a finished 25-hour build from zero."""
|
|
import subprocess as sp
|
|
import dataset_build
|
|
|
|
class Failed:
|
|
returncode = 255
|
|
stdout = b""
|
|
stderr = b"ssh: connect to host: No route to host"
|
|
|
|
monkeypatch.setattr(dataset_build.subprocess, "run", lambda *a, **k: Failed())
|
|
monkeypatch.setattr(dataset_build.gpubox, "ssh_argv", lambda host=None: ["ssh", host])
|
|
destination = dataset_build.Destination("nowhere:/tmp/x")
|
|
with pytest.raises(OSError, match="cannot list"):
|
|
destination.existing()
|
|
|
|
|
|
def test_remote_streams_use_the_authenticated_ssh_invocation(monkeypatch):
|
|
"""The host authenticates by password; a plain `ssh -o BatchMode=yes` cannot
|
|
connect at all, and failed as a broken pipe on the first shard."""
|
|
import dataset_build
|
|
|
|
monkeypatch.setattr(dataset_build.gpubox, "_password", lambda: "secret")
|
|
argv = dataset_build.gpubox.ssh_argv("user@host")
|
|
assert argv[0] == "sshpass" and "secret" in argv
|
|
assert argv[-1] == "user@host"
|
|
assert "BatchMode=yes" not in argv
|
|
|
|
|
|
class _Args:
|
|
def __init__(self, retries=2, retry_wait=1, workers=1):
|
|
self.retries, self.retry_wait, self.workers = retries, retry_wait, workers
|
|
|
|
|
|
def test_a_dropped_destination_is_waited_out_and_the_shard_retried(monkeypatch):
|
|
"""The host is on WiFi and dropped twice in three days, each time killing a
|
|
build partway through a shard. Resume was already correct; what was missing
|
|
was anyone to restart it."""
|
|
import dataset_build
|
|
|
|
calls = {"build": 0}
|
|
|
|
def flaky(conn, root, day, satellite, times, destination, workers):
|
|
calls["build"] += 1
|
|
if calls["build"] == 1:
|
|
raise BrokenPipeError(32, "Broken pipe")
|
|
return (123, 4, {}, [])
|
|
|
|
monkeypatch.setattr(dataset_build, "build_shard", flaky)
|
|
monkeypatch.setattr(dataset_build.time, "sleep", lambda s: None)
|
|
|
|
class Back:
|
|
def existing(self):
|
|
return set()
|
|
|
|
result = dataset_build.build_with_retry(
|
|
None, "/root", "2024-01-01", 16, [1], Back(), _Args())
|
|
assert result == (123, 4, {}, [])
|
|
assert calls["build"] == 2 # failed once, retried once
|
|
|
|
|
|
def test_retries_are_bounded_and_the_error_surfaces(monkeypatch):
|
|
"""A destination that never returns must eventually fail loudly, not spin."""
|
|
import dataset_build
|
|
|
|
def always_fails(*a, **k):
|
|
raise BrokenPipeError(32, "Broken pipe")
|
|
|
|
monkeypatch.setattr(dataset_build, "build_shard", always_fails)
|
|
monkeypatch.setattr(dataset_build.time, "sleep", lambda s: None)
|
|
|
|
class Back:
|
|
def existing(self):
|
|
return set()
|
|
|
|
with pytest.raises(BrokenPipeError):
|
|
dataset_build.build_with_retry(
|
|
None, "/root", "2024-01-01", 16, [1], Back(), _Args(retries=2))
|
|
|
|
|
|
def test_an_archive_read_failure_is_not_retried(monkeypatch):
|
|
"""Only the destination is retried. Spinning on a mount that cannot serve a
|
|
file would hide the problem instead of reporting it."""
|
|
import dataset_build
|
|
|
|
def bad_archive(*a, **k):
|
|
raise ValueError("corrupt FITS")
|
|
|
|
monkeypatch.setattr(dataset_build, "build_shard", bad_archive)
|
|
with pytest.raises(ValueError, match="corrupt FITS"):
|
|
dataset_build.build_with_retry(
|
|
None, "/root", "2024-01-01", 16, [1], None, _Args())
|
|
|
|
|
|
def test_a_build_queued_against_an_absent_host_waits_rather_than_crashing(monkeypatch):
|
|
"""The corpus is queued against a host that may be away when queueing happens.
|
|
Crashing at startup means the work waits for a person; waiting means it starts
|
|
itself the moment the host returns."""
|
|
import dataset_build
|
|
|
|
state = {"calls": 0}
|
|
|
|
class LateHost:
|
|
def existing(self):
|
|
state["calls"] += 1
|
|
if state["calls"] < 3:
|
|
raise OSError("No route to host")
|
|
return {"already.shard"}
|
|
|
|
monkeypatch.setattr(dataset_build.time, "sleep", lambda s: None)
|
|
got = dataset_build.wait_for_destination(LateHost(), seconds=600, poll=0)
|
|
assert got == {"already.shard"}
|
|
assert state["calls"] == 3
|
|
|
|
|
|
def test_waiting_for_a_destination_is_bounded(monkeypatch):
|
|
"""A host that never returns must surface the error, not wait forever."""
|
|
import dataset_build
|
|
|
|
class NeverHost:
|
|
def existing(self):
|
|
raise OSError("No route to host")
|
|
|
|
times = iter([0, 1, 10_000, 10_001, 10_002])
|
|
monkeypatch.setattr(dataset_build.time, "sleep", lambda s: None)
|
|
monkeypatch.setattr(dataset_build.time, "time", lambda: next(times))
|
|
with pytest.raises(OSError, match="No route to host"):
|
|
dataset_build.wait_for_destination(NeverHost(), seconds=60, poll=0)
|
|
|
|
|
|
def test_a_failed_shard_reaps_its_worker_pool(monkeypatch, tmp_path):
|
|
"""A link that drops mid-shard makes this run hundreds of times. `with Pool`
|
|
only terminates -- it never joins -- so every abandoned imap kept its buffered
|
|
4.9 MB blocks alive and the builder grew 260 MB -> 16.6 GB over 80 attempts."""
|
|
import dataset_build
|
|
|
|
reaped = {"terminate": 0, "join": 0, "closed": 0}
|
|
|
|
class FakePool:
|
|
def __init__(self, workers):
|
|
pass
|
|
|
|
def imap(self, fn, jobs, chunksize=1):
|
|
for job in jobs:
|
|
yield job[2], None, None
|
|
|
|
def terminate(self):
|
|
reaped["terminate"] += 1
|
|
|
|
def join(self):
|
|
reaped["join"] += 1
|
|
|
|
class Exploding:
|
|
def open(self, name, sealed=False):
|
|
raise BrokenPipeError(32, "Broken pipe")
|
|
|
|
monkeypatch.setattr(dataset_build, "Pool", FakePool)
|
|
monkeypatch.setattr(dataset_build, "slot_paths", lambda conn, sat, times: {})
|
|
|
|
with pytest.raises(BrokenPipeError):
|
|
dataset_build.build_shard(None, "/root", "2024-01-01", 16, [1, 2, 3],
|
|
Exploding(), 1)
|
|
assert reaped["terminate"] == 1, "pool must be terminated on failure"
|
|
assert reaped["join"] == 1, "terminate without join leaks the result buffers"
|
|
|
|
|
|
def _truncated_and_whole_shard(tmp_path):
|
|
"""Bytes of a real v2 shard, and a prefix of it cut short mid-record."""
|
|
import numpy as np
|
|
from suvi import dataset
|
|
block = np.full((6, 640, 640), 1234, dtype=np.uint16)
|
|
path = tmp_path / "whole.shard"
|
|
dataset.write_shard(str(path), "2024-01-01", 16, (94, 131, 171, 195, 284, 304),
|
|
{1700000000: block, 1700000240: block})
|
|
whole = path.read_bytes()
|
|
return whole[: len(whole) // 3], whole
|
|
|
|
|
|
def test_the_remote_commit_refuses_a_stream_cut_short(tmp_path):
|
|
"""The real failure: the local builder died, ssh forwarded a clean EOF, the
|
|
remote `cat` exited 0, and `mv` committed a 33 MB fragment of a 1.4 GB day
|
|
under its final name. Runs the exact remote shell command through bash."""
|
|
import subprocess
|
|
import dataset_build
|
|
|
|
short, whole = _truncated_and_whole_shard(tmp_path)
|
|
dest = dataset_build.Destination(f"host:{tmp_path / 'remote'}")
|
|
|
|
cut = dataset_build._RemoteWrite(dest, "day_g16.shard", sealed=True)
|
|
result = subprocess.run(["bash", "-c", cut.command], input=short,
|
|
stderr=subprocess.PIPE)
|
|
assert result.returncode != 0
|
|
assert b"truncated" in result.stderr
|
|
assert not (tmp_path / "remote" / "day_g16.shard").exists()
|
|
assert (tmp_path / "remote" / "day_g16.shard.partial").exists()
|
|
|
|
ok = dataset_build._RemoteWrite(dest, "day_g16.shard", sealed=True)
|
|
result = subprocess.run(["bash", "-c", ok.command], input=whole)
|
|
assert result.returncode == 0
|
|
assert (tmp_path / "remote" / "day_g16.shard").read_bytes() == whole
|
|
|
|
|
|
def test_unsealed_remote_files_commit_without_a_trailer(tmp_path):
|
|
"""manifest.json and friends do not end in the shard magic and must still land."""
|
|
import subprocess
|
|
import dataset_build
|
|
|
|
dest = dataset_build.Destination(f"host:{tmp_path / 'remote'}")
|
|
put = dataset_build._RemoteWrite(dest, "manifest.json")
|
|
assert subprocess.run(["bash", "-c", put.command], input=b'{"a": 1}').returncode == 0
|
|
assert (tmp_path / "remote" / "manifest.json").read_bytes() == b'{"a": 1}'
|
|
|
|
|
|
def test_a_local_sealed_write_cut_short_is_not_committed(tmp_path):
|
|
import dataset_build
|
|
|
|
short, whole = _truncated_and_whole_shard(tmp_path)
|
|
dest = dataset_build.Destination(str(tmp_path / "local"))
|
|
with pytest.raises(OSError, match="truncated"):
|
|
with dest.open("day_g16.shard", sealed=True) as handle:
|
|
handle.write(short)
|
|
assert not (tmp_path / "local" / "day_g16.shard").exists()
|
|
assert not (tmp_path / "local" / "day_g16.shard.partial").exists()
|
|
|
|
with dest.open("day_g16.shard", sealed=True) as handle:
|
|
handle.write(whole)
|
|
assert (tmp_path / "local" / "day_g16.shard").read_bytes() == whole
|