wehub-resource-sync
79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
import logging
|
|
import os
|
|
import subprocess
|
|
from functools import lru_cache
|
|
|
|
import docker
|
|
import pytest
|
|
import requests
|
|
from packaging.version import Version
|
|
|
|
import mlflow
|
|
|
|
TEST_IMAGE_NAME = "test_image"
|
|
MLFLOW_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
|
|
RESOURCE_DIR = os.path.join(MLFLOW_ROOT, "tests", "resources", "dockerfile")
|
|
|
|
docker_client = docker.from_env()
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clean_up_docker():
|
|
yield
|
|
|
|
# Get all containers using the test image
|
|
containers = docker_client.containers.list(filters={"ancestor": TEST_IMAGE_NAME})
|
|
for container in containers:
|
|
container.remove(force=True)
|
|
|
|
# Clean up the image
|
|
try:
|
|
docker_client.images.remove(TEST_IMAGE_NAME, force=True)
|
|
except docker.errors.ImageNotFound:
|
|
pass
|
|
|
|
# Clean up the build cache and volumes
|
|
try:
|
|
subprocess.check_call(["docker", "builder", "prune", "-a", "-f"])
|
|
except subprocess.CalledProcessError as e:
|
|
_logger.warning("Failed to clean up docker system: %s", e)
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_released_mlflow_version():
|
|
url = "https://pypi.org/pypi/mlflow/json"
|
|
response = requests.get(url)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
versions = [
|
|
v for v in map(Version, data["releases"]) if not (v.is_devrelease or v.is_prerelease)
|
|
]
|
|
return str(max(versions))
|
|
|
|
|
|
def save_model_with_latest_mlflow_version(flavor, extra_pip_requirements=None, **kwargs):
|
|
"""
|
|
Save a model with overriding MLflow version from dev version to the latest released version.
|
|
By default a model is saved with the dev version of MLflow, which is not available on PyPI.
|
|
Usually we can be workaround this by adding --serve-wheel flag that starts local PyPI server,
|
|
however, this doesn't work when installing dependencies inside Docker container. Hence, this
|
|
function uses `extra_pip_requirements` to save the model with the latest released MLflow.
|
|
"""
|
|
latest_mlflow_version = get_released_mlflow_version()
|
|
if flavor == "langchain":
|
|
kwargs["pip_requirements"] = [
|
|
f"mlflow[gateway]=={latest_mlflow_version}",
|
|
"langchain<1.1.0",
|
|
]
|
|
else:
|
|
extra_pip_requirements = extra_pip_requirements or []
|
|
extra_pip_requirements.append(f"mlflow=={latest_mlflow_version}")
|
|
if flavor == "lightgbm":
|
|
# Adding pyarrow < 18 to prevent pip installation resolution conflicts.
|
|
extra_pip_requirements.append("pyarrow<18")
|
|
kwargs["extra_pip_requirements"] = extra_pip_requirements
|
|
flavor_module = getattr(mlflow, flavor)
|
|
flavor_module.save_model(**kwargs)
|