voxkit.storage.models

  1"""Specialized CRUD operations for managing models within the VoxKit storage system.
  2
  3Directory Structure
  4-------------------
  5Each model follows a hierarchical structure:
  6
  7    chosen_engine/
  8    ├── train/
  9    │   ├── model_id_1/
 10    │   │   ├── entrypoint.model      # Model file
 11    │   │   ├── data/                 # Data directory
 12    │   │   ├── eval/                 # Evaluation directory
 13    │   │   ├── train/                # Training directory
 14    │   │   └── voxkit_model.json     # Model metadata
 15    │   ├── model_id_2/
 16    │   │   └── ...
 17    │   └── ...
 18
 19API
 20---
 21- **create_model**: Create a new model entry in storage
 22- **get_model_metadata**: Retrieve metadata for a specific model
 23- **update_model_metadata**: Update the status or details of an existing model
 24- **list_models**: List all models for a given engine
 25- **delete_model**: Remove a model from storage
 26- **import_models**: Import models from an external directory
 27
 28Notes
 29-----
 30- All paths are managed using pathlib
 31- Engine-specific branching may be necessary to bridge different model formats
 32- Error handling only exposes user-friendly messages
 33- Model IDs are generated using unique timestamps with microsecond precision
 34"""
 35
 36import json
 37import logging
 38import shutil
 39from pathlib import Path
 40from typing import Literal, Tuple, TypedDict
 41
 42from voxkit.storage.utils import generate_unique_id, get_storage_root, readable_from_unique_id
 43
 44from .constants import MODELS_ROOT
 45
 46logger = logging.getLogger(__name__)
 47
 48
 49class ModelMetadata(TypedDict):
 50    """Model metadata structure.
 51
 52    Attributes:
 53        name: Human-readable name of the model.
 54        engine_id: Identifier of the engine this model belongs to (e.g., "mfa").
 55        model_path: Path to the model entrypoint file.
 56        data_path: Path to the model's data directory.
 57        eval_path: Path to the model's evaluation directory.
 58        train_path: Path to the model's training artifacts directory.
 59        download_date: Human-readable download/creation timestamp.
 60        id: Unique identifier (timestamp with microsecond precision).
 61    """
 62
 63    name: str
 64    engine_id: str
 65    model_path: Path
 66    data_path: Path
 67    eval_path: Path
 68    train_path: Path
 69    download_date: str
 70    id: str
 71
 72
 73def _get_model_root(engine_id: str, model_id: str) -> Path | None:
 74    """Get the root directory for a specific model.
 75
 76    Args:
 77        engine_id: Identifier of the engine the model belongs to
 78        model_id: Identifier of the model
 79
 80    Returns:
 81        Path to the model root directory or None if not found
 82    """
 83    model_root = Path(f"{get_storage_root()}/{engine_id}/{MODELS_ROOT}/{model_id}")
 84    if model_root.exists():
 85        return model_root
 86    return None
 87
 88
 89def _get_models_root(engine_id: str) -> Path | None:
 90    """Get the root directory for storing models for a given engine.
 91
 92    Args:
 93        engine_id: Identifier of the engine
 94
 95    Returns:
 96        Path to the models root directory or None if not found
 97    """
 98    models_root = Path(f"{get_storage_root()}/{engine_id}/{MODELS_ROOT}")
 99    if models_root.exists():
