File size: 4,914 Bytes
ec0a9aa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
#!/usr/bin/env python3
"""
Replace symlinks under datasets/single_arm with real files.

- Regular files (mp4, parquet): group by resolved target; move target to the first
  link path, then hardlink remaining paths (same inode, no symlink).
- meta/ directory (4 symlinks → one source): hardlink the 5 meta files into each
  layout path so there is no symlink and data is not duplicated.

Run from repo root:
  python scripts/materialize_single_arm_symlinks.py [--dry-run]

After a full run without --dry-run, datasets/droid_1.0.1_20chunks will no longer
contain the moved video/parquet files (they live under single_arm). Re-extract
from upstream if you need droid again.
"""
from __future__ import annotations

import argparse
import os
import sys
from collections import defaultdict
from pathlib import Path


def collect_file_symlinks(base: Path) -> list[Path]:
    out: list[Path] = []
    for root, _dirs, files in os.walk(base, followlinks=False):
        for name in files:
            p = Path(root) / name
            if p.is_symlink():
                out.append(p)
    # directory symlinks appear as "dirs" in walk — handle meta separately
    return out


def iter_dir_symlinks_only(base: Path) -> list[Path]:
    out: list[Path] = []
    for root, dirs, _files in os.walk(base, followlinks=False):
        # do not descend into symlinked dirs
        for name in list(dirs):
            p = Path(root) / name
            if p.is_symlink():
                out.append(p)
                dirs.remove(name)  # do not traverse
    return out


def materialize_file_group(target: Path, links: list[Path], dry_run: bool) -> None:
    links = sorted(links, key=lambda x: str(x))
    primary = links[0]
    rest = links[1:]
    if not target.exists():
        raise FileNotFoundError(f"Missing target {target}")
    if not target.is_file():
        raise IsADirectoryError(f"Expected file: {target}")

    if dry_run:
        print(f"[dry-run] move {target} -> {primary}; hardlink -> {len(rest)} path(s)")
        return

    os.unlink(primary)
    os.replace(target, primary)
    for p in rest:
        if p.exists() or p.is_symlink():
            os.unlink(p)
        os.link(primary, p)


def materialize_meta(meta_links: list[Path], source_meta: Path, dry_run: bool) -> None:
    """All meta symlinks must point to the same source_meta directory."""
    meta_links = sorted(meta_links, key=lambda x: str(x))
    first = meta_links[0]
    sources = sorted(source_meta.iterdir(), key=lambda x: x.name)
    for s in sources:
        if not s.is_file():
            raise ValueError(f"unexpected non-file in meta: {s}")

    if dry_run:
        print(f"[dry-run] meta: hardlink {len(sources)} files into {len(meta_links)} dirs")
        return

    if not first.is_symlink():
        raise RuntimeError(f"expected symlink at {first}")
    os.unlink(first)
    first.mkdir(parents=True)
    for s in sources:
        os.link(s, first / s.name)

    for other in meta_links[1:]:
        if not other.is_symlink():
            raise RuntimeError(f"expected symlink at {other}")
        os.unlink(other)
        other.mkdir(parents=True)
        for s in sources:
            os.link(first / s.name, other / s.name)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--root",
        type=Path,
        default=Path(__file__).resolve().parents[1] / "datasets" / "single_arm",
        help="single_arm dataset root",
    )
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()
    base: Path = args.root.resolve()
    if not base.is_dir():
        print(f"Not a directory: {base}", file=sys.stderr)
        return 1

    # --- file symlinks ---
    by_target: dict[Path, list[Path]] = defaultdict(list)
    for p in collect_file_symlinks(base):
        if not p.is_symlink():
            continue
        tgt = p.resolve()
        if tgt.is_file():
            by_target[tgt].append(p)

    # Deterministic order: process targets sorted by string
    for tgt in sorted(by_target.keys(), key=str):
        materialize_file_group(tgt, by_target[tgt], args.dry_run)

    # --- meta directory symlinks (only those under single_arm) ---
    dir_links = [p for p in iter_dir_symlinks_only(base) if p.name == "meta" and p.is_symlink()]
    if not dir_links:
        return 0
    resolved = {p.resolve() for p in dir_links}
    if len(resolved) != 1:
        print(
            "Refusing meta handling: multiple distinct targets: "
            + ", ".join(str(x) for x in sorted(resolved, key=str)),
            file=sys.stderr,
        )
        return 1
    source_meta = next(iter(resolved))
    if not source_meta.is_dir():
        print(f"meta target is not a directory: {source_meta}", file=sys.stderr)
        return 1
    materialize_meta(dir_links, source_meta, args.dry_run)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())