File size: 5,461 Bytes
e8b6587 | 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 151 152 153 154 155 156 157 158 159 | from __future__ import annotations
import json
import os
import zipfile
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Iterable
@dataclass(frozen=True)
class ProbeMember:
"""One file included in a returned Probe bundle."""
role: str
path: Path
arcname: str | None = None
required: bool = True
def _normalized_members(members: Iterable[ProbeMember]) -> list[ProbeMember]:
out: list[ProbeMember] = []
seen_arcnames: set[str] = set()
for member in members:
path = Path(member.path)
if not path.exists():
if member.required:
raise FileNotFoundError(f"Required Probe member is missing: {path}")
continue
if not path.is_file():
raise ValueError(f"Probe member must be a file: {path}")
arcname = str(member.arcname or path.name).strip()
if not arcname or Path(arcname).is_absolute() or ".." in Path(arcname).parts:
raise ValueError(f"Unsafe Probe archive name: {arcname!r}")
if arcname in seen_arcnames:
raise ValueError(f"Duplicate Probe archive name: {arcname}")
seen_arcnames.add(arcname)
out.append(ProbeMember(role=str(member.role), path=path, arcname=arcname, required=member.required))
return out
def build_probe_zip(
output_path: str | Path,
*,
candidate_id: str,
request_id: str,
status: str,
members: Iterable[ProbeMember],
sha256_file: Callable[[str | Path], str],
metadata: dict[str, Any] | None = None,
) -> Path:
"""Build one self-describing Probe ZIP without leaking absolute runtime paths."""
output = Path(output_path)
output.parent.mkdir(parents=True, exist_ok=True)
normalized = _normalized_members(members)
manifest_files = []
for member in normalized:
assert member.arcname is not None
manifest_files.append(
{
"role": member.role,
"name": member.arcname,
"size_bytes": int(member.path.stat().st_size),
"sha256": sha256_file(member.path),
}
)
manifest = {
"schema_version": "ltx25-probe-bundle-v1",
"candidate_id": str(candidate_id),
"request_id": str(request_id),
"status": str(status).upper(),
"created_utc": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
"files": manifest_files,
"metadata": dict(metadata or {}),
}
tmp = output.with_name(f".{output.name}.{os.getpid()}.tmp")
try:
with zipfile.ZipFile(tmp, mode="w", compression=zipfile.ZIP_DEFLATED, compresslevel=6) as archive:
archive.writestr("probe_manifest.json", json.dumps(manifest, indent=2, sort_keys=True) + "\n")
for member in normalized:
assert member.arcname is not None
archive.write(member.path, arcname=member.arcname)
tmp.replace(output)
finally:
try:
tmp.unlink(missing_ok=True)
except Exception:
pass
return output
def build_request_probe(
paths: Any,
*,
candidate_id: str,
status: str,
sha256_file: Callable[[str | Path], str],
include_run_info: bool = True,
include_diagnostics: bool = True,
include_log: bool = True,
extra_members: Iterable[ProbeMember] = (),
metadata: dict[str, Any] | None = None,
) -> Path:
"""Bundle the ordinary request evidence files owned by one RequestPaths object."""
members: list[ProbeMember] = []
if include_run_info:
members.append(ProbeMember("run_info", Path(paths.run_info), required=False))
if include_diagnostics:
members.append(ProbeMember("diagnostics", Path(paths.diagnostics), required=False))
if include_log:
members.append(ProbeMember("request_log", Path(paths.log), required=False))
members.extend(extra_members)
return build_probe_zip(
paths.probe,
candidate_id=candidate_id,
request_id=paths.request_id,
status=status,
members=members,
sha256_file=sha256_file,
metadata=metadata,
)
def append_probe_path(existing: Any, new_path: str | Path, *, limit: int | None = None) -> list[str]:
"""Append one Probe bundle path while preserving existing gr.File-style history."""
values: list[str] = []
if isinstance(existing, (list, tuple)):
values.extend(str(item) for item in existing if item)
elif existing:
values.append(str(existing))
new_value = str(new_path)
values = [item for item in values if item != new_value]
values.append(new_value)
if limit is not None and int(limit) > 0:
values = values[-int(limit) :]
return values
def list_session_probe_paths(worker_root: str | Path, session_id: str, *, limit: int | None = None) -> list[str]:
"""Return existing request Probe ZIPs for one UUID session, oldest-to-newest."""
session_root = Path(worker_root) / str(session_id or "").strip().lower()
if not session_root.is_dir():
return []
found = [path for path in session_root.glob("*/ltx25_probe_*.zip") if path.is_file()]
found.sort(key=lambda path: (path.stat().st_mtime_ns, path.name))
values = [str(path) for path in found]
if limit is not None and int(limit) > 0:
values = values[-int(limit) :]
return values
|