def run(iterative_process: tff.templates.IterativeProcess, client_datasets_fn: Callable[([int], List[tf.data.Dataset])], validation_fn: Callable[([Any, int], Dict[(str, float)])], total_rounds: int, experiment_name: str, test_fn: Optional[Callable[([Any], Dict[(str, float)])]]=None, root_output_dir: Optional[str]='/tmp/fed_opt', rounds_per_eval: Optional[int]=1, rounds_per_checkpoint: Optional[int]=50, rounds_per_profile: Optional[int]=0):
    'Runs federated training for a given `tff.templates.IterativeProcess`.\n\n  We assume that the iterative process has the following functional type\n  signatures:\n\n    *   `initialize`: `( -> S@SERVER)` where `S` represents the server state.\n    *   `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S`\n        represents the server state, `{B*}` represents the client datasets,\n        and `T` represents a python `Mapping` object.\n\n  The iterative process must also have a callable attribute `get_model_weights`\n  that takes as input the state of the iterative process, and returns a\n  `tff.learning.ModelWeights` object.\n\n  Args:\n    iterative_process: A `tff.templates.IterativeProcess` instance to run.\n    client_datasets_fn: Function accepting an integer argument (the round\n      number) and returning a list of client datasets to use as federated data\n      for that round.\n    validation_fn: A callable accepting a `tff.learning.ModelWeights` and the\n      current round number, and returning a dict of evaluation metrics. Used to\n      compute validation metrics throughout the training process.\n    total_rounds: The number of federated training rounds to perform.\n    experiment_name: The name of the experiment being run. This will be appended\n      to the `root_output_dir` for purposes of writing outputs.\n    test_fn: An optional callable accepting a `tff.learning.ModelWeights` and\n      returning a dict of test set metrics. Used to compute test metrics at the\n      end of the training process.\n    root_output_dir: The name of the root output directory for writing\n      experiment outputs.\n    rounds_per_eval: How often to compute validation metrics.\n    rounds_per_checkpoint: How often to checkpoint the iterative process state.\n      If you expect the job to restart frequently, this should be small. If no\n      interruptions are expected, this can be made larger.\n    rounds_per_profile: Experimental setting. If set to a value greater than 0,\n      this dictates how often a TensorFlow profiler is run.\n\n  Returns:\n    The final `state` of the iterative process after training.\n  '
    _check_iterative_process_compatibility(iterative_process)
    if (not callable(client_datasets_fn)):
        raise TypeError('client_datasets_fn should be callable.')
    if (not callable(validation_fn)):
        raise TypeError('validation_fn should be callable.')
    if ((test_fn is not None) and (not callable(test_fn))):
        raise TypeError('test_fn should be callable.')
    logging.info('Starting iterative_process training loop...')
    initial_state = iterative_process.initialize()
    (checkpoint_mngr, metrics_mngr, tb_mngr, profiler) = _setup_outputs(root_output_dir, experiment_name, rounds_per_profile)
    logging.info('Asking checkpoint manager to load checkpoint.')
    (state, round_num) = checkpoint_mngr.load_latest_checkpoint(initial_state)
    if (state is None):
        logging.info('Initializing experiment from scratch.')
        state = initial_state
        round_num = 0
    else:
        logging.info('Restarted from checkpoint round %d', round_num)
        round_num += 1
    metrics_mngr.clear_metrics(round_num)
    current_model = iterative_process.get_model_weights(state)
    loop_start_time = time.time()
    loop_start_round = round_num
    while (round_num < total_rounds):
        data_prep_start_time = time.time()
        federated_train_data = client_datasets_fn(round_num)
        train_metrics = {'prepare_datasets_secs': (time.time() - data_prep_start_time)}
        training_start_time = time.time()
        prev_model = current_model
        try:
            with profiler(round_num):
                (state, round_metrics) = iterative_process.next(state, federated_train_data)
        except (tf.errors.FailedPreconditionError, tf.errors.NotFoundError, tf.errors.InternalError) as e:
            logging.warning('Caught %s exception while running round %d:\n\t%s', type(e), round_num, e)
            continue
        current_model = iterative_process.get_model_weights(state)
        train_metrics['training_secs'] = (time.time() - training_start_time)
        train_metrics['model_delta_l2_norm'] = _compute_numpy_l2_difference(current_model, prev_model)
        train_metrics['client_drift'] = state.client_drift
        train_metrics.update(round_metrics)
        loop_time = (time.time() - loop_start_time)
        loop_rounds = ((round_num - loop_start_round) + 1)
        logging.info('Round {:2d}, {:.2f}s per round in average.'.format(round_num, (loop_time / loop_rounds)))
        if (((round_num % rounds_per_checkpoint) == 0) or (round_num == (total_rounds - 1))):
            save_checkpoint_start_time = time.time()
            checkpoint_mngr.save_checkpoint(state, round_num)
            train_metrics['save_checkpoint_secs'] = (time.time() - save_checkpoint_start_time)
        metrics = {'train': train_metrics}
        if ((round_num % rounds_per_eval) == 0):
            evaluate_start_time = time.time()
            validation_metrics = validation_fn(current_model, round_num)
            validation_metrics['evaluate_secs'] = (time.time() - evaluate_start_time)
            metrics['eval'] = validation_metrics
        _write_metrics(metrics_mngr, tb_mngr, metrics, round_num)
        round_num += 1
    metrics = {}
    evaluate_start_time = time.time()
    validation_metrics = validation_fn(current_model, round_num)
    validation_metrics['evaluate_secs'] = (time.time() - evaluate_start_time)
    metrics['eval'] = validation_metrics
    if test_fn:
        test_start_time = time.time()
        test_metrics = test_fn(current_model)
        test_metrics['evaluate_secs'] = (time.time() - test_start_time)
        metrics['test'] = test_metrics
    _write_metrics(metrics_mngr, tb_mngr, metrics, total_rounds)
    return state
