143 lines
3.8 KiB
Python
143 lines
3.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Runs preprocessing subprocess for datasets.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
from pathlib import Path
|
|
from collections.abc import Callable
|
|
|
|
logger = logging.getLogger("fastvideo.ui.preprocess_runner")
|
|
|
|
|
|
def build_preprocess_args(
|
|
dataset_id: str,
|
|
raw_path: str,
|
|
output_dir: str,
|
|
workload_type: str,
|
|
model_path: str,
|
|
dataset_type: str = "merged",
|
|
num_gpus: int = 1,
|
|
) -> list[str]:
|
|
"""Build CLI args for v1_preprocessing_new."""
|
|
return [
|
|
"--model-path",
|
|
model_path,
|
|
"--mode",
|
|
"preprocess",
|
|
"--workload-type",
|
|
workload_type,
|
|
"--preprocess.video_loader_type",
|
|
"torchvision",
|
|
"--preprocess.dataset_type",
|
|
dataset_type,
|
|
"--preprocess.dataset_path",
|
|
raw_path,
|
|
"--preprocess.dataset_output_dir",
|
|
output_dir,
|
|
"--preprocess.preprocess_video_batch_size",
|
|
"2",
|
|
"--preprocess.dataloader_num_workers",
|
|
"0",
|
|
"--preprocess.max_height",
|
|
"480",
|
|
"--preprocess.max_width",
|
|
"832",
|
|
"--preprocess.num_frames",
|
|
"77",
|
|
"--preprocess.train_fps",
|
|
"16",
|
|
"--preprocess.samples_per_file",
|
|
"8",
|
|
"--preprocess.flush_frequency",
|
|
"8",
|
|
"--preprocess.video_length_tolerance_range",
|
|
"5",
|
|
]
|
|
|
|
|
|
def run_preprocess(
|
|
dataset_id: str,
|
|
raw_path: str,
|
|
output_dir: str,
|
|
workload_type: str,
|
|
model_path: str,
|
|
dataset_type: str,
|
|
num_gpus: int,
|
|
log_file_path: str,
|
|
on_status_change: Callable[[str, str | None], None],
|
|
stop_event: threading.Event,
|
|
) -> None:
|
|
"""Run preprocessing subprocess. Calls on_status_change(status, error)."""
|
|
repo_root = Path(__file__).resolve().parent.parent
|
|
preprocess_module = "fastvideo.pipelines.preprocess.v1_preprocessing_new"
|
|
|
|
args = build_preprocess_args(
|
|
dataset_id=dataset_id,
|
|
raw_path=raw_path,
|
|
output_dir=output_dir,
|
|
workload_type=workload_type,
|
|
model_path=model_path,
|
|
dataset_type=dataset_type,
|
|
num_gpus=num_gpus,
|
|
)
|
|
|
|
cmd = [
|
|
sys.executable,
|
|
"-m",
|
|
"torch.distributed.run",
|
|
"--nproc_per_node",
|
|
str(num_gpus),
|
|
"--nnodes",
|
|
"1",
|
|
"-m",
|
|
preprocess_module,
|
|
] + args
|
|
|
|
env = os.environ.copy()
|
|
env["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
|
try:
|
|
on_status_change("preprocessing", None)
|
|
with open(log_file_path, "w", encoding="utf-8") as log_file:
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
cwd=str(repo_root),
|
|
env=env,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
bufsize=1,
|
|
)
|
|
assert proc.stdout is not None
|
|
for line in iter(proc.stdout.readline, ""):
|
|
if stop_event.is_set():
|
|
proc.terminate()
|
|
try:
|
|
proc.wait(timeout=30)
|
|
except subprocess.TimeoutExpired:
|
|
proc.kill()
|
|
on_status_change("stopped", "Preprocessing stopped by user")
|
|
return
|
|
line = line.rstrip()
|
|
if line:
|
|
log_file.write(line + "\n")
|
|
log_file.flush()
|
|
|
|
proc.wait()
|
|
exit_code = proc.returncode or 0
|
|
|
|
if exit_code == 0:
|
|
on_status_change("ready", None)
|
|
else:
|
|
on_status_change("failed", f"Preprocessing exited with code {exit_code}")
|
|
except Exception as exc:
|
|
on_status_change("failed", f"{type(exc).__name__}: {exc}")
|
|
logger.exception("Preprocessing failed for dataset %s", dataset_id)
|