[𝘀𝗽𝗿] initial version

Created using spr 1.3.4

Author: boomanaiden154

compiler_opt/es/blackbox_evaluator.py

@@ -45,10 +45,18 @@ def get_results(
   def set_baseline(self) -> None:
     raise NotImplementedError()
 
-  @abc.abstractmethod
   def get_rewards(
       self, results: List[concurrent.futures.Future]) -> List[Optional[float]]:
-    raise NotImplementedError()
+    rewards = [None] * len(results)
+
+    for i in range(len(results)):
+      if not results[i].exception():
+        rewards[i] = results[i].result()
+      else:
+        logging.info('Error retrieving result from future: %s',
+                     str(results[i].exception()))
+
+    return rewards
 
 
 @gin.configurable
@@ -96,15 +104,58 @@ def get_results(
   def set_baseline(self) -> None:
     pass
 
-  def get_rewards(
-      self, results: List[concurrent.futures.Future]) -> List[Optional[float]]:
-    rewards = [None] * len(results)
 
-    for i in range(len(results)):
-      if not results[i].exception():
-        rewards[i] = results[i].result()
-      else:
-        logging.info('Error retrieving result from future: %s',
-                     str(results[i].exception()))
+@gin.configurable
+class TraceBlackboxEvaluator(BlackboxEvaluator):
+  """A blackbox evaluator that utilizes trace based cost modelling."""
 
-    return rewards
+  def __init__(self, train_corpus: corpus.Corpus,
+               est_type: blackbox_optimizers.EstimatorType, bb_trace_path: str,
+               function_index_path: str):
+    self._train_corpus = train_corpus
+    self._est_type = est_type
+    self._bb_trace_path = bb_trace_path
+    self._function_index_path = function_index_path
+
+    self._has_baseline = False
+    self._baseline = 1
+
+  def get_results(
+      self, pool: FixedWorkerPool, perturbations: List[policy_saver.Policy]
+  ) -> List[concurrent.futures.Future]:
+    job_args = []
+    for perturbation in perturbations:
+      job_args.append(
+          (self._train_corpus.module_specs, self._function_index_path,
+           self._bb_trace_path, perturbation))
+
+    _, futures = buffered_scheduler.schedule_on_worker_pool(
+        action=lambda w, v: w.compile_corpus_and_evaluate(
+            v[0], v[1], v[2], v[3]),
+        jobs=job_args,
+        worker_pool=pool)
+    concurrent.futures.wait(
+        futures, return_when=concurrent.futures.ALL_COMPLETED)
+    return futures
+
+  def set_baseline(self, pool: FixedWorkerPool) -> None:
+    if self._has_baseline:
+      return
+
+    job_args = [(
+        self._train_corpus.module_specs,
+        self._function_index_path,
+        self._bb_trace_path,
+        None,
+    )]
+
+    _, futures = buffered_scheduler.schedule_on_worker_pool(
+        action=lambda w, v: w.compile_corpus_and_evaluate(
+            v[0], v[1], v[2], v[3]),
+        jobs=job_args,
+        worker_pool=pool)
+
+    concurrent.futures.wait(
+        futures, return_when=concurrent.futures.ALL_COMPLETED)
+    self._baseline = futures[0].result()
+    self._has_baseline = True

compiler_opt/es/blackbox_evaluator_test.py

@@ -41,7 +41,7 @@ def test_sampling_get_results(self):
       self.assertSequenceAlmostEqual([result.result() for result in results],
                                      [1.0, 1.0, 1.0])
 
-  def test_sampling_get_rewards(self):
+  def test_get_rewards(self):
     f1 = concurrent.futures.Future()
     f1.set_exception(None)
     f2 = concurrent.futures.Future()
@@ -50,3 +50,27 @@ def test_sampling_get_rewards(self):
     evaluator = blackbox_evaluator.SamplingBlackboxEvaluator(None, 5, 5, None)
     rewards = evaluator.get_rewards(results)
     self.assertEqual(rewards, [None, 2])
+
+  def test_trace_get_results(self):
+    with local_worker_manager.LocalWorkerPoolManager(
+        blackbox_test_utils.ESTraceWorker, count=3, arg='', kwarg='') as pool:
+      perturbations = [b'00', b'01', b'10']
+      test_corpus = corpus.create_corpus_for_testing(
+          location=self.create_tempdir(),
+          elements=[corpus.ModuleSpec(name='name1', size=1)])
+      evaluator = blackbox_evaluator.TraceBlackboxEvaluator(
+          test_corpus, 5, 'fake_bb_trace_path', 'fake_function_index_path')
+      results = evaluator.get_results(pool, perturbations)
+      self.assertSequenceAlmostEqual([result.result() for result in results],
+                                     [1.0, 1.0, 1.0])
+
+  def test_trace_set_baseline(self):
+    with local_worker_manager.LocalWorkerPoolManager(
+        blackbox_test_utils.ESTraceWorker, count=1, arg='', kwarg='') as pool:
+      test_corpus = corpus.create_corpus_for_testing(
+          location=self.create_tempdir(),
+          elements=[corpus.ModuleSpec(name='name1', size=1)])
+      evaluator = blackbox_evaluator.TraceBlackboxEvaluator(
+          test_corpus, 5, 'fake_bb_trace_path', 'fake_function_index_path')
+      evaluator.set_baseline(pool)
+      self.assertAlmostEqual(evaluator._baseline, 10)

compiler_opt/es/blackbox_test_utils.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 """Test facilities for Blackbox classes."""
 
-from typing import List
+from typing import List, Collection, Optional
 
 import gin
 
@@ -41,3 +41,26 @@ def compile(self, policy: policy_saver.Policy,
       return self.function_value
     else:
       return 0.0
+
+
+class ESTraceWorker(worker.Worker):
+  """Temporary placeholder worker.
+  
+  This is a test worker for TraceBlackboxEvaluator that expects a slightly
+  different interface than other workers.
+  """
+
+  def __init__(self, arg, *, kwarg):
+    del arg  # Unused.
+    del kwarg  # Unused.
+    self._function_value = 0.0
+
+  def compile_corpus_and_evaluate(
+      self, modules: Collection[corpus.ModuleSpec], function_index_path: str,
+      bb_trace_path: str,
+      tflite_policy: Optional[policy_saver.Policy]) -> float:
+    if modules and function_index_path and bb_trace_path and tflite_policy:
+      self._function_value += 1
+      return self._function_value
+    else:
+      return 10