Spaces:
Running
Running
File size: 7,329 Bytes
edae745 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | """内容指纹:识别"换个包名重传同一份数据"。
实测背景:采集端把已上传的录制重新推一遍时包名会变(本地去重标记丢失
或包名规则变化),``session.py`` 的 incoming_name 判断不出来,同一份数据
被追加成新 episode(UMIGripper_AI 出现 4 组字节级相同的批次)。这里的
指纹必须与包名无关,且对"同一次录制的两份拷贝"稳定命中。
"""
import json
import tempfile
import unittest
from pathlib import Path
from unittest import mock
import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
from app import batch_fingerprint
def _write_episode(root: Path, episode_index: int, video_bytes: bytes,
frames: int = 4) -> None:
"""写一个最小 canonical 布局:一个视频 + 一个 data parquet。"""
videos = root / "videos" / "observation.images.gripper_rgb" / "chunk-000"
videos.mkdir(parents=True, exist_ok=True)
(videos / f"episode_{episode_index:06d}.mp4").write_bytes(video_bytes)
data = root / "data" / "chunk-000"
data.mkdir(parents=True, exist_ok=True)
table = pa.table({
"episode_index": pa.array([episode_index] * frames, pa.int64()),
"frame_index": pa.array(list(range(frames)), pa.int64()),
"observation.value": pa.array(
[[float(i), float(i) + 1.0] for i in range(frames)],
pa.list_(pa.float32()),
),
})
pq.write_table(table, data / f"episode_{episode_index:06d}.parquet")
class FingerprintTests(unittest.TestCase):
def setUp(self):
self._tmp = tempfile.TemporaryDirectory()
self.root = Path(self._tmp.name)
# 索引落到临时目录,别碰真实 state/
self._index_patch = mock.patch.object(
batch_fingerprint, "INDEX_PATH", self.root / "fingerprints.json")
self._index_patch.start()
def tearDown(self):
self._index_patch.stop()
self._tmp.cleanup()
def test_same_recording_different_episode_index_matches(self):
"""同一次录制的两份拷贝,只有 episode_index 不同 → 指纹相同。"""
a = self.root / "a"
b = self.root / "b"
_write_episode(a, 3, b"IDENTICAL-VIDEO-BYTES")
_write_episode(b, 4, b"IDENTICAL-VIDEO-BYTES")
self.assertIsNotNone(batch_fingerprint.compute_fingerprint(a))
self.assertEqual(
batch_fingerprint.compute_fingerprint(a),
batch_fingerprint.compute_fingerprint(b),
)
def test_different_video_differs(self):
a = self.root / "a"
b = self.root / "b"
_write_episode(a, 0, b"VIDEO-ONE")
_write_episode(b, 0, b"VIDEO-TWO")
self.assertNotEqual(
batch_fingerprint.compute_fingerprint(a),
batch_fingerprint.compute_fingerprint(b),
)
def test_different_frame_data_differs(self):
"""视频相同但帧数据不同 → 不同指纹(不能只比视频)。"""
a = self.root / "a"
b = self.root / "b"
_write_episode(a, 0, b"SAME-VIDEO")
_write_episode(b, 0, b"SAME-VIDEO")
# 只改 b 的一帧数据
path = b / "data" / "chunk-000" / "episode_000000.parquet"
table = pq.read_table(path).to_pydict()
table["observation.value"] = [[9.0, 9.0] for _ in table["frame_index"]]
pq.write_table(pa.table(table), path)
self.assertNotEqual(
batch_fingerprint.compute_fingerprint(a),
batch_fingerprint.compute_fingerprint(b),
)
def test_missing_data_returns_none(self):
"""没有 parquet 就算不出指纹 → None(调用方按新批次放行)。"""
empty = self.root / "empty"
empty.mkdir()
self.assertIsNone(batch_fingerprint.compute_fingerprint(empty))
def test_record_and_lookup_roundtrip(self):
fingerprint = "v:test=abc;d:def"
batch_fingerprint.record_fingerprint("UMIGripper_AI", fingerprint, "ep_000003")
self.assertEqual(
batch_fingerprint.lookup_duplicate(
"UMIGripper_AI", fingerprint, {"ep_000003"}),
"ep_000003",
)
def test_lookup_ignores_deleted_episode(self):
"""索引指向的 episode 已删除 → 不判重(允许重新上传)。"""
fingerprint = "v:test=abc;d:def"
batch_fingerprint.record_fingerprint("UMIGripper_AI", fingerprint, "ep_000003")
self.assertIsNone(
batch_fingerprint.lookup_duplicate(
"UMIGripper_AI", fingerprint, {"ep_000004"}),
)
def test_lookup_none_fingerprint_is_safe(self):
self.assertIsNone(
batch_fingerprint.lookup_duplicate("p", None, {"a"}))
def test_record_drops_stale_entry_for_same_episode(self):
"""同名重传换了内容:旧指纹不该继续指向这个 episode。"""
batch_fingerprint.record_fingerprint("p", "old-content", "ep_1")
batch_fingerprint.record_fingerprint("p", "new-content", "ep_1")
self.assertIsNone(
batch_fingerprint.lookup_duplicate("p", "old-content", {"ep_1"}))
self.assertEqual(
batch_fingerprint.lookup_duplicate("p", "new-content", {"ep_1"}),
"ep_1",
)
def test_forget_fingerprint(self):
batch_fingerprint.record_fingerprint("p", "abc", "ep_1")
batch_fingerprint.forget_fingerprint("ep_1")
self.assertIsNone(
batch_fingerprint.lookup_duplicate("p", "abc", {"ep_1"}))
def test_corrupt_index_is_tolerated(self):
batch_fingerprint.INDEX_PATH.write_text("{not json", encoding="utf-8")
self.assertIsNone(
batch_fingerprint.lookup_duplicate("p", "abc", {"ep_1"}))
batch_fingerprint.record_fingerprint("p", "abc", "ep_1")
self.assertEqual(
batch_fingerprint.lookup_duplicate("p", "abc", {"ep_1"}), "ep_1")
def test_index_is_valid_json_with_episode_map(self):
batch_fingerprint.record_fingerprint("proj", "fp1", "ep_1")
data = json.loads(batch_fingerprint.INDEX_PATH.read_text(encoding="utf-8"))
self.assertEqual(data["proj"]["fp1"], "ep_1")
class ParquetHashTests(unittest.TestCase):
def test_episode_index_column_is_excluded(self):
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
table = pa.table({
"episode_index": pa.array([7], pa.int64()),
"frame_index": pa.array([0], pa.int64()),
"v": pa.array([[1.0, 2.0]], pa.list_(pa.float32())),
})
one = root / "one.parquet"
pq.write_table(table, one)
other = root / "other.parquet"
pq.write_table(
table.set_column(0, "episode_index",
pa.array([99], pa.int64())), other)
self.assertEqual(
batch_fingerprint._parquet_content_hash(one),
batch_fingerprint._parquet_content_hash(other),
)
def test_unreadable_parquet_returns_none(self):
with tempfile.TemporaryDirectory() as tmp:
bad = Path(tmp) / "bad.parquet"
bad.write_bytes(b"not a parquet file")
self.assertIsNone(batch_fingerprint._parquet_content_hash(bad))
if __name__ == "__main__":
unittest.main()
|