Add core and nlu additional params to training endpoint

changelog/4596.improvement.md

@@ -0,0 +1,3 @@
+Add `augmentation` and `num_threads` arguments to API `POST /model/train`
+
+Fix boolean casting issue for `force_training` and `save_to_default_model_directory` arguments

docs/static/spec/rasa.yml

@@ -464,6 +464,18 @@ paths:
           type: boolean
           default: False
           description: Force a model training even if the data has not changed
+      - in: query
+        name: augmentation
+        schema:
+          type: string
+          default: 50
+          description: How much data augmentation to use during training
+      - in: query
+        name: num_threads
+        schema:
+          type: string
+          default: 1
+          description: Maximum amount of threads to use when training
       - $ref: '#/components/parameters/callback_url'
       requestBody:
         required: true

rasa/server.py

@@ -699,7 +699,6 @@ async def status(request: Request):
     @ensure_loaded_agent(app)
     async def retrieve_tracker(request: Request, conversation_id: Text):
         """Get a dump of a conversation's tracker including its events."""
-
         verbosity = event_verbosity_parameter(request, EventVerbosity.AFTER_RESTART)
         until_time = rasa.utils.endpoints.float_arg(request, "until")
 
@@ -1469,7 +1468,7 @@ def _training_payload_from_json(
     model_output_directory = str(temp_dir)
     if request_payload.get(
         "save_to_default_model_directory",
-        request.args.get("save_to_default_model_directory", True),
+        rasa.utils.endpoints.bool_arg(request, "save_to_default_model_directory", True),
     ):
         model_output_directory = DEFAULT_MODELS_PATH
 
@@ -1479,8 +1478,10 @@ def _training_payload_from_json(
         training_files=str(temp_dir),
         output=model_output_directory,
         force_training=request_payload.get(
-            "force", request.args.get("force_training", False)
+            "force", rasa.utils.endpoints.bool_arg(request, "force_training", False)
         ),
+        core_additional_arguments=_extract_core_additional_arguments(request),
+        nlu_additional_arguments=_extract_nlu_additional_arguments(request),
     )
 
 
@@ -1532,15 +1533,17 @@ def _training_payload_from_yaml(
     rasa.shared.utils.io.write_text_file(decoded, training_data)
 
     model_output_directory = str(temp_dir)
-    if request.args.get("save_to_default_model_directory", True):
+    if rasa.utils.endpoints.bool_arg(request, "save_to_default_model_directory", True):
         model_output_directory = DEFAULT_MODELS_PATH
 
     return dict(
         domain=str(training_data),
         config=str(training_data),
         training_files=str(temp_dir),
         output=model_output_directory,
-        force_training=request.args.get("force_training", False),
+        force_training=rasa.utils.endpoints.bool_arg(request, "force_training", False),
+        core_additional_arguments=_extract_core_additional_arguments(request),
+        nlu_additional_arguments=_extract_nlu_additional_arguments(request),
     )
 
 
@@ -1554,3 +1557,17 @@ def _validate_yaml_training_payload(yaml_text: Text) -> None:
             f"The request body does not contain valid YAML. Error: {e}",
             help_url=DOCS_URL_TRAINING_DATA,
         )
+
+
+def _extract_core_additional_arguments(request: Request) -> Dict:
+    return {
+        "augmentation_factor": rasa.utils.endpoints.int_arg(
+            request, "augmentation", 50
+        ),
+    }
+
+
+def _extract_nlu_additional_arguments(request: Request) -> Dict:
+    return {
+        "num_threads": rasa.utils.endpoints.int_arg(request, "num_threads", 1),
+    }

rasa/utils/endpoints.py

@@ -203,19 +203,35 @@ def bool_arg(request: Request, name: Text, default: bool = True) -> bool:
     """Return a passed boolean argument of the request or a default.
 
     Checks the `name` parameter of the request if it contains a valid
-    boolean value. If not, `default` is returned."""
+    boolean value. If not, `default` is returned.
 
-    return request.args.get(name, str(default)).lower() == "true"
+    Args:
+        request: Sanic request.
+        name: Name of argument.
+        default: Default value for `name` argument.
+
+    Returns:
+        A bool value if `name` is a valid boolean, `default` otherwise.
+    """
+    return str(request.args.get(name, default)).lower() == "true"
 
 
 def float_arg(
     request: Request, key: Text, default: Optional[float] = None
 ) -> Optional[float]:
     """Return a passed argument cast as a float or None.
 
-    Checks the `name` parameter of the request if it contains a valid
-    float value. If not, `None` is returned."""
+    Checks the `key` parameter of the request if it contains a valid
+    float value. If not, `default` is returned.
+
+    Args:
+        request: Sanic request.
+        key: Name of argument.
+        default: Default value for `key` argument.
 
+    Returns:
+        A float value if `key` is a valid float, `default` otherwise.
+    """
     arg = request.args.get(key, default)
 
     if arg is default:
@@ -226,3 +242,31 @@ def float_arg(
     except (ValueError, TypeError):
         logger.warning(f"Failed to convert '{arg}' to float.")
         return default
+
+
+def int_arg(
+    request: Request, key: Text, default: Optional[int] = None
+) -> Optional[int]:
+    """Return a passed argument cast as an int or None.
+
+    Checks the `key` parameter of the request if it contains a valid
+    int value. If not, `default` is returned.
+
+    Args:
+        request: Sanic request.
+        key: Name of argument.
+        default: Default value for `key` argument.
+
+    Returns:
+        An int value if `key` is a valid integer, `default` otherwise.
+    """
+    arg = request.args.get(key, default)
+
+    if arg is default:
+        return arg
+
+    try:
+        return int(str(arg))
+    except (ValueError, TypeError):
+        logger.warning(f"Failed to convert '{arg}' to int.")
+        return default

tests/test_server.py

@@ -686,6 +686,40 @@ def test_training_payload_from_yaml_save_to_default_model_directory(
     assert payload.get("output") == expected
 
 
+@pytest.mark.parametrize(
+    "headers, expected",
+    [
+        ({}, {"augmentation_factor": 50}),
+        ({"augmentation": "25"}, {"augmentation_factor": 25}),
+    ],
+)
+def test_training_payload_from_yaml_core_arguments(
+    headers: Dict, expected: bool, tmp_path: Path
+):
+    request = Mock()
+    request.body = b""
+    request.args = headers
+
+    payload = rasa.server._training_payload_from_yaml(request, tmp_path)
+    assert payload.get("core_additional_arguments") == expected
+
+
+@pytest.mark.parametrize(
+    "headers, expected",
+    [({}, {"num_threads": 1}), ({"num_threads": "2"}, {"num_threads": 2})],
+)
+def test_training_payload_from_yaml_nlu_arguments(
+    headers: Dict, expected: bool, tmp_path: Path
+):
+    request = Mock()
+    request.body = b""
+    request.args = headers
+
+    payload = rasa.server._training_payload_from_yaml(request, tmp_path)
+    assert payload.get("nlu_additional_arguments") == expected
+
+
+@pytest.mark.trains_model
 async def test_train_missing_config(rasa_app: SanicASGITestClient):
     payload = dict(domain="domain data", config=None)
 

tests/utils/test_endpoints.py

@@ -1,5 +1,6 @@
 import logging
 from typing import Text
+from unittest.mock import Mock
 
 import pytest
 from aioresponses import aioresponses
@@ -144,3 +145,44 @@ def test_read_endpoint_config(filename: Text, endpoint_type: Text):
 def test_read_endpoint_config_not_found(filename: Text, endpoint_type: Text):
     conf = endpoint_utils.read_endpoint_config(filename, endpoint_type)
     assert conf is None
+
+
+@pytest.mark.parametrize(
+    "value, default, expected_result",
+    [
+        (None, True, True),
+        (False, True, False),
+        ("false", True, False),
+        ("true", False, True),
+    ],
+)
+def test_bool_arg(value, default, expected_result):
+    request = Mock()
+    request.args = {}
+    if value is not None:
+        request.args = {"key": value}
+    assert endpoint_utils.bool_arg(request, "key", default) == expected_result
+
+
+@pytest.mark.parametrize(
+    "value, default, expected_result",
+    [(None, 0.5, 0.5), (0.5, None, 0.5), ("0.5", 0, 0.5), ("a", 0.5, 0.5)],
+)
+def test_float_arg(value, default, expected_result):
+    request = Mock()
+    request.args = {}
+    if value is not None:
+        request.args = {"key": value}
+    assert endpoint_utils.float_arg(request, "key", default) == expected_result
+
+
+@pytest.mark.parametrize(
+    "value, default, expected_result",
+    [(None, 0, 0), (1, 0, 1), ("1", 0, 1), ("a", 0, 0)],
+)
+def test_int_arg(value, default, expected_result):
+    request = Mock()
+    request.args = {}
+    if value is not None:
+        request.args = {"key": value}
+    assert endpoint_utils.int_arg(request, "key", default) == expected_result