| """Only enable a prebuilt worker after a complete, measured acceptance run.""" |
|
|
| import hashlib |
| import json |
| import math |
| import re |
| from pathlib import Path |
|
|
| from flora_captioner import FLORA |
| from huggingface_hub.utils import validate_repo_id |
|
|
| IDENTITY_FILES = { |
| "worker.py", |
| "subset.py", |
| "flora_captioner.py", |
| "flora_acceleration.py", |
| "flora_batcher.py", |
| "flora_speed_entry.py", |
| "studio_runtime.py", |
| "pipeline-version.json", |
| } |
| STUDIO_RUNTIME_FILES = { |
| "Dockerfile", |
| "worker-requirements.txt", |
| "flora-requirements.txt", |
| "flora-acceleration-requirements.txt", |
| "build_causal_conv1d.py", |
| "bootstrap.py", |
| "flora_studio_entry.py", |
| } |
|
|
|
|
| def validated_worker_image(captioner, features, path=None): |
| if captioner != "flora" or set(features) != {"Subtasks"}: |
| return None |
| path = path or Path(__file__).with_name("flora-worker-image.json") |
| if not path.exists(): |
| return None |
| profile = json.loads(path.read_text()) |
| if type(profile.get("batch_size", 4)) is not int or profile.get("batch_size", 4) not in (1, 2, 4, 8): |
| raise ValueError("Accelerated Flora batch size is not validated") |
| if ( |
| type(profile.get("episode_parallelism", 1)) is not int |
| or not 1 <= profile.get("episode_parallelism", 1) <= 8 |
| ): |
| raise ValueError("Accelerated Flora episode parallelism is not validated") |
| if profile.get("generation_batch_size", profile.get("batch_size", 4)) not in (1, 2, 4, 8, 16, 32, 64): |
| raise ValueError("Accelerated Flora GPU batch size is not validated") |
| accepted = profile.get("acceptance", {}) |
| seconds = accepted.get("end_to_end_seconds") |
| if ( |
| accepted.get("episodes") != 25 |
| or accepted.get("independent") is not True |
| or accepted.get("fused_kernels") is not True |
| or not isinstance(seconds, (int, float)) |
| or isinstance(seconds, bool) |
| or not math.isfinite(seconds) |
| or not 0 < seconds < 1500 |
| or profile.get("captioner_profile") != FLORA |
| ): |
| raise ValueError("Accelerated Flora worker has not passed its acceptance run") |
| validate_repo_id(profile["repo"]) |
| source = profile.get("image_source", "external") |
| if source not in {"external", "studio"}: |
| raise ValueError("Unknown Flora image source") |
| if source == "external" and not re.fullmatch(r"[a-f0-9]{40}", profile["revision"]): |
| raise ValueError("Accelerated Flora image revision must be pinned") |
| hashes = profile.get("expected_code", {}) |
| expected_files = IDENTITY_FILES | ({"flora_vision.py"} if profile.get("vision_frame_batching") else set()) |
| if set(hashes) != expected_files or any( |
| not re.fullmatch(r"[a-f0-9]{64}", value) for value in hashes.values() |
| ): |
| raise ValueError("Accelerated Flora code checksums are incomplete") |
| if source == "studio": |
| runtime_files = profile.get("runtime_files", {}) |
| if set(runtime_files) != STUDIO_RUNTIME_FILES or any( |
| not re.fullmatch(r"[a-f0-9]{64}", value) for value in runtime_files.values() |
| ): |
| raise ValueError("Studio's CUDA build checksums are incomplete") |
| if set(profile.get("kernel_versions", {})) != { |
| "torch", |
| "transformers", |
| "flash-linear-attention", |
| "causal-conv1d", |
| }: |
| raise ValueError("Studio's CUDA dependency versions are incomplete") |
| |
| |
| for name, expected in {**hashes, **runtime_files}.items(): |
| if ( |
| Path(name).name != name |
| or hashlib.sha256((path.parent / name).read_bytes()).hexdigest() != expected |
| ): |
| raise ValueError("Studio's Flora runtime changed; validate it before launching") |
| return profile |
|
|
