190 cleanup extra docker images created by mlflow

mlops/Experiment.py

@@ -9,6 +9,7 @@
 import docker
 import mlflow
 from git import Repo
+from mlflow.utils.git_utils import get_git_commit
 
 from mlops import LOG_FILE
 from mlops.ProjectFile import ProjectFile
@@ -267,22 +268,96 @@ 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 <none>:<none> 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')
 
         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_replaced_image(self, client, old_image):
+        """
+        Removes the previous project image after a rebuild (rebuild_docker) has replaced it.
+
+        Rebuilding <experiment_name>:latest moves the tag to the new image, and the old one is left with no name,
+        showing as <none>:<none>. 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 <experiment_name>:latest, and if it has another tag it is kept.
+
+        :param client: docker client
+        :param old_image: the <experiment_name>: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 (<experiment_name>:latest).
+
+        mlflow tags this per-run image <experiment_name>:<first 7 chars of git commit>. 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 <none>:<none>.
+
+        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 <experiment_name>: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}')

tests/mlops/test_Experiment.py

@@ -1,8 +1,11 @@
 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
 
@@ -64,6 +67,39 @@ 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_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
@@ -75,3 +111,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']