Merge branch 'main' into users/boomanaiden154/bc-percentage-correct

.github/workflows/main.yml

@@ -14,12 +14,12 @@ on:
 
 jobs:
   LicenseCheck:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v4
       - run: ./check-license.sh
   Envvars:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     outputs:
       version_matrix: ${{ steps.set-output.outputs.version_matrix }}
       do_cache: ${{ steps.set-output.outputs.do_cache }}
@@ -28,25 +28,25 @@ jobs:
         run: |
           if [ -z $ACT ]
           then
-            _ver="['3.8', '3.9', '3.10']"
+            _ver="['3.10','3.11']"
             _cache="1"
           else
-            # 3.8 instead of '3.8' to make github act work.
-            _ver="[3.8]"
+            # 3.10 instead of '3.10' to make github act work.
+            _ver="[3.10,3.11]"
             _cache="0"
           fi
           echo "version_matrix=$_ver" >> $GITHUB_OUTPUT
           echo "do_cache=$_cache" >> $GITHUB_OUTPUT
   Checks:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     needs: [Envvars]
     strategy:
       matrix:
         python-version: ${{ fromJSON(needs.Envvars.outputs.version_matrix) }}
         task: [Typecheck, Lint, Yapf, Test]
         include:
           - task: Typecheck
-            cmd: pytype -j auto --overriding-parameter-count-checks .
+            cmd: pytype -j auto .
           - task: Lint
             cmd: pylint --rcfile .pylintrc --recursive yes .
           - task: Yapf

.gitignore

@@ -1,3 +1,4 @@
 __pycache__
 .vscode
 .idea
+*.swp

.pylintrc

@@ -104,12 +104,14 @@ disable=abstract-method,
         no-self-use,
 	no-value-for-parameter, # gin causes false positives
         nonzero-method,
+        not-callable, # lots of false positives
         oct-method,
         old-division,
         old-ne-operator,
         old-octal-literal,
         old-raise-syntax,
         parameter-unpacking,
+        possibly-used-before-assignment, # false positives with control flow
         print-statement,
         raising-string,
         range-builtin-not-iterating,
@@ -132,6 +134,7 @@ disable=abstract-method,
         too-many-instance-attributes,
         too-many-locals,
         too-many-nested-blocks,
+        too-many-positional-arguments,
         too-many-public-methods,
         too-many-return-statements,
         too-many-statements,
@@ -140,6 +143,7 @@ disable=abstract-method,
         unicode-builtin,
         unnecessary-pass,
         unpacking-in-except,
+        used-before-assignment,
         useless-else-on-loop,
         useless-object-inheritance,
         useless-suppression,
@@ -425,7 +429,7 @@ valid-metaclass-classmethod-first-arg=mcs
 
 # Exceptions that will emit a warning when being caught. Defaults to
 # "Exception"
-overgeneral-exceptions=StandardError,
-                       Exception,
-                       BaseException
+overgeneral-exceptions=builtins.StandardError,
+                       builtins.Exception,
+                       builtins.BaseException
 

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.9
+FROM python:3.10
 
 WORKDIR /ml-compiler-opt
 COPY . .

Pipfile

@@ -4,19 +4,20 @@ verify_ssl = true
 name = "pypi"
 
 [packages]
-absl-py = "==1.0.0"
+absl-py = "==2.1.0"
 gin-config = "==0.5.0"
-psutil = "==5.9.0"
-tf-agents = "==0.16.0"
-tensorflow = "==2.12.0"
-dm-reverb = "==0.11.0"
+psutil = "==6.1.0"
+tf-agents = "==0.19.0"
+tensorflow = "==2.15.0"
+dm-reverb = "==0.14.0"
 
 [dev-packages]
-pylint = "==2.14.1"
-pytest = "==7.1.2"
-pytype = "==2022.06.06"
-yapf = "==0.32.0"
+pylint = "==3.3.2"
+pytest = "==8.3.4"
+pytype = "==2024.10.11"
+yapf = "==0.43.0"
 
 [ci]
-joblib = "==1.1.1"
-shap = "==0.41.0"
+joblib = "==1.4.2"
+shap = "==0.46.0"
+

compiler_opt/distributed/buffered_scheduler_test.py

@@ -63,16 +63,18 @@ def square(self, the_value, extra_factor=1):
                            [x * x for x in range(10)])
 
       _, futures = buffered_scheduler.schedule_on_worker_pool(
-          lambda w, v: w.square(**v), [dict(the_value=v) for v in range(10)],
-          pool)
+          lambda w, v: w.square(**v), [{
+              'the_value': v
+          } for v in range(10)], pool)
       worker.wait_for(futures)
       self.assertListEqual([f.result() for f in futures],
                            [x * x for x in range(10)])
 
       # same idea, but mix some kwargs
       _, futures = buffered_scheduler.schedule_on_worker_pool(
-          lambda w, v: w.square(v[0], **v[1]),
-          [(v, dict(extra_factor=10)) for v in range(10)], pool)
+          lambda w, v: w.square(v[0], **v[1]), [(v, {
+              'extra_factor': 10
+          }) for v in range(10)], pool)
       worker.wait_for(futures)
       self.assertListEqual([f.result() for f in futures],
                            [x * x * 10 for x in range(10)])

