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