100        return models_root
101    return None
102
103
104def create_model(
105    engine_id: str, model_name: str, source_path: Path | str | None = None
106) -> tuple[Literal[True], ModelMetadata] | tuple[Literal[False], str]:
107    """Create a new model entry in the storage.
108
109    Creates a new model directory structure with subdirectories for data, evaluation,
110    and training artifacts. Generates a unique ID and creates a metadata file.
111    Optionally copies model files from a source path.
112
113    Args:
114        engine_id: Identifier of the engine the model belongs to
115        model_name: Human-readable name for the model
116        source_path: Optional path to source model files (.zip, .model, or directory)
117
118    Returns:
119        Tuple of (True, ModelMetadata) on success or (False, error_message) on failure
120
121    Raises:
122        FileNotFoundError: If the engine's model root does not exist
123        Exception: If directory creation or metadata writing fails
124
125    Notes:
126        - Model paths in metadata are stored as strings for JSON serialization
127        - Automatically cleans up partially created directories on failure
128        - Creates four directories: model root, data, eval, and train
129        - If source_path is a .zip file, copies as entrypoint.zip
130        - If source_path is a .model file or directory, copies via copytree
131    """
132
133    engine_models_root = Path(f"{get_storage_root()}/{engine_id}/{MODELS_ROOT}")
134    if not engine_models_root.exists():
135        return False, f"Unsupported engine_id: {engine_id}"
136
137    now = generate_unique_id()
138    model_root = Path(f"{engine_models_root}/{now}")
139    logger.debug("Creating model at: %s", model_root)
140
141    try:
142        model_path = model_root / "entrypoint.model"
143        data_path = model_root / "data"
144        eval_path = model_root / "eval"
145        train_path = model_root / "train"
146
147        humandate = readable_from_unique_id(now)
148        model_metadata = ModelMetadata(
149            name=model_name or f"Model_{now}",
150            engine_id=engine_id,
151            model_path=model_path.with_suffix(".model"),
152            data_path=data_path,
153            eval_path=eval_path,
154            train_path=train_path,
155            download_date=humandate,
156            id=now,
157        )
158
159        # Create model directories
160        model_path.mkdir(parents=True, exist_ok=False)
161        data_path.mkdir(parents=True, exist_ok=False)
162        eval_path.mkdir(parents=True, exist_ok=False)
163        train_path.mkdir(parents=True, exist_ok=False)
164        metadata_path = model_root / "voxkit_model.json"
165
166        # Copy source files if provided
167        if source_path is not None:
168            source_path = Path(source_path)
169            if not source_path.exists():
170                raise FileNotFoundError(f"Source path does not exist: {source_path}")
171
172            if str(source_path).endswith(".zip"):
173                # Copy zip file as entrypoint.zip
174                dest_file = model_root / "entrypoint.zip"
175                shutil.copy2(source_path, dest_file)
176                model_metadata["model_path"] = dest_file
177            else:
178                # Copy directory or .model file
179                shutil.copytree(source_path, model_path, dirs_exist_ok=True)
180
181        # Convert Path objects to strings for JSON serialization
182        json_metadata = {k: str(v) if isinstance(v, Path) else v for k, v in model_metadata.items()}
183
184        # Create metadata file and write metadata
185        with open(metadata_path, "w", encoding="utf-8") as f:
186            json.dump(json_metadata, f, indent=4)
187
188        return True, model_metadata
189
190    except Exception as e:
191        logger.exception("Exception occurred during model creation")
192        # Clean up partially created model directory
193        if model_root and model_root.exists():
194            shutil.rmtree(model_root)
195
196        return False, f"Failed to create model: {e}"
197
198
199def update_model_metadata(engine_id: str, model_id: str, updates: dict) -> Tuple[bool, str]:
200    """Update metadata for an existing model.
201
202    Updates fields in the model's metadata file. Only fields present in the
203    metadata are updated; unknown fields are ignored. Values are converted to
204    strings before writing.
205
206    Args:
207        engine_id: Identifier of the engine the model belongs to
208        model_id: Identifier of the model to update
209        updates: Dictionary of fields to update in the model metadata
210
211    Returns:
212        Tuple of (True, success_message) on success or (False, error_message) on failure
213
214    Raises:
215        FileNotFoundError: If the model is not found
216        Exception: If metadata file cannot be read or written
217    """
218    model_root = _get_model_root(engine_id, model_id)
219    if not model_root:
220        return False, f"Model '{model_id}' for engine '{engine_id}' not found"
221
222    metadata_path = Path(model_root) / "voxkit_model.json"
223    try:
224        with open(metadata_path, "r", encoding="utf-8") as f:
225            metadata = json.load(f)
226
227        # Update fields
228        for key, value in updates.items():
229            if key in metadata:
230                metadata[key] = str(value)
231
232        with open(metadata_path, "w", encoding="utf-8") as f:
233            json.dump(metadata, f, indent=4)
234
235        return True, "Model metadata updated successfully."
236
237    except Exception:
238        logger.exception("Exception occurred during model metadata update")
239        return False, "Failed to update model metadata."
240
241
242def list_models(engine_id: str) -> list[ModelMetadata]:
243    """List available models for the given engine.
244
245    Scans the engine's model directory and collects metadata from all subdirectories
246    containing valid voxkit_model.json files. Creates the models directory if it
247    doesn't exist.
248
249    Args:
250        engine_id: Identifier of the engine to list models for
251
252    Returns:
253        List of ModelMetadata dictionaries (empty list if none found)
254
255    Notes:
256        - Skips directories with invalid or missing metadata files
257        - Returns empty list on error
258        - Automatically creates models directory if missing
259    """
260    try:
261        models_root = Path(f"{get_storage_root()}/{engine_id}/{MODELS_ROOT}")
262        if not models_root.exists():
263            models_root.mkdir(parents=True, exist_ok=True)
264            return []
265
266        models_found = []
267        for dir in models_root.iterdir():
268            if dir.is_dir():
269                metadata_path = dir / "voxkit_model.json"
270                if metadata_path.exists():
271                    try:
272                        with open(metadata_path, "r", encoding="utf-8") as f:
273                            metadata = json.load(f)
274                            models_found.append(metadata)
275                    except json.JSONDecodeError as e:
276                        logger.warning("Skipping invalid JSON in %s: %s", metadata_path, e)
277                        continue
278        return models_found
279
280    except Exception:
281        logger.exception("Error listing models")
282        return []
283
284
285def get_model_metadata(engine_id: str, model_id: str) -> ModelMetadata:
286    """Get metadata for a specific model by its ID.
287
288    Retrieves the model metadata from the voxkit_model.json file in the
289    model's directory.
290
291    Args:
292        engine_id: Identifier of the engine the model belongs to
293        model_id: Identifier of the model
294
295    Returns:
296        ModelMetadata dictionary
297
298    Raises:
299        FileNotFoundError: If the model or metadata file is not found
300        JSONDecodeError: If the metadata file is malformed
301    """
302    model_root = _get_model_root(engine_id, model_id)
303    if not model_root:
304        raise FileNotFoundError(f"Model '{model_id}' for engine '{engine_id}' not found")
305    metadata_path = Path(model_root) / "voxkit_model.json"
306    if not metadata_path.exists():
307        raise FileNotFoundError(f"Metadata file not found for model '{model_id}'")
308    with open(metadata_path, "r", encoding="utf-8") as f:
309        metadata: ModelMetadata = json.load(f)
310        return metadata
311
312
313def download_and_copy_huggingface_model(
314    model_path: str,
315    destination: str,
316) -> str | None:
317    """
318    Download model from HuggingFace and copy actual files to destination.
319    Follows symlinks to get real model files (like git clone behavior).
320
321    Args:
322        model_path: HuggingFace model path (e.g., 'pkadambi/Wav2TextGrid')
323        destination: Where to copy the model files
324
325    Returns:
326        Destination path if successful, None if failed
327    """
328    try:
329        from huggingface_hub import snapshot_download
330        from huggingface_hub.utils import HfHubHTTPError, RepositoryNotFoundError
331
332        # Validate model path format
333        if not model_path or "/" not in model_path:
334            logger.error("Invalid model path format: %s", model_path)
335            return None
336
337        # Download to HF cache (returns path to snapshot with symlinks)
338        cache_snapshot_path = snapshot_download(
339            repo_id=model_path,
340        )
341
342        logger.debug("Downloaded to cache: %s", cache_snapshot_path)
343
344        # Create destination directory
345        dest_path = Path(destination).expanduser()
346        dest_path.mkdir(parents=True, exist_ok=True)
347
348        # Copy all files, following symlinks (like git clone)
349        cache_path = Path(cache_snapshot_path)
350        for item in cache_path.iterdir():
351            if item.name.startswith("."):
352                # Skip .gitattributes and other hidden files if desired
353                continue
354
355            if item.is_symlink() or item.is_file():
356                # Resolve symlink to get actual file, then copy
357                actual_file = item.resolve()
358                dest_file = dest_path / item.name
359                shutil.copy2(actual_file, dest_file)
360                logger.debug("Copied: %s", item.name)
361            elif item.is_dir():
362                # Recursively copy directories
363                shutil.copytree(
364                    item,
365                    dest_path / item.name,
366                    symlinks=False,  # Follow symlinks
367                    dirs_exist_ok=True,
368                )
369
370        logger.info("Successfully copied model to: %s", dest_path)
371        return str(dest_path)
372
373    except RepositoryNotFoundError:
374        logger.error("Model not found: %s", model_path)
375        return None
376
377    except HfHubHTTPError:
378        logger.exception("HTTP error downloading model %s", model_path)
379        return None
380
381    except Exception:
382        logger.exception("Error downloading model %s", model_path)
383        return None
384
385
386def delete_model(engine_id: str, model_id: str) -> Tuple[bool, str]:
387    """Delete a model given its engine ID and model ID.
388
389    Permanently removes the model directory and all its contents, including
390    the model file, metadata, and associated data/eval/train directories.
391
392    Args:
393        engine_id: Identifier of the engine the model belongs to
394        model_id: Identifier of the model to delete
395
396    Returns:
397        Tuple of (True, success_message) on success or (False, error_message) on failure
398
399    Raises:
400        Exception: If the directory cannot be removed
401
402    Notes:
403        - This operation is irreversible
404        - Removes the entire model directory tree
405        - Validates that engine_id and model_id are not empty before proceeding
406    """
407
408    if not engine_id or not model_id:
409        return False, "Engine ID and Model ID cannot be empty."
410
411    logger.debug("Attempting to delete model: engine_id=%s, model_id=%s", engine_id, model_id)
412    model_path = _get_model_root(engine_id, model_id)
413
414    if not model_path:
415        return False, f"Model {model_id} not found"
416
417    logger.debug("Deleting model at path: %s", model_path)
418    shutil.rmtree(model_path)
419    return True, "Model deleted successfully."
420
421
422def import_models(engine_id, new_models_root: Path) -> Tuple[bool, str]:
423    """Import models into the storage system.
424
425    This function imports models from an external directory into VoxKit storage. Each model
426    must have a valid voxkit_model.json metadata file. The function validates metadata,
427    generates new IDs, updates paths, and copies the model to storage.
428
429    Args:
430        engine_id: Identifier of the engine
431        new_models_root: Source directory containing models to import
432
433    Returns:
434        Tuple of (True, success_message) on success or (False, error_message) on failure
435
436    Raises:
437        Exception: If model validation or copy operations fail
438    """
439    try:
440        for source_model_path in new_models_root.iterdir():
441            if source_model_path.is_dir():
442                try:
443                    # Check for voxkit_model.json file
444                    metadata_path = Path(source_model_path / "voxkit_model.json")
445                    if not metadata_path.exists():
446                        return False, f"{source_model_path.name} (missing metadata file)"
447
448                    metadata = None
449                    # Read json metadata
450
451                    with open(metadata_path, "r", encoding="utf-8") as f:
452                        metadata = json.load(f)
453
454                    if metadata is None:
455                        return False, f"{source_model_path.name} (invalid metadata file)"
456
457                    engine_models_root = get_storage_root() / engine_id
458                    if not engine_models_root.exists():
459                        engine_models_root.mkdir(parents=True, exist_ok=False)
460
461                    model_id = generate_unique_id()
462
463                    if engine_id != metadata["engine_id"]:
464                        return False, f"{source_model_path.name} (engine_id mismatch)"
465
466                    entrypoint_name = Path(metadata["model_path"]).name
467                    if not entrypoint_name:
468                        return False, f"{source_model_path.name} (invalid model_path in metadata)"
469
470                    dest_model_entrypoint = (
471                        engine_models_root / MODELS_ROOT / model_id / entrypoint_name
472                    )
473                    new_metadata = ModelMetadata(
474                        name=metadata["name"],
475                        engine_id=metadata["engine_id"],
476                        model_path=Path(dest_model_entrypoint),
477                        data_path=Path(engine_models_root / MODELS_ROOT / model_id / "data"),
478                        eval_path=Path(engine_models_root / MODELS_ROOT / model_id / "eval"),
479                        train_path=Path(engine_models_root / MODELS_ROOT / model_id / "train"),
480                        download_date=readable_from_unique_id(model_id),
481                        id=model_id,
482                    )
483
484                    # Copy model directory to storage
485                    dest_path = engine_models_root / MODELS_ROOT / model_id
486
487                    shutil.copytree(source_model_path, dest_path, dirs_exist_ok=True)
488
489                    # Convert Path objects to strings for JSON serialization
490                    json_metadata = {
491                        k: str(v) if isinstance(v, Path) else v for k, v in new_metadata.items()
492                    }
493
494                    # Overwrite metadata file with new IDs and paths
495                    new_metadata_path = dest_path / "voxkit_model.json"
496
497                    with open(new_metadata_path, "w", encoding="utf-8") as f:
498                        json.dump(json_metadata, f, indent=4)
499
500                except Exception as e:
501                    return False, f"{source_model_path.name} (error: {str(e)})"
502
503        return True, f"Models imported successfully from: {new_models_root}"
504
505    except Exception as e:
506        return False, f"Failed to import model: {str(e)}"