qwer

Author: jstourac

tests/containers/base_image_test.py

@@ -189,6 +189,69 @@ def test_pip_install_cowsay_runs(self, image: str):
         finally:
             docker_utils.NotebookContainer(container).stop(timeout=0)
 
+    # There are two ways how the image is being updated
+    # 1. A change to the image is being done (e.g. package update, Dockerfile update etc.). In this case, we need to check the size of the build image that contains these updates.
+    # 2. A change is done to the params.env file where we update manifest references to a new image. In this case we shall check that the updated reference matches our size constraints
+    size_treshold: int = 100 # in MBs
+    percent_treshold: int = 10
+    def test_image_size_change(self, image: str):
+        """Checks the image size didn't change extensively - treshold is 10% or XXX MB."""
+
+        # Map of image label names with expected size in MBs.
+        expected_image_name_size_map = {
+            "odh-notebook-base-centos-stream9-python-3.11": 1,
+            "odh-notebook-base-ubi9-python-3.11": 1,
+            "odh-notebook-cuda-c9s-python-3.11": 1,
+            "odh-notebook-cuda-ubi9-python-3.11": 1,
+            "odh-notebook-jupyter-datascience-ubi9-python-3.11": 1,
+            "odh-notebook-jupyter-datascience-ubi9-python-3.9": 2681,  # TODO
+            "odh-notebook-jupyter-minimal-ubi9-python-3.11": 1,
+            "odh-notebook-jupyter-pytorch-ubi9-python-3.11": 1,
+            "odh-notebook-cuda-jupyter-tensorflow-ubi9-python-3.11": 1,
+            "odh-notebook-jupyter-trustyai-ubi9-python-3.11": 1,
+            "odh-notebook-jupyter-rocm-pytorch-ubi9-python-3.11": 1,
+            "odh-notebook-jupyter-rocm-tensorflow-ubi9-python-3.11": 1,
+            "odh-notebook-rstudio-server-c9s-python-3.11": 1,
+            "odh-notebook-runtime-datascience-ubi9-python-3.11": 1,
+            "odh-notebook-runtime-minimal-ubi9-python-3.11": 1,
+            "odh-notebook-runtime-pytorch-ubi9-python-3.11": 1,
+            "odh-notebook-cuda-runtime-tensorflow-ubi9-python-3.11": 1,
+            "odh-notebook-runtime-rocm-pytorch-ubi9-python-3.11": 1,
+            "odh-notebook-rocm-runtime-tensorflow-ubi9-python-3.11": 1,
+            "odh-notebook-code-server-ubi9-python-3.11": 1,
+            "odh-notebook-rocm-python-3.11": 1,
+        }
+
+        client = testcontainers.core.container.DockerClient()
+
+        image_metadata = client.client.images.get(image)
+        actual_img_size = image_metadata.attrs["Size"]
+        actual_img_size = round(actual_img_size / 1024 / 1024)
+        logging.info(f"The size of the image is {actual_img_size} MBs.")
+
+        # print(image_metadata)
+
+        import docker
+        try:
+            image_metadata = client.client.images.get(image)
+        except docker.errors.ImageNotFound:
+            image_metadata = client.client.images.pull(image)
+            assert isinstance(image_metadata, docker.models.images.Image)
+
+        img_label_name = image_metadata.labels["name"]
+        if img_label_name in expected_image_name_size_map:
+            expected_img_size = expected_image_name_size_map[img_label_name]
+            logging.debug(f"Expected size of the image is {expected_img_size} MBs.")
+        else:
+            pytest.fail(f"Image name label '{img_label_name}' is not in the expected image size map {expected_image_name_size_map}")
+
+        # Check the size change constraints now
+        # 1. Percentual size change
+        percent_change = actual_img_size / expected_img_size * 100 - 100
+        assert percent_change < self.percent_treshold, f"Image size changed by {percent_change}% (expected: {expected_img_size} MB; actual: {actual_img_size} MB; treshold: {self.percent_treshold}%)."
+        # 2. Absolute size change
+        size_difference = actual_img_size - expected_img_size
+        assert size_difference < self.size_treshold, f"Image size changed by {size_difference} MB (expected: {expected_img_size} MB; actual: {actual_img_size} MB; treshold: {self.size_treshold} MB)."
 
 def encode_python_function_execution_command_interpreter(python: str, function: Callable[..., Any], *args: list[Any]) -> list[str]:
     """Returns a cli command that will run the given Python function encoded inline.