Runs federated training for a given `tff.templates.IterativeProcess`.

We assume that the iterative process has the following functional type
signatures:

  *   `initialize`: `( -> S@SERVER)` where `S` represents the server state.
  *   `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S`
      represents the server state, `{B*}` represents the client datasets,
      and `T` represents a python `Mapping` object.

The iterative process must also have a callable attribute `get_model_weights`
that takes as input the state of the iterative process, and returns a
`tff.learning.ModelWeights` object.

Args:
  iterative_process: A `tff.templates.IterativeProcess` instance to run.
  client_datasets_fn: Function accepting an integer argument (the round
    number) and returning a list of client datasets to use as federated data
    for that round.
  validation_fn: A callable accepting a `tff.learning.ModelWeights` and the
    current round number, and returning a dict of evaluation metrics. Used to
    compute validation metrics throughout the training process.
  total_rounds: The number of federated training rounds to perform.
  experiment_name: The name of the experiment being run. This will be appended
    to the `root_output_dir` for purposes of writing outputs.
  test_fn: An optional callable accepting a `tff.learning.ModelWeights` and
    returning a dict of test set metrics. Used to compute test metrics at the
    end of the training process.
  root_output_dir: The name of the root output directory for writing
    experiment outputs.
  rounds_per_eval: How often to compute validation metrics.
  rounds_per_checkpoint: How often to checkpoint the iterative process state.
    If you expect the job to restart frequently, this should be small. If no
    interruptions are expected, this can be made larger.
  rounds_per_profile: Experimental setting. If set to a value greater than 0,
    this dictates how often a TensorFlow profiler is run.

Returns:
  The final `state` of the iterative process after training.
