seai2526-uniba-TheClouds/Code-Comment-Classification-Api

1"""Synchronise champion MLflow models from the remote registry to the local filesystem."""2 3import logging4import os5from pathlib import Path6import shutil7 8import mlflow9from mlflow.tracking import MlflowClient10 11logger = logging.getLogger(__name__)12LANGUAGES = ("python", "java", "pharo")13 14 15def _get_mlflow_client() -> MlflowClient:16    """Return an MLflow client configured from environment variables.17 18    If ``MLFLOW_TRACKING_URI`` is defined, it is passed to19    :func:`mlflow.set_tracking_uri`. Authentication (for example on DagsHub)20    is handled by MLflow itself via the standard environment variables21    ``MLFLOW_TRACKING_USERNAME`` and ``MLFLOW_TRACKING_PASSWORD``.22    """23    tracking_uri = os.getenv("MLFLOW_TRACKING_URI")24    if tracking_uri:25        mlflow.set_tracking_uri(tracking_uri)26    return MlflowClient()27 28 29def _find_champion_version_for_language(30    client: MlflowClient,31    lang: str,32):33    """Return the champion model version for the given language, if any.34 35    The function searches all registered models and looks for models whose name36    starts with ``"<lang>-"`` (for example ``"python-transformer"``). For each37    matching model it tries to resolve the alias ``"<lang>-champion"`` using38    :meth:`MlflowClient.get_model_version_by_alias`.39 40    Args:41        client: Initialised MLflow client.42        lang: Language identifier, such as ``"python"``, ``"java"`` or43            ``"pharo"``.44 45    Returns:46        The matching :class:`mlflow.entities.model_registry.ModelVersion` if a47        champion is found, otherwise ``None``.48 49    """50    alias_name = f"{lang}-champion"51    prefix = f"{lang}-"52 53    # Get all registered models and filter by language prefix.54    for rm in client.search_registered_models():55        model_name = rm.name56        if not model_name.startswith(prefix):57            continue58 59        try:60            mv = client.get_model_version_by_alias(61                name=model_name,62                alias=alias_name,63            )64            logger.info(65                "Found champion model for %s: %s (version %s)",66                lang,67                model_name,68                mv.version,69            )70            return mv71        except Exception:  # noqa: BLE00172            logger.info("Alias not defined for model %s, trying next one.", model_name)73            continue74 75    logger.warning("No champion model found for %s.", lang)76    return None77 78 79def sync_best_models_to_disk(80    models_root: str | Path = "models",81    api_subdir: str = "api",82) -> None:83    """Download champion models from MLflow and write them to disk.84 85    For each language in :data:`LANGUAGES`, this function looks up the model86    version with alias ``"<lang>-champion"`` and downloads its artifacts. After87    download, the directory structure is normalised so that the final layout is:88 89    .. code-block:: text90 91        models/92          <api_subdir>/93            python/94              <model_type>/95                ...96            java/97              <model_type>/98                ...99            pharo/100              <model_type>/101                ...102 103    For transformer models logged via ``mlflow.transformers``, the inner104    ``model/`` directory is flattened so that the Hugging Face files105    (``config.json``, ``model.safetensors``, ``tokenizer.json``, and so on)106    live directly under ``<model_type>/``.107 108    Args:109        models_root: Base directory under which models are written. Can be a110            string or :class:`pathlib.Path`. Defaults to ``"models"``.111        api_subdir: Optional subdirectory appended under ``models_root`` (for112            example ``"api"``). If empty, models are stored directly under113            ``models_root``.114 115    Raises:116        OSError: If creating directories, moving files, or removing directories117            fails at the OS level.118 119    """120    client = _get_mlflow_client()121 122    root = Path(models_root)123    if api_subdir:124        root = root / api_subdir125    root.mkdir(parents=True, exist_ok=True)126    logger.info("Syncing best models to: %s", root.resolve())127 128    for lang in LANGUAGES:129        mv = _find_champion_version_for_language(client, lang)130        if mv is None:131            continue132 133        model_name = mv.name134        try:135            lang_from_name, model_type = model_name.split("-", 1)136        except ValueError:137            logger.error("Unexpected model name format: %s", model_name)138            continue139 140        if lang_from_name != lang:141            logger.warning(142                "Language mismatch for model %s: expected %s, got %s",143                model_name,144                lang,145                lang_from_name,146            )147 148        dest_dir = root / lang / model_type149        if dest_dir.exists():150            shutil.rmtree(dest_dir)151        dest_dir.mkdir(parents=True, exist_ok=True)152 153        logger.info(154            "Downloading model '%s' version %s to %s...",155            model_name,156            mv.version,157            dest_dir.resolve(),158        )159 160        try:161            # Download the artifact (for example ".../java_transformer_model").162            downloaded_path = Path(163                mlflow.artifacts.download_artifacts(164                    artifact_uri=mv.source,165                    dst_path=str(dest_dir),166                ),167            )168 169            # For transformer models logged with mlflow.transformers, artifacts170            # are stored under an inner "model/" directory.171            model_subdir = downloaded_path / "model"172            if model_subdir.is_dir():173                # Move the contents of "model" directly into dest_dir.174                for item in model_subdir.iterdir():175                    shutil.move(str(item), dest_dir / item.name)176 177                # Remove the wrapper directory (with MLmodel, conda.yaml, etc.).178                if downloaded_path != dest_dir:179                    shutil.rmtree(downloaded_path)180 181        except Exception as e:182            logger.error(183                "Failed to download/reshape model '%s' version %s: %s",184                model_name,185                mv.version,186                e,187            )188 189 190if __name__ == "__main__":191    logging.basicConfig(level=logging.INFO)192    sync_best_models_to_disk()193