TF2 environment settings

changelog/5230.feature.rst

@@ -0,0 +1,16 @@
+Refactor how GPU and CPU environments are configured for TensorFlow 2.0
+
+Please refer to the `documentation <https://rasa.com/docs/rasa/api/tensorflow_usage.html>`_ to understand
+which environment variables to set in what scenarios. A couple of examples are shown below as well:
+
+.. code-block:: python
+
+    # This specifies to use 1024 MB of memory from GPU with logical ID 0 and 2048 MB of memory from GPU with logical ID 1
+    TF_GPU_MEMORY_ALLOC="0:1024, 1:2048"
+
+    # Specifies that atmost 3 CPU threads can be used to parallelize multiple non-blocking operations
+    TF_INTER_OP_PARALLELISM_THREADS="3"
+
+    # Specifies that atmost 2 CPU threads can be used to parallelize a particular operation.
+    TF_INTRA_OP_PARALLELISM_THREADS="2"
+

docs/api/tensorflow_usage.rst

@@ -0,0 +1,55 @@
+:desc: Find out how to configure your environment for efficient usage of TensorFlow inside Rasa Open Source
+
+.. _tensorflow_usage:
+
+Setting up the TensorFlow Runtime
+=================================
+
+TensorFlow allows setting the runtime environment via
+`TF Config submodule <https://www.tensorflow.org/api_docs/python/tf/config>`_. Rasa Open Source supports a smaller subset of these
+configuration options and makes appropriate calls to the ``tf.config`` submodule.
+This smaller subset comprises of configurations that developers frequently use with Rasa Open Source.
+All configuration options are specified using environment variables as shown in subsequent sections.
+
+Optimizing CPU Performance
+--------------------------
+
+Parallelizing one operation
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Set ``TF_INTRA_OP_PARALLELISM_THREADS`` as an environment variable to specify maximum number of threads that can be used
+to parallelize the execution of one operation. If left unspecified, this value defaults to ``0`` which means TensorFlow should
+pick an appropriate value depending on the system configuration.
+
+
+Parallelizing multiple operations
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Set ``TF_INTER_OP_PARALLELISM_THREADS`` as an environment variable to specify maximum number of threads that can be used
+to parallelize the execution of multiple **non-blocking** operations. If left unspecified, this value defaults to ``0``
+which means TensorFlow should pick an appropriate value depending on the system configuration.
+
+
+Optimizing GPU Performance
+--------------------------
+
+Limiting GPU memory growth
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+TensorFlow by default blocks all the available GPU memory for the running process. This can be limiting if you are running
+multiple TensorFlow processes and want to distribute memory across them. To prevent this,
+set an environment variable ``TF_FORCE_GPU_ALLOW_GROWTH`` to ``True``.
+
+
+Restricting absolute GPU memory available
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Often, a developer wants to limit the absolute amount of GPU memory that can be used by a process.
+
+For example, you may have two visible GPUs(``GPU:0`` and ``GPU:1``) and you want to allocate 1024 MB from the first GPU
+and 2048 MB from the second GPU.
+You can do so by setting an environment variable as ``TF_GPU_MEMORY_ALLOC="0:1024, 1:2048"``.
+
+Another scenario can be where you have access to 2 GPUs(``GPU:0`` and ``GPU:1``) but you would like to use only the second
+GPU.
+``TF_GPU_MEMORY_ALLOC="1:2048"`` would make 2048 MB of memory availble from GPU 1.

docs/index.rst

@@ -93,6 +93,7 @@ Understand messages, hold conversations, and connect to messaging channels and A
    api/lock-stores
    api/training-data-importers
    api/core-featurization
+   api/tensorflow_usage
    migration-guide
    changelog
 

rasa/__main__.py

@@ -8,6 +8,7 @@
 from rasa.cli.arguments.default_arguments import add_logging_options
 from rasa.cli.utils import parse_last_positional_argument_as_model_path
 from rasa.utils.common import set_log_level
+import rasa.utils.tensorflow.environment as tf_env
 
 logger = logging.getLogger(__name__)
 
@@ -68,6 +69,8 @@ def main() -> None:
     )
     set_log_level(log_level)
 
+    tf_env.setup_tf_environment()
+
     # insert current path in syspath so custom modules are found
     sys.path.insert(1, os.getcwd())
 

rasa/constants.py

@@ -59,3 +59,7 @@
 
 DEFAULT_SESSION_EXPIRATION_TIME_IN_MINUTES = 60
 DEFAULT_CARRY_OVER_SLOTS_TO_NEW_SESSION = True
+
+ENV_GPU_CONFIG = "TF_GPU_MEMORY_ALLOC"
+ENV_CPU_INTER_OP_CONFIG = "TF_INTER_OP_PARALLELISM_THREADS"
+ENV_CPU_INTRA_OP_CONFIG = "TF_INTRA_OP_PARALLELISM_THREADS"

