Convert MNIST PyTorch and MNIST Ray test cases to go

tests/go.mod

@@ -1,9 +1,10 @@
-module github.com/opendatahub-io/distributed-workloads/tests/new-tests
+module github.com/opendatahub-io/distributed-workloads/tests
 
 require (
 	github.com/onsi/gomega v1.27.10
 	github.com/openshift/api v0.0.0-20230718161610-2a3e8b481cec
-	github.com/project-codeflare/codeflare-operator v0.2.3
+	github.com/project-codeflare/codeflare-operator v0.2.4-0.20230913142530-526bb53289e1
+	github.com/project-codeflare/multi-cluster-app-dispatcher v1.34.1
 	github.com/ray-project/kuberay/ray-operator v0.0.0-20230908233208-a8f730e5a2b6
 	k8s.io/api v0.27.2
 	k8s.io/apimachinery v0.27.2
@@ -30,7 +31,6 @@ require (
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/openshift/client-go v0.0.0-20230718165156-6014fb98e86a // indirect
-	github.com/project-codeflare/multi-cluster-app-dispatcher v1.34.1 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/stretchr/testify v1.8.4 // indirect
 	golang.org/x/net v0.12.0 // indirect

tests/go.sum

@@ -107,8 +107,8 @@ github.com/openshift/client-go v0.0.0-20230718165156-6014fb98e86a h1:ZKewwwEIURD
 github.com/openshift/client-go v0.0.0-20230718165156-6014fb98e86a/go.mod h1:EjhPQjEm8HM3GThz5ywNGLEec1P1IjTn08kwzdvupvA=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/project-codeflare/codeflare-operator v0.2.3 h1:aVsJD519hBjFoftSFlVjEzySM7JjoKgRGhxLwy88YKE=
-github.com/project-codeflare/codeflare-operator v0.2.3/go.mod h1:6J91NMtSthXp/gFTl1CDlyJo+rhBPQ+jc5OagdEqaVk=
+github.com/project-codeflare/codeflare-operator v0.2.4-0.20230913142530-526bb53289e1 h1:1iBWbUlDja0qpMnyH+u8uL0qEjlxMfldFxygFgyViqU=
+github.com/project-codeflare/codeflare-operator v0.2.4-0.20230913142530-526bb53289e1/go.mod h1:6J91NMtSthXp/gFTl1CDlyJo+rhBPQ+jc5OagdEqaVk=
 github.com/project-codeflare/multi-cluster-app-dispatcher v1.34.1 h1:ZNQ/JPdjS6CtaAzt6SNqaoWcpwS1PyVdgZlmIYikPLI=
 github.com/project-codeflare/multi-cluster-app-dispatcher v1.34.1/go.mod h1:Yge6GRNpO9YIDfeL+XOcCE9xbmfCTD5C1h5dlW87mxQ=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=

tests/integration/mcad_ray_test.go

@@ -0,0 +1,100 @@
+/*
+Copyright 2023.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package integration
+
+import (
+	"testing"
+
+	. "github.com/onsi/gomega"
+	cfosupport "github.com/project-codeflare/codeflare-operator/test/support"
+	mcadv1beta1 "github.com/project-codeflare/multi-cluster-app-dispatcher/pkg/apis/controller/v1beta1"
+	rayv1alpha1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1alpha1"
+
+	corev1 "k8s.io/api/core/v1"
+	rbacv1 "k8s.io/api/rbac/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	support "github.com/opendatahub-io/distributed-workloads/tests/integration/support"
+)
+
+func TestMCADRay(t *testing.T) {
+	test := cfosupport.With(t)
+
+	test.T().Skip("Requires https://github.com/project-codeflare/codeflare-sdk/issues/190")
+
+	// Create a namespace
+	namespace := test.NewTestNamespace()
+
+	// Test configuration
+	jupyterNotebookConfigMapFileName := "mnist_ray_mini.ipynb"
+	config := &corev1.ConfigMap{
+		TypeMeta: metav1.TypeMeta{
+			APIVersion: corev1.SchemeGroupVersion.String(),
+			Kind:       "ConfigMap",
+		},
+		ObjectMeta: metav1.ObjectMeta{
+			Name: "notebooks-ray",
+		},
+		BinaryData: map[string][]byte{
+			// MNIST MCAD Notebook
+			jupyterNotebookConfigMapFileName: ReadFile(test, "resources/mnist_ray_mini.ipynb"),
+			"mnist.py":                       ReadFile(test, "resources/mnist.py"),
+			"requirements.txt":               ReadFile(test, "resources/requirements.txt"),
+		},
+		Immutable: cfosupport.Ptr(true),
+	}
+	config, err := test.Client().Core().CoreV1().ConfigMaps(namespace.Name).Create(test.Ctx(), config, metav1.CreateOptions{})
+	test.Expect(err).NotTo(HaveOccurred())
+	test.T().Logf("Created ConfigMap %s/%s successfully", config.Namespace, config.Name)
+
+	// Create RBAC, retrieve token for user with limited rights
+	policyRules := []rbacv1.PolicyRule{
+		{
+			Verbs:     []string{"get", "create", "delete", "list", "patch", "update"},
+			APIGroups: []string{mcadv1beta1.GroupName},
+			Resources: []string{"appwrappers"},
+		},
+		{
+			Verbs:     []string{"get", "list"},
+			APIGroups: []string{rayv1alpha1.GroupVersion.Group},
+			Resources: []string{"rayclusters", "rayclusters/status"},
+		},
+		{
+			Verbs:     []string{"get", "list"},
+			APIGroups: []string{"route.openshift.io"},
+			Resources: []string{"routes"},
+		},
+	}
+	token := support.CreateTestRBAC(test, namespace, policyRules)
+
+	// Create Notebook CR
+	support.CreateNotebook(test, namespace, token, config.Name, jupyterNotebookConfigMapFileName)
+
+	// Make sure the AppWrapper is created and running
+	test.Eventually(cfosupport.AppWrappers(test, namespace), cfosupport.TestTimeoutLong).
+		Should(
+			And(
+				HaveLen(1),
+				ContainElement(WithTransform(cfosupport.AppWrapperName, HavePrefix("mnistjob"))),
+				ContainElement(WithTransform(cfosupport.AppWrapperState, Equal(mcadv1beta1.AppWrapperStateActive))),
+			),
+		)
+
+	// Make sure the AppWrapper finishes and is deleted
+	test.Eventually(cfosupport.AppWrappers(test, namespace), cfosupport.TestTimeoutLong).
+		Should(HaveLen(0))
+}

tests/integration/pytorch_mcad_test.go

@@ -0,0 +1,91 @@
+/*
+Copyright 2023.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package integration
+
+import (
+	"testing"
+
+	. "github.com/onsi/gomega"
+	cfosupport "github.com/project-codeflare/codeflare-operator/test/support"
+	mcadv1beta1 "github.com/project-codeflare/multi-cluster-app-dispatcher/pkg/apis/controller/v1beta1"
+
+	corev1 "k8s.io/api/core/v1"
+	rbacv1 "k8s.io/api/rbac/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	support "github.com/opendatahub-io/distributed-workloads/tests/integration/support"
+)
+
+func TestMnistPyTorchMCAD(t *testing.T) {
+	test := cfosupport.With(t)
+
+	// Create a namespace
+	namespace := test.NewTestNamespace()
+
+	// Test configuration
+	jupyterNotebookConfigMapFileName := "mnist_mcad_mini.ipynb"
+	config := &corev1.ConfigMap{
+		TypeMeta: metav1.TypeMeta{
+			APIVersion: corev1.SchemeGroupVersion.String(),
+			Kind:       "ConfigMap",
+		},
+		ObjectMeta: metav1.ObjectMeta{
+			Name: "notebooks-mcad",
+		},
+		BinaryData: map[string][]byte{
+			// MNIST MCAD Notebook
+			jupyterNotebookConfigMapFileName: ReadFile(test, "resources/mnist_mcad_mini.ipynb"),
+		},
+		Immutable: cfosupport.Ptr(true),
+	}
+	config, err := test.Client().Core().CoreV1().ConfigMaps(namespace.Name).Create(test.Ctx(), config, metav1.CreateOptions{})
+	test.Expect(err).NotTo(HaveOccurred())
+	test.T().Logf("Created ConfigMap %s/%s successfully", config.Namespace, config.Name)
+
+	// Create RBAC, retrieve token for user with limited rights
+	policyRules := []rbacv1.PolicyRule{
+		{
+			Verbs:     []string{"get", "create", "delete", "list", "patch", "update"},
+			APIGroups: []string{mcadv1beta1.GroupName},
+			Resources: []string{"appwrappers"},
+		},
+		// Needed for job.logs()
+		{
+			Verbs:     []string{"get"},
+			APIGroups: []string{""},
+			Resources: []string{"pods/log"},
+		},
+	}
+	token := support.CreateTestRBAC(test, namespace, policyRules)
+
+	// Create Notebook CR
+	support.CreateNotebook(test, namespace, token, config.Name, jupyterNotebookConfigMapFileName)
+
+	// Make sure the AppWrapper is created and running
+	test.Eventually(cfosupport.AppWrappers(test, namespace), cfosupport.TestTimeoutLong).
+		Should(
+			And(
+				HaveLen(1),
+				ContainElement(WithTransform(cfosupport.AppWrapperName, HavePrefix("mnistjob"))),
+				ContainElement(WithTransform(cfosupport.AppWrapperState, Equal(mcadv1beta1.AppWrapperStateActive))),
+			),
+		)
+
+	// Make sure the AppWrapper finishes and is deleted
+	test.Eventually(cfosupport.AppWrappers(test, namespace), cfosupport.TestTimeoutLong).
+		Should(HaveLen(0))
+}

tests/integration/ray_test.go

@@ -38,6 +38,10 @@ func TestRayCluster(t *testing.T) {
 	test := support.With(t)
 	test.T().Parallel()
 
+	// This test is unstable. It seems that RayJob CR sometimes trigger 2 jobs in Ray, causing confusion in KubeRay operator.
+	// Needs to be checked with newer KubeRay version. If still unstable then it needs to be reported.
+	test.T().Skip("Requires https://github.com/opendatahub-io/distributed-workloads/issues/65")
+
 	// Create a namespace
 	namespace := test.NewTestNamespace()
 

tests/integration/resources/mnist.py

@@ -149,7 +149,7 @@ def test_dataloader(self):
 trainer = Trainer(
     accelerator="auto",
     # devices=1 if torch.cuda.is_available() else None,  # limiting got iPython runs
-    max_epochs=5,
+    max_epochs=2,
     callbacks=[TQDMProgressBar(refresh_rate=20)],
     num_nodes=int(os.environ.get("GROUP_WORLD_SIZE", 1)),
     devices=int(os.environ.get("LOCAL_WORLD_SIZE", 1)),

tests/integration/resources/mnist_mcad_mini.ipynb

@@ -0,0 +1,93 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b55bc3ea-4ce3-49bf-bb1f-e209de8ca47a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Import pieces from codeflare-sdk\n",
+    "from codeflare_sdk.job.jobs import DDPJobDefinition\n",
+    "from time import sleep"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "47ca5c15",
+   "metadata": {
+    "tags": ["parameters"]
+   },
+   "outputs": [],
+   "source": [
+    "#parameters\n",
+    "namespace = \"default\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "26b21373",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "job = DDPJobDefinition(name=\"mnistjob\", script=\"mnist.py\", scheduler_args={\"namespace\": namespace}, j=\"1x1\", gpu=0, cpu=1, memMB=2000, image=\"quay.io/project-codeflare/mnist-job-test:v0.0.1\").submit()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d24e9f95",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "finished = False\n",
+    "while not finished:\n",
+    "    sleep(1)\n",
+    "    try:\n",
+    "        finished = (\"Epoch 4: 100%\" in job.logs())\n",
+    "    except:\n",
+    "        finished = False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f078b7cd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "job.cancel()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.13"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "f9f85f796d01129d0dd105a088854619f454435301f6ffec2fea96ecbd9be4ac"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}