code: def save(
      self,
      step,
      bucket,
      path,
      aux=None,
      init=False,
      overwrite=False,
      keep_n=3,
      delete_old=True,
  ):
      assert path
      client = storage.Client()
  
      if aux is None:
          aux = {}
  
      if init:
          # check existing checkpoint folder does not exist, and delete it if it does
          for blob in client.list_blobs(bucket, prefix=f"{path}/"):
              assert overwrite
              # print(f"deleting {blob.name}")
              assert path in blob.name
              blob.delete()
  
          # create metadata file
          with open(f"gs://{bucket}/{path}/meta.json", "w") as f:
              json.dump({"step": 0, "checkpoints": [], "aux": {}}, f)
  
      # do sharded checkpoint writing
      start = time.time()
      res = []
  
      if self.version == 1:
          for shard_id, node in zip(
              range(self.mp), itertools.cycle(self.nodes)
          ):
              res.append(
                  node.write_ckpt.remote(
                      f"gs://{bucket}/{path}/step_{step}/", shard_id
                  )
              )
      elif self.version == 2:
          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
      meta["aux"] = all_aux
  
      with open(f"gs://{bucket}/{path}/meta.json", "w") as f:
          json.dump(meta, f)
  
file_path: mesh-transformer-jax/mesh_transformer/TPU_cluster.py
imports: import itertools
  import json
  import time
  import ray
  from typing import Callable
  import numpy
  from mesh_transformer.train_actor import NetworkRunner
  from google.cloud import storage
  from smart_open import open
  from func_timeout import func_set_timeout
label_desc: Use image node_id 7 for calling the TPUCluster obj's underlying member method code with example usage: obj.save(step, bucket, path, aux, init, overwrite, keep_n, delete_old) without return types