compiler_opt/distributed/worker_test.py

@@ -35,8 +35,11 @@ def test_gin_args(self):
         gin.bind_parameter('_test.SomeType.argument', 42)
       real_args = worker.get_full_worker_args(
           SomeType, more_args=2, even_more_args='hi')
-      self.assertDictEqual(real_args,
-                           dict(argument=42, more_args=2, even_more_args='hi'))
+      self.assertDictEqual(real_args, {
+          'argument': 42,
+          'more_args': 2,
+          'even_more_args': 'hi'
+      })
 
 
 if __name__ == '__main__':

compiler_opt/es/blackbox_learner.py

@@ -32,6 +32,10 @@
 from compiler_opt.rl import policy_saver
 from compiler_opt.es import blackbox_evaluator  # pylint: disable=unused-import
 
+# Pytype cannot pick up the pyi file for tensorflow.summary. Disable the error
+# here as these errors are false positives.
+# pytype: disable=pyi-error
+
 # If less than 40% of requests succeed, skip the step.
 _SKIP_STEP_SUCCESS_RATIO = 0.4
 
@@ -168,7 +172,7 @@ def _get_perturbations(self) -> List[npt.NDArray[np.float32]]:
     rng = np.random.default_rng(seed=self._seed)
     for _ in range(self._config.total_num_perturbations):
       perturbations.append(
-          rng.normal(size=(len(self._model_weights))) *
+          rng.normal(size=len(self._model_weights)) *
           self._config.precision_parameter)
     return perturbations
 

compiler_opt/es/blackbox_optimizers.py

@@ -983,16 +983,14 @@ def trust_region_test(self, current_input: FloatArray,
                      str(tr_imp_ratio))
       if should_reject:
         self.radius *= self.params['reject_factor']
-        if self.radius < self.params['minimum_radius']:
-          self.radius = self.params['minimum_radius']
+        self.radius = max(self.radius, self.params['minimum_radius'])
         self.is_returned_step = True
         print('Step rejected. Shrink: ' + str(self.radius) + log_message)
         return False
       else:  # accept step
         if should_shrink:
           self.radius *= self.params['shrink_factor']
-          if self.radius < self.params['minimum_radius']:
-            self.radius = self.params['minimum_radius']
+          self.radius = max(self.radius, self.params['minimum_radius'])
           print('Shrink: ' + str(self.radius) + log_message)
         elif should_grow:
           self.radius *= self.params['grow_factor']

compiler_opt/es/es_trainer_lib.py

@@ -70,6 +70,7 @@
 @gin.configurable
 def train(additional_compilation_flags=(),
           delete_compilation_flags=(),
+          replace_compilation_flags=(),
           worker_class=None):
   """Train with ES."""
 
