def train(self):
        experiment_fn = self._generate_experiment_fn()
        hparams = HParams(**self.customer_params)

        learn_runner.run(experiment_fn,
                         run_config=self._build_run_config(),
                         hparams=hparams)

 def test_fail_invalid_hparams_type(self):
   run_config = run_config_lib.RunConfig(model_dir=_MODIR_DIR)
   with self.assertRaisesRegexp(ValueError, _INVALID_HPARAMS_ERR_MSG):
     learn_runner.run(build_experiment_for_run_config,
                      run_config=run_config,
                      schedule="local_run",
                      hparams=["hparams"])

 def test_fail_output_dir_and_run_config_are_both_set(self):
   with self.assertRaisesRegexp(
       ValueError, _CANNOT_SET_BOTH_OUTPUT_DIR_AND_CONFIG_MSG):
     learn_runner.run(build_experiment,
                      output_dir=_MODIR_DIR,
                      schedule="simple_task",
                      run_config=run_config_lib.RunConfig())

def main(argv=None):
  """Run a Tensorflow model on the Criteo dataset."""
  env = json.loads(os.environ.get('TF_CONFIG', '{}'))
  # First find out if there's a task value on the environment variable.
  # If there is none or it is empty define a default one.
  task_data = env.get('task') or {'type': 'master', 'index': 0}
  argv = sys.argv if argv is None else argv
  args = create_parser().parse_args(args=argv[1:])

  trial = task_data.get('trial')
  if trial is not None:
    output_dir = os.path.join(args.output_path, trial)
  else:
    output_dir = args.output_path

  # Do only evaluation if instructed so, or call Experiment's run.
  if args.eval_only_summary_filename:
    experiment = get_experiment_fn(args)(output_dir)
    # Note that evaluation here will appear as 'one_pass' in tensorboard.
    results = experiment.evaluate(delay_secs=0)
    # Converts numpy types to native types for json dumps.
    json_out = json.dumps(
        {key: value.tolist() for key, value in results.iteritems()})
    with tf.Session():
      tf.write_file(args.eval_only_summary_filename, json_out).run()
  else:
    learn_runner.run(experiment_fn=get_experiment_fn(args),
                     output_dir=output_dir)

def main(argv=None):
  args = parse_arguments(sys.argv if argv is None else argv)

  tf.logging.set_verbosity(tf.logging.INFO)
  learn_runner.run(
      experiment_fn=get_experiment_fn(args),
      output_dir=args.job_dir)

def run(data_dir, model, output_dir, train_steps, eval_steps, schedule):
  """Runs an Estimator locally or distributed.

  Args:
    data_dir: The directory the data can be found in.
    model: The name of the model to use.
    output_dir: The directory to store outputs in.
    train_steps: The number of steps to run training for.
    eval_steps: The number of steps to run evaluation for.
    schedule: (str) The schedule to run. The value here must
      be the name of one of Experiment's methods.
  """
  exp_fn = make_experiment_fn(
      data_dir=data_dir,
      model_name=model,
      train_steps=train_steps,
      eval_steps=eval_steps)

  # Create hparams and run_config
  run_config = create_run_config(output_dir)
  hparams = create_hparams(
      FLAGS.hparams_set, data_dir, passed_hparams=FLAGS.hparams)

  if is_chief():
    save_metadata(output_dir, hparams)

  learn_runner.run(
      experiment_fn=exp_fn,
      schedule=schedule,
      run_config=run_config,
      hparams=hparams)

def main(argv=None):
  """Run a Tensorflow model on the Iris dataset."""
  args = parse_arguments(sys.argv if argv is None else argv)

  tf.logging.set_verbosity(tf.logging.INFO)
  learn_runner.run(
      experiment_fn=get_experiment_fn(args),
      output_dir=args.job_dir)

