def create_sstable_builder(
    result_dir: str, *, num_shards: int = 8, sort: bool = False
):
  """Creates a context manager for the sstable writer."""
  shard_spec = os.path.join(result_dir, f'p{jax.process_index()}@{num_shards}')
  shard_paths = gfile.GenerateShardedFilenames(shard_spec)
  # Remove partial files from previous run.
  for path in shard_paths:
    if gfile.Exists(path):
      gfile.Remove(path)
  sharder = sstable_sharder.SSTable_SharderFactory.New(
      'SSTable_KeyFingerprintSharder')
  if sort:
    return sharded_builder.SortingShardedBuilder(shard_paths, sharder)
  else:
    return sharded_builder.ShardedBuilder(shard_paths, sharder)
Creates a context manager for the sstable writer.
videogvt/train_lib/train_utils.py