File size: 6,375 Bytes
5d08972 | 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 | #!/usr/bin/env python3
"""Pick the first 256 episodes in true generation order and lay out the work for both
downstream stages.
Every worker writes each success to tmp/date_..._time_HH_MM_SS.hdf5 before merge_all_hdf5
sorts those files by timestamp into demo_0..demo_N and deletes the folder. snapshot_times.sh
captured the filenames while they existed, so joining them back gives each merged demo its
wall-clock time -- and since all 32 workers share one machine clock, a global ordering.
Writes:
selection.json the ordered 256 with their timestamps
render_plan.json for render_from_plan.py (preview mp4s, original asset)
gen/shard_K/... per-shard gen trees for render_to_gr00t.py (dataset, swapped mesh)
"""
import json, os, re, glob, sys, collections, datetime
import h5py
S = "/tmp/claude-1001/-home-nvidia/68c27ffb-706b-4cb3-8bf2-4618a200d378/scratchpad"
NAT = f"{S}/natural256"
SRC = f"{S}/big/src150_fix.hdf5"
N_KEEP = int(os.environ.get("N_KEEP", 256))
N_CONV_SHARDS = int(os.environ.get("N_CONV_SHARDS", 24))
CATEGORY = {
"donut_5": "donut", "Jar023": "jar", "MeasuringCup009": "measuring_cup",
"SoapDispenser010": "soap_dispenser", "steak_8": "steak",
"SyrupBottle005": "syrup_bottle", "teapot_6": "teapot", "teapot_7": "teapot",
}
# ---- source demo index -> (demo key, the object that source episode had) ----
src = {}
with h5py.File(SRC, "r") as f:
keys = [k.decode() for k in f["mask/valid"][:]]
for i, k in enumerate(keys):
meta = json.loads(f[f"data/{k}"].attrs["ep_meta"])
cat = next((c["info"]["cat"] for c in meta.get("object_cfgs", [])
if c.get("name") == "obj"), "unknown")
src[i] = (k, cat)
# ---- timestamps captured before the tmp folders were deleted ----
# One line per poll, so the same file appears many times and its mtime drifts while it is
# still being written -- dedupe on the name and order by the timestamp the name itself
# carries, which is fixed at creation.
seen = collections.defaultdict(dict)
with open(f"{NAT}/episode_times.tsv") as f:
for line in f:
parts = line.rstrip("\n").split("\t")
if len(parts) != 3:
continue
tmpdir, fname, _ = parts
m = re.match(r"date_(\d+)_(\d+)_(\d+)_time_(\d+)_(\d+)_(\d+)\.hdf5", fname)
if not m:
continue
worker = os.path.basename(os.path.dirname(os.path.dirname(tmpdir)))
mo, dy, yr, hh, mm, ss = (int(x) for x in m.groups())
seen[worker][fname] = datetime.datetime(
yr, mo, dy, hh, mm, ss, tzinfo=datetime.timezone.utc).timestamp()
times = {w: sorted((v, k) for k, v in d.items()) for w, d in seen.items()}
rows, missing = [], []
for wdir in sorted(glob.glob(f"{NAT}/nat_*")) + sorted(glob.glob(f"{NAT}/r2_*")):
worker = os.path.basename(wdir)
model = re.sub(r"^(nat|r2)_|_w\d+$", "", worker)
hp = f"{wdir}/{worker}/demo.hdf5"
if not os.path.isdir(wdir) or not os.path.exists(hp):
continue
with h5py.File(hp, "r") as f:
demos = sorted(f["data"].keys(), key=lambda x: int(x.split("_")[1]))
stamps = list(times.get(worker, []))
# The poller can miss the very last success: the worker merges and deletes tmp/
# immediately after writing it. Only the tail can ever be missing (earlier files sat
# there for the whole run), so pad it with the merged file's own mtime -- the moment
# the worker finished, which bounds that episode from above.
if len(stamps) < len(demos):
end = os.path.getmtime(hp)
pad = len(demos) - len(stamps)
missing.append(f"{worker}: padded {pad} with run end")
stamps += [(end, "<run end>")] * pad
for i, dk in enumerate(demos):
g = f[f"data/{dk}"]
si = int(g["src_demo_inds"][0]) if "src_demo_inds" in g else -1
sk, scat = src.get(si, (f"idx{si}", "unknown"))
t = stamps[i][0] if i < len(stamps) else float("inf")
rows.append(dict(t=t, worker=worker, run_dir=f"{wdir}/{worker}", demo=dk,
target=model, source_demo=sk, source_object=scat))
rows.sort(key=lambda r: (r["t"], r["worker"], r["demo"]))
print(f"총 성공 {len(rows)}개" + (f" (경고: {missing})" if missing else ""))
for i, r in enumerate(rows):
r["order"] = i
kept = rows[:N_KEEP]
taken = collections.Counter()
for r in kept:
base = f"{r['source_demo'].replace('demo_', 'ep')}_{r['source_object']}_{r['target']}"
taken[base] += 1
r["name"] = base if taken[base] == 1 else f"{base}_{taken[base]}"
json.dump(kept, open(f"{NAT}/selection.json", "w"), indent=1)
json.dump([{k: r[k] for k in ("run_dir", "demo", "target", "name")} for r in kept],
open(f"{NAT}/render_plan.json", "w"), indent=1)
# ---- conversion shards: contiguous chunks, each its own GEN_DIR tree ----
by_target = collections.defaultdict(list)
for r in kept:
by_target[r["target"]].append(r)
chunks = []
per = max(1, -(-len(kept) // N_CONV_SHARDS))
for model, rs in by_target.items():
for i in range(0, len(rs), per):
chunks.append((model, rs[i:i + per]))
for k, (model, rs) in enumerate(chunks):
name = f"{CATEGORY[model]}_{model}"
d = f"{NAT}/gen/shard_{k:02d}/{name}/{name}_gen"
os.makedirs(d, exist_ok=True)
out = f"{d}/demo.hdf5"
if os.path.exists(out):
os.remove(out)
with h5py.File(out, "w") as fo:
grp = fo.create_group("data")
total, env_args = 0, None
by_run = collections.defaultdict(list)
for r in rs:
by_run[r["run_dir"]].append(r)
j = 0
for run_dir, items in by_run.items():
with h5py.File(f"{run_dir}/demo.hdf5", "r") as fi:
if env_args is None:
env_args = fi["data"].attrs["env_args"]
for r in items:
fi.copy(f"data/{r['demo']}", grp, name=f"demo_{j}")
total += int(fi[f"data/{r['demo']}"].attrs["num_samples"])
j += 1
grp.attrs["total"] = total
grp.attrs["env_args"] = env_args
print(f"선택 {len(kept)}개 -> render_plan.json, 변환 샤드 {len(chunks)}개 "
f"(샤드당 최대 {per})")
c = collections.Counter(r["target"] for r in kept)
print("\n물체별 분포:")
for k, v in c.most_common():
print(f" {k:18s} {v}")
|