# Path: mesh_transformer/train_actor.py
# class NetworkRunner(object):
#     def __init__(self, mesh_shape, network_builder):
#         self.mesh_shape = mesh_shape
#         self.network_builder = network_builder
# 
#         self.input_q = Queue(maxsize=1)
#         self.output_q = Queue(maxsize=1)
# 
#     def run(self):
#         print(f"jax runtime initialization starting")
#         import jax
#         from jax.experimental.maps import thread_resources, ResourceEnv, Mesh
#         import haiku as hk
#         # jax.experimental.maps.EXPERIMENTAL_SPMD_LOWERING = True
# 
#         thread_resources.env = ResourceEnv(Mesh(np.empty((), dtype=object), ()), ())
# 
#         start = time.time()
#         jax.devices()
# 
#         import warnings
#         warnings.filterwarnings("ignore")
#         warnings.filterwarnings("ignore", category=ResourceWarning)
# 
#         if jax.host_id() == 0:
#             warnings.filterwarnings("default")
# 
#         head_print(f"jax devices: {jax.device_count()}")
#         head_print(f"jax runtime initialized in {time.time() - start:.06}s")
#         devices = np.array(jax.devices()).reshape(self.mesh_shape)
# 
#         with jax.experimental.maps.mesh(devices, ('dp', 'mp')):
#             start = time.time()
#             network = self.network_builder()
#             head_print(f"Initialized in {time.time() - start:.06}s")
# 
#             while True:
#                 operation, input = self.input_q.get()
#                 if operation == "train":
#                     self.output_q.put(network.train(input))
#                 elif operation == "eval":
#                     self.output_q.put(network.eval(input))
#                 elif operation == "generate":
#                     self.output_q.put(network.generate(*input))
#                 elif operation == "write_ckpt":
#                     path, shard = input
#                     network.write_ckpt(path, shard)
#                     self.output_q.put(None)
#                 elif operation == "load_ckpt":
#                     network.load_ckpt(input)
#                     self.output_q.put(network.state["step"][0])
#                 elif operation == "get_params":
#                     self.output_q.put(hk.data_structures.tree_size(network.state['params']))
#                 elif operation == "move_params":
#                     # only needed for inference, otherwise first train step does this
#                     local_shards = max(jax.local_device_count() // self.mesh_shape[1], 1)
# 
#                     # delete the optimizer states otherwise it OOMs for some reason
#                     # TODO: use ShardedDeviceArray or something to get around this for bigger models
#                     del network.state["opt_state"]
#                     network.state = network.move_xmap(network.state, np.zeros(local_shards))
#                     self.output_q.put(None)
#                 else:
#                     raise Exception("Not implemented")
# 
#     def get_params(self):
#         self.input_q.put(("get_params", None))
#         return self.output_q.get()
# 
#     def train(self, sample):
#         self.input_q.put(("train", sample))
#         return self.output_q.get()
# 
#     def eval(self, sample):
#         self.input_q.put(("eval", sample))
#         return self.output_q.get()
# 
#     def generate(self, input):
#         self.input_q.put(("generate", input))
#         return self.output_q.get()
# 
#     def write_ckpt(self, path, shard):
#         self.input_q.put(("write_ckpt", (path, shard)))
#         return self.output_q.get()
# 
#     def load_ckpt(self, path):
#         self.input_q.put(("load_ckpt", path))
#         return self.output_q.get()
# 
#     def move_params(self):
#         self.input_q.put(("move_params", None))
#         return self.output_q.get()


import itertools
import json
import time
import ray
import numpy as np
from typing import Callable
from mesh_transformer.train_actor import NetworkRunner
from google.cloud import storage
from smart_open import open
from func_timeout import func_set_timeout
            for node in self.nodes:
                res.append(node.write_ckpt.remote(f"gs://{bucket}/{path}/step_{step}", 0))

        ray.get(res)
        print(f"Wrote checkpoint in {time.time() - start:.06}s")

        with open(f"gs://{bucket}/{path}/meta.json", "r") as f:
            meta = json.load(f)

        meta["step"] = step
        meta["checkpoints"].append(step)
        all_aux = meta.get("aux", {})

        while len(meta["checkpoints"]) > keep_n:
            ckpt_to_delete = meta["checkpoints"].pop(0)

            try:
                del all_aux[str(ckpt_to_delete)]
            except:
                print(f"failed to delete the aux state for {step}")

            if delete_old:
                print(f"deleting checkpoint {ckpt_to_delete}")
                for blob in client.list_blobs(bucket, prefix=f"{path}/step_{ckpt_to_delete}/"):
                    # print(f"deleting {blob.name}")
                    assert path in blob.name
                    blob.delete()
            else:
                print(f"keeping checkpoint {ckpt_to_delete}")


