« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
1# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Implementation of Cluster Resolvers for TF_CONFIG Environment Variables."""
18import json
19import os
21from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
22from tensorflow.python.training.server_lib import ClusterSpec
23from tensorflow.python.util.tf_export import tf_export
25_TF_CONFIG_ENV = 'TF_CONFIG'
26_SESSION_MASTER_KEY = 'session_master'
27_RPC_LAYER_KEY = 'rpc_layer'
28_TASK_KEY = 'task'
31def format_master_url(master, rpc_layer=None):
32 if rpc_layer:
33 return '%s://%s' % (rpc_layer, master)
34 else:
35 return master
38def _load_tf_config():
39 return json.loads(os.environ.get(_TF_CONFIG_ENV, '{}'))
42def _get_value_in_tfconfig(key, default=None):
43 tf_config = _load_tf_config()
44 return tf_config[key] if key in tf_config else default
47@tf_export('distribute.cluster_resolver.TFConfigClusterResolver')
48class TFConfigClusterResolver(ClusterResolver):
49 """Implementation of a ClusterResolver which reads the TF_CONFIG EnvVar.
51 This is an implementation of cluster resolvers when using TF_CONFIG to set
52 information about the cluster. The cluster spec returned will be
53 initialized from the TF_CONFIG environment variable.
55 An example to set TF_CONFIG is:
57 ```Python
58 os.environ['TF_CONFIG'] = json.dumps({
59 'cluster': {
60 'worker': ["localhost:12345", "localhost:23456"]
61 },
62 'task': {'type': 'worker', 'index': 0}
63 })
64 ```
66 However, sometimes the container orchestration framework will set TF_CONFIG
67 for you. In this case, you can just create an instance without passing in any
68 arguments. You can find an example here to let Kuburnetes set TF_CONFIG for
69 you: https://github.com/tensorflow/ecosystem/tree/master/kubernetes. Then you
70 can use it with `tf.distribute.Strategy` as:
72 ```Python
73 # `TFConfigClusterResolver` is already the default one in the following
74 # strategy.
75 strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
76 cluster_resolver=TFConfigClusterResolver())
77 ```
78 """
80 def __init__(self,
81 task_type=None,
82 task_id=None,
83 rpc_layer=None,
84 environment=None):
85 """Creates a new TFConfigClusterResolver.
87 Args:
88 task_type: (String, optional) Overrides the task type specified in the
89 TF_CONFIG environment variable.
90 task_id: (Integer, optional) Overrides the task index specified in the
91 TF_CONFIG environment variable.
92 rpc_layer: (String, optional) Overrides the rpc layer TensorFlow uses.
93 environment: (String, optional) Overrides the environment TensorFlow
94 operates in.
95 """
96 self._task_type = task_type
97 self._task_id = task_id
98 self._rpc_layer = rpc_layer
99 self._environment = environment
101 @property
102 def task_type(self):
103 if self._task_type is None:
104 task_info = _get_value_in_tfconfig(_TASK_KEY, {})
105 return str(task_info['type']) if 'type' in task_info else None
106 else:
107 return str(self._task_type)
109 @property
110 def task_id(self):
111 if self._task_id is None:
112 task_info = _get_value_in_tfconfig(_TASK_KEY, {})
113 return int(task_info['index']) if 'index' in task_info else None
114 else:
115 return int(self._task_id)
117 @task_type.setter
118 def task_type(self, task_type):
119 self._task_type = task_type
121 @task_id.setter
122 def task_id(self, task_id):
123 self._task_id = task_id
125 @property
126 def environment(self):
127 return self._environment
129 @property
130 def rpc_layer(self):
131 if self._rpc_layer is None:
132 return _get_value_in_tfconfig(_RPC_LAYER_KEY)
133 else:
134 return self._rpc_layer
136 @rpc_layer.setter
137 def rpc_layer(self, rpc_layer):
138 self._rpc_layer = rpc_layer
140 def num_accelerators(self,
141 task_type=None,
142 task_id=None,
143 config_proto=None):
144 task_type = self.task_type if task_type is None else task_type
145 task_id = self.task_id if task_id is None else task_id
146 return super(TFConfigClusterResolver, self).num_accelerators(
147 task_type, task_id, config_proto)
149 def cluster_spec(self):
150 """Returns a ClusterSpec based on the TF_CONFIG environment variable.
152 Returns:
153 A ClusterSpec with information from the TF_CONFIG environment variable.
154 """
155 tf_config = _load_tf_config()
156 if 'cluster' not in tf_config:
157 return ClusterSpec({})
158 return ClusterSpec(tf_config['cluster'])
160 def master(self, task_type=None, task_id=None, rpc_layer=None):
161 """Returns the master address to use when creating a TensorFlow session.
163 Note: this is only useful for TensorFlow 1.x.
165 Args:
166 task_type: (String, optional) Overrides and sets the task_type of the
167 master.
168 task_id: (Integer, optional) Overrides and sets the task id of the
169 master.
170 rpc_layer: (String, optional) Overrides and sets the protocol over which
171 TensorFlow nodes communicate with each other.
173 Returns:
174 The address of the master.
176 Raises:
177 RuntimeError: If the task_type or task_id is not specified and the
178 `TF_CONFIG` environment variable does not contain a task section.
179 """
181 # If `session_master` is set, just use that.
182 session_master = _get_value_in_tfconfig(_SESSION_MASTER_KEY)
183 if session_master is not None:
184 return session_master
186 # Return an empty string if we are the only job in the ClusterSpec.
187 cluster_spec = self.cluster_spec()
188 if (not cluster_spec.jobs or
189 (len(cluster_spec.jobs) == 1 and
190 len(cluster_spec.job_tasks(cluster_spec.jobs[0])) == 1)):
191 return ''
193 # We try to auto-detect the task type and id, but uses the user-supplied one
194 # where available
195 task_type = task_type if task_type is not None else self.task_type
196 task_id = task_id if task_id is not None else self.task_id
197 rpc_layer = rpc_layer if rpc_layer is not None else self.rpc_layer
199 return format_master_url(cluster_spec.task_address(task_type, task_id),
200 rpc_layer)