Make notebook iamge configurable based on static values for RHOAI and…

… ODH Signed-off-by: Jakub Stejskal <[email protected]>
skodjob · Jan 11, 2024 · 0164f93 · 0164f93
1 parent c346bc7
commit 0164f93
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 6 deletions.
diff --git a/src/main/java/io/odh/test/framework/logs/LogCollector.java b/src/main/java/io/odh/test/framework/logs/LogCollector.java
@@ -58,6 +58,7 @@ private static void saveClusterState(Path logpath) throws IOException {
         Files.writeString(logpath.resolve("dsc.yml"), cmdClient.exec(false, false, "get", "dsc", "-o", "yaml").out());
         Files.writeString(logpath.resolve("dsci.yml"), cmdClient.exec(false, false, "get", "dsci", "-o", "yaml").out());
         Files.writeString(logpath.resolve("subscriptions.yml"), cmdClient.exec(false, false, "get", "subscriptions.operators.coreos.com", "--all-namespaces", "-o", "yaml").out());
+        Files.writeString(logpath.resolve("notebooks.yml"), cmdClient.exec(false, false, "get", "notebook", "--all-namespaces", "-o", "yaml").out());
         kube.listPodsByPrefixInName(OdhConstants.BUNDLE_OPERATOR_NAMESPACE, "opendatahub-operator-controller-manager").forEach(pod -> {
             writeLogsFromPods(logpath, pod);
         });

diff --git a/src/main/java/io/odh/test/framework/manager/resources/NotebookResource.java b/src/main/java/io/odh/test/framework/manager/resources/NotebookResource.java
@@ -8,6 +8,7 @@
 import io.fabric8.kubernetes.client.dsl.MixedOperation;
 import io.fabric8.kubernetes.client.dsl.Resource;
 import io.fabric8.openshift.client.OpenShiftClient;
+import io.odh.test.Environment;
 import io.odh.test.OdhConstants;
 import io.odh.test.TestUtils;
 import io.odh.test.framework.manager.ResourceManager;
@@ -16,11 +17,28 @@
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.Map;
+import java.util.Objects;
+
 import org.apache.commons.io.IOUtils;
 
 
 public class NotebookResource implements ResourceType<Notebook> {
 
+    private static final String REGISTRY_PATH = "image-registry.openshift-image-registry.svc:5000";
+    public static final String PYTORCH_IMAGE = "pytorch";
+    public static final String PYTORCH_2023_2_TAG = "2023.2";
+    private static final Map<String, String> ODH_IMAGES_MAP;
+
+    static {
+        ODH_IMAGES_MAP = Map.<String, String>of(PYTORCH_IMAGE, "jupyter-pytorch-notebook");
+    }
+
+    private static final Map<String, String> RHOAI_IMAGES_MAP;
+
+    static {
+        RHOAI_IMAGES_MAP = Map.<String, String>of(PYTORCH_IMAGE, "pytorch");
+    }
     private static final String NOTEBOOK_TEMPLATE_PATH = "notebook.yaml";
     @Override
     public String getKind() {
@@ -56,7 +74,7 @@ public static MixedOperation<Notebook, KubernetesResourceList<Notebook>, Resourc
         return ResourceManager.getClient().getClient().resources(Notebook.class);
     }
 
-    public static Notebook loadDefaultNotebook(String namespace, String name) throws IOException {
+    public static Notebook loadDefaultNotebook(String namespace, String name, String image) throws IOException {
         InputStream is = TestUtils.getFileFromResourceAsStream(NOTEBOOK_TEMPLATE_PATH);
         String notebookString = IOUtils.toString(is, "UTF-8");
         notebookString = notebookString.replace("my-project", namespace).replace("my-workbench", name);
@@ -66,7 +84,17 @@ public static Notebook loadDefaultNotebook(String namespace, String name) throws
         // Set correct username
         String username = ResourceManager.getKubeCmdClient().getUsername().strip();
         notebookString = notebookString.replace("odh_user", username);
+        // Replace image
+        notebookString = notebookString.replace("notebook_image_placeholder", image);
 
         return TestUtils.configFromYaml(notebookString, Notebook.class);
     }
+
+    public static String getNotebookImage(String imageName, String imageTag) {
+        if (Objects.equals(Environment.PRODUCT, Environment.PRODUCT_DEFAULT)) {
+            return REGISTRY_PATH + "/" + OdhConstants.CONTROLLERS_NAMESPACE + "/" + ODH_IMAGES_MAP.get(imageName) + ":" + imageTag;
+        } else {
+            return REGISTRY_PATH + "/" + OdhConstants.CONTROLLERS_NAMESPACE + "/" + RHOAI_IMAGES_MAP.get(imageName) + ":" + imageTag;
+        }
+    }
 }
diff --git a/src/main/resources/notebook.yaml b/src/main/resources/notebook.yaml
@@ -27,8 +27,8 @@ spec:
                                   --ServerApp.quit_button=False
                                   --ServerApp.tornado_settings={"user":"odh_user","hub_host":"odh_dashboard_route","hub_prefix":"/projects/my-project"}
             - name: JUPYTER_IMAGE
-              value: image-registry.openshift-image-registry.svc:5000/opendatahub/jupyter-pytorch-notebook:2023.2
-          image: image-registry.openshift-image-registry.svc:5000/opendatahub/jupyter-pytorch-notebook:2023.2
+              value: notebook_image_placeholder
+          image: notebook_image_placeholder
           imagePullPolicy: Always
           livenessProbe:
             failureThreshold: 3

diff --git a/src/test/java/io/odh/test/e2e/standard/NotebookST.java b/src/test/java/io/odh/test/e2e/standard/NotebookST.java
@@ -70,7 +70,8 @@ void testCreateSimpleNotebook() throws IOException {
                 .build();
         ResourceManager.getInstance().createResourceWithoutWait(pvc);
 
-        Notebook notebook = new NotebookBuilder(NotebookResource.loadDefaultNotebook(NTB_NAMESPACE, NTB_NAME)).build();
+        String notebookImage = NotebookResource.getNotebookImage(NotebookResource.PYTORCH_IMAGE, NotebookResource.PYTORCH_2023_2_TAG);
+        Notebook notebook = new NotebookBuilder(NotebookResource.loadDefaultNotebook(NTB_NAMESPACE, NTB_NAME, notebookImage)).build();
         ResourceManager.getInstance().createResourceWithoutWait(notebook);
 
         LabelSelector lblSelector = new LabelSelectorBuilder()

diff --git a/src/test/java/io/odh/test/e2e/upgrade/UpgradeAbstract.java b/src/test/java/io/odh/test/e2e/upgrade/UpgradeAbstract.java
@@ -104,9 +104,10 @@ public void deployNotebook(String namespace, String name) throws IOException {
                 .build();
         ResourceManager.getInstance().createResourceWithoutWait(pvc);
 
-        Notebook notebook = new NotebookBuilder(NotebookResource.loadDefaultNotebook(namespace, name)).build();
+        String notebookImage = NotebookResource.getNotebookImage(NotebookResource.PYTORCH_IMAGE, NotebookResource.PYTORCH_2023_2_TAG);
+        Notebook notebook = new NotebookBuilder(NotebookResource.loadDefaultNotebook(namespace, name, notebookImage)).build();
         if (!Environment.PRODUCT.equals(Environment.PRODUCT_DEFAULT)) {
-            notebook = new NotebookBuilder(NotebookResource.loadDefaultNotebook(namespace, name))
+            notebook = new NotebookBuilder(NotebookResource.loadDefaultNotebook(namespace, name, notebookImage))
                     .editSpec()
                     .editNotebookspecTemplate()
                     .editOrNewSpec()