r'''Update parameters from cluster_spec. If task_type, task_id or cluster_spec is None, these arguments will not be changed. Args: task_type: (Optional.) name of current job. `localhost` by default. task_id: (Optional.) index of current task. 0 by defaul
(self, task_type=None, task_id=None, cluster_spec=None,
num_gpus=None)
| 166 | return self._num_gpus |
| 167 | |
| 168 | def _update(self, task_type=None, task_id=None, cluster_spec=None, |
| 169 | num_gpus=None): |
| 170 | r'''Update parameters from cluster_spec. |
| 171 | |
| 172 | If task_type, task_id or cluster_spec is None, these arguments will not be |
| 173 | changed. |
| 174 | |
| 175 | Args: |
| 176 | task_type: (Optional.) name of current job. `localhost` by default. |
| 177 | task_id: (Optional.) index of current task. 0 by default. |
| 178 | cluster_spec: (Optional.) ClusterSpec object. |
| 179 | ''' |
| 180 | tf_config = None |
| 181 | try: |
| 182 | tf_config = self.get_tf_config() |
| 183 | except: # pylint: disable=bare-except |
| 184 | pass |
| 185 | if tf_config: |
| 186 | self._task_type = tf_config.task_type |
| 187 | self._task_id = tf_config.task_id |
| 188 | self._cluster_spec = server_lib.ClusterSpec(tf_config.cluster) |
| 189 | else: |
| 190 | self._task_type = 'localhost' |
| 191 | self._task_id = 0 |
| 192 | self._cluster_spec = None |
| 193 | if task_type: |
| 194 | self._task_type = task_type |
| 195 | if self._task_type not in ('localhost', 'chief', 'worker'): |
| 196 | logging.info('No valid configuration for non-worker roles') |
| 197 | return |
| 198 | |
| 199 | if task_id: |
| 200 | self._task_id = task_id |
| 201 | if cluster_spec: |
| 202 | self._cluster_spec = cluster_spec |
| 203 | if self._cluster_spec: |
| 204 | self._cluster_spec = multi_worker_util.normalize_cluster_spec( |
| 205 | self._cluster_spec) |
| 206 | self._is_chief = False |
| 207 | try: |
| 208 | self._is_chief = multi_worker_util.is_chief( |
| 209 | self._cluster_spec, self._task_type, self._task_id) |
| 210 | except: # pylint: disable=bare-except |
| 211 | pass |
| 212 | if num_gpus: |
| 213 | self._num_gpus = num_gpus |
| 214 | elif not self._num_gpus: |
| 215 | num_gpus = 0 |
| 216 | num_gpus_config = config_pb2.ConfigProto() |
| 217 | num_gpus_config.inter_op_parallelism_threads = 1 |
| 218 | num_gpus_config.intra_op_parallelism_threads = 1 |
| 219 | num_gpus_config.gpu_options.allow_growth = True |
| 220 | for device in device_lib.list_local_devices(num_gpus_config): |
| 221 | if device.device_type == 'GPU': |
| 222 | num_gpus += 1 |
| 223 | self._num_gpus = num_gpus |
| 224 | self._default_device = ( |
| 225 | f'/job:{self._task_type}/replica:0/task:{self._task_id}') |