wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/data/experimental/service/BUILD
+100
View File
@@ -0,0 +1,100 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test", "tf_python_pybind_extension")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//tensorflow:internal"],
licenses = ["notice"],
)
tf_python_pybind_extension(
name = "_pywrap_server_lib",
srcs = ["server_lib_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_server_lib.pyi",
],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:dispatcher_client",
"//tensorflow/core/data/service:grpc_util",
"//tensorflow/core/data/service:server_lib",
"//tensorflow/core/data/service:server_lib_headers_lib",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/python/lib/core:pybind11_lib",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@pybind11",
"@pybind11_protobuf//pybind11_protobuf:native_proto_caster",
"@xla//third_party/python_runtime:headers",
],
)
py_library(
name = "server_lib",
srcs = ["server_lib.py"],
strict_deps = True,
visibility = [
"//visibility:public",
],
deps = [
":_pywrap_server_lib",
":_pywrap_utils_exp",
"//tensorflow/core:protos_all_py",
"//tensorflow/python:pywrap_tensorflow",
"//tensorflow/python/util:tf_export",
],
)
tf_py_strict_test(
name = "server_lib_test",
srcs = ["server_lib_test.py"],
deps = [
":server_lib",
"//tensorflow/python/framework:errors",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/profiler:profiler_client",
],
)
tf_python_pybind_extension(
name = "_pywrap_utils_exp",
srcs = ["utils_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_utils_exp.pyi",
],
deps = [
"//tensorflow/core/data/service:py_utils",
"//tensorflow/python/lib/core:pybind11_lib",
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
tf_python_pybind_extension(
name = "_pywrap_snapshot_utils",
srcs = ["snapshot_utils_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_snapshot_utils.pyi",
],
deps = [
"//tensorflow/core/data/service/snapshot:path_utils",
"//tensorflow/python/lib/core:pybind11_lib",
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
py_library(
name = "service",
srcs = ["__init__.py"],
strict_deps = True,
deps = [
":server_lib",
"//tensorflow/python/data/experimental/ops:data_service_ops",
],
)
tensorflow/python/data/experimental/service/__init__.py
+426
View File
@@ -0,0 +1,426 @@
# Copyright 2020 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.
# ==============================================================================
"""API for using the tf.data service.
This module contains:
1. tf.data server implementations for running the tf.data service.
2. APIs for registering datasets with the tf.data service and reading from
the registered datasets.
The tf.data service provides the following benefits:
- Horizontal scaling of tf.data input pipeline processing to solve input
bottlenecks.
- Data coordination for distributed training. Coordinated reads
enable all replicas to train on similar-length examples across each global
training step, improving step times in synchronous training.
- Dynamic balancing of data across training replicas.
>>> dispatcher = tf.data.experimental.service.DispatchServer()
>>> dispatcher_address = dispatcher.target.split("://")[1]
>>> worker = tf.data.experimental.service.WorkerServer(
... tf.data.experimental.service.WorkerConfig(
... dispatcher_address=dispatcher_address))
>>> dataset = tf.data.Dataset.range(10)
>>> dataset = dataset.apply(tf.data.experimental.service.distribute(
... processing_mode=tf.data.experimental.service.ShardingPolicy.OFF,
... service=dispatcher.target))
>>> print([a.item() for a in dataset.as_numpy_iterator()])
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
## Setup
This section goes over how to set up the tf.data service.
### Run tf.data servers
The tf.data service consists of one dispatch server and `n` worker servers.
tf.data servers should be brought up alongside your training jobs, then brought
down when the jobs are finished.
Use `tf.data.experimental.service.DispatchServer` to start a dispatch server,
and `tf.data.experimental.service.WorkerServer` to start worker servers. Servers
can be run in the same process for testing purposes, or scaled up on separate
machines.
See https://github.com/tensorflow/ecosystem/tree/master/data_service for an
example of using Google Kubernetes Engine (GKE) to manage the tf.data service.
Note that the server implementation in
[tf_std_data_server.py](https://github.com/tensorflow/ecosystem/blob/master/data_service/tf_std_data_server.py)
is not GKE-specific, and can be used to run the tf.data service in other
contexts.
### Custom ops
If your dataset uses custom ops, these ops need to be made available to tf.data
servers by calling
[load_op_library](https://www.tensorflow.org/api_docs/python/tf/load_op_library)
from the dispatcher and worker processes at startup.
## Usage
Users interact with tf.data service by programmatically registering their
datasets with tf.data service, then creating datasets that read from the
registered datasets. The
[register_dataset](https://www.tensorflow.org/api_docs/python/tf/data/experimental/service/register_dataset)
function registers a dataset, then the
[from_dataset_id](https://www.tensorflow.org/api_docs/python/tf/data/experimental/service/from_dataset_id)
function creates a new dataset which reads from the registered dataset.
The
[distribute](https://www.tensorflow.org/api_docs/python/tf/data/experimental/service/distribute)
function wraps `register_dataset` and `from_dataset_id` into a single convenient
transformation which registers its input dataset and then reads from it.
`distribute` enables tf.data service to be used with a one-line code change.
However, it assumes that the dataset is created and consumed by the same entity
and this assumption might not always be valid or desirable. In particular, in
certain scenarios, such as distributed training, it might be desirable to
decouple the creation and consumption of the dataset (via `register_dataset`
and `from_dataset_id` respectively) to avoid having to create the dataset on
each of the training workers.
### Example
#### `distribute`
To use the `distribute` transformation, apply the transformation after the
prefix of your input pipeline that you would like to be executed using tf.data
service (typically at the end).
```
dataset = ... # Define your dataset here.
# Move dataset processing from the local machine to the tf.data service
dataset = dataset.apply(
tf.data.experimental.service.distribute(
processing_mode=tf.data.experimental.service.ShardingPolicy.OFF,
service=FLAGS.tf_data_service_address,
job_name="shared_job"))
# Any transformations added after `distribute` will be run on the local machine.
dataset = dataset.prefetch(1)
```
The above code will create a tf.data service "job", which iterates through the
dataset to generate data. To share the data from a job across multiple clients
(e.g. when using TPUStrategy or MultiWorkerMirroredStrategy), set a common
`job_name` across all clients.
#### `register_dataset` and `from_dataset_id`
`register_dataset` registers a dataset with the tf.data service, returning a
dataset id for the registered dataset. `from_dataset_id` creates a dataset that
reads from the registered dataset. These APIs can be used to reduce dataset
building time for distributed training. Instead of building the dataset on all
training workers, we can build the dataset just once and then register the
dataset using `register_dataset`. Then all workers can call `from_dataset_id`
without needing to build the dataset themselves.
```
dataset = ... # Define your dataset here.
dataset_id = tf.data.experimental.service.register_dataset(
service=FLAGS.tf_data_service_address,
dataset=dataset)
# Use `from_dataset_id` to create per-worker datasets.
per_worker_datasets = {}
for worker in workers:
per_worker_datasets[worker] = tf.data.experimental.service.from_dataset_id(
processing_mode=tf.data.experimental.service.ShardingPolicy.OFF,
service=FLAGS.tf_data_service_address,
dataset_id=dataset_id,
job_name="shared_job")
```
### Processing Modes
`processing_mode` specifies how to shard a dataset among tf.data service
workers. tf.data service supports `OFF`, `DYNAMIC`, `FILE`, `DATA`,
`FILE_OR_DATA`, `HINT` sharding policies.
OFF: No sharding will be performed. The entire input dataset will be processed
independently by each of the tf.data service workers. For this reason, it is
important to shuffle data (e.g. filenames) non-deterministically, so that each
worker will process the elements of the dataset in a different order. This mode
can be used to distribute datasets that aren't splittable.
If a worker is added or restarted during ShardingPolicy.OFF processing, the
worker will instantiate a new copy of the dataset and begin producing data from
the beginning.
#### Dynamic Sharding
DYNAMIC: In this mode, tf.data service divides the dataset into two components:
a source component that generates "splits" such as filenames, and a processing
component that takes splits and outputs dataset elements. The source component
is executed in a centralized fashion by the tf.data service dispatcher, which
generates different splits of input data. The processing component is executed
in a parallel fashion by the tf.data service workers, each operating on a
different set of input data splits.
For example, consider the following dataset:
```
dataset = tf.data.Dataset.from_tensor_slices(filenames)
dataset = dataset.interleave(TFRecordDataset)
dataset = dataset.map(preprocess_fn)
dataset = dataset.batch(batch_size)
dataset = dataset.apply(
tf.data.experimental.service.distribute(
processing_mode=tf.data.experimental.service.ShardingPolicy.DYNAMIC,
...))
```
The `from_tensor_slices` will be run on the dispatcher, while the `interleave`,
`map`, and `batch` will be run on tf.data service workers. The workers will pull
filenames from the dispatcher for processing. To process a dataset with
dynamic sharding, the dataset must have a splittable source, and all of
its transformations must be compatible with splitting. While most sources and
transformations support splitting, there are exceptions, such as custom datasets
which may not implement the splitting API. Please file a Github issue if you
would like to use distributed epoch processing for a currently unsupported
dataset source or transformation.
If no workers are restarted during training, dynamic sharding mode will visit
every example exactly once. If workers are restarted during training, the splits
they were processing will not be fully visited. The dispatcher maintains a
cursor through the dataset's splits. Assuming fault tolerance is enabled (See
"Fault Tolerance" below), the dispatcher will store cursor state in write-ahead
logs so that the cursor can be restored in case the dispatcher is restarted
mid-training. This provides an at-most-once visitation guarantee in the presence
of server restarts.
#### Static Sharding
The following are static sharding policies. The semantics are similar to
`tf.data.experimental.AutoShardPolicy`. These policies require:
* The tf.data service cluster is configured with a fixed list of workers
in DispatcherConfig.
* Each client only reads from the local tf.data service worker.
If a worker is restarted while performing static sharding, the worker will
begin processing its shard again from the beginning.
FILE: Shards by input files (i.e. each worker will get a fixed set of files to
process). When this option is selected, make sure that there is at least as
many files as workers. If there are fewer input files than workers, a runtime
error will be raised.
DATA: Shards by elements produced by the dataset. Each worker will process the
whole dataset and discard the portion that is not for itself. Note that for
this mode to correctly partition the dataset elements, the dataset needs to
produce elements in a deterministic order.
FILE_OR_DATA: Attempts FILE-based sharding, falling back to DATA-based
sharding on failure.
HINT: Looks for the presence of `shard(SHARD_HINT, ...)` which is treated as a
placeholder to replace with `shard(num_workers, worker_index)`.
For backwards compatibility, `processing_mode` may also be set to the strings
`"parallel_epochs"` or `"distributed_epoch"`, which are respectively equivalent
to `ShardingPolicy.OFF` and `ShardingPolicy.DYNAMIC`.
### Coordinated Data Read
By default, when multiple consumers read from the same job, they receive data on
a first-come first-served basis. In some use cases, it is advantageous to
coordinate the consumers. At each step, consumers read data from the same
worker.
For example, the tf.data service can be used to coordinate example sizes across
a cluster during synchronous training, so that during each step all replicas
train on similar-sized elements. To achieve this, define a dataset which
generates rounds of `num_consumers` consecutive similar-sized batches, then
enable coordinated reads by setting `consumer_index` and `num_consumers`.
NOTE: To keep consumers in sync, coordinated reads require that the dataset have
infinite cardinality. You can get this by adding `.repeat()` at the end of the
dataset definition.
### Jobs
A tf.data service "job" refers to the process of reading from a dataset managed
by the tf.data service, using one or more data consumers. Jobs are created when
iterating over datasets that read from tf.data service. The data produced by a
job is determined by (1) dataset associated with the job and (2) the job's
processing mode. For example, if a job is created for the dataset
`Dataset.range(5)`, and the processing mode is `ShardingPolicy.OFF`, each
tf.data worker will produce the elements `{0, 1, 2, 3, 4}` for the job,
resulting in the
job producing `5 * num_workers` elements. If the processing mode is
`ShardingPolicy.DYNAMIC`, the job will only produce `5` elements.
One or more consumers can consume data from a job. By default, jobs are
"anonymous", meaning that only the consumer which created the job can read from
it. To share the output of a job across multiple consumers, you can set a common
`job_name`.
### Fault Tolerance
By default, the tf.data dispatch server stores its state in-memory, making it a
single point of failure during training. To avoid this, pass
`fault_tolerant_mode=True` when creating your `DispatchServer`. Dispatcher
fault tolerance requires `work_dir` to be configured and accessible from the
dispatcher both before and after restart (e.g. a GCS path). With fault tolerant
mode enabled, the dispatcher will journal its state to the work directory so
that no state is lost when the dispatcher is restarted.
WorkerServers may be freely restarted, added, or removed during training. At
startup, workers will register with the dispatcher and begin processing all
outstanding jobs from the beginning.
### Usage with tf.distribute
tf.distribute is the TensorFlow API for distributed training. There are
several ways to use tf.data with tf.distribute:
`strategy.experimental_distribute_dataset`,
`strategy.distribute_datasets_from_function`, and (for PSStrategy)
`coordinator.create_per_worker_dataset`. The following sections give code
examples for each.
In general we recommend using
`tf.data.experimental.service.{register_dataset,from_dataset_id}` over
`tf.data.experimental.service.distribute` for two reasons:
- The dataset only needs to be constructed and optimized once, instead of once
per worker. This can significantly reduce startup time, because the current
`experimental_distribute_dataset` and `distribute_datasets_from_function`
implementations create and optimize worker datasets sequentially.
- If a dataset depends on lookup tables or variables that are only present on
one host, the dataset needs to be registered from that host. Typically this
only happens when resources are placed on the chief or worker 0. Registering
the dataset from the chief will avoid issues with depending on remote
resources.
#### strategy.experimental_distribute_dataset
Nothing special is required when using
`strategy.experimental_distribute_dataset`, just apply `register_dataset` and
`from_dataset_id` as above, making sure to specify a `job_name` so that all
workers consume from the same tf.data service job.
```
dataset = ... # Define your dataset here.
dataset_id = tf.data.experimental.service.register_dataset(
service=FLAGS.tf_data_service_address,
dataset=dataset)
dataset = tf.data.experimental.service.from_dataset_id(
processing_mode=tf.data.experimental.service.ShardingPolicy.OFF,
service=FLAGS.tf_data_service_address,
dataset_id=dataset_id,
job_name="shared_job")
dataset = strategy.experimental_distribute_dataset(dataset)
```
#### strategy.distribute_datasets_from_function
First, make sure the dataset produced by the `dataset_fn` does not depend on the
`input_context` for the training worker on which it is run. Instead of each
worker building its own (sharded) dataset, one worker should register an
unsharded dataset, and the remaining workers should consume data from that
dataset.
```
dataset = dataset_fn()
dataset_id = tf.data.experimental.service.register_dataset(
service=FLAGS.tf_data_service_address,
dataset=dataset)
def new_dataset_fn(input_context):
del input_context
return tf.data.experimental.service.from_dataset_id(
processing_mode=tf.data.experimental.service.ShardingPolicy.OFF,
service=FLAGS.tf_data_service_address,
dataset_id=dataset_id,
job_name="shared_job")
dataset = strategy.distribute_datasets_from_function(new_dataset_fn)
```
#### coordinator.create_per_worker_dataset
`create_per_worker_dataset` works the same as
`distribute_datasets_from_function`.
```
dataset = dataset_fn()
dataset_id = tf.data.experimental.service.register_dataset(
service=FLAGS.tf_data_service_address,
dataset=dataset)
def new_dataset_fn(input_context):
del input_context
return tf.data.experimental.service.from_dataset_id(
processing_mode=tf.data.experimental.service.ShardingPolicy.OFF,
service=FLAGS.tf_data_service_address,
dataset_id=dataset_id,
job_name="shared_job")
dataset = coordinator.create_per_worker_dataset(new_dataset_fn)
```
### Sharing tf.data service with concurrent trainers
If you run multiple trainers concurrently using the same training data, it could
save resources to cache the data in one tf.data service cluster and share the
cluster with the trainers. For example, if you use Vizier to tune
hyperparameters, the Vizier jobs can run concurrently and share one tf.data
service cluster.
To enable this feature, each trainer needs to generate a unique trainer ID, and
you pass the trainer ID to `tf.data.experimental.service.distribute`. Once a job
has consumed data, the data remains in the cache and is re-used by jobs with
different `trainer_id`s. Requests with the same `trainer_id` do not re-use data.
For example:
```
dataset = expensive_computation()
dataset = dataset.apply(tf.data.experimental.service.distribute(
processing_mode=tf.data.experimental.service.ShardingPolicy.OFF,
service=FLAGS.tf_data_service_address,
job_name="job",
cross_trainer_cache=data_service_ops.CrossTrainerCache(
trainer_id=trainer_id())))
```
tf.data service uses a sliding-window cache to store shared data. When one
trainer consumes data, the data remains in the cache. When other trainers need
data, they can get data from the cache instead of repeating the expensive
computation. The cache has a bounded size, so some workers may not read the full
dataset. To ensure all the trainers get sufficient training data, we require the
input dataset to be infinite. This can be achieved, for example, by repeating
the dataset and performing random augmentation on the training instances.
## Limitations
- Python-based data processing: Datasets which use Python-based data processing
(e.g. `tf.py_function`, `tf.numpy_function`, or
`tf.data.Dataset.from_generator`) are currently not supported.
- Non-Serializable Resources: Datasets may only depend on TF resources that
support serialization. Serialization is currently supported for lookup
tables and variables. If your dataset depends on a TF resource that cannot be
serialized, please file a Github issue.
- Remote Resources: If a dataset depends on a resource, the dataset must be
registered from the same process that created the resource (e.g. the "chief"
job of ParameterServerStrategy).
"""
from tensorflow.python.data.experimental.ops.data_service_ops import distribute
from tensorflow.python.data.experimental.ops.data_service_ops import from_dataset_id
from tensorflow.python.data.experimental.ops.data_service_ops import register_dataset
from tensorflow.python.data.experimental.ops.data_service_ops import ShardingPolicy
from tensorflow.python.data.experimental.service.server_lib import DispatcherConfig
from tensorflow.python.data.experimental.service.server_lib import DispatchServer
from tensorflow.python.data.experimental.service.server_lib import WorkerConfig
from tensorflow.python.data.experimental.service.server_lib import WorkerServer
tensorflow/python/data/experimental/service/_pywrap_server_lib.pyi
+52
View File
@@ -0,0 +1,52 @@
# Copyright 2023 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.
# ==============================================================================
class DispatchGrpcDataServer:
def __init__(self, *args, **kwargs) -> None: ...
def bound_port(self) -> int: ...
def join(self) -> None: ...
def num_workers(self) -> int: ...
def snapshot_streams(self, *args, **kwargs): ...
def start(self) -> Status: ...
def stop(self) -> None: ...
class SnapshotStreamInfoWrapper:
def __init__(self) -> None: ...
@property
def index(self) -> int: ...
@property
def state(self) -> int: ...
class SnapshotTaskProgressWrapper:
def __init__(self) -> None: ...
@property
def completed(self) -> bool: ...
@property
def snapshot_task_base_path(self) -> bytes: ...
@property
def snapshot_task_stream_index(self) -> int: ...
class WorkerGrpcDataServer:
def __init__(self, *args, **kwargs) -> None: ...
def bound_port(self) -> int: ...
def join(self) -> None: ...
def num_tasks(self) -> int: ...
def snapshot_task_progresses(self, *args, **kwargs): ...
def start(self) -> Status: ...
def stop(self) -> None: ...
def TF_DATA_GetDataServiceMetadataByID(*args, **kwargs): ...
def TF_DATA_NewDispatchServer(arg0: str) -> DispatchGrpcDataServer: ...
def TF_DATA_NewWorkerServer(arg0: str) -> WorkerGrpcDataServer: ...
tensorflow/python/data/experimental/service/_pywrap_snapshot_utils.pyi
+19
View File
@@ -0,0 +1,19 @@
# Copyright 2023 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.
# ==============================================================================
def TF_DATA_CommittedChunksDirectory(arg0: str) -> str: ...
def TF_DATA_SnapshotDoneFilePath(arg0: str) -> str: ...
def TF_DATA_SnapshotErrorFilePath(arg0: str) -> str: ...
def TF_DATA_SnapshotMetadataFilePath(arg0: str) -> str: ...
tensorflow/python/data/experimental/service/_pywrap_utils_exp.pyi
+16
View File
@@ -0,0 +1,16 @@
# Copyright 2023 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.
# ==============================================================================
def TF_DATA_DefaultProtocol() -> str: ...
tensorflow/python/data/experimental/service/server_lib.py
+478
View File
@@ -0,0 +1,478 @@
# Copyright 2020 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.
# ==============================================================================
"""A Python interface for creating dataset servers."""
import collections
from typing import Iterable
# pylint: disable=invalid-import-order,g-bad-import-order, unused-import
from tensorflow.core.protobuf import service_config_pb2
from tensorflow.python import pywrap_tensorflow
from tensorflow.python.data.experimental.service import _pywrap_server_lib
from tensorflow.python.data.experimental.service import _pywrap_utils_exp
from tensorflow.python.util.tf_export import tf_export
def _get_time_or_placeholder(value) -> int:
"""Modifies time-based config values to account for special behaviors."""
# Servers interpret time values of 0 to mean "choose a reasonable
# default". However, the Python API uses `None` for this, and allows 0 as a
# normal value. To account for this, if a user explicitly configures the
# interval/timeout to 0, we interpret it to mean "a very small number", and
# replace it with 1.
if value == 0:
return 1
# `None` indicates that the user wants to leave the behavior to the runtime.
if value is None:
return 0
return value
@tf_export("data.experimental.service.DispatcherConfig")
class DispatcherConfig(
collections.namedtuple(
"DispatcherConfig",
[
"port",
"protocol",
"work_dir",
"fault_tolerant_mode",
"worker_addresses",
"job_gc_check_interval_ms",
"job_gc_timeout_ms",
"worker_timeout_ms",
"worker_max_concurrent_snapshots",
],
)
):
"""Configuration class for tf.data service dispatchers.
Fields:
port: Specifies the port to bind to. A value of 0 indicates that the server
may bind to any available port.
protocol: The protocol to use for communicating with the tf.data service,
e.g. "grpc".
work_dir: A directory to store dispatcher state in. This
argument is required for the dispatcher to be able to recover from
restarts.
fault_tolerant_mode: Whether the dispatcher should write its state to a
journal so that it can recover from restarts. Dispatcher state, including
registered datasets and created jobs, is synchronously written to the
journal before responding to RPCs. If `True`, `work_dir` must also be
specified.
worker_addresses: If the job uses auto-sharding, it needs to specify a fixed
list of worker addresses that will register with the dispatcher. The
worker addresses should be in the format `"host"` or `"host:port"`, where
`"port"` is an integer, named port, or `%port%` to match any port.
job_gc_check_interval_ms: How often the dispatcher should scan through to
delete old and unused jobs, in milliseconds. If not set, the runtime will
select a reasonable default. A higher value will reduce load on the
dispatcher, while a lower value will reduce the time it takes for the
dispatcher to garbage collect expired jobs.
job_gc_timeout_ms: How long a job needs to be unused before it becomes a
candidate for garbage collection, in milliseconds. A value of -1 indicates
that jobs should never be garbage collected. If not set, the runtime will
select a reasonable default. A higher value will cause jobs to stay around
longer with no consumers. This is useful if there is a large gap in
time between when consumers read from the job. A lower value will reduce
the time it takes to reclaim the resources from expired jobs.
worker_timeout_ms: How long to wait for a worker to heartbeat before
considering it missing. If not set, the runtime will select a reasonable
default.
worker_max_concurrent_snapshots: The maximum number of snapshots a worker
can concurrently process.
"""
def __new__(
cls,
port=0,
protocol=None,
work_dir=None,
fault_tolerant_mode=False,
worker_addresses=None,
job_gc_check_interval_ms=None,
job_gc_timeout_ms=None,
worker_timeout_ms=None,
worker_max_concurrent_snapshots=0,
):
if protocol is None:
protocol = _pywrap_utils_exp.TF_DATA_DefaultProtocol()
job_gc_check_interval_ms = _get_time_or_placeholder(
job_gc_check_interval_ms)
job_gc_timeout_ms = _get_time_or_placeholder(job_gc_timeout_ms)
return super().__new__(
cls,
port,
protocol,
work_dir,
fault_tolerant_mode,
worker_addresses,
job_gc_check_interval_ms,
job_gc_timeout_ms,
worker_timeout_ms,
worker_max_concurrent_snapshots,
)
@tf_export("data.experimental.service.DispatchServer", v1=[])
class DispatchServer:
"""An in-process tf.data service dispatch server.
A `tf.data.experimental.service.DispatchServer` coordinates a cluster of
`tf.data.experimental.service.WorkerServer`s. When the workers start, they
register themselves with the dispatcher.
>>> dispatcher = tf.data.experimental.service.DispatchServer()
>>> dispatcher_address = dispatcher.target.split("://")[1]
>>> worker = tf.data.experimental.service.WorkerServer(
... tf.data.experimental.service.WorkerConfig(
... dispatcher_address=dispatcher_address))
>>> dataset = tf.data.Dataset.range(10)
>>> dataset = dataset.apply(tf.data.experimental.service.distribute(
... processing_mode="parallel_epochs", service=dispatcher.target))
>>> [a.item() for a in dataset.as_numpy_iterator()]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
When starting a dedicated tf.data dispatch process, use join() to block
after starting up the server, until the server terminates.
```
dispatcher = tf.data.experimental.service.DispatchServer(
tf.data.experimental.service.DispatcherConfig(port=5050))
dispatcher.join()
```
Call stop() to gracefully terminate the dispatcher. The server automatically
stops when all reference to it have been deleted.
To start a `DispatchServer` in fault-tolerant mode, set `work_dir` and
`fault_tolerant_mode` like below:
```
dispatcher = tf.data.experimental.service.DispatchServer(
tf.data.experimental.service.DispatcherConfig(
port=5050,
work_dir="gs://my-bucket/dispatcher/work_dir",
fault_tolerant_mode=True))
```
"""
def __init__(self, config=None, start=True):
"""Creates a new dispatch server.
Args:
config: (Optional.) A `tf.data.experimental.service.DispatcherConfig`
configuration. If `None`, the dispatcher will use default configuration
values.
start: (Optional.) Boolean, indicating whether to start the server after
creating it. Defaults to True.
"""
config = config or DispatcherConfig()
if config.fault_tolerant_mode and not config.work_dir:
raise ValueError(
"Cannot enable fault tolerant mode without configuring a work dir. "
"Make sure to set `work_dir` in the `config` object passed to "
"`DispatcherServer`.")
self._config = config
if isinstance(config, service_config_pb2.DispatcherConfig):
config_proto = config
else:
config_proto = service_config_pb2.DispatcherConfig(
port=config.port,
protocol=config.protocol,
work_dir=config.work_dir,
fault_tolerant_mode=config.fault_tolerant_mode,
worker_addresses=config.worker_addresses,
job_gc_check_interval_ms=config.job_gc_check_interval_ms,
job_gc_timeout_ms=config.job_gc_timeout_ms,
worker_timeout_ms=config.worker_timeout_ms,
worker_max_concurrent_snapshots=config.worker_max_concurrent_snapshots
)
self._server = _pywrap_server_lib.TF_DATA_NewDispatchServer(
config_proto.SerializeToString())
if start:
self._server.start()
def start(self):
"""Starts this server.
>>> dispatcher = tf.data.experimental.service.DispatchServer(start=False)
>>> dispatcher.start()
Raises:
tf.errors.OpError: Or one of its subclasses if an error occurs while
starting the server.
"""
self._server.start()
def join(self) -> None:
"""Blocks until the server has shut down.
This is useful when starting a dedicated dispatch process.
```
dispatcher = tf.data.experimental.service.DispatchServer(
tf.data.experimental.service.DispatcherConfig(port=5050))
dispatcher.join()
```
Raises:
tf.errors.OpError: Or one of its subclasses if an error occurs while
joining the server.
"""
self._server.join()
def stop(self) -> None:
"""Stops the server.
Raises:
tf.errors.OpError: Or one of its subclasses if an error occurs while
stopping the server.
"""
self._stop()
@property
def target(self) -> str:
"""Returns a target that can be used to connect to the server.
>>> dispatcher = tf.data.experimental.service.DispatchServer()
>>> dataset = tf.data.Dataset.range(10)
>>> dataset = dataset.apply(tf.data.experimental.service.distribute(
... processing_mode="parallel_epochs", service=dispatcher.target))
The returned string will be in the form protocol://address, e.g.
"grpc://localhost:5050".
"""
return "{0}://localhost:{1}".format(self._config.protocol,
self._server.bound_port())
def _stop(self) -> None:
"""Stops the server.
Raises:
tf.errors.OpError: Or one of its subclasses if an error occurs while
stopping the server.
"""
self._server.stop()
def __del__(self) -> None:
self._stop()
@property
def _address(self) -> str:
"""Returns the address of the server.
The returned string will be in the form address:port, e.g. "localhost:1000".
"""
return "localhost:{0}".format(self._server.bound_port())
def _num_workers(self) -> int:
"""Returns the number of workers registered with the dispatcher."""
return self._server.num_workers()
def _snapshot_streams(
self, path) -> Iterable[_pywrap_server_lib.SnapshotStreamInfoWrapper]:
"""Returns information about all the streams for a snapshot."""
return self._server.snapshot_streams(path)
@tf_export("data.experimental.service.WorkerConfig")
class WorkerConfig(
collections.namedtuple("WorkerConfig", [
"dispatcher_address", "worker_address", "port", "protocol",
"heartbeat_interval_ms", "dispatcher_timeout_ms",
"data_transfer_protocol", "data_transfer_address"
])):
"""Configuration class for tf.data service dispatchers.
Fields:
dispatcher_address: Specifies the address of the dispatcher.
worker_address: Specifies the address of the worker server. This address is
passed to the dispatcher so that the dispatcher can tell clients how to
connect to this worker.
port: Specifies the port to bind to. A value of 0 indicates that the worker
can bind to any available port.
protocol: A string indicating the protocol to be used by the worker to
connect to the dispatcher. E.g. "grpc".
heartbeat_interval_ms: How often the worker should heartbeat to the
dispatcher, in milliseconds. If not set, the runtime will select a
reasonable default. A higher value will reduce the load on the dispatcher,
while a lower value will reduce the time it takes to reclaim resources
from finished jobs.
dispatcher_timeout_ms: How long, in milliseconds, to retry requests to the
dispatcher before giving up and reporting an error. Defaults to 1 hour.
data_transfer_protocol: A string indicating the protocol to be used by the
worker to transfer data to the client. E.g. "grpc".
data_transfer_address: A string indicating the data transfer address of the
worker server.
"""
def __new__(cls,
dispatcher_address,
worker_address=None,
port=0,
protocol=None,
heartbeat_interval_ms=None,
dispatcher_timeout_ms=None,
data_transfer_protocol=None,
data_transfer_address=None):
if worker_address is None:
worker_address = "localhost:%port%"
if protocol is None:
protocol = _pywrap_utils_exp.TF_DATA_DefaultProtocol()
if data_transfer_address is None:
data_transfer_address = "localhost:%dts_port%"
heartbeat_interval_ms = _get_time_or_placeholder(heartbeat_interval_ms)
dispatcher_timeout_ms = _get_time_or_placeholder(dispatcher_timeout_ms)
return super(WorkerConfig,
cls).__new__(cls, dispatcher_address, worker_address, port,
protocol, heartbeat_interval_ms,
dispatcher_timeout_ms, data_transfer_protocol,
data_transfer_address)
@tf_export("data.experimental.service.WorkerServer", v1=[])
class WorkerServer:
"""An in-process tf.data service worker server.
A `tf.data.experimental.service.WorkerServer` performs `tf.data.Dataset`
processing for user-defined datasets, and provides the resulting elements over
RPC. A worker is associated with a single
`tf.data.experimental.service.DispatchServer`.
>>> dispatcher = tf.data.experimental.service.DispatchServer()
>>> dispatcher_address = dispatcher.target.split("://")[1]
>>> worker = tf.data.experimental.service.WorkerServer(
... tf.data.experimental.service.WorkerConfig(
... dispatcher_address=dispatcher_address))
>>> dataset = tf.data.Dataset.range(10)
>>> dataset = dataset.apply(tf.data.experimental.service.distribute(
... processing_mode="parallel_epochs", service=dispatcher.target))
>>> [a.item() for a in dataset.as_numpy_iterator()]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
When starting a dedicated tf.data worker process, use join() to block
after starting up the worker, until the worker terminates.
```
worker = tf.data.experimental.service.WorkerServer(
port=5051, dispatcher_address="localhost:5050")
worker.join()
```
Call stop() to gracefully terminate the worker. The worker automatically stops
when all reference to it have been deleted.
"""
def __init__(self, config, start=True):
"""Creates a new worker server.
Args:
config: A `tf.data.experimental.service.WorkerConfig` configuration.
start: (Optional.) Boolean, indicating whether to start the server after
creating it. Defaults to True.
"""
if config.dispatcher_address is None:
raise ValueError(
"Must specify a `dispatcher_address` in the `config` passed "
"to `WorkerServer`.")
if isinstance(config, service_config_pb2.WorkerConfig):
config_proto = config
else:
config_proto = service_config_pb2.WorkerConfig(
dispatcher_address=config.dispatcher_address,
worker_address=config.worker_address,
port=config.port,
protocol=config.protocol,
heartbeat_interval_ms=config.heartbeat_interval_ms,
dispatcher_timeout_ms=config.dispatcher_timeout_ms,
data_transfer_protocol=config.data_transfer_protocol,
data_transfer_address=config.data_transfer_address)
self._server = _pywrap_server_lib.TF_DATA_NewWorkerServer(
config_proto.SerializeToString())
if start:
self._server.start()
def start(self) -> None:
"""Starts this server.
Raises:
tf.errors.OpError: Or one of its subclasses if an error occurs while
starting the server.
"""
self._server.start()
def join(self) -> None:
"""Blocks until the server has shut down.
This is useful when starting a dedicated worker process.
```
worker_server = tf.data.experimental.service.WorkerServer(
port=5051, dispatcher_address="localhost:5050")
worker_server.join()
```
This method currently blocks forever.
Raises:
tf.errors.OpError: Or one of its subclasses if an error occurs while
joining the server.
"""
self._server.join()
def stop(self) -> None:
"""Stops the server.
Raises:
tf.errors.OpError: Or one of its subclasses if an error occurs while
stopping the server.
"""
self._stop()
def _stop(self) -> None:
"""Stops the server.
Raises:
tf.errors.OpError: Or one of its subclasses if an error occurs while
stopping the server.
"""
self._server.stop()
def __del__(self) -> None:
self._stop()
@property
def _address(self) -> str:
"""Returns the address of the server.
The returned string will be in the form address:port, e.g. "localhost:1000".
"""
return "localhost:{0}".format(self._server.bound_port())
def _num_tasks(self) -> int:
"""Returns the number of tasks currently being executed on the worker."""
return self._server.num_tasks()
def _snapshot_task_progresses(
self) -> Iterable[_pywrap_server_lib.SnapshotTaskProgressWrapper]:
"""Returns the progresses of the snapshot tasks currently being executed.
Returns:
An `Iterable[common_pb2.SnapshotTaskProgress]`.
"""
return self._server.snapshot_task_progresses()
tensorflow/python/data/experimental/service/server_lib_test.py
+177
View File
@@ -0,0 +1,177 @@
# Copyright 2020 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.
# ==============================================================================
"""Tests for tf.data service server lib."""
import logging
import tempfile
import threading
import unittest
from tensorflow.python.data.experimental.service import server_lib
from tensorflow.python.platform import test
from tensorflow.python.profiler import profiler_client
_portpicker_import_error = None
try:
import portpicker # pylint: disable=g-import-not-at-top
except ImportError as _error: # pylint: disable=invalid-name
_portpicker_import_error = _error
portpicker = None
ASSIGNED_PORTS = set()
lock = threading.Lock()
def pick_unused_port():
"""Returns an unused and unassigned local port."""
if _portpicker_import_error:
raise _portpicker_import_error # pylint: disable=raising-bad-type
global ASSIGNED_PORTS
with lock:
while True:
try:
port = portpicker.pick_unused_port()
except portpicker.NoFreePortFoundError:
raise unittest.SkipTest("Flakes in portpicker library do not represent "
"TensorFlow errors.")
if port > 10000 and port not in ASSIGNED_PORTS:
ASSIGNED_PORTS.add(port)
logging.info("Using local port %r", port)
return port
class ServerLibTest(test.TestCase):
def testStartDispatcher(self):
dispatcher = server_lib.DispatchServer(start=False)
dispatcher.start()
def testStartDispatcherWithPortConfig(self):
port = pick_unused_port()
config = server_lib.DispatcherConfig(port=port)
dispatcher = server_lib.DispatchServer(config=config, start=True)
self.assertEqual(dispatcher.target, "grpc://localhost:{}".format(port))
def testStartDispatcherWithWorkDirConfig(self):
temp_dir = tempfile.mkdtemp()
config = server_lib.DispatcherConfig(work_dir=temp_dir)
dispatcher = server_lib.DispatchServer( # pylint: disable=unused-variable
config=config, start=True)
def testStartDispatcherWithFaultTolerantConfig(self):
temp_dir = tempfile.mkdtemp()
config = server_lib.DispatcherConfig(
work_dir=temp_dir, fault_tolerant_mode=True)
dispatcher = server_lib.DispatchServer( # pylint: disable=unused-variable
config=config, start=True)
def testStartDispatcherWithWrongFaultTolerantConfig(self):
config = server_lib.DispatcherConfig(fault_tolerant_mode=True)
error = "Cannot enable fault tolerant mode without configuring a work dir"
with self.assertRaisesRegex(ValueError, error):
dispatcher = server_lib.DispatchServer( # pylint: disable=unused-variable
config=config, start=True)
def testMultipleStartDispatcher(self):
dispatcher = server_lib.DispatchServer(start=True)
dispatcher.start()
def testStartWorker(self):
dispatcher = server_lib.DispatchServer()
worker = server_lib.WorkerServer(
server_lib.WorkerConfig(dispatcher._address), start=False)
worker.start()
def testStartWorkerWithPortConfig(self):
dispatcher = server_lib.DispatchServer()
port = pick_unused_port()
worker = server_lib.WorkerServer(
server_lib.WorkerConfig(dispatcher._address, port=port), start=True)
self.assertEqual(worker._address, "localhost:{}".format(port))
def testMultipleStartWorker(self):
dispatcher = server_lib.DispatchServer()
worker = server_lib.WorkerServer(
server_lib.WorkerConfig(dispatcher._address), start=True)
worker.start()
def testStopDispatcher(self):
dispatcher = server_lib.DispatchServer()
dispatcher.stop()
dispatcher.stop()
def testStopWorker(self):
dispatcher = server_lib.DispatchServer()
worker = server_lib.WorkerServer(
server_lib.WorkerConfig(dispatcher._address))
worker.stop()
worker.stop()
def testStopStartDispatcher(self):
dispatcher = server_lib.DispatchServer()
dispatcher.stop()
with self.assertRaisesRegex(
RuntimeError, "Server cannot be started after it has been stopped"):
dispatcher.start()
def testStopStartWorker(self):
dispatcher = server_lib.DispatchServer()
worker = server_lib.WorkerServer(
server_lib.WorkerConfig(dispatcher._address))
worker.stop()
with self.assertRaisesRegex(
RuntimeError, "Server cannot be started after it has been stopped"):
worker.start()
def testJoinDispatcher(self):
dispatcher = server_lib.DispatchServer()
dispatcher.stop()
dispatcher.join()
def testJoinWorker(self):
dispatcher = server_lib.DispatchServer()
worker = server_lib.WorkerServer(
server_lib.WorkerConfig(dispatcher._address))
worker.stop()
worker.join()
def testDispatcherNumWorkers(self):
dispatcher = server_lib.DispatchServer()
self.assertEqual(0, dispatcher._num_workers())
worker1 = server_lib.WorkerServer( # pylint: disable=unused-variable
server_lib.WorkerConfig(dispatcher._address))
self.assertEqual(1, dispatcher._num_workers())
worker2 = server_lib.WorkerServer( # pylint: disable=unused-variable
server_lib.WorkerConfig(dispatcher._address))
self.assertEqual(2, dispatcher._num_workers())
def testProfileWorker(self):
dispatcher = server_lib.DispatchServer()
worker = server_lib.WorkerServer(
server_lib.WorkerConfig(dispatcher._address)
)
# Test the profilers are successfully started and connected to profiler
# service on the worker. Since there is no op running, it is expected to
# return RuntimeError with no trace events collected string.
with self.assertRaises(RuntimeError) as error:
profiler_client.trace(worker._address, tempfile.mkdtemp(), duration_ms=10)
self.assertStartsWith(
str(error.exception), "UNAVAILABLE: No trace event was collected"
)
if __name__ == "__main__":
test.main()
tensorflow/python/data/experimental/service/server_lib_wrapper.cc
+188
View File
@@ -0,0 +1,188 @@
/* Copyright 2020 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.
==============================================================================*/
#include <cstdint>
#include <limits>
#include <memory>
#include <string>
#include <vector>
#include "Python.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "pybind11/chrono.h" // from @pybind11
#include "pybind11/complex.h" // from @pybind11
#include "pybind11/detail/common.h" // from @pybind11
#include "pybind11/functional.h" // from @pybind11
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/pytypes.h" // from @pybind11
#include "pybind11/stl.h" // from @pybind11
#include "pybind11_protobuf/native_proto_caster.h" // from @pybind11_protobuf
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/dispatcher_client.h"
#include "tensorflow/core/data/service/grpc_util.h"
#include "tensorflow/core/data/service/server_lib.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
#include "tensorflow/python/lib/core/pybind11_lib.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
namespace py = pybind11;
PYBIND11_MODULE(_pywrap_server_lib, m, pybind11::mod_gil_not_used()) {
pybind11_protobuf::ImportNativeProtoCasters();
py::class_<tensorflow::data::DispatchGrpcDataServer>(m,
"DispatchGrpcDataServer")
.def("start", &tensorflow::data::DispatchGrpcDataServer::Start)
.def("stop", &tensorflow::data::DispatchGrpcDataServer::Stop)
.def("join", &tensorflow::data::DispatchGrpcDataServer::Join,
py::call_guard<py::gil_scoped_release>())
.def("bound_port", &tensorflow::data::DispatchGrpcDataServer::BoundPort)
.def("num_workers",
[](tensorflow::data::DispatchGrpcDataServer* server) -> int {
int num_workers;
absl::Status status = server->NumWorkers(&num_workers);
tensorflow::MaybeRaiseFromStatus(status);
return num_workers;
})
.def("snapshot_streams",
[](tensorflow::data::DispatchGrpcDataServer* server,
const std::string& path)
-> std::vector<tensorflow::data::SnapshotStreamInfoWrapper> {
std::vector<tensorflow::data::SnapshotStreamInfoWrapper> streams;
absl::Status status = server->SnapshotStreams(path, &streams);
tensorflow::MaybeRaiseFromStatus(status);
return streams;
});
py::class_<tensorflow::data::WorkerGrpcDataServer>(m, "WorkerGrpcDataServer")
.def("start", &tensorflow::data::WorkerGrpcDataServer::Start)
.def("stop", &tensorflow::data::WorkerGrpcDataServer::Stop)
.def("join", &tensorflow::data::WorkerGrpcDataServer::Join,
py::call_guard<py::gil_scoped_release>())
.def("bound_port", &tensorflow::data::WorkerGrpcDataServer::BoundPort)
.def("num_tasks",
[](tensorflow::data::WorkerGrpcDataServer* server) -> int {
int num_tasks;
absl::Status status = server->NumTasks(&num_tasks);
tensorflow::MaybeRaiseFromStatus(status);
return num_tasks;
})
.def("snapshot_task_progresses",
[](tensorflow::data::WorkerGrpcDataServer* server)
-> std::vector<tensorflow::data::SnapshotTaskProgressWrapper> {
std::vector<tensorflow::data::SnapshotTaskProgressWrapper>
snapshot_task_progresses;
absl::Status status =
server->SnapshotTaskProgresses(&snapshot_task_progresses);
tensorflow::MaybeRaiseFromStatus(status);
return snapshot_task_progresses;
});
m.def(
"TF_DATA_NewDispatchServer",
[](std::string serialized_dispatcher_config)
-> std::unique_ptr<tensorflow::data::DispatchGrpcDataServer> {
tensorflow::data::experimental::DispatcherConfig config;
if (!config.ParseFromString(serialized_dispatcher_config)) {
tensorflow::MaybeRaiseFromStatus(absl::InvalidArgumentError(
"Failed to deserialize dispatcher config."));
}
std::unique_ptr<tensorflow::data::DispatchGrpcDataServer> server;
absl::Status status =
tensorflow::data::NewDispatchServer(config, server);
tensorflow::MaybeRaiseFromStatus(status);
return server;
},
py::return_value_policy::reference);
m.def(
"TF_DATA_NewWorkerServer",
[](std::string serialized_worker_config)
-> std::unique_ptr<tensorflow::data::WorkerGrpcDataServer> {
tensorflow::data::experimental::WorkerConfig config;
if (!config.ParseFromString(serialized_worker_config)) {
tensorflow::MaybeRaiseFromStatus(absl::InvalidArgumentError(
"Failed to deserialize worker config."));
}
std::unique_ptr<tensorflow::data::WorkerGrpcDataServer> server;
absl::Status status = tensorflow::data::NewWorkerServer(config, server);
tensorflow::MaybeRaiseFromStatus(status);
return server;
},
py::return_value_policy::reference);
m.def(
"TF_DATA_GetDataServiceMetadataByID",
[](std::string dataset_id, const std::string& address,
const std::string& protocol) -> tensorflow::data::DataServiceMetadata {
tensorflow::data::DataServiceMetadata metadata;
tensorflow::data::DataServiceDispatcherClient client(address, protocol);
int64_t deadline_micros = std::numeric_limits<int64_t>::max();
absl::Status status;
Py_BEGIN_ALLOW_THREADS;
status = tensorflow::data::grpc_util::Retry(
[&]() {
return client.GetDataServiceMetadata(dataset_id, metadata);
},
/*description=*/
absl::StrCat("Get data service metadata for dataset ", dataset_id,
" from dispatcher at ", address),
deadline_micros);
Py_END_ALLOW_THREADS;
tensorflow::MaybeRaiseFromStatus(status);
return metadata;
},
py::return_value_policy::reference);
py::class_<tensorflow::data::SnapshotTaskProgressWrapper>
snapshot_task_progress_wrapper(m, "SnapshotTaskProgressWrapper");
snapshot_task_progress_wrapper.def(py::init<>())
.def_property_readonly(
"snapshot_task_base_path",
[](const tensorflow::data::SnapshotTaskProgressWrapper&
snapshot_task_progress_wrapper) -> py::bytes {
return snapshot_task_progress_wrapper.snapshot_task_base_path;
})
.def_property_readonly(
"snapshot_task_stream_index",
[](const tensorflow::data::SnapshotTaskProgressWrapper&
snapshot_task_progress_wrapper) -> int {
return snapshot_task_progress_wrapper.snapshot_task_stream_index;
})
.def_property_readonly(
"completed",
[](const tensorflow::data::SnapshotTaskProgressWrapper&
snapshot_task_progress_wrapper) -> bool {
return snapshot_task_progress_wrapper.completed;
});
py::class_<tensorflow::data::SnapshotStreamInfoWrapper>
snapshot_stream_info_wrapper(m, "SnapshotStreamInfoWrapper");
snapshot_stream_info_wrapper.def(py::init<>())
.def_property_readonly(
"index",
[](const tensorflow::data::SnapshotStreamInfoWrapper&
snapshot_stream_info_wrapper) -> int {
return snapshot_stream_info_wrapper.index;
})
.def_property_readonly(
"state",
[](const tensorflow::data::SnapshotStreamInfoWrapper&
snapshot_stream_info_wrapper) -> int {
return snapshot_stream_info_wrapper.state;
});
};
tensorflow/python/data/experimental/service/snapshot_utils_wrapper.cc
+39
View File
@@ -0,0 +1,39 @@
/* Copyright 2023 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.
==============================================================================*/
#include <string>
#include "Python.h"
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/core/data/service/snapshot/path_utils.h"
#include "tensorflow/python/lib/core/pybind11_lib.h"
PYBIND11_MODULE(_pywrap_snapshot_utils, m) {
m.def("TF_DATA_SnapshotDoneFilePath",
[](const std::string& snapshot_path) -> std::string {
return tensorflow::data::SnapshotDoneFilePath(snapshot_path);
});
m.def("TF_DATA_SnapshotErrorFilePath",
[](const std::string& snapshot_path) -> std::string {
return tensorflow::data::SnapshotErrorFilePath(snapshot_path);
});
m.def("TF_DATA_SnapshotMetadataFilePath",
[](const std::string& snapshot_path) -> std::string {
return tensorflow::data::SnapshotMetadataFilePath(snapshot_path);
});
m.def("TF_DATA_CommittedChunksDirectory",
[](const std::string& snapshot_path) -> std::string {
return tensorflow::data::CommittedChunksDirectory(snapshot_path);
});
};
tensorflow/python/data/experimental/service/utils_wrapper.cc
+26
View File
@@ -0,0 +1,26 @@
/* Copyright 2021 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.
==============================================================================*/
#include <string>
#include "Python.h"
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/core/data/service/py_utils.h"
#include "tensorflow/python/lib/core/pybind11_lib.h"
PYBIND11_MODULE(_pywrap_utils_exp, m, pybind11::mod_gil_not_used()) {
m.def("TF_DATA_DefaultProtocol",
[]() -> std::string { return tensorflow::data::DefaultProtocol(); });
};