def main(unused_argv):
  tf.flags.mark_flag_as_required('model_dir')
  tf.flags.mark_flag_as_required('pipeline_config_path')
  config = tf.contrib.learn.RunConfig(model_dir=FLAGS.model_dir)
  learn_runner.run(
      experiment_fn=build_experiment_fn(FLAGS.num_train_steps,
                                        FLAGS.num_eval_steps),
      run_config=config,
      hparams=model_hparams.create_hparams())

  def test_fail_not_experiment(self):
    def _experiment_fn(run_config, hparams):
      del run_config, hparams  # unused.
      return "not experiment"

    run_config = run_config_lib.RunConfig(model_dir=_MODIR_DIR)
    with self.assertRaisesRegexp(TypeError, _NOT_EXP_TYPE_MSG):
      learn_runner.run(_experiment_fn,
                       run_config=run_config,
                       schedule="simple_task")

def main(argv=None):
  args = parse_arguments(sys.argv if argv is None else argv)
  local_analysis(args)
  set_logging_level(args)
  # Supress TensorFlow Debugging info.
  os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

  learn_runner.run(
      experiment_fn=get_experiment_fn(args),
      output_dir=args.job_dir)

def main(_argv):
  """The entrypoint for the script"""

  # Parse YAML FLAGS
  FLAGS.hooks = _maybe_load_yaml(FLAGS.hooks)
  FLAGS.metrics = _maybe_load_yaml(FLAGS.metrics)
  FLAGS.model_params = _maybe_load_yaml(FLAGS.model_params)
  FLAGS.input_pipeline_train = _maybe_load_yaml(FLAGS.input_pipeline_train)
  FLAGS.input_pipeline_dev = _maybe_load_yaml(FLAGS.input_pipeline_dev)

  # Load flags from config file
  final_config = {}
  if FLAGS.config_paths:
    for config_path in FLAGS.config_paths.split(","):
      config_path = config_path.strip()
      if not config_path:
        continue
      config_path = os.path.abspath(config_path)
      tf.logging.info("Loading config from %s", config_path)
      with gfile.GFile(config_path.strip()) as config_file:
        config_flags = yaml.load(config_file)
        final_config = _deep_merge_dict(final_config, config_flags)

  tf.logging.info("Final Config:\n%s", yaml.dump(final_config))

  # Merge flags with config values
  for flag_key, flag_value in final_config.items():
    if hasattr(FLAGS, flag_key) and isinstance(getattr(FLAGS, flag_key), dict):
      merged_value = _deep_merge_dict(flag_value, getattr(FLAGS, flag_key))
      setattr(FLAGS, flag_key, merged_value)
    elif hasattr(FLAGS, flag_key):
      setattr(FLAGS, flag_key, flag_value)
    else:
      tf.logging.warning("Ignoring config flag: %s", flag_key)

  if FLAGS.save_checkpoints_secs is None \
    and FLAGS.save_checkpoints_steps is None:
    FLAGS.save_checkpoints_secs = 600
    tf.logging.info("Setting save_checkpoints_secs to %d",
                    FLAGS.save_checkpoints_secs)

  if not FLAGS.output_dir:
    FLAGS.output_dir = tempfile.mkdtemp()

  if not FLAGS.input_pipeline_train:
    raise ValueError("You must specify input_pipeline_train")

  if not FLAGS.input_pipeline_dev:
    raise ValueError("You must specify input_pipeline_dev")

  learn_runner.run(
      experiment_fn=create_experiment,
      output_dir=FLAGS.output_dir,
      schedule=FLAGS.schedule)

  def test_basic_run_config_uid_check(self):
    expected_run_config = run_config_lib.RunConfig(model_dir=_MODIR_DIR)

    def _experiment_fn(run_config, hparams):
      del run_config, hparams  # unused.
      # Explicitly use a new run_config.
      new_config = run_config_lib.RunConfig(model_dir=_MODIR_DIR + "/123")

      return TestExperiment(config=new_config)

    with self.assertRaisesRegexp(RuntimeError, _RUN_CONFIG_UID_CHECK_ERR_MSG):
      learn_runner.run(experiment_fn=_experiment_fn,
                       run_config=expected_run_config)

