Fix kfto-sdk-minist-test notebook to update local-queue usage flow

Author: abhijeet-dhumal

tests/kfto/kfto_mnist_sdk_test.go

@@ -91,7 +91,8 @@ func runMnistSDK(t *testing.T, trainingImage string) {
 
 	clusterQueue := CreateKueueClusterQueue(test, cqSpec)
 	defer test.Client().Kueue().KueueV1beta1().ClusterQueues().Delete(test.Ctx(), clusterQueue.Name, metav1.DeleteOptions{})
-	CreateKueueLocalQueue(test, namespace.Name, clusterQueue.Name, AsDefaultQueue)
+
+	localQueue := CreateKueueLocalQueue(test, namespace.Name, clusterQueue.Name, AsDefaultQueue)
 
 	jupyterNotebook := string(readFile(test, "resources/mnist_kfto.ipynb"))
 	requirements := readFile(test, "resources/requirements.txt")
@@ -107,8 +108,8 @@ func runMnistSDK(t *testing.T, trainingImage string) {
 		"-c",
 		fmt.Sprintf("pip install papermill && papermill /opt/app-root/notebooks/%s"+
 			" /opt/app-root/src/mnist-kfto-out.ipynb -p namespace %s -p openshift_api_url %s"+
-			" -p token %s -p num_gpus %d -p training_image %s --log-output && sleep infinity",
-			jupyterNotebookConfigMapFileName, namespace.Name, GetOpenShiftApiUrl(test), userToken, 0, trainingImage)}
+			" -p token %s -p num_gpus %d -p training_image %s -p localQueue %s --log-output && sleep infinity",
+			jupyterNotebookConfigMapFileName, namespace.Name, GetOpenShiftApiUrl(test), userToken, 0, trainingImage, localQueue.Name)}
 
 	// Create PVC for Notebook
 	notebookPVC := CreatePersistentVolumeClaim(test, namespace.Name, "10Gi", AccessModes(corev1.ReadWriteOnce))

tests/kfto/resources/mnist_kfto.ipynb

@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 42,
+   "execution_count": null,
    "id": "b55bc3ea-4ce3-49bf-bb1f-e209de8ca47a",
    "metadata": {
     "tags": []
@@ -21,15 +21,12 @@
    "source": [
     "import sys\n",
     "sys.path.append(\"../notebooks\")  # needed to make kfto_sdk_mnist\n",
-    "from kfto_sdk_mnist import train_func\n",
-    "from kubeflow.training import TrainingClient\n",
-    "from kubernetes import client as c\n",
-    "import time"
+    "from kfto_sdk_mnist import train_func "
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 41,
+   "execution_count": null,
    "id": "72dd1751",
    "metadata": {
     "tags": [
@@ -42,54 +39,9 @@
     "num_gpus = \"${num_gpus}\"\n",
     "openshift_api_url = \"${api_url}\"\n",
     "namespace = \"${namespace}\"\n",
-    "token = \"${password}\"\n",
-    "training_image= \"${training_image}\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "aadaa4a7",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def GetDefaultLocalQueue(namespace: str):\n",
-    "    \"\"\"\n",
-    "    Fetches the LocalQueue in the given namespace whose annotation\n",
-    "    \"kueue.x-k8s.io/default-queue\" == \"true\". Returns the dict\n",
-    "    for the queue, or None.\n",
-    "    \"\"\"\n",
-    "    group   = \"kueue.x-k8s.io\"\n",
-    "    version = \"v1beta1\"\n",
-    "    plural  = \"localqueues\"\n",
-    "\n",
-    "    conf = c.Configuration()\n",
-    "    conf.host = openshift_api_url\n",
-    "    conf.verify_ssl = False\n",
-    "    conf.api_key = {\"authorization\": f\"Bearer {token}\"}\n",
-    "\n",
-    "    api_client = c.ApiClient(configuration=conf)\n",
-    "    api = c.CustomObjectsApi(api_client)\n",
-    "\n",
-    "    resp = api.list_namespaced_custom_object(\n",
-    "        group=group, version=version, namespace=namespace, plural=plural\n",
-    "    )\n",
-    "\n",
-    "    default_q = None\n",
-    "    for item in resp.get(\"items\", []):\n",
-    "        ann = item.get(\"metadata\", {}).get(\"annotations\") or {}\n",
-    "        if ann.get(\"kueue.x-k8s.io/default-queue\") == \"true\":\n",
-    "            if default_q is not None:\n",
-    "                raise RuntimeError(\n",
-    "                    f\"multiple LocalQueues annotated as default in {namespace}: \"\n",
-    "                    f\"{default_q['metadata']['name']} and {item['metadata']['name']}\"\n",
-    "                )\n",
-    "            default_q = item['metadata']['name']\n",
-    "\n",
-    "    if default_q is None:\n",
-    "        raise RuntimeError(f\"no LocalQueue annotated as default in namespace {namespace}\")\n",
-    "\n",
-    "    return default_q"
+    "token = \"${token}\"\n",
+    "training_image= \"${training_image}\"\n",
+    "localQueue = \"${localQueue}\""
    ]
   },
   {
@@ -99,13 +51,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "api_key = {\"authorization\": f\"Bearer {token}\"}\n",
-    "# config = c.Configuration(host=openshift_api_url, api_key=token)\n",
-    "# config.verify_ssl = False\n",
-    "tc = TrainingClient()\n",
+    "from kubernetes import client\n",
+    "from kubeflow.training import TrainingClient\n",
     "\n",
-    "# get default local queue\n",
-    "default_local_queue=GetDefaultLocalQueue(namespace)"
+    "configuration = client.Configuration()\n",
+    "configuration.host = openshift_api_url\n",
+    "configuration.api_key = {\"authorization\": f\"Bearer {token}\"}\n",
+    "configuration.verify_ssl = False\n",
+    "api_client = client.ApiClient(configuration)\n",
+    "client = TrainingClient(client_configuration=api_client.configuration)"
    ]
   },
   {
@@ -116,7 +70,7 @@
    "outputs": [],
    "source": [
     "import os\n",
-    "tc.create_job(\n",
+    "client.create_job(\n",
     "   name=\"pytorch-ddp\",\n",
     "   namespace=namespace,\n",
     "   train_func=train_func,\n",
@@ -132,7 +86,7 @@
     "      \"PIP_TRUSTED_HOST\": os.environ.get(\"PIP_TRUSTED_HOST\")\n",
     "   },\n",
     "   labels={\n",
-    "       \"kueue.x-k8s.io/queue-name\": default_local_queue,\n",
+    "       \"kueue.x-k8s.io/queue-name\": localQueue,\n",
     "   }\n",
     ")"
    ]
@@ -144,7 +98,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "while not tc.is_job_succeeded(name=\"pytorch-ddp\", namespace=namespace): \n",
+    "import time\n",
+    "while not client.is_job_succeeded(name=\"pytorch-ddp\", namespace=namespace): \n",
     "    time.sleep(1)\n",
     "print(\"PytorchJob Succeeded!\")"
    ]