Add coordinator.py for scheduling suggestions to workers

iris/algorithms/ars_algorithm.py

@@ -15,10 +15,8 @@
 """Algorithm class for Augmented Random Search Blackbox algorithm."""
 
 import collections
-import datetime
 import math
 import pathlib
-import pickle as pkl
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 from absl import logging
@@ -32,6 +30,7 @@
 import jax.numpy as jnp
 import numpy as np
 
+
 PRNGKey = jax.Array
 
 _DUMMY_REWARD = -1_000_000_000.0
@@ -540,13 +539,7 @@ def maybe_save_custom_checkpoint(self,
         per_agent_state["obs_norm_state"]["n"] = obs_norm_state["n"]
       agent_checkpoint_path = f"{checkpoint_path}_agent_{i}"
       logging.info("Saving agent checkpoints to %s...", agent_checkpoint_path)
-      self.save_checkpoint_internal(
-          agent_checkpoint_path, per_agent_state
-      self.save_checkpoint_oss(agent_checkpoint_path, per_agent_state)
-
-  def save_checkpoint_oss(self, checkpoint_path: str, state: Any) -> None:
-    with open(checkpoint_path, "wb") as f:
-      pkl.dump(state, f)
+      checkpoint_util.save_checkpoint(agent_checkpoint_path, per_agent_state)
 
   def split_and_save_checkpoint(self, checkpoint_path: str) -> None:
     state = checkpoint_util.load_checkpoint_state(checkpoint_path)

iris/algorithms/multi_agent_ars_algorithm_test.py

@@ -478,9 +478,7 @@ def test_split_checkpoint(self):
                 'n': 5,
             },
         }]
-    algo.save_checkpoint_internal(
-        full_path, state
-    algo.save_checkpoint_oss(full_path, state)
+    checkpoint_util.save_checkpoint(full_path, state)
     algo.split_and_save_checkpoint(checkpoint_path=full_path)
     for i in range(3):
       agent_checkpoint_path = f'{full_path}_agent_{i}'

iris/checkpoint_util.py

@@ -14,16 +14,28 @@
 
 """Utilities for interacting with BBv2 checkpoints."""
 
+import datetime
+import pathlib
 import pickle as pkl
 from typing import Any
 
 from absl import logging
 
 
-def load_checkpoint_state_oss(path: str) -> dict[str, Any]:
+def load_checkpoint_state_oss(path: str | pathlib.Path) -> dict[str, Any]:
   """Loads a checkpoint state from a given path."""
   with open(path, "rb") as f:
     state = pkl.load(f)
   return state
 
+
 load_checkpoint_state = load_checkpoint_state_oss
+
+
+def save_checkpoint_oss(path: str | pathlib.Path, state: dict[str, Any]):
+  """Saves a checkpoint state to a given path."""
+  with open(path, "wb") as f:
+    pkl.dump(state, f)
+
+
+save_checkpoint = save_checkpoint_oss