Add core and nlu additional params to training endpoint

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

@@ -69,7 +69,7 @@
 from rasa.core.utils import AvailableEndpoints
 from rasa.nlu.emulators.no_emulator import NoEmulator
 from rasa.nlu.test import run_evaluation, CVEvaluationResult
-from rasa.utils.endpoints import EndpointConfig
+from rasa.utils.endpoints import EndpointConfig, bool_arg, float_arg, int_arg
 
 if TYPE_CHECKING:
     from ssl import SSLContext
@@ -699,9 +699,8 @@ 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")
+        until_time = float_arg(request, "until")
 
         tracker = await app.agent.create_processor().fetch_tracker_with_initial_session(
             conversation_id
@@ -745,9 +744,7 @@ async def append_events(request: Request, conversation_id: Text):
 
                 output_channel = _get_output_channel(request, tracker)
 
-                if rasa.utils.endpoints.bool_arg(
-                    request, EXECUTE_SIDE_EFFECTS_QUERY_KEY, False
-                ):
+                if bool_arg(request, EXECUTE_SIDE_EFFECTS_QUERY_KEY, False):
                     await processor.execute_side_effects(
                         events, tracker, output_channel
                     )
@@ -823,10 +820,8 @@ async def replace_events(request: Request, conversation_id: Text):
     @ensure_conversation_exists()
     async def retrieve_story(request: Request, conversation_id: Text):
         """Get an end-to-end story corresponding to this conversation."""
-        until_time = rasa.utils.endpoints.float_arg(request, "until")
-        fetch_all_sessions = rasa.utils.endpoints.bool_arg(
-            request, "all_sessions", default=False
-        )
+        until_time = float_arg(request, "until")
+        fetch_all_sessions = bool_arg(request, "all_sessions", default=False)
 
         try:
             stories = get_test_stories(
@@ -1091,7 +1086,7 @@ async def evaluate_stories(
 
         test_data = _test_data_file_from_payload(request, temporary_directory, ".md")
 
-        use_e2e = rasa.utils.endpoints.bool_arg(request, "e2e", default=False)
+        use_e2e = bool_arg(request, "e2e", default=False)
 
         try:
             evaluation = await test(
@@ -1469,7 +1464,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),
+        bool_arg(request, "save_to_default_model_directory", True),
     ):
         model_output_directory = DEFAULT_MODELS_PATH
 
@@ -1479,8 +1474,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", bool_arg(request, "force_training", False)
         ),
+        core_additional_arguments=_extract_core_additional_arguments(request),
+        nlu_additional_arguments=_extract_nlu_additional_arguments(request),
     )
 
 
@@ -1532,15 +1529,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 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=bool_arg(request, "force_training", False),
+        core_additional_arguments=_extract_core_additional_arguments(request),
+        nlu_additional_arguments=_extract_nlu_additional_arguments(request),
     )
 
 
@@ -1554,3 +1553,15 @@ 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": int_arg(request, "augmentation", 50),
+    }
+
+
+def _extract_nlu_additional_arguments(request: Request) -> Dict:
+    return {
+        "num_threads": int_arg(request, "num_threads", 1),
+    }

rasa/utils/endpoints.py

@@ -203,19 +203,19 @@ 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."""
-
-    return request.args.get(name, str(default)).lower() == "true"
+    boolean value. If not, `default` is returned.
+    """
+    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.
+    """
     arg = request.args.get(key, default)
 
     if arg is default:
@@ -226,3 +226,23 @@ 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.
+    """
+    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

@@ -704,6 +704,39 @@ 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)