def main(args):
  env = json.loads(os.environ.get('TF_CONFIG', '{}'))

  # Print the job data as provided by the service.
  logging.info('Original job data: %s', env.get('job', {}))

  # First find out if there's a task value on the environment variable.
  # If there is none or it is empty define a default one.
  task_data = env.get('task', {'type': 'master', 'index': 0})
  trial = task_data.get('trial')
  if trial is not None:
    args.output_path = os.path.join(args.output_path, trial)

  learn_runner.run(make_experiment_fn(args), args.output_path)

def main(argv):
    parser = argparse.ArgumentParser()
    # You must accept a --job-dir argument when running on Cloud ML Engine. It specifies where checkpoints should be saved.
    # You can define additional user arguments which will have to be specified after an empty arg -- on the command line:
    # gcloud ml-engine jobs submit training jobXXX --job-dir=... --ml-engine-args -- --user-args
    parser.add_argument('--job-dir', default="checkpoints", help='GCS or local path where to store training checkpoints')
    args = parser.parse_args()
    arguments = args.__dict__
    arguments['data'] = "data" # Hard-coded here: training data will be downloaded to folder 'data'.

    # learn_runner needs an experiment function with a single parameter: the output directory.
    # Here we pass additional command line arguments through a closure.
    output_dir = arguments.pop('job_dir')
    experiment_fn = lambda output_dir: experiment_fn_with_params(output_dir, **arguments)
    learn_runner.run(experiment_fn, output_dir)

  def test_fail_invalid_experiment_config_type(self):
    expected_run_config = run_config_lib.RunConfig(model_dir=_MODIR_DIR)

    def _experiment_fn(run_config, hparams):
      del run_config, hparams  # unused.
      # Explicitly use a new run_config without `uid` method.
      new_config = core_run_config_lib.RunConfig(
          model_dir=_MODIR_DIR + "/123")

      return TestExperiment(config=new_config)

    with self.assertRaisesRegexp(RuntimeError,
                                 _MISSING_RUN_CONFIG_UID_ERR_MSG):
      learn_runner.run(experiment_fn=_experiment_fn,
                       run_config=expected_run_config)

 def test_schedule_from_config_runs_train_on_worker(self):
   config = run_config.RunConfig(
       job_name="worker", cluster_spec=build_distributed_cluster_spec())
   self.assertEqual(
       "train",
       learn_runner.run(lambda output_dir: TestExperiment(config=config),
                        output_dir="/tmp"))

 def test_no_schedule_and_non_distributed_runs_train_and_evaluate(self):
   config = run_config.RunConfig(
       cluster_spec=build_non_distributed_cluster_spec())
   self.assertEqual(
       "train_and_evaluate",
       learn_runner.run(lambda output_dir: TestExperiment(config=config),
                        output_dir="/tmp"))

 def test_schedule_from_tf_config_runs_train_on_worker(self):
     os.environ["TF_CONFIG"] = json.dumps(
         {"cluster": build_distributed_cluster_spec(), "task": {"type": tf.contrib.learn.TaskType.WORKER}}
     )
     # RunConfig constructor will set job_name from TF_CONFIG.
     config = run_config.RunConfig()
     self.assertEqual("train", learn_runner.run(lambda output_dir: TestExperiment(config=config), output_dir="/tmp"))

 def test_schedule_from_tf_config_runs_serve_on_ps(self):
     tf_config = {"cluster": build_distributed_cluster_spec(), "task": {"type": tf.contrib.learn.TaskType.PS}}
     with patch.dict("os.environ", {"TF_CONFIG": json.dumps(tf_config)}):
         config = run_config.RunConfig()
         self.assertEqual(
             "run_std_server", learn_runner.run(lambda output_dir: TestExperiment(config=config), output_dir="/tmp")
         )

 def test_run_with_custom_schedule(self):
   run_config = run_config_lib.RunConfig(model_dir=_MODIR_DIR)
   self.assertEqual(
       "simple_task, default=None.",
       learn_runner.run(build_experiment_for_run_config,
                        run_config=run_config,
                        schedule="simple_task"))