rasa/utils/tensorflow/environment.py

@@ -0,0 +1,129 @@
+import logging
+import os
+from typing import Text, Dict
+import typing
+
+if typing.TYPE_CHECKING:
+    from tensorflow import config as tf_config
+
+import rasa.utils.common as rasa_utils
+from rasa.constants import (
+    ENV_GPU_CONFIG,
+    ENV_CPU_INTER_OP_CONFIG,
+    ENV_CPU_INTRA_OP_CONFIG,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def setup_gpu_environment() -> None:
+    """Set configuration for a GPU environment based on the environment variable set"""
+
+    gpu_memory_config = os.getenv(ENV_GPU_CONFIG)
+    if gpu_memory_config:
+
+        # Import from tensorflow only if necessary(environment variable was set)
+        from tensorflow import config as tf_config
+
+        parsed_gpu_config = parse_gpu_config(gpu_memory_config)
+        physical_gpus = tf_config.list_physical_devices("GPU")
+
+        # Logic taken from https://www.tensorflow.org/guide/gpu
+        if physical_gpus:
+            for gpu_id, gpu_id_memory in parsed_gpu_config.items():
+
+                allocate_gpu_memory(physical_gpus[gpu_id], gpu_id_memory)
+
+        else:
+            rasa_utils.raise_warning(
+                f"You have an environment variable '{ENV_GPU_CONFIG}' set but no GPUs were detected to configure"
+            )
+
+
+def allocate_gpu_memory(
+    gpu_instance: "tf_config.PhysicalDevice", logical_memory: int
+) -> None:
+
+    from tensorflow import config as tf_config
+
+    try:
+        tf_config.experimental.set_virtual_device_configuration(
+            gpu_instance,
+            [
+                tf_config.experimental.VirtualDeviceConfiguration(
+                    memory_limit=logical_memory
+                )
+            ],
+        )
+
+    except RuntimeError:
+        # Add a helper explanation where the error comes from
+        raise RuntimeError(
+            "Error while setting up tensorflow environment. "
+            "Virtual devices must be set before GPUs have been initialized"
+        )
+
+
+def parse_gpu_config(gpu_memory_config: Text) -> Dict[int, int]:
+    """Parse GPU configuration variable from a string to a dict"""
+
+    # gpu_config is of format "gpu_id_1:gpu_id_1_memory, gpu_id_2: gpu_id_2_memory"
+    # Parse it and store in a dictionary
+    parsed_gpu_config = {}
+
+    try:
+        for instance in gpu_memory_config.split(","):
+            instance_gpu_id, instance_gpu_mem = instance.split(":")
+            instance_gpu_id = int(instance_gpu_id)
+            instance_gpu_mem = int(instance_gpu_mem)
+
+            parsed_gpu_config[instance_gpu_id] = instance_gpu_mem
+    except ValueError:
+        # Add a helper explanation
+        raise ValueError(
+            f"Error parsing GPU configuration. Please cross-check the format of '{ENV_GPU_CONFIG}'"
+        )
+
+    return parsed_gpu_config
+
+
+def setup_cpu_environment() -> None:
+    """Set configuration for the CPU environment based on the environment variable set"""
+
+    inter_op_parallel_threads = os.getenv(ENV_CPU_INTER_OP_CONFIG)
+    intra_op_parallel_threads = os.getenv(ENV_CPU_INTRA_OP_CONFIG)
+
+    if not inter_op_parallel_threads and not intra_op_parallel_threads:
+        return
+
+    from tensorflow import config as tf_config
+
+    if inter_op_parallel_threads:
+
+        try:
+            inter_op_parallel_threads = int(inter_op_parallel_threads.strip())
+        except ValueError:
+            raise ValueError(
+                f"Error parsing the environment variable '{ENV_CPU_INTER_OP_CONFIG}'. Please "
+                f"cross-check the value"
+            )
+
+        tf_config.threading.set_inter_op_parallelism_threads(inter_op_parallel_threads)
+
+    if intra_op_parallel_threads:
+
+        try:
+            intra_op_parallel_threads = int(intra_op_parallel_threads.strip())
+        except ValueError:
+            raise ValueError(
+                f"Error parsing the environment variable '{ENV_CPU_INTRA_OP_CONFIG}'. Please "
+                f"cross-check the value"
+            )
+
+        tf_config.threading.set_intra_op_parallelism_threads(intra_op_parallel_threads)
+
+
+def setup_tf_environment() -> None:
+
+    setup_cpu_environment()
+    setup_gpu_environment()

rasa/utils/train_utils.py

@@ -31,15 +31,6 @@
 logger = logging.getLogger(__name__)
 
 
-def load_tf_config(config: Dict[Text, Any]) -> Optional[tf.compat.v1.ConfigProto]:
-    """Prepare `tf.compat.v1.ConfigProto` for training"""
-
-    if config.get("tf_config") is not None:
-        return tf.compat.v1.ConfigProto(**config.pop("tf_config"))
-    else:
-        return None
-
-
 def normalize(values: np.ndarray, ranking_length: Optional[int] = 0) -> np.ndarray:
     """Normalizes an array of positive numbers over the top `ranking_length` values.
     Other values will be set to 0.

tests/core/test_policies.py

@@ -52,35 +52,6 @@
 from tests.core.utilities import get_tracker, read_dialogue_file, user_uttered
 
 
-def tf_defaults():
-    return {
-        "tf_config": {
-            "device_count": {"CPU": 4},
-            # tell tf.Session to use CPU limit, if you have
-            # more CPU, you can increase this value appropriately
-            "inter_op_parallelism_threads": 0,
-            # the number of threads in the thread pool available
-            # for each process for blocking operation nodes set to 0
-            # to allow the system to select the appropriate value.
-            "intra_op_parallelism_threads": 0,  # tells the degree of thread
-            # parallelism of the tf.Session operation.
-            # the smaller the value, the less reuse the thread will have
-            # and the more likely it will use more CPU cores.
-            # if the value is 0,
-            # tensorflow will automatically select an appropriate value.
-            "gpu_options": {"allow_growth": True}
-            # if set True, will try to allocate
-            # as much GPU memory as possible to support running
-        }
-    }
-
-
-def session_config():
-    import tensorflow as tf
-
-    return tf.ConfigProto(**tf_defaults()["tf_config"])
-
-
 async def train_trackers(domain, augmentation_factor=20):
     return await training.load_data(
         DEFAULT_STORIES_FILE, domain, augmentation_factor=augmentation_factor
@@ -185,18 +156,6 @@ def test_persist_and_load_empty_policy(self, tmpdir):
         loaded = empty_policy.__class__.load(tmpdir.strpath)
         assert loaded is not None
 
-    # TODO test tf config
-    # def test_tf_config(self, trained_policy, tmpdir):
-    #     if hasattr(trained_policy, "session"):
-    #         import tensorflow as tf
-    #
-    #         # noinspection PyProtectedMember
-    #         assert trained_policy.session._config == tf.Session()._config
-    #         trained_policy.persist(tmpdir.strpath)
-    #         loaded = trained_policy.__class__.load(tmpdir.strpath)
-    #         # noinspection PyProtectedMember
-    #         assert loaded.session._config == tf.Session()._config
-
     @staticmethod
     def _get_next_action(policy, events, domain):
         tracker = get_tracker(events)
@@ -212,22 +171,6 @@ def create_policy(self, featurizer, priority):
         return p
 
 
-class TestKerasPolicyWithTfConfig(PolicyTestCollection):
-    def create_policy(self, featurizer, priority):
-        p = KerasPolicy(featurizer, priority, **tf_defaults())
-        return p
-
-    # TODO fix and test tf config
-    @pytest.mark.skip(reason="We need to fix tf.config!")
-    def test_tf_config(self, trained_policy, tmpdir):
-        # noinspection PyProtectedMember
-        assert trained_policy.session._config == session_config()
-        trained_policy.persist(tmpdir.strpath)
-        loaded = trained_policy.__class__.load(tmpdir.strpath)
-        # noinspection PyProtectedMember
-        assert loaded.session._config == session_config()
-
-
 class TestSklearnPolicy(PolicyTestCollection):
     def create_policy(self, featurizer, priority, **kwargs):
         p = SklearnPolicy(featurizer, priority, **kwargs)
@@ -529,22 +472,6 @@ def create_policy(self, featurizer, priority):
         return p
 
 
-class TestTEDPolicyWithTfConfig(TestTEDPolicy):
-    def create_policy(self, featurizer, priority):
-        p = TEDPolicy(featurizer=featurizer, priority=priority, **tf_defaults())
-        return p
-
-    # TODO test tf config
-    @pytest.mark.skip(reason="Fix tf config.")
-    def test_tf_config(self, trained_policy, tmpdir):
-        # noinspection PyProtectedMember
-        assert trained_policy.session._config == session_config()
-        trained_policy.persist(tmpdir.strpath)
-        loaded = trained_policy.__class__.load(tmpdir.strpath)
-        # noinspection PyProtectedMember
-        assert loaded.session._config == session_config()
-
-
 class TestMemoizationPolicy(PolicyTestCollection):
     def create_policy(self, featurizer, priority):
         max_history = None

tests/utils/test_tf_environment.py

@@ -0,0 +1,11 @@
+import pytest
+from rasa.utils.tensorflow.environment import parse_gpu_config
+
+
+@pytest.mark.parametrize(
+    "gpu_config_string, parsed_gpu_config",
+    [("0: 1024", {0: 1024}), ("0:1024, 1:2048", {0: 1024, 1: 2048})],
+)
+def test_gpu_config_parser(gpu_config_string, parsed_gpu_config):
+
+    assert parse_gpu_config(gpu_config_string) == parsed_gpu_config