Fix issues on hello-world TF2 notebook

NVIDIA · Apr 12, 2024 · f4df50f · f4df50f
1 parent 8dcac91
commit f4df50f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 49 deletions.
diff --git a/examples/hello-world/hello-tf2/jobs/hello-tf2/app/custom/trainer.py b/examples/hello-world/hello-tf2/jobs/hello-tf2/app/custom/trainer.py
@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
+os.enviro["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
+
 import numpy as np
 import tensorflow as tf
 from tf2_net import Net

diff --git a/examples/hello-world/hello_world.ipynb b/examples/hello-world/hello_world.ipynb
@@ -738,8 +738,9 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
-   "id": "425c7f1d-7cb6-4602-bb88-4b87ff517529",
+   "id": "696a384e-f6da-4044-a7b9-db4c464f52ac",
    "metadata": {},
    "source": [
     "#### Running Tensorflow on local host with GPU \n",
@@ -748,53 +749,13 @@
     "We are running with 1 server, 2 sites in a local machine, which means three process involved for this federated training. \n",
     "If the local host has GPU, you might enter OOM error, due to the way Tensorflow consumes GPU memory. By default, TensorFlow maps nearly all of the GPU memory of all GPUs (subject to CUDA_VISIBLE_DEVICES) visible to the process. If one has multiple process, some of the process will be OOM. To avoid multiple processes grabbing all GPU memory in TF, use the options described in [Limiting GPU memory growth]( https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth). \n",
     "\n",
-    "In our cases,  we prefer that the process only allocates a subset of the available memory, or to only grow the memory usage as is needed by the process. TensorFlow provides two methods to control this. \n",
+    "In our cases,  we prefer that the process only allocates a subset of the available memory, or to only grow the memory usage as is needed by the process. TensorFlow provides two methods to control this, as described in the above link.\n",
+    "\n",
+    "In this example, we explictly set the environment varialble `TF_FORCE_GPU_ALLOW_GROWTH` to `true` at the very beginning of the trainer.py file, which runs in the clients and will allocate GPU memory for training.  With the env var been set, TF will not grab the entire GPU memory and will not cause GPU OOM error when running POC on local host.\n",
+    "\n",
+    "Note that setting the env var `TF_FORCE_GPU_ALLOW_GROWTH` inside this notebook takes no effect because the clients of POC have already started and their env vars are set at the starting time.\n",
+    "\n",
     "\n",
-    "The First method is set the environmental variable TF_FORCE_GPU_ALLOW_GROWTH to true. This configuration is platform specific. \n",
-    "The 2nd method is using the piece of code below"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "026ce33a-90b1-4ef7-8c7c-f25722f2d2ae",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "%env TF_FORCE_GPU_ALLOW_GROWTH=true"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "91cf84ec-57f1-439f-b72b-b33fea1a7f7c",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "import tensorflow as tf\n",
-    "gpus = tf.config.list_physical_devices('GPU')\n",
-    "if gpus:\n",
-    "  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU\n",
-    "  try:\n",
-    "    tf.config.set_logical_device_configuration(\n",
-    "        gpus[0],\n",
-    "        [tf.config.LogicalDeviceConfiguration(memory_limit=1024)])\n",
-    "    logical_gpus = tf.config.list_logical_devices('GPU')\n",
-    "    print(len(gpus), \"Physical GPUs,\", len(logical_gpus), \"Logical GPUs\")\n",
-    "  except RuntimeError as e:\n",
-    "    # Virtual devices must be set before GPUs have been initialized\n",
-    "    print(e)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "0c3f1149-d54d-4f3b-b497-c06deba22fef",
-   "metadata": {},
-   "source": [
     "### 1. Submit job using FLARE API\n",
     "\n",
     "Starting a FLARE API session and submit the hello-tf2 job\n",
@@ -833,7 +794,7 @@
    },
    "outputs": [],
    "source": [
-    "! tail -100 /tmp/nvflare/poc/server/log.txt"
+    "! tail -100 /tmp/nvflare/poc/example_project/prod_00/server/log.txt"
    ]
   },
   {
@@ -977,7 +938,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.17"
+   "version": "3.8.15"
   },
   "vscode": {
    "interpreter": {