Connects to the given cluster. Will make devices on the cluster available to use. Note that calling this more than once will work, but will invalidate any tensor handles on the old remote devices. If the given local job name is not present in the cluster specification, it will be automat
(cluster_spec_or_resolver,
job_name="localhost",
task_index=0,
protocol=None)
| 75 | |
| 76 | @tf_export("config.experimental_connect_to_cluster") |
| 77 | def connect_to_cluster(cluster_spec_or_resolver, |
| 78 | job_name="localhost", |
| 79 | task_index=0, |
| 80 | protocol=None): |
| 81 | """Connects to the given cluster. |
| 82 | |
| 83 | Will make devices on the cluster available to use. Note that calling this more |
| 84 | than once will work, but will invalidate any tensor handles on the old remote |
| 85 | devices. |
| 86 | |
| 87 | If the given local job name is not present in the cluster specification, it |
| 88 | will be automatically added, using an unused port on the localhost. |
| 89 | |
| 90 | Args: |
| 91 | cluster_spec_or_resolver: A `ClusterSpec` or `ClusterResolver` describing |
| 92 | the cluster. |
| 93 | job_name: The name of the local job. |
| 94 | task_index: The local task index. |
| 95 | protocol: The communication protocol, such as `"grpc"`. If unspecified, will |
| 96 | use the default from `python/platform/remote_utils.py`. |
| 97 | """ |
| 98 | protocol = protocol or remote_utils.get_default_communication_protocol() |
| 99 | if isinstance(cluster_spec_or_resolver, server_lib.ClusterSpec): |
| 100 | cluster_spec = cluster_spec_or_resolver |
| 101 | elif isinstance(cluster_spec_or_resolver, cluster_resolver.ClusterResolver): |
| 102 | cluster_spec = cluster_spec_or_resolver.cluster_spec() |
| 103 | else: |
| 104 | raise ValueError( |
| 105 | "`cluster_spec_or_resolver` must be a `ClusterSpec` or a " |
| 106 | "`ClusterResolver`.") |
| 107 | |
| 108 | cluster_def = cluster_spec.as_cluster_def() |
| 109 | |
| 110 | # Automatically add local job, if not part of the cluster spec. |
| 111 | if job_name not in cluster_spec.jobs: |
| 112 | local_port = pywrap_tensorflow.TF_PickUnusedPortOrDie() |
| 113 | job_def = cluster_def.job.add() |
| 114 | job_def.name = job_name |
| 115 | # TODO(fishx): Update this to make sure remote worker has valid ip address |
| 116 | # to connect with local. |
| 117 | job_def.tasks[0] = "localhost:{}".format(local_port) |
| 118 | |
| 119 | server_def = ServerDef( |
| 120 | cluster=cluster_def, job_name=job_name, task_index=task_index, |
| 121 | protocol=protocol) |
| 122 | |
| 123 | # TODO(nareshmodi): Make this default since it works in more situations. |
| 124 | os.environ["TF_EAGER_REMOTE_USE_SEND_TENSOR_RPC"] = "1" |
| 125 | context.set_server_def(server_def) |
| 126 | |
| 127 | |
| 128 | def _strip_prefix(s, prefix): |
no test coverage detected