[two_step_routing] More precise and predictable controls for internal…

…Parameters.

Instead of just taking the internal parameters from the input request, allow
specifying overrides for each stage of the request with clearly defined
override chains.

python/gmpro/two_step_routing/_shared.py

@@ -45,6 +45,27 @@ class IntegrationMode(enum.Enum):
 class Options:
   """Options for the two-step planner.
 
+  When determining the `internalParameters` attribute of the request in each
+  phase, the library uses the last value of the chain that is provided and not
+  None. When no value is provided/all are None, `internalParameters` is not used
+  in this phase:
+  - initial local request:
+    1. Options.local_internal_parameters.
+  - initial global request:
+    1. `internalParameters` from the input request,
+    2. Options.global_internal_parameters.
+  - local_refinement_request:
+    1. Options.local_internal_parameters,
+    2. Options.local_refinement_internal_parameters.
+  - global_refinement_requet:
+    1. `internalParameters` from the input request,
+    2. Options.global_internal_parameters,
+    3. Options.global_refinement_internal_parameters.
+  Also note that `two_step_routing_main.py` adds one additional override for the
+  last global (refinement) request that is sent to the server. This override
+  can't be done from the library because it doesn't know the final number of
+  refinement phases.
+
   Attributes:
     local_model_grouping: The grouping strategy used in the local model.
     initial_local_model_grouping: The grouping strategy used in the initial
@@ -70,6 +91,16 @@ class Options:
       used while computing route for the transition. These fields are extensions
       to the CFR API, and converting a JSON response with these fields to the
       CFR proto may fail.
+    local_internal_parameters: An override of the `internalParameters` attribute
+      for local requests. When None, this override is not used.
+    global_internal_parameters: An override of the `internalParameters`
+      attribute for global requests. When None, this override is not used.
+    local_refinement_internal_parameters: An override of the
+      `internalParameters` attribute for local refinement requests. When None,
+      the override is not used.
+    global_refinement_internal_parameters: An override of the
+      `internalParameters` attribute for global refinement requests. When None,
+      the override is not used.
   """
 
   initial_local_model_grouping: _parking.InitialLocalModelGrouping
@@ -85,6 +116,32 @@ class Options:
   use_deprecated_fields: bool = True
   travel_mode_in_merged_transitions: bool = False
 
