229 lines
9.4 KiB
Python
229 lines
9.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Bounded byte transport for the Trainlog generation namespace on rclone.
|
|
|
|
Business validation, causality, import and acknowledgement remain owned by the
|
|
existing generation services. This adapter only moves immutable bytes and
|
|
mutable run-correlated coordination objects.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
from pathlib import Path, PurePosixPath
|
|
|
|
COORDINATION = (
|
|
"android-peer-v1.json",
|
|
"android-generation-v1.json",
|
|
"android-consumption-ack-v1.json",
|
|
"android-archive-acknowledgements-v1.json",
|
|
"android-generation-error-v1.json",
|
|
"request-v1.json",
|
|
"desktop-archive-acknowledgements-v1.json",
|
|
"desktop-consumption-ack-v1.json",
|
|
"desktop-generation-v1.json",
|
|
)
|
|
MAX_FILE = 64 * 1024 * 1024
|
|
MAX_GENERATION = 256 * 1024 * 1024
|
|
|
|
|
|
class DriveTransportError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def safe_relative(value: str) -> Path:
|
|
candidate = PurePosixPath(value)
|
|
if (
|
|
not value
|
|
or value.startswith("/")
|
|
or "\\" in value
|
|
or "\x00" in value
|
|
or any(part in ("", ".", "..") for part in candidate.parts)
|
|
or len(candidate.parts) > 4
|
|
):
|
|
raise DriveTransportError("unsafe Drive object path")
|
|
return Path(*candidate.parts)
|
|
|
|
|
|
def digest(path: Path) -> str:
|
|
value = hashlib.sha256()
|
|
with path.open("rb") as stream:
|
|
while chunk := stream.read(1024 * 1024):
|
|
value.update(chunk)
|
|
return value.hexdigest()
|
|
|
|
|
|
class RcloneDrive:
|
|
def __init__(self, executable: str, remote: str):
|
|
if not remote or "\n" in remote or "\r" in remote or "Trainlog/AI" in remote:
|
|
raise DriveTransportError("invalid or AI Drive namespace")
|
|
self.executable = executable
|
|
self.remote = remote.rstrip("/")
|
|
|
|
def target(self, relative: str) -> str:
|
|
safe_relative(relative)
|
|
return f"{self.remote}/{relative}"
|
|
|
|
def run(self, *arguments: str, timeout: int = 60) -> subprocess.CompletedProcess[bytes]:
|
|
try:
|
|
return subprocess.run(
|
|
[self.executable, *arguments],
|
|
stdin=subprocess.DEVNULL,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
timeout=timeout,
|
|
check=False,
|
|
)
|
|
except (OSError, subprocess.TimeoutExpired) as error:
|
|
raise DriveTransportError(f"Drive transport unavailable: {error}") from error
|
|
|
|
def download(self, relative: str, destination: Path, optional: bool = False) -> bool:
|
|
destination.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
with tempfile.TemporaryDirectory(prefix=".trainlog-drive-pull-", dir=destination.parent) as raw:
|
|
temporary = Path(raw) / "object"
|
|
result = self.run("copyto", self.target(relative), str(temporary))
|
|
if result.returncode != 0:
|
|
message = result.stderr.decode("utf-8", "replace")[:1024]
|
|
if optional and any(token in message.lower() for token in ("not found", "doesn't exist", "directory not found")):
|
|
return False
|
|
raise DriveTransportError(message.strip() or "Drive download failed")
|
|
if not temporary.is_file() or temporary.stat().st_size > MAX_FILE:
|
|
raise DriveTransportError("Drive object exceeds local bound")
|
|
os.replace(temporary, destination)
|
|
return True
|
|
|
|
def upload_verified(self, source: Path, relative: str) -> None:
|
|
if source.is_symlink() or not source.is_file() or source.stat().st_size > MAX_FILE:
|
|
raise DriveTransportError("invalid Drive upload source")
|
|
result = self.run("copyto", str(source), self.target(relative))
|
|
if result.returncode != 0:
|
|
raise DriveTransportError(result.stderr.decode("utf-8", "replace")[:1024].strip())
|
|
with tempfile.TemporaryDirectory(prefix="trainlog-drive-verify-") as raw:
|
|
verification = Path(raw) / "object"
|
|
self.download(relative, verification)
|
|
if verification.stat().st_size != source.stat().st_size or digest(verification) != digest(source):
|
|
raise DriveTransportError("Drive upload verification failed")
|
|
|
|
|
|
def generation_from_reference(root: Path, name: str) -> str | None:
|
|
path = root / name
|
|
if not path.is_file() or path.stat().st_size > 65536:
|
|
return None
|
|
value = json.loads(path.read_text(encoding="utf-8"))
|
|
relative = value.get("relative_path")
|
|
if not isinstance(relative, str):
|
|
raise DriveTransportError("invalid generation reference")
|
|
safe = safe_relative(relative)
|
|
expected_root = (
|
|
"android-objects" if name == "android-generation-v1.json" else "desktop-objects"
|
|
)
|
|
# CONTRACT: references point only at the immutable generation namespace
|
|
# owned by their producer. The generation identity is the third path
|
|
# component; the middle component is the literal ``generations``.
|
|
if (
|
|
len(safe.parts) != 3
|
|
or safe.parts[0] != expected_root
|
|
or safe.parts[1] != "generations"
|
|
or safe.parts[2] != value.get("generation_id")
|
|
):
|
|
raise DriveTransportError("uncorrelated generation reference")
|
|
return relative
|
|
|
|
|
|
def pull_generation(drive: RcloneDrive, root: Path, relative: str) -> None:
|
|
generation_path = safe_relative(relative)
|
|
producer_root = generation_path.parts[0]
|
|
generation_relative = Path(*generation_path.parts[1:])
|
|
destination = root / relative
|
|
# WHY: polling must remain bounded even when Drive retains many immutable
|
|
# generations. CONTRACT: os.replace below makes a generation visible
|
|
# locally only after every listed artifact has downloaded. INVARIANT: a
|
|
# committed local generation is immutable and is validated by the shared
|
|
# consumer before import; an incomplete pre-existing path is never hidden.
|
|
if destination.is_dir() and (destination / "manifest.json").is_file():
|
|
return
|
|
if destination.exists():
|
|
raise DriveTransportError("incomplete local Drive generation")
|
|
manifest_relative = f"{relative}/manifest.json"
|
|
with tempfile.TemporaryDirectory(prefix=".trainlog-drive-generation-", dir=root) as raw:
|
|
staging = Path(raw) / relative
|
|
drive.download(manifest_relative, staging / "manifest.json")
|
|
manifest = json.loads((staging / "manifest.json").read_text(encoding="utf-8"))
|
|
artifacts = manifest.get("artifacts")
|
|
if not isinstance(artifacts, list) or len(artifacts) > 32:
|
|
raise DriveTransportError("invalid Drive generation manifest")
|
|
total = (staging / "manifest.json").stat().st_size
|
|
for artifact in artifacts:
|
|
filename = artifact.get("filename") if isinstance(artifact, dict) else None
|
|
if not isinstance(filename, str):
|
|
raise DriveTransportError("artifact outside Drive generation")
|
|
artifact_path = safe_relative(filename)
|
|
if artifact_path.parent != generation_relative:
|
|
raise DriveTransportError("artifact outside Drive generation")
|
|
artifact_destination = staging / artifact_path.name
|
|
drive.download(f"{producer_root}/{filename}", artifact_destination)
|
|
total += artifact_destination.stat().st_size
|
|
if total > MAX_GENERATION:
|
|
raise DriveTransportError("Drive generation exceeds bound")
|
|
destination.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
os.replace(staging, destination)
|
|
|
|
|
|
def pull(drive: RcloneDrive, root: Path) -> None:
|
|
root.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
for name in COORDINATION:
|
|
drive.download(name, root / name, optional=True)
|
|
for name in ("android-generation-v1.json", "desktop-generation-v1.json"):
|
|
relative = generation_from_reference(root, name)
|
|
if relative is not None:
|
|
pull_generation(drive, root, relative)
|
|
|
|
|
|
def ordered_files(root: Path) -> list[Path]:
|
|
files = [path for path in root.rglob("*") if path.is_file()]
|
|
if any(path.is_symlink() for path in files):
|
|
raise DriveTransportError("Drive outbox contains a symbolic link")
|
|
# CONTRACT: a generation is consumable only after its manifest exists;
|
|
# coordination references are published after the referenced directory.
|
|
return sorted(
|
|
files,
|
|
key=lambda path: (
|
|
2 if path.name.endswith("generation-v1.json") else
|
|
1 if path.name == "manifest.json" else 0,
|
|
path.relative_to(root).as_posix(),
|
|
),
|
|
)
|
|
|
|
|
|
def push(drive: RcloneDrive, root: Path) -> None:
|
|
if not root.is_dir():
|
|
raise DriveTransportError("Drive outbox is unavailable")
|
|
for source in ordered_files(root):
|
|
drive.upload_verified(source, source.relative_to(root).as_posix())
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("operation", choices=("pull", "push"))
|
|
parser.add_argument("remote")
|
|
parser.add_argument("root", type=Path)
|
|
parser.add_argument("--rclone", default=os.environ.get("TRAINLOG_RCLONE", "rclone"))
|
|
args = parser.parse_args()
|
|
drive = RcloneDrive(args.rclone, args.remote)
|
|
if args.operation == "pull":
|
|
pull(drive, args.root)
|
|
else:
|
|
push(drive, args.root)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
raise SystemExit(main())
|
|
except (DriveTransportError, json.JSONDecodeError, OSError) as error:
|
|
print(f"Drive transport failed: {error}", file=os.sys.stderr)
|
|
raise SystemExit(1) from None
|