ls-algorithm-plugin-sdk 0.3.3__tar.gz → 0.3.5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (51) hide show
  1. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/PKG-INFO +12 -1
  2. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/README.md +10 -0
  3. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/pyproject.toml +4 -2
  4. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/cli.py +4 -0
  5. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/cli_impl/configure.py +12 -0
  6. ls_algorithm_plugin_sdk-0.3.5/src/algorithm_plugin_sdk/cli_impl/cut.py +551 -0
  7. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/cli_impl/serve.py +6 -0
  8. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/deployment.py +7 -0
  9. ls_algorithm_plugin_sdk-0.3.5/src/algorithm_plugin_sdk/log_manager.py +125 -0
  10. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/service.py +80 -8
  11. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/ls_algorithm_plugin_sdk.egg-info/PKG-INFO +12 -1
  12. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/ls_algorithm_plugin_sdk.egg-info/SOURCES.txt +3 -0
  13. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/ls_algorithm_plugin_sdk.egg-info/requires.txt +1 -0
  14. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/tests/test_configure.py +8 -0
  15. ls_algorithm_plugin_sdk-0.3.5/tests/test_cut.py +104 -0
  16. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/tests/test_deployment.py +4 -0
  17. ls_algorithm_plugin_sdk-0.3.5/tests/test_serve.py +73 -0
  18. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/tests/test_service.py +172 -0
  19. ls_algorithm_plugin_sdk-0.3.3/src/algorithm_plugin_sdk/log_manager.py +0 -76
  20. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/setup.cfg +0 -0
  21. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/__init__.py +0 -0
  22. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/algorithm.py +0 -0
  23. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/cli_impl/__init__.py +0 -0
  24. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/cli_impl/parsing.py +0 -0
  25. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/cli_impl/run.py +0 -0
  26. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/context.py +0 -0
  27. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/errors.py +0 -0
  28. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/examples/__init__.py +0 -0
  29. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/examples/example_algorithm.py +0 -0
  30. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/examples/simulated_algorithm.py +0 -0
  31. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/gpu_isolation.py +0 -0
  32. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/loader.py +0 -0
  33. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/models.py +0 -0
  34. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/registration.py +0 -0
  35. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/release.py +0 -0
  36. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/runner.py +0 -0
  37. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/webui/__init__.py +0 -0
  38. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/webui/app.css +0 -0
  39. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/webui/app.js +0 -0
  40. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/webui/index.html +0 -0
  41. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/algorithm_plugin_sdk/webui_app.py +0 -0
  42. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/ls_algorithm_plugin_sdk.egg-info/dependency_links.txt +0 -0
  43. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/ls_algorithm_plugin_sdk.egg-info/entry_points.txt +0 -0
  44. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/src/ls_algorithm_plugin_sdk.egg-info/top_level.txt +0 -0
  45. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/tests/test_cli.py +0 -0
  46. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/tests/test_gpu_isolation.py +0 -0
  47. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/tests/test_models.py +0 -0
  48. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/tests/test_registration.py +0 -0
  49. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/tests/test_release.py +0 -0
  50. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/tests/test_runner.py +0 -0
  51. {ls_algorithm_plugin_sdk-0.3.3 → ls_algorithm_plugin_sdk-0.3.5}/tests/test_sdk_automation.py +0 -0
@@ -1,11 +1,12 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ls-algorithm-plugin-sdk
3
- Version: 0.3.3
3
+ Version: 0.3.5
4
4
  Summary: Protocol-independent runtime SDK for dataset algorithms
5
5
  Author: Ling Robotics
6
6
  License: Proprietary
7
7
  Requires-Python: >=3.10
8
8
  Description-Content-Type: text/markdown
9
+ Requires-Dist: pyarrow>=14
9
10
  Provides-Extra: service
10
11
  Requires-Dist: fastapi<1,>=0.110; extra == "service"
11
12
  Requires-Dist: pydantic<3,>=2.0; extra == "service"
@@ -40,6 +41,16 @@ algorithm-plugin run camera-space-mano \
40
41
 
