testcontainers / testcontainers/testcontainers-python

New Container: MFflow

Open
#649 1 comment 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

🚀 enhancement
Dominant language
Python
Stars
2.3k
Forks
386
Avg merge
4h 40m
Merged PRs (30d)
1

Description

What is the new container you'd like to have?

I would like to provide a new testcontainer for mlflow, a tool for managing your machine learning model life cycle.

Why not just use a generic container for this?

I added handy utilites to get the url to track to and to get the client to interact with the container directly.

The implementation would look like:

import logging

import requests
from mlflow import MlflowClient
from testcontainers.core.container import DockerContainer
from testcontainers.core.waiting_utils import wait_container_is_ready

logger = logging.getLogger(__name__)


class MFlowContainer(DockerContainer):
    """Test container for MLflow.

    Args:
        image: the image to use. Change if you need different version.
        port: the internal port to use.  The exposed port is assigned automatically.
        cmd: the command to run. Defaults to "mlflow server". If you want to use the ui for debugging and testing use "mlflow ui".
    """

    def __init__(
        self, image: str = "ghcr.io/mlflow/mlflow:v2.14.1", port: int = 5000, cmd: str = "mlflow server"
    ) -> None:
        super().__init__(image=image)
        self.port = port
        self.with_exposed_ports(self.port)
        self.cmd = cmd

    def _configure(self) -> None:
        self.with_env("MLFLOW_PORT", str(self.port))
        self.with_env("MLFLOW_HOST", "0.0.0.0")
        self.with_command(self.cmd)

    def get_url(self) -> str:
        """Returns the url of the container.

        Returns:
            The url. Use to track to.
        """
        return f"http://{self.get_container_host_ip()}:{self.get_exposed_port(self.port)}"

    @wait_container_is_ready(requests.exceptions.ConnectionError, requests.exceptions.ReadTimeout)
    def _readiness_probe(self) -> None:
        # https://mlflow.org/docs/latest/deployment/deploy-model-locally.html?highlight=health
        response = requests.get(f"{self.get_url()}/health", timeout=1)
        response.raise_for_status()

    def get_client(self) -> MlflowClient:
        """Returns the MlflowClient of the container.

        Can be used for testing.

        """
        return MlflowClient(self.get_url())

    def start(self) -> "MFlowContainer":
        self._configure()
        super().start()
        self._readiness_probe()
        return self

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

No repository files or tests are named. Start by locating existing Testcontainers Python container integrations and their tests, then use the proposed MLflow container entry points as a guide; done means the MLflow service becomes ready and its URL and client can be used in tests.

Written by the indexing model from the issue text.

Assessment

Tech stack
docker, python
Domain
machine-learning, testing
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.