docs: def batch_predict_async(training_dir, prediction_input_file, output_dir,
                          mode, batch_size=16, shard_files=True, output_format='csv', cloud=False):
    """Local and cloud batch prediction.
  
    Args:
      training_dir: The output folder of training.
      prediction_input_file: csv file pattern to a file. File must be on GCS if
          running cloud prediction
      output_dir: output location to save the results. Must be a GSC path if
          running cloud prediction.
      mode: 'evaluation' or 'prediction'. If 'evaluation', the input data must
          contain a target column. If 'prediction', the input data must not
          contain a target column.
      batch_size: Int. How many instances to run in memory at once. Larger values
          mean better performace but more memeory consumed.
      shard_files: If False, the output files are not shardded.
      output_format: csv or json. Json file are json-newlined.
      cloud: If ture, does cloud batch prediction. If False, runs batch prediction
          locally.
  
    Returns:
      A google.datalab.utils.Job object that can be used to query state from or wait.
    """
    import google.datalab.utils as du
    with warnings.catch_warnings():
      warnings.simplefilter("ignore")
      if cloud:
        runner_results = cloud_batch_predict(training_dir, prediction_input_file, output_dir, mode,
                                             batch_size, shard_files, output_format)
        job = du.DataflowJob(runner_results)
      else:
        runner_results = local_batch_predict(training_dir, prediction_input_file, output_dir, mode,
                                             batch_size, shard_files, output_format)
        job = du.LambdaJob(lambda: runner_results.wait_until_finish(), job_id=None)
  
    return job