move DataClassJSONEncoder to a more common place for future reuse (#144)

* test model test

* yapf

* pylint

* move DataClassJSONEncoder to a common place so that we can re-use it

* yapf

Author: yundiqian

compiler_opt/rl/compilation_runner.py

@@ -16,7 +16,6 @@
 
 import abc
 import dataclasses
-import json
 import os
 import signal
 import subprocess
@@ -50,14 +49,6 @@ class RewardStat:
   moving_average_reward: float
 
 
-class DataClassJSONEncoder(json.JSONEncoder):
-
-  def default(self, o):
-    if dataclasses.is_dataclass(o):
-      return dataclasses.asdict(o)
-    return super().default(o)
-
-
 def _overwrite_trajectory_reward(sequence_example: tf.train.SequenceExample,
                                  reward: float) -> tf.train.SequenceExample:
   """Overwrite the reward in the trace (sequence_example) with the given one.

compiler_opt/rl/constant.py

@@ -14,8 +14,10 @@
 # limitations under the License.
 """Constants for policy training."""
 
+import dataclasses
 import enum
 import gin
+import json
 
 BASE_DIR = 'compiler_opt/rl'
 BASE_MODULE_DIR = 'compiler_opt.rl'
@@ -33,3 +35,11 @@ class AgentName(enum.Enum):
   BEHAVIORAL_CLONE = 0
   DQN = 1
   PPO = 2
+
+
+class DataClassJSONEncoder(json.JSONEncoder):
+
+  def default(self, o):
+    if dataclasses.is_dataclass(o):
+      return dataclasses.asdict(o)
+    return super().default(o)

compiler_opt/rl/train_locally.py

@@ -152,8 +152,7 @@ def sequence_example_iterator_fn(seq_ex: List[str]):
       logging.info('Last iteration took: %f', t2 - t1)
       t1 = t2
       with tf.io.gfile.GFile(reward_stat_map_path, 'w') as f:
-        json.dump(
-            reward_stat_map, f, cls=compilation_runner.DataClassJSONEncoder)
+        json.dump(reward_stat_map, f, cls=constant.DataClassJSONEncoder)
 
       policy_path = os.path.join(root_dir, 'policy',
                                  str(llvm_trainer.global_step_numpy()))