control-service: ability to set builder image per python version

projects/control-service/projects/helm_charts/pipelines-control-service/values.yaml

@@ -279,7 +279,8 @@ deploymentDataJobBaseImage:
 
 ## The map of python versions for data job deployments supported by the Control Service.
 ## All values should be strings, and follow the format
-## {"pythonVersion": {"baseImage": "<repository>:<tag>", "vdkImage": "<registry>/<repository>:<tag>"}}
+## {"pythonVersion": {"baseImage": "<repository>:<tag>", "vdkImage": "<registry>/<repository>:<tag>",
+## "builderImage": "<registry>/<repository>:<tag>"}}
 ##
 ## The value of baseImage will be used to create the image, where data job would be run.
 ## On top of it the job source and its dependencies are installed for each job.
@@ -292,11 +293,16 @@ deploymentDataJobBaseImage:
 ## Only the installed python modules (vdk and its dependencies) will be used from the image.
 ## Everything else is effectively discarded since another image is used as base during execution.
 ##
+## The value of builderImage is the image name of Job Builder, which will be used for data jobs building. \
+## The presence of this parameter is OPTIONAL; if it is not provided, the Control Service will automatically default
+## to the deploymentBuilderImage.
+##
 ## Example:
 ## deploymentSupportedPythonVersions:
 ##   3.7:
 ##     baseImage: "registry.hub.docker.com/versatiledatakit/data-job-base-python-3.7:latest"
 ##     vdkImage: "registry.hub.docker.com/versatiledatakit/quickstart-vdk:release"
+##     builderImage: "registry.hub.docker.com/versatiledatakit/job-builder:latest"
 deploymentSupportedPythonVersions:
   3.9:
     baseImage: "registry.hub.docker.com/versatiledatakit/data-job-base-python-3.9:latest"

projects/control-service/projects/pipelines_control_service/src/integration-test/java/com/vmware/taurus/datajobs/it/TestJobImageBuilderDynamicVdkImageIT.java

