Merge remote-tracking branch 'upstream/main'

Author: dchourasia

tests/kfto/resources/mnist_kfto.ipynb

@@ -52,6 +52,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import os\n",
     "tc.create_job(\n",
     "   name=\"pytorch-ddp\",\n",
     "   namespace=namespace,\n",
@@ -60,7 +61,12 @@
     "   resources_per_worker={\"gpu\": num_gpus},\n",
     "   base_image=training_image,\n",
     "   packages_to_install=[\"torchvision==0.19.0\",\"minio==7.2.13\"],\n",
-    "   env_vars={\"NCCL_DEBUG\": \"INFO\", \"TORCH_DISTRIBUTED_DEBUG\": \"DETAIL\"}\n",
+    "   env_vars={\n",
+    "      \"NCCL_DEBUG\": \"INFO\", \n",
+    "      \"TORCH_DISTRIBUTED_DEBUG\": \"DETAIL\", \n",
+    "      \"PIP_INDEX_URL\": os.environ.get(\"PIP_INDEX_URL\"),\n",
+    "      \"PIP_TRUSTED_HOST\": os.environ.get(\"PIP_TRUSTED_HOST\")\n",
+    "   }\n",
     ")"
    ]
   },