def run(iterative_process: tff.templates.IterativeProcess, client_datasets_fn: Callable[([int], List[tf.data.Dataset])], validation_fn: Callable[([Any, int], Dict[(str, float)])], total_rounds: int, experiment_name: str, test_fn: Optional[Callable[([Any], Dict[(str, float)])]]=None, root_output_dir: Optional[str]='/tmp/fed_opt', rounds_per_eval: Optional[int]=1, rounds_per_checkpoint: Optional[int]=50, rounds_per_profile: Optional[int]=0):
    'Runs federated training for a given `tff.templates.IterativeProcess`.\n\n  We assume that the iterative process has the following functional type\n  signatures:\n\n    *   `initialize`: `( -> S@SERVER)` where `S` represents the server state.\n    *   `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S`\n        represents the server state, `{B*}` represents the client datasets,\n        and `T` represents a python `Mapping` object.\n\n  The iterative process must also have a callable attribute `get_model_weights`\n  that takes as input the state of the iterative process, and returns a\n  `tff.learning.ModelWeights` object.\n\n  Args:\n    iterative_process: A `tff.templates.IterativeProcess` instance to run.\n    client_datasets_fn: Function accepting an integer argument (the round\n      number) and returning a list of client datasets to use as federated data\n      for that round.\n    validation_fn: A callable accepting a `tff.learning.ModelWeights` and the\n      current round number, and returning a dict of evaluation metrics. Used to\n      compute validation metrics throughout the training process.\n    total_rounds: The number of federated training rounds to perform.\n    experiment_name: The name of the experiment being run. This will be appended\n      to the `root_output_dir` for purposes of writing outputs.\n    test_fn: An optional callable accepting a `tff.learning.ModelWeights` and\n      returning a dict of test set metrics. Used to compute test metrics at the\n      end of the training process.\n    root_output_dir: The name of the root output directory for writing\n      experiment outputs.\n    rounds_per_eval: How often to compute validation metrics.\n    rounds_per_checkpoint: How often to checkpoint the iterative process state.\n      If you expect the job to restart frequently, this should be small. If no\n      interruptions are expected, this can be made larger.\n    rounds_per_profile: Experimental setting. If set to a value greater than 0,\n      this dictates how often a TensorFlow profiler is run.\n\n  Returns:\n    The final `state` of the iterative process after training.\n  '
    _check_iterative_process_compatibility(iterative_process)
    if (not callable(client_datasets_fn)):
        raise TypeError('client_datasets_fn should be callable.')
    if (not callable(validation_fn)):
        raise TypeError('validation_fn should be callable.')
    if ((test_fn is not None) and (not callable(test_fn))):
        raise TypeError('test_fn should be callable.')
    logging.info('Starting iterative_process training loop...')
    initial_state = iterative_process.initialize()
    (checkpoint_mngr, metrics_mngr, tb_mngr, profiler) = _setup_outputs(root_output_dir, experiment_name, rounds_per_profile)
    logging.info('Asking checkpoint manager to load checkpoint.')
    (state, round_num) = checkpoint_mngr.load_latest_checkpoint(initial_state)
    if (state is None):
        logging.info('Initializing experiment from scratch.')
        state = initial_state
        round_num = 0
    else:
        logging.info('Restarted from checkpoint round %d', round_num)
        round_num += 1
    metrics_mngr.clear_metrics(round_num)
    current_model = iterative_process.get_model_weights(state)
    loop_start_time = time.time()
    loop_start_round = round_num
    while (round_num < total_rounds):
        data_prep_start_time = time.time()
        federated_train_data = client_datasets_fn(round_num)
        train_metrics = {'prepare_datasets_secs': (time.time() - data_prep_start_time)}
        training_start_time = time.time()
        prev_model = current_model
        try:
            with profiler(round_num):
                (state, round_metrics) = iterative_process.next(state, federated_train_data)
        except (tf.errors.FailedPreconditionError, tf.errors.NotFoundError, tf.errors.InternalError) as e:
            logging.warning('Caught %s exception while running round %d:\n\t%s', type(e), round_num, e)
            continue
        current_model = iterative_process.get_model_weights(state)
        train_metrics['training_secs'] = (time.time() - training_start_time)
        train_metrics['model_delta_l2_norm'] = _compute_numpy_l2_difference(current_model, prev_model)
        train_metrics['client_drift'] = state.client_drift
        train_metrics.update(round_metrics)
        loop_time = (time.time() - loop_start_time)
        loop_rounds = ((round_num - loop_start_round) + 1)
        logging.info('Round {:2d}, {:.2f}s per round in average.'.format(round_num, (loop_time / loop_rounds)))
        if (((round_num % rounds_per_checkpoint) == 0) or (round_num == (total_rounds - 1))):
            save_checkpoint_start_time = time.time()
            checkpoint_mngr.save_checkpoint(state, round_num)
            train_metrics['save_checkpoint_secs'] = (time.time() - save_checkpoint_start_time)
        metrics = {'train': train_metrics}
        if ((round_num % rounds_per_eval) == 0):
            evaluate_start_time = time.time()
            validation_metrics = validation_fn(current_model, round_num)
            validation_metrics['evaluate_secs'] = (time.time() - evaluate_start_time)
            metrics['eval'] = validation_metrics
        _write_metrics(metrics_mngr, tb_mngr, metrics, round_num)
        round_num += 1
    metrics = {}
    evaluate_start_time = time.time()
    validation_metrics = validation_fn(current_model, round_num)
    validation_metrics['evaluate_secs'] = (time.time() - evaluate_start_time)
    metrics['eval'] = validation_metrics
    if test_fn:
        test_start_time = time.time()
        test_metrics = test_fn(current_model)
        test_metrics['evaluate_secs'] = (time.time() - test_start_time)
        metrics['test'] = test_metrics
    _write_metrics(metrics_mngr, tb_mngr, metrics, total_rounds)
    return state