1484 lines
59 KiB
Python
1484 lines
59 KiB
Python
import base64
|
|
import gzip
|
|
import hashlib
|
|
import io
|
|
import ipaddress
|
|
import json
|
|
import mimetypes
|
|
import os
|
|
import re
|
|
import socket
|
|
import sqlite3
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
import uuid
|
|
from contextlib import closing, contextmanager
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Callable, Optional
|
|
from zoneinfo import ZoneInfo
|
|
|
|
|
|
PENDING = "pending"
|
|
UPLOADING = "uploading"
|
|
UPLOADED = "uploaded"
|
|
CLOUD_COMPLETED = "cloud_completed"
|
|
CLOUD_FAILED = "cloud_failed"
|
|
NO_ARTIFACTS = "no_artifacts"
|
|
SOURCE_MISSING = "source_missing"
|
|
TERMINAL_ARTIFACT_STATES = (UPLOADED, CLOUD_FAILED, SOURCE_MISSING)
|
|
RETRY_DELAYS = (30, 120)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CloudArchiveConfig:
|
|
enabled: bool
|
|
server_id: str
|
|
region: str
|
|
endpoint: str
|
|
bucket: str
|
|
prefix: str
|
|
public_base_url: str
|
|
max_attempts: int = 3
|
|
upload_concurrency: int = 2
|
|
remote_max_bytes: int = 20 * 1024 * 1024 * 1024
|
|
manifest_max_bytes: int = 2 * 1024 * 1024 * 1024
|
|
staging_dir: str = "cloud_staging"
|
|
checkpoint_dir: str = "cloud_upload_checkpoints"
|
|
|
|
def validate(self) -> None:
|
|
if not self.enabled:
|
|
return
|
|
required = {
|
|
"cloud_server_id": self.server_id,
|
|
"cloud_oss_region": self.region,
|
|
"cloud_oss_endpoint": self.endpoint,
|
|
"cloud_oss_bucket": self.bucket,
|
|
"cloud_oss_prefix": self.prefix,
|
|
"cloud_public_base_url": self.public_base_url,
|
|
}
|
|
missing = [name for name, value in required.items() if not str(value or "").strip()]
|
|
if missing:
|
|
raise ValueError("Missing cloud archive settings: " + ", ".join(missing))
|
|
if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,63}", self.server_id):
|
|
raise ValueError("cloud_server_id contains unsupported characters")
|
|
|
|
def media_key(self, prompt_id: str, ordinal: int, filename: str, finished_at: float) -> str:
|
|
date = datetime.fromtimestamp(float(finished_at or time.time()), ZoneInfo("Asia/Shanghai"))
|
|
safe_name = sanitize_filename(filename)
|
|
return (
|
|
f"{self.prefix}/public/servers/{self.server_id}/"
|
|
f"{date:%Y/%m/%d}/{prompt_id}/{int(ordinal):03d}-{safe_name}"
|
|
)
|
|
|
|
def manifest_key(self, prompt_id: str) -> str:
|
|
return f"{self.prefix}/private/jobs/{prompt_id}/manifest.json.gz"
|
|
|
|
def oss_uri(self, key: str) -> str:
|
|
return f"oss://{self.bucket}/{key}"
|
|
|
|
def public_url(self, key: str) -> str:
|
|
return f"{self.public_base_url}/{urllib.parse.quote(key, safe='/')}"
|
|
|
|
|
|
def sanitize_filename(filename: str) -> str:
|
|
name = os.path.basename(str(filename or "artifact"))
|
|
name = re.sub(r"[\x00-\x1f\x7f/\\]", "_", name).strip(" .")
|
|
if not name:
|
|
name = "artifact"
|
|
if len(name.encode("utf-8")) <= 220:
|
|
return name
|
|
stem, suffix = os.path.splitext(name)
|
|
suffix = suffix[:32]
|
|
while stem and len((stem + suffix).encode("utf-8")) > 220:
|
|
stem = stem[:-1]
|
|
return (stem or "artifact") + suffix
|
|
|
|
|
|
def _safe_join(root: str, subfolder: str, filename: str) -> Optional[str]:
|
|
root_path = os.path.realpath(root)
|
|
candidate = os.path.realpath(os.path.join(root_path, str(subfolder or ""), filename))
|
|
try:
|
|
if os.path.commonpath((root_path, candidate)) != root_path:
|
|
return None
|
|
except ValueError:
|
|
return None
|
|
return candidate
|
|
|
|
|
|
def _walk_outputs(value, visit: Callable[[object, str], None], path: str = "outputs") -> None:
|
|
visit(value, path)
|
|
if isinstance(value, dict):
|
|
for key, child in value.items():
|
|
_walk_outputs(child, visit, f"{path}/{key}")
|
|
elif isinstance(value, (list, tuple)):
|
|
for index, child in enumerate(value):
|
|
_walk_outputs(child, visit, f"{path}/{index}")
|
|
|
|
|
|
def discover_artifacts(history_item: dict, output_dir: str, temp_dir: str) -> list[dict]:
|
|
discovered = []
|
|
seen = set()
|
|
|
|
def add_local(path: str, filename: str, reference: dict, output_path: str) -> None:
|
|
if not path:
|
|
return
|
|
normalized = os.path.realpath(path)
|
|
key = ("local", os.path.normcase(normalized))
|
|
if key in seen:
|
|
return
|
|
seen.add(key)
|
|
discovered.append(
|
|
{
|
|
"source_kind": "local_output",
|
|
"source_path": normalized,
|
|
"source_url": None,
|
|
"filename": sanitize_filename(filename or os.path.basename(normalized)),
|
|
"mime_type": mimetypes.guess_type(filename or normalized)[0],
|
|
"reference": reference,
|
|
"output_path": output_path,
|
|
}
|
|
)
|
|
|
|
def visit(value, output_path: str) -> None:
|
|
if isinstance(value, dict):
|
|
filename = str(value.get("filename") or "")
|
|
explicit = str(value.get("fullpath") or value.get("file_path") or "")
|
|
if explicit:
|
|
add_local(explicit, filename or os.path.basename(explicit), value, output_path)
|
|
elif filename and filename == os.path.basename(filename):
|
|
kind = str(value.get("type") or "output").lower()
|
|
root = temp_dir if kind == "temp" else output_dir
|
|
candidate = _safe_join(root, str(value.get("subfolder") or ""), filename)
|
|
if candidate:
|
|
add_local(candidate, filename, value, output_path)
|
|
elif isinstance(value, str) and value.startswith(("http://", "https://")):
|
|
normalized = value
|
|
key = ("remote", normalized)
|
|
if key in seen:
|
|
return
|
|
parsed = urllib.parse.urlsplit(value)
|
|
filename = sanitize_filename(os.path.basename(parsed.path) or "remote-artifact")
|
|
seen.add(key)
|
|
discovered.append(
|
|
{
|
|
"source_kind": "remote_url",
|
|
"source_path": None,
|
|
"source_url": normalized,
|
|
"filename": filename,
|
|
"mime_type": mimetypes.guess_type(filename)[0],
|
|
"reference": None,
|
|
"output_path": output_path,
|
|
}
|
|
)
|
|
|
|
_walk_outputs((history_item or {}).get("outputs", {}), visit)
|
|
return discovered
|
|
|
|
|
|
def _blob(value: bytes) -> dict:
|
|
raw = bytes(value)
|
|
return {
|
|
"$sqlite_type": "blob",
|
|
"encoding": "base64",
|
|
"byte_length": len(raw),
|
|
"sha256": hashlib.sha256(raw).hexdigest(),
|
|
"data": base64.b64encode(raw).decode("ascii"),
|
|
}
|
|
|
|
|
|
def json_value(value):
|
|
if isinstance(value, (bytes, bytearray, memoryview)):
|
|
return _blob(bytes(value))
|
|
return value
|
|
|
|
|
|
def row_dict(row) -> Optional[dict]:
|
|
if row is None:
|
|
return None
|
|
return {key: json_value(row[key]) for key in row.keys()}
|
|
|
|
|
|
def _asset_ids(value) -> set[str]:
|
|
found = set()
|
|
if isinstance(value, dict):
|
|
asset_id = str(value.get("id") or "")
|
|
if asset_id:
|
|
found.add(asset_id)
|
|
for child in value.values():
|
|
found.update(_asset_ids(child))
|
|
elif isinstance(value, list):
|
|
for child in value:
|
|
found.update(_asset_ids(child))
|
|
return found
|
|
|
|
|
|
class CloudArchiveStore:
|
|
def __init__(self, database: str, busy_timeout_ms: int = 15000):
|
|
self.database = os.fspath(database)
|
|
self.busy_timeout_ms = max(1000, int(busy_timeout_ms))
|
|
self._initialize()
|
|
|
|
def _connect(self, readonly: bool = False) -> sqlite3.Connection:
|
|
database = self.database
|
|
if readonly:
|
|
database = f"file:{os.path.abspath(database)}?mode=ro"
|
|
connection = sqlite3.connect(
|
|
database,
|
|
timeout=self.busy_timeout_ms / 1000,
|
|
isolation_level=None,
|
|
uri=readonly,
|
|
)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute(f"PRAGMA busy_timeout={self.busy_timeout_ms}")
|
|
return connection
|
|
|
|
def _initialize(self) -> None:
|
|
with closing(self._connect()) as connection:
|
|
connection.execute("PRAGMA journal_mode=WAL")
|
|
connection.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS scheduler_cloud_tasks (
|
|
prompt_id TEXT PRIMARY KEY,
|
|
server_id TEXT NOT NULL,
|
|
owner_id TEXT,
|
|
worker_port INTEGER,
|
|
generation_status TEXT NOT NULL,
|
|
cloud_status TEXT NOT NULL,
|
|
artifact_count INTEGER NOT NULL DEFAULT 0,
|
|
uploaded_count INTEGER NOT NULL DEFAULT 0,
|
|
failed_count INTEGER NOT NULL DEFAULT 0,
|
|
manifest_object_key TEXT,
|
|
manifest_oss_uri TEXT,
|
|
manifest_size_bytes INTEGER,
|
|
manifest_sha256 TEXT,
|
|
manifest_upload_status TEXT NOT NULL DEFAULT 'pending',
|
|
manifest_attempt_count INTEGER NOT NULL DEFAULT 0,
|
|
manifest_next_attempt_at REAL,
|
|
manifest_lease_owner TEXT,
|
|
manifest_lease_expires_at REAL,
|
|
last_error TEXT,
|
|
created_at REAL NOT NULL,
|
|
updated_at REAL NOT NULL,
|
|
completed_at REAL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS scheduler_cloud_tasks_status
|
|
ON scheduler_cloud_tasks(cloud_status, updated_at);
|
|
|
|
CREATE TABLE IF NOT EXISTS scheduler_cloud_artifacts (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
prompt_id TEXT NOT NULL,
|
|
server_id TEXT NOT NULL,
|
|
ordinal INTEGER NOT NULL,
|
|
source_kind TEXT NOT NULL,
|
|
source_path TEXT,
|
|
source_url TEXT,
|
|
output_path TEXT,
|
|
reference_data TEXT,
|
|
filename TEXT NOT NULL,
|
|
mime_type TEXT,
|
|
size_bytes INTEGER,
|
|
sha256 TEXT,
|
|
oss_object_key TEXT NOT NULL,
|
|
oss_uri TEXT,
|
|
public_url TEXT,
|
|
etag TEXT,
|
|
oss_crc64 TEXT,
|
|
upload_status TEXT NOT NULL DEFAULT 'pending',
|
|
attempt_count INTEGER NOT NULL DEFAULT 0,
|
|
next_attempt_at REAL,
|
|
lease_owner TEXT,
|
|
lease_expires_at REAL,
|
|
last_error TEXT,
|
|
created_at REAL NOT NULL,
|
|
updated_at REAL NOT NULL,
|
|
uploaded_at REAL,
|
|
UNIQUE(prompt_id, ordinal)
|
|
);
|
|
CREATE INDEX IF NOT EXISTS scheduler_cloud_artifacts_claim
|
|
ON scheduler_cloud_artifacts(upload_status, next_attempt_at, id);
|
|
CREATE INDEX IF NOT EXISTS scheduler_cloud_artifacts_prompt
|
|
ON scheduler_cloud_artifacts(prompt_id, ordinal);
|
|
"""
|
|
)
|
|
|
|
def enqueue_task(
|
|
self,
|
|
config: CloudArchiveConfig,
|
|
prompt_id: str,
|
|
owner_id: str,
|
|
worker_port: int,
|
|
generation_status: str,
|
|
finished_at: float,
|
|
artifacts: list[dict],
|
|
) -> None:
|
|
now = time.time()
|
|
manifest_key = config.manifest_key(prompt_id)
|
|
with closing(self._connect()) as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
existing = connection.execute(
|
|
"SELECT server_id FROM scheduler_cloud_tasks WHERE prompt_id = ?",
|
|
(prompt_id,),
|
|
).fetchone()
|
|
if existing and existing["server_id"] != config.server_id:
|
|
raise RuntimeError("prompt_id already belongs to another cloud server_id")
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO scheduler_cloud_tasks(
|
|
prompt_id, server_id, owner_id, worker_port, generation_status,
|
|
cloud_status, artifact_count, manifest_object_key,
|
|
manifest_oss_uri, manifest_upload_status,
|
|
manifest_next_attempt_at, created_at, updated_at
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(prompt_id) DO UPDATE SET
|
|
owner_id=excluded.owner_id,
|
|
worker_port=excluded.worker_port,
|
|
generation_status=excluded.generation_status,
|
|
artifact_count=excluded.artifact_count,
|
|
manifest_object_key=excluded.manifest_object_key,
|
|
manifest_oss_uri=excluded.manifest_oss_uri,
|
|
updated_at=excluded.updated_at
|
|
""",
|
|
(
|
|
prompt_id,
|
|
config.server_id,
|
|
owner_id or "",
|
|
int(worker_port or 0),
|
|
generation_status,
|
|
PENDING if artifacts else NO_ARTIFACTS,
|
|
len(artifacts),
|
|
manifest_key,
|
|
config.oss_uri(manifest_key),
|
|
PENDING,
|
|
now,
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
for ordinal, artifact in enumerate(artifacts, 1):
|
|
key = config.media_key(
|
|
prompt_id, ordinal, artifact["filename"], finished_at
|
|
)
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO scheduler_cloud_artifacts(
|
|
prompt_id, server_id, ordinal, source_kind, source_path,
|
|
source_url, output_path, reference_data, filename, mime_type,
|
|
oss_object_key, oss_uri, public_url, upload_status,
|
|
next_attempt_at, created_at, updated_at
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(prompt_id, ordinal) DO UPDATE SET
|
|
source_path=excluded.source_path,
|
|
source_url=excluded.source_url,
|
|
output_path=excluded.output_path,
|
|
reference_data=excluded.reference_data,
|
|
filename=excluded.filename,
|
|
mime_type=excluded.mime_type,
|
|
oss_object_key=excluded.oss_object_key,
|
|
oss_uri=excluded.oss_uri,
|
|
public_url=excluded.public_url,
|
|
updated_at=excluded.updated_at
|
|
""",
|
|
(
|
|
prompt_id,
|
|
config.server_id,
|
|
ordinal,
|
|
artifact["source_kind"],
|
|
artifact.get("source_path"),
|
|
artifact.get("source_url"),
|
|
artifact.get("output_path"),
|
|
json.dumps(artifact.get("reference"), ensure_ascii=False),
|
|
artifact["filename"],
|
|
artifact.get("mime_type"),
|
|
key,
|
|
config.oss_uri(key),
|
|
config.public_url(key),
|
|
PENDING,
|
|
now,
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
connection.execute("COMMIT")
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def claim_artifact(self, claimant: str, concurrency: int, lease_seconds: int = 1800):
|
|
now = time.time()
|
|
with closing(self._connect()) as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
active_artifacts = connection.execute(
|
|
"""
|
|
SELECT COUNT(*) FROM scheduler_cloud_artifacts
|
|
WHERE upload_status = ? AND lease_expires_at > ?
|
|
""",
|
|
(UPLOADING, now),
|
|
).fetchone()[0]
|
|
active_manifests = connection.execute(
|
|
"""
|
|
SELECT COUNT(*) FROM scheduler_cloud_tasks
|
|
WHERE manifest_upload_status = ? AND manifest_lease_expires_at > ?
|
|
""",
|
|
(UPLOADING, now),
|
|
).fetchone()[0]
|
|
if int(active_artifacts) + int(active_manifests) >= int(concurrency):
|
|
connection.execute("COMMIT")
|
|
return None
|
|
row = connection.execute(
|
|
"""
|
|
SELECT * FROM scheduler_cloud_artifacts
|
|
WHERE (
|
|
upload_status = ?
|
|
OR (upload_status = ? AND COALESCE(lease_expires_at, 0) <= ?)
|
|
)
|
|
AND COALESCE(next_attempt_at, 0) <= ?
|
|
ORDER BY id
|
|
LIMIT 1
|
|
""",
|
|
(PENDING, UPLOADING, now, now),
|
|
).fetchone()
|
|
if not row:
|
|
connection.execute("COMMIT")
|
|
return None
|
|
connection.execute(
|
|
"""
|
|
UPDATE scheduler_cloud_artifacts
|
|
SET upload_status=?, attempt_count=attempt_count+1,
|
|
lease_owner=?, lease_expires_at=?, updated_at=?
|
|
WHERE id=?
|
|
""",
|
|
(UPLOADING, claimant, now + lease_seconds, now, row["id"]),
|
|
)
|
|
connection.execute(
|
|
"UPDATE scheduler_cloud_tasks SET cloud_status=?, updated_at=? WHERE prompt_id=?",
|
|
(UPLOADING, now, row["prompt_id"]),
|
|
)
|
|
claimed = connection.execute(
|
|
"SELECT * FROM scheduler_cloud_artifacts WHERE id=?", (row["id"],)
|
|
).fetchone()
|
|
connection.execute("COMMIT")
|
|
return dict(claimed)
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def artifact_succeeded(self, artifact_id: int, result: dict) -> None:
|
|
now = time.time()
|
|
with closing(self._connect()) as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
row = connection.execute(
|
|
"SELECT prompt_id FROM scheduler_cloud_artifacts WHERE id=?",
|
|
(artifact_id,),
|
|
).fetchone()
|
|
if not row:
|
|
return
|
|
connection.execute(
|
|
"""
|
|
UPDATE scheduler_cloud_artifacts
|
|
SET upload_status=?, mime_type=?, size_bytes=?, sha256=?, etag=?, oss_crc64=?,
|
|
lease_owner=NULL, lease_expires_at=NULL, next_attempt_at=NULL,
|
|
last_error=NULL, uploaded_at=?, updated_at=?
|
|
WHERE id=?
|
|
""",
|
|
(
|
|
UPLOADED,
|
|
result.get("mime_type"),
|
|
result.get("size_bytes"),
|
|
result.get("sha256"),
|
|
result.get("etag"),
|
|
result.get("oss_crc64"),
|
|
now,
|
|
now,
|
|
artifact_id,
|
|
),
|
|
)
|
|
self._refresh_counts(connection, row["prompt_id"], now)
|
|
connection.execute("COMMIT")
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def artifact_failed(self, artifact_id: int, error: str, max_attempts: int) -> None:
|
|
now = time.time()
|
|
with closing(self._connect()) as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
row = connection.execute(
|
|
"SELECT prompt_id, attempt_count FROM scheduler_cloud_artifacts WHERE id=?",
|
|
(artifact_id,),
|
|
).fetchone()
|
|
if not row:
|
|
return
|
|
attempts = int(row["attempt_count"])
|
|
terminal = attempts >= int(max_attempts)
|
|
delay = RETRY_DELAYS[min(max(0, attempts - 1), len(RETRY_DELAYS) - 1)]
|
|
connection.execute(
|
|
"""
|
|
UPDATE scheduler_cloud_artifacts
|
|
SET upload_status=?, next_attempt_at=?, lease_owner=NULL,
|
|
lease_expires_at=NULL, last_error=?, updated_at=?
|
|
WHERE id=?
|
|
""",
|
|
(
|
|
CLOUD_FAILED if terminal else PENDING,
|
|
None if terminal else now + delay,
|
|
str(error),
|
|
now,
|
|
artifact_id,
|
|
),
|
|
)
|
|
self._refresh_counts(connection, row["prompt_id"], now)
|
|
connection.execute("COMMIT")
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def artifact_source_missing(self, artifact_id: int, error: str) -> None:
|
|
now = time.time()
|
|
with closing(self._connect()) as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
row = connection.execute(
|
|
"SELECT prompt_id FROM scheduler_cloud_artifacts WHERE id=?",
|
|
(artifact_id,),
|
|
).fetchone()
|
|
if not row:
|
|
connection.execute("COMMIT")
|
|
return
|
|
connection.execute(
|
|
"""
|
|
UPDATE scheduler_cloud_artifacts
|
|
SET upload_status=?, next_attempt_at=NULL, lease_owner=NULL,
|
|
lease_expires_at=NULL, last_error=?, updated_at=?
|
|
WHERE id=?
|
|
""",
|
|
(SOURCE_MISSING, str(error), now, artifact_id),
|
|
)
|
|
self._refresh_counts(connection, row["prompt_id"], now)
|
|
connection.execute("COMMIT")
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def renew_artifact_lease(
|
|
self, artifact_id: int, claimant: str, lease_seconds: int = 1800
|
|
) -> bool:
|
|
now = time.time()
|
|
with closing(self._connect()) as connection:
|
|
changed = connection.execute(
|
|
"""
|
|
UPDATE scheduler_cloud_artifacts
|
|
SET lease_expires_at=?, updated_at=?
|
|
WHERE id=? AND upload_status=? AND lease_owner=?
|
|
""",
|
|
(now + lease_seconds, now, artifact_id, UPLOADING, claimant),
|
|
).rowcount
|
|
return bool(changed)
|
|
|
|
@staticmethod
|
|
def _refresh_counts(connection, prompt_id: str, now: float) -> None:
|
|
counts = connection.execute(
|
|
"""
|
|
SELECT COUNT(*) total,
|
|
SUM(CASE WHEN upload_status='uploaded' THEN 1 ELSE 0 END) uploaded,
|
|
SUM(CASE WHEN upload_status IN ('cloud_failed','source_missing') THEN 1 ELSE 0 END) failed
|
|
FROM scheduler_cloud_artifacts WHERE prompt_id=?
|
|
""",
|
|
(prompt_id,),
|
|
).fetchone()
|
|
connection.execute(
|
|
"""
|
|
UPDATE scheduler_cloud_tasks
|
|
SET artifact_count=?, uploaded_count=?, failed_count=?, updated_at=?
|
|
WHERE prompt_id=?
|
|
""",
|
|
(
|
|
int(counts["total"] or 0),
|
|
int(counts["uploaded"] or 0),
|
|
int(counts["failed"] or 0),
|
|
now,
|
|
prompt_id,
|
|
),
|
|
)
|
|
|
|
def claim_manifest(
|
|
self,
|
|
claimant: str,
|
|
max_attempts: int,
|
|
concurrency: int,
|
|
lease_seconds: int = 1800,
|
|
):
|
|
now = time.time()
|
|
with closing(self._connect()) as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
active_artifacts = connection.execute(
|
|
"""
|
|
SELECT COUNT(*) FROM scheduler_cloud_artifacts
|
|
WHERE upload_status = ? AND lease_expires_at > ?
|
|
""",
|
|
(UPLOADING, now),
|
|
).fetchone()[0]
|
|
active_manifests = connection.execute(
|
|
"""
|
|
SELECT COUNT(*) FROM scheduler_cloud_tasks
|
|
WHERE manifest_upload_status = ? AND manifest_lease_expires_at > ?
|
|
""",
|
|
(UPLOADING, now),
|
|
).fetchone()[0]
|
|
if int(active_artifacts) + int(active_manifests) >= int(concurrency):
|
|
connection.execute("COMMIT")
|
|
return None
|
|
row = connection.execute(
|
|
"""
|
|
SELECT task.*
|
|
FROM scheduler_cloud_tasks task
|
|
WHERE (
|
|
task.manifest_upload_status = ?
|
|
OR (task.manifest_upload_status = ?
|
|
AND COALESCE(task.manifest_lease_expires_at, 0) <= ?)
|
|
)
|
|
AND task.manifest_attempt_count < ?
|
|
AND COALESCE(task.manifest_next_attempt_at, 0) <= ?
|
|
AND NOT EXISTS (
|
|
SELECT 1 FROM scheduler_cloud_artifacts artifact
|
|
WHERE artifact.prompt_id=task.prompt_id
|
|
AND artifact.upload_status NOT IN ('uploaded','cloud_failed','source_missing')
|
|
)
|
|
ORDER BY task.created_at
|
|
LIMIT 1
|
|
""",
|
|
(PENDING, UPLOADING, now, int(max_attempts), now),
|
|
).fetchone()
|
|
if not row:
|
|
connection.execute("COMMIT")
|
|
return None
|
|
connection.execute(
|
|
"""
|
|
UPDATE scheduler_cloud_tasks
|
|
SET manifest_upload_status=?, manifest_attempt_count=manifest_attempt_count+1,
|
|
manifest_lease_owner=?, manifest_lease_expires_at=?, updated_at=?
|
|
WHERE prompt_id=?
|
|
""",
|
|
(UPLOADING, claimant, now + lease_seconds, now, row["prompt_id"]),
|
|
)
|
|
claimed = connection.execute(
|
|
"SELECT * FROM scheduler_cloud_tasks WHERE prompt_id=?",
|
|
(row["prompt_id"],),
|
|
).fetchone()
|
|
connection.execute("COMMIT")
|
|
return dict(claimed)
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def manifest_succeeded(self, prompt_id: str, result: dict) -> None:
|
|
now = time.time()
|
|
with closing(self._connect()) as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
task = connection.execute(
|
|
"SELECT artifact_count, failed_count FROM scheduler_cloud_tasks WHERE prompt_id=?",
|
|
(prompt_id,),
|
|
).fetchone()
|
|
if not task:
|
|
return
|
|
if int(task["artifact_count"] or 0) == 0:
|
|
status = NO_ARTIFACTS
|
|
elif int(task["failed_count"] or 0) > 0:
|
|
status = CLOUD_FAILED
|
|
else:
|
|
status = CLOUD_COMPLETED
|
|
connection.execute(
|
|
"""
|
|
UPDATE scheduler_cloud_tasks
|
|
SET cloud_status=?, manifest_upload_status=?, manifest_size_bytes=?,
|
|
manifest_sha256=?, manifest_next_attempt_at=NULL,
|
|
manifest_lease_owner=NULL, manifest_lease_expires_at=NULL,
|
|
last_error=NULL, completed_at=?, updated_at=?
|
|
WHERE prompt_id=?
|
|
""",
|
|
(
|
|
status,
|
|
UPLOADED,
|
|
result.get("size_bytes"),
|
|
result.get("sha256"),
|
|
now,
|
|
now,
|
|
prompt_id,
|
|
),
|
|
)
|
|
connection.execute("COMMIT")
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def manifest_failed(self, prompt_id: str, error: str, max_attempts: int) -> None:
|
|
now = time.time()
|
|
with closing(self._connect()) as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
task = connection.execute(
|
|
"SELECT manifest_attempt_count FROM scheduler_cloud_tasks WHERE prompt_id=?",
|
|
(prompt_id,),
|
|
).fetchone()
|
|
if not task:
|
|
return
|
|
attempts = int(task["manifest_attempt_count"])
|
|
terminal = attempts >= int(max_attempts)
|
|
delay = RETRY_DELAYS[min(max(0, attempts - 1), len(RETRY_DELAYS) - 1)]
|
|
connection.execute(
|
|
"""
|
|
UPDATE scheduler_cloud_tasks
|
|
SET cloud_status=?, manifest_upload_status=?, manifest_next_attempt_at=?,
|
|
manifest_lease_owner=NULL, manifest_lease_expires_at=NULL,
|
|
last_error=?, updated_at=?, completed_at=?
|
|
WHERE prompt_id=?
|
|
""",
|
|
(
|
|
CLOUD_FAILED if terminal else UPLOADING,
|
|
"failed" if terminal else PENDING,
|
|
None if terminal else now + delay,
|
|
str(error),
|
|
now,
|
|
now if terminal else None,
|
|
prompt_id,
|
|
),
|
|
)
|
|
connection.execute("COMMIT")
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def renew_manifest_lease(
|
|
self, prompt_id: str, claimant: str, lease_seconds: int = 1800
|
|
) -> bool:
|
|
now = time.time()
|
|
with closing(self._connect()) as connection:
|
|
changed = connection.execute(
|
|
"""
|
|
UPDATE scheduler_cloud_tasks
|
|
SET manifest_lease_expires_at=?, updated_at=?
|
|
WHERE prompt_id=? AND manifest_upload_status=?
|
|
AND manifest_lease_owner=?
|
|
""",
|
|
(now + lease_seconds, now, prompt_id, UPLOADING, claimant),
|
|
).rowcount
|
|
return bool(changed)
|
|
|
|
def task(self, prompt_id: str) -> Optional[dict]:
|
|
with closing(self._connect(readonly=True)) as connection:
|
|
row = connection.execute(
|
|
"SELECT * FROM scheduler_cloud_tasks WHERE prompt_id=?", (prompt_id,)
|
|
).fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def artifacts(self, prompt_id: str) -> list[dict]:
|
|
with closing(self._connect(readonly=True)) as connection:
|
|
rows = connection.execute(
|
|
"SELECT * FROM scheduler_cloud_artifacts WHERE prompt_id=? ORDER BY ordinal",
|
|
(prompt_id,),
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
def reset_failed(self, prompt_id: str) -> bool:
|
|
now = time.time()
|
|
with closing(self._connect()) as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
changed = connection.execute(
|
|
"""
|
|
UPDATE scheduler_cloud_artifacts
|
|
SET upload_status=?, attempt_count=0, next_attempt_at=?,
|
|
lease_owner=NULL, lease_expires_at=NULL, last_error=NULL, updated_at=?
|
|
WHERE prompt_id=? AND upload_status IN (?, ?)
|
|
""",
|
|
(PENDING, now, now, prompt_id, CLOUD_FAILED, SOURCE_MISSING),
|
|
).rowcount
|
|
task_changed = connection.execute(
|
|
"""
|
|
UPDATE scheduler_cloud_tasks
|
|
SET cloud_status=?, manifest_upload_status=?, manifest_attempt_count=0,
|
|
manifest_next_attempt_at=?, manifest_lease_owner=NULL,
|
|
manifest_lease_expires_at=NULL, last_error=NULL, completed_at=NULL,
|
|
updated_at=? WHERE prompt_id=?
|
|
""",
|
|
(PENDING, PENDING, now, now, prompt_id),
|
|
).rowcount
|
|
connection.execute("COMMIT")
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
return bool(changed or task_changed)
|
|
|
|
|
|
class OssV2Backend:
|
|
def __init__(self, config: CloudArchiveConfig):
|
|
config.validate()
|
|
try:
|
|
import alibabacloud_oss_v2 as oss
|
|
except ImportError as error:
|
|
raise RuntimeError("alibabacloud-oss-v2 is not installed") from error
|
|
provider = oss.credentials.EnvironmentVariableCredentialsProvider()
|
|
sdk_config = oss.Config(
|
|
region=config.region,
|
|
endpoint=config.endpoint,
|
|
signature_version="v4",
|
|
credentials_provider=provider,
|
|
retry_max_attempts=3,
|
|
connect_timeout=10,
|
|
readwrite_timeout=120,
|
|
)
|
|
self.oss = oss
|
|
self.client = oss.Client(sdk_config)
|
|
os.makedirs(config.checkpoint_dir, exist_ok=True)
|
|
self.uploader = self.client.uploader(
|
|
parallel_num=3,
|
|
enable_checkpoint=True,
|
|
checkpoint_dir=config.checkpoint_dir,
|
|
)
|
|
self.config = config
|
|
|
|
def upload_file(
|
|
self,
|
|
key: str,
|
|
path: str,
|
|
content_type: str,
|
|
content_encoding: str = None,
|
|
cache_control: str = None,
|
|
sha256: str = None,
|
|
private: bool = False,
|
|
) -> dict:
|
|
metadata = {"server-id": self.config.server_id}
|
|
if sha256:
|
|
metadata["sha256"] = sha256
|
|
request = self.oss.PutObjectRequest(
|
|
bucket=self.config.bucket,
|
|
key=key,
|
|
content_type=content_type or "application/octet-stream",
|
|
content_encoding=content_encoding,
|
|
cache_control=cache_control,
|
|
metadata=metadata,
|
|
object_acl="private" if private else None,
|
|
)
|
|
result = self.uploader.upload_file(request, path)
|
|
head = self.client.head_object(
|
|
self.oss.HeadObjectRequest(bucket=self.config.bucket, key=key)
|
|
)
|
|
expected = os.path.getsize(path)
|
|
if int(head.content_length or -1) != expected:
|
|
raise RuntimeError(
|
|
f"OSS object length mismatch: expected {expected}, got {head.content_length}"
|
|
)
|
|
return {
|
|
"etag": getattr(result, "etag", None) or getattr(head, "etag", None),
|
|
"oss_crc64": getattr(result, "hash_crc64", None)
|
|
or getattr(head, "hash_crc64", None),
|
|
"content_length": int(head.content_length),
|
|
}
|
|
|
|
|
|
def _validate_public_http_url(url: str) -> None:
|
|
parsed = urllib.parse.urlsplit(url)
|
|
if parsed.scheme not in ("http", "https") or not parsed.hostname:
|
|
raise ValueError("Only absolute HTTP(S) artifact URLs are supported")
|
|
try:
|
|
addresses = socket.getaddrinfo(parsed.hostname, parsed.port or 443, type=socket.SOCK_STREAM)
|
|
except socket.gaierror as error:
|
|
raise ValueError(f"Unable to resolve artifact host: {error}") from error
|
|
for address in addresses:
|
|
ip = ipaddress.ip_address(address[4][0])
|
|
if not ip.is_global:
|
|
raise ValueError("Remote artifact URL resolves to a non-public address")
|
|
|
|
|
|
class _SafeRedirectHandler(urllib.request.HTTPRedirectHandler):
|
|
def redirect_request(self, req, fp, code, msg, headers, newurl):
|
|
_validate_public_http_url(newurl)
|
|
return super().redirect_request(req, fp, code, msg, headers, newurl)
|
|
|
|
|
|
def download_remote(url: str, destination: str, max_bytes: int) -> tuple[str, int]:
|
|
_validate_public_http_url(url)
|
|
opener = urllib.request.build_opener(_SafeRedirectHandler())
|
|
request = urllib.request.Request(url, headers={"User-Agent": "ComfyUI-Account-Manager/3"})
|
|
os.makedirs(os.path.dirname(destination), exist_ok=True)
|
|
with opener.open(request, timeout=60) as response:
|
|
content_type = str(response.headers.get("Content-Type") or "").split(";", 1)[0]
|
|
allowed = content_type.startswith(("image/", "video/", "audio/")) or content_type in (
|
|
"application/octet-stream",
|
|
"application/zip",
|
|
)
|
|
if content_type and not allowed:
|
|
raise ValueError(f"Remote artifact content type is not media: {content_type}")
|
|
declared = int(response.headers.get("Content-Length") or 0)
|
|
if declared > max_bytes:
|
|
raise ValueError("Remote artifact exceeds cloud_remote_max_bytes")
|
|
written = 0
|
|
with open(destination, "wb") as output:
|
|
while True:
|
|
chunk = response.read(1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
written += len(chunk)
|
|
if written > max_bytes:
|
|
raise ValueError("Remote artifact exceeds cloud_remote_max_bytes")
|
|
output.write(chunk)
|
|
return content_type or "application/octet-stream", written
|
|
|
|
|
|
def file_sha256(path: str) -> str:
|
|
digest = hashlib.sha256()
|
|
with open(path, "rb") as source:
|
|
while True:
|
|
chunk = source.read(4 * 1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
class ManifestExporter:
|
|
SCHEDULER_TABLES = (
|
|
"scheduler_jobs",
|
|
"scheduler_api_logs",
|
|
"scheduler_assets",
|
|
"scheduler_job_logs",
|
|
"scheduler_workers",
|
|
"scheduler_cloud_tasks",
|
|
"scheduler_cloud_artifacts",
|
|
)
|
|
|
|
def __init__(self, scheduler_database: str, history_database: str, project_dir: str):
|
|
self.scheduler_database = scheduler_database
|
|
self.history_database = history_database
|
|
self.project_dir = project_dir
|
|
|
|
@staticmethod
|
|
def _schema(connection, table: str) -> dict:
|
|
sql_row = connection.execute(
|
|
"SELECT sql FROM sqlite_master WHERE type='table' AND name=?", (table,)
|
|
).fetchone()
|
|
columns = connection.execute(f"PRAGMA table_info({table})").fetchall()
|
|
return {
|
|
"create_sql": sql_row[0] if sql_row else None,
|
|
"columns": [row_dict(row) for row in columns],
|
|
}
|
|
|
|
def _commit(self) -> str:
|
|
try:
|
|
return subprocess.check_output(
|
|
["git", "-C", self.project_dir, "rev-parse", "--short", "HEAD"],
|
|
text=True,
|
|
timeout=3,
|
|
stderr=subprocess.DEVNULL,
|
|
).strip()
|
|
except Exception:
|
|
return "unknown"
|
|
|
|
def build(self, prompt_id: str, server_id: str, desired_status: str) -> dict:
|
|
with closing(sqlite3.connect(
|
|
f"file:{os.path.abspath(self.scheduler_database)}?mode=ro", uri=True
|
|
)) as scheduler:
|
|
scheduler.row_factory = sqlite3.Row
|
|
scheduler.execute("BEGIN")
|
|
job = scheduler.execute(
|
|
"SELECT * FROM scheduler_jobs WHERE prompt_id=?", (prompt_id,)
|
|
).fetchone()
|
|
api_logs = scheduler.execute(
|
|
"SELECT * FROM scheduler_api_logs WHERE prompt_id=? ORDER BY sequence",
|
|
(prompt_id,),
|
|
).fetchall()
|
|
job_logs = scheduler.execute(
|
|
"SELECT * FROM scheduler_job_logs WHERE prompt_id=?", (prompt_id,)
|
|
).fetchall()
|
|
cloud_task = scheduler.execute(
|
|
"SELECT * FROM scheduler_cloud_tasks WHERE prompt_id=?", (prompt_id,)
|
|
).fetchone()
|
|
cloud_artifacts = scheduler.execute(
|
|
"SELECT * FROM scheduler_cloud_artifacts WHERE prompt_id=? ORDER BY ordinal",
|
|
(prompt_id,),
|
|
).fetchall()
|
|
|
|
history_row = None
|
|
history_parsed = None
|
|
history_schema = {"create_sql": None, "columns": []}
|
|
if os.path.isfile(self.history_database):
|
|
with closing(sqlite3.connect(
|
|
f"file:{os.path.abspath(self.history_database)}?mode=ro", uri=True
|
|
)) as history:
|
|
history.row_factory = sqlite3.Row
|
|
history.execute("BEGIN")
|
|
history_row = history.execute(
|
|
"SELECT * FROM history WHERE prompt_id=?", (prompt_id,)
|
|
).fetchone()
|
|
history_schema = self._schema(history, "history")
|
|
if history_row:
|
|
try:
|
|
history_parsed = json.loads(history_row["data"])
|
|
except (TypeError, json.JSONDecodeError):
|
|
history_parsed = None
|
|
|
|
related_assets = {}
|
|
for row in scheduler.execute(
|
|
"SELECT * FROM scheduler_assets WHERE prompt_id=?", (prompt_id,)
|
|
).fetchall():
|
|
related_assets[row["asset_id"]] = {
|
|
"relation_basis": ["scheduler_assets.prompt_id"],
|
|
"row": row_dict(row),
|
|
}
|
|
for asset_id in _asset_ids((history_parsed or {}).get("outputs", {})):
|
|
row = scheduler.execute(
|
|
"SELECT * FROM scheduler_assets WHERE asset_id=?", (asset_id,)
|
|
).fetchone()
|
|
if not row:
|
|
continue
|
|
item = related_assets.setdefault(
|
|
asset_id, {"relation_basis": [], "row": row_dict(row)}
|
|
)
|
|
if "history.outputs.asset_id" not in item["relation_basis"]:
|
|
item["relation_basis"].append("history.outputs.asset_id")
|
|
|
|
worker = None
|
|
if job and job["worker_port"]:
|
|
worker_row = scheduler.execute(
|
|
"SELECT * FROM scheduler_workers WHERE port=?", (job["worker_port"],)
|
|
).fetchone()
|
|
if worker_row:
|
|
worker = {
|
|
"snapshot_kind": "snapshot_at_export",
|
|
"row": row_dict(worker_row),
|
|
}
|
|
|
|
schemas = {
|
|
table: self._schema(scheduler, table) for table in self.SCHEDULER_TABLES
|
|
}
|
|
schemas["history"] = history_schema
|
|
|
|
task_data = row_dict(cloud_task)
|
|
if task_data:
|
|
task_data["cloud_status"] = desired_status
|
|
task_data["manifest_upload_status"] = UPLOADED
|
|
artifact_rows = [row_dict(row) for row in cloud_artifacts]
|
|
artifacts = [
|
|
{
|
|
"ordinal": row["ordinal"],
|
|
"source_kind": row["source_kind"],
|
|
"original_filename": row["filename"],
|
|
"content_type": row["mime_type"],
|
|
"size_bytes": row["size_bytes"],
|
|
"sha256": row["sha256"],
|
|
"oss_object_key": row["oss_object_key"],
|
|
"oss_uri": row["oss_uri"],
|
|
"public_url": row["public_url"],
|
|
"upload_status": row["upload_status"],
|
|
"attempt_count": row["attempt_count"],
|
|
"last_error": row["last_error"],
|
|
}
|
|
for row in artifact_rows
|
|
]
|
|
manifest = {
|
|
"schema_version": 1,
|
|
"export": {
|
|
"prompt_id": prompt_id,
|
|
"server_id": server_id,
|
|
"exported_at": datetime.now(ZoneInfo("Asia/Shanghai")).isoformat(),
|
|
"source_scheduler_db": os.path.basename(self.scheduler_database),
|
|
"source_history_db": os.path.basename(self.history_database),
|
|
"plugin_commit": self._commit(),
|
|
"generation_status": job["status"] if job else None,
|
|
"cloud_status": desired_status,
|
|
},
|
|
"database_schema": schemas,
|
|
"records": {
|
|
"scheduler_jobs": row_dict(job),
|
|
"scheduler_api_logs": [row_dict(row) for row in api_logs],
|
|
"scheduler_assets": list(related_assets.values()),
|
|
"scheduler_job_logs": [row_dict(row) for row in job_logs],
|
|
"scheduler_workers": worker,
|
|
"history": {
|
|
"row": row_dict(history_row),
|
|
"parsed_data": history_parsed,
|
|
"restore_source": "row.data",
|
|
}
|
|
if history_row
|
|
else None,
|
|
"scheduler_cloud_tasks": task_data,
|
|
"scheduler_cloud_artifacts": artifact_rows,
|
|
},
|
|
"artifacts": artifacts,
|
|
"integrity": {
|
|
"manifest_content_sha256": None,
|
|
"record_counts": {
|
|
"scheduler_jobs": 1 if job else 0,
|
|
"scheduler_api_logs": len(api_logs),
|
|
"scheduler_assets": len(related_assets),
|
|
"scheduler_job_logs": len(job_logs),
|
|
"scheduler_workers": 1 if worker else 0,
|
|
"history": 1 if history_row else 0,
|
|
"scheduler_cloud_tasks": 1 if cloud_task else 0,
|
|
"scheduler_cloud_artifacts": len(cloud_artifacts),
|
|
},
|
|
"artifact_count": len(artifacts),
|
|
"artifact_uploaded_count": sum(
|
|
1 for row in artifact_rows if row["upload_status"] == UPLOADED
|
|
),
|
|
"artifact_failed_count": sum(
|
|
1
|
|
for row in artifact_rows
|
|
if row["upload_status"] in (CLOUD_FAILED, SOURCE_MISSING)
|
|
),
|
|
},
|
|
}
|
|
canonical = json.dumps(
|
|
manifest, ensure_ascii=False, sort_keys=True, separators=(",", ":")
|
|
).encode("utf-8")
|
|
manifest["integrity"]["manifest_content_sha256"] = hashlib.sha256(
|
|
canonical
|
|
).hexdigest()
|
|
return manifest
|
|
|
|
def write_gzip(self, manifest: dict, path: str, max_bytes: int) -> dict:
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
# Do not let gzip derive its embedded FNAME from the staging path. The
|
|
# staging file deliberately ends in `.gz.tmp`; archive tools on Windows
|
|
# prefer the embedded name and would otherwise extract that temporary
|
|
# filename instead of the user-facing `manifest.json`.
|
|
with open(path, "wb") as raw_output:
|
|
with gzip.GzipFile(
|
|
filename="manifest.json",
|
|
mode="wb",
|
|
fileobj=raw_output,
|
|
compresslevel=6,
|
|
mtime=0,
|
|
) as gzip_output:
|
|
with io.TextIOWrapper(gzip_output, encoding="utf-8") as output:
|
|
json.dump(
|
|
manifest,
|
|
output,
|
|
ensure_ascii=False,
|
|
sort_keys=True,
|
|
indent=2,
|
|
)
|
|
output.write("\n")
|
|
size = os.path.getsize(path)
|
|
if size > int(max_bytes):
|
|
raise ValueError("manifest_size_limit_exceeded")
|
|
return {"size_bytes": size, "sha256": file_sha256(path)}
|
|
|
|
|
|
class CloudArchiveManager:
|
|
def __init__(
|
|
self,
|
|
config: CloudArchiveConfig,
|
|
scheduler_database: str,
|
|
history_database: str,
|
|
output_dir: str,
|
|
temp_dir: str,
|
|
project_dir: str,
|
|
logger=None,
|
|
backend=None,
|
|
start_worker: bool = True,
|
|
):
|
|
config.validate()
|
|
self.config = config
|
|
self.store = CloudArchiveStore(scheduler_database)
|
|
self.scheduler_database = scheduler_database
|
|
self.history_database = history_database
|
|
self.output_dir = output_dir
|
|
self.temp_dir = temp_dir
|
|
self.project_dir = project_dir
|
|
self.logger = logger
|
|
self._backend = backend
|
|
self._stop = threading.Event()
|
|
self._thread = None
|
|
self._credentials_warning_logged = False
|
|
self._next_reconcile_at = 0.0
|
|
self._archive_started_at = time.time()
|
|
self.claimant = f"{config.server_id}:{os.getpid()}:{uuid.uuid4().hex[:8]}"
|
|
self.exporter = ManifestExporter(
|
|
scheduler_database, history_database, project_dir
|
|
)
|
|
os.makedirs(config.staging_dir, exist_ok=True)
|
|
if config.enabled and start_worker:
|
|
self._thread = threading.Thread(
|
|
target=self._run,
|
|
name=f"cloud-archive-{config.server_id}-{os.getpid()}",
|
|
daemon=True,
|
|
)
|
|
self._thread.start()
|
|
|
|
@property
|
|
def backend(self):
|
|
if self._backend is None:
|
|
self._backend = OssV2Backend(self.config)
|
|
return self._backend
|
|
|
|
def enqueue_completion(
|
|
self,
|
|
prompt_id: str,
|
|
owner_id: str,
|
|
worker_port: int,
|
|
generation_status: str,
|
|
history_item: dict,
|
|
finished_at: float = None,
|
|
) -> int:
|
|
artifacts = discover_artifacts(history_item, self.output_dir, self.temp_dir)
|
|
self.store.enqueue_task(
|
|
self.config,
|
|
prompt_id,
|
|
owner_id,
|
|
worker_port,
|
|
generation_status,
|
|
finished_at or time.time(),
|
|
artifacts,
|
|
)
|
|
return len(artifacts)
|
|
|
|
def stop(self) -> None:
|
|
self._stop.set()
|
|
if self._thread:
|
|
self._thread.join(timeout=5)
|
|
|
|
def _log_exception(self, message: str, *args) -> None:
|
|
if self.logger:
|
|
self.logger.exception(message, *args)
|
|
|
|
def _run(self) -> None:
|
|
while not self._stop.wait(1):
|
|
try:
|
|
if time.time() >= self._next_reconcile_at:
|
|
self._reconcile_terminal_jobs()
|
|
self._next_reconcile_at = time.time() + 30
|
|
if not (
|
|
os.getenv("OSS_ACCESS_KEY_ID")
|
|
and os.getenv("OSS_ACCESS_KEY_SECRET")
|
|
):
|
|
if not self._credentials_warning_logged and self.logger:
|
|
self.logger.error(
|
|
"Cloud archive is enabled but OSS credentials are not loaded; "
|
|
"pending rows will wait without consuming attempts"
|
|
)
|
|
self._credentials_warning_logged = True
|
|
self._stop.wait(10)
|
|
continue
|
|
self._credentials_warning_logged = False
|
|
artifact = self.store.claim_artifact(
|
|
self.claimant, self.config.upload_concurrency
|
|
)
|
|
if artifact:
|
|
self._upload_artifact(artifact)
|
|
continue
|
|
task = self.store.claim_manifest(
|
|
self.claimant,
|
|
self.config.max_attempts,
|
|
self.config.upload_concurrency,
|
|
)
|
|
if task:
|
|
self._upload_manifest(task)
|
|
except Exception:
|
|
self._log_exception("Unexpected cloud archive worker error")
|
|
self._stop.wait(2)
|
|
|
|
def _reconcile_terminal_jobs(self, limit: int = 20) -> int:
|
|
with closing(sqlite3.connect(self.scheduler_database)) as scheduler:
|
|
scheduler.row_factory = sqlite3.Row
|
|
jobs = scheduler.execute(
|
|
"""
|
|
SELECT job.prompt_id, job.owner_id, job.worker_port, job.status,
|
|
COALESCE(job.finished_at, job.submitted_at) finished_at
|
|
FROM scheduler_jobs job
|
|
LEFT JOIN scheduler_cloud_tasks cloud
|
|
ON cloud.prompt_id=job.prompt_id
|
|
WHERE cloud.prompt_id IS NULL
|
|
AND job.status IN ('completed','failed','cancelled','worker_lost')
|
|
AND COALESCE(job.finished_at, job.submitted_at) >= ?
|
|
ORDER BY COALESCE(job.finished_at, job.submitted_at), job.sequence
|
|
LIMIT ?
|
|
""",
|
|
(self._archive_started_at, max(1, int(limit))),
|
|
).fetchall()
|
|
if not jobs:
|
|
return 0
|
|
|
|
history_by_prompt = {}
|
|
if os.path.isfile(self.history_database):
|
|
with closing(sqlite3.connect(self.history_database)) as history:
|
|
for job in jobs:
|
|
row = history.execute(
|
|
"SELECT data FROM history WHERE prompt_id=?",
|
|
(job["prompt_id"],),
|
|
).fetchone()
|
|
if not row:
|
|
continue
|
|
try:
|
|
parsed = json.loads(row[0])
|
|
except (TypeError, json.JSONDecodeError):
|
|
parsed = {}
|
|
history_by_prompt[job["prompt_id"]] = (
|
|
parsed if isinstance(parsed, dict) else {}
|
|
)
|
|
|
|
registered = 0
|
|
for job in jobs:
|
|
item = history_by_prompt.get(job["prompt_id"], {})
|
|
artifacts = discover_artifacts(item, self.output_dir, self.temp_dir)
|
|
self.store.enqueue_task(
|
|
self.config,
|
|
job["prompt_id"],
|
|
job["owner_id"] or "",
|
|
int(job["worker_port"] or 0),
|
|
job["status"],
|
|
float(job["finished_at"] or time.time()),
|
|
artifacts,
|
|
)
|
|
registered += 1
|
|
return registered
|
|
|
|
def _artifact_staging_path(self, artifact: dict) -> str:
|
|
return os.path.join(
|
|
self.config.staging_dir,
|
|
self.config.server_id,
|
|
artifact["prompt_id"],
|
|
f"{int(artifact['ordinal']):03d}-{sanitize_filename(artifact['filename'])}",
|
|
)
|
|
|
|
@contextmanager
|
|
def _lease_heartbeat(self, renew: Callable[[], bool]):
|
|
stopped = threading.Event()
|
|
|
|
def run():
|
|
while not stopped.wait(60):
|
|
try:
|
|
if not renew():
|
|
return
|
|
except Exception:
|
|
self._log_exception("Failed to renew cloud archive lease")
|
|
|
|
thread = threading.Thread(target=run, name="cloud-archive-lease", daemon=True)
|
|
thread.start()
|
|
try:
|
|
yield
|
|
finally:
|
|
stopped.set()
|
|
thread.join(timeout=2)
|
|
|
|
def _upload_artifact(self, artifact: dict) -> None:
|
|
try:
|
|
path = artifact.get("source_path")
|
|
mime_type = artifact.get("mime_type") or "application/octet-stream"
|
|
if artifact["source_kind"] == "remote_url":
|
|
path = self._artifact_staging_path(artifact)
|
|
if not os.path.isfile(path):
|
|
download_path = path + ".download.tmp"
|
|
try:
|
|
mime_type, _ = download_remote(
|
|
artifact["source_url"],
|
|
download_path,
|
|
self.config.remote_max_bytes,
|
|
)
|
|
os.replace(download_path, path)
|
|
except Exception:
|
|
try:
|
|
os.remove(download_path)
|
|
except OSError:
|
|
pass
|
|
raise
|
|
if not path or not os.path.isfile(path):
|
|
raise FileNotFoundError(f"Artifact source is missing: {path}")
|
|
size = os.path.getsize(path)
|
|
digest = file_sha256(path)
|
|
with self._lease_heartbeat(
|
|
lambda: self.store.renew_artifact_lease(
|
|
artifact["id"], self.claimant
|
|
)
|
|
):
|
|
result = self.backend.upload_file(
|
|
artifact["oss_object_key"], path, mime_type, sha256=digest
|
|
)
|
|
if os.path.getsize(path) != size or file_sha256(path) != digest:
|
|
raise RuntimeError("Artifact source changed during upload")
|
|
self.store.artifact_succeeded(
|
|
artifact["id"],
|
|
{
|
|
"source_path": path,
|
|
"mime_type": mime_type,
|
|
"size_bytes": size,
|
|
"sha256": digest,
|
|
"etag": result.get("etag"),
|
|
"oss_crc64": result.get("oss_crc64"),
|
|
},
|
|
)
|
|
if artifact["source_kind"] == "remote_url":
|
|
try:
|
|
os.remove(path)
|
|
except OSError:
|
|
pass
|
|
except FileNotFoundError as error:
|
|
self.store.artifact_source_missing(artifact["id"], str(error))
|
|
self._log_exception(
|
|
"Cloud artifact source is missing for %s #%s",
|
|
artifact["prompt_id"],
|
|
artifact["ordinal"],
|
|
)
|
|
except Exception as error:
|
|
self.store.artifact_failed(
|
|
artifact["id"], str(error), self.config.max_attempts
|
|
)
|
|
self._log_exception(
|
|
"Cloud artifact upload failed for %s #%s",
|
|
artifact["prompt_id"],
|
|
artifact["ordinal"],
|
|
)
|
|
|
|
def _desired_status(self, task: dict) -> str:
|
|
if int(task.get("artifact_count") or 0) == 0:
|
|
return NO_ARTIFACTS
|
|
if int(task.get("failed_count") or 0) > 0:
|
|
return CLOUD_FAILED
|
|
return CLOUD_COMPLETED
|
|
|
|
def _upload_manifest(self, task: dict) -> None:
|
|
prompt_id = task["prompt_id"]
|
|
staging = os.path.join(
|
|
self.config.staging_dir,
|
|
self.config.server_id,
|
|
prompt_id,
|
|
"manifest.json.gz.tmp",
|
|
)
|
|
try:
|
|
desired = self._desired_status(task)
|
|
manifest = self.exporter.build(prompt_id, self.config.server_id, desired)
|
|
file_result = self.exporter.write_gzip(
|
|
manifest, staging, self.config.manifest_max_bytes
|
|
)
|
|
with self._lease_heartbeat(
|
|
lambda: self.store.renew_manifest_lease(prompt_id, self.claimant)
|
|
):
|
|
self.backend.upload_file(
|
|
task["manifest_object_key"],
|
|
staging,
|
|
"application/gzip",
|
|
cache_control="no-store",
|
|
sha256=file_result["sha256"],
|
|
private=True,
|
|
)
|
|
self.store.manifest_succeeded(prompt_id, file_result)
|
|
try:
|
|
os.remove(staging)
|
|
except OSError:
|
|
pass
|
|
except Exception as error:
|
|
self.store.manifest_failed(
|
|
prompt_id, str(error), self.config.max_attempts
|
|
)
|
|
self._log_exception("Cloud manifest upload failed for %s", prompt_id)
|