Add file size conversion and enhance Civitai job processing

- Introduced the `file_size_bytes` function to convert Civitai model sizes from kilobytes to bytes, improving data handling.
- Updated the `seed_civitai` function to include file size in job definitions, enhancing model processing efficiency.
- Enhanced the `should_skip` function to utilize expected size for faster decision-making during job processing.
- Added tests for new functionality, ensuring robustness in file size handling and job processing logic.
This commit is contained in:
Leonid Pershin
2026-08-21 10:52:44 +03:00
parent 6cdd6ecfa1
commit ef743a6e6d
6 changed files with 260 additions and 14 deletions
+47 -5
View File
@@ -1,7 +1,12 @@
from pathlib import Path
from gpu_rent.civitai import pick_preview_image
from gpu_rent.remote.civitai_fetch import fmt_bytes, progress_line, should_skip
from gpu_rent.civitai import file_size_bytes, pick_preview_image
from gpu_rent.remote.civitai_fetch import (
fmt_bytes,
progress_line,
should_skip,
write_sha_sidecar,
)
def test_should_skip_missing(tmp_path: Path):
@@ -17,21 +22,52 @@ def test_should_skip_exists_no_sha(tmp_path: Path):
assert "уже есть" in reason
def test_should_skip_sha_match(tmp_path: Path):
def test_should_skip_size_ok_fast(tmp_path: Path):
"""Re-up must not hash multi-GB weights when size matches."""
dest = tmp_path / "big.safetensors"
data = b"x" * 10_000
dest.write_bytes(data)
expect = "deadbeef" * 8
skip, reason = should_skip(dest, expect, expect_size=len(data))
assert skip is True
assert "size ok" in reason
# sidecar written for next run
assert (tmp_path / "big.safetensors.sha256").is_file()
def test_should_skip_sha_sidecar(tmp_path: Path):
dest = tmp_path / "a.safetensors"
dest.write_bytes(b"weights")
expect = "abcd" * 16
write_sha_sidecar(dest, expect)
skip, reason = should_skip(dest, expect, expect_size=7)
assert skip is True
assert "sidecar" in reason
def test_should_skip_sha_match_slow(tmp_path: Path):
dest = tmp_path / "a.safetensors"
dest.write_bytes(b"weights")
import hashlib
digest = hashlib.sha256(b"weights").hexdigest()
skip, reason = should_skip(dest, digest)
skip, reason = should_skip(dest, digest, verify_sha=True)
assert skip is True
assert "sha ok" in reason
def test_should_not_skip_bad_size(tmp_path: Path):
dest = tmp_path / "a.safetensors"
dest.write_bytes(b"weights")
skip, reason = should_skip(dest, "deadbeef" * 8, expect_size=99999)
assert skip is False
assert "размер" in reason
def test_should_not_skip_bad_sha(tmp_path: Path):
dest = tmp_path / "a.safetensors"
dest.write_bytes(b"weights")
skip, reason = should_skip(dest, "deadbeef" * 8)
skip, reason = should_skip(dest, "deadbeef" * 8, verify_sha=True)
assert skip is False
assert "перекачиваю" in reason
@@ -80,3 +116,9 @@ def test_pick_preview_image_png_and_default():
] == ".preview.jpg"
assert pick_preview_image({"images": []}) is None
assert pick_preview_image({}) is None
def test_file_size_bytes():
assert file_size_bytes({"sizeKB": 1.0}) == 1024
assert file_size_bytes({"sizeKB": 6775430.35}) == int(round(6775430.35 * 1024))
assert file_size_bytes({}) is None