From d7ec514c2f191b27bbb185de00909da9dd61de27 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mika=C3=ABl=20Simard?= Date: Fri, 2 Oct 2026 14:03:50 +0100 Subject: [PATCH 1/4] Clean per-run image built by mlflow Updated the run() method in a try/finally statement that cleans the per-run image, with appropriate logging. --- mlops/Experiment.py | 53 +++++++++++++++++++++++++++++++++++++++------ 1 file changed, 46 insertions(+), 7 deletions(-) diff --git a/mlops/Experiment.py b/mlops/Experiment.py index ddff906..59b38f6 100644 --- a/mlops/Experiment.py +++ b/mlops/Experiment.py @@ -11,6 +11,7 @@ import mlflow from git import Repo from minio import Minio +from mlflow.utils.git_utils import get_git_commit from mlops import LOG_FILE from mlops.ProjectFile import ProjectFile @@ -320,12 +321,50 @@ def run(self, **kwargs): logger.debug(f'Artifact URI: {mlflow.get_artifact_uri()}') logger.debug(f'Project URI: {self.project_path}') - mlflow.run(uri=self.project_path, - experiment_id=self.experiment_id, - env_manager='local', - build_image=True, - # revisit this in the future, mlflow 2.0 changed this behaviour, can we be more efficient? - docker_args=docker_args_default - ) + try: + mlflow.run(uri=self.project_path, + experiment_id=self.experiment_id, + env_manager='local', + build_image=True, + # revisit this in the future, mlflow 2.0 changed this behaviour, can we be more efficient? + docker_args=docker_args_default + ) + finally: + self.remove_run_image(client) mlflow.log_artifact(LOG_FILE) + + def remove_run_image(self, client): + """ + Removes the per-run image that mlflow builds on top of the project image (:latest). + + mlflow tags this per-run image :. Without cleanup it is + left behind after every run. If a later run reuses the same tag, the tag moves to the new image and + the old one is left with no name, showing as :. + + As such, cleanup is useful so you don't end up with a lot of untagged images. + + :param client: docker client + :return: + """ + + git_commit = get_git_commit(self.project_path) + + if not git_commit: + # without a commit mlflow tags its image :latest, i.e. the project image itself, so nothing to do + logger.debug('No git commit found, skipping removal of mlflow run image') + return + + run_image = f'{self.experiment_name}:{git_commit[:7]}' + + try: + client.images.remove(run_image) + logger.info(f'Removed mlflow run image: {run_image}') + except docker.errors.ImageNotFound: + logger.warning(f'mlflow run image {run_image} not found, nothing removed. This is expected if the run ' + f'failed before mlflow built its image, in which case the run\'s own error explains why. ' + f'Otherwise if mlflow has been upgraded, it may have modified how it names run images:' + f' check `docker images {self.experiment_name}`, remove any leftover image manually, ' + f'and report this as an MLOps issue.') + except docker.errors.APIError as e: + logger.warning(f'Could not remove mlflow run image {run_image}: {e}') From 8c3b8c094a2be4152458c789b44942e64a24aa3a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mika=C3=ABl=20Simard?= Date: Fri, 2 Oct 2026 14:08:31 +0100 Subject: [PATCH 2/4] Unit tests for per-run image removal test_remove_run_image proves the method removes the right name and never touches the project image. test_run -> now also checks if the try/finally statement calls the cleanup and that mlflow really builds and named its image the way we expect it. --- tests/mlops/test_Experiment.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/mlops/test_Experiment.py b/tests/mlops/test_Experiment.py index 5020898..3586209 100644 --- a/tests/mlops/test_Experiment.py +++ b/tests/mlops/test_Experiment.py @@ -3,6 +3,7 @@ import os import docker +from mlflow.utils.git_utils import get_git_commit from mlops.Experiment import Experiment @@ -68,6 +69,20 @@ def test_build_experiment_image_subprocess(self): images_list = [img['RepoTags'][0] for img in client.api.images() if img['RepoTags']] assert self.experiment.experiment_name + ':latest' in images_list + def test_remove_run_image(self): + client = docker.from_env() + self.experiment.build_experiment_image_subprocess(context_path='tests/data', + dockerfile_path=self.experiment.project_path + '/Dockerfile') + # tag the project image the same way mlflow tags its per-run image + run_tag = get_git_commit(self.experiment.project_path)[:7] + client.images.get(self.experiment.experiment_name + ':latest').tag(self.experiment.experiment_name, run_tag) + + self.experiment.remove_run_image(client) + + tags = [tag for img in client.images.list(name=self.experiment.experiment_name) for tag in img.tags] + assert f'{self.experiment.experiment_name}:{run_tag}' not in tags + assert self.experiment.experiment_name + ':latest' in tags + def test_run(self, capsys): """ this test will fail locally, setup git in test data dir is done in github actions @@ -79,3 +94,8 @@ def test_run(self, capsys): self.experiment.run() captured = capsys.readouterr() assert 'succeeded' in captured.err + + # only the project image should remain, the mlflow per-run image is removed + client = docker.from_env() + tags = [tag for img in client.images.list(name=self.experiment.experiment_name) for tag in img.tags] + assert tags == [self.experiment.experiment_name + ':latest'] From d562adf73de3a103bf9419e1f78b9273a4713363 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mika=C3=ABl=20Simard?= Date: Fri, 2 Oct 2026 14:33:43 +0100 Subject: [PATCH 3/4] Remove previous docker image if rebuiliding If rebuilding docker image, then remove the previous one as it will be left here as : after the rebuild. That's an extra case that could leave old images around. --- mlops/Experiment.py | 40 ++++++++++++++++++++++++++++++++++++++-- 1 file changed, 38 insertions(+), 2 deletions(-) diff --git a/mlops/Experiment.py b/mlops/Experiment.py index 59b38f6..03690f6 100644 --- a/mlops/Experiment.py +++ b/mlops/Experiment.py @@ -311,10 +311,18 @@ def run(self, **kwargs): logger.info('Checking for existing image') client = docker.from_env() - images = [str(img['RepoTags']) for img in client.api.images()] - if all([(self.experiment_name + ':latest') not in item for item in images]) or rebuild_docker: + try: + old_image = client.images.get(self.experiment_name + ':latest') + except docker.errors.ImageNotFound: + old_image = None + + if old_image is None or rebuild_docker: logger.info(f'No existing image found, rebuild flag {rebuild_docker}') self.build_experiment_image_subprocess(context_path=self.project_path) + + # if rebuilding docker image, then remove the previous one as it will be left here as : after the rebuild. + if old_image is not None: + self.remove_replaced_image(client, old_image) else: logger.info(f'Found existing project image: {self.experiment_name}:latest') @@ -334,6 +342,34 @@ def run(self, **kwargs): mlflow.log_artifact(LOG_FILE) + def remove_replaced_image(self, client, old_image): + """ + Removes the previous project image after a rebuild (rebuild_docker) has replaced it. + + Rebuilding :latest moves the tag to the new image, and the old one is left with no name, + showing as :. The old image is only removed if it has no name left: if the build failed or + produced the same image (cached) it is still :latest, and if it has another tag it is kept. + + :param client: docker client + :param old_image: the :latest image from before the rebuild + :return: + """ + + try: + old_image.reload() + except docker.errors.ImageNotFound: + return + + if old_image.tags: + logger.debug(f'Previous project image {old_image.short_id} still tagged {old_image.tags}, keeping it') + return + + try: + client.images.remove(old_image.id) + logger.info(f'Removed previous project image: {old_image.short_id}') + except docker.errors.APIError as e: + logger.warning(f'Could not remove previous project image {old_image.short_id}: {e}') + def remove_run_image(self, client): """ Removes the per-run image that mlflow builds on top of the project image (:latest). From cf17dffd868698e1d7bb68a4af886e8f89c169cf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mika=C3=ABl=20Simard?= Date: Fri, 2 Oct 2026 14:44:07 +0100 Subject: [PATCH 4/4] Unit test for remove image after docker rebuild --- tests/mlops/test_Experiment.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/mlops/test_Experiment.py b/tests/mlops/test_Experiment.py index 3586209..a1f4271 100644 --- a/tests/mlops/test_Experiment.py +++ b/tests/mlops/test_Experiment.py @@ -1,8 +1,10 @@ import configparser +import io import logging import os import docker +import pytest from mlflow.utils.git_utils import get_git_commit from mlops.Experiment import Experiment @@ -83,6 +85,25 @@ def test_remove_run_image(self): assert f'{self.experiment.experiment_name}:{run_tag}' not in tags assert self.experiment.experiment_name + ':latest' in tags + def test_remove_replaced_image(self): + client = docker.from_env() + # separate image name so test_project:latest (used by test_run) is not replaced. + # client.images.build(fileobj...) is the python version of docker build with the Dockerfile written inline. + name = 'mlops_test_replaced_image' + old_image, _ = client.images.build(fileobj=io.BytesIO(b'FROM scratch\nLABEL build=old\n'), tag=name, rm=True) + # rebuild that produces a different image, moving the tag and leaving old_image with no name + new_image, _ = client.images.build(fileobj=io.BytesIO(b'FROM scratch\nLABEL build=new\n'), tag=name, rm=True) + + self.experiment.remove_replaced_image(client, old_image) + with pytest.raises(docker.errors.ImageNotFound): + client.images.get(old_image.id) + + # an image that still has its tag (cached or failed rebuild) is kept + self.experiment.remove_replaced_image(client, new_image) + assert client.images.get(name + ':latest').id == new_image.id + + client.images.remove(name + ':latest') + def test_run(self, capsys): """ this test will fail locally, setup git in test data dir is done in github actions