368 lines
16 KiB
Python
368 lines
16 KiB
Python
import gzip
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from contextlib import closing
|
|
from pathlib import Path
|
|
|
|
|
|
ROOT = Path(__file__).parents[1]
|
|
|
|
|
|
def load_module(name, path):
|
|
spec = importlib.util.spec_from_file_location(name, path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
cloud = load_module("test_cloud_archive_module", ROOT / "utils" / "cloud_archive.py")
|
|
scheduler_module = load_module(
|
|
"test_cloud_scheduler_store", ROOT / "utils" / "scheduler_store.py"
|
|
)
|
|
history_module = load_module(
|
|
"test_cloud_history_store", ROOT / "utils" / "history_store.py"
|
|
)
|
|
|
|
|
|
class FakeBackend:
|
|
def __init__(self, fail=False):
|
|
self.fail = fail
|
|
self.objects = {}
|
|
self.calls = []
|
|
|
|
def upload_file(
|
|
self,
|
|
key,
|
|
path,
|
|
content_type,
|
|
content_encoding=None,
|
|
cache_control=None,
|
|
sha256=None,
|
|
private=False,
|
|
):
|
|
self.calls.append(
|
|
(key, content_type, content_encoding, cache_control, sha256, private)
|
|
)
|
|
if self.fail:
|
|
raise RuntimeError("simulated OSS outage")
|
|
with open(path, "rb") as source:
|
|
data = source.read()
|
|
self.objects[key] = data
|
|
return {"etag": "fake-etag", "oss_crc64": "123", "content_length": len(data)}
|
|
|
|
|
|
class CloudArchiveTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.root = Path(self.temp.name)
|
|
self.output = self.root / "output"
|
|
self.temp_output = self.root / "temp"
|
|
self.output.mkdir()
|
|
self.temp_output.mkdir()
|
|
self.scheduler_db = self.root / "scheduler.sqlite3"
|
|
self.history_db = self.root / "history.sqlite3"
|
|
self.scheduler = scheduler_module.SchedulerStore(os.fspath(self.scheduler_db))
|
|
self.history = history_module.HistoryStore(os.fspath(self.history_db))
|
|
self.config = cloud.CloudArchiveConfig(
|
|
enabled=True,
|
|
server_id="seetacloud-5090x2",
|
|
region="cn-hangzhou",
|
|
endpoint="oss-cn-hangzhou.aliyuncs.com",
|
|
bucket="goumee-coze",
|
|
prefix="Goumee-ComfyUI-Server-Data",
|
|
public_base_url="https://goumee-coze.oss-cn-hangzhou.aliyuncs.com",
|
|
staging_dir=os.fspath(self.root / "staging"),
|
|
checkpoint_dir=os.fspath(self.root / "checkpoints"),
|
|
)
|
|
|
|
def tearDown(self):
|
|
self.temp.cleanup()
|
|
|
|
def add_job(self, prompt_id="prompt-1", status="completed", worker_port=6009):
|
|
item = (0.0, prompt_id, {"1": {"class_type": "Test"}}, {}, ["1"], {})
|
|
self.scheduler.enqueue(item, "owner-1", 6006)
|
|
with closing(sqlite3.connect(self.scheduler_db)) as connection:
|
|
with connection:
|
|
connection.execute(
|
|
"""
|
|
UPDATE scheduler_jobs
|
|
SET status=?, worker_port=?, started_at=100, finished_at=200,
|
|
payload=?, error=?
|
|
WHERE prompt_id=?
|
|
""",
|
|
(
|
|
status,
|
|
worker_port,
|
|
'{"raw":"payload"}',
|
|
"failure detail" if status == "failed" else None,
|
|
prompt_id,
|
|
),
|
|
)
|
|
|
|
def history_item(self, filename=None, asset_id=None):
|
|
output = {}
|
|
if filename:
|
|
output.update({"filename": filename, "subfolder": "", "type": "output"})
|
|
if asset_id:
|
|
output["id"] = asset_id
|
|
return {
|
|
"outputs": {"1": {"images": [output] if output else []}},
|
|
"status": {"status_str": "success", "completed": True},
|
|
"user_id": "owner-1",
|
|
}
|
|
|
|
def manager(self, backend=None):
|
|
return cloud.CloudArchiveManager(
|
|
self.config,
|
|
os.fspath(self.scheduler_db),
|
|
os.fspath(self.history_db),
|
|
os.fspath(self.output),
|
|
os.fspath(self.temp_output),
|
|
os.fspath(ROOT),
|
|
backend=backend or FakeBackend(),
|
|
start_worker=False,
|
|
)
|
|
|
|
def test_migration_is_idempotent_and_duplicate_enqueue_is_idempotent(self):
|
|
self.add_job()
|
|
store = cloud.CloudArchiveStore(os.fspath(self.scheduler_db))
|
|
cloud.CloudArchiveStore(os.fspath(self.scheduler_db))
|
|
artifact = {
|
|
"source_kind": "local_output",
|
|
"source_path": os.fspath(self.output / "a.png"),
|
|
"source_url": None,
|
|
"filename": "a.png",
|
|
"mime_type": "image/png",
|
|
"reference": {},
|
|
"output_path": "outputs/1",
|
|
}
|
|
store.enqueue_task(self.config, "prompt-1", "owner-1", 6009, "completed", 200, [artifact])
|
|
store.enqueue_task(self.config, "prompt-1", "owner-1", 6009, "completed", 200, [artifact])
|
|
self.assertEqual(1, len(store.artifacts("prompt-1")))
|
|
|
|
def test_discovery_only_walks_outputs_and_supports_fullpath_and_remote_url(self):
|
|
local = self.output / "result.mp4"
|
|
local.write_bytes(b"video")
|
|
history = {
|
|
"prompt": {"input": "https://example.com/input.mp4"},
|
|
"outputs": {
|
|
"1": {"videos": [{"fullpath": os.fspath(local)}]},
|
|
"2": {"url": "https://cdn.example.com/final.mp4"},
|
|
},
|
|
}
|
|
found = cloud.discover_artifacts(
|
|
history, os.fspath(self.output), os.fspath(self.temp_output)
|
|
)
|
|
self.assertEqual(["local_output", "remote_url"], [row["source_kind"] for row in found])
|
|
self.assertNotIn("input.mp4", json.dumps(found))
|
|
|
|
def test_public_and_private_object_keys_are_strictly_separated(self):
|
|
media = self.config.media_key("abc", 1, "a b.mp4", 1787639123)
|
|
manifest = self.config.manifest_key("abc")
|
|
self.assertIn("/public/servers/seetacloud-5090x2/", media)
|
|
self.assertEqual(
|
|
"Goumee-ComfyUI-Server-Data/private/jobs/abc/manifest.json.gz", manifest
|
|
)
|
|
self.assertNotIn("private", self.config.public_url(media))
|
|
|
|
def test_local_media_and_gzip_manifest_upload_complete_without_deleting_source(self):
|
|
self.add_job()
|
|
source = self.output / "result.png"
|
|
source.write_bytes(b"png-data")
|
|
item = self.history_item("result.png")
|
|
self.history.save("prompt-1", item, 100)
|
|
backend = FakeBackend()
|
|
manager = self.manager(backend)
|
|
self.assertEqual(1, manager.enqueue_completion("prompt-1", "owner-1", 6009, "completed", item, 200))
|
|
|
|
artifact = manager.store.claim_artifact("worker", 2)
|
|
manager._upload_artifact(artifact)
|
|
task = manager.store.claim_manifest("worker", 3, 2)
|
|
manager._upload_manifest(task)
|
|
|
|
current = manager.store.task("prompt-1")
|
|
self.assertEqual(cloud.CLOUD_COMPLETED, current["cloud_status"])
|
|
self.assertEqual(cloud.UPLOADED, current["manifest_upload_status"])
|
|
self.assertTrue(source.exists())
|
|
manifest_bytes = backend.objects[self.config.manifest_key("prompt-1")]
|
|
manifest = json.loads(gzip.decompress(manifest_bytes))
|
|
self.assertEqual("prompt-1", manifest["export"]["prompt_id"])
|
|
self.assertEqual("uploaded", manifest["artifacts"][0]["upload_status"])
|
|
self.assertEqual("application/gzip", backend.calls[-1][1])
|
|
self.assertIsNone(backend.calls[-1][2])
|
|
self.assertEqual("no-store", backend.calls[-1][3])
|
|
self.assertTrue(backend.calls[-1][5])
|
|
decoded_text = gzip.decompress(manifest_bytes).decode("utf-8")
|
|
self.assertIn('\n "artifacts":', decoded_text)
|
|
self.assertTrue(decoded_text.endswith("\n"))
|
|
self.assertTrue(manifest_bytes[3] & 0x08)
|
|
embedded_name = manifest_bytes[10:].split(b"\0", 1)[0]
|
|
self.assertEqual(b"manifest.json", embedded_name)
|
|
|
|
def test_empty_failed_task_still_uploads_private_manifest(self):
|
|
self.add_job(status="failed")
|
|
item = self.history_item()
|
|
self.history.save("prompt-1", item, 100)
|
|
backend = FakeBackend()
|
|
manager = self.manager(backend)
|
|
manager.enqueue_completion("prompt-1", "owner-1", 6009, "failed", item, 200)
|
|
task = manager.store.claim_manifest("worker", 3, 2)
|
|
manager._upload_manifest(task)
|
|
self.assertEqual(cloud.NO_ARTIFACTS, manager.store.task("prompt-1")["cloud_status"])
|
|
self.assertIn(self.config.manifest_key("prompt-1"), backend.objects)
|
|
|
|
def test_missing_historical_output_becomes_source_missing(self):
|
|
self.add_job()
|
|
manager = self.manager(FakeBackend())
|
|
item = self.history_item("missing.png")
|
|
manager.enqueue_completion("prompt-1", "owner-1", 6009, "completed", item, 200)
|
|
artifact = manager.store.claim_artifact("worker", 2)
|
|
manager._upload_artifact(artifact)
|
|
row = manager.store.artifacts("prompt-1")[0]
|
|
self.assertEqual(cloud.SOURCE_MISSING, row["upload_status"])
|
|
|
|
def test_exactly_three_outer_failures_are_terminal(self):
|
|
self.add_job()
|
|
source = self.output / "result.png"
|
|
source.write_bytes(b"png")
|
|
manager = self.manager(FakeBackend(fail=True))
|
|
manager.enqueue_completion(
|
|
"prompt-1", "owner-1", 6009, "completed", self.history_item("result.png"), 200
|
|
)
|
|
for attempt in range(3):
|
|
with closing(sqlite3.connect(self.scheduler_db)) as connection:
|
|
with connection:
|
|
connection.execute(
|
|
"UPDATE scheduler_cloud_artifacts SET next_attempt_at=0 WHERE prompt_id='prompt-1'"
|
|
)
|
|
artifact = manager.store.claim_artifact("worker", 2)
|
|
self.assertIsNotNone(artifact)
|
|
manager._upload_artifact(artifact)
|
|
row = manager.store.artifacts("prompt-1")[0]
|
|
self.assertEqual(3, row["attempt_count"])
|
|
self.assertEqual(cloud.CLOUD_FAILED, row["upload_status"])
|
|
|
|
def test_exactly_three_manifest_failures_make_task_cloud_failed(self):
|
|
self.add_job(status="failed")
|
|
manager = self.manager(FakeBackend(fail=True))
|
|
manager.enqueue_completion("prompt-1", "owner-1", 6009, "failed", {}, 200)
|
|
for _ in range(3):
|
|
with closing(sqlite3.connect(self.scheduler_db)) as connection:
|
|
with connection:
|
|
connection.execute(
|
|
"UPDATE scheduler_cloud_tasks SET manifest_next_attempt_at=0 WHERE prompt_id='prompt-1'"
|
|
)
|
|
task = manager.store.claim_manifest("worker", 3, 2)
|
|
self.assertIsNotNone(task)
|
|
manager._upload_manifest(task)
|
|
task = manager.store.task("prompt-1")
|
|
self.assertEqual(3, task["manifest_attempt_count"])
|
|
self.assertEqual("failed", task["manifest_upload_status"])
|
|
self.assertEqual(cloud.CLOUD_FAILED, task["cloud_status"])
|
|
|
|
def test_expired_artifact_lease_can_be_taken_over(self):
|
|
self.add_job()
|
|
source = self.output / "result.png"
|
|
source.write_bytes(b"png")
|
|
manager = self.manager(FakeBackend())
|
|
manager.enqueue_completion(
|
|
"prompt-1", "owner-1", 6009, "completed", self.history_item("result.png"), 200
|
|
)
|
|
first = manager.store.claim_artifact("worker-a", 2, lease_seconds=0)
|
|
second = manager.store.claim_artifact("worker-b", 2)
|
|
self.assertEqual(first["id"], second["id"])
|
|
self.assertEqual(2, second["attempt_count"])
|
|
|
|
def test_same_artifact_reference_can_be_archived_for_two_prompts(self):
|
|
self.add_job("prompt-1")
|
|
self.add_job("prompt-2")
|
|
source = self.output / "shared.png"
|
|
source.write_bytes(b"shared")
|
|
item = self.history_item("shared.png", asset_id="asset-shared")
|
|
manager = self.manager(FakeBackend())
|
|
manager.enqueue_completion("prompt-1", "owner-1", 6009, "completed", item, 200)
|
|
manager.enqueue_completion("prompt-2", "owner-1", 6009, "completed", item, 200)
|
|
first = manager.store.artifacts("prompt-1")[0]
|
|
second = manager.store.artifacts("prompt-2")[0]
|
|
self.assertNotEqual(first["oss_object_key"], second["oss_object_key"])
|
|
self.assertIn("/prompt-1/", first["oss_object_key"])
|
|
self.assertIn("/prompt-2/", second["oss_object_key"])
|
|
|
|
def test_manifest_preserves_blob_raw_history_and_overwritten_asset_relation(self):
|
|
self.add_job()
|
|
item = self.history_item(asset_id="asset-shared")
|
|
self.history.save("prompt-1", item, 100)
|
|
raw_blob = b"\x00\xffapi-body"
|
|
with closing(sqlite3.connect(self.scheduler_db)) as connection:
|
|
with connection:
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO scheduler_api_logs(
|
|
prompt_id, recorded_at, client, method, url, request_headers,
|
|
request_body, response_status, response_headers, response_body, error
|
|
) VALUES (?, 1, 'test', 'POST', 'https://api.example', '{}', ?, 200, '{}', ?, NULL)
|
|
""",
|
|
("prompt-1", raw_blob, sqlite3.Binary(b"response")),
|
|
)
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO scheduler_assets(
|
|
asset_id, owner_id, prompt_id, worker_port, file_path, data, updated_at
|
|
) VALUES ('asset-shared', 'owner-1', 'different-prompt', 6009, '', '{}', 1)
|
|
"""
|
|
)
|
|
manager = self.manager(FakeBackend())
|
|
manager.enqueue_completion("prompt-1", "owner-1", 6009, "completed", item, 200)
|
|
manifest = manager.exporter.build("prompt-1", self.config.server_id, cloud.NO_ARTIFACTS)
|
|
body = manifest["records"]["scheduler_api_logs"][0]["request_body"]
|
|
self.assertEqual("blob", body["$sqlite_type"])
|
|
self.assertEqual(hashlib.sha256(raw_blob).hexdigest(), body["sha256"])
|
|
self.assertEqual(
|
|
["history.outputs.asset_id"],
|
|
manifest["records"]["scheduler_assets"][0]["relation_basis"],
|
|
)
|
|
history_record = manifest["records"]["history"]
|
|
self.assertIsInstance(history_record["row"]["data"], str)
|
|
self.assertEqual("asset-shared", history_record["parsed_data"]["outputs"]["1"]["images"][0]["id"])
|
|
self.assertNotIn("sqlite_sequence", manifest["database_schema"])
|
|
|
|
def test_manifest_canonical_content_hash_is_reproducible(self):
|
|
self.add_job()
|
|
manager = self.manager(FakeBackend())
|
|
manager.enqueue_completion("prompt-1", "owner-1", 6009, "completed", {}, 200)
|
|
manifest = manager.exporter.build("prompt-1", self.config.server_id, cloud.NO_ARTIFACTS)
|
|
expected = manifest["integrity"]["manifest_content_sha256"]
|
|
manifest["integrity"]["manifest_content_sha256"] = None
|
|
canonical = json.dumps(
|
|
manifest, ensure_ascii=False, sort_keys=True, separators=(",", ":")
|
|
).encode("utf-8")
|
|
self.assertEqual(expected, hashlib.sha256(canonical).hexdigest())
|
|
|
|
def test_manifest_does_not_export_oss_environment_credentials(self):
|
|
self.add_job()
|
|
os.environ["OSS_ACCESS_KEY_ID"] = "SHOULD_NOT_APPEAR"
|
|
os.environ["OSS_ACCESS_KEY_SECRET"] = "SHOULD_NOT_APPEAR_EITHER"
|
|
try:
|
|
manager = self.manager(FakeBackend())
|
|
manager.enqueue_completion("prompt-1", "owner-1", 6009, "completed", {}, 200)
|
|
encoded = json.dumps(
|
|
manager.exporter.build(
|
|
"prompt-1", self.config.server_id, cloud.NO_ARTIFACTS
|
|
)
|
|
)
|
|
finally:
|
|
os.environ.pop("OSS_ACCESS_KEY_ID", None)
|
|
os.environ.pop("OSS_ACCESS_KEY_SECRET", None)
|
|
self.assertNotIn("SHOULD_NOT_APPEAR", encoded)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|