41
42
  `--gpu-ids` 和 `--parameters` 可省略,也可用 `--request-file` 读取完整请求 JSON。
42
43
 
44
+ ## 数据集裁剪
45
+
46
+ 将每个 episode 裁剪为最多指定分钟数,并同步更新 Parquet、视频、传感器 CSV 和元数据:
47
+
48
+ ```bash
49
+ algorithm-plugin cut /data/input_dataset /data/output_dataset --minutes 1.5
50
+ ```
51
+
52
+ 支持 LeRobot v2/v3 的 `episode_*`、`file-*` 和多数据源目录布局。视频裁剪需要系统已安装 `ffmpeg` 与 `ffprobe`。
53
+
43
54
  ## 配置现有仓库
44
55
 
45
56
  目标算法仓库必须满足:
@@ -25,6 +25,16 @@ algorithm-plugin run camera-space-mano \
25
25
 
26
26
  `--gpu-ids` 和 `--parameters` 可省略,也可用 `--request-file` 读取完整请求 JSON。
27
27
 
28
+ ## 数据集裁剪
29
+
30
+ 将每个 episode 裁剪为最多指定分钟数,并同步更新 Parquet、视频、传感器 CSV 和元数据:
31
+
32
+ ```bash
33
+ algorithm-plugin cut /data/input_dataset /data/output_dataset --minutes 1.5
34
+ ```
35
+
36
+ 支持 LeRobot v2/v3 的 `episode_*`、`file-*` 和多数据源目录布局。视频裁剪需要系统已安装 `ffmpeg` 与 `ffprobe`。
37
+
28
38
  ## 配置现有仓库
29
39
 
30
40
  目标算法仓库必须满足:
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "ls-algorithm-plugin-sdk"
7
- version = "0.3.3"
7
+ version = "0.3.5"
8
8
  description = "Protocol-independent runtime SDK for dataset algorithms"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.10"
@@ -12,7 +12,9 @@ license = {text = "Proprietary"}
12
12
  authors = [
13
13
  {name = "Ling Robotics"},
14
14
  ]
15
- dependencies = []
15
+ dependencies = [
16
+ "pyarrow>=14",
17
+ ]
16
18
 
17
19
  [project.optional-dependencies]
18
20
  service = [
@@ -5,6 +5,7 @@ import logging
5
5
  import sys
6
6
 
7
7
  from .cli_impl import configure as configure_command
8
+ from .cli_impl import cut as cut_command
8
9
  from .cli_impl import run as run_command
9
10
  from .cli_impl import serve as serve_command
10
11
 
@@ -22,6 +23,9 @@ def build_parser() -> argparse.ArgumentParser:
22
23
  configure_command.configure_parser(
23
24
  subs.add_parser("configure", help="configure an existing Algorithm repository for systemd")
24
25
  )
26
+ cut_command.configure_parser(
27
+ subs.add_parser("cut", help="cut each dataset episode to a maximum duration")
28
+ )
25
29
  return parser
26
30
 
27
31
 
@@ -62,6 +62,16 @@ def configure_parser(parser: argparse.ArgumentParser) -> None:
62
62
  help="maximum concurrent executions (default: GPU count, or 1)",
63
63
  )
64
64
  parser.add_argument("--scratch-dir")
