Use fondant dev image if fondant dev version is installed

scripts/build_base_image.sh

@@ -4,46 +4,30 @@ set -e
 function usage {
   echo "Usage: $0 [options]"
   echo "Options:"
-  echo "  -t, --tag <value>        Tag to add to image, repeatable"
+  echo "  -t, --tag <value>        Set the tag (default: latest)"
   echo "  -h, --help               Display this help message"
 }
 
 # Parse the arguments
 while [[ "$#" -gt 0 ]]; do case $1 in
-  -t |--tag) tags+=("$2"); shift;;
+  -t|--tag) tags+=("$2"); shift;;
   -h|--help) usage; exit;;
   *) echo "Unknown parameter passed: $1"; exit 1;;
 esac; shift; done
 
-# Check for required argument
-if [ -z "${tags}" ]; then
-  echo "Error: tag parameter is required"
-  usage
-  exit 1
-fi
-
 # Supported Python versions
 python_versions=("3.8" "3.9" "3.10" "3.11")
 
-
-for python_version in "${python_versions[@]}"; do
+for tag in "${tags[@]}"; do
+  for python_version in "${python_versions[@]}"; do
     BASENAME=fondant
-
-    image_tags=()
-    echo "Tagging image with following tags:"
-    for tag in "${tags[@]}"; do
-      image_tags+=("${tag}-py${python_version}")
-    done
-
+    IMAGE_TAG=${tag}-py${python_version}
     full_image_names=()
 
     # create repo if not exists
     aws ecr-public describe-repositories --region us-east-1 --repository-names ${BASENAME} || aws ecr-public create-repository --region us-east-1 --repository-name ${BASENAME}
-
-    for image_tag in "${image_tags[@]}"; do
-      full_image_names+=("public.ecr.aws/fndnt/${BASENAME}:${image_tag}")
-      full_image_names+=("fndnt/${BASENAME}:${image_tag}")
-    done
+    full_image_names+=("public.ecr.aws/fndnt/${BASENAME}:${IMAGE_TAG}")
+    full_image_names+=("fndnt/${BASENAME}:${IMAGE_TAG}")
 
     # Add argument for each tag
     for image_name in "${full_image_names[@]}" ; do
@@ -57,7 +41,8 @@ for python_version in "${python_versions[@]}"; do
     # Build docker images and push to docker hub
     docker build --push "${args[@]}" \
     --build-arg="PYTHON_VERSION=${python_version}" \
-    --build-arg="FONDANT_VERSION=${tag}" \
+    --build-arg="FONDANT_VERSION=${tags[0]}" \
     -f "images/Dockerfile" \
     .
+  done
 done

src/fondant/pipeline/lightweight_component.py

@@ -6,7 +6,7 @@
 import typing as t
 from dataclasses import asdict, dataclass
 from functools import wraps
-from importlib.metadata import version
+from importlib import metadata
 
 from fondant.component import BaseComponent, Component
 
@@ -52,7 +52,9 @@ def resolve_fndnt_base_image():
         else:
             python_version = f"{MAX_PYTHON_VERSION[0]}.{MAX_PYTHON_VERSION[1]}"
 
-        fondant_version = version("fondant")
+        fondant_version = metadata.version("fondant")
+        if fondant_version == "0.1.dev0":
+            fondant_version = "dev"
         basename = "fndnt/fondant"
         return f"{basename}:{fondant_version}-py{python_version}"
 

tests/pipeline/test_pipeline.py

@@ -1,7 +1,5 @@
 """Fondant pipelines test."""
 import copy
-import sys
-from importlib.metadata import version
 from pathlib import Path
 
 import dask.dataframe as dd
@@ -13,7 +11,13 @@
 from fondant.core.component_spec import ComponentSpec
 from fondant.core.exceptions import InvalidPipelineDefinition
 from fondant.core.schema import Field, Type
-from fondant.pipeline import ComponentOp, Pipeline, Resources, lightweight_component
+from fondant.pipeline import (
+    ComponentOp,
+    Image,
+    Pipeline,
+    Resources,
+    lightweight_component,
+)
 
 valid_pipeline_path = Path(__file__).parent / "examples/pipelines/valid_pipeline"
 invalid_pipeline_path = Path(__file__).parent / "examples/pipelines/invalid_pipeline"
@@ -85,16 +89,10 @@ def load(self) -> dd.DataFrame:
             )
             return dd.from_pandas(df, npartitions=1)
 
