code: def write_ckpt_v2(model_state, dir):
      start = time.time()
      if jax.host_id() == 0:
          param_map = tree_leaves_with_names(model_state["params"])
          opt_map = tree_leaves_with_names(model_state["opt_state"])
  
          meta = {
              "total_hosts": jax.host_count(),
              "step": int(model_state["step"]),
              "param_order": [
                  param_map[id(i)]
                  for i in jax.tree_leaves(model_state["params"])
              ],
              "opt_order": [
                  opt_map[id(i)]
                  for i in jax.tree_leaves(model_state["opt_state"])
              ],
          }
  
          print("step:", model_state["step"])
          with open(dir + "/meta.json", "w") as f:
              json.dump(meta, f)
          print(f"meta written in {time.time() - start:.06}s")
  
      start = time.time()
      parallel_write(
          jax.tree_flatten(model_state["params"])[0],
          dir + f"/params/shard_{jax.host_id()}.npz",
      )
      head_print(f"params written in {time.time() - start:.06}s")
  
      start = time.time()
      parallel_write(
          jax.tree_flatten(model_state["opt_state"])[0],
          dir + f"/opt_state/shard_{jax.host_id()}.npz",
      )
      head_print(f"opt_state written in {time.time() - start:.06}s")
  
imports: import functools
  import io
  import json
  import time
  import jax
  import jax.numpy
  import numpy
  import multiprocessing
  import ray
  from smart_open import open
  from mesh_transformer.util import head_print
label_desc: Use image node_id 14 for calling a global function with example usage: write_ckpt_v2(model_state, dir) without return types