Add SDK version rollback code. (#1634)

Author: Laren-AWS

aws_sagemaker_studio/sagemaker_neo_compilation_jobs/deploy_tensorflow_model_on_Inf1_instance/tensorflow_distributed_mnist_neo_inf1_studio.ipynb

@@ -34,6 +34,9 @@
    "source": [
     "import sagemaker\n",
     "if sagemaker.__version__ >= '2':\n",
+    "    orig_sm_version = sagemaker.__version__\n",
+    "    with open('orig_sm_version.txt', \"w\") as f:\n",
+    "        f.write(orig_sm_version)\n",
     "    %pip install \"sagemaker>=1.14.2,<2\"\n",
     "\n",
     "if sagemaker.__version__ >= '2':\n",
@@ -313,6 +316,36 @@
    "source": [
     "sagemaker_session.delete_endpoint(optimized_predictor.endpoint)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Rollback the SageMaker Python SDK version"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# rollback the SageMaker Python SDK to the kernel's original version\n",
+    "if os.path.exists('orig_sm_version.txt'):\n",
+    "    with open('orig_sm_version.txt', 'r') as f:\n",
+    "        orig_sm_version = f.read()\n",
+    "    print(f\"Original version: {orig_sm_version}\")\n",
+    "    print(f\"Current version: {sagemaker.__version__}\")\n",
+    "    %pip install sagemaker=={orig_sm_version}\n",
+    "    os.remove('orig_sm_version.txt')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Restart the kernel to run the updated version of the SDK."
+   ]
   }
  ],
  "metadata": {