training-operator test: Update image (#70)

Update the image to use one from katib repository, which is 2.75GB instead of 10GB
that the previous one was.

Closes #69

tests/notebooks/training/training-integration.ipynb

@@ -415,10 +415,13 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "PYTORCHJOB_NAME = \"pytorch-dist-mnist-gloo\"\n",
+ "PYTORCHJOB_NAME = \"pytorch-mnist-gloo\"\n",
  "PYTORCHJOB_CONTAINER = \"pytorch\"\n",
- "PYTORCHJOB_IMAGE = \"kubeflow/pytorch-dist-mnist:v1-3a360ba\"\n",
- "# The image above should be updated with each release with the latest available in the registry."
+ "PYTORCHJOB_IMAGE = \"kubeflowkatib/pytorch-mnist-cpu:v0.16.0\"\n",
+ "# The image above should be updated with each release with the corresponding Katib version used in CKF release.\n",
+ "# Note that instead of using the [image from training-operator repository](https://github.com/kubeflow/training-operator/blob/master/examples/pytorch/mnist/Dockerfile),\n",
+ "# the one [from Katib](https://github.com/kubeflow/katib/blob/master/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.cpu) is being used\n",
+ "# due to the large size of the first one."
  ]
  },
  {
@@ -430,7 +433,8 @@
  "container = V1Container(\n",
  " name=PYTORCHJOB_CONTAINER,\n",
  " image=PYTORCHJOB_IMAGE,\n",
- " args=[\"--backend\", \"gloo\"],\n",
+ " args=[\"--backend\", \"gloo\", \"--epochs\", \"2\"],\n",
+ " # Passing `epochs`argument since kubeflowkatib image defaults to 10.\n",
  ")\n",
  "\n",
  "replica_spec = V1ReplicaSpec(\n",