@@ -46,6 +46,12 @@
 @TestPropertySource(
     properties = {
       "datajobs.control.k8s.k8sSupportsV1CronJob=true",
+      "datajobs.deployment.supportedPythonVersions={3.8: {"
+          + "vdkImage: 'registry.hub.docker.com/versatiledatakit/quickstart-vdk:pre-release', "
+          + "baseImage: 'versatiledatakit/data-job-base-python-3.8:latest', "
+          + "builderImage: 'ghcr.io/versatile-data-kit-dev/versatiledatakit/job-builder:1.3.3'}}",
+      "datajobs.deployment.defaultPythonVersion=3.8",
+      "datajobs.builder.image="
     })
 @SpringBootTest(
     webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,

projects/control-service/projects/pipelines_control_service/src/main/java/com/vmware/taurus/service/deploy/JobImageBuilder.java

@@ -185,14 +185,15 @@ public boolean buildImage(
             registryPassword,
             builderAwsSessionToken);
     var envs = getBuildParameters(dataJob, jobDeployment);
+    String builderImage = supportedPythonVersions.getBuilderImage(jobDeployment.getPythonVersion());
 
     log.info(
         "Creating builder job {} for data job version {}",
         builderJobName,
         jobDeployment.getGitCommitSha());
     controlKubernetesService.createJob(
         builderJobName,
-        dockerRegistryService.builderImage(),
+        builderImage,
         false,
         false,
         envs,

projects/control-service/projects/pipelines_control_service/src/main/java/com/vmware/taurus/service/deploy/SupportedPythonVersions.java

@@ -23,24 +23,30 @@
 @Slf4j
 public class SupportedPythonVersions {
 
-  private static String BASE_IMAGE = "baseImage";
+  protected static String BASE_IMAGE = "baseImage";
 
-  private static String VDK_IMAGE = "vdkImage";
+  protected static String VDK_IMAGE = "vdkImage";
+
+  protected static String BUILDER_IMAGE = "builderImage";
 
   @Value("#{${datajobs.deployment.supportedPythonVersions:{}}}")
   private Map<String, Map<String, String>> supportedPythonVersions;
 
   @Value("${datajobs.deployment.defaultPythonVersion}")
   private String defaultPythonVersion;
 
+  private final DockerRegistryService dockerRegistryService;
+
   /**
    * Check if the pythonVersion passed by the user is supported by the Control Service.
    *
    * @param pythonVersion python version passed by the user.
    * @return true if the version is supported, and false otherwise.
    */
   public boolean isPythonVersionSupported(String pythonVersion) {
-    return !supportedPythonVersions.isEmpty() && supportedPythonVersions.containsKey(pythonVersion);
+    return supportedPythonVersions != null
+        && !supportedPythonVersions.isEmpty()
+        && supportedPythonVersions.containsKey(pythonVersion);
   }
 
   /**
@@ -66,7 +72,7 @@ public Set<String> getSupportedPythonVersions() {
    * @return a string of the data job base image.
    */
   public String getJobBaseImage(String pythonVersion) {
-    if (supportedPythonVersions != null && isPythonVersionSupported(pythonVersion)) {
+    if (isPythonVersionSupported(pythonVersion)) {
       return supportedPythonVersions.get(pythonVersion).get(BASE_IMAGE);
     } else {
       log.warn(
@@ -100,6 +106,17 @@ public String getVdkImage(String pythonVersion) {
     }
   }
 
+  public String getBuilderImage(String pythonVersion) {
+    if (isPythonVersionSupported(pythonVersion)) {
+      return getBuilderImage(supportedPythonVersions.get(pythonVersion));
+    } else {
+      log.warn(
+          "An issue with the passed pythonVersion or supportedPythonVersions configuration has"
+              + " occurred. Returning default builder image");
+      return getBuilderImage(supportedPythonVersions.get(defaultPythonVersion));
+    }
+  }
+
   public String getDefaultVdkImage() {
     return supportedPythonVersions.get(defaultPythonVersion).get(VDK_IMAGE);
   }
@@ -113,4 +130,8 @@ public String getDefaultVdkImage() {
   public String getDefaultPythonVersion() {
     return defaultPythonVersion;
   }
+
+  private String getBuilderImage(Map<String, String> supportedPythonVersion) {
+    return supportedPythonVersion.getOrDefault(BUILDER_IMAGE, dockerRegistryService.builderImage());
+  }
 }

projects/control-service/projects/pipelines_control_service/src/test/java/com/vmware/taurus/service/deploy/JobImageBuilderTest.java

@@ -86,12 +86,12 @@ public void setUp() {
 
   @Test
   public void buildImage_notExist_success() throws InterruptedException, ApiException, IOException {
-    when(dockerRegistryService.builderImage()).thenReturn(TEST_BUILDER_IMAGE_NAME);
     when(kubernetesService.listJobs()).thenReturn(Collections.emptySet());
     var builderJobResult =
         new KubernetesService.JobStatusCondition(true, "type", "test-reason", "test-message", 0);
     when(kubernetesService.watchJob(any(), anyInt(), any())).thenReturn(builderJobResult);
     when(supportedPythonVersions.getJobBaseImage(any())).thenReturn("python:3.7-slim");
+    when(supportedPythonVersions.getBuilderImage(any())).thenReturn(TEST_BUILDER_IMAGE_NAME);
 
     JobDeployment jobDeployment = new JobDeployment();
     jobDeployment.setDataJobName(TEST_JOB_NAME);
@@ -129,13 +129,13 @@ public void buildImage_builderRunning_oldBuilderDeleted()
       throws InterruptedException, ApiException, IOException {
     when(dockerRegistryService.dataJobImageExists(eq(TEST_IMAGE_NAME), Mockito.any()))
         .thenReturn(false);
-    when(dockerRegistryService.builderImage()).thenReturn(TEST_BUILDER_IMAGE_NAME);
     when(kubernetesService.listJobs())
         .thenReturn(Set.of(TEST_BUILDER_IMAGE_NAME), Collections.emptySet());
     var builderJobResult =
         new KubernetesService.JobStatusCondition(true, "type", "test-reason", "test-message", 0);
     when(kubernetesService.watchJob(any(), anyInt(), any())).thenReturn(builderJobResult);
     when(supportedPythonVersions.getJobBaseImage(any())).thenReturn("python:3.7-slim");
+    when(supportedPythonVersions.getBuilderImage(any())).thenReturn(TEST_BUILDER_IMAGE_NAME);
 
     JobDeployment jobDeployment = new JobDeployment();
     jobDeployment.setDataJobName(TEST_JOB_NAME);
@@ -207,13 +207,13 @@ public void buildImage_imageExists_buildSkipped()
   @Test
   public void buildImage_jobFailed_failure()
       throws InterruptedException, ApiException, IOException {
-    when(dockerRegistryService.builderImage()).thenReturn(TEST_BUILDER_IMAGE_NAME);
     when(kubernetesService.listJobs()).thenReturn(Collections.emptySet());
     var builderJobResult =
         new KubernetesService.JobStatusCondition(false, "type", "test-reason", "test-message", 0);
     when(kubernetesService.watchJob(any(), anyInt(), any())).thenReturn(builderJobResult);
     when(kubernetesService.getPodLogs(TEST_BUILDER_JOB_NAME)).thenReturn(TEST_BUILDER_LOGS);
     when(supportedPythonVersions.getJobBaseImage(any())).thenReturn("python:3.7-slim");
+    when(supportedPythonVersions.getBuilderImage(any())).thenReturn(TEST_BUILDER_IMAGE_NAME);
 
     JobDeployment jobDeployment = new JobDeployment();
     jobDeployment.setDataJobName(TEST_JOB_NAME);
@@ -258,13 +258,12 @@ public void buildImage_jobFailed_failure()
   public void
       buildImage_deploymentDataJobBaseImageNullAndSupportedPythonVersions_shouldCreateCronjobUsingSupportedPythonVersions()
           throws InterruptedException, ApiException, IOException {
-
-    when(dockerRegistryService.builderImage()).thenReturn(TEST_BUILDER_IMAGE_NAME);
     when(kubernetesService.listJobs()).thenReturn(Collections.emptySet());
     var builderJobResult =
         new KubernetesService.JobStatusCondition(true, "type", "test-reason", "test-message", 0);
     when(kubernetesService.watchJob(any(), anyInt(), any())).thenReturn(builderJobResult);
     when(supportedPythonVersions.getJobBaseImage("3.11")).thenReturn("test-base-image");
+    when(supportedPythonVersions.getBuilderImage(any())).thenReturn(TEST_BUILDER_IMAGE_NAME);
 
     JobDeployment jobDeployment = new JobDeployment();
     jobDeployment.setDataJobName(TEST_JOB_NAME);

projects/control-service/projects/pipelines_control_service/src/test/java/com/vmware/taurus/service/deploy/SupportedPythonVersionsTest.java

@@ -9,21 +9,26 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.InjectMocks;
+import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.springframework.test.util.ReflectionTestUtils;
 
 import java.util.Map;
 import java.util.Set;
 
+import static com.vmware.taurus.service.deploy.SupportedPythonVersions.*;
+import static org.mockito.Mockito.when;
+
 @ExtendWith(MockitoExtension.class)
 public class SupportedPythonVersionsTest {
   private static final String SUPPORTED_PYTHON_VERSIONS = "supportedPythonVersions";
-  private static final String BASE_IMAGE = "baseImage";
-  private static final String VDK_IMAGE = "vdkImage";
+
   private static final String DEFAULT_PYTHON_VERSION = "defaultPythonVersion";
 
   @InjectMocks private SupportedPythonVersions supportedPythonVersions;
 
+  @Mock private DockerRegistryService dockerRegistryService;
+
   @Test
   public void isPythonVersionSupported_noSupportedVersions() {
     ReflectionTestUtils.setField(supportedPythonVersions, SUPPORTED_PYTHON_VERSIONS, Map.of());
@@ -73,6 +78,130 @@ public void getDefaultPythonVersion() {
     Assertions.assertEquals(res, supportedPythonVersions.getDefaultPythonVersion());
   }
 
+  @Test
+  public void
+      getBuilderImage_notSupportedPythonVersionAndProvidedBuilderImage_shouldReturnBuilderImageForDefaultPythonVersion() {
+    var supportedVersions =
+        Map.of(
+            "3.7",
+            Map.of(
+                BASE_IMAGE,
+                "python:3.7-slim",
+                VDK_IMAGE,
+                "test_vdk_image_3.7",
+                BUILDER_IMAGE,
+                "test_builder_image_3.7"),
+            "3.8",
+            Map.of(
+                BASE_IMAGE,
+                "python:3.8-slim",
+                VDK_IMAGE,
+                "test_vdk_image_3.8",
+                BUILDER_IMAGE,
+                "test_builder_image_3.8"),
+            "3.9",
+            Map.of(BASE_IMAGE, "python:3.9-slim", VDK_IMAGE, "test_vdk_image_3.9"),
+            BUILDER_IMAGE,
+            "test_builder_image_3.9");
+    ReflectionTestUtils.setField(
+        supportedPythonVersions, SUPPORTED_PYTHON_VERSIONS, supportedVersions);
+    ReflectionTestUtils.setField(supportedPythonVersions, DEFAULT_PYTHON_VERSION, "3.7");
+    when(dockerRegistryService.builderImage()).thenReturn("default_builder_image");
+
+    Assertions.assertEquals(
+        "test_builder_image_3.7", supportedPythonVersions.getBuilderImage("3.11"));
+  }
+
+  @Test
+  public void
+      getBuilderImage_notSupportedPythonVersionAndNotProvidedBuilderImage_shouldReturnDefaultBuilderImage() {
+    var supportedVersions =
+        Map.of(
+            "3.7",
+            Map.of(BASE_IMAGE, "python:3.7-slim", VDK_IMAGE, "test_vdk_image_3.7"),
+            "3.8",
+            Map.of(
+                BASE_IMAGE,
+                "python:3.8-slim",
+                VDK_IMAGE,
+                "test_vdk_image_3.8",
+                BUILDER_IMAGE,
+                "test_builder_image_3.8"),
+            "3.9",
+            Map.of(BASE_IMAGE, "python:3.9-slim", VDK_IMAGE, "test_vdk_image_3.9"),
+            BUILDER_IMAGE,
+            "test_builder_image_3.9");
+    ReflectionTestUtils.setField(
+        supportedPythonVersions, SUPPORTED_PYTHON_VERSIONS, supportedVersions);
+    ReflectionTestUtils.setField(supportedPythonVersions, DEFAULT_PYTHON_VERSION, "3.7");
+    when(dockerRegistryService.builderImage()).thenReturn("default_builder_image");
+
+    Assertions.assertEquals(
+        "default_builder_image", supportedPythonVersions.getBuilderImage("3.11"));
+  }
+
+  @Test
+  public void
+      getBuilderImage_supportedPythonVersionAndNotProvidedBuilderImage_shouldReturnDefaultBuilderImage() {
+    var supportedVersions =
+        Map.of(
+            "3.7",
+            Map.of(BASE_IMAGE, "python:3.7-slim", VDK_IMAGE, "test_vdk_image_3.7"),
+            "3.8",
+            Map.of(
+                BASE_IMAGE,
+                "python:3.8-slim",
+                VDK_IMAGE,
+                "test_vdk_image_3.8",
+                BUILDER_IMAGE,
+                "test_builder_image_3.8"),
+            "3.9",
+            Map.of(BASE_IMAGE, "python:3.9-slim", VDK_IMAGE, "test_vdk_image_3.9"),
+            BUILDER_IMAGE,
+            "test_builder_image_3.9");
+    ReflectionTestUtils.setField(
+        supportedPythonVersions, SUPPORTED_PYTHON_VERSIONS, supportedVersions);
+    ReflectionTestUtils.setField(supportedPythonVersions, DEFAULT_PYTHON_VERSION, "3.7");
+    when(dockerRegistryService.builderImage()).thenReturn("default_builder_image");
+
+    Assertions.assertEquals(
+        "default_builder_image", supportedPythonVersions.getBuilderImage("3.7"));
+  }
+
+  @Test
+  public void
+      getBuilderImage_supportedPythonVersionAndProvidedBuilderImage_shouldReturnBuilderImageForProvidedPythonVersion() {
+    var supportedVersions =
+        Map.of(
+            "3.7",
+            Map.of(
+                BASE_IMAGE,
+                "python:3.7-slim",
+                VDK_IMAGE,
+                "test_vdk_image_3.7",
+                BUILDER_IMAGE,
+                "test_builder_image_3.7"),
+            "3.8",
+            Map.of(
+                BASE_IMAGE,
+                "python:3.8-slim",
+                VDK_IMAGE,
+                "test_vdk_image_3.8",
+                BUILDER_IMAGE,
+                "test_builder_image_3.8"),
+            "3.9",
+            Map.of(BASE_IMAGE, "python:3.9-slim", VDK_IMAGE, "test_vdk_image_3.9"),
+            BUILDER_IMAGE,
+            "test_builder_image_3.9");
+    ReflectionTestUtils.setField(
+        supportedPythonVersions, SUPPORTED_PYTHON_VERSIONS, supportedVersions);
+    ReflectionTestUtils.setField(supportedPythonVersions, DEFAULT_PYTHON_VERSION, "3.7");
+    when(dockerRegistryService.builderImage()).thenReturn("default_builder_image");
+
+    Assertions.assertEquals(
+        "test_builder_image_3.8", supportedPythonVersions.getBuilderImage("3.8"));
+  }
+
   @Test
   public void getJobBaseImage_defaultImage() {
     var supportedVersions = generateSupportedPythonVersionsConf();