Genericize interactive loop away from AxClient (#1143)

Summary:
Pull Request resolved: #1143

Move the interactive loop away from requiring AxClient. Use new generic infra to replicate behavior of previous diff in a new method named `interactive_optimize_with_client`.

Reviewed By: lena-kashtelyan

Differential Revision: D39545664

fbshipit-source-id: af1b923b17770b159fc374359a3e94a42f7f02ba

Author: mpolson64

ax/service/interactive_loop.py

@@ -7,7 +7,7 @@
 from logging import Logger
 from queue import Queue
 from threading import Event, Lock, Thread
-from typing import Callable, Tuple
+from typing import Any, Callable, Dict, Tuple
 
 from ax.core.types import TEvaluationOutcome, TParameterization
 
@@ -21,11 +21,15 @@
 IDLE_SLEEP_SEC = 0.1
 
 
-# TODO[mpolson64] Create `optimize` method that constructs its own ax_client
-def optimize_with_client(
-    ax_client: AxClient,
+def interactive_optimize(
     num_trials: int,
     candidate_queue_maxsize: int,
+    # Callable[[Queue[Tuple[TParameterization, int]], Event, ...], None]
+    candidate_generator_function: Callable[..., None],
+    candidate_generator_kwargs: Dict[str, Any],
+    # Callable[[Queue[Tuple[int, TEvaluationOutcome]], Event, ...], None]
+    data_attacher_function: Callable[..., None],
+    data_attacher_kwargs: Dict[str, Any],
     elicitation_function: Callable[[TParameterization], TEvaluationOutcome],
 ) -> None:
     """
@@ -40,27 +44,23 @@ def optimize_with_client(
     only be used in contexts where it is necessary for the user to not experience any
     "lag" while candidates are being generated.
 
-    Extract results of the experiment from the AxClient passed in.
-
-    The basic structure is as follows: One thread tries for a lock on the AxClient,
-    generates a candidate, and enqueues it to a candidate queue. Another thread tries
-    for the same lock, takes all the trial outcomes in the outcome queue, and attaches
-    them to the AxClient. The main thread pops a candidate off the candidate queue,
-    elicits response from the user, and puts the response onto the outcome queue.
-
     Args:
-        ax_client: An AxClient properly configured for the experiment intended to be
-            run. Construct and configure this before passing in.
         num_trials: The total number of trials to be run.
         candidate_queue_maxsize: The maximum number of candidates to pregenerate.
+        candidate_generator_function: A function taking in a queue and event that
+            enqueues candidates (generated by any means). See
+            `ax_client_candidate_generator` for an example.
+        candidate_generator_kwargs: kwargs to be passed into
+            `candidate_generator_function` when it is spawned as a thread.
+        data_attacher_function: A function taking in a queue and event that attaches
+            observations to Ax. See `ax_client_data_attacher` for an example.
+        data_attacher_kwargs: kwargs to be passed into `data_attacher_function` when
+            it is spawned as a thread.
         elicitation_function: Function from parameterization (as returned by
-        `AxClient.get_next_trial`) to outcome (as expected by
-        `AxClient.complete_trial`).
+            `AxClient.get_next_trial`) to outcome (as expected by
+            `AxClient.complete_trial`).
     """
 
-    # Construct a lock to ensure only one thread my access the AxClient at any moment
-    ax_client_lock = Lock()
-
     # Construct queues to buffer inputs and outputs of the AxClient
     candidate_queue: "Queue[Tuple[TParameterization, int]]" = Queue(
         maxsize=candidate_queue_maxsize
@@ -74,18 +74,21 @@ def optimize_with_client(
     # Construct threads to run candidate thread-safe pregeneration and thread-safe
     # data attaching respectively
     candidate_generator_thread = Thread(
-        target=_candidate_generator,
+        target=candidate_generator_function,
         args=(
-            ax_client,
-            ax_client_lock,
-            candidate_generator_stop_event,
             candidate_queue,
+            candidate_generator_stop_event,
             num_trials,
         ),
+        kwargs=candidate_generator_kwargs,
     )
     data_attacher_thread = Thread(
-        target=_data_attacher,
-        args=(ax_client, ax_client_lock, data_attacher_stop_event, data_queue),
+        target=data_attacher_function,
+        args=(
+            data_queue,
+            data_attacher_stop_event,
+        ),
+        kwargs=data_attacher_kwargs,
     )
 
     candidate_generator_thread.start()
@@ -105,12 +108,43 @@ def optimize_with_client(
     data_attacher_stop_event.set()
 
 
-def _candidate_generator(
+def interactive_optimize_with_client(
     ax_client: AxClient,
-    lock: Lock,
-    stop_event: Event,
+    num_trials: int,
+    candidate_queue_maxsize: int,
+    elicitation_function: Callable[[TParameterization], TEvaluationOutcome],
+) -> None:
+    """
+    Implementation of `interactive_loop` using the AxClient. Extract results of the
+    experiment from the AxClient passed in.
+
+    The basic structure is as follows: One thread tries for a lock on the AxClient,
+    generates a candidate, and enqueues it to a candidate queue. Another thread tries
+    for the same lock, takes all the trial outcomes in the outcome queue, and attaches
+    them to the AxClient. The main thread pops a candidate off the candidate queue,
+    elicits response from the user, and puts the response onto the outcome queue.
+    """
+
+    # Construct a lock to ensure only one thread my access the AxClient at any moment
+    ax_client_lock = Lock()
+
+    interactive_optimize(
+        num_trials=num_trials,
+        candidate_queue_maxsize=candidate_queue_maxsize,
+        candidate_generator_function=ax_client_candidate_generator,
+        candidate_generator_kwargs={"ax_client": ax_client, "lock": ax_client_lock},
+        data_attacher_function=ax_client_data_attacher,
+        data_attacher_kwargs={"ax_client": ax_client, "lock": ax_client_lock},
+        elicitation_function=elicitation_function,
+    )
+
+
+def ax_client_candidate_generator(
     queue: "Queue[Tuple[TParameterization, int]]",
+    stop_event: Event,
     num_trials: int,
+    ax_client: AxClient,
+    lock: Lock,
 ) -> None:
     """Thread-safe method for generating the next trial from the AxClient and
     enqueueing it to the candidate queue. The number of candidates pre-generated is
@@ -138,11 +172,11 @@ def _candidate_generator(
         time.sleep(IDLE_SLEEP_SEC)
 
 
-def _data_attacher(
+def ax_client_data_attacher(
+    queue: "Queue[Tuple[int, TEvaluationOutcome]]",
+    stop_event: Event,
     ax_client: AxClient,
     lock: Lock,
-    stop_event: Event,
-    queue: "Queue[Tuple[int, TEvaluationOutcome]]",
 ) -> None:
     """Thread-safe method for attaching evaluation outcomes to the AxClient from the
     outcome queue. If the AxClient's lock is acquired all data in the outcome queue

ax/service/tests/test_interactive_loop.py

@@ -12,7 +12,7 @@
 import numpy as np
 from ax.core.types import TEvaluationOutcome
 from ax.service.ax_client import AxClient, TParameterization
-from ax.service.interactive_loop import optimize_with_client
+from ax.service.interactive_loop import interactive_optimize_with_client
 from ax.utils.common.testutils import TestCase
 from ax.utils.measurement.synthetic_functions import hartmann6
 from ax.utils.testing.mock import fast_botorch_optimize
@@ -49,7 +49,7 @@ def _elicit(parameterization: TParameterization) -> TEvaluationOutcome:
             minimize=True,
         )
 
-        optimize_with_client(
+        interactive_optimize_with_client(
             ax_client=ax_client,
             num_trials=15,
             candidate_queue_maxsize=3,
@@ -90,7 +90,7 @@ def _elicit(parameterization: TParameterization) -> TEvaluationOutcome:
         ax_client.generation_strategy._steps[0].max_parallelism = 1
 
         with self.assertLogs(logger="ax", level=WARN) as logger:
-            optimize_with_client(
+            interactive_optimize_with_client(
                 ax_client=ax_client,
                 num_trials=3,
                 candidate_queue_maxsize=3,