# Path: mesh_transformer/train_actor.py
# class NetworkRunner(object):
#     def __init__(self, mesh_shape, network_builder):
#         self.mesh_shape = mesh_shape
#         self.network_builder = network_builder
# 
#         self.input_q = Queue(maxsize=1)
#         self.output_q = Queue(maxsize=1)
# 
#     def run(self):
#         print(f"jax runtime initialization starting")
#         import jax
#         from jax.experimental.maps import thread_resources, ResourceEnv, Mesh
#         import haiku as hk
#         # jax.experimental.maps.EXPERIMENTAL_SPMD_LOWERING = True
# 
#         thread_resources.env = ResourceEnv(Mesh(np.empty((), dtype=object), ()), ())
# 
#         start = time.time()
#         jax.devices()
# 
#         import warnings
#         warnings.filterwarnings("ignore")
#         warnings.filterwarnings("ignore", category=ResourceWarning)
# 
#         if jax.host_id() == 0:
#             warnings.filterwarnings("default")
# 
#         head_print(f"jax devices: {jax.device_count()}")
#         head_print(f"jax runtime initialized in {time.time() - start:.06}s")
#         devices = np.array(jax.devices()).reshape(self.mesh_shape)
# 
#         with jax.experimental.maps.mesh(devices, ('dp', 'mp')):
#             start = time.time()
#             network = self.network_builder()
#             head_print(f"Initialized in {time.time() - start:.06}s")
# 
#             while True:
#                 operation, input = self.input_q.get()
#                 if operation == "train":
#                     self.output_q.put(network.train(input))
#                 elif operation == "eval":
#                     self.output_q.put(network.eval(input))
#                 elif operation == "generate":
#                     self.output_q.put(network.generate(*input))
#                 elif operation == "write_ckpt":
#                     path, shard = input
#                     network.write_ckpt(path, shard)
#                     self.output_q.put(None)
#                 elif operation == "load_ckpt":
#                     network.load_ckpt(input)
#                     self.output_q.put(network.state["step"][0])
#                 elif operation == "get_params":
#                     self.output_q.put(hk.data_structures.tree_size(network.state['params']))
#                 elif operation == "move_params":
#                     # only needed for inference, otherwise first train step does this
#                     local_shards = max(jax.local_device_count() // self.mesh_shape[1], 1)
# 
#                     # delete the optimizer states otherwise it OOMs for some reason
#                     # TODO: use ShardedDeviceArray or something to get around this for bigger models
#                     del network.state["opt_state"]
#                     network.state = network.move_xmap(network.state, np.zeros(local_shards))
#                     self.output_q.put(None)
#                 else:
#                     raise Exception("Not implemented")
# 
#     def get_params(self):
#         self.input_q.put(("get_params", None))
#         return self.output_q.get()
# 
#     def train(self, sample):
#         self.input_q.put(("train", sample))
#         return self.output_q.get()
# 
#     def eval(self, sample):
#         self.input_q.put(("eval", sample))
#         return self.output_q.get()
# 
#     def generate(self, input):
#         self.input_q.put(("generate", input))
#         return self.output_q.get()
# 
#     def write_ckpt(self, path, shard):
#         self.input_q.put(("write_ckpt", (path, shard)))
#         return self.output_q.get()
# 
#     def load_ckpt(self, path):
#         self.input_q.put(("load_ckpt", path))
#         return self.output_q.get()
# 
#     def move_params(self):
#         self.input_q.put(("move_params", None))
#         return self.output_q.get()


# Path: mesh_transformer/TPU_cluster.py
import itertools
import json
import time
import ray
import numpy as np
from typing import Callable
from mesh_transformer.train_actor import NetworkRunner
from google.cloud import storage
from smart_open import open
from func_timeout import func_set_timeout

            for node in self.nodes:
                res.append(node.write_ckpt.remote(f"gs://{bucket}/{path}/step_{step}", 0))

        ray.get(res)
        print(f"Wrote checkpoint in {time.time() - start:.06}s")

        with open(f"gs://{bucket}/{path}/meta.json", "r") as f:
            meta = json.load(f)

        meta["step"] = step
        meta["checkpoints"].append(step)
        all_aux = meta.get("aux", {})

        while len(meta["checkpoints"]) > keep_n:
            ckpt_to_delete = meta["checkpoints"].pop(0)

            try:
                del all_aux[str(ckpt_to_delete)]
            except:
                print(f"failed to delete the aux state for {step}")

            if delete_old:
                print(f"deleting checkpoint {ckpt_to_delete}")
                for blob in client.list_blobs(bucket, prefix=f"{path}/step_{ckpt_to_delete}/"):
                    # print(f"deleting {blob.name}")
                    assert path in blob.name
                    blob.delete()
            else:
                print(f"keeping checkpoint {ckpt_to_delete}")


        all_aux[step] = aux