65
+ parser.add_argument(
66
+ "--local-scratch",
67
+ action="store_true",
68
+ help="override request scratchRoot with <repository>/.scratch",
69
+ )
70
+ parser.add_argument(
71
+ "--local-output",
72
+ action="store_true",
73
+ help="override request output paths with <repository>/.output/<job-id>",
74
+ )
65
75
  parser.add_argument(
66
76
  "--gpu-ids",
67
77
  type=gpu_ids,
@@ -409,6 +419,8 @@ def generate_config(
409
419
  webui=args.webui,
410
420
  max_concurrency=max_concurrency,
411
421
  scratch_dir=args.scratch_dir,
422
+ local_scratch=args.local_scratch,
423
+ local_output=args.local_output,
412
424
  token=args.service_token,
413
425
  registration=registration_from_args(args, repository),
414
426
  gpu_ids=args.gpu_ids,
@@ -0,0 +1,551 @@
1
+ from __future__ import annotations
2
+
3
+ import argparse
4
+ import csv
5
+ import json
6
+ import re
7
+ import shutil
8
+ import subprocess
9
+ from collections import defaultdict
10
+ from pathlib import Path
11
+ from typing import Any
12
+
13
+ import pyarrow as pa
14
+ import pyarrow.parquet as pq
15
+
16
+
17
+ EPISODE_FILE_PATTERN = re.compile(r"episode[-_](\d+)")
18
+ FILE_PATTERN = re.compile(r"file-(\d+)")
19
+ VIDEO_SUFFIXES = {".avi", ".mkv", ".mov", ".mp4", ".webm"}
20
+
21
+
22
+ def positive_minutes(value: str) -> float:
23
+ try:
24
+ minutes = float(value)
25
+ except ValueError as exc:
26
+ raise argparse.ArgumentTypeError("minutes must be a number") from exc
27
+ if minutes <= 0:
28
+ raise argparse.ArgumentTypeError("minutes must be greater than zero")
29
+ return minutes
30
+
31
+
32
+ def configure_parser(parser: argparse.ArgumentParser) -> None:
33
+ parser.add_argument("input_dataset", type=Path)
34
+ parser.add_argument("output_path", type=Path)
35
+ parser.add_argument(
36
+ "--minutes",
37
+ type=positive_minutes,
38
+ required=True,
39
+ help="maximum duration retained from the start of each episode",
40
+ )
41
+ parser.set_defaults(handler=execute)
42
+
43
+
44
+ def _source_name(path: Path, data_root: Path) -> str:
45
+ relative = path.relative_to(data_root)
46
+ parts = []
47
+ for part in relative.parts[:-1]:
48
+ if part.startswith("chunk-"):
49
+ break
50
+ parts.append(part)
51
+ return "/".join(parts)
52
+
53
+
54
+ def _replace_column(table: pa.Table, name: str, values: list[int]) -> pa.Table:
55
+ index = table.column_names.index(name)
56
+ return table.set_column(
57
+ index,
58
+ name,
59
+ pa.array(values, type=table.schema.field(name).type),
60
+ )
61
+
62
+
63
+ def _read_json_lines(path: Path) -> list[dict[str, Any]]:
64
+ if not path.is_file():
65
+ return []
66
+ return [
67
+ json.loads(line)
68
+ for line in path.read_text(encoding="utf-8").splitlines()
69
+ if line.strip()
70
+ ]
71
+
72
+
73
+ def _write_json_lines(path: Path, rows: list[dict[str, Any]]) -> None:
74
+ path.write_text(
75
+ "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
76
+ encoding="utf-8",
77
+ )
78
+
79
+
80
+ def _episode_tables(root: Path) -> list[tuple[Path, pa.Table]]:
81
+ result = []
82
+ for path in sorted((root / "meta" / "episodes").glob("**/*.parquet")):
83
+ result.append((path, pq.read_table(path)))
84
+ return result
85
+
86
+
87
+ def _episode_rows(tables: list[tuple[Path, pa.Table]]) -> list[dict[str, Any]]:
88
+ return [row for _, table in tables for row in table.to_pylist()]
89
+
90
+
91
+ def _table_time_values(table: pa.Table, fps: float) -> list[float]:
92
+ if "timestamp" in table.column_names:
93
+ return [float(value) for value in table["timestamp"].to_pylist()]
94
+ if "frame_index" in table.column_names:
95
+ return [float(value) / fps for value in table["frame_index"].to_pylist()]
96
+ return [index / fps for index in range(table.num_rows)]
97
+
98
+
99
+ def _candidate_lengths(
100
+ paths: list[Path], data_root: Path, seconds: float, fps: float
101
+ ) -> dict[int, int]:
102
+ counts: dict[str, dict[int, int]] = defaultdict(lambda: defaultdict(int))
103
+ first_timestamps: dict[tuple[str, int], float] = {}
104
+ for path in paths:
105
+ schema = pq.read_schema(path)
106
+ if "episode_index" not in schema.names:
107
+ raise ValueError(f"data parquet lacks episode_index: {path}")
108
+ columns = ["episode_index"]
109
+ for candidate in ("timestamp", "frame_index"):
110
+ if candidate in schema.names:
111
+ columns.append(candidate)
112
+ break
113
+ table = pq.read_table(path, columns=columns)
114
+ source = _source_name(path, data_root)
115
+ episodes = [int(value) for value in table["episode_index"].to_pylist()]
116
+ timestamps = _table_time_values(table, fps)
117
+ for episode, timestamp in zip(episodes, timestamps):
118
+ key = (source, episode)
119
+ first = first_timestamps.setdefault(key, timestamp)
120
+ if timestamp <= first + seconds + 1e-9:
121
+ counts[source][episode] += 1
122
+
123
+ episodes = {episode for values in counts.values() for episode in values}
124
+ if not episodes:
125
+ raise ValueError("no episode rows found in data parquet files")
126
+ lengths = {
127
+ episode: min(
128
+ values[episode] for values in counts.values() if episode in values
129
+ )
130
+ for episode in episodes
131
+ }
132
+ return lengths
133
+
134
+
135
+ def _probe_frames(path: Path) -> int:
136
+ command = [
137
+ "ffprobe",
138
+ "-v",
139
+ "error",
140
+ "-select_streams",
141
+ "v:0",
142
+ "-show_entries",
143
+ "stream=nb_frames",
144
+ "-of",
145
+ "json",
146
+ str(path),
147
+ ]
148
+ try:
149
+ result = subprocess.run(
150
+ command, check=True, capture_output=True, text=True
151
+ )
152
+ stream = json.loads(result.stdout)["streams"][0]
153
+ value = stream.get("nb_frames")
154
+ if value not in (None, "N/A"):
155
+ return int(value)
156
+
157
+ command.insert(command.index("-show_entries"), "-count_frames")
158
+ command[command.index("stream=nb_frames")] = "stream=nb_read_frames"
159
+ result = subprocess.run(
160
+ command, check=True, capture_output=True, text=True
161
+ )
162
+ value = json.loads(result.stdout)["streams"][0].get("nb_read_frames")
163
+ if value not in (None, "N/A"):
164
+ return int(value)
165
+ except (subprocess.CalledProcessError, KeyError, ValueError, json.JSONDecodeError) as exc:
166
+ raise RuntimeError(f"cannot inspect video frames: {path}") from exc
167
+ raise RuntimeError(f"video does not report a frame count: {path}")
168
+
169
+
170
+ def _chunk_index(path: Path) -> int | None:
171
+ for part in path.parts:
172
+ if part.startswith("chunk-"):
173
+ try:
174
+ return int(part.removeprefix("chunk-"))
175
+ except ValueError:
176
+ return None
177
+ return None
178
+
179
+
180
+ def _media_episodes(path: Path, root: Path, rows: list[dict[str, Any]]) -> list[int]:
181
+ direct = EPISODE_FILE_PATTERN.search(path.stem)
182
+ if direct:
183
+ return [int(direct.group(1))]
184
+
185
+ file_match = FILE_PATTERN.search(path.stem)
186
+ if not file_match:
187
+ raise ValueError(f"cannot determine episode for media file: {path}")
188
+ file_index = int(file_match.group(1))
189
+ chunk_index = _chunk_index(path)
190
+ relative = path.relative_to(root).as_posix()
191
+ kind = relative.split("/", 1)[0]
192
+ matches = []
193
+ for row in rows:
194
+ episode = int(row.get("episode_index", -1))
195
+ for key, value in row.items():
196
+ if not key.startswith(f"{kind}/") or not key.endswith("/file_index"):
197
+ continue
198
+ feature = key[len(kind) + 1 : -len("/file_index")]
199
+ chunk_key = key[: -len("file_index")] + "chunk_index"
200
+ if (
201
+ int(value) == file_index
202
+ and (chunk_index is None or int(row.get(chunk_key, -1)) == chunk_index)
203
+ and feature in relative
204
+ ):
205
+ matches.append(episode)
206
+ break
207
+ return sorted(set(matches)) or [file_index]
208
+
209
+
210
+ def _copy_static_tree(source: Path, output: Path) -> None:
211
+ skipped = {"data", "videos", "depths", "sensors"}
212
+
213
+ def ignore(directory: str, names: list[str]) -> set[str]:
214
+ return skipped.intersection(names) if Path(directory) == source else set()
215
+
216
+ shutil.copytree(source, output, ignore=ignore)
217
+
218
+
219
+ def _write_data(
220
+ paths: list[Path],
221
+ source_root: Path,
222
+ output_root: Path,
223
+ limits: dict[int, int],
224
+ fps: float,
225
+ ) -> tuple[dict[str, dict[int, dict[str, float | int]]], dict[str, int]]:
226
+ seen: dict[str, dict[int, int]] = defaultdict(lambda: defaultdict(int))
227
+ global_index: dict[str, int] = defaultdict(int)
228
+ metrics: dict[str, dict[int, dict[str, float | int]]] = defaultdict(dict)
229
+
230
+ for path in paths:
231
+ table = pq.read_table(path)
232
+ source = _source_name(path, source_root)
233
+ episodes = [int(value) for value in table["episode_index"].to_pylist()]
234
+ timestamps = _table_time_values(table, fps)
235
+ keep = []
236
+ frame_indexes = []
237
+ indexes = []
238
+ for row_index, (episode, timestamp) in enumerate(zip(episodes, timestamps)):
239
+ ordinal = seen[source][episode]
240
+ seen[source][episode] += 1
241
+ if ordinal >= limits[episode]:
242
+ continue
243
+ keep.append(row_index)
244
+ frame_indexes.append(ordinal)
245
+ indexes.append(global_index[source])
246
+ global_index[source] += 1
247
+ episode_metrics = metrics[source].setdefault(
248
+ episode,
249
+ {
250
+ "from_index": indexes[-1],
251
+ "to_index": indexes[-1] + 1,
252
+ "from_timestamp": timestamp,
253
+ "to_timestamp": timestamp,
254
+ "length": 0,
255
+ },
256
+ )
257
+ episode_metrics["to_index"] = indexes[-1] + 1
258
+ episode_metrics["to_timestamp"] = timestamp
259
+ episode_metrics["length"] = int(episode_metrics["length"]) + 1
260
+
261
+ trimmed = table.take(pa.array(keep, type=pa.int64()))
262
+ if "frame_index" in trimmed.column_names:
263
+ trimmed = _replace_column(trimmed, "frame_index", frame_indexes)
264
+ if "index" in trimmed.column_names:
265
+ trimmed = _replace_column(trimmed, "index", indexes)
266
+ destination = output_root / path.relative_to(source_root.parent)
267
+ destination.parent.mkdir(parents=True, exist_ok=True)
268
+ pq.write_table(trimmed, destination)
269
+ return metrics, dict(global_index)
270
+
271
+
272
+ def _copy_or_cut_video(source: Path, destination: Path, frames: int) -> None:
273
+ source_frames = _probe_frames(source)
274
+ destination.parent.mkdir(parents=True, exist_ok=True)
275
+ if source_frames <= frames:
276
+ shutil.copy2(source, destination)
277
+ return
278
+ depth_video = "depth" in source.as_posix().lower()
279
+ codec = ["-c:v", "libx264rgb", "-crf", "0"] if depth_video else [
280
+ "-c:v",
281
+ "libx264",
282
+ "-crf",
283
+ "20",
284
+ ]
285
+ subprocess.run(
286
+ [
287
+ "ffmpeg",
288
+ "-y",
289
+ "-loglevel",
290
+ "error",
291
+ "-i",
292
+ str(source),
293
+ "-map",
294
+ "0:v:0",
295
+ "-frames:v",
296
+ str(frames),
297
+ *codec,
298
+ "-preset",
299
+ "veryfast",
300
+ "-an",
301
+ str(destination),
302
+ ],
303
+ check=True,
304
+ )
305
+ actual = _probe_frames(destination)
306
+ if actual != frames:
307
+ raise RuntimeError(
308
+ f"video frame count mismatch for {destination}: expected {frames}, got {actual}"
309
+ )
310
+
311
+
312
+ def _sensor_episode(path: Path, rows: list[dict[str, Any]], root: Path) -> int:
313
+ match = EPISODE_FILE_PATTERN.search(path.stem)
314
+ if match:
315
+ return int(match.group(1))
316
+ relative = path.relative_to(root).as_posix()
317
+ matches = {
318
+ int(row["episode_index"])
319
+ for row in rows
320
+ for key, value in row.items()
321
+ if key.startswith("sensors/") and key.endswith("/path") and value == relative
322
+ }
323
+ if len(matches) != 1:
324
+ raise ValueError(f"cannot determine episode for sensor file: {path}")
325
+ return matches.pop()
326
+
327
+
328
+ def _copy_sensors(
329
+ source: Path,
330
+ output: Path,
331
+ episode_rows: list[dict[str, Any]],
332
+ end_times: dict[int, float],
333
+ ) -> dict[str, dict[str, float | int]]:
334
+ metrics = {}
335
+ sensor_root = source / "sensors"
336
+ if not sensor_root.exists():
337
+ return metrics
338
+ for path in sorted(sensor_root.glob("**/*")):
339
+ if not path.is_file():
340
+ continue
341
+ destination = output / path.relative_to(source)
342
+ destination.parent.mkdir(parents=True, exist_ok=True)
343
+ if path.suffix.lower() != ".csv":
344
+ shutil.copy2(path, destination)
345
+ continue
346
+ episode = _sensor_episode(path, episode_rows, source)
347
+ with path.open(newline="", encoding="utf-8") as stream:
348
+ reader = csv.DictReader(stream)
349
+ rows = list(reader)
350
+ fieldnames = reader.fieldnames or []
351
+ if "timestamp" in fieldnames and rows:
352
+ first_timestamp = float(rows[0]["timestamp"])
353
+ maximum = first_timestamp + end_times[episode] + 1e-9
354
+ rows = [row for row in rows if float(row["timestamp"]) <= maximum]
355
+ if "imu_index" in fieldnames:
356
+ for index, row in enumerate(rows):
357
+ row["imu_index"] = str(index)
358
+ with destination.open("w", newline="", encoding="utf-8") as stream:
359
+ writer = csv.DictWriter(stream, fieldnames=fieldnames)
360
+ writer.writeheader()
361
+ writer.writerows(rows)
362
+ relative = path.relative_to(source).as_posix()
363
+ metrics[relative] = {
364
+ "from_index": 0,
365
+ "to_index": len(rows),
366
+ "from_timestamp": float(rows[0]["timestamp"]) if rows and "timestamp" in fieldnames else 0.0,
367
+ "to_timestamp": float(rows[-1]["timestamp"]) if rows and "timestamp" in fieldnames else 0.0,
368
+ }
369
+ return metrics
370
+
371
+
372
+ def _update_episode_row(
373
+ row: dict[str, Any],
374
+ limits: dict[int, int],
375
+ offsets: dict[int, int],
376
+ source_metrics: dict[str, dict[int, dict[str, float | int]]],
377
+ sensor_metrics: dict[str, dict[str, float | int]],
378
+ end_times: dict[int, float],
379
+ ) -> None:
380
+ episode = int(row["episode_index"])
381
+ length = limits[episode]
382
+ for key in list(row):
383
+ if key == "length" or key.endswith("_length"):
384
+ row[key] = length
385
+ for key in ("duration", "duration_s"):
386
+ if key in row:
387
+ row[key] = min(float(row[key]), end_times[episode])
388
+ if "dataset_from_index" in row:
389
+ row["dataset_from_index"] = offsets[episode]
390
+ if "dataset_to_index" in row:
391
+ row["dataset_to_index"] = offsets[episode] + length
392
+
393
+ for source, episodes in source_metrics.items():
394
+ values = episodes.get(episode)
395
+ if values is None:
396
+ continue
397
+ prefix = f"data/{source}/" if source else "data/"
398
+ for field in ("from_index", "to_index", "from_timestamp", "to_timestamp"):
399
+ key = prefix + field
400
+ if key in row:
401
+ row[key] = values[field]
402
+
403
+ for key, value in list(row.items()):
404
+ if key.startswith(("videos/", "depths/")) and key.endswith("/to_timestamp"):
405
+ prefix = key[: -len("to_timestamp")]
406
+ row[key] = float(row.get(prefix + "from_timestamp", 0.0)) + end_times[episode]
407
+ if key.startswith("sensors/") and key.endswith("/path") and value in sensor_metrics:
408
+ prefix = key[: -len("path")]
409
+ for field, metric_value in sensor_metrics[value].items():
410
+ metric_key = prefix + field
411
+ if metric_key in row:
412
+ row[metric_key] = metric_value
413
+
414
+
415
+ def cut_dataset(input_dataset: Path, output_path: Path, minutes: float) -> None:
416
+ source = input_dataset.expanduser().resolve()
417
+ output = output_path.expanduser().resolve()
418
+ if not source.is_dir():
419
+ raise NotADirectoryError(f"input dataset does not exist: {source}")
420
+ if output.exists():
421
+ raise FileExistsError(f"output path already exists: {output}")
422
+ if output.is_relative_to(source):
423
+ raise ValueError("output path must be outside the input dataset")
424
+ data_root = source / "data"
425
+ data_paths = sorted(data_root.glob("**/*.parquet"))
426
+ if not data_paths:
427
+ raise ValueError(f"no data parquet files found under {data_root}")
428
+
429
+ info_path = source / "meta" / "info.json"
430
+ info = json.loads(info_path.read_text(encoding="utf-8"))
431
+ fps = float(info.get("fps", 30.0))
432
+ limits = _candidate_lengths(data_paths, data_root, minutes * 60, fps)
433
+ episode_tables = _episode_tables(source)
434
+ episode_rows = _episode_rows(episode_tables)
435
+
436
+ media_paths = sorted(
437
+ path
438
+ for directory in (source / "videos", source / "depths")
439
+ if directory.exists()
440
+ for path in directory.glob("**/*")
441
+ if path.is_file() and path.suffix.lower() in VIDEO_SUFFIXES
442
+ )
443
+ media_episodes = {}
444
+ for path in media_paths:
445
+ episodes = _media_episodes(path, source, episode_rows)
446
+ if len(episodes) != 1:
447
+ raise ValueError(
448
+ f"media file contains multiple episodes and cannot be cut safely: {path}"
449
+ )
450
+ episode = episodes[0]
451
+ if episode not in limits:
452
+ raise ValueError(f"media references unknown episode {episode}: {path}")
453
+ media_episodes[path] = episode
454
+ limits[episode] = min(limits[episode], _probe_frames(path))
455
+
456
+ if any(length <= 0 for length in limits.values()):
457
+ raise ValueError("the requested cut produced an empty episode")
458
+
459
+ try:
460
+ _copy_static_tree(source, output)
461
+ source_metrics, source_totals = _write_data(
462
+ data_paths, data_root, output, limits, fps
463
+ )
464
+ end_times = {
465
+ episode: min(
466
+ float(episodes[episode]["to_timestamp"])
467
+ - float(episodes[episode]["from_timestamp"])
468
+ for episodes in source_metrics.values()
469
+ if episode in episodes
470
+ )
471
+ for episode in limits
472
+ }
473
+
474
+ for path, episode in media_episodes.items():
475
+ _copy_or_cut_video(
476
+ path, output / path.relative_to(source), limits[episode]
477
+ )
478
+ for directory in (source / "videos", source / "depths"):
479
+ if not directory.exists():
480
+ continue
481
+ for path in directory.glob("**/*"):
482
+ if path.is_file() and path not in media_episodes:
483
+ destination = output / path.relative_to(source)
484
+ destination.parent.mkdir(parents=True, exist_ok=True)
485
+ shutil.copy2(path, destination)
486
+ sensor_metrics = _copy_sensors(
487
+ source, output, episode_rows, end_times
488
+ )
489
+
490
+ offsets = {}
491
+ offset = 0
492
+ for episode in sorted(limits):
493
+ offsets[episode] = offset
494
+ offset += limits[episode]
495
+
496
+ for original_path, table in episode_tables:
497
+ rows = table.to_pylist()
498
+ for row in rows:
499
+ _update_episode_row(
500
+ row,
501
+ limits,
502
+ offsets,
503
+ source_metrics,
504
+ sensor_metrics,
505
+ end_times,
506
+ )
507
+ destination = output / original_path.relative_to(source)
508
+ pq.write_table(pa.Table.from_pylist(rows, schema=table.schema), destination)
509
+
510
+ jsonl_path = output / "meta" / "episodes.jsonl"
511
+ jsonl_rows = _read_json_lines(jsonl_path)
512
+ for row in jsonl_rows:
513
+ _update_episode_row(
514
+ row,
515
+ limits,
516
+ offsets,
517
+ source_metrics,
518
+ sensor_metrics,
519
+ end_times,
520
+ )
521
+ if jsonl_path.exists():
522
+ _write_json_lines(jsonl_path, jsonl_rows)
523
+
524
+ info["total_frames"] = sum(limits.values())
525
+ if isinstance(info.get("source_frame_counts"), dict):
526
+ info["source_frame_counts"] = {
527
+ source_name: source_totals.get(source_name, 0)
528
+ for source_name in info["source_frame_counts"]
529
+ }
530
+ (output / "meta" / "info.json").write_text(
531
+ json.dumps(info, ensure_ascii=False, indent=2) + "\n",
532
+ encoding="utf-8",
533
+ )
534
+ except Exception:
535
+ shutil.rmtree(output, ignore_errors=True)
536
+ raise
537
+
538
+
539
+ def execute(args: argparse.Namespace) -> int:
540
+ cut_dataset(args.input_dataset, args.output_path, args.minutes)
541
+ print(
542
+ json.dumps(
543
+ {
544
+ "input": str(args.input_dataset),
545
+ "output": str(args.output_path),
546
+ "minutes": args.minutes,
547
+ },
548
+ ensure_ascii=False,
549
+ )
550
+ )
551
+ return 0
@@ -81,6 +81,8 @@ def configure_parser(parser: argparse.ArgumentParser) -> None:
81
81
  def execute(args: argparse.Namespace) -> int:
82
82
  config: DeploymentConfig | None = None
83
83
  service: dict[str, object] | None = None
84
+ local_scratch = False
85
+ local_output = False
84
86
  if args.config:
85
87
  # Apply the process mask before refreshing/loading the algorithm. A
86
88
  # repository module may import CUDA during module import.
@@ -91,6 +93,8 @@ def execute(args: argparse.Namespace) -> int:
91
93
  args.port = int(service["port"])
92
94
  args.max_concurrent_executions = int(service["maxConcurrency"])
93
95
  args.scratch_dir = service.get("scratchDir")
96
+ local_scratch = bool(service.get("localScratch", False))
97
+ local_output = bool(service.get("localOutput", False))
94
98
  args.gpu_ids = list(service.get("gpuIds", []))
95
99
  args.webui = bool(service.get("webui"))
96
100
  values = config.process_environment()
@@ -120,6 +124,8 @@ def execute(args: argparse.Namespace) -> int:
120
124
  ),
121
125
  max_concurrent_executions=args.max_concurrent_executions,
122
126
  scratch_dir=args.scratch_dir,
127
+ local_scratch=local_scratch,
128
+ local_output=local_output,
123
129
  gpu_ids=args.gpu_ids,
124
130
  runner_factory=create_runner,
125
131
  max_attempts=args.max_attempts,