Returns the global id of the given task type in a cluster.
(cluster_spec, task_type, task_id, chief_task_type)
| 57 | |
| 58 | |
| 59 | def _get_global_id(cluster_spec, task_type, task_id, chief_task_type): |
| 60 | """Returns the global id of the given task type in a cluster.""" |
| 61 | if not task_type: |
| 62 | return 0 |
| 63 | |
| 64 | # Sort task names in cluster by "chief"/"master", "evaluator", "worker" |
| 65 | # and "ps". More details can be found at the documentation of |
| 66 | # `tf.estimator.RunConfig.global_id_in_cluster`. |
| 67 | task_type_ordered_list = [] |
| 68 | if chief_task_type in cluster_spec.jobs: |
| 69 | task_type_ordered_list = [chief_task_type] |
| 70 | task_type_ordered_list.extend([ |
| 71 | t for t in sorted(cluster_spec.jobs) if t != chief_task_type and t != PS |
| 72 | ]) |
| 73 | if PS in cluster_spec.jobs: |
| 74 | task_type_ordered_list.append(PS) |
| 75 | |
| 76 | # Find the right global_id for current task. |
| 77 | next_global_id = 0 |
| 78 | for t in task_type_ordered_list: |
| 79 | if t == task_type: |
| 80 | return next_global_id + task_id |
| 81 | # `cluster_spec.job_tasks` returns all task addresses of type `t`. |
| 82 | next_global_id += len(cluster_spec.job_tasks(t)) |
| 83 | |
| 84 | # It is unexpected that it passes through all task_types in |
| 85 | # `task_type_ordered_list`. |
| 86 | raise RuntimeError('Internal Error: `task_type` ({}) is not in ' |
| 87 | 'cluster_spec ({}).'.format(task_type, cluster_spec)) |
| 88 | |
| 89 | |
| 90 | def _init_run_config_from_worker_context(config, worker_context): |
no test coverage detected