@@ -113,11 +114,11 @@ def train(additional_compilation_flags=(),
   logging.info("Parameter dimension: %s", initial_parameters.shape)
   logging.info("Initial parameters: %s", initial_parameters)
 
-  cps = corpus.create_corpus_for_testing(
-      location=_TRAIN_CORPORA.value,
-      elements=[corpus.ModuleSpec(name="smth", size=1)],
+  cps = corpus.Corpus(
+      data_path=_TRAIN_CORPORA.value,
       additional_flags=additional_compilation_flags,
-      delete_flags=delete_compilation_flags)
+      delete_flags=delete_compilation_flags,
+      replace_flags=replace_compilation_flags)
 
   # Construct policy saver
   saved_policy = policy_utils.create_actor_policy(greedy=True)

compiler_opt/rl/compilation_runner.py

@@ -122,15 +122,15 @@ def __init__(self):
     Exception.__init__(self)
 
 
-def kill_process_ignore_exceptions(p: 'subprocess.Popen[bytes]'):
+def kill_process_ignore_exceptions(p: 'subprocess.Popen[bytes]'):  # pylint: disable=useless-return
   # kill the process and ignore exceptions. Exceptions would be thrown if the
   # process has already been killed/finished (which is inherently in a race
   # condition with us killing it)
   try:
     p.kill()
     p.wait()
   finally:
-    return  # pylint: disable=lost-exception
+    return  # pylint: disable=lost-exception,return-in-finally
 
 
 class WorkerCancellationManager:

compiler_opt/rl/distributed/agent.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 """PPO agent definition and utility functions."""
 
-from typing import Optional, Text, Tuple
+from typing import Optional, Tuple
 
 from absl import logging
 
@@ -67,7 +67,7 @@ def __init__(self,
                train_step_counter: Optional[tf.Variable] = None,
                aggregate_losses_across_replicas=False,
                loss_scaling_factor=1.,
-               name: Optional[Text] = 'PPOClipAgent'):
+               name: Optional[str] = 'PPOClipAgent'):
     """Creates a PPO Agent implementing the clipped probability ratios.
 
     Args:

compiler_opt/rl/distributed/learner.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 """Utility to create MLGO policy learner."""
 
-from typing import Callable, List, Optional, Text, Tuple
+from typing import Callable, List, Optional, Tuple
 
 from absl import logging
 
@@ -47,7 +47,7 @@ class MLGOPPOLearner(object):
   """
 
   def __init__(self,
-               root_dir: Text,
+               root_dir: str,
                train_step: tf.Variable,
                model_id: tf.Variable,
                agent: ppo_agent.PPOAgent,

compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py

@@ -320,8 +320,10 @@ def __init__(
       max_horizon_to_explore=np.inf,
       explore_on_features: Optional[Dict[str, Callable[[tf.Tensor],
                                                        bool]]] = None,
-      obs_action_specs: Optional[Tuple[time_step.TimeStep,
-                                       tensor_spec.BoundedTensorSpec,]] = None,
+      obs_action_specs: Optional[Tuple[
+          time_step.TimeStep,
+          tensor_spec.BoundedTensorSpec,
+      ]] = None,
       reward_key: str = '',
       **kwargs,
   ):
@@ -748,8 +750,10 @@ def __init__(
       exploration_policy_paths: Optional[str] = None,
       explore_on_features: Optional[Dict[str, Callable[[tf.Tensor],
                                                        bool]]] = None,
-      obs_action_specs: Optional[Tuple[time_step.TimeStep,
-                                       tensor_spec.BoundedTensorSpec,]] = None,
+      obs_action_specs: Optional[Tuple[
+          time_step.TimeStep,
+          tensor_spec.BoundedTensorSpec,
+      ]] = None,
       base_path: Optional[str] = None,
       partitions: List[float] = [
           0.,
@@ -943,9 +947,7 @@ def gen_trajectories(
     )
     not_done = result_futures
     succeeded_idx = 0
-    succeeded: List[Optional[Tuple[Tuple[int, ProfilingDictValueType,
-                                         ProfilingDictValueType],
-                                   tf.train.SequenceExample]]] = []
+    succeeded = []
 
     for written_files_idx in range(num_output_files):
       written_per_file = 0

compiler_opt/rl/local_data_collector.py

@@ -105,13 +105,12 @@ def _schedule_jobs(self, policy: policy_saver.Policy, model_id: int,
     # by now, all the pending work, which was signaled to cancel, must've
     # finished
     self._join_pending_jobs()
-    jobs = [
-        dict(
-            loaded_module_spec=loaded_module_spec,
-            policy=policy,
-            reward_stat=self._reward_stat_map[loaded_module_spec.name],
-            model_id=model_id) for loaded_module_spec in sampled_modules
-    ]
+    jobs = [{
+        'loaded_module_spec': loaded_module_spec,
+        'policy': policy,
+        'reward_stat': self._reward_stat_map[loaded_module_spec.name],
+        'model_id': model_id
+    } for loaded_module_spec in sampled_modules]
 
     (self._workers,
      self._current_futures) = buffered_scheduler.schedule_on_worker_pool(

compiler_opt/rl/policy_saver.py

@@ -163,8 +163,8 @@ def __init__(self, policy_dict: Dict[str, tf_policy.TFPolicy]):
     self._policy_saver_dict: Dict[str, Tuple[
         policy_saver.PolicySaver, tf_policy.TFPolicy]] = {
             policy_name: (policy_saver.PolicySaver(
-                policy, batch_size=1, use_nest_path_signatures=False), policy)
-            for policy_name, policy in policy_dict.items()
+                policy, batch_size=1, use_nest_path_signatures=False), policy
+                         ) for policy_name, policy in policy_dict.items()
         }
 
   def _write_output_signature(

compiler_opt/rl/regalloc/regalloc_network.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 """Actor network for Register Allocation."""
 
-from typing import Optional, Sequence, Callable, Text, Any
+from typing import Optional, Sequence, Callable, Any
 
 import gin
 import tensorflow as tf
@@ -78,7 +78,7 @@ def __init__(
       kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
       batch_squash: bool = True,
       dtype: tf.DType = tf.float32,
-      name: Text = 'RegAllocNetwork'):
+      name: str = 'RegAllocNetwork'):
     """Creates an instance of `RegAllocNetwork`.
 
     Args:

compiler_opt/rl/train_bc.py

@@ -21,6 +21,10 @@
 from absl import logging
 import gin
 
+# Pytype cannot pick up the pyi file for tensorflow.summary. Disable the error
+# here as these errors are false positives.
+# pytype: disable=pyi-error
+
 # <Internal> Using XM - flags.  # pylint: disable=unused-import
 from compiler_opt.rl import agent_config
 from compiler_opt.rl import data_reader

compiler_opt/rl/train_locally.py

@@ -159,8 +159,8 @@ def sequence_example_iterator_fn(seq_ex: List[str]):
 
     # Repeat for num_policy_iterations iterations.
     t1 = time.time()
-    while (llvm_trainer.global_step_numpy() <
-           num_policy_iterations * num_iterations):
+    while (llvm_trainer.global_step_numpy()
+           < num_policy_iterations * num_iterations):
       t2 = time.time()
       logging.info('Last iteration took: %f', t2 - t1)
       t1 = t2