1071 lines
40 KiB
Python
Executable File
1071 lines
40 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Pull completed TG3 data-collection episodes from the robot atomically."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import ctypes
|
|
import errno
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import re
|
|
import shlex
|
|
import shutil
|
|
import signal
|
|
import stat
|
|
import subprocess
|
|
import time
|
|
import uuid
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Iterator, Sequence
|
|
|
|
|
|
SAFE_EPISODE_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$")
|
|
SAFE_SHA256 = re.compile(r"^[0-9a-f]{64}$")
|
|
VERIFIED_RECEIPT_NAME = "VERIFIED"
|
|
VERIFIED_RECEIPT_VERSION = 1
|
|
FIXED_REMOTE_READY = "/home/nvidia/tg3_data_collection/ready"
|
|
DEFAULT_REMOTE_DELETE_HELPER = (
|
|
"/home/nvidia/tg3_local_teleop/delete_ready_episode.py"
|
|
)
|
|
AT_FDCWD = -100
|
|
RENAME_NOREPLACE = 1
|
|
|
|
|
|
def _fsync_directory(path: Path) -> None:
|
|
flags = os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0)
|
|
descriptor = os.open(path, flags)
|
|
try:
|
|
if not stat.S_ISDIR(os.fstat(descriptor).st_mode):
|
|
raise ValueError(f"not a directory: {path}")
|
|
os.fsync(descriptor)
|
|
finally:
|
|
os.close(descriptor)
|
|
|
|
|
|
def atomic_rename_noreplace(source: Path, destination: Path) -> None:
|
|
"""Atomically publish a directory without replacing a colliding target."""
|
|
|
|
libc = ctypes.CDLL(None, use_errno=True)
|
|
renameat2 = getattr(libc, "renameat2", None)
|
|
if renameat2 is not None:
|
|
renameat2.argtypes = [
|
|
ctypes.c_int,
|
|
ctypes.c_char_p,
|
|
ctypes.c_int,
|
|
ctypes.c_char_p,
|
|
ctypes.c_uint,
|
|
]
|
|
renameat2.restype = ctypes.c_int
|
|
result = renameat2(
|
|
AT_FDCWD,
|
|
os.fsencode(source),
|
|
AT_FDCWD,
|
|
os.fsencode(destination),
|
|
RENAME_NOREPLACE,
|
|
)
|
|
if result != 0:
|
|
error_number = ctypes.get_errno()
|
|
raise OSError(
|
|
error_number,
|
|
os.strerror(error_number),
|
|
str(destination),
|
|
)
|
|
return
|
|
# Linux targets used by TG3 expose renameat2. This fallback remains safe
|
|
# against the supported sync processes because they share the flock.
|
|
if os.path.lexists(destination):
|
|
raise FileExistsError(errno.EEXIST, "destination already exists", destination)
|
|
os.rename(source, destination)
|
|
|
|
|
|
def atomic_write_json(
|
|
path: Path,
|
|
payload: dict[str, Any],
|
|
*,
|
|
durable: bool = False,
|
|
) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = path.parent / f".{path.name}.{uuid.uuid4().hex}.tmp"
|
|
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
|
|
descriptor = os.open(temporary, flags, 0o600)
|
|
try:
|
|
data = (json.dumps(payload, ensure_ascii=False, indent=2) + "\n").encode(
|
|
"utf-8"
|
|
)
|
|
with os.fdopen(descriptor, "wb", closefd=False) as stream:
|
|
stream.write(data)
|
|
stream.flush()
|
|
if durable:
|
|
os.fsync(descriptor)
|
|
except Exception:
|
|
try:
|
|
temporary.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
raise
|
|
finally:
|
|
os.close(descriptor)
|
|
try:
|
|
os.replace(temporary, path)
|
|
if durable:
|
|
_fsync_directory(path.parent)
|
|
finally:
|
|
try:
|
|
temporary.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
|
|
descriptor = os.open(path, flags)
|
|
try:
|
|
if not stat.S_ISREG(os.fstat(descriptor).st_mode):
|
|
raise ValueError(f"not a regular file: {path}")
|
|
stream = os.fdopen(descriptor, "rb", closefd=False)
|
|
with stream:
|
|
for block in iter(lambda: stream.read(1024 * 1024), b""):
|
|
digest.update(block)
|
|
finally:
|
|
os.close(descriptor)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _regular_file(path: Path, description: str) -> None:
|
|
try:
|
|
metadata = path.lstat()
|
|
except FileNotFoundError as exc:
|
|
raise ValueError(f"missing {description}: {path}") from exc
|
|
if not stat.S_ISREG(metadata.st_mode):
|
|
raise ValueError(f"{description} is not a regular file: {path}")
|
|
|
|
|
|
def fsync_episode_tree(directory: Path) -> None:
|
|
"""Reject links/special entries, fsync every file, then every directory."""
|
|
|
|
try:
|
|
root_metadata = directory.lstat()
|
|
except FileNotFoundError as exc:
|
|
raise ValueError(f"episode directory is missing: {directory}") from exc
|
|
if not stat.S_ISDIR(root_metadata.st_mode):
|
|
raise ValueError(f"episode path is not a real directory: {directory}")
|
|
|
|
directories: list[Path] = []
|
|
for root, child_directories, child_files in os.walk(
|
|
directory, topdown=True, followlinks=False
|
|
):
|
|
root_path = Path(root)
|
|
directories.append(root_path)
|
|
for name in child_directories:
|
|
child = root_path / name
|
|
if not stat.S_ISDIR(child.lstat().st_mode):
|
|
raise ValueError(f"unsafe non-directory entry: {child}")
|
|
for name in child_files:
|
|
child = root_path / name
|
|
metadata = child.lstat()
|
|
if not stat.S_ISREG(metadata.st_mode):
|
|
raise ValueError(f"unsafe non-regular episode entry: {child}")
|
|
flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
|
|
descriptor = os.open(child, flags)
|
|
try:
|
|
if not stat.S_ISREG(os.fstat(descriptor).st_mode):
|
|
raise ValueError(f"unsafe episode file: {child}")
|
|
os.fsync(descriptor)
|
|
finally:
|
|
os.close(descriptor)
|
|
for child in reversed(directories):
|
|
_fsync_directory(child)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ManifestDocument:
|
|
raw: bytes
|
|
payload: dict[str, Any]
|
|
sha256: str
|
|
|
|
@classmethod
|
|
def from_bytes(cls, raw: bytes, expected_episode: str) -> "ManifestDocument":
|
|
try:
|
|
text = raw.decode("utf-8")
|
|
except UnicodeDecodeError as exc:
|
|
raise ValueError(f"manifest is not UTF-8: {exc}") from exc
|
|
return cls(
|
|
raw=raw,
|
|
payload=parse_manifest(text, expected_episode),
|
|
sha256=hashlib.sha256(raw).hexdigest(),
|
|
)
|
|
|
|
|
|
def safe_episode_name(value: str) -> bool:
|
|
return bool(SAFE_EPISODE_NAME.fullmatch(value)) and value not in (".", "..")
|
|
|
|
|
|
def safe_relative_path(value: object) -> Path:
|
|
if not isinstance(value, str) or not value:
|
|
raise ValueError("manifest file path must be a non-empty string")
|
|
path = Path(value)
|
|
if path.is_absolute() or ".." in path.parts:
|
|
raise ValueError(f"unsafe manifest file path: {value!r}")
|
|
return path
|
|
|
|
|
|
def parse_manifest(raw: str, expected_episode: str) -> dict[str, Any]:
|
|
payload = json.loads(raw)
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("manifest root must be an object")
|
|
if payload.get("state") != "complete":
|
|
raise ValueError("remote episode is not complete")
|
|
if payload.get("episode_id") != expected_episode:
|
|
raise ValueError("manifest episode_id does not match directory name")
|
|
files = payload.get("files")
|
|
if not isinstance(files, list) or not files:
|
|
raise ValueError("manifest files must be a non-empty list")
|
|
saw_mcap = False
|
|
seen_paths: set[str] = set()
|
|
for entry in files:
|
|
if not isinstance(entry, dict):
|
|
raise ValueError("manifest file entry must be an object")
|
|
relative = safe_relative_path(entry.get("path"))
|
|
normalized = relative.as_posix()
|
|
if normalized in seen_paths:
|
|
raise ValueError(f"duplicate manifest file path: {relative}")
|
|
seen_paths.add(normalized)
|
|
saw_mcap = saw_mcap or relative.suffix == ".mcap"
|
|
size = entry.get("size")
|
|
digest = entry.get("sha256")
|
|
if isinstance(size, bool) or not isinstance(size, int) or size < 0:
|
|
raise ValueError(f"invalid size for {relative}")
|
|
if not (
|
|
isinstance(digest, str)
|
|
and len(digest) == 64
|
|
and all(character in "0123456789abcdef" for character in digest)
|
|
):
|
|
raise ValueError(f"invalid sha256 for {relative}")
|
|
if not saw_mcap:
|
|
raise ValueError("manifest contains no MCAP file")
|
|
return payload
|
|
|
|
|
|
def validate_episode_dir(directory: Path, manifest: dict[str, Any]) -> None:
|
|
ready = directory / "READY"
|
|
_regular_file(ready, "READY marker")
|
|
_regular_file(directory / "manifest.json", "manifest")
|
|
listed_payloads: set[str] = set()
|
|
for entry in manifest["files"]:
|
|
relative = safe_relative_path(entry["path"])
|
|
listed_payloads.add(relative.as_posix())
|
|
path = directory / relative
|
|
_regular_file(path, f"copied file {relative}")
|
|
if path.stat().st_size != entry["size"]:
|
|
raise ValueError(f"size mismatch: {relative}")
|
|
if sha256_file(path) != entry["sha256"]:
|
|
raise ValueError(f"sha256 mismatch: {relative}")
|
|
# A resumed staging directory must never smuggle an old bag segment into a
|
|
# newly verified episode. Logs/control markers may be unlisted, but every
|
|
# MCAP and rosbag metadata payload must be closed over by the manifest.
|
|
for root, _directories, files in os.walk(directory, followlinks=False):
|
|
root_path = Path(root)
|
|
for name in files:
|
|
relative = (root_path / name).relative_to(directory).as_posix()
|
|
if (name.endswith(".mcap") or name == "metadata.yaml") and (
|
|
relative not in listed_payloads
|
|
):
|
|
raise ValueError(f"unlisted bag payload: {relative}")
|
|
|
|
|
|
class CommandRunner:
|
|
def run(
|
|
self,
|
|
command: Sequence[str],
|
|
*,
|
|
timeout: float,
|
|
) -> subprocess.CompletedProcess[str]:
|
|
return subprocess.run(
|
|
list(command),
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=timeout,
|
|
)
|
|
|
|
|
|
class DataGetSync:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
remote: str,
|
|
remote_ready: str,
|
|
destination: Path,
|
|
status_file: Path,
|
|
runner: CommandRunner | None = None,
|
|
ssh_timeout_s: float = 8.0,
|
|
reserve_bytes: int = 1_073_741_824,
|
|
verify_existing: bool = False,
|
|
delete_remote_after_sync: bool = False,
|
|
remote_delete_helper: str = DEFAULT_REMOTE_DELETE_HELPER,
|
|
remote_delete_timeout_s: float = 600.0,
|
|
) -> None:
|
|
self.remote = remote
|
|
self.remote_ready = remote_ready.rstrip("/")
|
|
self.destination = destination
|
|
self.status_file = status_file
|
|
self.runner = runner or CommandRunner()
|
|
self.ssh_timeout_s = ssh_timeout_s
|
|
self.reserve_bytes = reserve_bytes
|
|
self.verify_existing = verify_existing
|
|
self.delete_remote_after_sync = delete_remote_after_sync
|
|
self.remote_delete_helper = remote_delete_helper
|
|
self.remote_delete_timeout_s = remote_delete_timeout_s
|
|
if self.delete_remote_after_sync and self.remote_ready != FIXED_REMOTE_READY:
|
|
raise ValueError(
|
|
"automatic deletion is restricted to the fixed Nvidia ready root: "
|
|
f"{FIXED_REMOTE_READY}"
|
|
)
|
|
if self.delete_remote_after_sync and not self.remote_delete_helper.startswith("/"):
|
|
raise ValueError("remote delete helper path must be absolute")
|
|
if (
|
|
not math.isfinite(self.remote_delete_timeout_s)
|
|
or self.remote_delete_timeout_s <= 0.0
|
|
):
|
|
raise ValueError("remote delete timeout must be positive and finite")
|
|
self.stop_requested = False
|
|
self.started_at = time.time()
|
|
self.sync_count = 0
|
|
self.delete_count = 0
|
|
self.last_episode: str | None = None
|
|
self.last_deleted_episode: str | None = None
|
|
self.last_error = ""
|
|
self.last_delete_error = ""
|
|
self.pending_remote_cleanup = 0
|
|
self._verified_this_process: set[tuple[str, str]] = set()
|
|
|
|
@property
|
|
def ssh_base(self) -> list[str]:
|
|
return [
|
|
"ssh",
|
|
"-o",
|
|
"BatchMode=yes",
|
|
"-o",
|
|
"StrictHostKeyChecking=yes",
|
|
"-o",
|
|
f"ConnectTimeout={max(1, int(self.ssh_timeout_s))}",
|
|
self.remote,
|
|
]
|
|
|
|
def request_stop(self, _signum: int, _frame: object) -> None:
|
|
self.stop_requested = True
|
|
|
|
@property
|
|
def remote_identity(self) -> str:
|
|
return f"{self.remote}:{self.remote_ready}"
|
|
|
|
def _write_status(self, state: str) -> None:
|
|
atomic_write_json(
|
|
self.status_file,
|
|
{
|
|
"state": state,
|
|
"remote": self.remote,
|
|
"remote_ready": self.remote_ready,
|
|
"destination": str(self.destination),
|
|
"sync_count": self.sync_count,
|
|
"delete_count": self.delete_count,
|
|
"last_episode": self.last_episode,
|
|
"last_deleted_episode": self.last_deleted_episode,
|
|
"last_error": self.last_error,
|
|
"last_delete_error": self.last_delete_error,
|
|
"pending_remote_cleanup": self.pending_remote_cleanup,
|
|
"uptime_s": round(time.time() - self.started_at, 1),
|
|
"updated_unix_s": time.time(),
|
|
},
|
|
)
|
|
|
|
def _remote_command(self, command: str, *, timeout: float | None = None) -> str:
|
|
result = self.runner.run(
|
|
[*self.ssh_base, command],
|
|
timeout=(self.ssh_timeout_s + 2.0 if timeout is None else timeout),
|
|
)
|
|
if result.returncode != 0:
|
|
detail = result.stderr.strip() or result.stdout.strip()
|
|
raise RuntimeError(f"remote command failed: {detail}")
|
|
return result.stdout
|
|
|
|
@contextmanager
|
|
def process_lock(self) -> Iterator[None]:
|
|
"""Serialize the daemon and any manual ``--once`` invocation."""
|
|
|
|
self.destination.mkdir(parents=True, exist_ok=True)
|
|
lock_path = self.destination / ".data_get_sync.lock"
|
|
flags = os.O_RDWR | os.O_CREAT | getattr(os, "O_NOFOLLOW", 0)
|
|
try:
|
|
descriptor = os.open(lock_path, flags, 0o600)
|
|
except OSError as exc:
|
|
raise RuntimeError(f"cannot open safe sync lock {lock_path}: {exc}") from exc
|
|
try:
|
|
if not stat.S_ISREG(os.fstat(descriptor).st_mode):
|
|
raise RuntimeError(f"sync lock is not a regular file: {lock_path}")
|
|
try:
|
|
fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
except BlockingIOError as exc:
|
|
raise RuntimeError(
|
|
"another tg3 data sync process already holds the destination lock"
|
|
) from exc
|
|
os.ftruncate(descriptor, 0)
|
|
os.write(descriptor, f"{os.getpid()}\n".encode("ascii"))
|
|
yield
|
|
finally:
|
|
os.close(descriptor)
|
|
|
|
def list_remote_episodes(self) -> list[str]:
|
|
root = shlex.quote(self.remote_ready)
|
|
output = self._remote_command(
|
|
f"if test -d {root}; then "
|
|
f"find {root} -mindepth 1 -maxdepth 1 -type d -printf '%f\\n'; "
|
|
"fi"
|
|
)
|
|
names = sorted({line.strip() for line in output.splitlines() if line.strip()})
|
|
unsafe = [name for name in names if not safe_episode_name(name)]
|
|
if unsafe:
|
|
raise RuntimeError(f"remote returned unsafe episode names: {unsafe!r}")
|
|
return names
|
|
|
|
def get_remote_manifest(self, episode: str) -> ManifestDocument:
|
|
if not safe_episode_name(episode):
|
|
raise ValueError(f"unsafe episode name: {episode!r}")
|
|
episode_path = f"{self.remote_ready}/{episode}"
|
|
manifest_path = f"{episode_path}/manifest.json"
|
|
ready_path = f"{episode_path}/READY"
|
|
raw = self._remote_command(
|
|
f"test -d {shlex.quote(episode_path)} && "
|
|
f"test ! -L {shlex.quote(episode_path)} && "
|
|
f"test -f {shlex.quote(ready_path)} && "
|
|
f"test ! -L {shlex.quote(ready_path)} && "
|
|
f"test -f {shlex.quote(manifest_path)} && "
|
|
f"test ! -L {shlex.quote(manifest_path)} && "
|
|
f"cat {shlex.quote(manifest_path)}"
|
|
)
|
|
return ManifestDocument.from_bytes(raw.encode("utf-8"), episode)
|
|
|
|
def _enough_local_space(self, manifest: dict[str, Any]) -> bool:
|
|
required = sum(int(entry["size"]) for entry in manifest["files"])
|
|
free = shutil.disk_usage(self.destination).free
|
|
return free >= required + self.reserve_bytes
|
|
|
|
@staticmethod
|
|
def _manifest_totals(manifest: dict[str, Any]) -> tuple[int, int]:
|
|
return (
|
|
len(manifest["files"]),
|
|
sum(int(entry["size"]) for entry in manifest["files"]),
|
|
)
|
|
|
|
@staticmethod
|
|
def _read_manifest_document(directory: Path, episode: str) -> ManifestDocument:
|
|
manifest_path = directory / "manifest.json"
|
|
_regular_file(manifest_path, "manifest")
|
|
return ManifestDocument.from_bytes(manifest_path.read_bytes(), episode)
|
|
|
|
def _receipt_payload(
|
|
self,
|
|
final: Path,
|
|
document: ManifestDocument,
|
|
*,
|
|
delete_state: str,
|
|
attempts: int = 0,
|
|
last_error: str = "",
|
|
deleted_at_unix_s: float | None = None,
|
|
next_retry_unix_s: float = 0.0,
|
|
retry_delay_s: float = 0.0,
|
|
) -> dict[str, Any]:
|
|
file_count, total_bytes = self._manifest_totals(document.payload)
|
|
remote_delete: dict[str, Any] = {
|
|
"state": delete_state,
|
|
"attempts": attempts,
|
|
"last_error": last_error,
|
|
"next_retry_unix_s": next_retry_unix_s,
|
|
"retry_delay_s": retry_delay_s,
|
|
}
|
|
if deleted_at_unix_s is not None:
|
|
remote_delete["deleted_at_unix_s"] = deleted_at_unix_s
|
|
return {
|
|
"schema_version": VERIFIED_RECEIPT_VERSION,
|
|
"state": "VERIFIED",
|
|
"episode_id": document.payload["episode_id"],
|
|
"manifest_sha256": document.sha256,
|
|
"remote_identity": self.remote_identity,
|
|
"remote": self.remote,
|
|
"remote_ready": self.remote_ready,
|
|
"local_final": str(final.resolve(strict=True)),
|
|
"verified_files": file_count,
|
|
"verified_bytes": total_bytes,
|
|
"verified_at_unix_s": time.time(),
|
|
"remote_delete": remote_delete,
|
|
}
|
|
|
|
@staticmethod
|
|
def _receipt_path(final: Path) -> Path:
|
|
return final / VERIFIED_RECEIPT_NAME
|
|
|
|
def _read_receipt(self, final: Path) -> dict[str, Any] | None:
|
|
path = self._receipt_path(final)
|
|
try:
|
|
metadata = path.lstat()
|
|
except FileNotFoundError:
|
|
return None
|
|
if not stat.S_ISREG(metadata.st_mode):
|
|
raise RuntimeError(f"VERIFIED receipt is not a regular file: {path}")
|
|
try:
|
|
payload = json.loads(path.read_text(encoding="utf-8"))
|
|
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
|
raise RuntimeError(f"invalid VERIFIED receipt {path}: {exc}") from exc
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError(f"invalid VERIFIED receipt root: {path}")
|
|
return payload
|
|
|
|
def _validate_receipt(
|
|
self,
|
|
final: Path,
|
|
document: ManifestDocument,
|
|
receipt: dict[str, Any],
|
|
) -> None:
|
|
file_count, total_bytes = self._manifest_totals(document.payload)
|
|
expected = {
|
|
"schema_version": VERIFIED_RECEIPT_VERSION,
|
|
"state": "VERIFIED",
|
|
"episode_id": document.payload["episode_id"],
|
|
"manifest_sha256": document.sha256,
|
|
"remote_identity": self.remote_identity,
|
|
"remote": self.remote,
|
|
"remote_ready": self.remote_ready,
|
|
"local_final": str(final.resolve(strict=True)),
|
|
"verified_files": file_count,
|
|
"verified_bytes": total_bytes,
|
|
}
|
|
for key, value in expected.items():
|
|
if receipt.get(key) != value:
|
|
raise RuntimeError(f"VERIFIED receipt mismatch for {key}")
|
|
remote_delete = receipt.get("remote_delete")
|
|
if not isinstance(remote_delete, dict):
|
|
raise RuntimeError("VERIFIED receipt has no remote_delete state")
|
|
if remote_delete.get("state") not in ("retained", "pending", "deleted"):
|
|
raise RuntimeError("VERIFIED receipt has invalid remote_delete state")
|
|
attempts = remote_delete.get("attempts")
|
|
if isinstance(attempts, bool) or not isinstance(attempts, int) or attempts < 0:
|
|
raise RuntimeError("VERIFIED receipt has invalid delete attempts")
|
|
for key in ("next_retry_unix_s", "retry_delay_s"):
|
|
value = remote_delete.get(key, 0.0)
|
|
if (
|
|
isinstance(value, bool)
|
|
or not isinstance(value, (int, float))
|
|
or not math.isfinite(float(value))
|
|
or float(value) < 0.0
|
|
):
|
|
raise RuntimeError(f"VERIFIED receipt has invalid {key}")
|
|
|
|
def _write_receipt(self, final: Path, receipt: dict[str, Any]) -> None:
|
|
atomic_write_json(self._receipt_path(final), receipt, durable=True)
|
|
|
|
@staticmethod
|
|
def _receipt_delete_state(receipt: dict[str, Any]) -> str:
|
|
remote_delete = receipt.get("remote_delete")
|
|
if not isinstance(remote_delete, dict):
|
|
raise RuntimeError("VERIFIED receipt has no remote_delete state")
|
|
state = remote_delete.get("state")
|
|
if state not in ("retained", "pending", "deleted"):
|
|
raise RuntimeError("VERIFIED receipt has invalid remote_delete state")
|
|
return str(state)
|
|
|
|
@staticmethod
|
|
def _delete_retry_due(receipt: dict[str, Any], now: float | None = None) -> bool:
|
|
remote_delete = receipt["remote_delete"]
|
|
next_retry = float(remote_delete.get("next_retry_unix_s", 0.0))
|
|
return (time.time() if now is None else now) >= next_retry
|
|
|
|
def _deep_verify_local(
|
|
self,
|
|
final: Path,
|
|
document: ManifestDocument,
|
|
) -> None:
|
|
validate_episode_dir(final, document.payload)
|
|
# Besides making the data durable, this traversal rejects symlinks in
|
|
# intermediate directories that a leaf-only manifest check cannot see.
|
|
fsync_episode_tree(final)
|
|
self._verified_this_process.add(
|
|
(document.payload["episode_id"], document.sha256)
|
|
)
|
|
|
|
def _ensure_existing_verified(
|
|
self,
|
|
final: Path,
|
|
remote_document: ManifestDocument,
|
|
) -> tuple[ManifestDocument, dict[str, Any]]:
|
|
try:
|
|
final_metadata = final.lstat()
|
|
except FileNotFoundError as exc:
|
|
raise RuntimeError(f"existing destination disappeared: {final}") from exc
|
|
if not stat.S_ISDIR(final_metadata.st_mode):
|
|
raise RuntimeError(f"existing destination is not a real directory: {final}")
|
|
|
|
local_document = self._read_manifest_document(
|
|
final, remote_document.payload["episode_id"]
|
|
)
|
|
if (
|
|
local_document.sha256 != remote_document.sha256
|
|
or local_document.payload != remote_document.payload
|
|
):
|
|
raise RuntimeError(
|
|
"existing local episode does not match the current remote manifest"
|
|
)
|
|
|
|
try:
|
|
receipt = self._read_receipt(final)
|
|
except RuntimeError:
|
|
# A malformed regular receipt is not a credential. The full tree
|
|
# verification below must succeed before it can be replaced. A
|
|
# symlink receipt is rejected by fsync_episode_tree.
|
|
receipt = None
|
|
receipt_trusted = False
|
|
if receipt is not None:
|
|
state = self._receipt_delete_state(receipt)
|
|
if (
|
|
state == "pending"
|
|
and receipt.get("remote_identity") != self.remote_identity
|
|
):
|
|
raise RuntimeError(
|
|
"pending remote deletion belongs to a different remote identity"
|
|
)
|
|
try:
|
|
self._validate_receipt(final, local_document, receipt)
|
|
receipt_trusted = True
|
|
except RuntimeError:
|
|
receipt = None
|
|
|
|
if receipt_trusted:
|
|
if self.verify_existing:
|
|
self._deep_verify_local(final, local_document)
|
|
return local_document, receipt
|
|
|
|
# Missing/malformed/stale non-pending receipts are never deletion
|
|
# credentials. Rebuild one only after a current full payload hash.
|
|
if self.verify_existing or not receipt_trusted:
|
|
self._deep_verify_local(final, local_document)
|
|
|
|
desired_state = "pending" if self.delete_remote_after_sync else "retained"
|
|
receipt = self._receipt_payload(
|
|
final, local_document, delete_state=desired_state
|
|
)
|
|
self._write_receipt(final, receipt)
|
|
return local_document, receipt
|
|
|
|
def _set_delete_state(
|
|
self,
|
|
final: Path,
|
|
document: ManifestDocument,
|
|
receipt: dict[str, Any],
|
|
*,
|
|
state: str,
|
|
attempts: int,
|
|
last_error: str,
|
|
deleted_at_unix_s: float | None = None,
|
|
next_retry_unix_s: float = 0.0,
|
|
retry_delay_s: float = 0.0,
|
|
) -> dict[str, Any]:
|
|
updated = self._receipt_payload(
|
|
final,
|
|
document,
|
|
delete_state=state,
|
|
attempts=attempts,
|
|
last_error=last_error,
|
|
deleted_at_unix_s=deleted_at_unix_s,
|
|
next_retry_unix_s=next_retry_unix_s,
|
|
retry_delay_s=retry_delay_s,
|
|
)
|
|
self._write_receipt(final, updated)
|
|
return updated
|
|
|
|
def _invoke_delete_helper(
|
|
self,
|
|
episode: str,
|
|
manifest_sha256: str,
|
|
) -> dict[str, Any]:
|
|
command = " ".join(
|
|
shlex.quote(part)
|
|
for part in (self.remote_delete_helper, episode, manifest_sha256)
|
|
)
|
|
output = self._remote_command(
|
|
command, timeout=self.remote_delete_timeout_s
|
|
)
|
|
lines = [line for line in output.splitlines() if line.strip()]
|
|
if not lines:
|
|
raise RuntimeError("remote delete helper returned no acknowledgement")
|
|
try:
|
|
acknowledgement = json.loads(lines[-1])
|
|
except json.JSONDecodeError as exc:
|
|
raise RuntimeError("remote delete helper returned invalid JSON") from exc
|
|
if not isinstance(acknowledgement, dict):
|
|
raise RuntimeError("remote delete helper acknowledgement is not an object")
|
|
if acknowledgement.get("state") not in (
|
|
"deleted",
|
|
"resumed_delete",
|
|
"already_absent",
|
|
):
|
|
raise RuntimeError(
|
|
f"remote delete helper refused: {acknowledgement!r}"
|
|
)
|
|
if acknowledgement.get("episode_id") != episode:
|
|
raise RuntimeError("remote delete acknowledgement episode mismatch")
|
|
if acknowledgement.get("manifest_sha256") != manifest_sha256:
|
|
raise RuntimeError("remote delete acknowledgement manifest mismatch")
|
|
return acknowledgement
|
|
|
|
def _delete_verified_remote(
|
|
self,
|
|
final: Path,
|
|
document: ManifestDocument,
|
|
receipt: dict[str, Any],
|
|
) -> bool:
|
|
self._validate_receipt(final, document, receipt)
|
|
# A VERIFIED receipt is a durable recovery cursor, not proof that the
|
|
# local payload has remained intact. Re-hash on every real helper call,
|
|
# including retries after an earlier remote deletion failure.
|
|
self._deep_verify_local(final, document)
|
|
|
|
attempts = int(receipt["remote_delete"]["attempts"]) + 1
|
|
retry_delay_s = min(900.0, 30.0 * (2 ** min(attempts - 1, 5)))
|
|
next_retry_unix_s = time.time() + retry_delay_s
|
|
# This durable pending write is the recovery cursor if the SSH ACK is
|
|
# lost after the robot has already renamed or removed its staging copy.
|
|
pending = self._set_delete_state(
|
|
final,
|
|
document,
|
|
receipt,
|
|
state="pending",
|
|
attempts=attempts,
|
|
last_error="",
|
|
next_retry_unix_s=next_retry_unix_s,
|
|
retry_delay_s=retry_delay_s,
|
|
)
|
|
try:
|
|
self._invoke_delete_helper(document.payload["episode_id"], document.sha256)
|
|
self._set_delete_state(
|
|
final,
|
|
document,
|
|
pending,
|
|
state="deleted",
|
|
attempts=attempts,
|
|
last_error="",
|
|
deleted_at_unix_s=time.time(),
|
|
next_retry_unix_s=0.0,
|
|
retry_delay_s=0.0,
|
|
)
|
|
except Exception as exc:
|
|
detail = str(exc)
|
|
try:
|
|
self._set_delete_state(
|
|
final,
|
|
document,
|
|
pending,
|
|
state="pending",
|
|
attempts=attempts,
|
|
last_error=detail,
|
|
next_retry_unix_s=next_retry_unix_s,
|
|
retry_delay_s=retry_delay_s,
|
|
)
|
|
except Exception as receipt_exc:
|
|
detail += f"; could not persist pending receipt: {receipt_exc}"
|
|
self.last_delete_error = detail
|
|
return False
|
|
self.delete_count += 1
|
|
self.last_deleted_episode = document.payload["episode_id"]
|
|
self.last_delete_error = ""
|
|
return True
|
|
|
|
def _pending_receipts(
|
|
self,
|
|
) -> tuple[
|
|
list[tuple[Path, ManifestDocument, dict[str, Any]]],
|
|
list[str],
|
|
int,
|
|
]:
|
|
pending: list[tuple[Path, ManifestDocument, dict[str, Any]]] = []
|
|
errors: list[str] = []
|
|
pending_count = 0
|
|
if not self.destination.is_dir():
|
|
return pending, errors, pending_count
|
|
for final in sorted(self.destination.iterdir(), key=lambda path: path.name):
|
|
if not safe_episode_name(final.name):
|
|
continue
|
|
if not stat.S_ISDIR(final.lstat().st_mode):
|
|
continue
|
|
try:
|
|
receipt = self._read_receipt(final)
|
|
if receipt is None:
|
|
continue
|
|
# Historical retained/deleted receipts may legitimately name a
|
|
# previous robot IP. They are not pending work and must not be
|
|
# validated against (or block) the current remote identity.
|
|
if self._receipt_delete_state(receipt) != "pending":
|
|
continue
|
|
pending_count += 1
|
|
document = self._read_manifest_document(final, final.name)
|
|
self._validate_receipt(final, document, receipt)
|
|
except Exception as exc:
|
|
errors.append(f"{final.name}: invalid pending receipt: {exc}")
|
|
continue
|
|
pending.append((final, document, receipt))
|
|
return pending, errors, pending_count
|
|
|
|
def _retry_pending_deletions(self, exclude: set[str]) -> list[str]:
|
|
if not self.delete_remote_after_sync:
|
|
return []
|
|
pending, errors, _pending_count = self._pending_receipts()
|
|
for final, document, receipt in pending:
|
|
episode = document.payload["episode_id"]
|
|
if episode in exclude:
|
|
continue
|
|
if not self._delete_retry_due(receipt):
|
|
continue
|
|
if not self._delete_verified_remote(final, document, receipt):
|
|
errors.append(f"{episode}: {self.last_delete_error}")
|
|
return errors
|
|
|
|
def sync_episode(self, episode: str) -> bool:
|
|
if not safe_episode_name(episode):
|
|
raise ValueError(f"unsafe episode name: {episode!r}")
|
|
remote_document = self.get_remote_manifest(episode)
|
|
final = self.destination / episode
|
|
copied = False
|
|
if os.path.lexists(final):
|
|
local_document, receipt = self._ensure_existing_verified(
|
|
final, remote_document
|
|
)
|
|
if self.delete_remote_after_sync and self._delete_retry_due(receipt):
|
|
self._delete_verified_remote(final, local_document, receipt)
|
|
return False
|
|
|
|
if not self._enough_local_space(remote_document.payload):
|
|
raise RuntimeError("not enough local disk space for episode")
|
|
|
|
incoming_root = self.destination / ".incoming"
|
|
try:
|
|
incoming_root_metadata = incoming_root.lstat()
|
|
except FileNotFoundError:
|
|
incoming_root.mkdir(mode=0o700)
|
|
incoming_root_metadata = incoming_root.lstat()
|
|
if not stat.S_ISDIR(incoming_root_metadata.st_mode):
|
|
raise RuntimeError(f"incoming root is not a real directory: {incoming_root}")
|
|
incoming = incoming_root / f"{episode}.{remote_document.sha256}.partial"
|
|
if os.path.lexists(incoming):
|
|
if not stat.S_ISDIR(incoming.lstat().st_mode):
|
|
raise RuntimeError(f"incoming episode collision: {incoming}")
|
|
else:
|
|
incoming.mkdir(mode=0o700)
|
|
remote_source = f"{self.remote}:{self.remote_ready}/{episode}/"
|
|
ssh_transport = (
|
|
"ssh -o BatchMode=yes "
|
|
"-o StrictHostKeyChecking=yes "
|
|
f"-o ConnectTimeout={max(1, int(self.ssh_timeout_s))}"
|
|
)
|
|
result = self.runner.run(
|
|
[
|
|
"rsync",
|
|
"-a",
|
|
"--partial",
|
|
"--delete-delay",
|
|
"--protect-args",
|
|
"-e",
|
|
ssh_transport,
|
|
remote_source,
|
|
str(incoming) + "/",
|
|
],
|
|
timeout=24 * 60 * 60,
|
|
)
|
|
if result.returncode != 0:
|
|
detail = result.stderr.strip() or result.stdout.strip()
|
|
raise RuntimeError(f"rsync failed: {detail}")
|
|
|
|
copied_manifest_path = incoming / "manifest.json"
|
|
_regular_file(copied_manifest_path, "copied manifest")
|
|
copied_document = ManifestDocument.from_bytes(
|
|
copied_manifest_path.read_bytes(), episode
|
|
)
|
|
if (
|
|
copied_document.sha256 != remote_document.sha256
|
|
or copied_document.payload != remote_document.payload
|
|
):
|
|
raise RuntimeError("remote manifest changed during transfer")
|
|
validate_episode_dir(incoming, copied_document.payload)
|
|
fsync_episode_tree(incoming)
|
|
if os.path.lexists(final):
|
|
raise RuntimeError(f"destination collision before publish: {final}")
|
|
atomic_rename_noreplace(incoming, final)
|
|
_fsync_directory(incoming_root)
|
|
_fsync_directory(self.destination)
|
|
self._verified_this_process.add((episode, copied_document.sha256))
|
|
receipt = self._receipt_payload(
|
|
final,
|
|
copied_document,
|
|
delete_state=("pending" if self.delete_remote_after_sync else "retained"),
|
|
)
|
|
self._write_receipt(final, receipt)
|
|
self.sync_count += 1
|
|
self.last_episode = episode
|
|
copied = True
|
|
if self.delete_remote_after_sync and self._delete_retry_due(receipt):
|
|
self._delete_verified_remote(final, copied_document, receipt)
|
|
return copied
|
|
|
|
def _run_once_unlocked(self) -> int:
|
|
self.destination.mkdir(parents=True, exist_ok=True)
|
|
incoming_root = self.destination / ".incoming"
|
|
incoming_root.mkdir(parents=True, exist_ok=True)
|
|
if not stat.S_ISDIR(incoming_root.lstat().st_mode):
|
|
raise RuntimeError(f"incoming root is not a real directory: {incoming_root}")
|
|
copied = 0
|
|
errors: list[str] = []
|
|
processed: set[str] = set()
|
|
try:
|
|
episodes = self.list_remote_episodes()
|
|
for episode in episodes:
|
|
processed.add(episode)
|
|
try:
|
|
if self.sync_episode(episode):
|
|
copied += 1
|
|
except Exception as exc:
|
|
# One damaged historical episode must not starve newer
|
|
# ready data. Keep its error visible and continue.
|
|
errors.append(f"{episode}: {exc}")
|
|
delete_errors = self._retry_pending_deletions(processed)
|
|
if self.delete_remote_after_sync:
|
|
pending, pending_scan_errors, pending_count = self._pending_receipts()
|
|
delete_errors.extend(pending_scan_errors)
|
|
else:
|
|
pending, pending_count = [], 0
|
|
self.pending_remote_cleanup = pending_count
|
|
if pending and not self.last_delete_error:
|
|
pending_details = [
|
|
str(receipt["remote_delete"].get("last_error", ""))
|
|
for _final, _document, receipt in pending
|
|
if receipt["remote_delete"].get("last_error")
|
|
]
|
|
self.last_delete_error = "; ".join(pending_details)
|
|
if delete_errors and not self.last_delete_error:
|
|
self.last_delete_error = "; ".join(delete_errors)
|
|
if errors:
|
|
raise RuntimeError("; ".join(errors))
|
|
self.last_error = ""
|
|
if self.pending_remote_cleanup:
|
|
self._write_status("delete_pending")
|
|
else:
|
|
self.last_delete_error = ""
|
|
self._write_status("idle")
|
|
except Exception as exc:
|
|
self.last_error = str(exc)
|
|
self._write_status("error")
|
|
raise
|
|
return copied
|
|
|
|
def run_once(self) -> int:
|
|
with self.process_lock():
|
|
return self._run_once_unlocked()
|
|
|
|
def run_forever(self, poll_seconds: float) -> None:
|
|
self.destination.mkdir(parents=True, exist_ok=True)
|
|
with self.process_lock():
|
|
while not self.stop_requested:
|
|
try:
|
|
self._run_once_unlocked()
|
|
except Exception:
|
|
pass
|
|
deadline = time.monotonic() + poll_seconds
|
|
while not self.stop_requested and time.monotonic() < deadline:
|
|
time.sleep(min(0.2, max(0.0, deadline - time.monotonic())))
|
|
self._write_status("stopped")
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--remote", default="nvidia@192.168.41.2")
|
|
parser.add_argument(
|
|
"--remote-ready", default=FIXED_REMOTE_READY
|
|
)
|
|
parser.add_argument(
|
|
"--destination",
|
|
type=Path,
|
|
default=Path(
|
|
"/home/ps/Desktop/TG3_TS1P_OmniSocket_Teleop/Data_Get"
|
|
),
|
|
)
|
|
parser.add_argument("--status-file", type=Path)
|
|
parser.add_argument("--poll-seconds", type=float, default=2.0)
|
|
parser.add_argument(
|
|
"--verify-existing",
|
|
action="store_true",
|
|
help="rehash already published local episodes (slow; intended for audits)",
|
|
)
|
|
parser.add_argument(
|
|
"--delete-remote-after-sync",
|
|
action="store_true",
|
|
help=(
|
|
"after durable local verification, use the fixed Nvidia helper to "
|
|
"remove only the matching remote ready episode"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--remote-delete-helper",
|
|
default=DEFAULT_REMOTE_DELETE_HELPER,
|
|
help="absolute path of the fixed-root deletion helper on Nvidia",
|
|
)
|
|
parser.add_argument(
|
|
"--remote-delete-timeout-seconds",
|
|
type=float,
|
|
default=600.0,
|
|
help="timeout for one remote recursive cleanup (default: 600)",
|
|
)
|
|
parser.add_argument("--once", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
if args.poll_seconds <= 0.0:
|
|
raise SystemExit("--poll-seconds must be positive")
|
|
destination = args.destination.resolve()
|
|
status_file = args.status_file or destination / "sync_status.json"
|
|
syncer = DataGetSync(
|
|
remote=args.remote,
|
|
remote_ready=args.remote_ready,
|
|
destination=destination,
|
|
status_file=status_file,
|
|
verify_existing=args.verify_existing,
|
|
delete_remote_after_sync=args.delete_remote_after_sync,
|
|
remote_delete_helper=args.remote_delete_helper,
|
|
remote_delete_timeout_s=args.remote_delete_timeout_seconds,
|
|
)
|
|
signal.signal(signal.SIGINT, syncer.request_stop)
|
|
signal.signal(signal.SIGTERM, syncer.request_stop)
|
|
if args.once:
|
|
try:
|
|
syncer.run_once()
|
|
except Exception as exc:
|
|
print(f"data sync failed: {exc}")
|
|
return 1
|
|
if syncer.pending_remote_cleanup:
|
|
print(
|
|
"data sync completed locally, but remote cleanup remains pending: "
|
|
f"{syncer.last_delete_error}"
|
|
)
|
|
return 2
|
|
return 0
|
|
syncer.run_forever(args.poll_seconds)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|