| # Copyright 2017 The TensorFlow Authors. All Rights Reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # ============================================================================== |
| """Implementation of ClusterResolvers for GCE instance groups.""" |
| |
| from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver |
| from tensorflow.python.training.server_lib import ClusterSpec |
| from tensorflow.python.util.tf_export import tf_export |
| |
| |
| _GOOGLE_API_CLIENT_INSTALLED = True |
| try: |
| from googleapiclient import discovery # pylint: disable=g-import-not-at-top |
| from oauth2client.client import GoogleCredentials # pylint: disable=g-import-not-at-top |
| except ImportError: |
| _GOOGLE_API_CLIENT_INSTALLED = False |
| |
| |
| @tf_export('distribute.cluster_resolver.GCEClusterResolver') |
| class GCEClusterResolver(ClusterResolver): |
| """ClusterResolver for Google Compute Engine. |
| |
| This is an implementation of cluster resolvers for the Google Compute Engine |
| instance group platform. By specifying a project, zone, and instance group, |
| this will retrieve the IP address of all the instances within the instance |
| group and return a ClusterResolver object suitable for use for distributed |
| TensorFlow. |
| |
| Note: this cluster resolver cannot retrieve `task_type`, `task_id` or |
| `rpc_layer`. To use it with some distribution strategies like |
| `tf.distribute.experimental.MultiWorkerMirroredStrategy`, you will need to |
| specify `task_type` and `task_id` in the constructor. |
| |
| Usage example with tf.distribute.Strategy: |
| |
| ```Python |
| # On worker 0 |
| cluster_resolver = GCEClusterResolver("my-project", "us-west1", |
| "my-instance-group", |
| task_type="worker", task_id=0) |
| strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy( |
| cluster_resolver=cluster_resolver) |
| |
| # On worker 1 |
| cluster_resolver = GCEClusterResolver("my-project", "us-west1", |
| "my-instance-group", |
| task_type="worker", task_id=1) |
| strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy( |
| cluster_resolver=cluster_resolver) |
| ``` |
| """ |
| |
| def __init__(self, |
| project, |
| zone, |
| instance_group, |
| port, |
| task_type='worker', |
| task_id=0, |
| rpc_layer='grpc', |
| credentials='default', |
| service=None): |
| """Creates a new GCEClusterResolver object. |
| |
| This takes in a few parameters and creates a GCEClusterResolver project. It |
| will then use these parameters to query the GCE API for the IP addresses of |
| each instance in the instance group. |
| |
| Args: |
| project: Name of the GCE project. |
| zone: Zone of the GCE instance group. |
| instance_group: Name of the GCE instance group. |
| port: Port of the listening TensorFlow server (default: 8470) |
| task_type: Name of the TensorFlow job this GCE instance group of VM |
| instances belong to. |
| task_id: The task index for this particular VM, within the GCE |
| instance group. In particular, every single instance should be assigned |
| a unique ordinal index within an instance group manually so that they |
| can be distinguished from each other. |
| rpc_layer: The RPC layer TensorFlow should use to communicate across |
| instances. |
| credentials: GCE Credentials. If nothing is specified, this defaults to |
| GoogleCredentials.get_application_default(). |
| service: The GCE API object returned by the googleapiclient.discovery |
| function. (Default: discovery.build('compute', 'v1')). If you specify a |
| custom service object, then the credentials parameter will be ignored. |
| |
| Raises: |
| ImportError: If the googleapiclient is not installed. |
| """ |
| self._project = project |
| self._zone = zone |
| self._instance_group = instance_group |
| self._task_type = task_type |
| self._task_id = task_id |
| self._rpc_layer = rpc_layer |
| self._port = port |
| self._credentials = credentials |
| |
| if credentials == 'default': |
| if _GOOGLE_API_CLIENT_INSTALLED: |
| self._credentials = GoogleCredentials.get_application_default() |
| |
| if service is None: |
| if not _GOOGLE_API_CLIENT_INSTALLED: |
| raise ImportError('googleapiclient must be installed before using the ' |
| 'GCE cluster resolver') |
| self._service = discovery.build( |
| 'compute', 'v1', |
| credentials=self._credentials) |
| else: |
| self._service = service |
| |
| def cluster_spec(self): |
| """Returns a ClusterSpec object based on the latest instance group info. |
| |
| This returns a ClusterSpec object for use based on information from the |
| specified instance group. We will retrieve the information from the GCE APIs |
| every time this method is called. |
| |
| Returns: |
| A ClusterSpec containing host information retrieved from GCE. |
| """ |
| request_body = {'instanceState': 'RUNNING'} |
| request = self._service.instanceGroups().listInstances( |
| project=self._project, |
| zone=self._zone, |
| instanceGroups=self._instance_group, |
| body=request_body, |
| orderBy='name') |
| |
| worker_list = [] |
| |
| while request is not None: |
| response = request.execute() |
| |
| items = response['items'] |
| for instance in items: |
| instance_name = instance['instance'].split('/')[-1] |
| |
| instance_request = self._service.instances().get( |
| project=self._project, |
| zone=self._zone, |
| instance=instance_name) |
| |
| if instance_request is not None: |
| instance_details = instance_request.execute() |
| ip_address = instance_details['networkInterfaces'][0]['networkIP'] |
| instance_url = '%s:%s' % (ip_address, self._port) |
| worker_list.append(instance_url) |
| |
| request = self._service.instanceGroups().listInstances_next( |
| previous_request=request, |
| previous_response=response) |
| |
| worker_list.sort() |
| return ClusterSpec({self._task_type: worker_list}) |
| |
| def master(self, task_type=None, task_id=None, rpc_layer=None): |
| task_type = task_type if task_type is not None else self._task_type |
| task_id = task_id if task_id is not None else self._task_id |
| |
| if task_type is not None and task_id is not None: |
| master = self.cluster_spec().task_address(task_type, task_id) |
| if rpc_layer or self._rpc_layer: |
| return '%s://%s' % (rpc_layer or self._rpc_layer, master) |
| else: |
| return master |
| |
| return '' |
| |
| @property |
| def task_type(self): |
| return self._task_type |
| |
| @property |
| def task_id(self): |
| return self._task_id |
| |
| @task_type.setter |
| def task_type(self, task_type): |
| raise RuntimeError( |
| 'You cannot reset the task_type of the GCEClusterResolver after it has ' |
| 'been created.') |
| |
| @task_id.setter |
| def task_id(self, task_id): |
| self._task_id = task_id |
| |
| @property |
| def rpc_layer(self): |
| return self._rpc_layer |
| |
| @rpc_layer.setter |
| def rpc_layer(self, rpc_layer): |
| self._rpc_layer = rpc_layer |