-    basename = "fndnt/fondant"
-    fondant_version = version("fondant")
-    python_version = sys.version_info
-    python_version = f"{python_version.major}.{python_version.minor}"
-    fondant_image_name = f"{basename}:{fondant_version}-py{python_version}"
-
     component = ComponentOp.from_ref(Foo, produces={"bar": pa.string()})
     assert component.component_spec._specification == {
         "name": "foo",
-        "image": fondant_image_name,
+        "image": Image.resolve_fndnt_base_image(),
         "description": "lightweight component",
         "consumes": {"additionalProperties": True},
         "produces": {"additionalProperties": True},

tests/pipeline/test_python_component.py

@@ -2,27 +2,19 @@
 import re
 import sys
 import textwrap
-from importlib.metadata import version
+from dataclasses import dataclass
+from unittest import mock
 
 import dask.dataframe as dd
 import pandas as pd
 import pyarrow as pa
 import pytest
 from fondant.component import DaskLoadComponent, PandasTransformComponent
 from fondant.core.exceptions import InvalidLightweightComponent
-from fondant.pipeline import Pipeline, lightweight_component
+from fondant.pipeline import Image, Pipeline, lightweight_component
 from fondant.pipeline.compiler import DockerCompiler
 
 
-@pytest.fixture()
-def default_fondant_image():
-    basename = "fndnt/fondant"
-    fondant_version = version("fondant")
-    python_version = sys.version_info
-    python_version = f"{python_version.major}.{python_version.minor}"
-    return f"{basename}:{fondant_version}-py{python_version}"
-
-
 def test_build_python_script():
     @lightweight_component()
     class CreateData(DaskLoadComponent):
@@ -62,7 +54,7 @@ def load(self) -> dd.DataFrame:
     )
 
 
-def test_lightweight_component_sdk(default_fondant_image, caplog):
+def test_lightweight_component_sdk(caplog):
     pipeline = Pipeline(
         name="dummy-pipeline",
         base_path="./data",
@@ -139,7 +131,7 @@ def transform(self, dataframe: pd.DataFrame) -> pd.DataFrame:
     assert operation_spec_dict == {
         "specification": {
             "name": "addn",
-            "image": default_fondant_image,
+            "image": Image.resolve_fndnt_base_image(),
             "description": "lightweight component",
             "consumes": {"additionalProperties": True},
             "produces": {"additionalProperties": True},
@@ -268,7 +260,7 @@ def load(self) -> int:
         CreateData(produces={}, consumes={})
 
 
-def test_lightweight_component_decorator_without_parentheses(default_fondant_image):
+def test_lightweight_component_decorator_without_parentheses():
     @lightweight_component
     class CreateData(DaskLoadComponent):
         def load(self) -> dd.DataFrame:
@@ -290,11 +282,49 @@ def load(self) -> dd.DataFrame:
     assert operation_spec_without_image == {
         "specification": {
             "name": "createdata",
-            "image": default_fondant_image,
+            "image": Image.resolve_fndnt_base_image(),
             "description": "lightweight component",
             "consumes": {"additionalProperties": True},
             "produces": {"additionalProperties": True},
         },
         "consumes": {},
         "produces": {},
     }
+
+
+@dataclass
+class MockSysVersionInfo:
+    major: int
+    minor: int
+    micro: int
+
+    def __lt__(self, other):
+        return (self.major, self.minor, self.micro) <= other
+
+    def __ge__(self, other):
+        return other < (self.major, self.minor, self.micro)
+
+
+def test_fndnt_base_image_resolution():
+    # Base image version is set to python version
+    with mock.patch.object(sys, "version_info", MockSysVersionInfo(3, 10, 0)):
+        base_image_name = Image.resolve_fndnt_base_image()
+        assert base_image_name == "fndnt/fondant:dev-py3.10"
+
+    # Local python version is not supported
+    with mock.patch.object(sys, "version_info", MockSysVersionInfo(3, 12, 0)):
+        base_image_name = Image.resolve_fndnt_base_image()
+        assert base_image_name == "fndnt/fondant:dev-py3.11"
+
+    with mock.patch.object(sys, "version_info", MockSysVersionInfo(3, 7, 0)):
+        base_image_name = Image.resolve_fndnt_base_image()
+        assert base_image_name == "fndnt/fondant:dev-py3.11"
+
+    with mock.patch.object(
+        sys,
+        "version_info",
+        MockSysVersionInfo(3, 9, 0),
+    ), mock.patch("importlib.metadata.version") as mock_call:
+        mock_call.return_value = "0.9"
+        base_image_name = Image.resolve_fndnt_base_image()
+        assert base_image_name == "fndnt/fondant:0.9-py3.9"