Initializes RunConfig for distribution strategies.
(config, tf_config)
| 125 | |
| 126 | |
| 127 | def init_run_config(config, tf_config): |
| 128 | """Initializes RunConfig for distribution strategies.""" |
| 129 | # pylint: disable=protected-access |
| 130 | if (config._experimental_distribute and |
| 131 | config._experimental_distribute.train_distribute): |
| 132 | if config._train_distribute: |
| 133 | raise ValueError('Either `train_distribute` or' |
| 134 | '`experimental_distribute.train_distribute` can be set.') |
| 135 | config._train_distribute = config._experimental_distribute.train_distribute |
| 136 | |
| 137 | if (config._experimental_distribute and |
| 138 | config._experimental_distribute.eval_distribute): |
| 139 | if config._eval_distribute: |
| 140 | raise ValueError('Either `eval_distribute` or' |
| 141 | '`experimental_distribute.eval_distribute` can be set.') |
| 142 | config._eval_distribute = config._experimental_distribute.eval_distribute |
| 143 | |
| 144 | cluster_spec = server_lib.ClusterSpec(tf_config.get('cluster', {})) |
| 145 | config._init_distributed_setting_from_environment_var({}) |
| 146 | |
| 147 | # Use distribute coordinator with STANDALONE_CLIENT mode if |
| 148 | # `experimental_distribute.remote_cluster` is set. |
| 149 | if (config._train_distribute and config._experimental_distribute and |
| 150 | config._experimental_distribute.remote_cluster): |
| 151 | if cluster_spec: |
| 152 | raise ValueError('Cannot set both "cluster_spec" of TF_CONFIG and ' |
| 153 | '`experimental_distribute.remote_cluster`') |
| 154 | config._distribute_coordinator_mode = dc.CoordinatorMode.STANDALONE_CLIENT |
| 155 | config._cluster_spec = config._experimental_distribute.remote_cluster |
| 156 | logging.info('RunConfig initialized for Distribute Coordinator with ' |
| 157 | 'STANDALONE_CLIENT mode') |
| 158 | return |
| 159 | |
| 160 | # Don't use distribute coordinator if it is local training or cluster has a |
| 161 | # MASTER job or `train_distribute` is not specifed. |
| 162 | if (not cluster_spec or 'master' in cluster_spec.jobs or |
| 163 | not config._train_distribute): |
| 164 | config._distribute_coordinator_mode = None |
| 165 | config._init_distributed_setting_from_environment_var(tf_config) |
| 166 | config._maybe_overwrite_session_config_for_distributed_training() |
| 167 | logging.info('Not using Distribute Coordinator.') |
| 168 | return |
| 169 | |
| 170 | # Use distribute coordinator with INDEPENDENT_WORKER mode otherwise. |
| 171 | assert tf_config |
| 172 | |
| 173 | # Set the cluster_spec only since the distributed setting will come from |
| 174 | # distribute coordinator. |
| 175 | config._cluster_spec = cluster_spec |
| 176 | config._distribute_coordinator_mode = dc.CoordinatorMode.INDEPENDENT_WORKER |
| 177 | logging.info('RunConfig initialized for Distribute Coordinator with ' |
| 178 | 'INDEPENDENT_WORKER mode') |
| 179 | |
| 180 | |
| 181 | def should_run_distribute_coordinator(config): |