From bcc41219f7245cb70568d06de02ceed3e7ff2d99 Mon Sep 17 00:00:00 2001 From: "LocalAI [bot]" <139863280+localai-bot@users.noreply.github.com> Date: Wed, 15 Jul 2026 01:09:33 +0200 Subject: [PATCH] feat: materialize Hugging Face model artifacts (#10825) * feat(config): add model artifact source contract Assisted-by: Codex:GPT-5 [Codex] * feat(downloader): add authenticated raw-byte progress Assisted-by: Codex:GPT-5 [Codex] * feat(huggingface): resolve immutable snapshot manifests Assisted-by: Codex:GPT-5 [Codex] * feat(models): add artifact storage primitives Assisted-by: Codex:GPT-5 [Codex] * feat(models): materialize pinned Hugging Face snapshots Assisted-by: Codex:GPT-5 [Codex] * feat(models): bind managed snapshots at runtime Assisted-by: Codex:GPT-5 [Codex] * feat(gallery): materialize model artifacts during install Assisted-by: Codex:GPT-5 [Codex] * feat(gallery): declare managed Hugging Face artifacts Assisted-by: Codex:GPT-5 [Codex] * feat(models): preload managed model artifacts Assisted-by: Codex:GPT-5 [Codex] * fix(gallery): retain shared artifact caches on delete Assisted-by: Codex:GPT-5 [Codex] * feat(models): report artifact acquisition progress Assisted-by: Codex:GPT-5 [Codex] * refactor(backends): load managed models from ModelFile Assisted-by: Codex:GPT-5 [Codex] * refactor(backends): load staged speech model snapshots Assisted-by: Codex:GPT-5 [Codex] * refactor(backends): use staged snapshots in engine backends Assisted-by: Codex:GPT-5 [Codex] * test(distributed): cover staged artifact snapshots Assisted-by: Codex:GPT-5 [Codex] * docs: explain managed model artifacts Assisted-by: Codex:GPT-5 [Codex] * docs: add product design context Assisted-by: Codex:GPT-5 [Codex] * feat(ui): show model artifact download progress Assisted-by: Codex:GPT-5 [Codex] * Eagerly materialize Hugging Face artifacts Materialize HF-backed model references as managed GGUF artifacts during load, with lazy download retained only as fallback. Assisted-by: Codex:GPT-5 [shell] * Refactor HF downloads through a shared executor Assisted-by: Codex:GPT-5 [shell] * drop Signed-off-by: Ettore Di Giacinto --------- Signed-off-by: Ettore Di Giacinto Co-authored-by: Ettore Di Giacinto --- .impeccable/live/config.json | 6 + backend/python/common/model_utils.py | 28 ++ backend/python/common/model_utils_test.py | 81 ++++ backend/python/diffusers/backend.py | 99 ++--- backend/python/fish-speech/backend.py | 16 +- backend/python/liquid-audio/backend.py | 18 +- backend/python/nemo/backend.py | 17 +- backend/python/qwen-asr/backend.py | 6 +- backend/python/qwen-tts/backend.py | 19 +- backend/python/sglang/backend.py | 7 +- backend/python/transformers/backend.py | 30 +- backend/python/vllm-omni/backend.py | 15 +- backend/python/vllm/backend.py | 8 +- backend/python/voxcpm/backend.py | 8 +- core/application/application.go | 6 +- .../model_artifact_materializer_test.go | 55 +++ core/application/startup.go | 6 +- core/backend/llm.go | 2 +- core/backend/model_artifact_options_test.go | 47 +++ core/backend/options.go | 34 +- core/cli/models.go | 16 +- core/cli/run.go | 8 + .../application_artifact_materializer_test.go | 34 ++ core/config/application_config.go | 41 +- core/config/meta/registry.go | 8 + core/config/model_artifact_binding.go | 110 ++++++ core/config/model_artifact_binding_test.go | 50 +++ core/config/model_artifact_inference.go | 61 +++ core/config/model_config.go | 29 +- core/config/model_config_loader.go | 269 +++++++++---- core/config/model_config_loader_test.go | 205 ++++++++++ core/config/model_config_test.go | 50 +++ core/gallery/importers/helpers.go | 47 +++ core/gallery/importers/helpers_test.go | 50 +++ core/gallery/importers/importers.go | 4 + core/gallery/model_artifacts.go | 94 +++++ core/gallery/model_artifacts_test.go | 240 ++++++++++++ core/gallery/models.go | 184 +++++---- .../e2e/model-artifact-operation.spec.js | 40 ++ .../react-ui/src/components/OperationsBar.jsx | 25 +- core/http/routes/ui_api.go | 18 +- core/http/routes/ui_api_operations_test.go | 38 ++ core/services/distributed/gallery.go | 21 +- core/services/galleryop/artifact_progress.go | 65 ++++ .../galleryop/artifact_progress_test.go | 33 ++ .../galleryop/cancellable_persist_test.go | 12 +- .../galleryop/distributed_sync_test.go | 6 + core/services/galleryop/managers_local.go | 11 +- .../model_artifact_materializer_test.go | 51 +++ core/services/galleryop/models.go | 34 +- core/services/galleryop/operation.go | 12 + core/services/galleryop/service.go | 17 +- core/services/modeladmin/remote_sync_test.go | 26 ++ core/services/nodes/router_dirstage_test.go | 46 +++ core/startup/model_preload.go | 10 +- docs/content/advanced/model-configuration.md | 56 +++ docs/content/features/distributed-mode.md | 17 + docs/content/features/model-gallery.md | 19 + gallery/index.yaml | 128 +++++++ gallery/liquid-audio.yaml | 6 + pkg/downloader/auth_progress_test.go | 173 +++++++++ pkg/downloader/download_plan.go | 39 ++ pkg/downloader/plan_test.go | 44 +++ pkg/downloader/progress.go | 43 ++- pkg/downloader/uri.go | 76 +++- pkg/downloader/uri_test.go | 13 +- pkg/huggingface-api/client.go | 12 + pkg/huggingface-api/snapshot.go | 239 ++++++++++++ pkg/huggingface-api/snapshot_test.go | 149 ++++++++ pkg/model/initializers.go | 6 +- pkg/model/loader.go | 19 +- pkg/model/loader_options.go | 7 + pkg/model/loader_test.go | 16 + pkg/modelartifacts/manifest.go | 47 +++ pkg/modelartifacts/materializer.go | 359 ++++++++++++++++++ pkg/modelartifacts/materializer_test.go | 153 ++++++++ .../modelartifacts_suite_test.go | 13 + pkg/modelartifacts/path.go | 86 +++++ pkg/modelartifacts/path_test.go | 109 ++++++ pkg/modelartifacts/progress.go | 40 ++ pkg/modelartifacts/reference.go | 82 ++++ pkg/modelartifacts/types.go | 160 ++++++++ pkg/modelartifacts/types_test.go | 70 ++++ 83 files changed, 4315 insertions(+), 339 deletions(-) create mode 100644 .impeccable/live/config.json create mode 100644 backend/python/common/model_utils.py create mode 100644 backend/python/common/model_utils_test.py create mode 100644 core/application/model_artifact_materializer_test.go create mode 100644 core/backend/model_artifact_options_test.go create mode 100644 core/config/application_artifact_materializer_test.go create mode 100644 core/config/model_artifact_binding.go create mode 100644 core/config/model_artifact_binding_test.go create mode 100644 core/config/model_artifact_inference.go create mode 100644 core/gallery/model_artifacts.go create mode 100644 core/gallery/model_artifacts_test.go create mode 100644 core/http/react-ui/e2e/model-artifact-operation.spec.js create mode 100644 core/services/galleryop/artifact_progress.go create mode 100644 core/services/galleryop/artifact_progress_test.go create mode 100644 core/services/galleryop/model_artifact_materializer_test.go create mode 100644 pkg/downloader/auth_progress_test.go create mode 100644 pkg/downloader/download_plan.go create mode 100644 pkg/downloader/plan_test.go create mode 100644 pkg/huggingface-api/snapshot.go create mode 100644 pkg/huggingface-api/snapshot_test.go create mode 100644 pkg/modelartifacts/manifest.go create mode 100644 pkg/modelartifacts/materializer.go create mode 100644 pkg/modelartifacts/materializer_test.go create mode 100644 pkg/modelartifacts/modelartifacts_suite_test.go create mode 100644 pkg/modelartifacts/path.go create mode 100644 pkg/modelartifacts/path_test.go create mode 100644 pkg/modelartifacts/progress.go create mode 100644 pkg/modelartifacts/reference.go create mode 100644 pkg/modelartifacts/types.go create mode 100644 pkg/modelartifacts/types_test.go diff --git a/.impeccable/live/config.json b/.impeccable/live/config.json new file mode 100644 index 000000000..53a0a6fe9 --- /dev/null +++ b/.impeccable/live/config.json @@ -0,0 +1,6 @@ +{ + "files": ["core/http/react-ui/index.html"], + "insertBefore": "", + "commentSyntax": "html", + "cspChecked": true +} diff --git a/backend/python/common/model_utils.py b/backend/python/common/model_utils.py new file mode 100644 index 000000000..6c72ece72 --- /dev/null +++ b/backend/python/common/model_utils.py @@ -0,0 +1,28 @@ +import os + + +def resolve_model_reference(request, default=""): + model_file = (getattr(request, "ModelFile", "") or "").strip() + if model_file and os.path.exists(model_file): + return model_file, True + model = (getattr(request, "Model", "") or "").strip() + return model or default, False + + +def require_snapshot_file(model_ref, suffix): + if os.path.isfile(model_ref) and model_ref.endswith(suffix): + return model_ref + if not os.path.isdir(model_ref): + raise ValueError(f"model snapshot does not exist: {model_ref}") + + matches = [] + for root, directories, files in os.walk(model_ref): + directories.sort() + for name in sorted(files): + if name.endswith(suffix): + matches.append(os.path.join(root, name)) + if len(matches) != 1: + raise ValueError( + f"model snapshot must contain exactly one {suffix} file; found {len(matches)}" + ) + return matches[0] diff --git a/backend/python/common/model_utils_test.py b/backend/python/common/model_utils_test.py new file mode 100644 index 000000000..844eb972d --- /dev/null +++ b/backend/python/common/model_utils_test.py @@ -0,0 +1,81 @@ +import ast +import os +import tempfile +import unittest +from types import SimpleNamespace + +from model_utils import require_snapshot_file, resolve_model_reference + + +class ResolveModelReferenceTest(unittest.TestCase): + def test_existing_model_file_wins(self): + with tempfile.TemporaryDirectory() as directory: + request = SimpleNamespace(Model="owner/repo", ModelFile=directory) + self.assertEqual(resolve_model_reference(request), (directory, True)) + + def test_missing_model_file_preserves_legacy_repo_fallback(self): + request = SimpleNamespace(Model=" owner/repo ", ModelFile="/does/not/exist") + self.assertEqual(resolve_model_reference(request), ("owner/repo", False)) + + def test_empty_request_uses_explicit_default(self): + request = SimpleNamespace(Model=" ", ModelFile="") + self.assertEqual( + resolve_model_reference(request, "default/repo"), + ("default/repo", False), + ) + + def test_requires_the_only_matching_snapshot_file(self): + with tempfile.TemporaryDirectory() as directory: + checkpoint = os.path.join(directory, "model.nemo") + with open(checkpoint, "wb"): + pass + self.assertEqual(require_snapshot_file(directory, ".nemo"), checkpoint) + + def test_rejects_ambiguous_packaged_files(self): + with tempfile.TemporaryDirectory() as directory: + for name in ("first.nemo", "second.nemo"): + with open(os.path.join(directory, name), "wb"): + pass + with self.assertRaisesRegex(ValueError, "exactly one"): + require_snapshot_file(directory, ".nemo") + + +class MigratedBackendSourceTest(unittest.TestCase): + def test_general_backends_use_shared_resolution(self): + common_dir = os.path.dirname(__file__) + python_root = os.path.dirname(common_dir) + for backend in ( + "transformers", + "diffusers", + "qwen-asr", + "fish-speech", + "nemo", + "voxcpm", + "qwen-tts", + "liquid-audio", + "vllm", + "vllm-omni", + "sglang", + ): + with self.subTest(backend=backend): + with open( + os.path.join(python_root, backend, "backend.py"), encoding="utf-8" + ) as backend_file: + tree = ast.parse(backend_file.read()) + imports = { + alias.name + for node in ast.walk(tree) + if isinstance(node, ast.ImportFrom) and node.module == "model_utils" + for alias in node.names + } + calls = { + node.func.id + for node in ast.walk(tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + self.assertIn("resolve_model_reference", imports) + self.assertIn("resolve_model_reference", calls) + + +if __name__ == "__main__": + unittest.main() diff --git a/backend/python/diffusers/backend.py b/backend/python/diffusers/backend.py index 0299b9b4c..cb43c88c0 100755 --- a/backend/python/diffusers/backend.py +++ b/backend/python/diffusers/backend.py @@ -25,6 +25,7 @@ import grpc sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors +from model_utils import resolve_model_reference # Import dynamic loader for pipeline discovery @@ -215,7 +216,7 @@ def get_scheduler(name: str, config: dict = {}): # Implement the BackendServicer class with the service methods class BackendServicer(backend_pb2_grpc.BackendServicer): - def _load_pipeline(self, request, modelFile, fromSingleFile, torchType, variant, device_map=None): + def _load_pipeline(self, request, model_ref, from_single_file, local_only, torchType, variant, device_map=None): """ Load a diffusers pipeline dynamically using the dynamic loader. @@ -225,8 +226,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): Args: request: The gRPC request containing pipeline configuration - modelFile: Path to the model file (for single file loading) - fromSingleFile: Whether to use from_single_file() vs from_pretrained() + model_ref: Repository ID or local model file/directory + from_single_file: Whether to use from_single_file() vs from_pretrained() + local_only: Whether Hub-backed loaders must remain offline torchType: The torch dtype to use variant: Model variant (e.g., "fp16") device_map: Device mapping strategy (e.g., "auto" for multi-GPU) @@ -251,7 +253,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): dtype = torch.bfloat16 bfl_repo = os.environ.get("BFL_REPO", "ChuckMcSneed/FLUX.1-dev") - transformer = FluxTransformer2DModel.from_single_file(modelFile, torch_dtype=dtype, device_map=device_map) + transformer = FluxTransformer2DModel.from_single_file(model_ref, torch_dtype=dtype, device_map=device_map) quantize(transformer, weights=qfloat8) freeze(transformer) text_encoder_2 = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype, device_map=device_map) @@ -269,17 +271,19 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): # WanPipeline - requires special VAE with float32 dtype if pipeline_type == "WanPipeline": vae = AutoencoderKLWan.from_pretrained( - request.Model, + model_ref, subfolder="vae", torch_dtype=torch.float32, - device_map=device_map + device_map=device_map, + local_files_only=local_only, ) pipe = load_diffusers_pipeline( class_name="WanPipeline", - model_id=request.Model, + model_id=model_ref, vae=vae, torch_dtype=torchType, - device_map=device_map + device_map=device_map, + local_files_only=local_only, ) self.txt2vid = True return pipe @@ -287,17 +291,19 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): # WanImageToVideoPipeline - requires special VAE with float32 dtype if pipeline_type == "WanImageToVideoPipeline": vae = AutoencoderKLWan.from_pretrained( - request.Model, + model_ref, subfolder="vae", torch_dtype=torch.float32, - device_map=device_map + device_map=device_map, + local_files_only=local_only, ) pipe = load_diffusers_pipeline( class_name="WanImageToVideoPipeline", - model_id=request.Model, + model_id=model_ref, vae=vae, torch_dtype=torchType, - device_map=device_map + device_map=device_map, + local_files_only=local_only, ) self.img2vid = True return pipe @@ -306,10 +312,11 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): if pipeline_type == "SanaPipeline": pipe = load_diffusers_pipeline( class_name="SanaPipeline", - model_id=request.Model, + model_id=model_ref, variant="bf16", torch_dtype=torch.bfloat16, - device_map=device_map + device_map=device_map, + local_files_only=local_only, ) pipe.vae.to(torch.bfloat16) pipe.text_encoder.to(torch.bfloat16) @@ -320,9 +327,10 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): self.txt2vid = True pipe = load_diffusers_pipeline( class_name="DiffusionPipeline", - model_id=request.Model, + model_id=model_ref, torch_dtype=torchType, - device_map=device_map + device_map=device_map, + local_files_only=local_only, ) return pipe @@ -331,10 +339,11 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): self.img2vid = True pipe = load_diffusers_pipeline( class_name="StableVideoDiffusionPipeline", - model_id=request.Model, + model_id=model_ref, torch_dtype=torchType, variant=variant, - device_map=device_map + device_map=device_map, + local_files_only=local_only, ) if not DISABLE_CPU_OFFLOAD: pipe.enable_model_cpu_offload() @@ -346,8 +355,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): self.ltx2_pipeline = True # Check if loading from single file (GGUF) - if fromSingleFile and LTX2VideoTransformer3DModel is not None: - _, single_file_ext = os.path.splitext(modelFile) + if from_single_file and LTX2VideoTransformer3DModel is not None: + _, single_file_ext = os.path.splitext(model_ref) if single_file_ext == ".gguf": # Load transformer from single GGUF file with quantization transformer_kwargs = {} @@ -355,7 +364,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): transformer_kwargs["quantization_config"] = quantization_config transformer = LTX2VideoTransformer3DModel.from_single_file( - modelFile, + model_ref, config=request.Model, # Use request.Model as the config/model_id subfolder="transformer", device_map=device_map, @@ -374,7 +383,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): # Single file but not GGUF - use standard single file loading pipe = load_diffusers_pipeline( class_name="LTX2ImageToVideoPipeline", - model_id=modelFile, + model_id=model_ref, from_single_file=True, torch_dtype=torchType, device_map=device_map, @@ -383,10 +392,11 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): # Standard loading from pretrained pipe = load_diffusers_pipeline( class_name="LTX2ImageToVideoPipeline", - model_id=request.Model, + model_id=model_ref, torch_dtype=torchType, variant=variant, - device_map=device_map + device_map=device_map, + local_files_only=local_only, ) if not DISABLE_CPU_OFFLOAD: @@ -399,8 +409,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): self.ltx2_pipeline = True # Check if loading from single file (GGUF) - if fromSingleFile and LTX2VideoTransformer3DModel is not None: - _, single_file_ext = os.path.splitext(modelFile) + if from_single_file and LTX2VideoTransformer3DModel is not None: + _, single_file_ext = os.path.splitext(model_ref) if single_file_ext == ".gguf": # Load transformer from single GGUF file with quantization transformer_kwargs = {} @@ -408,7 +418,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): transformer_kwargs["quantization_config"] = quantization_config transformer = LTX2VideoTransformer3DModel.from_single_file( - modelFile, + model_ref, config=request.Model, # Use request.Model as the config/model_id subfolder="transformer", device_map=device_map, @@ -427,7 +437,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): # Single file but not GGUF - use standard single file loading pipe = load_diffusers_pipeline( class_name="LTX2Pipeline", - model_id=modelFile, + model_id=model_ref, from_single_file=True, torch_dtype=torchType, device_map=device_map, @@ -436,10 +446,11 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): # Standard loading from pretrained pipe = load_diffusers_pipeline( class_name="LTX2Pipeline", - model_id=request.Model, + model_id=model_ref, torch_dtype=torchType, variant=variant, - device_map=device_map + device_map=device_map, + local_files_only=local_only, ) if not DISABLE_CPU_OFFLOAD: @@ -455,12 +466,13 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): load_kwargs = {"torch_dtype": torchType} # Add variant if not loading from single file - if not fromSingleFile and variant: + if not from_single_file and variant: load_kwargs["variant"] = variant # Add use_safetensors for from_pretrained - if not fromSingleFile: + if not from_single_file: load_kwargs["use_safetensors"] = SAFETENSORS + load_kwargs["local_files_only"] = local_only # Add device_map for multi-GPU support (when TensorParallelSize > 1) if device_map: @@ -473,8 +485,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): try: pipe = load_diffusers_pipeline( class_name=effective_pipeline_type, - model_id=modelFile if fromSingleFile else request.Model, - from_single_file=fromSingleFile, + model_id=model_ref, + from_single_file=from_single_file, **load_kwargs ) except Exception as e: @@ -539,8 +551,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): print(f"Options: {self.options}", file=sys.stderr) - local = False - modelFile = request.Model + model_ref, local_only = resolve_model_reference(request) self.cfg_scale = 7 self.PipelineType = request.PipelineType @@ -555,13 +566,10 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): if request.CLIPSubfolder != "": clipsubfolder = request.CLIPSubfolder - # Check if ModelFile exists - if request.ModelFile != "": - if os.path.exists(request.ModelFile): - local = True - modelFile = request.ModelFile - - fromSingleFile = request.Model.startswith("http") or request.Model.startswith("/") or local + from_single_file = os.path.isfile(model_ref) or ( + not local_only + and (request.Model.startswith("http") or request.Model.startswith("/")) + ) self.img2vid = False self.txt2vid = False self.ltx2_pipeline = False @@ -579,8 +587,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): # Special cases that require custom initialization are handled first self.pipe = self._load_pipeline( request=request, - modelFile=modelFile, - fromSingleFile=fromSingleFile, + model_ref=model_ref, + from_single_file=from_single_file, + local_only=local_only, torchType=torchType, variant=variant, device_map=device_map diff --git a/backend/python/fish-speech/backend.py b/backend/python/fish-speech/backend.py index 9be8264a0..138396599 100644 --- a/backend/python/fish-speech/backend.py +++ b/backend/python/fish-speech/backend.py @@ -22,6 +22,7 @@ import grpc sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors +from model_utils import resolve_model_reference @@ -169,22 +170,17 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): else None ) - # Get model path from request - model_path = request.Model - if not model_path: - model_path = "fishaudio/s2-pro" - - # If model_path looks like a HuggingFace repo ID (e.g. "fishaudio/fish-speech-1.5"), - # download it locally first since fish-speech expects a local directory - if "/" in model_path and not os.path.exists(model_path): + model_path, local_only = resolve_model_reference( + request, "fishaudio/s2-pro" + ) + if not local_only and "/" in model_path: from huggingface_hub import snapshot_download print( - f"Downloading model from HuggingFace: {model_path}", + f"Downloading legacy unmanaged model from HuggingFace: {model_path}", file=sys.stderr, ) model_path = snapshot_download(repo_id=model_path) - print(f"Model downloaded to: {model_path}", file=sys.stderr) # Determine precision if device in ("mps", "cpu"): diff --git a/backend/python/liquid-audio/backend.py b/backend/python/liquid-audio/backend.py index 0a64ecb3e..bcc9165c1 100644 --- a/backend/python/liquid-audio/backend.py +++ b/backend/python/liquid-audio/backend.py @@ -23,6 +23,7 @@ import grpc sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors # noqa: E402 +from model_utils import resolve_model_reference # noqa: E402 from python_utils import parse_options # noqa: E402 import backend_pb2 # noqa: E402 @@ -172,19 +173,14 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): "fp32": torch.float32, }.get(dtype_name, torch.bfloat16) - # request.Model holds the raw `parameters.model` value (an HF - # repo id like "LiquidAI/LFM2.5-Audio-1.5B"); request.ModelFile - # is LocalAI's ModelPath-prefixed local copy that exists only - # when the gallery supplied a `files:` list. Mirror the - # transformers/vibevoice convention: prefer the repo id and - # only switch to the local path if it's been staged on disk. - model_id = request.Model - if not model_id: - model_id = request.ModelFile + # `_local_only` distinguishes a staged snapshot from the legacy + # repository fallback. liquid_audio has no matching Hub keyword, + # so the compatibility shim below enforces local loading instead. + model_id, _local_only = resolve_model_reference( + request, "LiquidAI/LFM2.5-Audio-1.5B" + ) if not model_id: return backend_pb2.Result(success=False, message="No model identifier provided") - if request.ModelFile and os.path.isdir(request.ModelFile): - model_id = request.ModelFile self.model_id = model_id # Pure fine-tune jobs don't need an in-memory inference model — the diff --git a/backend/python/nemo/backend.py b/backend/python/nemo/backend.py index a5c30694e..8b56f8d72 100644 --- a/backend/python/nemo/backend.py +++ b/backend/python/nemo/backend.py @@ -17,6 +17,7 @@ import grpc sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors +from model_utils import require_snapshot_file, resolve_model_reference @@ -70,11 +71,21 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): value = value.lower() == "true" self.options[key] = value - model_name = request.Model or "nvidia/parakeet-tdt-0.6b-v3" + model_ref, local_only = resolve_model_reference( + request, "nvidia/parakeet-tdt-0.6b-v3" + ) try: - print(f"Loading NEMO ASR model from {model_name}", file=sys.stderr) - self.model = nemo_asr.models.ASRModel.from_pretrained(model_name=model_name) + print(f"Loading NEMO ASR model from {model_ref}", file=sys.stderr) + if local_only: + checkpoint = require_snapshot_file(model_ref, ".nemo") + self.model = nemo_asr.models.ASRModel.restore_from( + restore_path=checkpoint + ) + else: + self.model = nemo_asr.models.ASRModel.from_pretrained( + model_name=model_ref + ) print("NEMO ASR model loaded successfully", file=sys.stderr) except Exception as err: print(f"[ERROR] LoadModel failed: {err}", file=sys.stderr) diff --git a/backend/python/qwen-asr/backend.py b/backend/python/qwen-asr/backend.py index 196f8f439..d576200a4 100644 --- a/backend/python/qwen-asr/backend.py +++ b/backend/python/qwen-asr/backend.py @@ -17,6 +17,7 @@ import grpc sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors +from model_utils import resolve_model_reference @@ -70,7 +71,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): value = value.lower() == "true" self.options[key] = value - model_path = request.Model or "Qwen/Qwen3-ASR-1.7B" + model_path, local_only = resolve_model_reference( + request, "Qwen/Qwen3-ASR-1.7B" + ) default_dtype = torch.bfloat16 if self.device == "cuda" else torch.float32 load_dtype = default_dtype if "torch_dtype" in self.options: @@ -117,6 +120,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): if attn_implementation: forced_aligner_kwargs["attn_implementation"] = attn_implementation load_kwargs["forced_aligner_kwargs"] = forced_aligner_kwargs + load_kwargs["local_files_only"] = local_only try: print(f"Loading Qwen3-ASR from {model_path}", file=sys.stderr) diff --git a/backend/python/qwen-tts/backend.py b/backend/python/qwen-tts/backend.py index 06b8084e0..18a3b1303 100644 --- a/backend/python/qwen-tts/backend.py +++ b/backend/python/qwen-tts/backend.py @@ -26,6 +26,7 @@ import grpc sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors +from model_utils import resolve_model_reference @@ -190,19 +191,23 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): ) print(traceback.format_exc(), file=sys.stderr) - # Get model path from request - model_path = request.Model - if not model_path: - model_path = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" + model_path, _local_only = resolve_model_reference( + request, "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" + ) + detection_ref = request.Model or model_path # Determine model type from model path or options self.model_type = self.options.get("model_type", None) if not self.model_type: - if "CustomVoice" in model_path: + if "CustomVoice" in detection_ref: self.model_type = "CustomVoice" - elif "VoiceDesign" in model_path: + elif "VoiceDesign" in detection_ref: self.model_type = "VoiceDesign" - elif "Base" in model_path or "0.6B" in model_path or "1.7B" in model_path: + elif ( + "Base" in detection_ref + or "0.6B" in detection_ref + or "1.7B" in detection_ref + ): self.model_type = "Base" # VoiceClone model else: # Default to CustomVoice diff --git a/backend/python/sglang/backend.py b/backend/python/sglang/backend.py index 0d38c6b7d..063079304 100644 --- a/backend/python/sglang/backend.py +++ b/backend/python/sglang/backend.py @@ -41,6 +41,7 @@ import grpc sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors +from model_utils import resolve_model_reference # sglang imports. Engine is the stable public entry point; parser modules # are wrapped in try/except so older / leaner installs that omit them @@ -173,7 +174,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): return backend_pb2.Reply(message=bytes("OK", 'utf-8')) async def LoadModel(self, request, context): - engine_kwargs = {"model_path": request.Model} + model_ref, local_only = resolve_model_reference(request) + engine_kwargs = {"model_path": model_ref} if request.Quantization: engine_kwargs["quantization"] = request.Quantization @@ -228,8 +230,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): if HAS_TRANSFORMERS: try: self.tokenizer = AutoTokenizer.from_pretrained( - request.Model, + model_ref, trust_remote_code=bool(request.TrustRemoteCode), + local_files_only=local_only, ) except Exception as err: print(f"AutoTokenizer load failed (non-fatal): {err!r}", file=sys.stderr) diff --git a/backend/python/transformers/backend.py b/backend/python/transformers/backend.py index 4b251c7f1..30d6334ba 100644 --- a/backend/python/transformers/backend.py +++ b/backend/python/transformers/backend.py @@ -19,6 +19,7 @@ import grpc sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors +from model_utils import resolve_model_reference import torch import torch.cuda @@ -61,11 +62,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): return backend_pb2.Reply(message=bytes("OK", 'utf-8')) def LoadModel(self, request, context): - model_name = request.Model - - # Check to see if the Model exists in the filesystem already. - if os.path.exists(request.ModelFile): - model_name = request.ModelFile + model_name, local_only = resolve_model_reference(request) compute = torch.float16 if request.F16Memory == True: @@ -150,7 +147,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): device_map=device_map, load_in_4bit=xpu_4bit, load_in_8bit=xpu_8bit, - torch_dtype=compute) + torch_dtype=compute, + local_files_only=local_only) elif request.Type == "OVModelForCausalLM": from optimum.intel.openvino import OVModelForCausalLM from openvino.runtime import Core @@ -171,7 +169,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): compile=True, trust_remote_code=request.TrustRemoteCode, ov_config=ovconfig, - device=device_map) + device=device_map, + local_files_only=local_only) self.OV = True elif request.Type == "OVModelForFeatureExtraction": from optimum.intel.openvino import OVModelForFeatureExtraction @@ -194,11 +193,16 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): trust_remote_code=request.TrustRemoteCode, ov_config=ovconfig, export=True, - device=device_map) + device=device_map, + local_files_only=local_only) self.OV = True elif request.Type == "SentenceTransformer": autoTokenizer = False - self.model = SentenceTransformer(model_name, trust_remote_code=request.TrustRemoteCode) + self.model = SentenceTransformer( + model_name, + trust_remote_code=request.TrustRemoteCode, + local_files_only=local_only, + ) self.SentenceTransformer = True elif request.Type == "TokenClassification": # NER / PII tagging via HuggingFace's token-classification @@ -213,6 +217,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): aggregation_strategy="simple", device=0 if self.CUDA else -1, trust_remote_code=request.TrustRemoteCode, + model_kwargs={"local_files_only": local_only}, ) self.TokenClassification = True else: @@ -231,6 +236,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): quantization_config=quantization, device_map=device_map, torch_dtype=compute, + local_files_only=local_only, ) # Try to load a processor (needed for TTS/audio models) @@ -238,6 +244,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): self.processor = AutoProcessor.from_pretrained( model_name, trust_remote_code=request.TrustRemoteCode, + local_files_only=local_only, ) self.GenericTTS = True print(f"Loaded processor for {model_name}", file=sys.stderr) @@ -252,7 +259,10 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): self.max_tokens = self.options.get("max_new_tokens", 512) if autoTokenizer: - self.tokenizer = AutoTokenizer.from_pretrained(model_name) + self.tokenizer = AutoTokenizer.from_pretrained( + model_name, + local_files_only=local_only, + ) self.XPU = False if XPU and self.OV == False: diff --git a/backend/python/vllm-omni/backend.py b/backend/python/vllm-omni/backend.py index e92b30c53..cc426e144 100644 --- a/backend/python/vllm-omni/backend.py +++ b/backend/python/vllm-omni/backend.py @@ -33,6 +33,7 @@ import grpc sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors +from model_utils import resolve_model_reference from vllm_utils import parse_options, messages_to_dicts, setup_parsers @@ -166,12 +167,16 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): # Parse options from request.Options using shared helper self.options = parse_options(request.Options) opts = self.options + model_ref, _local_only = resolve_model_reference(request) + detection_ref = request.Model or model_ref print(f"Options: {self.options}", file=sys.stderr) # Detect model type - self.model_name = request.Model - self.model_type = request.Type if request.Type else self._detect_model_type(request.Model) + self.model_name = model_ref + self.model_type = ( + request.Type if request.Type else self._detect_model_type(detection_ref) + ) print(f"Detected model type: {self.model_type}", file=sys.stderr) # Build DiffusionParallelConfig if diffusion model (image or video) @@ -205,9 +210,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): } # Base Omni initialization parameters - omni_kwargs = { - "model": request.Model, - } + omni_kwargs = {"model": model_ref} # Add diffusion-specific parameters (image/video models) if self.model_type in ["image", "video"]: @@ -251,7 +254,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): try: from vllm.transformers_utils.tokenizer import get_tokenizer self.tokenizer = get_tokenizer( - request.Model, + model_ref, trust_remote_code=opts.get("trust_remote_code", False), ) except Exception as e: diff --git a/backend/python/vllm/backend.py b/backend/python/vllm/backend.py index 7d85daf94..8a42e7bcb 100644 --- a/backend/python/vllm/backend.py +++ b/backend/python/vllm/backend.py @@ -21,6 +21,7 @@ import grpc sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors +from model_utils import resolve_model_reference from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.async_llm_engine import AsyncLLMEngine @@ -197,9 +198,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): Returns: backend_pb2.Result: The load model result. """ - engine_args = AsyncEngineArgs( - model=request.Model, - ) + model_ref, _local_only = resolve_model_reference(request) + engine_args = AsyncEngineArgs(model=model_ref) if request.Quantization != "": engine_args.quantization = request.Quantization @@ -260,7 +260,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): tokenizer = self.llm.tokenizer if tokenizer is None: tokenizer = get_tokenizer( - request.Model, + model_ref, trust_remote_code=bool(request.TrustRemoteCode), truncation_side="left", ) diff --git a/backend/python/voxcpm/backend.py b/backend/python/voxcpm/backend.py index 0c6d73545..84c0fae6e 100644 --- a/backend/python/voxcpm/backend.py +++ b/backend/python/voxcpm/backend.py @@ -21,6 +21,7 @@ import grpc sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors +from model_utils import resolve_model_reference def is_float(s): @@ -99,10 +100,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): value = value.lower() == "true" self.options[key] = value - # Get model path from request - model_path = request.Model - if not model_path: - model_path = "openbmb/VoxCPM1.5" + model_path, _local_only = resolve_model_reference( + request, "openbmb/VoxCPM1.5" + ) try: print(f"Loading model from {model_path}", file=sys.stderr) diff --git a/core/application/application.go b/core/application/application.go index 0d83c1a81..ab7f34a41 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -116,7 +116,11 @@ func newApplication(appConfig *config.ApplicationConfig) *Application { ml.SetLoadObserver(corebackend.ModelLoadTraceObserver(appConfig)) app := &Application{ - backendLoader: config.NewModelConfigLoader(appConfig.SystemState.Model.ModelsPath), + backendLoader: config.NewModelConfigLoader( + appConfig.SystemState.Model.ModelsPath, + config.WithArtifactMaterializer(appConfig.ModelArtifactMaterializer), + config.WithPreloadDisplay(appConfig.ModelPreloadRenderMode, appConfig.DisableModelPreloadColor), + ), modelLoader: ml, applicationConfig: appConfig, templatesEvaluator: templates.NewEvaluator(appConfig.SystemState.Model.ModelsPath), diff --git a/core/application/model_artifact_materializer_test.go b/core/application/model_artifact_materializer_test.go new file mode 100644 index 000000000..f6ad40a69 --- /dev/null +++ b/core/application/model_artifact_materializer_test.go @@ -0,0 +1,55 @@ +package application + +import ( + "context" + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/modelartifacts" + "github.com/mudler/LocalAI/pkg/system" +) + +type applicationLoaderMaterializer struct { + seen []modelartifacts.Spec +} + +func (f *applicationLoaderMaterializer) Ensure(_ context.Context, _ string, spec modelartifacts.Spec) (modelartifacts.Result, error) { + f.seen = append(f.seen, spec) + spec.Resolved = &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + } + return modelartifacts.Result{Spec: spec}, nil +} + +var _ = Describe("application model artifact wiring", func() { + It("injects the controller materializer into the model config loader", func() { + modelsPath := GinkgoT().TempDir() + state, err := system.GetSystemState( + system.WithModelPath(modelsPath), + system.WithBackendPath(GinkgoT().TempDir()), + ) + Expect(err).NotTo(HaveOccurred()) + materializer := &applicationLoaderMaterializer{} + app := newApplication(&config.ApplicationConfig{ + Context: context.Background(), + SystemState: state, + ModelArtifactMaterializer: materializer, + }) + Expect(os.WriteFile(filepath.Join(modelsPath, "managed.yaml"), []byte(` +name: managed +backend: transformers +artifacts: + - source: {type: huggingface, repo: owner/repo} +parameters: {model: owner/repo} +`), 0644)).To(Succeed()) + Expect(app.ModelConfigLoader().LoadModelConfigsFromPath(modelsPath)).To(Succeed()) + Expect(app.ModelConfigLoader().PreloadWithContext(context.Background(), modelsPath)).To(Succeed()) + Expect(materializer.seen).To(HaveLen(1)) + }) +}) diff --git a/core/application/startup.go b/core/application/startup.go index 47445e117..a1c6b8a09 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -414,18 +414,18 @@ func New(opts ...config.AppOption) (*Application, error) { } } - if err := application.ModelConfigLoader().Preload(options.SystemState.Model.ModelsPath); err != nil { + if err := application.ModelConfigLoader().PreloadWithContext(options.Context, options.SystemState.Model.ModelsPath); err != nil { xlog.Error("error downloading models", "error", err) } if options.PreloadJSONModels != "" { - if err := galleryop.ApplyGalleryFromString(options.SystemState, application.ModelLoader(), options.EnforcePredownloadScans, options.AutoloadBackendGalleries, options.Galleries, options.BackendGalleries, options.PreloadJSONModels, options.RequireBackendIntegrity); err != nil { + if err := galleryop.ApplyGalleryFromString(options.SystemState, application.ModelLoader(), options.EnforcePredownloadScans, options.AutoloadBackendGalleries, options.Galleries, options.BackendGalleries, options.PreloadJSONModels, options.RequireBackendIntegrity, gallery.WithArtifactMaterializer(options.ModelArtifactMaterializer)); err != nil { return nil, err } } if options.PreloadModelsFromPath != "" { - if err := galleryop.ApplyGalleryFromFile(options.SystemState, application.ModelLoader(), options.EnforcePredownloadScans, options.AutoloadBackendGalleries, options.Galleries, options.BackendGalleries, options.PreloadModelsFromPath, options.RequireBackendIntegrity); err != nil { + if err := galleryop.ApplyGalleryFromFile(options.SystemState, application.ModelLoader(), options.EnforcePredownloadScans, options.AutoloadBackendGalleries, options.Galleries, options.BackendGalleries, options.PreloadModelsFromPath, options.RequireBackendIntegrity, gallery.WithArtifactMaterializer(options.ModelArtifactMaterializer)); err != nil { return nil, err } } diff --git a/core/backend/llm.go b/core/backend/llm.go index 409e02e5d..6d56f40f8 100644 --- a/core/backend/llm.go +++ b/core/backend/llm.go @@ -109,7 +109,7 @@ func ModelInference(ctx context.Context, s string, messages schema.Messages, ima if !slices.Contains(modelNames, modelName) { utils.ResetDownloadTimers() // if we failed to load the model, we try to download it - err := gallery.InstallModelFromGallery(ctx, o.Galleries, o.BackendGalleries, o.SystemState, loader, modelName, gallery.GalleryModel{}, utils.DisplayDownloadFunction, o.EnforcePredownloadScans, o.AutoloadBackendGalleries, o.RequireBackendIntegrity) + err := gallery.InstallModelFromGallery(ctx, o.Galleries, o.BackendGalleries, o.SystemState, loader, modelName, gallery.GalleryModel{}, utils.DisplayDownloadFunction, o.EnforcePredownloadScans, o.AutoloadBackendGalleries, o.RequireBackendIntegrity, gallery.WithArtifactMaterializer(o.ModelArtifactMaterializer)) if err != nil { xlog.Error("failed to install model from gallery", "error", err, "model", modelFile) //return nil, err diff --git a/core/backend/model_artifact_options_test.go b/core/backend/model_artifact_options_test.go new file mode 100644 index 000000000..177e89825 --- /dev/null +++ b/core/backend/model_artifact_options_test.go @@ -0,0 +1,47 @@ +package backend + +import ( + "encoding/json" + "os" + "path/filepath" + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/modelartifacts" +) + +var _ = Describe("managed model runtime options", func() { + It("estimates weights from the committed manifest without repository fallback", func() { + const cacheKey = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" + modelsPath := GinkgoT().TempDir() + spec := modelartifacts.Spec{ + Name: "model", Target: "model", + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", Revision: "main"}, + Resolved: &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: cacheKey, + }, + } + relative, err := modelartifacts.RelativeSnapshotPath(cacheKey) + Expect(err).NotTo(HaveOccurred()) + manifest := modelartifacts.Manifest{ + Version: modelartifacts.ManifestVersion, + Artifact: spec, + Files: []modelartifacts.ManifestFile{{ + Path: "weights/model.safetensors", Size: 12, SHA256: strings.Repeat("a", 64), + }}, + } + encoded, err := json.Marshal(manifest) + Expect(err).NotTo(HaveOccurred()) + manifestPath := filepath.Join(modelsPath, filepath.Dir(relative), "manifest.json") + Expect(os.MkdirAll(filepath.Dir(manifestPath), 0750)).To(Succeed()) + Expect(os.WriteFile(manifestPath, encoded, 0644)).To(Succeed()) + + cfg := config.ModelConfig{Artifacts: []modelartifacts.Spec{spec}} + Expect(estimateModelSizeBytes(cfg, modelsPath)).To(Equal(int64(12))) + }) +}) diff --git a/core/backend/options.go b/core/backend/options.go index 29b84d924..5dd2d1ad5 100644 --- a/core/backend/options.go +++ b/core/backend/options.go @@ -12,9 +12,10 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/trace" - pb "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/LocalAI/pkg/downloader" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/vram" "github.com/mudler/xlog" ) @@ -93,8 +94,9 @@ func recordModelLoadFailure(appConfig *config.ApplicationConfig, modelName, back func estimateModelSizeBytes(c config.ModelConfig, modelsPath string) int64 { seen := make(map[string]bool) input := vram.ModelEstimateInput{} + managedPrimary := len(c.Artifacts) > 0 && c.Artifacts[0].Resolved != nil - addFile := func(uri string) { + addFile := func(uri string, size int64) { if !vram.IsWeightFile(uri) { return } @@ -106,7 +108,7 @@ func estimateModelSizeBytes(c config.ModelConfig, modelsPath string) int64 { return } seen[resolved] = true - input.Files = append(input.Files, vram.FileInput{URI: resolved}) + input.Files = append(input.Files, vram.FileInput{URI: resolved, Size: size}) } // tryHFRepo resolves any huggingface:// or hf:// URI to an HTTPS URL and @@ -123,13 +125,25 @@ func estimateModelSizeBytes(c config.ModelConfig, modelsPath string) int64 { for _, f := range c.DownloadFiles { uriStr := string(f.URI) - addFile(uriStr) - tryHFRepo(uriStr) + addFile(uriStr, 0) + if !managedPrimary { + tryHFRepo(uriStr) + } + } + if managedPrimary { + relative := c.ModelFileName() + manifest, err := modelartifacts.ReadManifest(filepath.Join(modelsPath, filepath.Dir(relative), "manifest.json")) + if err == nil { + for _, file := range manifest.Files { + addFile(filepath.Join(relative, filepath.FromSlash(file.Path)), file.Size) + } + } + } else { + addFile(c.Model, 0) + tryHFRepo(c.Model) } - addFile(c.Model) - tryHFRepo(c.Model) if c.MMProj != "" { - addFile(c.MMProj) + addFile(c.MMProj, 0) } if len(input.Files) == 0 && input.HFRepo == "" { @@ -153,6 +167,10 @@ func ModelOptions(c config.ModelConfig, so *config.ApplicationConfig, opts ...mo model.WithContext(so.Context), model.WithModelID(c.ModelID()), } + managedPrimary := len(c.Artifacts) > 0 && c.Artifacts[0].Resolved != nil + if managedPrimary { + defOpts = append(defOpts, model.WithModelFile(c.ModelFileName())) + } threads := 1 diff --git a/core/cli/models.go b/core/cli/models.go index 76d23aebd..a95462da1 100644 --- a/core/cli/models.go +++ b/core/cli/models.go @@ -14,6 +14,7 @@ import ( "github.com/mudler/LocalAI/core/startup" "github.com/mudler/LocalAI/pkg/downloader" "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/xlog" "github.com/schollz/progressbar/v3" @@ -24,6 +25,9 @@ type ModelsCMDFlags struct { BackendGalleries string `env:"LOCALAI_BACKEND_GALLERIES,BACKEND_GALLERIES" help:"JSON list of backend galleries" group:"backends" default:"${backends}"` ModelsPath string `env:"LOCALAI_MODELS_PATH,MODELS_PATH" type:"path" default:"${basepath}/models" help:"Path containing models used for inferencing" group:"storage"` BackendsPath string `env:"LOCALAI_BACKENDS_PATH,BACKENDS_PATH" type:"path" default:"${basepath}/backends" help:"Path containing backends used for inferencing" group:"storage"` + Color string `env:"COLOR" hidden:""` + NoColor string `env:"NO_COLOR" hidden:""` + HFToken string `env:"HF_TOKEN" hidden:""` } type ModelsList struct { @@ -80,10 +84,18 @@ func (mi *ModelsInstall) Run(ctx *cliContext.Context) error { return err } + artifactMaterializer := modelartifacts.NewDefaultManager( + modelartifacts.WithHuggingFaceToken(mi.HFToken), + ) galleryService := galleryop.NewGalleryService(&config.ApplicationConfig{ - SystemState: systemState, + SystemState: systemState, + ModelArtifactMaterializer: artifactMaterializer, + ModelPreloadRenderMode: mi.Color, + DisableModelPreloadColor: mi.NoColor != "", }, model.NewModelLoader(systemState)) - err = galleryService.Start(context.Background(), config.NewModelConfigLoader(mi.ModelsPath), systemState) + err = galleryService.Start(context.Background(), config.NewModelConfigLoader(mi.ModelsPath, + config.WithArtifactMaterializer(artifactMaterializer), + config.WithPreloadDisplay(mi.Color, mi.NoColor != "")), systemState) if err != nil { return err } diff --git a/core/cli/run.go b/core/cli/run.go index 542d441c1..b067098b1 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -16,6 +16,7 @@ import ( "github.com/mudler/LocalAI/core/http" "github.com/mudler/LocalAI/core/p2p" "github.com/mudler/LocalAI/internal" + "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/signals" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/xlog" @@ -28,6 +29,9 @@ import ( type RunCMD struct { ModelArgs []string `arg:"" optional:"" name:"models" help:"Model configuration URLs to load"` + Color string `env:"COLOR" hidden:""` + NoColor string `env:"NO_COLOR" hidden:""` + HFToken string `env:"HF_TOKEN" hidden:""` ExternalBackends []string `env:"LOCALAI_EXTERNAL_BACKENDS,EXTERNAL_BACKENDS" help:"A list of external backends to load from gallery on boot" group:"backends"` WebRTCNAT1To1IPs []string `env:"LOCALAI_WEBRTC_NAT_1TO1_IPS,WEBRTC_NAT_1TO1_IPS" help:"IPs advertised as the host ICE candidates for /v1/realtime WebRTC instead of every local interface. Set to the reachable host/LAN IP when running under Docker host networking or NAT, where pion otherwise offers unreachable bridge addresses and the connection drops after ICE consent checks fail." group:"api"` @@ -238,6 +242,10 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { opts := []config.AppOption{ config.WithContext(context.Background()), + config.WithModelArtifactMaterializer(modelartifacts.NewDefaultManager( + modelartifacts.WithHuggingFaceToken(r.HFToken), + )), + config.WithModelPreloadDisplay(r.Color, r.NoColor != ""), config.WithConfigFile(r.ModelsConfigFile), config.WithJSONStringPreload(r.PreloadModels), config.WithYAMLConfigPreload(r.PreloadModelsConfig), diff --git a/core/config/application_artifact_materializer_test.go b/core/config/application_artifact_materializer_test.go new file mode 100644 index 000000000..13a9d856b --- /dev/null +++ b/core/config/application_artifact_materializer_test.go @@ -0,0 +1,34 @@ +package config + +import ( + "context" + "reflect" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/modelartifacts" +) + +type applicationArtifactMaterializer struct{} + +func (*applicationArtifactMaterializer) Ensure(context.Context, string, modelartifacts.Spec) (modelartifacts.Result, error) { + return modelartifacts.Result{}, nil +} + +var _ = Describe("ApplicationConfig model artifact materializer", func() { + It("provides a default materializer", func() { + Expect(NewApplicationConfig().ModelArtifactMaterializer).NotTo(BeNil()) + }) + + It("accepts an injected materializer without exposing it to serialization", func() { + materializer := &applicationArtifactMaterializer{} + appConfig := NewApplicationConfig(WithModelArtifactMaterializer(materializer)) + Expect(appConfig.ModelArtifactMaterializer).To(BeIdenticalTo(materializer)) + + field, found := reflect.TypeOf(ApplicationConfig{}).FieldByName("ModelArtifactMaterializer") + Expect(found).To(BeTrue()) + Expect(field.Tag.Get("json")).To(Equal("-")) + Expect(field.Tag.Get("yaml")).To(Equal("-")) + }) +}) diff --git a/core/config/application_config.go b/core/config/application_config.go index 2c863beb6..047a16703 100644 --- a/core/config/application_config.go +++ b/core/config/application_config.go @@ -7,6 +7,7 @@ import ( "time" "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/LocalAI/pkg/xsysinfo" "github.com/mudler/xlog" @@ -18,6 +19,13 @@ type ApplicationConfig struct { SystemState *system.SystemState ExternalBackends []string + // ModelArtifactMaterializer is a controller-only acquisition capability. + // It is excluded from serialization so credentials captured by its concrete + // implementation cannot enter persisted settings or distributed payloads. + ModelArtifactMaterializer ArtifactMaterializer `json:"-" yaml:"-"` + ModelPreloadRenderMode string `json:"-" yaml:"-"` + DisableModelPreloadColor bool `json:"-" yaml:"-"` + // WebRTCNAT1To1IPs, when set, are advertised as the host ICE candidates for // /v1/realtime WebRTC instead of every local interface address. Needed when // the routable address differs from what pion gathers — e.g. Docker host @@ -246,9 +254,11 @@ type AppOption func(*ApplicationConfig) func NewApplicationConfig(o ...AppOption) *ApplicationConfig { opt := &ApplicationConfig{ - Context: context.Background(), - UploadLimitMB: 15, - Debug: true, + Context: context.Background(), + UploadLimitMB: 15, + Debug: true, + ModelArtifactMaterializer: modelartifacts.NewDefaultManager(), + ModelPreloadRenderMode: "dark", // Capture backend process stdout/stderr into the per-model // BackendLogStore by default so the UI "Backend Logs" page works out // of the box in single mode, matching worker/distributed mode (which @@ -256,9 +266,9 @@ func NewApplicationConfig(o ...AppOption) *ApplicationConfig { // toggle can still turn it off (a persisted false wins - see // loadRuntimeSettingsFromFile). EnableBackendLogging: true, - AgentJobRetentionDays: 30, // Default: 30 days - LRUEvictionMaxRetries: 30, // Default: 30 retries - LRUEvictionRetryInterval: 1 * time.Second, // Default: 1 second + AgentJobRetentionDays: 30, // Default: 30 days + LRUEvictionMaxRetries: 30, // Default: 30 retries + LRUEvictionRetryInterval: 1 * time.Second, // Default: 1 second ModelLoadFailureCooldown: 10 * time.Second, // Default: 10s base cooldown after a failed load // WatchDogInterval is intentionally left at the zero value here. // The startup loader applies a persisted runtime_settings.json value @@ -636,6 +646,25 @@ func WithContext(ctx context.Context) AppOption { } } +// WithModelArtifactMaterializer injects the controller-only model acquisition capability. +func WithModelArtifactMaterializer(materializer ArtifactMaterializer) AppOption { + return func(o *ApplicationConfig) { + if materializer != nil { + o.ModelArtifactMaterializer = materializer + } + } +} + +// WithModelPreloadDisplay configures terminal rendering for model preload output. +func WithModelPreloadDisplay(renderMode string, disableColor bool) AppOption { + return func(o *ApplicationConfig) { + if renderMode != "" { + o.ModelPreloadRenderMode = renderMode + } + o.DisableModelPreloadColor = disableColor + } +} + func WithYAMLConfigPreload(configFile string) AppOption { return func(o *ApplicationConfig) { o.PreloadModelsFromPath = configFile diff --git a/core/config/meta/registry.go b/core/config/meta/registry.go index e30bc7d33..2ad64f4d3 100644 --- a/core/config/meta/registry.go +++ b/core/config/meta/registry.go @@ -82,6 +82,14 @@ func DefaultRegistry() map[string]FieldMetaOverride { Options: ModalityOptions, Order: 8, }, + "artifacts": { + Section: "general", + Label: "Managed Model Artifacts", + Description: "Controller-managed model sources. LocalAI resolves, downloads, verifies, and binds these sources before a backend loads.", + Component: "json-editor", + Advanced: true, + Order: 9, + }, // --- LLM --- "context_size": { diff --git a/core/config/model_artifact_binding.go b/core/config/model_artifact_binding.go new file mode 100644 index 000000000..416ad3e73 --- /dev/null +++ b/core/config/model_artifact_binding.go @@ -0,0 +1,110 @@ +package config + +import ( + "fmt" + "os" + "path/filepath" + + "gopkg.in/yaml.v3" + + "github.com/mudler/LocalAI/pkg/modelartifacts" +) + +func persistArtifactBinding(fileName, modelName string, result modelartifacts.Result) error { + data, err := os.ReadFile(fileName) + if err != nil { + return err + } + var document yaml.Node + if err := yaml.Unmarshal(data, &document); err != nil { + return err + } + target, err := findModelMapping(&document, modelName) + if err != nil { + return err + } + artifactValue := &yaml.Node{} + encoded, err := yaml.Marshal([]modelartifacts.Spec{result.Spec}) + if err != nil { + return err + } + if err := yaml.Unmarshal(encoded, artifactValue); err != nil { + return err + } + setMappingValue(target, "artifacts", artifactValue.Content[0]) + updated, err := yaml.Marshal(&document) + if err != nil { + return err + } + return writeBindingAtomic(fileName, updated) +} + +func findModelMapping(document *yaml.Node, modelName string) (*yaml.Node, error) { + if len(document.Content) != 1 { + return nil, fmt.Errorf("invalid model configuration document") + } + root := document.Content[0] + if root.Kind == yaml.MappingNode { + return root, nil + } + if root.Kind == yaml.SequenceNode { + for _, candidate := range root.Content { + if candidate.Kind == yaml.MappingNode { + name := mappingValue(candidate, "name") + if name != nil && name.Value == modelName { + return candidate, nil + } + } + } + } + return nil, fmt.Errorf("model %q not found in configuration document", modelName) +} + +func mappingValue(mapping *yaml.Node, key string) *yaml.Node { + for index := 0; index+1 < len(mapping.Content); index += 2 { + if mapping.Content[index].Value == key { + return mapping.Content[index+1] + } + } + return nil +} + +func setMappingValue(mapping *yaml.Node, key string, value *yaml.Node) { + for index := 0; index+1 < len(mapping.Content); index += 2 { + if mapping.Content[index].Value == key { + mapping.Content[index+1] = value + return + } + } + mapping.Content = append(mapping.Content, + &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key}, value, + ) +} + +func writeBindingAtomic(fileName string, data []byte) error { + temporary, err := os.CreateTemp(filepath.Dir(fileName), ".artifact-binding-*") + if err != nil { + return err + } + temporaryName := temporary.Name() + defer func() { _ = os.Remove(temporaryName) }() + if err := temporary.Chmod(0600); err != nil { + _ = temporary.Close() + return err + } + if _, err := temporary.Write(data); err != nil { + _ = temporary.Close() + return err + } + if err := temporary.Sync(); err != nil { + _ = temporary.Close() + return err + } + if err := temporary.Close(); err != nil { + return err + } + if err := os.Chmod(temporaryName, 0644); err != nil { + return err + } + return os.Rename(temporaryName, fileName) +} diff --git a/core/config/model_artifact_binding_test.go b/core/config/model_artifact_binding_test.go new file mode 100644 index 000000000..c61228d04 --- /dev/null +++ b/core/config/model_artifact_binding_test.go @@ -0,0 +1,50 @@ +package config + +import ( + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/modelartifacts" +) + +var _ = Describe("artifact binding persistence", func() { + It("updates only the named model in a multi-config document", func() { + fileName := filepath.Join(GinkgoT().TempDir(), "models.yaml") + Expect(os.WriteFile(fileName, []byte(` +- name: managed + backend: transformers + unknown_extension: keep-me + artifacts: + - source: {type: huggingface, repo: owner/repo} + parameters: {model: owner/repo} +- name: sibling + backend: llama-cpp + sibling_only: true + parameters: {model: sibling.gguf} +`), 0644)).To(Succeed()) + result := modelartifacts.Result{ + RelativePath: ".artifacts/huggingface/0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef/snapshot", + Spec: modelartifacts.Spec{ + Name: "model", Target: "model", + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", Revision: "main"}, + Resolved: &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + }, + }, + } + Expect(persistArtifactBinding(fileName, "managed", result)).To(Succeed()) + updated, err := os.ReadFile(fileName) + Expect(err).NotTo(HaveOccurred()) + Expect(string(updated)).To(ContainSubstring("name: sibling")) + Expect(string(updated)).To(ContainSubstring("sibling_only: true")) + Expect(string(updated)).To(ContainSubstring("unknown_extension: keep-me")) + Expect(string(updated)).To(ContainSubstring("model: owner/repo")) + Expect(string(updated)).To(ContainSubstring("cache_key: 0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef")) + Expect(string(updated)).To(ContainSubstring("revision: 0123456789abcdef0123456789abcdef01234567")) + }) +}) diff --git a/core/config/model_artifact_inference.go b/core/config/model_artifact_inference.go new file mode 100644 index 000000000..af9c00d21 --- /dev/null +++ b/core/config/model_artifact_inference.go @@ -0,0 +1,61 @@ +package config + +import ( + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/mudler/LocalAI/pkg/modelartifacts" +) + +// PrimaryArtifactSpec returns the managed primary artifact to materialize for +// this config. The boolean return is false when the config should stay on the +// legacy path. +func (c ModelConfig) PrimaryArtifactSpec(modelsPath string) (modelartifacts.Spec, bool, bool, error) { + if len(c.Artifacts) > 0 { + return c.Artifacts[0], false, true, nil + } + + modelRef := strings.TrimSpace(c.Model) + if modelRef == "" { + return modelartifacts.Spec{}, false, false, nil + } + if len(c.DownloadFiles) > 0 { + return modelartifacts.Spec{}, false, false, nil + } + + if modelsPath != "" { + for _, candidate := range []string{ + modelRef, + filepath.Join(modelsPath, modelRef), + } { + if info, err := os.Stat(candidate); err == nil && (info.IsDir() || info.Mode().IsRegular()) { + return modelartifacts.Spec{}, false, false, nil + } + } + } + + spec, ok, err := modelartifacts.ParsePrimaryReference(modelRef) + if err != nil { + return modelartifacts.Spec{}, false, false, err + } + if !ok { + if strings.Contains(modelRef, "/") && !strings.HasPrefix(modelRef, ".") && !filepath.IsAbs(modelRef) { + repo, err := modelartifacts.CanonicalRepo(modelRef) + if err != nil { + if strings.Count(modelRef, "/") == 1 { + return modelartifacts.Spec{}, false, false, fmt.Errorf("invalid Hugging Face reference %q: %w", modelRef, err) + } + return modelartifacts.Spec{}, false, false, nil + } + return modelartifacts.Spec{ + Name: modelartifacts.TargetModel, + Target: modelartifacts.TargetModel, + Source: modelartifacts.Source{Type: modelartifacts.SourceTypeHuggingFace, Repo: repo}, + }, true, true, nil + } + return modelartifacts.Spec{}, false, false, nil + } + return spec, true, true, nil +} diff --git a/core/config/model_config.go b/core/config/model_config.go index 6541e2e3a..a71a2c683 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -13,6 +13,7 @@ import ( "github.com/mudler/LocalAI/core/services/routing/piipattern" "github.com/mudler/LocalAI/pkg/downloader" "github.com/mudler/LocalAI/pkg/functions" + "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/reasoning" "github.com/mudler/cogito" "gopkg.in/yaml.v3" @@ -40,7 +41,8 @@ type ModelConfig struct { modelConfigFile string `yaml:"-" json:"-"` modelTemplate string `yaml:"-" json:"-"` schema.PredictionOptions `yaml:"parameters,omitempty" json:"parameters,omitempty"` - Name string `yaml:"name,omitempty" json:"name,omitempty"` + Name string `yaml:"name,omitempty" json:"name,omitempty"` + Artifacts []modelartifacts.Spec `yaml:"artifacts,omitempty" json:"artifacts,omitempty"` // Alias, when set, makes this config a pure redirect: every request for // Name is served by the model named here. All other fields are ignored. @@ -1210,9 +1212,15 @@ func (c ModelConfig) ModelID() string { return c.Model } -// ModelFileName returns the filename of the model -// If the model is a URL, it will return the MD5 of the URL which is the filename +// ModelFileName returns the controller-managed snapshot when the model has a +// committed artifact, otherwise preserving the legacy URL/repository behavior. func (c *ModelConfig) ModelFileName() string { + if len(c.Artifacts) > 0 && c.Artifacts[0].Resolved != nil { + relative, err := modelartifacts.RelativeSnapshotPath(c.Artifacts[0].Resolved.CacheKey) + if err == nil { + return relative + } + } uri := downloader.URI(c.Model) if uri.LooksLikeURL() { f, _ := uri.FilenameFromUrl() @@ -1315,6 +1323,21 @@ func (cfg *ModelConfig) SetDefaults(opts ...ConfigLoaderOption) { } func (c *ModelConfig) Validate() (bool, error) { + if c.IsAlias() && len(c.Artifacts) > 0 { + return false, fmt.Errorf("alias model %q cannot declare artifacts", c.Name) + } + seenArtifacts := make(map[string]struct{}, len(c.Artifacts)) + for i, artifact := range c.Artifacts { + normalized, err := artifact.Normalize() + if err != nil { + return false, fmt.Errorf("artifact %d: %w", i, err) + } + if _, exists := seenArtifacts[normalized.Name]; exists { + return false, fmt.Errorf("duplicate artifact name %q", normalized.Name) + } + seenArtifacts[normalized.Name] = struct{}{} + } + // An alias is a pure redirect: validate only its own shape here. Target // existence and the no-chain rule need the full config set, so the loader // (load-time) and the create/swap endpoints enforce those. diff --git a/core/config/model_config_loader.go b/core/config/model_config_loader.go index 29319cbec..4423fddad 100644 --- a/core/config/model_config_loader.go +++ b/core/config/model_config_loader.go @@ -2,11 +2,13 @@ package config import ( "cmp" + "context" "errors" "fmt" "io/fs" "os" "path/filepath" + "reflect" "slices" "strings" "sync" @@ -14,22 +16,59 @@ import ( "github.com/charmbracelet/glamour" "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/pkg/downloader" + "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/utils" "github.com/mudler/xlog" "gopkg.in/yaml.v3" ) +// ArtifactMaterializer resolves and commits a model artifact into local storage. +type ArtifactMaterializer interface { + Ensure(context.Context, string, modelartifacts.Spec) (modelartifacts.Result, error) +} + +// ModelConfigLoaderOption customizes a ModelConfigLoader at construction time. +type ModelConfigLoaderOption func(*ModelConfigLoader) + +// WithArtifactMaterializer sets the controller-side artifact acquisition boundary. +func WithArtifactMaterializer(materializer ArtifactMaterializer) ModelConfigLoaderOption { + return func(loader *ModelConfigLoader) { + if materializer != nil { + loader.artifactMaterializer = materializer + } + } +} + +// WithPreloadDisplay configures terminal rendering without reading process state. +func WithPreloadDisplay(renderMode string, disableColor bool) ModelConfigLoaderOption { + return func(loader *ModelConfigLoader) { + if renderMode != "" { + loader.preloadRenderMode = renderMode + } + loader.disablePreloadColor = disableColor + } +} + type ModelConfigLoader struct { - configs map[string]ModelConfig - modelPath string + configs map[string]ModelConfig + modelPath string + artifactMaterializer ArtifactMaterializer + preloadRenderMode string + disablePreloadColor bool sync.Mutex } -func NewModelConfigLoader(modelPath string) *ModelConfigLoader { - return &ModelConfigLoader{ - configs: make(map[string]ModelConfig), - modelPath: modelPath, +func NewModelConfigLoader(modelPath string, options ...ModelConfigLoaderOption) *ModelConfigLoader { + loader := &ModelConfigLoader{ + configs: make(map[string]ModelConfig), + modelPath: modelPath, + artifactMaterializer: modelartifacts.NewDefaultManager(), + preloadRenderMode: "dark", } + for _, option := range options { + option(loader) + } + return loader } type LoadOptions struct { @@ -351,98 +390,170 @@ func (bcl *ModelConfigLoader) ValidateAliasTarget(cfg *ModelConfig) error { return nil } -// Preload prepare models if they are not local but url or huggingface repositories +type preloadWork struct { + key string + config ModelConfig +} + +// Preload prepares models if they are not local but URLs or Hugging Face repositories. func (bcl *ModelConfigLoader) Preload(modelPath string) error { + return bcl.PreloadWithContext(context.Background(), modelPath) +} + +// PreloadWithContext prepares remote model inputs while honoring cancellation. +func (bcl *ModelConfigLoader) PreloadWithContext(ctx context.Context, modelPath string) error { bcl.Lock() - defer bcl.Unlock() + work := make([]preloadWork, 0, len(bcl.configs)) + for key, config := range bcl.configs { + configCopy := config + configCopy.Artifacts = slices.Clone(config.Artifacts) + configCopy.DownloadFiles = slices.Clone(config.DownloadFiles) + work = append(work, preloadWork{key: key, config: configCopy}) + } + bcl.Unlock() status := func(fileName, current, total string, percent float64) { utils.DisplayDownloadFunction(fileName, current, total, percent) } xlog.Info("Preloading models", "path", modelPath) + for _, item := range work { + if err := ctx.Err(); err != nil { + return err + } + updated, artifactResult, err := bcl.preloadOne(ctx, modelPath, item.config, status) + if err != nil { + return err + } + if err := ctx.Err(); err != nil { + return err + } - renderMode := "dark" - if os.Getenv("COLOR") != "" { - renderMode = os.Getenv("COLOR") + bcl.Lock() + current, exists := bcl.configs[item.key] + if !exists || !reflect.DeepEqual(current, item.config) { + bcl.Unlock() + continue + } + if artifactResult != nil && bindingNeedsPersistence(current, *artifactResult) && current.modelConfigFile != "" { + modelartifacts.ReportProgress(ctx, modelartifacts.ProgressEvent{ + Phase: modelartifacts.PhasePersisting, + Artifact: artifactResult.Spec.Name, + }) + if err := persistArtifactBinding(current.modelConfigFile, current.Name, *artifactResult); err != nil { + bcl.Unlock() + return err + } + } + bcl.configs[item.key] = updated + bcl.Unlock() + bcl.displayPreloadedModel(updated) + } + return nil +} + +func (bcl *ModelConfigLoader) preloadOne( + ctx context.Context, + modelPath string, + config ModelConfig, + status func(string, string, string, float64), +) (ModelConfig, *modelartifacts.Result, error) { + updated := config + updated.Artifacts = slices.Clone(config.Artifacts) + tasks := make([]downloader.FileTask, 0, len(updated.DownloadFiles)) + for index, file := range updated.DownloadFiles { + if err := ctx.Err(); err != nil { + return ModelConfig{}, nil, err + } + xlog.Debug("Checking file exists and matches SHA", "filename", file.Filename) + if err := utils.VerifyPath(file.Filename, modelPath); err != nil { + return ModelConfig{}, nil, err + } + tasks = append(tasks, downloader.FileTask{ + URI: file.URI, + Destination: filepath.Join(modelPath, file.Filename), + SHA256: file.SHA256, + FileIndex: index, + TotalFiles: len(updated.DownloadFiles), + }) + } + if err := downloader.DownloadFilesWithContext(ctx, tasks, status); err != nil { + return ModelConfig{}, nil, err } + artifactSpec, inferred, managedPrimary, err := updated.PrimaryArtifactSpec(modelPath) + if err != nil { + return ModelConfig{}, nil, err + } + var artifactResult *modelartifacts.Result + if managedPrimary { + result, err := bcl.artifactMaterializer.Ensure(ctx, modelPath, artifactSpec) + if err != nil { + if inferred { + xlog.Warn("falling back to legacy model loading after artifact materialization failed", "model", updated.Name, "error", err) + managedPrimary = false + } else { + return ModelConfig{}, nil, err + } + } else { + next := []modelartifacts.Spec{result.Spec} + if len(updated.Artifacts) > 1 { + next = append(next, updated.Artifacts[1:]...) + } + updated.Artifacts = next + artifactResult = &result + } + } + + if !managedPrimary && updated.IsModelURL() { + modelFileName := updated.ModelFileName() + uri := downloader.URI(updated.Model) + if uri.ResolveURL() != updated.Model { + if _, err := os.Stat(filepath.Join(modelPath, modelFileName)); errors.Is(err, os.ErrNotExist) { + if err := uri.DownloadFileWithContext(ctx, filepath.Join(modelPath, modelFileName), "", 0, 0, status); err != nil { + return ModelConfig{}, nil, err + } + } + updated.Model = modelFileName + } + } + + if updated.IsMMProjURL() { + modelFileName := updated.MMProjFileName() + uri := downloader.URI(updated.MMProj) + if _, err := os.Stat(filepath.Join(modelPath, modelFileName)); errors.Is(err, os.ErrNotExist) { + if err := uri.DownloadFileWithContext(ctx, filepath.Join(modelPath, modelFileName), "", 0, 0, status); err != nil { + return ModelConfig{}, nil, err + } + } + updated.MMProj = modelFileName + } + return updated, artifactResult, nil +} + +func bindingNeedsPersistence(current ModelConfig, result modelartifacts.Result) bool { + return len(current.Artifacts) == 0 || !reflect.DeepEqual(current.Artifacts[0], result.Spec) +} + +func (bcl *ModelConfigLoader) displayPreloadedModel(config ModelConfig) { glamText := func(t string) { - out, err := glamour.Render(t, renderMode) - if err == nil && os.Getenv("NO_COLOR") == "" { + out, err := glamour.Render(t, bcl.preloadRenderMode) + if err == nil && !bcl.disablePreloadColor { fmt.Println(out) } else { fmt.Println(t) } } - for i, config := range bcl.configs { - - // Download files and verify their SHA - for i, file := range config.DownloadFiles { - xlog.Debug("Checking file exists and matches SHA", "filename", file.Filename) - - if err := utils.VerifyPath(file.Filename, modelPath); err != nil { - return err - } - // Create file path - filePath := filepath.Join(modelPath, file.Filename) - - if err := file.URI.DownloadFile(filePath, file.SHA256, i, len(config.DownloadFiles), status); err != nil { - return err - } - } - - // If the model is an URL, expand it, and download the file - if config.IsModelURL() { - modelFileName := config.ModelFileName() - uri := downloader.URI(config.Model) - if uri.ResolveURL() != config.Model { - // check if file exists - if _, err := os.Stat(filepath.Join(modelPath, modelFileName)); errors.Is(err, os.ErrNotExist) { - err := uri.DownloadFile(filepath.Join(modelPath, modelFileName), "", 0, 0, status) - if err != nil { - return err - } - } - - cc := bcl.configs[i] - c := &cc - c.PredictionOptions.Model = modelFileName - bcl.configs[i] = *c - } - } - - if config.IsMMProjURL() { - modelFileName := config.MMProjFileName() - uri := downloader.URI(config.MMProj) - // check if file exists - if _, err := os.Stat(filepath.Join(modelPath, modelFileName)); errors.Is(err, os.ErrNotExist) { - err := uri.DownloadFile(filepath.Join(modelPath, modelFileName), "", 0, 0, status) - if err != nil { - return err - } - } - - cc := bcl.configs[i] - c := &cc - c.MMProj = modelFileName - bcl.configs[i] = *c - } - - if bcl.configs[i].Name != "" { - glamText(fmt.Sprintf("**Model name**: _%s_", bcl.configs[i].Name)) - } - if bcl.configs[i].Description != "" { - //glamText("**Description**") - glamText(bcl.configs[i].Description) - } - if bcl.configs[i].Usage != "" { - //glamText("**Usage**") - glamText(bcl.configs[i].Usage) - } + if config.Name != "" { + glamText(fmt.Sprintf("**Model name**: _%s_", config.Name)) + } + if config.Description != "" { + glamText(config.Description) + } + if config.Usage != "" { + glamText(config.Usage) } - return nil } // MITMHostOwnership is the result of mapping intercept hosts to the diff --git a/core/config/model_config_loader_test.go b/core/config/model_config_loader_test.go index 06ab65a20..87807deec 100644 --- a/core/config/model_config_loader_test.go +++ b/core/config/model_config_loader_test.go @@ -1,10 +1,215 @@ package config import ( + "context" + "os" + "path/filepath" + . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/modelartifacts" ) +type preloadArtifactMaterializer struct { + result modelartifacts.Result + err error + seen chan modelartifacts.Spec + release <-chan struct{} +} + +func (f *preloadArtifactMaterializer) Ensure(ctx context.Context, _ string, spec modelartifacts.Spec) (modelartifacts.Result, error) { + if f.seen != nil { + f.seen <- spec + } + if f.release != nil { + select { + case <-ctx.Done(): + return modelartifacts.Result{}, ctx.Err() + case <-f.release: + } + } + return f.result, f.err +} + +var _ = Describe("ModelConfigLoader artifact preload", func() { + It("materializes and persists a source-only artifact binding", func() { + modelsPath := GinkgoT().TempDir() + configPath := filepath.Join(modelsPath, "managed.yaml") + Expect(os.WriteFile(configPath, []byte(` +name: managed +backend: transformers +unknown_extension: keep-me +artifacts: + - name: model + target: model + source: {type: huggingface, repo: owner/repo} +parameters: {model: owner/repo} +`), 0644)).To(Succeed()) + resolved := modelartifacts.Spec{ + Name: "model", Target: "model", + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", Revision: "main"}, + Resolved: &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + }, + } + fake := &preloadArtifactMaterializer{result: modelartifacts.Result{ + Spec: resolved, + RelativePath: ".artifacts/huggingface/0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef/snapshot", + }, seen: make(chan modelartifacts.Spec, 1)} + loader := NewModelConfigLoader(modelsPath, WithArtifactMaterializer(fake)) + Expect(loader.LoadModelConfigsFromPath(modelsPath)).To(Succeed()) + Expect(loader.PreloadWithContext(context.Background(), modelsPath)).To(Succeed()) + + loaded, found := loader.GetModelConfig("managed") + Expect(found).To(BeTrue()) + Expect(loaded.Model).To(Equal("owner/repo")) + Expect(loaded.ModelFileName()).To(Equal(fake.result.RelativePath)) + data, err := os.ReadFile(configPath) + Expect(err).NotTo(HaveOccurred()) + Expect(string(data)).To(ContainSubstring("unknown_extension: keep-me")) + Expect(string(data)).To(ContainSubstring("revision: 0123456789abcdef0123456789abcdef01234567")) + Expect(string(data)).To(ContainSubstring("cache_key: 0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef")) + Expect(string(data)).To(ContainSubstring("model: owner/repo")) + }) + + It("materializes a direct Hugging Face file reference", func() { + modelsPath := GinkgoT().TempDir() + configPath := filepath.Join(modelsPath, "hf-file.yaml") + Expect(os.WriteFile(configPath, []byte(` +name: hf-file +backend: transformers +parameters: + model: https://huggingface.co/nomic-ai/nomic-embed-text-v1.5-GGUF/resolve/main/nomic-embed-text-v1.5.f16.gguf +`), 0644)).To(Succeed()) + resolved := modelartifacts.Spec{ + Name: "model", + Target: "model", + Source: modelartifacts.Source{Type: "huggingface", Repo: "nomic-ai/nomic-embed-text-v1.5-GGUF", AllowPatterns: []string{"nomic-embed-text-v1.5.f16.gguf"}, Revision: "main"}, + Resolved: &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + }, + } + fake := &preloadArtifactMaterializer{result: modelartifacts.Result{ + Spec: resolved, + RelativePath: ".artifacts/huggingface/0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef/snapshot", + }, seen: make(chan modelartifacts.Spec, 1)} + loader := NewModelConfigLoader(modelsPath, WithArtifactMaterializer(fake)) + Expect(loader.LoadModelConfigsFromPath(modelsPath)).To(Succeed()) + Expect(loader.PreloadWithContext(context.Background(), modelsPath)).To(Succeed()) + + loaded, found := loader.GetModelConfig("hf-file") + Expect(found).To(BeTrue()) + Expect(loaded.Model).To(Equal("https://huggingface.co/nomic-ai/nomic-embed-text-v1.5-GGUF/resolve/main/nomic-embed-text-v1.5.f16.gguf")) + Expect(loaded.ModelFileName()).To(Equal(fake.result.RelativePath)) + Expect(fake.seen).To(Receive(SatisfyAll( + WithTransform(func(spec modelartifacts.Spec) string { return spec.Source.Repo }, Equal("nomic-ai/nomic-embed-text-v1.5-GGUF")), + WithTransform(func(spec modelartifacts.Spec) []string { return spec.Source.AllowPatterns }, Equal([]string{"nomic-embed-text-v1.5.f16.gguf"})), + ))) + }) + + It("falls back to the legacy path when inferred materialization fails", func() { + modelsPath := GinkgoT().TempDir() + configPath := filepath.Join(modelsPath, "hf-legacy.yaml") + Expect(os.WriteFile(configPath, []byte(` +name: hf-legacy +backend: transformers +parameters: + model: https://huggingface.co/nomic-ai/nomic-embed-text-v1.5-GGUF/resolve/main/nomic-embed-text-v1.5.f16.gguf +`), 0644)).To(Succeed()) + fake := &preloadArtifactMaterializer{err: context.Canceled, seen: make(chan modelartifacts.Spec, 1)} + loader := NewModelConfigLoader(modelsPath, WithArtifactMaterializer(fake)) + Expect(loader.LoadModelConfigsFromPath(modelsPath)).To(Succeed()) + Expect(loader.PreloadWithContext(context.Background(), modelsPath)).To(Succeed()) + + loaded, found := loader.GetModelConfig("hf-legacy") + Expect(found).To(BeTrue()) + Expect(loaded.Artifacts).To(BeEmpty()) + Expect(loaded.Model).To(Equal("https://huggingface.co/nomic-ai/nomic-embed-text-v1.5-GGUF/resolve/main/nomic-embed-text-v1.5.f16.gguf")) + Expect(fake.seen).To(Receive()) + }) + + It("does not hold the loader lock while materialization blocks", func() { + seen := make(chan modelartifacts.Spec, 1) + release := make(chan struct{}) + fake := &preloadArtifactMaterializer{seen: seen, release: release} + loader := NewModelConfigLoader(GinkgoT().TempDir(), WithArtifactMaterializer(fake)) + loader.Lock() + loader.configs["managed"] = ModelConfig{ + Name: "managed", + Artifacts: []modelartifacts.Spec{{ + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}, + }}, + } + loader.Unlock() + done := make(chan error, 1) + go func() { done <- loader.PreloadWithContext(context.Background(), loader.modelPath) }() + <-seen + lookupDone := make(chan struct{}) + go func() { + _, _ = loader.GetModelConfig("managed") + close(lookupDone) + }() + Eventually(lookupDone).Should(BeClosed()) + close(release) + Expect(<-done).NotTo(HaveOccurred()) + }) + + It("propagates preload cancellation", func() { + seen := make(chan modelartifacts.Spec, 1) + release := make(chan struct{}) + fake := &preloadArtifactMaterializer{seen: seen, release: release} + loader := NewModelConfigLoader(GinkgoT().TempDir(), WithArtifactMaterializer(fake)) + loader.Lock() + loader.configs["managed"] = ModelConfig{ + Name: "managed", + Artifacts: []modelartifacts.Spec{{ + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}, + }}, + } + loader.Unlock() + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- loader.PreloadWithContext(ctx, loader.modelPath) }() + <-seen + cancel() + Expect(<-done).To(MatchError(context.Canceled)) + }) + + It("does not overwrite a config changed during materialization", func() { + seen := make(chan modelartifacts.Spec, 1) + release := make(chan struct{}) + fake := &preloadArtifactMaterializer{ + seen: seen, release: release, + result: modelartifacts.Result{RelativePath: ".artifacts/huggingface/cached/snapshot"}, + } + loader := NewModelConfigLoader(GinkgoT().TempDir(), WithArtifactMaterializer(fake)) + loader.Lock() + loader.configs["managed"] = ModelConfig{ + Name: "managed", Description: "before", + Artifacts: []modelartifacts.Spec{{ + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}, + }}, + } + loader.Unlock() + done := make(chan error, 1) + go func() { done <- loader.PreloadWithContext(context.Background(), loader.modelPath) }() + <-seen + loader.UpdateModelConfig("managed", func(cfg *ModelConfig) { + cfg.Description = "changed concurrently" + }) + close(release) + Expect(<-done).NotTo(HaveOccurred()) + loaded, found := loader.GetModelConfig("managed") + Expect(found).To(BeTrue()) + Expect(loaded.Description).To(Equal("changed concurrently")) + }) +}) + var _ = Describe("ModelConfigLoader.GetModelsConflictingWith", func() { var bcl *ModelConfigLoader diff --git a/core/config/model_config_test.go b/core/config/model_config_test.go index 4d80e8734..c1e64858d 100644 --- a/core/config/model_config_test.go +++ b/core/config/model_config_test.go @@ -4,9 +4,12 @@ import ( "io" "net/http" "os" + "path/filepath" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/modelartifacts" "gopkg.in/yaml.v3" ) @@ -28,6 +31,40 @@ var _ = Describe("Test cases for config related functions", func() { }) }) + It("round-trips and validates a managed model artifact", func() { + raw := []byte(` +name: qwen-asr +backend: qwen-asr +artifacts: + - name: model + target: model + source: + type: huggingface + repo: Qwen/Qwen3-ASR-1.7B +parameters: + model: Qwen/Qwen3-ASR-1.7B +`) + var cfg ModelConfig + Expect(yaml.Unmarshal(raw, &cfg)).To(Succeed()) + Expect(cfg.Artifacts).To(HaveLen(1)) + valid, err := cfg.Validate() + Expect(err).NotTo(HaveOccurred()) + Expect(valid).To(BeTrue()) + }) + + It("derives a managed snapshot filename without replacing the logical model", func() { + const cacheKey = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" + cfg := ModelConfig{ + Artifacts: []modelartifacts.Spec{{ + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}, + Resolved: &modelartifacts.Resolved{CacheKey: cacheKey}, + }}, + } + cfg.Model = "owner/repo" + Expect(cfg.Model).To(Equal("owner/repo")) + Expect(cfg.ModelFileName()).To(Equal(filepath.Join(".artifacts", "huggingface", cacheKey, "snapshot"))) + }) + Context("Test Read configuration functions", func() { It("Test Validate", func() { tmp, err := os.CreateTemp("", "config.yaml") @@ -834,4 +871,17 @@ var _ = Describe("ModelConfig alias", func() { Expect(ok).To(BeFalse()) Expect(err).To(MatchError(ContainSubstring("pure redirect"))) }) + + It("rejects artifacts on alias configurations", func() { + cfg := ModelConfig{ + Name: "alias-name", + Alias: "target-name", + Artifacts: []modelartifacts.Spec{{ + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}, + }}, + } + valid, err := cfg.Validate() + Expect(valid).To(BeFalse()) + Expect(err).To(MatchError(ContainSubstring("alias"))) + }) }) diff --git a/core/gallery/importers/helpers.go b/core/gallery/importers/helpers.go index 03f90d963..5345f3bed 100644 --- a/core/gallery/importers/helpers.go +++ b/core/gallery/importers/helpers.go @@ -4,10 +4,57 @@ import ( "path/filepath" "strings" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/gallery" "github.com/mudler/LocalAI/pkg/downloader" hfapi "github.com/mudler/LocalAI/pkg/huggingface-api" + "github.com/mudler/LocalAI/pkg/modelartifacts" + "gopkg.in/yaml.v3" ) +var managedArtifactBackends = map[string]struct{}{ + "transformers": {}, "huggingface-embeddings": {}, "sentencetransformers": {}, + "transformers-musicgen": {}, "mamba": {}, "diffusers": {}, "qwen-asr": {}, + "fish-speech": {}, "nemo": {}, "voxcpm": {}, "qwen-tts": {}, + "liquid-audio": {}, "vllm": {}, "vllm-omni": {}, "sglang": {}, +} + +// AttachPrimaryArtifact adds the controller-managed source only when the +// importer selected the same repository and a migrated backend. +func AttachPrimaryArtifact(model gallery.ModelConfig, details Details) (gallery.ModelConfig, error) { + if len(model.Files) != 0 || details.HuggingFace == nil || details.HuggingFace.ModelID == "" { + return model, nil + } + var cfg config.ModelConfig + if err := yaml.Unmarshal([]byte(model.ConfigFile), &cfg); err != nil { + return gallery.ModelConfig{}, err + } + if _, supported := managedArtifactBackends[cfg.Backend]; !supported { + return model, nil + } + if len(cfg.Artifacts) != 0 || cfg.Model != details.HuggingFace.ModelID { + return model, nil + } + var document map[string]any + if err := yaml.Unmarshal([]byte(model.ConfigFile), &document); err != nil { + return gallery.ModelConfig{}, err + } + document["artifacts"] = []map[string]any{{ + "name": modelartifacts.TargetModel, + "target": modelartifacts.TargetModel, + "source": map[string]any{ + "type": modelartifacts.SourceTypeHuggingFace, + "repo": details.HuggingFace.ModelID, + }, + }} + encoded, err := yaml.Marshal(document) + if err != nil { + return gallery.ModelConfig{}, err + } + model.ConfigFile = string(encoded) + return model, nil +} + // LocalModelPath normalizes a model URI for backends that treat the model // field as a HuggingFace repo id or local filesystem path (mlx, mlx-vlm, // vllm, transformers, diffusers). A "file://" import URI is reduced to the diff --git a/core/gallery/importers/helpers_test.go b/core/gallery/importers/helpers_test.go index 390f5cb4e..5dafc77f3 100644 --- a/core/gallery/importers/helpers_test.go +++ b/core/gallery/importers/helpers_test.go @@ -4,8 +4,11 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/gallery" "github.com/mudler/LocalAI/core/gallery/importers" hfapi "github.com/mudler/LocalAI/pkg/huggingface-api" + "gopkg.in/yaml.v3" ) var _ = Describe("importer helpers", func() { @@ -104,3 +107,50 @@ var _ = Describe("importer helpers", func() { }) }) }) + +var _ = Describe("managed artifact attachment", func() { + details := importers.Details{ + URI: "https://huggingface.co/owner/repo", + HuggingFace: &hfapi.ModelDetails{ModelID: "owner/repo"}, + } + + It("attaches a source when the imported model is the discovered repository", func() { + input := gallery.ModelConfig{ConfigFile: "backend: transformers\nparameters:\n model: owner/repo\n"} + output, err := importers.AttachPrimaryArtifact(input, details) + Expect(err).NotTo(HaveOccurred()) + var cfg config.ModelConfig + Expect(yaml.Unmarshal([]byte(output.ConfigFile), &cfg)).To(Succeed()) + Expect(cfg.Artifacts).To(HaveLen(1)) + Expect(cfg.Artifacts[0].Source.Repo).To(Equal("owner/repo")) + }) + + It("preserves explicit gallery files", func() { + input := gallery.ModelConfig{ + ConfigFile: "parameters:\n model: owner/repo\n", + Files: []gallery.File{{Filename: "model.gguf", URI: "https://huggingface.co/owner/repo/resolve/main/model.gguf"}}, + } + output, err := importers.AttachPrimaryArtifact(input, details) + Expect(err).NotTo(HaveOccurred()) + Expect(output.ConfigFile).To(Equal(input.ConfigFile)) + }) + + It("does not attach a managed source to an unmigrated backend", func() { + input := gallery.ModelConfig{ConfigFile: "backend: custom-python\nparameters:\n model: owner/repo\n"} + output, err := importers.AttachPrimaryArtifact(input, details) + Expect(err).NotTo(HaveOccurred()) + Expect(output.ConfigFile).To(Equal(input.ConfigFile)) + }) + + DescribeTable("does not guess an unrelated or local source", + func(model string, candidate importers.Details) { + input := gallery.ModelConfig{ConfigFile: "parameters:\n model: " + model + "\n"} + output, err := importers.AttachPrimaryArtifact(input, candidate) + Expect(err).NotTo(HaveOccurred()) + var cfg config.ModelConfig + Expect(yaml.Unmarshal([]byte(output.ConfigFile), &cfg)).To(Succeed()) + Expect(cfg.Artifacts).To(BeEmpty()) + }, + Entry("different repo", "other/repo", details), + Entry("local file", "/models/repo", importers.Details{URI: "file:///models/repo"}), + ) +}) diff --git a/core/gallery/importers/importers.go b/core/gallery/importers/importers.go index 3fc40e02c..8c156144d 100644 --- a/core/gallery/importers/importers.go +++ b/core/gallery/importers/importers.go @@ -318,6 +318,10 @@ func DiscoverModelConfig(uri string, preferences json.RawMessage) (gallery.Model if err != nil { continue } + modelConfig, err = AttachPrimaryArtifact(modelConfig, details) + if err != nil { + return gallery.ModelConfig{}, fmt.Errorf("attach managed artifact: %w", err) + } break } } diff --git a/core/gallery/model_artifacts.go b/core/gallery/model_artifacts.go new file mode 100644 index 000000000..07924d9df --- /dev/null +++ b/core/gallery/model_artifacts.go @@ -0,0 +1,94 @@ +package gallery + +import ( + "context" + "fmt" + "os" + "path/filepath" + + "gopkg.in/yaml.v3" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/modelartifacts" + "github.com/mudler/xlog" +) + +type ArtifactMaterializer interface { + Ensure(context.Context, string, modelartifacts.Spec) (modelartifacts.Result, error) +} + +type installOptions struct { + materializer ArtifactMaterializer +} + +type InstallOption func(*installOptions) + +func WithArtifactMaterializer(materializer ArtifactMaterializer) InstallOption { + return func(options *installOptions) { + if materializer != nil { + options.materializer = materializer + } + } +} + +func applyInstallOptions(options ...InstallOption) installOptions { + result := installOptions{materializer: modelartifacts.NewDefaultManager()} + for _, option := range options { + option(&result) + } + return result +} + +func bindPrimaryArtifact(ctx context.Context, modelsPath string, typed *config.ModelConfig, configMap map[string]any, materializer ArtifactMaterializer, artifactSpec modelartifacts.Spec, inferred bool) (bool, error) { + result, err := materializer.Ensure(ctx, modelsPath, artifactSpec) + if err != nil { + if inferred { + xlog.Warn("falling back to legacy model loading after artifact materialization failed", "model", typed.Name, "error", err) + return false, nil + } + return false, fmt.Errorf("materialize primary model artifact: %w", err) + } + next := []modelartifacts.Spec{result.Spec} + if len(typed.Artifacts) > 1 { + next = append(next, typed.Artifacts[1:]...) + } + typed.Artifacts = next + artifactYAML, err := yaml.Marshal(typed.Artifacts) + if err != nil { + return false, err + } + var artifactValue any + if err := yaml.Unmarshal(artifactYAML, &artifactValue); err != nil { + return false, err + } + configMap["artifacts"] = artifactValue + return true, nil +} + +func writeModelConfigAtomic(fileName string, data []byte) error { + temporary, err := os.CreateTemp(filepath.Dir(fileName), ".model-config-*") + if err != nil { + return err + } + temporaryName := temporary.Name() + defer func() { _ = os.Remove(temporaryName) }() + if err := temporary.Chmod(0600); err != nil { + _ = temporary.Close() + return err + } + if _, err := temporary.Write(data); err != nil { + _ = temporary.Close() + return err + } + if err := temporary.Sync(); err != nil { + _ = temporary.Close() + return err + } + if err := temporary.Close(); err != nil { + return err + } + if err := os.Chmod(temporaryName, 0644); err != nil { + return err + } + return os.Rename(temporaryName, fileName) +} diff --git a/core/gallery/model_artifacts_test.go b/core/gallery/model_artifacts_test.go new file mode 100644 index 000000000..cd9d4cdd9 --- /dev/null +++ b/core/gallery/model_artifacts_test.go @@ -0,0 +1,240 @@ +package gallery_test + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gopkg.in/yaml.v3" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/pkg/downloader" + "github.com/mudler/LocalAI/pkg/modelartifacts" + "github.com/mudler/LocalAI/pkg/system" +) + +type fakeArtifactMaterializer struct { + result modelartifacts.Result + err error + seen []modelartifacts.Spec +} + +func (f *fakeArtifactMaterializer) Ensure(_ context.Context, _ string, spec modelartifacts.Spec) (modelartifacts.Result, error) { + f.seen = append(f.seen, spec) + return f.result, f.err +} + +var _ = Describe("gallery artifact installation", func() { + It("persists resolved state only after materialization succeeds", func() { + modelsPath := GinkgoT().TempDir() + state, err := system.GetSystemState(system.WithModelPath(modelsPath)) + Expect(err).NotTo(HaveOccurred()) + resolved := modelartifacts.Spec{ + Name: "model", Target: "model", + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", Revision: "main"}, + Resolved: &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + }, + } + fake := &fakeArtifactMaterializer{result: modelartifacts.Result{ + Spec: resolved, + RelativePath: ".artifacts/huggingface/0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef/snapshot", + }} + definition := &gallery.ModelConfig{Name: "managed", ConfigFile: ` +backend: transformers +artifacts: + - name: model + target: model + source: + type: huggingface + repo: owner/repo +parameters: + model: owner/repo +unknown_extension: + keep: true +`} + + installed, err := gallery.InstallModel(context.Background(), state, "managed", definition, nil, nil, false, + gallery.WithArtifactMaterializer(fake)) + Expect(err).NotTo(HaveOccurred()) + Expect(fake.seen).To(HaveLen(1)) + Expect(installed.Model).To(Equal("owner/repo")) + Expect(installed.ModelFileName()).To(Equal(fake.result.RelativePath)) + + data, err := os.ReadFile(filepath.Join(modelsPath, "managed.yaml")) + Expect(err).NotTo(HaveOccurred()) + var persisted map[string]any + Expect(yaml.Unmarshal(data, &persisted)).To(Succeed()) + Expect(persisted).To(HaveKey("unknown_extension")) + parameters := persisted["parameters"].(map[string]any) + Expect(parameters["model"]).To(Equal("owner/repo")) + artifacts := persisted["artifacts"].([]any) + resolvedMap := artifacts[0].(map[string]any)["resolved"].(map[string]any) + Expect(resolvedMap["revision"]).To(Equal(resolved.Resolved.Revision)) + }) + + It("does not replace an existing config when materialization fails", func() { + modelsPath := GinkgoT().TempDir() + state, err := system.GetSystemState(system.WithModelPath(modelsPath)) + Expect(err).NotTo(HaveOccurred()) + configPath := filepath.Join(modelsPath, "managed.yaml") + Expect(os.WriteFile(configPath, []byte("name: old\n"), 0644)).To(Succeed()) + fake := &fakeArtifactMaterializer{err: errors.New("acquisition failed")} + definition := &gallery.ModelConfig{Name: "managed", ConfigFile: ` +artifacts: + - source: {type: huggingface, repo: owner/repo} +parameters: {model: owner/repo} +`} + _, err = gallery.InstallModel(context.Background(), state, "managed", definition, nil, nil, false, + gallery.WithArtifactMaterializer(fake)) + Expect(err).To(MatchError(ContainSubstring("acquisition failed"))) + Expect(os.ReadFile(configPath)).To(Equal([]byte("name: old\n"))) + }) + + It("does not replace an existing config when a legacy file download fails", func() { + modelsPath := GinkgoT().TempDir() + state, err := system.GetSystemState(system.WithModelPath(modelsPath)) + Expect(err).NotTo(HaveOccurred()) + configPath := filepath.Join(modelsPath, "legacy.yaml") + Expect(os.WriteFile(configPath, []byte("name: old\n"), 0644)).To(Succeed()) + definition := &gallery.ModelConfig{ + Name: "legacy", + ConfigFile: "parameters: {model: owner/legacy}\n", + Files: []gallery.File{{ + Filename: "missing.bin", + URI: "file:///definitely-not-a-localai-test-file", + }}, + } + + _, err = gallery.InstallModel(context.Background(), state, "legacy", definition, nil, nil, false) + Expect(err).To(HaveOccurred()) + Expect(os.ReadFile(configPath)).To(Equal([]byte("name: old\n"))) + }) + + It("blocks an explicitly unsafe repository before materialization", func() { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + Expect(r.URL.Path).To(Equal("/api/models/owner/repo/scan")) + w.Header().Set("Content-Type", "application/json") + _, err := w.Write([]byte(`{"hasUnsafeFile":true}`)) + Expect(err).NotTo(HaveOccurred()) + })) + DeferCleanup(server.Close) + originalEndpoint := downloader.HF_ENDPOINT + downloader.HF_ENDPOINT = server.URL + DeferCleanup(func() { downloader.HF_ENDPOINT = originalEndpoint }) + + modelsPath := GinkgoT().TempDir() + state, err := system.GetSystemState(system.WithModelPath(modelsPath)) + Expect(err).NotTo(HaveOccurred()) + fake := &fakeArtifactMaterializer{err: errors.New("unsafe install must not materialize")} + definition := &gallery.ModelConfig{Name: "managed", ConfigFile: ` +artifacts: + - source: {type: huggingface, repo: owner/repo} +parameters: {model: owner/repo} +`} + + _, err = gallery.InstallModel(context.Background(), state, "managed", definition, nil, nil, true, + gallery.WithArtifactMaterializer(fake)) + Expect(err).To(MatchError(downloader.ErrUnsafeFilesFound)) + Expect(fake.seen).To(BeEmpty()) + }) + + It("falls back to the legacy path when inferred materialization fails", func() { + modelsPath := GinkgoT().TempDir() + state, err := system.GetSystemState(system.WithModelPath(modelsPath)) + Expect(err).NotTo(HaveOccurred()) + fake := &fakeArtifactMaterializer{err: errors.New("legacy install must not materialize")} + definition := &gallery.ModelConfig{Name: "legacy", ConfigFile: ` +backend: transformers +parameters: + model: owner/legacy +`} + installed, err := gallery.InstallModel( + context.Background(), state, "legacy", definition, nil, nil, false, + gallery.WithArtifactMaterializer(fake), + ) + Expect(err).NotTo(HaveOccurred()) + Expect(fake.seen).To(HaveLen(1)) + Expect(installed.Model).To(Equal("owner/legacy")) + Expect(installed.Artifacts).To(BeEmpty()) + Expect(filepath.Join(modelsPath, "legacy.yaml")).To(BeAnExistingFile()) + }) + + It("passes the controller materializer through named gallery installs", func() { + modelsPath := GinkgoT().TempDir() + state, err := system.GetSystemState(system.WithModelPath(modelsPath)) + Expect(err).NotTo(HaveOccurred()) + definitionPath := filepath.Join(modelsPath, "definition.yaml") + definition, err := yaml.Marshal(gallery.ModelConfig{Name: "managed", ConfigFile: ` +backend: transformers +artifacts: + - source: {type: huggingface, repo: owner/repo} +parameters: {model: owner/repo} +`}) + Expect(err).NotTo(HaveOccurred()) + Expect(os.WriteFile(definitionPath, definition, 0644)).To(Succeed()) + galleryPath := filepath.Join(modelsPath, "index.yaml") + index, err := yaml.Marshal([]gallery.GalleryModel{{Metadata: gallery.Metadata{ + Name: "managed", + URL: "file://" + definitionPath, + }}}) + Expect(err).NotTo(HaveOccurred()) + Expect(os.WriteFile(galleryPath, index, 0644)).To(Succeed()) + galleries := []config.Gallery{{Name: "test", URL: "file://" + galleryPath}} + fake := &fakeArtifactMaterializer{result: modelartifacts.Result{Spec: modelartifacts.Spec{ + Name: "model", Target: "model", + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", Revision: "main"}, + Resolved: &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + }, + }}} + + Expect(gallery.InstallModelFromGallery( + context.Background(), galleries, nil, state, nil, "test@managed", + gallery.GalleryModel{}, nil, false, false, false, + gallery.WithArtifactMaterializer(fake), + )).To(Succeed()) + Expect(fake.seen).To(HaveLen(1)) + }) +}) + +var _ = Describe("artifact cache deletion policy", func() { + It("deletes the installed config without treating its snapshot directory as a regular file", func() { + const cacheKey = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" + modelsPath := GinkgoT().TempDir() + state, err := system.GetSystemState(system.WithModelPath(modelsPath)) + Expect(err).NotTo(HaveOccurred()) + artifactRoot := filepath.Join(modelsPath, ".artifacts", "huggingface", cacheKey) + Expect(os.MkdirAll(filepath.Join(artifactRoot, "snapshot"), 0750)).To(Succeed()) + Expect(os.WriteFile(filepath.Join(artifactRoot, "manifest.json"), []byte("{}"), 0644)).To(Succeed()) + Expect(os.WriteFile(filepath.Join(modelsPath, "managed.yaml"), []byte(` +name: managed +artifacts: + - name: model + target: model + source: {type: huggingface, repo: owner/repo} + resolved: + endpoint: https://huggingface.co + revision: 0123456789abcdef0123456789abcdef01234567 + cache_key: 0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef +parameters: + model: owner/repo +`), 0644)).To(Succeed()) + + Expect(gallery.DeleteModelFromSystem(state, "managed")).To(Succeed()) + Expect(filepath.Join(modelsPath, "managed.yaml")).NotTo(BeAnExistingFile()) + Expect(artifactRoot).To(BeADirectory()) + Expect(filepath.Join(artifactRoot, "snapshot")).To(BeADirectory()) + Expect(filepath.Join(artifactRoot, "manifest.json")).To(BeAnExistingFile()) + }) +}) diff --git a/core/gallery/models.go b/core/gallery/models.go index 82a36d7a5..4d39cc0cc 100644 --- a/core/gallery/models.go +++ b/core/gallery/models.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "net/url" "os" "path/filepath" "slices" @@ -13,6 +14,7 @@ import ( lconfig "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/pkg/downloader" "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/LocalAI/pkg/utils" @@ -77,7 +79,7 @@ func InstallModelFromGallery( modelGalleries, backendGalleries []lconfig.Gallery, systemState *system.SystemState, modelLoader *model.ModelLoader, - name string, req GalleryModel, downloadStatus func(string, string, string, float64), enforceScan, automaticallyInstallBackend, requireBackendIntegrity bool) error { + name string, req GalleryModel, downloadStatus func(string, string, string, float64), enforceScan, automaticallyInstallBackend, requireBackendIntegrity bool, options ...InstallOption) error { applyModel := func(model *GalleryModel) error { name = strings.ReplaceAll(name, string(os.PathSeparator), "__") @@ -129,7 +131,7 @@ func InstallModelFromGallery( } } - installedModel, err := InstallModel(ctx, systemState, installName, &config, model.Overrides, downloadStatus, enforceScan) + installedModel, err := InstallModel(ctx, systemState, installName, &config, model.Overrides, downloadStatus, enforceScan, options...) if err != nil { return err } @@ -158,7 +160,9 @@ func InstallModelFromGallery( return applyModel(model) } -func InstallModel(ctx context.Context, systemState *system.SystemState, nameOverride string, config *ModelConfig, configOverrides map[string]any, downloadStatus func(string, string, string, float64), enforceScan bool) (*lconfig.ModelConfig, error) { +func InstallModel(ctx context.Context, systemState *system.SystemState, nameOverride string, galleryConfig *ModelConfig, configOverrides map[string]any, downloadStatus func(string, string, string, float64), enforceScan bool, options ...InstallOption) (*lconfig.ModelConfig, error) { + installOptions := applyInstallOptions(options...) + config := galleryConfig basePath := systemState.Model.ModelsPath // Create base path if it doesn't exist err := os.MkdirAll(basePath, 0750) @@ -170,59 +174,6 @@ func InstallModel(ctx context.Context, systemState *system.SystemState, nameOver xlog.Debug("Config overrides", "overrides", configOverrides) } - // Download files and verify their SHA - for i, file := range config.Files { - // Check for cancellation before each file - select { - case <-ctx.Done(): - return nil, ctx.Err() - default: - } - - xlog.Debug("Checking file exists and matches SHA", "filename", file.Filename) - - if err := utils.VerifyPath(file.Filename, basePath); err != nil { - return nil, err - } - - // Create file path - filePath := filepath.Join(basePath, file.Filename) - - if enforceScan { - scanResults, err := downloader.HuggingFaceScan(downloader.URI(file.URI)) - if err != nil && errors.Is(err, downloader.ErrUnsafeFilesFound) { - xlog.Error("Contains unsafe file(s)!", "model", config.Name, "clamAV", scanResults.ClamAVInfectedFiles, "pickles", scanResults.DangerousPickles) - return nil, err - } - } - uri := downloader.URI(file.URI) - if err := uri.DownloadFileWithContext(ctx, filePath, file.SHA256, i, len(config.Files), downloadStatus); err != nil { - return nil, err - } - } - - // Write prompt template contents to separate files - for _, template := range config.PromptTemplates { - if err := utils.VerifyPath(template.Name+".tmpl", basePath); err != nil { - return nil, err - } - // Create file path - filePath := filepath.Join(basePath, template.Name+".tmpl") - - // Create parent directory - err := os.MkdirAll(filepath.Dir(filePath), 0750) - if err != nil { - return nil, fmt.Errorf("failed to create parent directory for prompt template %q: %v", template.Name, err) - } - // Create and write file content - err = os.WriteFile(filePath, []byte(template.Content), 0644) - if err != nil { - return nil, fmt.Errorf("failed to write prompt template %q: %v", template.Name, err) - } - - xlog.Debug("Prompt template written", "template", template.Name) - } - name := config.Name if nameOverride != "" { name = nameOverride @@ -233,11 +184,11 @@ func InstallModel(ctx context.Context, systemState *system.SystemState, nameOver } modelConfig := lconfig.ModelConfig{} + configFilePath := filepath.Join(basePath, name+".yaml") + var updatedConfigYAML []byte + writeConfig := len(configOverrides) != 0 || len(config.ConfigFile) != 0 - // write config file - if len(configOverrides) != 0 || len(config.ConfigFile) != 0 { - configFilePath := filepath.Join(basePath, name+".yaml") - + if writeConfig { // Read and update config file as map[string]interface{} configMap := make(map[string]any) err = yaml.Unmarshal([]byte(config.ConfigFile), &configMap) @@ -253,8 +204,8 @@ func InstallModel(ctx context.Context, systemState *system.SystemState, nameOver } } - // Write updated config file - updatedConfigYAML, err := yaml.Marshal(configMap) + // Marshal the merged map for typed validation and default application. + updatedConfigYAML, err = yaml.Marshal(configMap) if err != nil { return nil, fmt.Errorf("failed to marshal updated config YAML: %v", err) } @@ -301,19 +252,110 @@ func InstallModel(ctx context.Context, systemState *system.SystemState, nameOver } } - // Re-marshal from configMap to preserve unknown fields - updatedConfigYAML, err = yaml.Marshal(configMap) - if err != nil { - return nil, fmt.Errorf("failed to marshal config with inference defaults: %v", err) - } - if valid, err := modelConfig.Validate(); !valid { return nil, fmt.Errorf("failed to validate updated config YAML: %v", err) } - err = os.WriteFile(configFilePath, updatedConfigYAML, 0644) + if len(config.Files) == 0 { + artifactSpec, inferred, hasArtifact, err := modelConfig.PrimaryArtifactSpec(basePath) + if err != nil { + return nil, err + } + if hasArtifact && enforceScan { + artifact, err := artifactSpec.Normalize() + if err != nil { + return nil, err + } + probe := downloader.URI(fmt.Sprintf("%s/%s/resolve/%s/.gitattributes", strings.TrimRight(downloader.HF_ENDPOINT, "/"), artifact.Source.Repo, url.PathEscape(artifact.Source.Revision))) + if _, err := downloader.HuggingFaceScan(probe); errors.Is(err, downloader.ErrUnsafeFilesFound) { + return nil, err + } + } + + if hasArtifact { + applied, err := bindPrimaryArtifact(ctx, basePath, &modelConfig, configMap, installOptions.materializer, artifactSpec, inferred) + if err != nil { + return nil, err + } + if applied { + // Re-marshal from configMap after artifact binding to preserve unknown fields. + updatedConfigYAML, err = yaml.Marshal(configMap) + if err != nil { + return nil, fmt.Errorf("failed to marshal config with inference defaults: %v", err) + } + } + } + } + } + + // Download files and verify their SHA + tasks := make([]downloader.FileTask, 0, len(config.Files)) + for i, file := range config.Files { + // Check for cancellation before each file + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + + xlog.Debug("Checking file exists and matches SHA", "filename", file.Filename) + + if err := utils.VerifyPath(file.Filename, basePath); err != nil { + return nil, err + } + + // Create file path + filePath := filepath.Join(basePath, file.Filename) + + if enforceScan { + scanResults, err := downloader.HuggingFaceScan(downloader.URI(file.URI)) + if err != nil && errors.Is(err, downloader.ErrUnsafeFilesFound) { + xlog.Error("Contains unsafe file(s)!", "model", config.Name, "clamAV", scanResults.ClamAVInfectedFiles, "pickles", scanResults.DangerousPickles) + return nil, err + } + } + tasks = append(tasks, downloader.FileTask{ + URI: downloader.URI(file.URI), + Destination: filePath, + SHA256: file.SHA256, + FileIndex: i, + TotalFiles: len(config.Files), + }) + } + if err := downloader.DownloadFilesWithContext(ctx, tasks, downloadStatus); err != nil { + return nil, err + } + + // Write prompt template contents to separate files + for _, template := range config.PromptTemplates { + if err := utils.VerifyPath(template.Name+".tmpl", basePath); err != nil { + return nil, err + } + // Create file path + filePath := filepath.Join(basePath, template.Name+".tmpl") + + // Create parent directory + err := os.MkdirAll(filepath.Dir(filePath), 0750) if err != nil { - return nil, fmt.Errorf("failed to write updated config file: %v", err) + return nil, fmt.Errorf("failed to create parent directory for prompt template %q: %v", template.Name, err) + } + // Create and write file content + err = os.WriteFile(filePath, []byte(template.Content), 0644) + if err != nil { + return nil, fmt.Errorf("failed to write prompt template %q: %v", template.Name, err) + } + + xlog.Debug("Prompt template written", "template", template.Name) + } + + if writeConfig { + if len(modelConfig.Artifacts) > 0 { + modelartifacts.ReportProgress(ctx, modelartifacts.ProgressEvent{ + Phase: modelartifacts.PhasePersisting, Artifact: modelConfig.Artifacts[0].Name, + }) + } + if err := writeModelConfigAtomic(configFilePath, updatedConfigYAML); err != nil { + return nil, fmt.Errorf("failed to atomically write updated config file: %w", err) } xlog.Debug("Written config file", "file", configFilePath) @@ -374,7 +416,7 @@ func listModelFiles(systemState *system.SystemState, name string) ([]string, err if err != nil { return nil, err } - if modelConfig.Model != "" { + if modelConfig.Model != "" && len(modelConfig.Artifacts) == 0 { additionalFiles = append(additionalFiles, modelConfig.ModelFileName()) } diff --git a/core/http/react-ui/e2e/model-artifact-operation.spec.js b/core/http/react-ui/e2e/model-artifact-operation.spec.js new file mode 100644 index 000000000..22cbfec5c --- /dev/null +++ b/core/http/react-ui/e2e/model-artifact-operation.spec.js @@ -0,0 +1,40 @@ +import { test, expect } from './coverage-fixtures.js' + +test('operations bar shows managed model acquisition phase and bytes', async ({ page }) => { + await page.route('**/api/operations', (route) => route.fulfill({ + contentType: 'application/json', + body: JSON.stringify({ + operations: [{ + id: 'qwen-asr', + name: 'qwen-asr', + fullName: 'qwen-asr', + jobID: 'artifact-job-123', + progress: 45, + taskType: 'installation', + isDeletion: false, + isBackend: false, + isQueued: false, + isCancelled: false, + cancellable: true, + phase: 'downloading', + currentBytes: 1073741824, + totalBytes: 4294967296, + }], + }), + })) + let cancelledPath = '' + await page.route('**/api/operations/artifact-job-123/cancel', (route) => { + cancelledPath = new URL(route.request().url()).pathname + return route.fulfill({ contentType: 'application/json', body: '{}' }) + }) + + await page.goto('/app/models') + const operation = page.locator('.operation-item').filter({ hasText: 'qwen-asr' }) + await expect(operation).toContainText('Downloading model files') + await expect(operation).toContainText('1 GB / 4 GB') + await expect(operation.locator('.operation-progress')).toHaveText('45%') + await expect(operation.locator('.operation-bar')).toHaveAttribute('style', /width: 45%/) + + await operation.getByTitle('Cancel').click() + expect(cancelledPath).toBe('/api/operations/artifact-job-123/cancel') +}) diff --git a/core/http/react-ui/src/components/OperationsBar.jsx b/core/http/react-ui/src/components/OperationsBar.jsx index 4d55f55b4..cea947fd8 100644 --- a/core/http/react-ui/src/components/OperationsBar.jsx +++ b/core/http/react-ui/src/components/OperationsBar.jsx @@ -1,5 +1,14 @@ import { useState } from 'react' import { useOperations } from '../hooks/useOperations' +import { formatBytes } from '../utils/format' + +const artifactPhaseLabels = { + resolving: 'Resolving model files', + downloading: 'Downloading model files', + verifying: 'Verifying model files', + committing: 'Finalizing model installation', + persisting: 'Saving model configuration', +} const nodeStatusLabels = { success: 'Done', @@ -26,6 +35,10 @@ export default function OperationsBar() { const nodes = Array.isArray(op.nodes) ? op.nodes : [] const canExpand = nodes.length > 1 const isOpen = !!expanded[key] + const phaseLabel = artifactPhaseLabels[op.phase] + const byteLabel = Number.isFinite(op.currentBytes) && Number.isFinite(op.totalBytes) && op.totalBytes > 0 + ? `${formatBytes(op.currentBytes)} / ${formatBytes(op.totalBytes)}` + : '' return (
@@ -68,7 +81,17 @@ export default function OperationsBar() { Cancelling... )} - {!op.error && op.message && !op.isQueued && !op.isCancelled && ( + {!op.error && phaseLabel && !op.isCancelled && ( + + {phaseLabel} + + )} + {!op.error && byteLabel && !op.isCancelled && ( + + {byteLabel} + + )} + {!op.error && op.message && !phaseLabel && !op.isQueued && !op.isCancelled && ( {op.message} diff --git a/core/http/routes/ui_api.go b/core/http/routes/ui_api.go index 218cf5dfb..e737e445e 100644 --- a/core/http/routes/ui_api.go +++ b/core/http/routes/ui_api.go @@ -173,6 +173,9 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model isCancelled := false isCancellable := false message := "" + phase := "" + currentBytes := int64(0) + totalBytes := int64(0) if status != nil { // Skip successfully completed operations @@ -189,6 +192,9 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model isCancelled = status.Cancelled isCancellable = status.Cancellable message = status.Message + phase = status.Phase + currentBytes = status.CurrentBytes + totalBytes = status.TotalBytes if isDeletion { taskType = "deletion" } @@ -257,6 +263,15 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model if scopedNodeID != "" { opData["nodeID"] = scopedNodeID } + if phase != "" { + opData["phase"] = phase + } + if currentBytes > 0 { + opData["currentBytes"] = currentBytes + } + if totalBytes > 0 { + opData["totalBytes"] = totalBytes + } if status != nil && status.Error != nil { opData["error"] = status.Error.Error() } @@ -909,7 +924,8 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model }) } - _, err = gallery.InstallModel(context.Background(), appConfig.SystemState, model.Name, &config, model.Overrides, nil, false) + _, err = gallery.InstallModel(context.Background(), appConfig.SystemState, model.Name, &config, model.Overrides, nil, false, + gallery.WithArtifactMaterializer(appConfig.ModelArtifactMaterializer)) if err != nil { return c.JSON(http.StatusInternalServerError, map[string]any{ "error": err.Error(), diff --git a/core/http/routes/ui_api_operations_test.go b/core/http/routes/ui_api_operations_test.go index a1a808209..83bd19717 100644 --- a/core/http/routes/ui_api_operations_test.go +++ b/core/http/routes/ui_api_operations_test.go @@ -13,6 +13,7 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/http/routes" "github.com/mudler/LocalAI/core/services/galleryop" + "github.com/mudler/LocalAI/pkg/system" ) // These specs guard the contract between the opcache (which stores @@ -152,4 +153,41 @@ var _ = Describe("/api/operations with node-scoped backend ops", func() { Expect(found).ToNot(HaveKey("nodeID"), "non-node-scoped ops must NOT carry a nodeID field") Expect(found["name"]).To(Equal("llama-cpp")) }) + + It("surfaces managed model artifact phase and byte counters", func() { + state, err := system.GetSystemState(system.WithModelPath(GinkgoT().TempDir())) + Expect(err).NotTo(HaveOccurred()) + appCfg := &config.ApplicationConfig{SystemState: state} + galleryService := galleryop.NewGalleryService(appCfg, nil) + opcache := galleryop.NewOpCache(galleryService) + jobID := "test-op-artifact-progress" + opcache.Set("qwen-asr", jobID) + galleryService.UpdateStatus(jobID, &galleryop.OpStatus{ + Phase: "downloading", CurrentBytes: 1024, TotalBytes: 4096, + Progress: 22.5, Message: "Downloading model file: weights.safetensors", Cancellable: true, + }) + + e := echo.New() + routes.RegisterUIAPIRoutes(e, nil, nil, appCfg, galleryService, opcache, &application.Application{}, noopMw) + req := httptest.NewRequest(http.MethodGet, "/api/operations", nil) + rec := httptest.NewRecorder() + e.ServeHTTP(rec, req) + Expect(rec.Code).To(Equal(http.StatusOK)) + + var envelope struct { + Operations []map[string]any `json:"operations"` + } + Expect(json.Unmarshal(rec.Body.Bytes(), &envelope)).To(Succeed()) + var found map[string]any + for _, op := range envelope.Operations { + if op["jobID"] == jobID { + found = op + break + } + } + Expect(found).ToNot(BeNil()) + Expect(found["phase"]).To(Equal("downloading")) + Expect(found["currentBytes"]).To(Equal(float64(1024))) + Expect(found["totalBytes"]).To(Equal(float64(4096))) + }) }) diff --git a/core/services/distributed/gallery.go b/core/services/distributed/gallery.go index d85fd76a8..de213c4f5 100644 --- a/core/services/distributed/gallery.go +++ b/core/services/distributed/gallery.go @@ -25,6 +25,9 @@ type GalleryOperationRecord struct { OpType string `gorm:"size:32" json:"op_type"` // "model_install", "model_delete", "backend_install" Status string `gorm:"size:32;default:pending" json:"status"` // pending, downloading, processing, completed, failed, cancelled Progress float64 `json:"progress"` // 0.0 to 1.0 + Phase string `gorm:"size:32" json:"phase,omitempty"` + CurrentBytes int64 `json:"current_bytes,omitempty"` + TotalBytes int64 `json:"total_bytes,omitempty"` Message string `gorm:"type:text" json:"message,omitempty"` Error string `gorm:"type:text" json:"error,omitempty"` FileName string `gorm:"size:512" json:"file_name,omitempty"` @@ -84,14 +87,26 @@ func (s *GalleryStore) Create(op *GalleryOperationRecord) error { // op as still cancellable — otherwise the column keeps its Create-time zero // value (false), the UI hides the cancel button, and the orphaned op can only // be dismissed by waiting for the 30-minute stale reaper. -func (s *GalleryStore) UpdateProgress(id string, progress float64, message, downloadedSize string, cancellable bool) error { - return s.db.Model(&GalleryOperationRecord{}).Where("id = ?", id).Updates(map[string]any{ +type OperationProgressDetails struct { + Phase string + CurrentBytes int64 + TotalBytes int64 +} + +func (s *GalleryStore) UpdateProgress(id string, progress float64, message, downloadedSize string, cancellable bool, details ...OperationProgressDetails) error { + updates := map[string]any{ "progress": progress, "message": message, "downloaded_file_size": downloadedSize, "cancellable": cancellable, "updated_at": time.Now(), - }).Error + } + if len(details) > 0 { + updates["phase"] = details[0].Phase + updates["current_bytes"] = details[0].CurrentBytes + updates["total_bytes"] = details[0].TotalBytes + } + return s.db.Model(&GalleryOperationRecord{}).Where("id = ?", id).Updates(updates).Error } // UpdateStatus updates the status of an operation. A terminal status is never diff --git a/core/services/galleryop/artifact_progress.go b/core/services/galleryop/artifact_progress.go new file mode 100644 index 000000000..d4e49be58 --- /dev/null +++ b/core/services/galleryop/artifact_progress.go @@ -0,0 +1,65 @@ +package galleryop + +import ( + "fmt" + "sync" + + "github.com/mudler/LocalAI/pkg/modelartifacts" +) + +type artifactProgressBridge struct { + mu sync.Mutex + last float64 + currentBytes int64 + totalBytes int64 + update func(*OpStatus) +} + +func newArtifactProgressBridge(update func(*OpStatus)) *artifactProgressBridge { + return &artifactProgressBridge{update: update} +} + +func (b *artifactProgressBridge) Sink(event modelartifacts.ProgressEvent) { + b.mu.Lock() + progress := b.last + message := "Preparing model files" + switch event.Phase { + case modelartifacts.PhaseResolving: + progress = max(progress, 0) + message = "Resolving model files" + case modelartifacts.PhaseDownloading: + if event.TotalBytes > 0 { + progress = max(progress, min(90, float64(event.CurrentBytes)*90/float64(event.TotalBytes))) + } + message = fmt.Sprintf("Downloading model file: %s", event.File) + case modelartifacts.PhaseVerifying: + progress = max(progress, 95) + message = "Verifying model files" + case modelartifacts.PhaseCommitting: + progress = max(progress, 99) + message = "Finalizing model installation" + case modelartifacts.PhasePersisting: + progress = max(progress, 99) + message = "Saving model configuration" + } + b.last = progress + b.currentBytes = max(b.currentBytes, event.CurrentBytes) + b.totalBytes = max(b.totalBytes, event.TotalBytes) + status := &OpStatus{ + Phase: string(event.Phase), Message: message, FileName: event.File, + Progress: progress, CurrentBytes: b.currentBytes, TotalBytes: b.totalBytes, + Cancellable: true, + } + update := b.update + b.mu.Unlock() + if update != nil { + update(status) + } +} + +func (b *artifactProgressBridge) ClampLegacy(progress float64) float64 { + b.mu.Lock() + defer b.mu.Unlock() + b.last = max(b.last, progress) + return b.last +} diff --git a/core/services/galleryop/artifact_progress_test.go b/core/services/galleryop/artifact_progress_test.go new file mode 100644 index 000000000..fa97eb3b9 --- /dev/null +++ b/core/services/galleryop/artifact_progress_test.go @@ -0,0 +1,33 @@ +package galleryop + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/modelartifacts" +) + +var _ = Describe("artifact operation progress", func() { + It("maps phases and never moves percentage backward", func() { + var statuses []*OpStatus + bridge := newArtifactProgressBridge(func(status *OpStatus) { + copy := *status + statuses = append(statuses, ©) + }) + bridge.Sink(modelartifacts.ProgressEvent{Phase: modelartifacts.PhaseResolving}) + bridge.Sink(modelartifacts.ProgressEvent{Phase: modelartifacts.PhaseDownloading, File: "model.bin", CurrentBytes: 50, TotalBytes: 100}) + bridge.Sink(modelartifacts.ProgressEvent{Phase: modelartifacts.PhaseDownloading, File: "model.bin", CurrentBytes: 10, TotalBytes: 100}) + bridge.Sink(modelartifacts.ProgressEvent{Phase: modelartifacts.PhaseVerifying, CurrentBytes: 100, TotalBytes: 100}) + bridge.Sink(modelartifacts.ProgressEvent{Phase: modelartifacts.PhaseCommitting, CurrentBytes: 100, TotalBytes: 100}) + bridge.Sink(modelartifacts.ProgressEvent{Phase: modelartifacts.PhasePersisting}) + + Expect(statuses).To(HaveLen(6)) + for index := 1; index < len(statuses); index++ { + Expect(statuses[index].Progress).To(BeNumerically(">=", statuses[index-1].Progress)) + } + Expect(statuses[1].Phase).To(Equal("downloading")) + Expect(statuses[1].CurrentBytes).To(Equal(int64(50))) + Expect(statuses[5].Progress).To(Equal(float64(99))) + Expect(statuses[5].CurrentBytes).To(Equal(int64(100))) + }) +}) diff --git a/core/services/galleryop/cancellable_persist_test.go b/core/services/galleryop/cancellable_persist_test.go index 6879dee92..524e42451 100644 --- a/core/services/galleryop/cancellable_persist_test.go +++ b/core/services/galleryop/cancellable_persist_test.go @@ -38,9 +38,12 @@ var _ = Describe("GalleryService cancellable persistence across restart", func() // Simulate a progress tick: the live path always marks installs // cancellable while they are downloading/processing. svc.UpdateStatus("op-inflight", &galleryop.OpStatus{ - Message: "downloading", - Progress: 25, - Cancellable: true, + Message: "downloading", + Progress: 25, + Phase: "downloading", + CurrentBytes: 123, + TotalBytes: 456, + Cancellable: true, }) // A fresh replica boots and hydrates from the store. @@ -52,5 +55,8 @@ var _ = Describe("GalleryService cancellable persistence across restart", func() Expect(st).ToNot(BeNil(), "the in-flight op must hydrate after a restart") Expect(st.Cancellable).To(BeTrue(), "a still-active install must rehydrate as cancellable so the admin can dismiss it") + Expect(st.Phase).To(Equal("downloading")) + Expect(st.CurrentBytes).To(Equal(int64(123))) + Expect(st.TotalBytes).To(Equal(int64(456))) }) }) diff --git a/core/services/galleryop/distributed_sync_test.go b/core/services/galleryop/distributed_sync_test.go index 71a96c7ae..b1c9a3120 100644 --- a/core/services/galleryop/distributed_sync_test.go +++ b/core/services/galleryop/distributed_sync_test.go @@ -122,6 +122,9 @@ var _ = Describe("OpStatus JSON wire format", func() { original := &galleryop.OpStatus{ Progress: 42.0, Message: "downloading", + Phase: "downloading", + CurrentBytes: 123, + TotalBytes: 456, GalleryElementName: "vllm", Error: errors.New("disk full"), Processed: true, @@ -134,6 +137,9 @@ var _ = Describe("OpStatus JSON wire format", func() { Expect(got.Error).ToNot(BeNil(), "the error must survive the round-trip — peer replicas need to surface the failure") Expect(got.Error.Error()).To(Equal("disk full")) Expect(got.Progress).To(Equal(42.0)) + Expect(got.Phase).To(Equal("downloading")) + Expect(got.CurrentBytes).To(Equal(int64(123))) + Expect(got.TotalBytes).To(Equal(int64(456))) Expect(got.GalleryElementName).To(Equal("vllm")) Expect(got.Processed).To(BeTrue()) }) diff --git a/core/services/galleryop/managers_local.go b/core/services/galleryop/managers_local.go index 338a097b4..7d6b75879 100644 --- a/core/services/galleryop/managers_local.go +++ b/core/services/galleryop/managers_local.go @@ -17,6 +17,7 @@ type LocalModelManager struct { enforcePredownloadScans bool automaticallyInstallBackend bool requireBackendIntegrity bool + artifactMaterializer config.ArtifactMaterializer } // NewLocalModelManager creates a LocalModelManager from the application config. @@ -27,6 +28,7 @@ func NewLocalModelManager(appConfig *config.ApplicationConfig, ml *model.ModelLo enforcePredownloadScans: appConfig.EnforcePredownloadScans, automaticallyInstallBackend: appConfig.AutoloadBackendGalleries, requireBackendIntegrity: appConfig.RequireBackendIntegrity, + artifactMaterializer: appConfig.ModelArtifactMaterializer, } } @@ -48,7 +50,8 @@ func (m *LocalModelManager) InstallModel(ctx context.Context, op *ManagementOp[g switch { case op.GalleryElement != nil: installedModel, err := gallery.InstallModel(ctx, m.systemState, op.GalleryElement.Name, - op.GalleryElement, op.Req.Overrides, progressCb, m.enforcePredownloadScans) + op.GalleryElement, op.Req.Overrides, progressCb, m.enforcePredownloadScans, + gallery.WithArtifactMaterializer(m.artifactMaterializer)) if err != nil { return err } @@ -61,10 +64,12 @@ func (m *LocalModelManager) InstallModel(ctx context.Context, op *ManagementOp[g case op.GalleryElementName != "": return gallery.InstallModelFromGallery(ctx, op.Galleries, op.BackendGalleries, m.systemState, m.modelLoader, op.GalleryElementName, op.Req, progressCb, - m.enforcePredownloadScans, m.automaticallyInstallBackend, m.requireBackendIntegrity) + m.enforcePredownloadScans, m.automaticallyInstallBackend, m.requireBackendIntegrity, + gallery.WithArtifactMaterializer(m.artifactMaterializer)) default: return installModelFromRemoteConfig(ctx, m.systemState, m.modelLoader, op.Req, - progressCb, m.enforcePredownloadScans, m.automaticallyInstallBackend, op.BackendGalleries, m.requireBackendIntegrity) + progressCb, m.enforcePredownloadScans, m.automaticallyInstallBackend, op.BackendGalleries, m.requireBackendIntegrity, + gallery.WithArtifactMaterializer(m.artifactMaterializer)) } } diff --git a/core/services/galleryop/model_artifact_materializer_test.go b/core/services/galleryop/model_artifact_materializer_test.go new file mode 100644 index 000000000..07df9e595 --- /dev/null +++ b/core/services/galleryop/model_artifact_materializer_test.go @@ -0,0 +1,51 @@ +package galleryop_test + +import ( + "context" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/core/services/galleryop" + "github.com/mudler/LocalAI/pkg/modelartifacts" + "github.com/mudler/LocalAI/pkg/system" +) + +type galleryOpArtifactMaterializer struct { + seen []modelartifacts.Spec +} + +func (f *galleryOpArtifactMaterializer) Ensure(_ context.Context, _ string, spec modelartifacts.Spec) (modelartifacts.Result, error) { + f.seen = append(f.seen, spec) + spec.Resolved = &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + } + return modelartifacts.Result{Spec: spec}, nil +} + +var _ = Describe("local model manager artifact materializer", func() { + It("passes the controller materializer to direct gallery installs", func() { + state, err := system.GetSystemState(system.WithModelPath(GinkgoT().TempDir())) + Expect(err).NotTo(HaveOccurred()) + materializer := &galleryOpArtifactMaterializer{} + manager := galleryop.NewLocalModelManager(&config.ApplicationConfig{ + SystemState: state, + ModelArtifactMaterializer: materializer, + }, nil) + definition := &gallery.ModelConfig{Name: "managed", ConfigFile: ` +backend: transformers +artifacts: + - source: {type: huggingface, repo: owner/repo} +parameters: {model: owner/repo} +`} + op := &galleryop.ManagementOp[gallery.GalleryModel, gallery.ModelConfig]{ + GalleryElement: definition, + } + Expect(manager.InstallModel(context.Background(), op, nil)).To(Succeed()) + Expect(materializer.seen).To(HaveLen(1)) + }) +}) diff --git a/core/services/galleryop/models.go b/core/services/galleryop/models.go index 3d50a330e..1c6e57c2a 100644 --- a/core/services/galleryop/models.go +++ b/core/services/galleryop/models.go @@ -11,6 +11,7 @@ import ( "github.com/mudler/LocalAI/core/gallery" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/LocalAI/pkg/utils" "gopkg.in/yaml.v3" @@ -52,6 +53,16 @@ func (g *GalleryService) modelHandler(op *ManagementOp[gallery.GalleryModel, gal g.UpdateStatus(op.ID, &OpStatus{Message: fmt.Sprintf("processing model: %s", op.GalleryElementName), Progress: 0, Cancellable: true}) + bridge := newArtifactProgressBridge(func(status *OpStatus) { + status.GalleryElementName = op.GalleryElementName + g.UpdateStatus(op.ID, status) + }) + operationCtx := op.Context + if operationCtx == nil { + operationCtx = context.Background() + } + operationCtx = modelartifacts.WithProgressSink(operationCtx, bridge.Sink) + // displayDownload displays the download progress progressCallback := func(fileName string, current string, total string, percentage float64) { // Check for cancellation during progress updates @@ -62,6 +73,7 @@ func (g *GalleryService) modelHandler(op *ManagementOp[gallery.GalleryModel, gal default: } } + percentage = bridge.ClampLegacy(percentage) g.UpdateStatus(op.ID, &OpStatus{Message: fmt.Sprintf(processingMessage, fileName, total, current), FileName: fileName, Progress: percentage, TotalFileSize: total, DownloadedFileSize: current, Cancellable: true}) utils.DisplayDownloadFunction(fileName, current, total, percentage) } @@ -70,7 +82,7 @@ func (g *GalleryService) modelHandler(op *ManagementOp[gallery.GalleryModel, gal if op.Delete { err = g.modelManager.DeleteModel(op.GalleryElementName) } else { - err = g.modelManager.InstallModel(op.Context, op, progressCallback) + err = g.modelManager.InstallModel(operationCtx, op, progressCallback) } if err != nil { // Check if error is due to cancellation @@ -107,7 +119,7 @@ func (g *GalleryService) modelHandler(op *ManagementOp[gallery.GalleryModel, gal return err } - err = cl.Preload(systemState.Model.ModelsPath) + err = cl.PreloadWithContext(operationCtx, systemState.Model.ModelsPath) if err != nil { return err } @@ -137,7 +149,7 @@ func (g *GalleryService) modelHandler(op *ManagementOp[gallery.GalleryModel, gal return nil } -func installModelFromRemoteConfig(ctx context.Context, systemState *system.SystemState, modelLoader *model.ModelLoader, req gallery.GalleryModel, downloadStatus func(string, string, string, float64), enforceScan, automaticallyInstallBackend bool, backendGalleries []config.Gallery, requireBackendIntegrity bool) error { +func installModelFromRemoteConfig(ctx context.Context, systemState *system.SystemState, modelLoader *model.ModelLoader, req gallery.GalleryModel, downloadStatus func(string, string, string, float64), enforceScan, automaticallyInstallBackend bool, backendGalleries []config.Gallery, requireBackendIntegrity bool, options ...gallery.InstallOption) error { config, err := gallery.GetGalleryConfigFromURLWithContext[gallery.ModelConfig](ctx, req.URL, systemState.Model.ModelsPath) if err != nil { return err @@ -145,7 +157,7 @@ func installModelFromRemoteConfig(ctx context.Context, systemState *system.Syste config.Files = append(config.Files, req.AdditionalFiles...) - installedModel, err := gallery.InstallModel(ctx, systemState, req.Name, &config, req.Overrides, downloadStatus, enforceScan) + installedModel, err := gallery.InstallModel(ctx, systemState, req.Name, &config, req.Overrides, downloadStatus, enforceScan, options...) if err != nil { return err } @@ -164,23 +176,23 @@ type galleryModel struct { ID string `json:"id"` } -func processRequests(systemState *system.SystemState, modelLoader *model.ModelLoader, enforceScan, automaticallyInstallBackend bool, galleries []config.Gallery, backendGalleries []config.Gallery, requests []galleryModel, requireBackendIntegrity bool) error { +func processRequests(systemState *system.SystemState, modelLoader *model.ModelLoader, enforceScan, automaticallyInstallBackend bool, galleries []config.Gallery, backendGalleries []config.Gallery, requests []galleryModel, requireBackendIntegrity bool, options ...gallery.InstallOption) error { ctx := context.Background() var err error for _, r := range requests { utils.ResetDownloadTimers() if r.ID == "" { - err = installModelFromRemoteConfig(ctx, systemState, modelLoader, r.GalleryModel, utils.DisplayDownloadFunction, enforceScan, automaticallyInstallBackend, backendGalleries, requireBackendIntegrity) + err = installModelFromRemoteConfig(ctx, systemState, modelLoader, r.GalleryModel, utils.DisplayDownloadFunction, enforceScan, automaticallyInstallBackend, backendGalleries, requireBackendIntegrity, options...) } else { err = gallery.InstallModelFromGallery( - ctx, galleries, backendGalleries, systemState, modelLoader, r.ID, r.GalleryModel, utils.DisplayDownloadFunction, enforceScan, automaticallyInstallBackend, requireBackendIntegrity) + ctx, galleries, backendGalleries, systemState, modelLoader, r.ID, r.GalleryModel, utils.DisplayDownloadFunction, enforceScan, automaticallyInstallBackend, requireBackendIntegrity, options...) } } return err } -func ApplyGalleryFromFile(systemState *system.SystemState, modelLoader *model.ModelLoader, enforceScan, automaticallyInstallBackend bool, galleries []config.Gallery, backendGalleries []config.Gallery, s string, requireBackendIntegrity bool) error { +func ApplyGalleryFromFile(systemState *system.SystemState, modelLoader *model.ModelLoader, enforceScan, automaticallyInstallBackend bool, galleries []config.Gallery, backendGalleries []config.Gallery, s string, requireBackendIntegrity bool, options ...gallery.InstallOption) error { dat, err := os.ReadFile(s) if err != nil { return err @@ -191,15 +203,15 @@ func ApplyGalleryFromFile(systemState *system.SystemState, modelLoader *model.Mo return err } - return processRequests(systemState, modelLoader, enforceScan, automaticallyInstallBackend, galleries, backendGalleries, requests, requireBackendIntegrity) + return processRequests(systemState, modelLoader, enforceScan, automaticallyInstallBackend, galleries, backendGalleries, requests, requireBackendIntegrity, options...) } -func ApplyGalleryFromString(systemState *system.SystemState, modelLoader *model.ModelLoader, enforceScan, automaticallyInstallBackend bool, galleries []config.Gallery, backendGalleries []config.Gallery, s string, requireBackendIntegrity bool) error { +func ApplyGalleryFromString(systemState *system.SystemState, modelLoader *model.ModelLoader, enforceScan, automaticallyInstallBackend bool, galleries []config.Gallery, backendGalleries []config.Gallery, s string, requireBackendIntegrity bool, options ...gallery.InstallOption) error { var requests []galleryModel err := json.Unmarshal([]byte(s), &requests) if err != nil { return err } - return processRequests(systemState, modelLoader, enforceScan, automaticallyInstallBackend, galleries, backendGalleries, requests, requireBackendIntegrity) + return processRequests(systemState, modelLoader, enforceScan, automaticallyInstallBackend, galleries, backendGalleries, requests, requireBackendIntegrity, options...) } diff --git a/core/services/galleryop/operation.go b/core/services/galleryop/operation.go index 1ac71dfd9..ed86f75a8 100644 --- a/core/services/galleryop/operation.go +++ b/core/services/galleryop/operation.go @@ -61,6 +61,9 @@ type OpStatus struct { Processed bool `json:"processed"` Message string `json:"message"` Progress float64 `json:"progress"` + Phase string `json:"phase,omitempty"` + CurrentBytes int64 `json:"current_bytes,omitempty"` + TotalBytes int64 `json:"total_bytes,omitempty"` TotalFileSize string `json:"file_size"` DownloadedFileSize string `json:"downloaded_size"` GalleryElementName string `json:"gallery_element_name"` @@ -89,6 +92,9 @@ type opStatusWire struct { Processed bool `json:"processed"` Message string `json:"message"` Progress float64 `json:"progress"` + Phase string `json:"phase,omitempty"` + CurrentBytes int64 `json:"current_bytes,omitempty"` + TotalBytes int64 `json:"total_bytes,omitempty"` TotalFileSize string `json:"file_size"` DownloadedFileSize string `json:"downloaded_size"` GalleryElementName string `json:"gallery_element_name"` @@ -104,6 +110,9 @@ func (o OpStatus) MarshalJSON() ([]byte, error) { Processed: o.Processed, Message: o.Message, Progress: o.Progress, + Phase: o.Phase, + CurrentBytes: o.CurrentBytes, + TotalBytes: o.TotalBytes, TotalFileSize: o.TotalFileSize, DownloadedFileSize: o.DownloadedFileSize, GalleryElementName: o.GalleryElementName, @@ -127,6 +136,9 @@ func (o *OpStatus) UnmarshalJSON(data []byte) error { o.Processed = w.Processed o.Message = w.Message o.Progress = w.Progress + o.Phase = w.Phase + o.CurrentBytes = w.CurrentBytes + o.TotalBytes = w.TotalBytes o.TotalFileSize = w.TotalFileSize o.DownloadedFileSize = w.DownloadedFileSize o.GalleryElementName = w.GalleryElementName diff --git a/core/services/galleryop/service.go b/core/services/galleryop/service.go index abe399088..771e22b9a 100644 --- a/core/services/galleryop/service.go +++ b/core/services/galleryop/service.go @@ -94,6 +94,15 @@ func (g *GalleryService) BackendManager() BackendManager { return g.backendManager } +// ModelArtifactMaterializer returns the controller-only acquisition capability +// used by startup paths that install gallery entries outside the operation loop. +func (g *GalleryService) ModelArtifactMaterializer() config.ArtifactMaterializer { + if g == nil || g.appConfig == nil { + return nil + } + return g.appConfig.ModelArtifactMaterializer +} + // SetNATSClient sets the NATS client for distributed progress publishing. // Accepting the wider MessagingClient (vs. plain Publisher) lets // SubscribeBroadcasts wire the wildcard subscriptions that keep peer @@ -167,7 +176,10 @@ func (g *GalleryService) UpdateStatus(s string, op *OpStatus) { xlog.Warn("Failed to persist gallery operation status", "op_id", s, "error", err) } } else { - if err := store.UpdateProgress(s, op.Progress, op.Message, op.DownloadedFileSize, op.Cancellable); err != nil { + if err := store.UpdateProgress(s, op.Progress, op.Message, op.DownloadedFileSize, op.Cancellable, + distributed.OperationProgressDetails{ + Phase: op.Phase, CurrentBytes: op.CurrentBytes, TotalBytes: op.TotalBytes, + }); err != nil { xlog.Warn("Failed to persist gallery operation progress", "op_id", s, "error", err) } } @@ -662,6 +674,9 @@ func (g *GalleryService) Hydrate() error { st := &OpStatus{ Message: op.Message, Progress: op.Progress, + Phase: op.Phase, + CurrentBytes: op.CurrentBytes, + TotalBytes: op.TotalBytes, FileName: op.FileName, TotalFileSize: op.TotalFileSize, DownloadedFileSize: op.DownloadedFileSize, diff --git a/core/services/modeladmin/remote_sync_test.go b/core/services/modeladmin/remote_sync_test.go index df4907a02..a01cb7b56 100644 --- a/core/services/modeladmin/remote_sync_test.go +++ b/core/services/modeladmin/remote_sync_test.go @@ -77,4 +77,30 @@ var _ = Describe("ApplyRemoteChange", func() { Expect(ok1).To(BeTrue()) Expect(ok2).To(BeTrue()) }) + + It("loads a peer-persisted artifact binding without materializing", func() { + const relative = ".artifacts/huggingface/0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef/snapshot" + writeYAML("peer-managed", map[string]any{ + "backend": "transformers", + "artifacts": []map[string]any{{ + "name": "model", "target": "model", + "source": map[string]any{"type": "huggingface", "repo": "owner/repo", "revision": "main"}, + "resolved": map[string]any{ + "endpoint": "https://huggingface.co", + "revision": "0123456789abcdef0123456789abcdef01234567", + "cache_key": "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + }, + }}, + "parameters": map[string]any{"model": "owner/repo"}, + }) + Expect(ApplyRemoteChange(loader, nil, dir, messaging.CacheInvalidateEvent{ + Element: "peer-managed", Op: "install", + })).To(Succeed()) + loaded, found := loader.GetModelConfig("peer-managed") + Expect(found).To(BeTrue()) + Expect(loaded.Model).To(Equal("owner/repo")) + Expect(loaded.ModelFileName()).To(Equal(relative)) + Expect(loaded.Artifacts).To(HaveLen(1)) + Expect(loaded.Artifacts[0].Resolved.CacheKey).To(HaveLen(64)) + }) }) diff --git a/core/services/nodes/router_dirstage_test.go b/core/services/nodes/router_dirstage_test.go index 19c36759f..fef138f48 100644 --- a/core/services/nodes/router_dirstage_test.go +++ b/core/services/nodes/router_dirstage_test.go @@ -4,10 +4,12 @@ import ( "context" "os" "path/filepath" + "strings" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/mudler/LocalAI/core/services/storage" pb "github.com/mudler/LocalAI/pkg/grpc/proto" ) @@ -61,4 +63,48 @@ var _ = Describe("stageModelFiles directory models", func() { Expect(staged).To(ConsistOf(weights, tokenizer)) Expect(staged).ToNot(ContainElement(modelDir)) }) + + It("stages a content-addressed Hugging Face snapshot with its relative tree intact", func() { + cacheKey := strings.Repeat("a", 64) + logicalModel := "owner/repo" + relativeSnapshot := filepath.Join(".artifacts", "huggingface", cacheKey, "snapshot") + snapshot := filepath.Join(tmp, "models", relativeSnapshot) + files := map[string]string{ + "config.json": "{}", + "model.safetensors.index.json": "{}", + "model-00001-of-00002.safetensors": "part-1", + "model-00002-of-00002.safetensors": "part-2", + filepath.Join("tokenizer", "vocab.json"): "{}", + } + for relative, contents := range files { + path := filepath.Join(snapshot, relative) + Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed()) + Expect(os.WriteFile(path, []byte(contents), 0o644)).To(Succeed()) + } + + opts := &pb.ModelOptions{Model: logicalModel, ModelFile: snapshot} + stagedOpts, err := router.stageModelFiles(context.Background(), node, opts, "managed-model") + Expect(err).NotTo(HaveOccurred()) + + stagedPaths := make([]string, 0, len(stager.ensureCalls)) + stagedKeys := make([]string, 0, len(stager.ensureCalls)) + for _, call := range stager.ensureCalls { + stagedPaths = append(stagedPaths, call.localPath) + stagedKeys = append(stagedKeys, call.key) + } + expectedPaths := make([]string, 0, len(files)) + expectedKeys := make([]string, 0, len(files)) + for relative := range files { + expectedPaths = append(expectedPaths, filepath.Join(snapshot, relative)) + expectedKeys = append(expectedKeys, storage.ModelKey(filepath.Join("managed-model", relative))) + } + Expect(stagedPaths).To(ConsistOf(expectedPaths)) + Expect(stagedKeys).To(ConsistOf(expectedKeys)) + remoteSnapshot := filepath.Join("/remote", storage.ModelKey("managed-model")) + Expect(stagedOpts.Model).To(Equal(logicalModel)) + Expect(stagedOpts.ModelFile).To(Equal(remoteSnapshot)) + Expect(stagedOpts.ModelPath).To(Equal(remoteSnapshot)) + Expect(opts.Model).To(Equal(logicalModel)) + Expect(opts.ModelFile).To(Equal(snapshot)) + }) }) diff --git a/core/startup/model_preload.go b/core/startup/model_preload.go index 0d74c4272..3ca375b04 100644 --- a/core/startup/model_preload.go +++ b/core/startup/model_preload.go @@ -24,9 +24,13 @@ import ( func InstallModels(ctx context.Context, galleryService *galleryop.GalleryService, galleries, backendGalleries []config.Gallery, systemState *system.SystemState, modelLoader *model.ModelLoader, enforceScan, autoloadBackendGalleries, requireBackendIntegrity bool, downloadStatus func(string, string, string, float64), models ...string) error { // create an error that groups all errors var err error + var installOptions []gallery.InstallOption + if galleryService != nil { + installOptions = append(installOptions, gallery.WithArtifactMaterializer(galleryService.ModelArtifactMaterializer())) + } for _, url := range models { // Check if it's a model gallery, or print a warning - e, found := installModel(ctx, galleries, backendGalleries, url, systemState, modelLoader, downloadStatus, enforceScan, autoloadBackendGalleries, requireBackendIntegrity) + e, found := installModel(ctx, galleries, backendGalleries, url, systemState, modelLoader, downloadStatus, enforceScan, autoloadBackendGalleries, requireBackendIntegrity, installOptions...) if e != nil && found { xlog.Error("[startup] failed installing model", "error", err, "model", url) err = errors.Join(err, e) @@ -82,7 +86,7 @@ func InstallModels(ctx context.Context, galleryService *galleryop.GalleryService return err } -func installModel(ctx context.Context, galleries, backendGalleries []config.Gallery, modelName string, systemState *system.SystemState, modelLoader *model.ModelLoader, downloadStatus func(string, string, string, float64), enforceScan, autoloadBackendGalleries, requireBackendIntegrity bool) (error, bool) { +func installModel(ctx context.Context, galleries, backendGalleries []config.Gallery, modelName string, systemState *system.SystemState, modelLoader *model.ModelLoader, downloadStatus func(string, string, string, float64), enforceScan, autoloadBackendGalleries, requireBackendIntegrity bool, options ...gallery.InstallOption) (error, bool) { models, err := gallery.AvailableGalleryModels(galleries, systemState) if err != nil { return err, false @@ -98,7 +102,7 @@ func installModel(ctx context.Context, galleries, backendGalleries []config.Gall } xlog.Info("installing model", "model", modelName, "license", model.License) - err = gallery.InstallModelFromGallery(ctx, galleries, backendGalleries, systemState, modelLoader, modelName, gallery.GalleryModel{}, downloadStatus, enforceScan, autoloadBackendGalleries, requireBackendIntegrity) + err = gallery.InstallModelFromGallery(ctx, galleries, backendGalleries, systemState, modelLoader, modelName, gallery.GalleryModel{}, downloadStatus, enforceScan, autoloadBackendGalleries, requireBackendIntegrity, options...) if err != nil { return err, true } diff --git a/docs/content/advanced/model-configuration.md b/docs/content/advanced/model-configuration.md index 1350e80e1..0662d7287 100644 --- a/docs/content/advanced/model-configuration.md +++ b/docs/content/advanced/model-configuration.md @@ -89,6 +89,62 @@ download_files: sha256: abc123... ``` +## Model artifacts + +The `artifacts` section makes installation of a Hugging Face model eager and +repeatable. LocalAI resolves the requested revision to an immutable commit, +downloads the selected repository files, and commits the complete snapshot +before the model installation succeeds. + +```yaml +artifacts: + - name: model + target: model + source: + type: huggingface + repo: Qwen/Qwen3-ASR-1.7B + revision: main + token_env: HF_TOKEN + resolved: + endpoint: https://huggingface.co + revision: 0123456789abcdef0123456789abcdef01234567 + cache_key: 0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef + +parameters: + model: Qwen/Qwen3-ASR-1.7B +``` + +Declare `source` when authoring a configuration. LocalAI owns the `resolved` +block and writes it after installation; do not choose its values manually. +For a public repository, omit `token_env`. For a private or gated repository, +set it to `HF_TOKEN` and provide that environment variable to the LocalAI +controller. + +| Field | Meaning | +|-------|---------| +| `name` | Logical artifact name; `model` for the initial primary artifact | +| `target` | Binding target; only `model` is supported initially | +| `source.type` | `huggingface` | +| `source.repo` | `owner/repository` or `hf://owner/repository` | +| `source.revision` | Branch, tag, or commit; defaults to `main` and resolves to a commit | +| `source.token_env` | Empty or `HF_TOKEN`; the secret value is never persisted | +| `source.allow_patterns` | Optional slash-separated glob allow-list | +| `source.ignore_patterns` | Optional slash-separated glob deny-list | +| `resolved` | Installer-owned immutable endpoint, revision, and cache key | + +Managed installation finishes only after every selected file is committed +locally. `parameters.model` remains the logical repository ID. Once +`resolved.cache_key` is present, LocalAI derives +`.artifacts/huggingface//snapshot` as the runtime `ModelFile`. +Configurations without `artifacts` keep the existing lazy repository-ID +behavior. + +The initially migrated backend families are `transformers` and its aliases, +`diffusers`, `qwen-asr`, `fish-speech`, `nemo`, `voxcpm`, `qwen-tts`, +`liquid-audio`, `vllm`, `vllm-omni`, and `sglang`. Automatic imports add +artifact declarations only for this set. Compatible external backends may opt +in by declaring the artifact explicitly. + ## Parameters Section The `parameters` section contains all OpenAI-compatible request parameters and model-specific options. diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index a64f6636f..fc3462412 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -142,6 +142,23 @@ Set `LOCALAI_DISTRIBUTED_SHARED_MODELS=true` (or `--distributed-shared-models`) This flag is a contract you assert: all nodes must mount identical paths. Leave it off (the default) when workers have independent models directories - the frontend stages files to them over HTTP (or S3) as described above. +### Model artifact staging + +For managed Hugging Face artifacts, the controller resolves the repository and +downloads every selected file. Workers receive the committed snapshot through +the existing directory stager. They never receive `HF_TOKEN` and do not contact +Hugging Face for managed artifacts. + +With `LOCALAI_DISTRIBUTED_SHARED_MODELS` enabled, workers use the shared +absolute snapshot path and skip transfer. Otherwise, the controller stages the +complete snapshot tree to each worker before loading the backend. + +{{% notice warning %}} +Every controller and worker must have enough disk space for its own snapshot +copy unless shared-models mode is enabled. Account for temporary partial files +during installation as well as the committed snapshot. +{{% /notice %}} + {{% notice warning %}} The worker HTTP file transfer server is authenticated by `LOCALAI_REGISTRATION_TOKEN`. If the token is **empty**, the server **fails open** — anyone who can reach the port gets read/write access to the worker's models/staging/data directories (a remote model-poisoning / exfiltration vector). The worker logs a loud warning at startup in this case. Always set `LOCALAI_REGISTRATION_TOKEN` in distributed mode, and set `LOCALAI_DISTRIBUTED_REQUIRE_AUTH=true` (frontend **and** workers) to make a missing token *or* missing NATS credentials a hard startup error rather than a silent fail-open. Firewall the file-transfer port (gRPC base − 1) so only the frontend can reach it. {{% /notice %}} diff --git a/docs/content/features/model-gallery.md b/docs/content/features/model-gallery.md index dc7d09493..fe86eca9d 100644 --- a/docs/content/features/model-gallery.md +++ b/docs/content/features/model-gallery.md @@ -151,6 +151,25 @@ where: - `bert-embeddings` is the model name in the gallery (read its [config here](https://github.com/mudler/LocalAI/tree/master/gallery/blob/main/bert-embeddings.yaml)). +### Artifact-backed models + +Gallery models with an `artifacts` declaration are fully materialized during +installation. Their operation progresses through these phases: + +```text +resolving -> downloading -> verifying -> committing -> persisting +``` + +The admin Operations Bar and `GET /api/operations` expose `currentBytes` and +`totalBytes` as raw transport bytes. Cancelling an active download leaves its +partial files in place so a retry can resume. A verification failure never +exposes a completed snapshot, while a retry or another installation reuses an +already verified content-addressed snapshot. + +Deleting a model configuration does not delete its content-addressed snapshot +bytes. This allows another configuration or a later reinstall to reuse the +cache; safe cache garbage collection is deferred. + ### How to install a model not part of a gallery If you don't want to set any gallery repository, you can still install models by loading a model configuration file. diff --git a/gallery/index.yaml b/gallery/index.yaml index ff7fb6509..b7018ddba 100644 --- a/gallery/index.yaml +++ b/gallery/index.yaml @@ -2424,6 +2424,12 @@ - torch_dtype:bf16 parameters: model: lodestones/Chroma1-HD + artifacts: + - name: model + target: model + source: + type: huggingface + repo: lodestones/Chroma1-HD step: 40 - name: nemotron-3-nano-omni-30b-a3b-reasoning-apex url: github:mudler/LocalAI/gallery/virtual.yaml@master @@ -2473,6 +2479,12 @@ temperature: 1 top_k: -1 top_p: 1 + artifacts: + - name: model + target: model + source: + type: huggingface + repo: mudler/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-APEX-GGUF template: use_tokenizer_template: true - name: carnice-v2-27b @@ -3890,6 +3902,14 @@ - transcript parameters: model: nvidia/parakeet-tdt-0.6b-v3 + artifacts: + - name: model + target: model + source: + type: huggingface + repo: nvidia/parakeet-tdt-0.6b-v3 + allow_patterns: + - "*.nemo" - name: voxtral-mini-4b-realtime url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -5080,6 +5100,12 @@ voice_cloning: true parameters: model: openbmb/VoxCPM1.5 + artifacts: + - name: model + target: model + source: + type: huggingface + repo: openbmb/VoxCPM1.5 - name: neutts-air url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -5124,6 +5150,12 @@ - image parameters: model: Tongyi-MAI/Z-Image-Turbo + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Tongyi-MAI/Z-Image-Turbo - name: vllm-omni-wan2.2-t2v url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -5146,6 +5178,12 @@ - video parameters: model: Wan-AI/Wan2.2-T2V-A14B-Diffusers + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Wan-AI/Wan2.2-T2V-A14B-Diffusers - name: vllm-omni-wan2.2-i2v url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -5168,6 +5206,12 @@ - video parameters: model: Wan-AI/Wan2.2-I2V-A14B-Diffusers + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Wan-AI/Wan2.2-I2V-A14B-Diffusers - name: longcat-video url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -5273,6 +5317,12 @@ - tts parameters: model: Qwen/Qwen3-Omni-30B-A3B-Instruct + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Qwen/Qwen3-Omni-30B-A3B-Instruct - name: vllm-omni-qwen3-tts-custom-voice url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -5296,6 +5346,12 @@ - tts parameters: model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice - name: ace-step-turbo url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -6027,6 +6083,12 @@ - torch_dtype:bf16 parameters: model: Tongyi-MAI/Z-Image + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Tongyi-MAI/Z-Image step: 35 - name: z-image-turbo-diffusers url: github:mudler/LocalAI/gallery/virtual.yaml@master @@ -6052,6 +6114,12 @@ - torch_dtype:bf16 parameters: model: Tongyi-MAI/Z-Image-Turbo + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Tongyi-MAI/Z-Image-Turbo step: 9 - name: glm-4.7-flash-derestricted url: github:mudler/LocalAI/gallery/virtual.yaml@master @@ -6110,6 +6178,12 @@ - tts parameters: model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice tts: voice: Aiden - name: qwen3-tts-0.6b-custom-voice @@ -6130,6 +6204,12 @@ - tts parameters: model: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice tts: voice: Aiden - name: fish-speech-s2-pro @@ -6153,6 +6233,12 @@ voice_cloning: true parameters: model: fishaudio/s2-pro + artifacts: + - name: model + target: model + source: + type: huggingface + repo: fishaudio/s2-pro - name: qwen3-asr-1.7b url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -6171,6 +6257,12 @@ - transcript parameters: model: Qwen/Qwen3-ASR-1.7B + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Qwen/Qwen3-ASR-1.7B - name: qwen3-asr-0.6b url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -6189,6 +6281,12 @@ - transcript parameters: model: Qwen/Qwen3-ASR-0.6B + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Qwen/Qwen3-ASR-0.6B - name: huihui-glm-4.7-flash-abliterated-i1 url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -7635,6 +7733,12 @@ - transcript parameters: model: Qwen/Qwen3-ASR-0.6B + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Qwen/Qwen3-ASR-0.6B - name: qwen3-asr-1.7b url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -7653,6 +7757,12 @@ - transcript parameters: model: Qwen/Qwen3-ASR-1.7B + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Qwen/Qwen3-ASR-1.7B - name: glm-ocr url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -8830,6 +8940,12 @@ - torch_dtype:bf16 parameters: model: Lightricks/LTX-2 + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Lightricks/LTX-2 - name: gpt-oss-20b url: github:mudler/LocalAI/gallery/harmony.yaml@master urls: @@ -9078,6 +9194,12 @@ name: dia parameters: model: nari-labs/Dia-1.6B-0626 + artifacts: + - name: model + target: model + source: + type: huggingface + repo: nari-labs/Dia-1.6B-0626 type: DiaForConditionalGeneration - name: outetts url: github:mudler/LocalAI/gallery/virtual.yaml@master @@ -35853,6 +35975,12 @@ - torch_dtype:bf16 parameters: model: Lightricks/LTX-2.3 + artifacts: + - name: model + target: model + source: + type: huggingface + repo: Lightricks/LTX-2.3 - <x-2-3-dev-ggml name: ltx-2.3-22b-dev-ggml url: github:mudler/LocalAI/gallery/ltx-ggml.yaml@master diff --git a/gallery/liquid-audio.yaml b/gallery/liquid-audio.yaml index cd9ead91b..c2e41b644 100644 --- a/gallery/liquid-audio.yaml +++ b/gallery/liquid-audio.yaml @@ -31,6 +31,12 @@ config_file: | - vad parameters: model: LiquidAI/LFM2.5-Audio-1.5B + artifacts: + - name: model + target: model + source: + type: huggingface + repo: LiquidAI/LFM2.5-Audio-1.5B # Special tokens emitted in the text track during interleaved generation. # Included so a future client-side parser can spot them; the LFM2 tool-call # format itself is auto-detected by the upstream llama.cpp parser when the diff --git a/pkg/downloader/auth_progress_test.go b/pkg/downloader/auth_progress_test.go new file mode 100644 index 000000000..debf45f9e --- /dev/null +++ b/pkg/downloader/auth_progress_test.go @@ -0,0 +1,173 @@ +package downloader_test + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/downloader" +) + +var _ = Describe("authenticated HTTP downloads", func() { + It("sends bearer auth to the origin and strips it at a cross-host redirect", func() { + var originAuth, cdnAuth string + cdn := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + cdnAuth = r.Header.Get("Authorization") + w.Header().Set("Content-Length", "5") + _, _ = w.Write([]byte("model")) + })) + DeferCleanup(cdn.Close) + + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + originAuth = r.Header.Get("Authorization") + http.Redirect(w, r, cdn.URL+"/weights", http.StatusTemporaryRedirect) + })) + DeferCleanup(origin.Close) + + target := filepath.Join(GinkgoT().TempDir(), "weights.bin") + err := downloader.URI(origin.URL+"/resolve").DownloadFileWithContext( + context.Background(), + target, + "", + 0, + 1, + nil, + downloader.WithBearerToken("secret-token"), + ) + Expect(err).NotTo(HaveOccurred()) + Expect(originAuth).To(Equal("Bearer secret-token")) + Expect(cdnAuth).To(BeEmpty()) + }) + + It("authenticates range probes and reports resumed bytes", func() { + var mu sync.Mutex + seen := make(map[string]string) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + seen[r.Method] = r.Header.Get("Authorization") + mu.Unlock() + if r.Method == http.MethodHead { + w.Header().Set("Accept-Ranges", "bytes") + w.Header().Set("Content-Length", "6") + return + } + Expect(r.Header.Get("Range")).To(Equal("bytes=3-")) + w.Header().Set("Content-Length", "3") + w.Header().Set("Content-Range", "bytes 3-5/6") + w.WriteHeader(http.StatusPartialContent) + _, _ = w.Write([]byte("def")) + })) + DeferCleanup(server.Close) + + target := filepath.Join(GinkgoT().TempDir(), "resume.bin") + Expect(os.WriteFile(target+".partial", []byte("abc"), 0o600)).To(Succeed()) + events := []downloader.TransferProgress{} + err := downloader.URI(server.URL+"/weights").DownloadFileWithContext( + context.Background(), + target, + "", + 0, + 1, + nil, + downloader.WithBearerToken("resume-token"), + downloader.WithTransferProgress(func(event downloader.TransferProgress) { + events = append(events, event) + }), + ) + Expect(err).NotTo(HaveOccurred()) + Expect(seen).To(HaveKeyWithValue(http.MethodHead, "Bearer resume-token")) + Expect(seen).To(HaveKeyWithValue(http.MethodGet, "Bearer resume-token")) + Expect(events).NotTo(BeEmpty()) + Expect(events[len(events)-1].Written).To(Equal(int64(6))) + Expect(events[len(events)-1].Total).To(Equal(int64(6))) + Expect(os.ReadFile(target)).To(Equal([]byte("abcdef"))) + }) + + It("rejects an origin that ignores a resume range", func() { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + w.Header().Set("Accept-Ranges", "bytes") + w.Header().Set("Content-Length", "6") + return + } + Expect(r.Header.Get("Range")).To(Equal("bytes=3-")) + w.Header().Set("Content-Length", "6") + _, _ = w.Write([]byte("abcdef")) + })) + DeferCleanup(server.Close) + + target := filepath.Join(GinkgoT().TempDir(), "ignored-range.bin") + Expect(os.WriteFile(target+".partial", []byte("abc"), 0o600)).To(Succeed()) + err := downloader.URI(server.URL).DownloadFileWithContext( + context.Background(), + target, + "", + 0, + 1, + nil, + ) + Expect(err).To(MatchError(ContainSubstring("status 200 instead of 206"))) + Expect(target).NotTo(BeAnExistingFile()) + Expect(target + ".partial").NotTo(BeAnExistingFile()) + }) + + It("creates resumable partial files with owner-only permissions", func() { + ctx, cancel := context.WithCancel(context.Background()) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Length", "5") + _, _ = w.Write([]byte("model")) + })) + DeferCleanup(server.Close) + + target := filepath.Join(GinkgoT().TempDir(), "private.bin") + err := downloader.URI(server.URL).DownloadFileWithContext( + ctx, + target, + "", + 0, + 1, + nil, + downloader.WithTransferProgress(func(downloader.TransferProgress) { + cancel() + }), + ) + Expect(errors.Is(err, context.Canceled)).To(BeTrue()) + info, statErr := os.Stat(target + ".partial") + Expect(statErr).NotTo(HaveOccurred()) + Expect(info.Mode().Perm()).To(Equal(os.FileMode(0o600))) + }) + + It("keeps the legacy total empty when the response length is unknown", func() { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + if flusher, ok := w.(http.Flusher); ok { + flusher.Flush() + } + _, _ = w.Write([]byte("model")) + })) + DeferCleanup(server.Close) + + target := filepath.Join(GinkgoT().TempDir(), "unknown-size.bin") + totals := []string{} + err := downloader.URI(server.URL).DownloadFileWithContext( + context.Background(), + target, + "", + 0, + 1, + func(_, _, total string, _ float64) { + totals = append(totals, total) + }, + ) + Expect(err).NotTo(HaveOccurred()) + Expect(totals).NotTo(BeEmpty()) + Expect(totals[len(totals)-1]).To(BeEmpty()) + }) +}) diff --git a/pkg/downloader/download_plan.go b/pkg/downloader/download_plan.go new file mode 100644 index 000000000..84f0a21a2 --- /dev/null +++ b/pkg/downloader/download_plan.go @@ -0,0 +1,39 @@ +package downloader + +import "context" + +// FileTask describes one download operation and an optional post-download +// hook that runs after the bytes are present on disk. Callers keep any +// higher-level commit logic outside this helper. +type FileTask struct { + URI URI + Destination string + SHA256 string + FileIndex int + TotalFiles int + AfterDownload func(string) error + Options []DownloadOption +} + +// DownloadFilesWithContext executes a set of file downloads sequentially. +// The helper centralizes the shared download path so callers only provide +// source/destination metadata and any post-download hook they need. +func DownloadFilesWithContext(ctx context.Context, tasks []FileTask, status func(string, string, string, float64), opts ...DownloadOption) error { + for i := range tasks { + task := tasks[i] + if err := ctx.Err(); err != nil { + return err + } + taskOpts := append([]DownloadOption{}, opts...) + taskOpts = append(taskOpts, task.Options...) + if err := task.URI.DownloadFileWithContext(ctx, task.Destination, task.SHA256, task.FileIndex, task.TotalFiles, status, taskOpts...); err != nil { + return err + } + if task.AfterDownload != nil { + if err := task.AfterDownload(task.Destination); err != nil { + return err + } + } + } + return nil +} diff --git a/pkg/downloader/plan_test.go b/pkg/downloader/plan_test.go new file mode 100644 index 000000000..9483664d2 --- /dev/null +++ b/pkg/downloader/plan_test.go @@ -0,0 +1,44 @@ +package downloader_test + +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/downloader" +) + +var _ = Describe("DownloadFilesWithContext", func() { + It("runs the post-download hook after fetching each file", func() { + payload := []byte("downloaded-bytes") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write(payload) + })) + DeferCleanup(server.Close) + + dest := filepath.Join(GinkgoT().TempDir(), "model.bin") + hookCalled := false + + err := downloader.DownloadFilesWithContext(context.Background(), []downloader.FileTask{{ + URI: downloader.URI(server.URL), + Destination: dest, + FileIndex: 0, + TotalFiles: 1, + AfterDownload: func(path string) error { + hookCalled = true + got, err := os.ReadFile(path) + Expect(err).NotTo(HaveOccurred()) + Expect(string(got)).To(Equal(string(payload))) + return nil + }, + }}, nil) + + Expect(err).NotTo(HaveOccurred()) + Expect(hookCalled).To(BeTrue()) + }) +}) diff --git a/pkg/downloader/progress.go b/pkg/downloader/progress.go index 05cd07a98..b3c3d8203 100644 --- a/pkg/downloader/progress.go +++ b/pkg/downloader/progress.go @@ -12,10 +12,34 @@ type progressWriter struct { totalFiles int written int64 downloadStatus func(string, string, string, float64) + transferSink TransferProgressSink hash hash.Hash ctx context.Context } +func (pw *progressWriter) report() { + if pw.transferSink != nil { + pw.transferSink(TransferProgress{ + FileName: pw.fileName, + Written: pw.written, + Total: pw.total, + }) + } + if pw.downloadStatus == nil { + return + } + percentage := float64(0) + total := "" + if pw.total > 0 { + percentage = float64(pw.written) / float64(pw.total) * 100 + total = formatBytes(pw.total) + if pw.totalFiles > 1 { + percentage = percentage/float64(pw.totalFiles) + float64(pw.fileNo)*100/float64(pw.totalFiles) + } + } + pw.downloadStatus(pw.fileName, formatBytes(pw.written), total, percentage) +} + func (pw *progressWriter) Write(p []byte) (n int, err error) { // Check for cancellation before writing if pw.ctx != nil { @@ -41,24 +65,7 @@ func (pw *progressWriter) Write(p []byte) (n int, err error) { } } - if pw.total > 0 { - percentage := float64(pw.written) / float64(pw.total) * 100 - if pw.totalFiles > 1 { - // This is a multi-file download - // so we need to adjust the percentage - // to reflect the progress of the whole download - // This is the file pw.fileNo (0-indexed) of pw.totalFiles files. We assume that - // the files before successfully downloaded. - percentage = percentage / float64(pw.totalFiles) - if pw.fileNo > 0 { - percentage += float64(pw.fileNo) * 100 / float64(pw.totalFiles) - } - } - //log.Debug().Msgf("Downloading %s: %s/%s (%.2f%%)", pw.fileName, formatBytes(pw.written), formatBytes(pw.total), percentage) - pw.downloadStatus(pw.fileName, formatBytes(pw.written), formatBytes(pw.total), percentage) - } else { - pw.downloadStatus(pw.fileName, formatBytes(pw.written), "", 0) - } + pw.report() return } diff --git a/pkg/downloader/uri.go b/pkg/downloader/uri.go index 41bdbe672..0db07ddd1 100644 --- a/pkg/downloader/uri.go +++ b/pkg/downloader/uri.go @@ -52,8 +52,21 @@ type ImageVerifier interface { VerifyImage(ctx context.Context, imageRef string) error } +// TransferProgress reports the raw byte counts for an HTTP download. +// Total is negative when the server does not advertise a response length. +type TransferProgress struct { + FileName string + Written int64 + Total int64 +} + +// TransferProgressSink receives raw byte progress updates for an HTTP download. +type TransferProgressSink func(TransferProgress) + type downloadOptions struct { - verifier ImageVerifier + verifier ImageVerifier + bearerToken string + transferProgress TransferProgressSink } // DownloadOption configures DownloadFileWithContext / DownloadFile. @@ -70,6 +83,17 @@ func WithImageVerifier(v ImageVerifier) DownloadOption { return func(o *downloadOptions) { o.verifier = v } } +// WithBearerToken authenticates HTTP download requests with a bearer token. +// The token is stripped if a request redirects to a different origin. +func WithBearerToken(token string) DownloadOption { + return func(o *downloadOptions) { o.bearerToken = token } +} + +// WithTransferProgress attaches a sink for raw HTTP download byte progress. +func WithTransferProgress(sink TransferProgressSink) DownloadOption { + return func(o *downloadOptions) { o.transferProgress = sink } +} + func applyDownloadOptions(opts []DownloadOption) downloadOptions { var o downloadOptions for _, fn := range opts { @@ -367,9 +391,28 @@ func calculateHashForPartialFile(file *os.File) (hash.Hash, error) { // deadline so large downloads are not truncated. var downloadClient = httpclient.New(httpclient.WithFollowRedirects()) -func (uri URI) checkSeverSupportsRangeHeader() (bool, error) { - url := uri.ResolveURL() - resp, err := downloadClient.Head(url) +func newDownloadRequest( + ctx context.Context, + method string, + rawURL string, + bearerToken string, +) (*http.Request, error) { + req, err := http.NewRequestWithContext(ctx, method, rawURL, nil) + if err != nil { + return nil, err + } + if bearerToken != "" { + req.Header.Set("Authorization", "Bearer "+bearerToken) + } + return req, nil +} + +func (uri URI) checkServerSupportsRangeHeader(ctx context.Context, bearerToken string) (bool, error) { + req, err := newDownloadRequest(ctx, http.MethodHead, uri.ResolveURL(), bearerToken) + if err != nil { + return false, err + } + resp, err := downloadClient.Do(req) if err != nil { return false, err } @@ -563,21 +606,22 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string xlog.Info("Downloading", "url", url) - req, err := http.NewRequestWithContext(ctx, "GET", url, nil) + req, err := newDownloadRequest(ctx, http.MethodGet, url, dopts.bearerToken) if err != nil { return fmt.Errorf("failed to create request for %q: %v", filePath, err) } // save partial download to dedicated file tmpFilePath := filePath + ".partial" + var startPos int64 tmpFileInfo, err := os.Stat(tmpFilePath) if err == nil && uri.LooksLikeHTTPURL() { - support, err := uri.checkSeverSupportsRangeHeader() + support, err := uri.checkServerSupportsRangeHeader(ctx, dopts.bearerToken) if err != nil { return fmt.Errorf("failed to check if uri server supports range header: %v", err) } if support { - startPos := tmpFileInfo.Size() + startPos = tmpFileInfo.Size() req.Header.Set("Range", fmt.Sprintf("bytes=%d-", startPos)) } else { err := removePartialFile(tmpFilePath) @@ -622,6 +666,15 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string } //defer resp.Body.Close() + if startPos > 0 && resp.StatusCode != http.StatusPartialContent { + _ = resp.Body.Close() + _ = removePartialFile(tmpFilePath) + return fmt.Errorf( + "resume request for %q returned status %d instead of 206", + filePath, + resp.StatusCode, + ) + } if resp.StatusCode >= 400 { return fmt.Errorf("failed to download url %q, invalid status code %d", url, resp.StatusCode) } @@ -633,7 +686,7 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string if DownloadStallTimeout > 0 { source = newIdleTimeoutReader(resp.Body, DownloadStallTimeout) } - contentLength = resp.ContentLength + contentLength = resp.ContentLength + startPos } defer source.Close() @@ -644,11 +697,14 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string } // Create and write file - outFile, err := os.OpenFile(tmpFilePath, os.O_APPEND|os.O_RDWR|os.O_CREATE, 0644) + outFile, err := os.OpenFile(tmpFilePath, os.O_APPEND|os.O_RDWR|os.O_CREATE, 0600) if err != nil { return fmt.Errorf("failed to create / open file %q: %v", tmpFilePath, err) } defer outFile.Close() + if err := outFile.Chmod(0600); err != nil { + return fmt.Errorf("failed to restrict partial file %q permissions: %v", tmpFilePath, err) + } hash, err := calculateHashForPartialFile(outFile) if err != nil { return fmt.Errorf("failed to calculate hash for partial file") @@ -659,7 +715,9 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string hash: hash, fileNo: fileN, totalFiles: total, + written: startPos, downloadStatus: downloadStatus, + transferSink: dopts.transferProgress, ctx: ctx, } diff --git a/pkg/downloader/uri_test.go b/pkg/downloader/uri_test.go index 0a25f82a1..9cb667b57 100644 --- a/pkg/downloader/uri_test.go +++ b/pkg/downloader/uri_test.go @@ -238,8 +238,17 @@ var _ = Describe("Download Test", func() { } } respData = mockData[startPos:endPos] - w.WriteHeader(http.StatusOK) - w.Write(respData) + w.Header().Set("Content-Length", strconv.Itoa(len(respData))) + if rangeString != "" { + w.Header().Set( + "Content-Range", + fmt.Sprintf("bytes %d-%d/%d", startPos, endPos-1, len(mockData)), + ) + w.WriteHeader(http.StatusPartialContent) + } else { + w.WriteHeader(http.StatusOK) + } + _, _ = w.Write(respData) })) mockServer.EnableHTTP2 = true mockServer.Start() diff --git a/pkg/huggingface-api/client.go b/pkg/huggingface-api/client.go index 175c38574..1d1c7ae3c 100644 --- a/pkg/huggingface-api/client.go +++ b/pkg/huggingface-api/client.go @@ -1,6 +1,7 @@ package hfapi import ( + "context" "encoding/json" "errors" "fmt" @@ -112,6 +113,17 @@ func NewClient() *Client { } } +func (c *Client) newRequest(ctx context.Context, method, rawURL, token string) (*http.Request, error) { + req, err := http.NewRequestWithContext(ctx, method, rawURL, nil) + if err != nil { + return nil, err + } + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + return req, nil +} + // SearchModels searches for models using the Hugging Face API func (c *Client) SearchModels(params SearchParams) ([]Model, error) { for attempt := 1; attempt <= c.maxRetries; attempt++ { diff --git a/pkg/huggingface-api/snapshot.go b/pkg/huggingface-api/snapshot.go new file mode 100644 index 000000000..de4e74730 --- /dev/null +++ b/pkg/huggingface-api/snapshot.go @@ -0,0 +1,239 @@ +package hfapi + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/url" + "path" + "regexp" + "sort" + "strings" +) + +const defaultMaxSnapshotFiles = 100_000 + +var immutableRevisionPattern = regexp.MustCompile(`^[0-9a-fA-F]{40}$`) + +type SnapshotRequest struct { + Repo string + Revision string + Token string + AllowPatterns []string + IgnorePatterns []string + MaxFiles int +} + +type Snapshot struct { + Endpoint string + Repo string + RequestedRevision string + ResolvedRevision string + Files []SnapshotFile +} + +type SnapshotFile struct { + Path string + Size int64 + BlobOID string + LFSOID string + XetHash string + URL string +} + +func FilterSnapshotFiles(files []SnapshotFile, allow, ignore []string) ([]SnapshotFile, error) { + selected := make([]SnapshotFile, 0, len(files)) + for _, file := range files { + allowed, err := matchesAny(file.Path, allow) + if err != nil { + return nil, err + } + if len(allow) > 0 && !allowed { + continue + } + ignored, err := matchesAny(file.Path, ignore) + if err != nil { + return nil, err + } + if ignored { + continue + } + selected = append(selected, file) + } + sort.Slice(selected, func(i, j int) bool { return selected[i].Path < selected[j].Path }) + return selected, nil +} + +func matchesAny(filePath string, patterns []string) (bool, error) { + for _, pattern := range patterns { + matched, err := path.Match(pattern, filePath) + if err != nil { + return false, fmt.Errorf("invalid Hub file pattern %q: %w", pattern, err) + } + if !matched && !strings.Contains(pattern, "/") { + matched, err = path.Match(pattern, path.Base(filePath)) + if err != nil { + return false, fmt.Errorf("invalid Hub file pattern %q: %w", pattern, err) + } + } + if matched { + return true, nil + } + } + return false, nil +} + +type revisionResponse struct { + SHA string `json:"sha"` +} + +type snapshotTreeItem struct { + Type string `json:"type"` + OID string `json:"oid"` + Size int64 `json:"size"` + Path string `json:"path"` + LFS *LFSInfo `json:"lfs,omitempty"` + XetHash string `json:"xetHash,omitempty"` +} + +func (c *Client) ResolveSnapshot(ctx context.Context, req SnapshotRequest) (Snapshot, error) { + if err := ctx.Err(); err != nil { + return Snapshot{}, err + } + if req.MaxFiles <= 0 { + req.MaxFiles = defaultMaxSnapshotFiles + } + endpoint := strings.TrimSuffix(c.baseURL, "/api/models") + parsed, err := url.Parse(endpoint) + if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" { + return Snapshot{}, fmt.Errorf("invalid Hugging Face endpoint") + } + endpoint = strings.TrimRight(parsed.String(), "/") + owner, repo, err := splitRepo(req.Repo) + if err != nil { + return Snapshot{}, err + } + revision := req.Revision + if revision == "" { + revision = "main" + } + metadataURL := fmt.Sprintf("%s/api/models/%s/%s/revision/%s", endpoint, url.PathEscape(owner), url.PathEscape(repo), url.PathEscape(revision)) + metadataReq, err := c.newRequest(ctx, http.MethodGet, metadataURL, req.Token) + if err != nil { + return Snapshot{}, err + } + var metadata revisionResponse + if err := c.decodeJSON(metadataReq, &metadata); err != nil { + return Snapshot{}, fmt.Errorf("resolve Hugging Face revision: %w", err) + } + if !immutableRevisionPattern.MatchString(metadata.SHA) { + return Snapshot{}, fmt.Errorf("Hub returned invalid immutable revision") + } + resolvedRevision := strings.ToLower(metadata.SHA) + + treeURL := fmt.Sprintf("%s/api/models/%s/%s/tree/%s?recursive=true&expand=true", endpoint, url.PathEscape(owner), url.PathEscape(repo), resolvedRevision) + items := make([]snapshotTreeItem, 0) + for treeURL != "" { + treeReq, err := c.newRequest(ctx, http.MethodGet, treeURL, req.Token) + if err != nil { + return Snapshot{}, err + } + page, next, err := c.decodeTreePage(treeReq) + if err != nil { + return Snapshot{}, err + } + if len(items)+len(page) > req.MaxFiles { + return Snapshot{}, fmt.Errorf("snapshot exceeds maximum file count %d", req.MaxFiles) + } + items = append(items, page...) + treeURL = next + } + + files := make([]SnapshotFile, 0, len(items)) + for _, item := range items { + if item.Type != "file" { + continue + } + file := SnapshotFile{Path: item.Path, Size: item.Size, BlobOID: item.OID, XetHash: item.XetHash} + if item.LFS != nil { + file.LFSOID = item.LFS.Oid + } + file.URL = fmt.Sprintf("%s/%s/%s/resolve/%s/%s", endpoint, url.PathEscape(owner), url.PathEscape(repo), resolvedRevision, escapeFilePath(item.Path)) + files = append(files, file) + } + files, err = FilterSnapshotFiles(files, req.AllowPatterns, req.IgnorePatterns) + if err != nil { + return Snapshot{}, err + } + return Snapshot{Endpoint: endpoint, Repo: req.Repo, RequestedRevision: revision, ResolvedRevision: resolvedRevision, Files: files}, nil +} + +func splitRepo(repo string) (string, string, error) { + parts := strings.Split(repo, "/") + if len(parts) != 2 || parts[0] == "" || parts[1] == "" { + return "", "", fmt.Errorf("Hugging Face repo must be owner/repo") + } + return parts[0], parts[1], nil +} + +func escapeFilePath(filePath string) string { + parts := strings.Split(filePath, "/") + for i := range parts { + parts[i] = url.PathEscape(parts[i]) + } + return strings.Join(parts, "/") +} + +func (c *Client) decodeJSON(req *http.Request, dst any) error { + resp, err := c.client.Do(req) + if err != nil { + return err + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("Hub request returned status %d", resp.StatusCode) + } + return json.NewDecoder(http.MaxBytesReader(nil, resp.Body, 8<<20)).Decode(dst) +} + +func (c *Client) decodeTreePage(req *http.Request) ([]snapshotTreeItem, string, error) { + resp, err := c.client.Do(req) + if err != nil { + return nil, "", err + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK { + return nil, "", fmt.Errorf("Hub tree request returned status %d", resp.StatusCode) + } + var page []snapshotTreeItem + if err := json.NewDecoder(http.MaxBytesReader(nil, resp.Body, 32<<20)).Decode(&page); err != nil { + return nil, "", err + } + next, err := nextPage(resp.Header.Get("Link"), req.URL) + if err != nil { + return nil, "", err + } + return page, next, nil +} + +func nextPage(link string, current *url.URL) (string, error) { + for _, entry := range strings.Split(link, ",") { + parts := strings.Split(entry, ";") + if len(parts) >= 2 && strings.TrimSpace(parts[1]) == `rel="next"` { + raw := strings.Trim(strings.TrimSpace(parts[0]), "<>") + next, err := url.Parse(raw) + if err != nil { + return "", fmt.Errorf("invalid Hub pagination link: %w", err) + } + if !next.IsAbs() { + next = current.ResolveReference(next) + } + if next.Scheme != current.Scheme || next.Host != current.Host { + return "", fmt.Errorf("cross-origin pagination link rejected") + } + return next.String(), nil + } + } + return "", nil +} diff --git a/pkg/huggingface-api/snapshot_test.go b/pkg/huggingface-api/snapshot_test.go new file mode 100644 index 000000000..0ce318988 --- /dev/null +++ b/pkg/huggingface-api/snapshot_test.go @@ -0,0 +1,149 @@ +package hfapi_test + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/downloader" + hfapi "github.com/mudler/LocalAI/pkg/huggingface-api" +) + +var _ = Describe("immutable snapshot resolution", func() { + It("pins the metadata sha and follows tree pagination", func() { + const commit = "0123456789abcdef0123456789abcdef01234567" + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + Expect(r.Header.Get("Authorization")).To(Equal("Bearer test-token")) + switch { + case strings.Contains(r.URL.Path, "/revision/main"): + _, _ = w.Write([]byte(`{"sha":"` + commit + `"}`)) + case strings.Contains(r.URL.Path, "/tree/"+commit) && r.URL.Query().Get("cursor") == "page-2": + _, _ = w.Write([]byte(`[{"type":"file","path":"tokenizer/tokenizer.json","size":7,"oid":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"}]`)) + case strings.Contains(r.URL.Path, "/tree/"+commit): + Expect(r.URL.Query().Get("recursive")).To(Equal("true")) + w.Header().Set("Link", fmt.Sprintf(`<%s%s?recursive=true&expand=true&cursor=page-2>; rel="next"`, server.URL, r.URL.Path)) + _, _ = w.Write([]byte(`[{"type":"file","path":"model.safetensors","size":11,"oid":"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb","lfs":{"oid":"cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc","size":11,"pointerSize":130}}]`)) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + DeferCleanup(server.Close) + + client := hfapi.NewClient() + client.SetBaseURL(server.URL + "/api/models") + snapshot, err := client.ResolveSnapshot(context.Background(), hfapi.SnapshotRequest{ + Repo: "owner/repo", + Revision: "main", + Token: "test-token", + MaxFiles: 10, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(snapshot.Endpoint).To(Equal(server.URL)) + Expect(snapshot.ResolvedRevision).To(Equal(commit)) + Expect(snapshot.Files).To(HaveLen(2)) + Expect(snapshot.Files[0].Path).To(Equal("model.safetensors")) + Expect(snapshot.Files[0].LFSOID).To(HaveLen(64)) + Expect(snapshot.Files[1].Path).To(Equal("tokenizer/tokenizer.json")) + Expect(snapshot.Files[0].URL).To(ContainSubstring("/resolve/" + commit + "/model.safetensors")) + }) + + It("applies deterministic allow and ignore filters", func() { + files := []hfapi.SnapshotFile{ + {Path: "config.json"}, + {Path: "nested/tokenizer.json"}, + {Path: "weights.bin"}, + {Path: "weights.safetensors"}, + } + filtered, err := hfapi.FilterSnapshotFiles(files, []string{"*.json", "*.safetensors"}, []string{"nested/*"}) + Expect(err).NotTo(HaveOccurred()) + Expect(filtered).To(ConsistOf( + hfapi.SnapshotFile{Path: "config.json"}, + hfapi.SnapshotFile{Path: "weights.safetensors"}, + )) + }) + + It("rejects a manifest larger than MaxFiles", func() { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.Contains(r.URL.Path, "/revision/") { + _, _ = w.Write([]byte(`{"sha":"0123456789abcdef0123456789abcdef01234567"}`)) + return + } + _, _ = w.Write([]byte(`[{"type":"file","path":"a","size":1},{"type":"file","path":"b","size":1}]`)) + })) + DeferCleanup(server.Close) + client := hfapi.NewClient() + client.SetBaseURL(server.URL + "/api/models") + _, err := client.ResolveSnapshot(context.Background(), hfapi.SnapshotRequest{Repo: "owner/repo", Revision: "main", MaxFiles: 1}) + Expect(err).To(MatchError(ContainSubstring("maximum file count"))) + }) + + It("honors a cancelled context before metadata access", func() { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := hfapi.NewClient().ResolveSnapshot(ctx, hfapi.SnapshotRequest{Repo: "owner/repo", Revision: "main"}) + Expect(err).To(MatchError(context.Canceled)) + }) + + It("rejects a cross-origin pagination link before sending credentials", func() { + const commit = "0123456789abcdef0123456789abcdef01234567" + requestedCrossOrigin := false + evil := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + requestedCrossOrigin = true + w.WriteHeader(http.StatusNoContent) + })) + DeferCleanup(evil.Close) + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.Contains(r.URL.Path, "/revision/") { + _, _ = w.Write([]byte(`{"sha":"` + commit + `"}`)) + return + } + w.Header().Set("Link", `<`+evil.URL+`/steal>; rel="next"`) + _, _ = w.Write([]byte(`[]`)) + })) + DeferCleanup(origin.Close) + client := hfapi.NewClient() + client.SetBaseURL(origin.URL + "/api/models") + _, err := client.ResolveSnapshot(context.Background(), hfapi.SnapshotRequest{ + Repo: "owner/repo", Revision: "main", Token: "test-token", + }) + Expect(err).To(MatchError(ContainSubstring("cross-origin pagination"))) + Expect(requestedCrossOrigin).To(BeFalse()) + }) + + It("downloads a pinned public Xet fixture through the ordinary resolve URL", Label("network"), func() { + if os.Getenv("LOCALAI_HF_XET_SMOKE") != "1" { + Skip("set LOCALAI_HF_XET_SMOKE=1 to run the public Hub compatibility gate") + } + client := hfapi.NewClient() + snapshot, err := client.ResolveSnapshot(context.Background(), hfapi.SnapshotRequest{ + Repo: "hf-internal-testing/tiny-random-t5-v1.1", + Revision: "main", + AllowPatterns: []string{"model.safetensors"}, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(snapshot.ResolvedRevision).To(MatchRegexp(`^[0-9a-f]{40}$`)) + Expect(snapshot.Files).To(HaveLen(1)) + file := snapshot.Files[0] + Expect(file.Path).To(Equal("model.safetensors")) + Expect(file.XetHash).NotTo(BeEmpty()) + target := filepath.Join(GinkgoT().TempDir(), file.Path) + var final downloader.TransferProgress + err = downloader.URI(file.URL).DownloadFileWithContext( + context.Background(), target, file.LFSOID, 0, 1, nil, + downloader.WithTransferProgress(func(event downloader.TransferProgress) { final = event }), + ) + Expect(err).NotTo(HaveOccurred()) + info, err := os.Stat(target) + Expect(err).NotTo(HaveOccurred()) + Expect(info.Size()).To(Equal(file.Size)) + Expect(final.Written).To(Equal(file.Size)) + }) +}) diff --git a/pkg/model/initializers.go b/pkg/model/initializers.go index 43273a662..7d9afca05 100644 --- a/pkg/model/initializers.go +++ b/pkg/model/initializers.go @@ -251,7 +251,11 @@ func (ml *ModelLoader) backendLoader(opts ...Option) (client grpc.Backend, err e backend = realBackend } - model, err := ml.LoadModel(o.modelID, o.model, ml.grpcModel(backend, o)) + modelFileName := o.modelFile + if modelFileName == "" { + modelFileName = o.model + } + model, err := ml.LoadModelWithFile(o.modelID, o.model, modelFileName, ml.grpcModel(backend, o)) if err != nil { // Defensive cleanup: the model usually wasn't registered yet (LoadModel // failed before that), so StopGRPC reporting "model not found" is the diff --git a/pkg/model/loader.go b/pkg/model/loader.go index 17420f3e2..f4bb4e7da 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -397,14 +397,21 @@ func (ml *ModelLoader) ListLoadedModels() []*Model { } func (ml *ModelLoader) LoadModel(modelID, modelName string, loader func(string, string, string) (*Model, error)) (*Model, error) { - return ml.loadModel(modelID, modelName, loader, true) + return ml.LoadModelWithFile(modelID, modelName, modelName, loader) } -// loadModel is the implementation behind LoadModel. checkCooldown gates fresh, +func (ml *ModelLoader) LoadModelWithFile(modelID, modelName, modelFileName string, loader func(string, string, string) (*Model, error)) (*Model, error) { + if modelFileName == "" { + modelFileName = modelName + } + return ml.loadModel(modelID, modelName, modelFileName, loader, true) +} + +// loadModel is the implementation behind LoadModelWithFile. checkCooldown gates fresh, // independent load triggers behind the per-model failure cooldown; it is set to // false for the coalesced retry of an in-flight burst (a follower whose leader // just failed), which is not a new trigger and should still get its one retry. -func (ml *ModelLoader) loadModel(modelID, modelName string, loader func(string, string, string) (*Model, error), checkCooldown bool) (*Model, error) { +func (ml *ModelLoader) loadModel(modelID, modelName, modelFileName string, loader func(string, string, string) (*Model, error), checkCooldown bool) (*Model, error) { ml.mu.Lock() distributed := ml.modelRouter != nil ml.mu.Unlock() @@ -430,7 +437,7 @@ func (ml *ModelLoader) loadModel(modelID, modelName string, loader func(string, // per inference. Trade-off: cross-frontend in-flight visibility // becomes eventually consistent, acceptable for 1-3 frontend // deployments. - modelFile := filepath.Join(ml.ModelPath, modelName) + modelFile := filepath.Join(ml.ModelPath, modelFileName) model, err := loader(modelID, modelName, modelFile) if err != nil { return nil, fmt.Errorf("failed to route model with internal loader: %s", err) @@ -484,7 +491,7 @@ func (ml *ModelLoader) loadModel(modelID, modelName string, loader func(string, } // If still not loaded, the other goroutine failed. Retry once as part of // this burst, bypassing the cooldown gate (we are not a new trigger). - return ml.loadModel(modelID, modelName, loader, false) + return ml.loadModel(modelID, modelName, modelFileName, loader, false) } // Mark this model as loading (create a channel that will be closed when done) @@ -501,7 +508,7 @@ func (ml *ModelLoader) loadModel(modelID, modelName string, loader func(string, }() // Load the model (this can take a long time, no lock held) - modelFile := filepath.Join(ml.ModelPath, modelName) + modelFile := filepath.Join(ml.ModelPath, modelFileName) xlog.Debug("Loading model in memory from file", "file", modelFile) model, err := loader(modelID, modelName, modelFile) diff --git a/pkg/model/loader_options.go b/pkg/model/loader_options.go index d247a00c8..9552c3c67 100644 --- a/pkg/model/loader_options.go +++ b/pkg/model/loader_options.go @@ -9,6 +9,7 @@ import ( type Options struct { backendString string model string + modelFile string modelID string context context.Context @@ -73,6 +74,12 @@ func WithModel(modelFile string) Option { } } +func WithModelFile(modelFile string) Option { + return func(o *Options) { + o.modelFile = modelFile + } +} + func WithLoadGRPCLoadModelOpts(opts *pb.ModelOptions) Option { return func(o *Options) { o.gRPCOptions = opts diff --git a/pkg/model/loader_test.go b/pkg/model/loader_test.go index e6a6148c3..1c8892bc3 100644 --- a/pkg/model/loader_test.go +++ b/pkg/model/loader_test.go @@ -73,6 +73,22 @@ var _ = Describe("ModelLoader", func() { }) Context("LoadModel", func() { + It("passes a logical model and managed model file independently", func() { + const relative = ".artifacts/huggingface/0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef/snapshot" + var receivedName, receivedFile string + mockModel = model.NewModel("managed", "test.model", nil) + mockModel.MarkHealthy() + mockLoader := func(_ string, modelName, modelFile string) (*model.Model, error) { + receivedName, receivedFile = modelName, modelFile + return mockModel, nil + } + + _, err := modelLoader.LoadModelWithFile("managed", "owner/repo", relative, mockLoader) + Expect(err).NotTo(HaveOccurred()) + Expect(receivedName).To(Equal("owner/repo")) + Expect(receivedFile).To(Equal(filepath.Join(modelPath, filepath.FromSlash(relative)))) + }) + It("should load a model and keep it in memory", func() { mockModel = model.NewModel("foo", "test.model", nil) mockModel.MarkHealthy() // skip gRPC health check (no real server) diff --git a/pkg/modelartifacts/manifest.go b/pkg/modelartifacts/manifest.go new file mode 100644 index 000000000..8e42f0cd7 --- /dev/null +++ b/pkg/modelartifacts/manifest.go @@ -0,0 +1,47 @@ +package modelartifacts + +import ( + "encoding/json" + "fmt" + "os" +) + +const ManifestVersion = 1 + +type Manifest struct { + Version int `json:"version"` + Artifact Spec `json:"artifact"` + Files []ManifestFile `json:"files"` +} + +type ManifestFile struct { + Path string `json:"path"` + Size int64 `json:"size"` + SHA256 string `json:"sha256"` + BlobOID string `json:"blob_oid,omitempty"` + LFSOID string `json:"lfs_oid,omitempty"` + XetHash string `json:"xet_hash,omitempty"` +} + +func ReadManifest(fileName string) (Manifest, error) { + data, err := os.ReadFile(fileName) + if err != nil { + return Manifest{}, err + } + var manifest Manifest + if err := json.Unmarshal(data, &manifest); err != nil { + return Manifest{}, err + } + if manifest.Version != ManifestVersion || manifest.Artifact.Resolved == nil { + return Manifest{}, fmt.Errorf("unsupported or incomplete artifact manifest") + } + for _, file := range manifest.Files { + if err := ValidateRelativeHubPath(file.Path); err != nil { + return Manifest{}, err + } + if file.Size < 0 || len(file.SHA256) != 64 { + return Manifest{}, fmt.Errorf("invalid manifest entry %q", file.Path) + } + } + return manifest, nil +} diff --git a/pkg/modelartifacts/materializer.go b/pkg/modelartifacts/materializer.go new file mode 100644 index 000000000..da9faa5d9 --- /dev/null +++ b/pkg/modelartifacts/materializer.go @@ -0,0 +1,359 @@ +package modelartifacts + +import ( + "context" + "crypto/sha1" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "os" + "path" + "path/filepath" + "strings" + "time" + + "github.com/gofrs/flock" + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/pkg/downloader" + hfapi "github.com/mudler/LocalAI/pkg/huggingface-api" +) + +type SnapshotResolver interface { + ResolveSnapshot(context.Context, hfapi.SnapshotRequest) (hfapi.Snapshot, error) +} + +type Manager struct { + resolver SnapshotResolver + huggingFaceToken string +} + +type ManagerOption func(*Manager) + +type Result struct { + Spec Spec + RelativePath string + Manifest Manifest + CacheHit bool +} + +func WithHuggingFaceToken(token string) ManagerOption { + return func(manager *Manager) { manager.huggingFaceToken = token } +} + +func NewManager(resolver SnapshotResolver, options ...ManagerOption) *Manager { + manager := &Manager{resolver: resolver} + for _, option := range options { + option(manager) + } + return manager +} + +func NewDefaultManager(options ...ManagerOption) *Manager { + client := hfapi.NewClient() + client.SetBaseURL(strings.TrimRight(downloader.HF_ENDPOINT, "/") + "/api/models") + return NewManager(client, options...) +} + +func committedResult(modelsPath string, spec Spec) (Result, bool) { + if spec.Resolved == nil || spec.Resolved.CacheKey == "" { + return Result{}, false + } + layout, err := LayoutFor(modelsPath, spec) + if err != nil { + return Result{}, false + } + manifest, err := ReadManifest(layout.Manifest) + if err != nil || manifest.Artifact.Resolved == nil || manifest.Artifact.Resolved.CacheKey != spec.Resolved.CacheKey { + return Result{}, false + } + specKey, err := CacheKey(spec) + if err != nil || specKey != spec.Resolved.CacheKey { + return Result{}, false + } + manifestKey, err := CacheKey(manifest.Artifact) + if err != nil || manifestKey != spec.Resolved.CacheKey || len(manifest.Files) == 0 { + return Result{}, false + } + for _, file := range manifest.Files { + info, err := os.Stat(filepath.Join(layout.Snapshot, filepath.FromSlash(file.Path))) + if err != nil || !info.Mode().IsRegular() || info.Size() != file.Size { + return Result{}, false + } + } + relative, err := RelativeSnapshotPath(spec.Resolved.CacheKey) + if err != nil { + return Result{}, false + } + return Result{Spec: spec, RelativePath: relative, Manifest: manifest, CacheHit: true}, true +} + +func (m *Manager) Ensure(ctx context.Context, modelsPath string, spec Spec) (Result, error) { + if err := ctx.Err(); err != nil { + return Result{}, err + } + normalized, err := spec.Normalize() + if err != nil { + return Result{}, err + } + if cached, ok := committedResult(modelsPath, normalized); ok { + return cached, nil + } + if m == nil || m.resolver == nil { + return Result{}, fmt.Errorf("artifact materializer has no snapshot resolver") + } + token := "" + if normalized.Source.TokenEnv == HuggingFaceTokenEnv { + token = m.huggingFaceToken + if token == "" { + return Result{}, fmt.Errorf("artifact requires non-empty %s", HuggingFaceTokenEnv) + } + } + revision := normalized.Source.Revision + if normalized.Resolved != nil { + revision = normalized.Resolved.Revision + } + ReportProgress(ctx, ProgressEvent{Phase: PhaseResolving, Artifact: normalized.Name}) + snapshot, err := m.resolver.ResolveSnapshot(ctx, hfapi.SnapshotRequest{ + Repo: normalized.Source.Repo, Revision: revision, Token: token, + AllowPatterns: normalized.Source.AllowPatterns, IgnorePatterns: normalized.Source.IgnorePatterns, + }) + if err != nil { + return Result{}, err + } + if normalized.Resolved != nil && (snapshot.Endpoint != normalized.Resolved.Endpoint || snapshot.ResolvedRevision != normalized.Resolved.Revision) { + return Result{}, fmt.Errorf("resolved artifact identity changed; reinstall the model") + } + normalized.Resolved = &Resolved{Endpoint: snapshot.Endpoint, Revision: snapshot.ResolvedRevision} + cacheKey, err := CacheKey(normalized) + if err != nil { + return Result{}, err + } + normalized.Resolved.CacheKey = cacheKey + layout, err := LayoutFor(modelsPath, normalized) + if err != nil { + return Result{}, err + } + if err := os.MkdirAll(filepath.Dir(layout.Lock), 0o750); err != nil { + return Result{}, err + } + artifactLock := flock.New(layout.Lock) + locked, err := artifactLock.TryLockContext(ctx, 100*time.Millisecond) + if err != nil { + return Result{}, err + } + if !locked { + if err := ctx.Err(); err != nil { + return Result{}, err + } + return Result{}, fmt.Errorf("artifact lock was not acquired") + } + defer func() { + if err := artifactLock.Unlock(); err != nil { + xlog.Warn("failed to unlock model artifact", "lock", layout.Lock, "error", err) + } + }() + if err := os.Chmod(layout.Lock, 0o600); err != nil { + return Result{}, err + } + if cached, ok := committedResult(modelsPath, normalized); ok { + return cached, nil + } + if err := removeInvalidFinal(layout); err != nil { + return Result{}, err + } + return m.materializeLocked(ctx, modelsPath, normalized, snapshot, token, layout) +} + +func removeInvalidFinal(layout Layout) error { + root, err := os.OpenRoot(layout.Root) + if err != nil { + return err + } + defer func() { _ = root.Close() }() + relative, err := filepath.Rel(layout.Root, layout.Final) + if err != nil || filepath.Dir(relative) != "huggingface" || !cacheKeyPattern.MatchString(filepath.Base(relative)) { + return fmt.Errorf("refusing to remove invalid artifact path %q", layout.Final) + } + return root.RemoveAll(relative) +} + +func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec Spec, snapshot hfapi.Snapshot, token string, layout Layout) (Result, error) { + if err := os.MkdirAll(layout.Partial, 0o750); err != nil { + return Result{}, err + } + root, err := os.OpenRoot(layout.Partial) + if err != nil { + return Result{}, err + } + rootClosed := false + defer func() { + if !rootClosed { + _ = root.Close() + } + }() + if err := root.MkdirAll(".downloads", 0o750); err != nil { + return Result{}, err + } + if err := root.MkdirAll("snapshot", 0o750); err != nil { + return Result{}, err + } + totalBytes := int64(0) + if len(snapshot.Files) == 0 { + return Result{}, fmt.Errorf("resolved snapshot contains no selected files") + } + seenPaths := make(map[string]struct{}, len(snapshot.Files)) + for _, file := range snapshot.Files { + if err := ValidateRelativeHubPath(file.Path); err != nil { + return Result{}, err + } + if file.Size < 0 || totalBytes > int64(^uint64(0)>>1)-file.Size { + return Result{}, fmt.Errorf("invalid aggregate snapshot size") + } + if _, exists := seenPaths[file.Path]; exists { + return Result{}, fmt.Errorf("duplicate Hub path %q", file.Path) + } + seenPaths[file.Path] = struct{}{} + totalBytes += file.Size + } + + manifest := Manifest{Version: ManifestVersion, Artifact: spec, Files: make([]ManifestFile, 0, len(snapshot.Files))} + completedBytes := int64(0) + tasks := make([]downloader.FileTask, 0, len(snapshot.Files)) + for index, file := range snapshot.Files { + if err := ctx.Err(); err != nil { + return Result{}, err + } + nameSum := sha256.Sum256([]byte(file.Path)) + blobRel := path.Join(".downloads", hex.EncodeToString(nameSum[:])) + blobAbs := filepath.Join(layout.Partial, filepath.FromSlash(blobRel)) + file := file + taskIndex := index + task := downloader.FileTask{ + URI: downloader.URI(file.URL), + Destination: blobAbs, + SHA256: file.LFSOID, + FileIndex: taskIndex, + TotalFiles: len(snapshot.Files), + Options: []downloader.DownloadOption{ + downloader.WithBearerToken(token), + downloader.WithTransferProgress(func(event downloader.TransferProgress) { + ReportProgress(ctx, ProgressEvent{ + Phase: PhaseDownloading, + Artifact: spec.Name, + File: file.Path, + CurrentBytes: completedBytes + event.Written, + TotalBytes: totalBytes, + CompletedFiles: taskIndex, + TotalFiles: len(snapshot.Files), + }) + }), + }, + AfterDownload: func(string) error { + ReportProgress(ctx, ProgressEvent{ + Phase: PhaseVerifying, + Artifact: spec.Name, + File: file.Path, + CurrentBytes: completedBytes + file.Size, + TotalBytes: totalBytes, + CompletedFiles: taskIndex, + TotalFiles: len(snapshot.Files), + }) + entry, err := verifyDownloadedFile(blobAbs, file) + if err != nil { + _ = root.Remove(blobRel) + return err + } + destination := path.Join("snapshot", file.Path) + if err := root.MkdirAll(path.Dir(destination), 0o750); err != nil { + return err + } + _ = root.Remove(destination) + if err := root.Rename(blobRel, destination); err != nil { + return err + } + manifest.Files = append(manifest.Files, entry) + completedBytes += file.Size + return nil + }, + } + tasks = append(tasks, task) + } + if err := downloader.DownloadFilesWithContext(ctx, tasks, nil); err != nil { + return Result{}, err + } + if err := root.RemoveAll(".downloads"); err != nil { + return Result{}, err + } + encoded, err := json.MarshalIndent(manifest, "", " ") + if err != nil { + return Result{}, err + } + if err := root.WriteFile("manifest.json.tmp", append(encoded, '\n'), 0o644); err != nil { + return Result{}, err + } + if err := root.Rename("manifest.json.tmp", "manifest.json"); err != nil { + return Result{}, err + } + if err := root.Close(); err != nil { + return Result{}, err + } + rootClosed = true + if err := os.MkdirAll(filepath.Dir(layout.Final), 0o750); err != nil { + return Result{}, err + } + ReportProgress(ctx, ProgressEvent{Phase: PhaseCommitting, Artifact: spec.Name, CurrentBytes: totalBytes, TotalBytes: totalBytes, CompletedFiles: len(snapshot.Files), TotalFiles: len(snapshot.Files)}) + if err := os.Rename(layout.Partial, layout.Final); err != nil { + return Result{}, err + } + relative, err := RelativeSnapshotPath(spec.Resolved.CacheKey) + if err != nil { + return Result{}, err + } + return Result{Spec: spec, RelativePath: relative, Manifest: manifest}, nil +} + +func verifyDownloadedFile(fileName string, source hfapi.SnapshotFile) (ManifestFile, error) { + file, err := os.Open(fileName) + if err != nil { + return ManifestFile{}, err + } + defer func() { _ = file.Close() }() + sha256Hash := sha256.New() + gitHash := sha1.New() + if _, err := fmt.Fprintf(gitHash, "blob %d%c", source.Size, byte(0)); err != nil { + return ManifestFile{}, err + } + if _, err := io.Copy(io.MultiWriter(sha256Hash, gitHash), file); err != nil { + return ManifestFile{}, err + } + info, err := file.Stat() + if err != nil { + return ManifestFile{}, err + } + if info.Size() != source.Size { + return ManifestFile{}, fmt.Errorf("size mismatch for %q", source.Path) + } + rawSHA256 := hex.EncodeToString(sha256Hash.Sum(nil)) + if source.LFSOID != "" { + if decoded, err := hex.DecodeString(source.LFSOID); err != nil || len(decoded) != sha256.Size { + return ManifestFile{}, fmt.Errorf("invalid LFS SHA-256 for %q", source.Path) + } + if !strings.EqualFold(rawSHA256, source.LFSOID) { + return ManifestFile{}, fmt.Errorf("LFS SHA-256 mismatch for %q", source.Path) + } + } else if source.BlobOID != "" { + if decoded, err := hex.DecodeString(source.BlobOID); err != nil || len(decoded) != sha1.Size { + return ManifestFile{}, fmt.Errorf("invalid Git blob OID for %q", source.Path) + } + if !strings.EqualFold(hex.EncodeToString(gitHash.Sum(nil)), source.BlobOID) { + return ManifestFile{}, fmt.Errorf("Git blob OID mismatch for %q", source.Path) + } + } + return ManifestFile{ + Path: source.Path, Size: source.Size, SHA256: rawSHA256, + BlobOID: source.BlobOID, LFSOID: source.LFSOID, XetHash: source.XetHash, + }, nil +} diff --git a/pkg/modelartifacts/materializer_test.go b/pkg/modelartifacts/materializer_test.go new file mode 100644 index 000000000..ea80c1331 --- /dev/null +++ b/pkg/modelartifacts/materializer_test.go @@ -0,0 +1,153 @@ +package modelartifacts_test + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync" + "sync/atomic" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + hfapi "github.com/mudler/LocalAI/pkg/huggingface-api" + "github.com/mudler/LocalAI/pkg/modelartifacts" +) + +type fakeSnapshotResolver struct { + mu sync.Mutex + snapshot hfapi.Snapshot + err error + calls int +} + +func (f *fakeSnapshotResolver) ResolveSnapshot(context.Context, hfapi.SnapshotRequest) (hfapi.Snapshot, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.calls++ + return f.snapshot, f.err +} + +func (f *fakeSnapshotResolver) callCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.calls +} + +var _ = Describe("controller artifact materializer", func() { + It("requires an injected token before resolving a private artifact", func() { + resolver := &fakeSnapshotResolver{} + _, err := modelartifacts.NewManager(resolver).Ensure(context.Background(), GinkgoT().TempDir(), + modelartifacts.Spec{Source: modelartifacts.Source{ + Type: modelartifacts.SourceTypeHuggingFace, Repo: "owner/private", TokenEnv: modelartifacts.HuggingFaceTokenEnv, + }}) + Expect(err).To(MatchError(ContainSubstring("non-empty HF_TOKEN"))) + Expect(resolver.callCount()).To(BeZero()) + }) + + It("downloads, commits, and reuses a pinned snapshot", func() { + weight := []byte("weight-bytes") + sum := sha256.Sum256(weight) + var requests atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + Expect(r.Header.Get("Authorization")).To(Equal("Bearer hf-secret")) + w.Header().Set("Content-Length", "12") + _, _ = w.Write(weight) + })) + DeferCleanup(server.Close) + + resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{ + Endpoint: "https://huggingface.co", Repo: "owner/repo", + RequestedRevision: "main", ResolvedRevision: "0123456789abcdef0123456789abcdef01234567", + Files: []hfapi.SnapshotFile{{Path: "nested/model.safetensors", Size: 12, LFSOID: hex.EncodeToString(sum[:]), URL: server.URL + "/model"}}, + }} + manager := modelartifacts.NewManager(resolver, modelartifacts.WithHuggingFaceToken("hf-secret")) + modelsPath := GinkgoT().TempDir() + spec := modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", TokenEnv: "HF_TOKEN"}} + + result, err := manager.Ensure(context.Background(), modelsPath, spec) + Expect(err).NotTo(HaveOccurred()) + Expect(result.CacheHit).To(BeFalse()) + Expect(result.Spec.Resolved.Revision).To(Equal("0123456789abcdef0123456789abcdef01234567")) + Expect(result.RelativePath).To(HavePrefix(".artifacts/huggingface/")) + Expect(os.ReadFile(filepath.Join(modelsPath, filepath.FromSlash(result.RelativePath), "nested", "model.safetensors"))).To(Equal(weight)) + + manifestBytes, err := os.ReadFile(filepath.Join(modelsPath, filepath.Dir(filepath.FromSlash(result.RelativePath)), "manifest.json")) + Expect(err).NotTo(HaveOccurred()) + Expect(string(manifestBytes)).NotTo(ContainSubstring("hf-secret")) + + cached, err := manager.Ensure(context.Background(), modelsPath, result.Spec) + Expect(err).NotTo(HaveOccurred()) + Expect(cached.CacheHit).To(BeTrue()) + Expect(requests.Load()).To(Equal(int32(1))) + Expect(resolver.callCount()).To(Equal(1)) + }) + + It("rejects a path escape before opening a destination", func() { + resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{ + Endpoint: "https://huggingface.co", Repo: "owner/repo", + ResolvedRevision: "0123456789abcdef0123456789abcdef01234567", + Files: []hfapi.SnapshotFile{{Path: "../escape", Size: 1, URL: "https://example.invalid/file"}}, + }} + modelsPath := GinkgoT().TempDir() + _, err := modelartifacts.NewManager(resolver).Ensure(context.Background(), modelsPath, + modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}}) + Expect(err).To(MatchError(ContainSubstring("unsafe Hub path"))) + Expect(filepath.Join(modelsPath, "escape")).NotTo(BeAnExistingFile()) + }) + + It("does not commit after cancellation", func() { + started := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + close(started) + <-r.Context().Done() + })) + DeferCleanup(server.Close) + resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{ + Endpoint: "https://huggingface.co", Repo: "owner/repo", + ResolvedRevision: "0123456789abcdef0123456789abcdef01234567", + Files: []hfapi.SnapshotFile{{Path: "model.bin", Size: 10, URL: server.URL}}, + }} + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + modelsPath := GinkgoT().TempDir() + go func() { + _, err := modelartifacts.NewManager(resolver).Ensure(ctx, modelsPath, + modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}}) + done <- err + }() + <-started + cancel() + Expect(<-done).To(MatchError(context.Canceled)) + entries, err := filepath.Glob(filepath.Join(modelsPath, ".artifacts", "huggingface", "*")) + Expect(err).NotTo(HaveOccurred()) + Expect(entries).To(BeEmpty()) + }) + + It("emits resolving, downloading, verifying, and committing phases", func() { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write([]byte("x")) })) + DeferCleanup(server.Close) + resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{ + Endpoint: "https://huggingface.co", Repo: "owner/repo", + ResolvedRevision: "0123456789abcdef0123456789abcdef01234567", + Files: []hfapi.SnapshotFile{{Path: "config.json", Size: 1, URL: server.URL}}, + }} + var phases []modelartifacts.Phase + ctx := modelartifacts.WithProgressSink(context.Background(), func(event modelartifacts.ProgressEvent) { + if len(phases) == 0 || phases[len(phases)-1] != event.Phase { + phases = append(phases, event.Phase) + } + }) + _, err := modelartifacts.NewManager(resolver).Ensure(ctx, GinkgoT().TempDir(), + modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}}) + Expect(err).NotTo(HaveOccurred()) + Expect(phases).To(Equal([]modelartifacts.Phase{ + modelartifacts.PhaseResolving, modelartifacts.PhaseDownloading, modelartifacts.PhaseVerifying, modelartifacts.PhaseCommitting, + })) + }) +}) diff --git a/pkg/modelartifacts/modelartifacts_suite_test.go b/pkg/modelartifacts/modelartifacts_suite_test.go new file mode 100644 index 000000000..e1ac6be89 --- /dev/null +++ b/pkg/modelartifacts/modelartifacts_suite_test.go @@ -0,0 +1,13 @@ +package modelartifacts_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestModelArtifacts(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Model Artifacts Suite") +} diff --git a/pkg/modelartifacts/path.go b/pkg/modelartifacts/path.go new file mode 100644 index 000000000..bde8fbbad --- /dev/null +++ b/pkg/modelartifacts/path.go @@ -0,0 +1,86 @@ +package modelartifacts + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "path/filepath" + "strings" +) + +type Layout struct { + Root string + Lock string + Partial string + PartialSnapshot string + Final string + Snapshot string + Manifest string +} + +func CacheKey(spec Spec) (string, error) { + normalized, err := spec.Normalize() + if err != nil { + return "", err + } + if normalized.Resolved == nil { + return "", fmt.Errorf("cache key requires resolved artifact state") + } + identity := struct { + Type string `json:"type"` + Endpoint string `json:"endpoint"` + Repo string `json:"repo"` + Revision string `json:"revision"` + AllowPatterns []string `json:"allow_patterns,omitempty"` + IgnorePatterns []string `json:"ignore_patterns,omitempty"` + }{ + Type: normalized.Source.Type, Endpoint: normalized.Resolved.Endpoint, + Repo: normalized.Source.Repo, Revision: normalized.Resolved.Revision, + AllowPatterns: normalized.Source.AllowPatterns, IgnorePatterns: normalized.Source.IgnorePatterns, + } + encoded, err := json.Marshal(identity) + if err != nil { + return "", err + } + sum := sha256.Sum256(encoded) + return hex.EncodeToString(sum[:]), nil +} + +func RelativeSnapshotPath(cacheKey string) (string, error) { + if !cacheKeyPattern.MatchString(cacheKey) { + return "", fmt.Errorf("invalid artifact cache key") + } + return filepath.Join(".artifacts", "huggingface", cacheKey, "snapshot"), nil +} + +func LayoutFor(modelsPath string, spec Spec) (Layout, error) { + if spec.Resolved == nil || !cacheKeyPattern.MatchString(spec.Resolved.CacheKey) { + return Layout{}, fmt.Errorf("layout requires a resolved cache key") + } + root := filepath.Join(modelsPath, ".artifacts") + partial := filepath.Join(root, ".partial", spec.Resolved.CacheKey) + final := filepath.Join(root, "huggingface", spec.Resolved.CacheKey) + return Layout{ + Root: root, + Lock: filepath.Join(root, ".locks", spec.Resolved.CacheKey+".lock"), + Partial: partial, + PartialSnapshot: filepath.Join(partial, "snapshot"), + Final: final, + Snapshot: filepath.Join(final, "snapshot"), + Manifest: filepath.Join(final, "manifest.json"), + }, nil +} + +func ValidateRelativeHubPath(candidate string) error { + if candidate == "" || filepath.IsAbs(candidate) || strings.ContainsAny(candidate, "\\\x00") { + return fmt.Errorf("unsafe Hub path %q", candidate) + } + parts := strings.Split(candidate, "/") + for _, part := range parts { + if part == "" || part == "." || part == ".." { + return fmt.Errorf("unsafe Hub path %q", candidate) + } + } + return nil +} diff --git a/pkg/modelartifacts/path_test.go b/pkg/modelartifacts/path_test.go new file mode 100644 index 000000000..df6aee577 --- /dev/null +++ b/pkg/modelartifacts/path_test.go @@ -0,0 +1,109 @@ +package modelartifacts_test + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/modelartifacts" +) + +var _ = Describe("artifact storage primitives", func() { + baseSpec := func() modelartifacts.Spec { + return modelartifacts.Spec{ + Name: "model", + Target: "model", + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", Revision: "main"}, + Resolved: &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + }, + } + } + + It("creates a stable path-safe cache identity", func() { + first, err := modelartifacts.CacheKey(baseSpec()) + Expect(err).NotTo(HaveOccurred()) + second := baseSpec() + second.Source.AllowPatterns = []string{"*.json", "*.safetensors"} + third := second + third.Source.AllowPatterns = []string{"*.safetensors", "*.json"} + secondKey, err := modelartifacts.CacheKey(second) + Expect(err).NotTo(HaveOccurred()) + thirdKey, err := modelartifacts.CacheKey(third) + Expect(err).NotTo(HaveOccurred()) + Expect(first).To(MatchRegexp(`^[0-9a-f]{64}$`)) + Expect(secondKey).To(Equal(thirdKey)) + Expect(secondKey).NotTo(Equal(first)) + }) + + It("places all state below the models artifact root", func() { + spec := baseSpec() + key, err := modelartifacts.CacheKey(spec) + Expect(err).NotTo(HaveOccurred()) + spec.Resolved.CacheKey = key + layout, err := modelartifacts.LayoutFor("/models", spec) + Expect(err).NotTo(HaveOccurred()) + Expect(layout.Final).To(Equal(filepath.Join("/models", ".artifacts", "huggingface", key))) + Expect(layout.Snapshot).To(Equal(filepath.Join(layout.Final, "snapshot"))) + relative, err := modelartifacts.RelativeSnapshotPath(key) + Expect(err).NotTo(HaveOccurred()) + Expect(relative).To(Equal(filepath.Join(".artifacts", "huggingface", key, "snapshot"))) + _, err = modelartifacts.RelativeSnapshotPath("sha256:" + key) + Expect(err).To(MatchError(ContainSubstring("invalid artifact cache key"))) + }) + + DescribeTable("rejects hostile Hub paths", + func(candidate string) { Expect(modelartifacts.ValidateRelativeHubPath(candidate)).NotTo(Succeed()) }, + Entry("absolute", "/etc/passwd"), + Entry("parent", "nested/../escape"), + Entry("backslash", `nested\escape`), + Entry("NUL", "bad\x00path"), + Entry("empty component", "nested//file"), + ) + + It("reads a complete versioned manifest", func() { + spec := baseSpec() + key, err := modelartifacts.CacheKey(spec) + Expect(err).NotTo(HaveOccurred()) + spec.Resolved.CacheKey = key + encoded, err := json.Marshal(modelartifacts.Manifest{ + Version: modelartifacts.ManifestVersion, + Artifact: spec, + Files: []modelartifacts.ManifestFile{{ + Path: "nested/model.safetensors", + Size: 11, + SHA256: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + }}, + }) + Expect(err).NotTo(HaveOccurred()) + fileName := filepath.Join(GinkgoT().TempDir(), "manifest.json") + Expect(os.WriteFile(fileName, encoded, 0o600)).To(Succeed()) + + manifest, err := modelartifacts.ReadManifest(fileName) + Expect(err).NotTo(HaveOccurred()) + Expect(manifest.Artifact.Resolved.CacheKey).To(Equal(key)) + Expect(manifest.Files).To(HaveLen(1)) + }) + + It("reports progress only through the request-scoped sink", func() { + event := modelartifacts.ProgressEvent{ + Phase: modelartifacts.PhaseDownloading, + Artifact: "model", + File: "weights.safetensors", + CurrentBytes: 7, + TotalBytes: 11, + } + var received []modelartifacts.ProgressEvent + ctx := modelartifacts.WithProgressSink(context.Background(), func(update modelartifacts.ProgressEvent) { + received = append(received, update) + }) + modelartifacts.ReportProgress(ctx, event) + modelartifacts.ReportProgress(context.Background(), event) + Expect(received).To(Equal([]modelartifacts.ProgressEvent{event})) + }) +}) diff --git a/pkg/modelartifacts/progress.go b/pkg/modelartifacts/progress.go new file mode 100644 index 000000000..b0dbf36b6 --- /dev/null +++ b/pkg/modelartifacts/progress.go @@ -0,0 +1,40 @@ +package modelartifacts + +import "context" + +type Phase string + +const ( + PhaseResolving Phase = "resolving" + PhaseDownloading Phase = "downloading" + PhaseVerifying Phase = "verifying" + PhaseCommitting Phase = "committing" + PhasePersisting Phase = "persisting" +) + +type ProgressEvent struct { + Phase Phase + Artifact string + File string + CurrentBytes int64 + TotalBytes int64 + CompletedFiles int + TotalFiles int +} + +type ProgressSink func(ProgressEvent) + +type progressSinkKey struct{} + +func WithProgressSink(ctx context.Context, sink ProgressSink) context.Context { + if sink == nil { + return ctx + } + return context.WithValue(ctx, progressSinkKey{}, sink) +} + +func ReportProgress(ctx context.Context, event ProgressEvent) { + if sink, ok := ctx.Value(progressSinkKey{}).(ProgressSink); ok && sink != nil { + sink(event) + } +} diff --git a/pkg/modelartifacts/reference.go b/pkg/modelartifacts/reference.go new file mode 100644 index 000000000..d5121db9e --- /dev/null +++ b/pkg/modelartifacts/reference.go @@ -0,0 +1,82 @@ +package modelartifacts + +import ( + "fmt" + "strings" +) + +var hfReferencePrefixes = []string{ + "https://huggingface.co/", + "huggingface://", + "hf://", + "hf.co/", +} + +// ParsePrimaryReference converts a Hugging Face repository or file reference +// into a managed artifact spec. It accepts repo roots like "owner/repo" and +// direct file references like "huggingface://owner/repo/path/to/model.gguf". +// The boolean return is false when the reference is not Hugging Face-shaped. +func ParsePrimaryReference(raw string) (Spec, bool, error) { + source, ok, err := ParsePrimarySource(raw) + if err != nil || !ok { + return Spec{}, ok, err + } + return Spec{ + Name: TargetModel, + Target: TargetModel, + Source: source, + }, true, nil +} + +// ParsePrimarySource converts a Hugging Face repository or file reference into +// a managed source definition. Direct file references are translated into a +// repo plus a single allow pattern so the materializer downloads only the +// selected file. +func ParsePrimarySource(raw string) (Source, bool, error) { + ref := strings.TrimSpace(raw) + if ref == "" { + return Source{}, false, nil + } + + stripped := ref + for _, prefix := range hfReferencePrefixes { + stripped = strings.TrimPrefix(stripped, prefix) + } + stripped = strings.TrimSuffix(stripped, "/") + + parts := strings.Split(stripped, "/") + if len(parts) < 2 { + return Source{}, false, nil + } + hadPrefix := ref != stripped + + repo, err := normalizeRepo(strings.Join(parts[:2], "/")) + if err != nil { + if ref != stripped { + return Source{}, false, fmt.Errorf("invalid Hugging Face reference %q: %w", raw, err) + } + return Source{}, false, nil + } + if !hadPrefix && len(parts) == 2 { + return Source{}, false, nil + } + + source := Source{ + Type: SourceTypeHuggingFace, + Repo: repo, + } + if len(parts) == 2 { + return source, true, nil + } + + if parts[2] == "resolve" { + if len(parts) < 5 { + return Source{}, false, fmt.Errorf("invalid Hugging Face file reference %q", raw) + } + source.AllowPatterns = []string{strings.Join(parts[4:], "/")} + return source, true, nil + } + + source.AllowPatterns = []string{strings.Join(parts[2:], "/")} + return source, true, nil +} diff --git a/pkg/modelartifacts/types.go b/pkg/modelartifacts/types.go new file mode 100644 index 000000000..774a55606 --- /dev/null +++ b/pkg/modelartifacts/types.go @@ -0,0 +1,160 @@ +package modelartifacts + +import ( + "fmt" + "net/url" + "path" + "regexp" + "slices" + "strings" +) + +const ( + SourceTypeHuggingFace = "huggingface" + TargetModel = "model" + HuggingFaceTokenEnv = "HF_TOKEN" +) + +var ( + commitPattern = regexp.MustCompile(`^[0-9a-f]{40}$`) + cacheKeyPattern = regexp.MustCompile(`^[0-9a-f]{64}$`) +) + +type Spec struct { + Name string `yaml:"name" json:"name"` + Target string `yaml:"target" json:"target"` + Source Source `yaml:"source" json:"source"` + Resolved *Resolved `yaml:"resolved,omitempty" json:"resolved,omitempty"` +} + +type Source struct { + Type string `yaml:"type" json:"type"` + Repo string `yaml:"repo" json:"repo"` + Revision string `yaml:"revision,omitempty" json:"revision,omitempty"` + TokenEnv string `yaml:"token_env,omitempty" json:"token_env,omitempty"` + AllowPatterns []string `yaml:"allow_patterns,omitempty" json:"allow_patterns,omitempty"` + IgnorePatterns []string `yaml:"ignore_patterns,omitempty" json:"ignore_patterns,omitempty"` +} + +type Resolved struct { + Endpoint string `yaml:"endpoint" json:"endpoint"` + Revision string `yaml:"revision" json:"revision"` + CacheKey string `yaml:"cache_key" json:"cache_key"` +} + +func (s Spec) Normalize() (Spec, error) { + if strings.TrimSpace(s.Name) == "" { + s.Name = TargetModel + } + if strings.TrimSpace(s.Target) == "" { + s.Target = TargetModel + } + s.Name = strings.TrimSpace(s.Name) + s.Target = strings.TrimSpace(s.Target) + s.Source.Type = strings.TrimSpace(s.Source.Type) + s.Source.TokenEnv = strings.TrimSpace(s.Source.TokenEnv) + + if s.Name != TargetModel || s.Target != TargetModel { + return Spec{}, fmt.Errorf("only artifact name/target %q is supported", TargetModel) + } + if s.Source.Type != SourceTypeHuggingFace { + return Spec{}, fmt.Errorf("unsupported artifact source type %q", s.Source.Type) + } + + repo, err := normalizeRepo(s.Source.Repo) + if err != nil { + return Spec{}, err + } + s.Source.Repo = repo + if strings.TrimSpace(s.Source.Revision) == "" { + s.Source.Revision = "main" + } else { + s.Source.Revision = strings.TrimSpace(s.Source.Revision) + } + if strings.ContainsRune(s.Source.Revision, '\x00') { + return Spec{}, fmt.Errorf("revision contains NUL") + } + if s.Source.TokenEnv != "" && s.Source.TokenEnv != HuggingFaceTokenEnv { + return Spec{}, fmt.Errorf("token_env must be empty or %s", HuggingFaceTokenEnv) + } + + for _, patterns := range [][]string{s.Source.AllowPatterns, s.Source.IgnorePatterns} { + for _, pattern := range patterns { + if err := validatePattern(pattern); err != nil { + return Spec{}, err + } + } + } + s.Source.AllowPatterns = slices.Clone(s.Source.AllowPatterns) + s.Source.IgnorePatterns = slices.Clone(s.Source.IgnorePatterns) + slices.Sort(s.Source.AllowPatterns) + slices.Sort(s.Source.IgnorePatterns) + + if s.Resolved != nil { + resolved := *s.Resolved + resolved.Endpoint, err = normalizeEndpoint(resolved.Endpoint) + if err != nil { + return Spec{}, err + } + resolved.Revision = strings.ToLower(strings.TrimSpace(resolved.Revision)) + if !commitPattern.MatchString(resolved.Revision) { + return Spec{}, fmt.Errorf("resolved revision must be 40 lowercase hexadecimal characters") + } + if resolved.CacheKey != "" && !cacheKeyPattern.MatchString(resolved.CacheKey) { + return Spec{}, fmt.Errorf("resolved cache key must be 64 lowercase hexadecimal characters") + } + s.Resolved = &resolved + } + + return s, nil +} + +func (s Spec) Validate() error { + normalized, err := s.Normalize() + if err != nil { + return err + } + if normalized.Resolved != nil && normalized.Resolved.CacheKey == "" { + return fmt.Errorf("resolved cache key is required in installed state") + } + return nil +} + +func normalizeRepo(raw string) (string, error) { + repo := strings.TrimSpace(raw) + for _, prefix := range []string{"https://huggingface.co/", "huggingface://", "hf://"} { + repo = strings.TrimPrefix(repo, prefix) + } + repo = strings.TrimSuffix(repo, "/") + parts := strings.Split(repo, "/") + if len(parts) != 2 || parts[0] == "" || parts[1] == "" || strings.ContainsAny(repo, "?#\\\x00") { + return "", fmt.Errorf("Hugging Face repo must be exactly owner/repo") + } + return repo, nil +} + +// CanonicalRepo normalizes a Hugging Face repository reference to owner/repo. +func CanonicalRepo(raw string) (string, error) { + return normalizeRepo(raw) +} + +func validatePattern(pattern string) error { + if pattern == "" || path.IsAbs(pattern) || strings.ContainsAny(pattern, "\\\x00") { + return fmt.Errorf("invalid artifact pattern %q", pattern) + } + for _, part := range strings.Split(pattern, "/") { + if part == ".." { + return fmt.Errorf("invalid artifact pattern %q", pattern) + } + } + return nil +} + +func normalizeEndpoint(raw string) (string, error) { + u, err := url.Parse(strings.TrimSpace(raw)) + if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" { + return "", fmt.Errorf("invalid resolved Hugging Face endpoint") + } + u.Path = strings.TrimRight(u.Path, "/") + return u.String(), nil +} diff --git a/pkg/modelartifacts/types_test.go b/pkg/modelartifacts/types_test.go new file mode 100644 index 000000000..e61fef317 --- /dev/null +++ b/pkg/modelartifacts/types_test.go @@ -0,0 +1,70 @@ +package modelartifacts_test + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/modelartifacts" +) + +var _ = Describe("artifact configuration", func() { + It("normalizes the supported primary Hugging Face source", func() { + spec, err := (modelartifacts.Spec{ + Source: modelartifacts.Source{ + Type: "huggingface", + Repo: "hf://Qwen/Qwen3-ASR-1.7B", + TokenEnv: "HF_TOKEN", + }, + }).Normalize() + Expect(err).NotTo(HaveOccurred()) + Expect(spec.Name).To(Equal("model")) + Expect(spec.Target).To(Equal("model")) + Expect(spec.Source.Repo).To(Equal("Qwen/Qwen3-ASR-1.7B")) + Expect(spec.Source.Revision).To(Equal("main")) + }) + + It("parses Hugging Face repo and file references into managed sources", func() { + repo, ok, err := modelartifacts.ParsePrimarySource("huggingface://Qwen/Qwen3-ASR-1.7B") + Expect(err).NotTo(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(repo.Type).To(Equal(modelartifacts.SourceTypeHuggingFace)) + Expect(repo.Repo).To(Equal("Qwen/Qwen3-ASR-1.7B")) + Expect(repo.AllowPatterns).To(BeEmpty()) + + file, ok, err := modelartifacts.ParsePrimarySource("https://huggingface.co/nomic-ai/nomic-embed-text-v1.5-GGUF/resolve/main/nomic-embed-text-v1.5.f16.gguf") + Expect(err).NotTo(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(file.Repo).To(Equal("nomic-ai/nomic-embed-text-v1.5-GGUF")) + Expect(file.AllowPatterns).To(Equal([]string{"nomic-embed-text-v1.5.f16.gguf"})) + }) + + It("ignores non-Hugging Face references", func() { + _, ok, err := modelartifacts.ParsePrimarySource("models/local-model.gguf") + Expect(err).NotTo(HaveOccurred()) + Expect(ok).To(BeFalse()) + }) + + DescribeTable("rejects unsafe or unsupported declarations", + func(spec modelartifacts.Spec, message string) { + Expect(spec.Validate()).To(MatchError(ContainSubstring(message))) + }, + Entry("unknown source", modelartifacts.Spec{Source: modelartifacts.Source{Type: "s3", Repo: "owner/repo"}}, "source type"), + Entry("secondary target", modelartifacts.Spec{Name: "controlnet", Target: "controlnet", Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}}, "target"), + Entry("malformed repo", modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo/file"}}, "owner/repo"), + Entry("unrelated secret", modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", TokenEnv: "AWS_SECRET_ACCESS_KEY"}}, "HF_TOKEN"), + Entry("parent filter", modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", AllowPatterns: []string{"../*.json"}}}, "pattern"), + Entry("prefixed cache key", modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}, Resolved: &modelartifacts.Resolved{Endpoint: "https://huggingface.co", Revision: "0123456789abcdef0123456789abcdef01234567", CacheKey: "sha256:bad"}}, "cache key"), + ) + + It("validates installed state", func() { + spec := modelartifacts.Spec{ + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}, + Resolved: &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + }, + } + Expect(spec.Validate()).To(Succeed()) + }) +})