Returns the canonnical job name to use to place TPU computations on. Args: master: A `string` representing the TensorFlow master to use. cluster_def: A ClusterDef object describing the TPU cluster. Returns: A string containing the job name, or None if no job should be specified.
(master, cluster_def)
| 185 | |
| 186 | |
| 187 | def master_job(master, cluster_def): |
| 188 | """Returns the canonnical job name to use to place TPU computations on. |
| 189 | |
| 190 | Args: |
| 191 | master: A `string` representing the TensorFlow master to use. |
| 192 | cluster_def: A ClusterDef object describing the TPU cluster. |
| 193 | |
| 194 | |
| 195 | Returns: |
| 196 | A string containing the job name, or None if no job should be specified. |
| 197 | |
| 198 | Raises: |
| 199 | ValueError: If the user needs to specify a tpu_job_name, because we are |
| 200 | unable to infer the job name automatically, or if the user-specified job |
| 201 | names are inappropriate. |
| 202 | """ |
| 203 | # If the user specifies the tpu_job_name, use that. |
| 204 | |
| 205 | if master in _LOCAL_MASTERS: |
| 206 | return None |
| 207 | |
| 208 | if (not cluster_def or not cluster_def.job): |
| 209 | return _DEFAULT_JOB_NAME |
| 210 | job_names = set([job.name for job in cluster_def.job]) |
| 211 | if _DEFAULT_JOB_NAME in job_names: |
| 212 | # b/37868888 tracks allowing ClusterSpec propagation to reuse job names. |
| 213 | raise ValueError('Currently, tpu_worker is not an allowed job name.') |
| 214 | if len(job_names) == 1: |
| 215 | return cluster_def.job[0].name |
| 216 | if len(job_names) == 2: |
| 217 | if _DEFAULT_COORDINATOR_JOB_NAME in job_names: |
| 218 | job_names.remove(_DEFAULT_COORDINATOR_JOB_NAME) |
| 219 | return job_names.pop() |
| 220 | # TODO(b/67716447): Include more sophisticated heuristics. |
| 221 | raise ValueError( |
| 222 | 'Could not infer TPU job name. Please specify a tpu_job_name as part ' |
| 223 | 'of your TPUConfig.') |