Merge pull request #52 from salesforce/numba-integration

colab tutor fix 2

tests/multiblocks_per_env/example_envs/pycuda_tests/test_tag_continuous_multiblocks.py

@@ -94,7 +94,7 @@ def __init__(self, *args, **kwargs):
             blocks_per_env=2,
             assert_equal_num_agents_per_block=False,
             num_episodes=2,
-            use_gpu_testing_mode=False,
+            use_gpu_testing_mode=True,
             testing_bin_filename="test_build_multiblocks.fatbin",
         )
 

tutorials/simple-end-to-end-example.ipynb

@@ -60,7 +60,7 @@
    },
    "outputs": [],
    "source": [
-    "! pip install --quiet \"rl-warp-drive>=1.6.5\" \"torch==1.10.*\" \"torchvision==0.11.*\" \"torchtext==0.11.*\""
+    "! pip install --quiet \"rl-warp-drive>=2.0\" \"torch==1.10.*\" \"torchvision==0.11.*\" \"torchtext==0.11.*\""
    ]
   },
   {

tutorials/tutorial-1.b-warp_drive_basics.ipynb

@@ -636,11 +636,11 @@
     "    num_runs=1,\n",
     "    num_envs=2,\n",
     "    num_agents=3,\n",
-    "    source_code=None,\n",
+    "    source_code_path=None,\n",
     "    episode_length=100,\n",
     "):\n",
     "\n",
-    "    assert source_code is not None\n",
+    "    assert source_code_path is not None\n",
     "\n",
     "    # Initialize the CUDA data manager\n",
     "    cuda_data_manager = NumbaDataManager(\n",
@@ -723,7 +723,7 @@
     "    times.update(\n",
     "        {\n",
     "            f\"envs={num_envs}, agents={num_agents}\": push_random_data_and_increment_timer(\n",
-    "                num_runs, num_envs, num_agents, source_code\n",
+    "                num_runs, num_envs, num_agents, source_code_path\n",
     "            )\n",
     "        }\n",
     "    )"

tutorials/tutorial-2.b-warp_drive_sampler.ipynb

@@ -135,7 +135,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "We first initialize the **NumbaDataManager** and **NumbaFunctionManager**. To illustrate the sampler, we first load a pre-compiled binary file called \"test_build.cubin\". Note that these low-level managers and modules will be hidden and called automatically by WarpDrive in any end-to-end training and simulation. In this and the next tutorials, we want to show how a few fundamental modules work and their performance, that is why some low-level APIs are called."
+    "We first initialize the **NumbaDataManager** and **NumbaFunctionManager**. To illustrate the sampler, we first load a \"test_build_sampler.py\". Note that these low-level managers and modules will be hidden and called automatically by WarpDrive in any end-to-end training and simulation. In this and the next tutorials, we want to show how a few fundamental modules work and their performance, that is why some low-level APIs such as \"import_numba_from_source_code()\" are called. "
    ]
   },
   {
@@ -157,7 +157,23 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "This `_NUMBA_FILEPATH` includes several Numba core services provided by WarpDrive. `test_build.py` collects those core services and include the backend source code for `NumbaSampleController`. Now we can use **NumbaFunctionManager** to load the source code."
+    "This `_NUMBA_FILEPATH` includes several Numba core services provided by WarpDrive.  Since Numba uses JIT compiler, we do not need to compile an executable before loading it as we did in the PyCUDA mode, however, in most cases, the source code still needs some environment configuration file to populate a few global environment constants and settings to the source code, for example, number of environment replicas. WarpDrive provides such template configure file and automatically update placeholders according to the current environment. (import_numba_env_config() is a low-level API, user will not need call those internal APIs directly for any WarpDrive end-to-end simulation and training.)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cuda_function_manager.import_numba_env_config(template_header_file=\"template_env_config.txt\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we can use **NumbaFunctionManager** to load the source code. In this demo, we use `test_build.py` which collects those core services and includes the backend source code for `NumbaSampleController`. "
    ]
   },
   {

tutorials/tutorial-5-training_with_warp_drive.ipynb

@@ -42,11 +42,19 @@
     "- Visualize the behavior using the trained policies.\n",
     "\n",
     "In case you haven't familiarized yourself with WarpDrive, please see the other tutorials we have prepared for you\n",
-    "- [WarpDrive basics](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-1-warp_drive_basics.ipynb)\n",
-    "- [WarpDrive sampler](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-2-warp_drive_sampler.ipynb)\n",
-    "- [WarpDrive reset and log controller](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-3-warp_drive_reset_and_log.ipynb)\n",
+    "1. [WarpDrive basics(intro and pycuda)](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-1.a-warp_drive_basics.ipynb)\n",
+    "2. [WarpDrive basics(numba)](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-1.b-warp_drive_basics.ipynb)\n",
+    "3. [WarpDrive sampler(pycuda)](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-2.a-warp_drive_sampler.ipynb)\n",
+    "4. [WarpDrive sampler(numba)](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-2.b-warp_drive_sampler.ipynb)\n",
+    "5. [WarpDrive resetter and logger](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-3-warp_drive_reset_and_log.ipynb)\n",
+    "\n",
+    "Please also see our tutorials\n",
+    "\n",
+    "6. [Create custom environments (pycuda)](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-4.a-create_custom_environments_pycuda.md)\n",
+    "\n",
+    "7. [Create custom environments (numba)](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-4.b-create_custom_environments_numba.md)\n",
     "\n",
-    "Please also see our [tutorial](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-4-create_custom_environments.md) on creating your own RL environment in CUDA C. Once you have your own environment in CUDA C, this tutorial explains how to integrate it with the WarpDrive framework to perform training."
+    "on creating your own RL environment in CUDA C or Numba. Once you have your own environment, this tutorial explains how to integrate it with the WarpDrive framework to perform training."
    ]
   },
   {
@@ -76,7 +84,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "! pip install --quiet \"torch==1.10.*\" \"torchvision==0.11.*\" \"torchtext==0.11.*\" \"ffmpeg-python\" "
+    "! pip install --quiet \"rl-warp-drive>=2.0\" \"torch==1.10.*\" \"torchvision==0.11.*\" \"torchtext==0.11.*\" \"ffmpeg-python\" "
    ]
   },
   {

warp_drive/managers/numba_managers/numba_function_manager.py

@@ -65,7 +65,9 @@ def __init__(
         print(f"function_manager: Setting Numba to use CUDA device {process_id}")
 
     def import_numba_from_source_code(
-        self, numba_path: str, default_functions_included: bool = True
+        self,
+        numba_path: str,
+        default_functions_included: bool = True,
     ):
         assert (
             self._NUMBA_module is None
@@ -76,6 +78,21 @@ def import_numba_from_source_code(
         if default_functions_included:
             self.initialize_default_functions()
 
+    def import_numba_env_config(
+        self,
+        template_header_file: str,
+        template_path: Optional[str] = None,
+    ):
+        if template_path is None:
+            template_path = f"{get_project_root()}/warp_drive/numba_includes"
+        update_env_header(
+            template_header_file=template_header_file,
+            path=template_path,
+            num_agents=self._num_agents,
+            num_envs=self._num_envs,
+            blocks_per_env=self._blocks_per_env,
+        )
+
     def dynamic_import_numba(
         self,
         env_name: str,