+  local_internal_parameters: str | None = None
+  global_internal_parameters: str | None = None
+  local_refinement_internal_parameters: str | None = None
+  global_refinement_internal_parameters: str | None = None
+
+
+def override_internal_parameters(
+    request: cfr_json.OptimizeToursRequest, *overrides: str | None
+) -> None:
+  """Overrides `internalParameters` in `request` with values.
+
+  Sets `internalParameters` to the last non-None value from `values`. When
+  `values` is empty or contains only None, the `internalParameters` attribute of
+  the request is left intact.
+
+  Args:
+    request: The request, in which the internal parameters are overridden.
+    *overrides: The override values, considered from first to last.
+  """
+  internal_parameters = request.get("internalParameters")
+  for value in overrides:
+    if value is not None:
+      internal_parameters = value
+  if internal_parameters is not None:
+    request["internalParameters"] = internal_parameters
+
 
 def copy_shared_options(
     from_request: cfr_json.OptimizeToursRequest,

python/gmpro/two_step_routing/_shared_test.py

@@ -7,7 +7,52 @@
 from . import _shared
 
 
+class OverrideInternalParametersTest(unittest.TestCase):
+  """Tests for override_internal_parameters."""
+
+  def test_no_initial_and_no_overrides(self):
+    request: cfr_json.OptimizeToursRequest = {}
+    _shared.override_internal_parameters(request)
+    self.assertEqual(request, {})
+    _shared.override_internal_parameters(request, None)
+    self.assertEqual(request, {})
+    _shared.override_internal_parameters(request, None, None, None)
+    self.assertEqual(request, {})
+
+  def test_no_initial_and_some_overrides(self):
+    test_cases = (
+        (("foo",), "foo"),
+        (("foo", None, "bar"), "bar"),
+        (("foo", None, "bar", None), "bar"),
+    )
+    for overrides, expected_value in test_cases:
+      with self.subTest(overrides=overrides):
+        request: cfr_json.OptimizeToursRequest = {}
+        _shared.override_internal_parameters(request, *overrides)
+        self.assertEqual(request, {"internalParameters": expected_value})
+
+  def test_some_initial_and_no_overrides(self):
+    request: cfr_json.OptimizeToursRequest = {"internalParameters": "foo"}
+    _shared.override_internal_parameters(request)
+    self.assertEqual(request, {"internalParameters": "foo"})
+    _shared.override_internal_parameters(request, None, None, None)
+    self.assertEqual(request, {"internalParameters": "foo"})
+
+  def test_some_initial_and_some_overrides(self):
+    test_cases = (
+        (("bar",), "bar"),
+        (("bar", None, "baz"), "baz"),
+        (("bar", None, "baz", None), "baz"),
+    )
+    for overrides, expected_value in test_cases:
+      with self.subTest(overrides=overrides):
+        request: cfr_json.OptimizeToursRequest = {"internalParameters": "foo"}
+        _shared.override_internal_parameters(request, *overrides)
+        self.assertEqual(request, {"internalParameters": expected_value})
+
+
 class CopySharedOptionsTest(unittest.TestCase):
+  """Tests for copy_shared_options."""
 
   def _run_copy_shared_options_test(
       self, from_request, to_request, expected_to_request

python/gmpro/two_step_routing/two_step_routing.py

@@ -255,6 +255,9 @@ def make_local_request(self) -> cfr_json.OptimizeToursRequest:
         "model": local_model,
     }
     _shared.copy_shared_options(from_request=self._request, to_request=request)
+    internal_parameters = self._options.local_internal_parameters
+    if internal_parameters is not None:
+      request["internalParameters"] = internal_parameters
     return request
 
   def make_global_request(
@@ -347,12 +350,11 @@ def make_global_request(
       consider_road_traffic = self._request.get("considerRoadTraffic")
       if consider_road_traffic is not None:
         request["considerRoadTraffic"] = consider_road_traffic
-    # TODO(ondrasej): Consider applying internal parameters also to the local
-    # request; potentially, add separate internal parameters for the local and
-    # the global models to the configuration of the planner.
-    internal_parameters = self._request.get("internalParameters")
-    if internal_parameters is not None:
-      request["internalParameters"] = internal_parameters
+    _shared.override_internal_parameters(
+        request,
+        self._request.get("internalParameters"),
+        self._options.global_internal_parameters,
+    )
     return request
 
   def make_local_refinement_request(
@@ -550,6 +552,11 @@ def make_local_refinement_request(
         "injectedFirstSolutionRoutes": refinement_injected_routes,
     }
     _shared.copy_shared_options(from_request=self._request, to_request=request)
+    _shared.override_internal_parameters(
+        request,
+        self._options.local_internal_parameters,
+        self._options.local_refinement_internal_parameters,
+    )
     return request
 
   def integrate_local_refinement(
@@ -961,6 +968,12 @@ def integrate(
         self._integrated_global_request["considerRoadTraffic"] = (
             consider_road_traffic
         )
+      _shared.override_internal_parameters(
+          self._integrated_global_request,
+          self._request.get("internalParameters"),
+          self._options.global_internal_parameters,
+          self._options.global_refinement_internal_parameters,
+      )
 
       _global_model.assert_routes_handle_same_shipments(
           self._global_response, {"routes": self._integrated_global_routes}

python/gmpro/two_step_routing/two_step_routing_main.py

@@ -43,6 +43,7 @@
 from ..json import cfr_api
 from ..json import cfr_json
 from ..json import io_utils
+from . import _shared
 from . import two_step_routing
 
 
@@ -73,6 +74,11 @@ class Flags(cfr_api.Flags):
   global_timeout: cfr_json.DurationString
   local_refinement_timeout: cfr_json.DurationString
   global_refinement_timeout: cfr_json.DurationString
+  local_internal_parameters: str | None
+  global_internal_parameters: str | None
+  local_refinement_internal_parameters: str | None
+  global_refinement_internal_parameters: str | None
+  last_global_internal_parameters: str | None
 
   @classmethod
   def add_arguments(cls, parser: argparse.ArgumentParser) -> None:
@@ -200,6 +206,61 @@ def add_arguments(cls, parser: argparse.ArgumentParser) -> None:
         default=False,
         action=argparse.BooleanOptionalAction,
     )
+    parser.add_argument(
+        "--local_internal_parameters",
+        help=(
+            "An override for the `internalParameters` attribute used in local"
+            " requests made by the two-step library. See the docstring of"
+            " _shared.Options for a description of the internal parameters"
+            " override logic."
+        ),
+        default=None,
+        action="store",
+    )
+    parser.add_argument(
+        "--global_internal_parameters",
+        help=(
+            "An override for the `internalParameters` attribute used in global"
+            " requests made by the two-step library. See the docstring of"
+            " _shared.Options for a description of the internal parameters"
+            " override logic."
+        ),
+        default=None,
+        action="store",
+    )
+    parser.add_argument(
+        "--local_refinement_internal_parameters",
+        help=(
+            "An override for the `internalParameters` attribute used in local"
+            " refinement requests made by the two-step library. See the"
+            " docstring of _shared.Options for a description of the internal"
+            " parameters override logic."
+        ),
+        default=None,
+        action="store",
+    )
+    parser.add_argument(
+        "--global_refinement_internal_parameters",
+        help=(
+            "An override for the `internalParameters` attribute used in global"
+            " refinement requests made by the two-step library. See the"
+            " docstring of _shared.Options for a description of the internal"
+            " parameters override logic."
+        ),
+        default=None,
+        action="store",
+    )
+    parser.add_argument(
+        "--last_global_internal_parameters",
+        help=(
+            "An override for the `internalParameters` attribute used in the"
+            " last global request sent to the server. See the docstring of"
+            " _shared.Options for a description of the internal parameters"
+            " override logic."
+        ),
+        default=None,
+        action="store",
+    )
 
 
 def _optimize_tours_and_write_response(
@@ -269,6 +330,10 @@ def _run_two_step_planner() -> None:
       local_model_vehicle_fixed_cost=flags.local_model_vehicle_fixed_cost,
       travel_mode_in_merged_transitions=flags.travel_mode_in_merged_transitions,
       use_deprecated_fields=not flags.no_deprecated_fields,
+      local_internal_parameters=flags.local_internal_parameters,
+      global_internal_parameters=flags.global_internal_parameters,
+      local_refinement_internal_parameters=flags.local_refinement_internal_parameters,
+      global_refinement_internal_parameters=flags.global_refinement_internal_parameters,
   )
   planner = two_step_routing.Planner(
       request_json, parking_locations, parking_for_shipment, options
@@ -307,11 +372,16 @@ def make_filename(stem, timeout_string=None):
   # global model. We will be injecting the solution from the base global model
   # into a refined global model, and for this to work correctly, we need to use
   # the same duration/distance matrices in both solves.
-  global_request_traffic_override = False if flags.num_refinements > 0 else None
+  is_last_global_cfr_request = False if flags.num_refinements > 0 else None
   global_request = planner.make_global_request(
       local_response,
-      consider_road_traffic_override=global_request_traffic_override,
+      consider_road_traffic_override=is_last_global_cfr_request,
   )
+  if is_last_global_cfr_request:
+    _shared.override_internal_parameters(
+        global_request,
+        flags.last_global_internal_parameters,
+    )
   global_request["searchMode"] = 2
   io_utils.write_json_to_file(
       make_filename("global_request", flags.local_timeout),
@@ -396,6 +466,10 @@ def make_filename(stem, timeout_string=None):
     if not is_last_global_cfr_request:
       # Override the live traffic option for all but the last global request.
       integrated_global_request["considerRoadTraffic"] = False
+    else:
+      _shared.override_internal_parameters(
+          integrated_global_request, flags.last_global_internal_parameters
+      )
 
     io_utils.write_json_to_file(
         make_filename("integrated_local_request"),