fix regalloc priority prediction lint (#159)

* lint

* more fix

Author: yundiqian

compiler_opt/rl/regalloc_priority/__init__.py

@@ -1,3 +1,18 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Implementation of the 'regalloc priority prediction' problem."""
 
 import gin
 
@@ -8,11 +23,12 @@
 
 @gin.register(module='configs')
 class RegallocPriorityConfig(problem_configuration.ProblemConfiguration):
-    def get_runner(self, *args, **kwargs):
-        return regalloc_priority_runner.RegAllocPriorityRunner(*args, **kwargs)
 
-    def get_signature_spec(self):
-        return config.get_regalloc_signature_spec()
+  def get_runner(self, *args, **kwargs):
+    return regalloc_priority_runner.RegAllocPriorityRunner(*args, **kwargs)
 
-    def get_preprocessing_layer_creator(self):
-        return config.get_observation_processing_layer_creator()
+  def get_signature_spec(self):
+    return config.get_regalloc_signature_spec()
+
+  def get_preprocessing_layer_creator(self):
+    return config.get_observation_processing_layer_creator()

compiler_opt/rl/regalloc_priority/config.py

@@ -1,3 +1,19 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""RegAlloc priority prediction training config."""
+
 import gin
 import tensorflow as tf
 from tf_agents.specs import tensor_spec
@@ -7,32 +23,31 @@
 
 @gin.configurable()
 def get_regalloc_signature_spec():
-    observation_spec = dict(
-        (key, tf.TensorSpec(dtype=tf.int64, shape=(), name=key))
-        for key in ('li_size', 'stage'))
-    observation_spec['weight'] = tf.TensorSpec(dtype=tf.float32, shape=(), name='weight')
+  observation_spec = dict(
+      (key, tf.TensorSpec(dtype=tf.int64, shape=(), name=key))
+      for key in ('li_size', 'stage'))
+  observation_spec['weight'] = tf.TensorSpec(
+      dtype=tf.float32, shape=(), name='weight')
 
-    reward_spec = tf.TensorSpec(dtype=tf.float32, shape=(), name='reward')
-    time_step_spec = time_step.time_step_spec(observation_spec, reward_spec)
+  reward_spec = tf.TensorSpec(dtype=tf.float32, shape=(), name='reward')
+  time_step_spec = time_step.time_step_spec(observation_spec, reward_spec)
 
-    action_spec = tensor_spec.TensorSpec(
-        dtype=tf.float32,
-        shape=(),
-        name='priority'
-    )
+  action_spec = tensor_spec.TensorSpec(
+      dtype=tf.float32, shape=(), name='priority')
 
-    return time_step_spec, action_spec
+  return time_step_spec, action_spec
 
 
 @gin.configurable
 def get_observation_processing_layer_creator():
-    def observation_processing_layer(obs_spec):
-        """Creates the layer to process observation given obs_spec."""
 
-        if obs_spec.name in ('li_size', 'stage', 'weight'):
-            return tf.keras.layers.Lambda(feature_ops.identity_fn)
+  def observation_processing_layer(obs_spec):
+    """Creates the layer to process observation given obs_spec."""
+
+    if obs_spec.name in ('li_size', 'stage', 'weight'):
+      return tf.keras.layers.Lambda(feature_ops.identity_fn)
 
-        # Make sure all features have a preprocessing function.
-        raise KeyError('Missing preprocessing function for some feature.')
+    # Make sure all features have a preprocessing function.
+    raise KeyError('Missing preprocessing function for some feature.')
 
-    return observation_processing_layer
+  return observation_processing_layer

compiler_opt/rl/regalloc_priority/regalloc_priority_runner.py

@@ -1,3 +1,19 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Module for collect data of regalloc priority prediction."""
+
 import gin
 import tensorflow as tf
 
@@ -6,65 +22,67 @@
 import os
 import tempfile
 from typing import Dict, Optional, Tuple
-from absl import logging
 
-from google.protobuf import struct_pb2
+# This is https://github.com/google/pytype/issues/764
+from google.protobuf import struct_pb2  # pytype: disable=pyi-error
 from compiler_opt.rl import compilation_runner
 
 
 @gin.configurable(module='runners')
 class RegAllocPriorityRunner(compilation_runner.CompilationRunner):
-    def _compile_fn(
-            self, file_paths: Tuple[str, ...], tf_policy_path: str, reward_only: bool,
-            cancellation_manager: Optional[
-                compilation_runner.WorkerCancellationManager]
-    ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
+  """Class for collecting data for regalloc-priority-prediction."""
+
+  def _compile_fn(
+      self, file_paths: Tuple[str, ...], tf_policy_path: str, reward_only: bool,
+      cancellation_manager: Optional[
+          compilation_runner.WorkerCancellationManager]
+  ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
 
-        file_paths = file_paths[0].replace('.bc', '')
-        working_dir = tempfile.mkdtemp()
+    file_paths = file_paths[0].replace('.bc', '')
+    working_dir = tempfile.mkdtemp()
 
-        log_path = os.path.join(working_dir, 'log')
-        output_native_path = os.path.join(working_dir, 'native')
+    log_path = os.path.join(working_dir, 'log')
+    output_native_path = os.path.join(working_dir, 'native')
 
-        result = {}
-        try:
-            command_line = []
-            if self._launcher_path:
-                command_line.append(self._launcher_path)
-            command_line.extend([self._clang_path] + [
-                                    '-c', file_paths, '-O3',
-                                    '-mllvm', '-regalloc-priority-training-log=' + log_path,
-                                    '-mllvm', '-regalloc-enable-priority-advisor=development',
-                                    '-o', output_native_path
-                                ])
+    result = {}
+    try:
+      command_line = []
+      if self._launcher_path:
+        command_line.append(self._launcher_path)
+      command_line.extend([self._clang_path] + [
+          '-c', file_paths, '-O3', '-mllvm', '-regalloc-priority-training-log='
+          + log_path, '-mllvm', '-regalloc-enable-priority-advisor=development',
+          '-o', output_native_path
+      ])
 
-            if tf_policy_path:
-                command_line.extend(['-mllvm', '-regalloc-priority-model=' + tf_policy_path])
-            compilation_runner.start_cancellable_process(command_line,
-                                                         self._compilation_timeout,
-                                                         cancellation_manager)
+      if tf_policy_path:
+        command_line.extend(
+            ['-mllvm', '-regalloc-priority-model=' + tf_policy_path])
+      compilation_runner.start_cancellable_process(command_line,
+                                                   self._compilation_timeout,
+                                                   cancellation_manager)
 
-            sequence_example = struct_pb2.Struct()
+      sequence_example = struct_pb2.Struct()
 
-            with io.open(log_path, 'rb') as f:
-                sequence_example.ParseFromString(f.read())
+      with io.open(log_path, 'rb') as f:
+        sequence_example.ParseFromString(f.read())
 
-            for key, value in sequence_example.fields.items():
-                e = tf.train.SequenceExample()
-                e.ParseFromString(base64.b64decode(value.string_value))
-                print(e)
-                if not e.HasField('feature_lists'):
-                    continue
-                r = (
-                    e.feature_lists.feature_list['reward'].feature[-1].float_list
-                        .value[0])
-                if reward_only:
-                    result[key] = (None, r)
-                else:
-                    del e.feature_lists.feature_list['reward']
-                    result[key] = (e, r)
+      for key, value in sequence_example.fields.items():
+        e = tf.train.SequenceExample()
+        e.ParseFromString(base64.b64decode(value.string_value))
+        print(e)
+        if not e.HasField('feature_lists'):
+          continue
+        r = (
+            e.feature_lists.feature_list['reward'].feature[-1].float_list
+            .value[0])
+        if reward_only:
+          result[key] = (None, r)
+        else:
+          del e.feature_lists.feature_list['reward']
+          result[key] = (e, r)
 
-        finally:
-            tf.io.gfile.rmtree(working_dir)
+    finally:
+      tf.io.gfile.rmtree(working_dir)
 
-        return result
+    return result