def __init__(self,
dataset_id,
processing_mode,
address,
element_spec,
protocol,
data_transfer_protocol,
job_name=None,
consumer_index=None,
num_consumers=None,
max_outstanding_requests=None,
task_refresh_interval_hint_ms=None):
"""Constructs a _DataServiceDatasetV2.
Args:
dataset_id: The dataset id for the dataset to read from.
processing_mode: A string specifying the policy for how data should be
processed by tf.data workers. Can be either "parallel_epochs" to have
each tf.data worker process a copy of the dataset, or
"distributed_epoch" to split a single iteration of the dataset across
all the workers.
address: The tf.data service address, e.g. "localhost:5000".
element_spec: The dataset element spec for the dataset to read from.
protocol: The protocol to use for communicating with the tf.data service,
e.g. "grpc".
data_transfer_protocol: The protocol to use for transferring data with the
tf.data service, e.g. "grpc".
job_name: (Optional.) The name of the job. This argument makes it possible
for multiple datasets to share the same job. The default behavior is
that the dataset creates anonymous, exclusively owned jobs.
consumer_index: (Optional.) The index of the consumer in the range from
`0` to `num_consumers`. Must be specified alongside `num_consumers`.
When specified, consumers will read from the job in a strict round-robin
order, instead of the default first-come-first-served order.
num_consumers: (Optional.) The number of consumers which will consume from
the job. Must be specified alongside `consumer_index`. When specified,
consumers will read from the job in a strict round-robin order, instead
of the default first-come-first-served order. When `num_consumers` is
specified, the dataset must have infinite cardinality to prevent a
producer from running out of data early and causing consumers to go out
of sync.
max_outstanding_requests: (Optional.) A limit on how many elements may be
requested at the same time. You can use this option to control the
amount of memory used, since `distribute` won't use more than
`element_size` * `max_outstanding_requests` of memory.
task_refresh_interval_hint_ms: (Optional.) A hint for how often to query
the dispatcher for task changes.
"""
if consumer_index is None != num_consumers is None:
raise ValueError(
"Must either set both consumer_index and num_consumers, or neither. ",
"consumer_index: ", consumer_index, ", num_consumers: ",
num_consumers)
if num_consumers is not None and job_name is None:
raise ValueError("job_name must be set when setting num_consumers")
if job_name is None:
job_name = ""
if max_outstanding_requests is None:
max_outstanding_requests = dataset_ops.AUTOTUNE
if task_refresh_interval_hint_ms is None:
task_refresh_interval_hint_ms = dataset_ops.AUTOTUNE
self._dataset_id = ops.convert_to_tensor(
dataset_id, dtype=dtypes.int64, name="dataset_id")
self._processing_mode = ops.convert_to_tensor(
processing_mode, dtype=dtypes.string, name="processing_mode")
self._address = ops.convert_to_tensor(
address, dtype=dtypes.string, name="address")
self._protocol = ops.convert_to_tensor(
protocol, dtype=dtypes.string, name="protocol")
self._job_name = ops.convert_to_tensor(
job_name, dtype=dtypes.string, name="job_name")
self._consumer_index = ops.convert_to_tensor(
-1 if consumer_index is None else consumer_index,
dtype=dtypes.int64,
name="consumer_index")
self._num_consumers = ops.convert_to_tensor(
-1 if num_consumers is None else num_consumers,
dtype=dtypes.int64,
name="num_consumers")
self._max_outstanding_requests = ops.convert_to_tensor(
max_outstanding_requests,
dtype=dtypes.int64,
name="max_outstanding_requests")
self._element_spec = element_spec
compat_kwargs = {}
if data_transfer_protocol is not None:
compat_kwargs["data_transfer_protocol"] = data_transfer_protocol
variant_tensor = gen_experimental_dataset_ops.data_service_dataset_v2(
dataset_id=self._dataset_id,
processing_mode=self._processing_mode,
address=self._address,
protocol=self._protocol,
job_name=self._job_name,
consumer_index=self._consumer_index,
num_consumers=self._num_consumers,
max_outstanding_requests=self._max_outstanding_requests,
task_refresh_interval_hint_ms=task_refresh_interval_hint_ms,
iteration_counter=gen_experimental_dataset_ops.dummy_iteration_counter(
),
**compat_kwargs,
**self._flat_structure)
super(_DataServiceDatasetV2, self).__init__(variant_tensor)
def __init__(self,
dataset_id,
processing_mode,
address,
element_spec,
protocol,
data_transfer_protocol,
job_name=None,
consumer_index=None,
num_consumers=None,
max_outstanding_requests=None,
task_refresh_interval_hint_ms=None,
target_workers="AUTO"):
"""Constructs a _DataServiceDatasetV2.
Args:
dataset_id: The dataset id for the dataset to read from.
processing_mode: A `tf.data.experimental.service.ShardingPolicy`
specifying how to shard the dataset among tf.data workers. See
`tf.data.experimental.service.ShardingPolicy` for details. 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`.
address: The tf.data service address, e.g. "localhost:5000".
element_spec: The dataset element spec for the dataset to read from.
protocol: The protocol to use for communicating with the tf.data service,
e.g. "grpc".
data_transfer_protocol: (Optional.) The protocol to use for transferring
data with the tf.data service. By default, data is transferred using
gRPC.
job_name: (Optional.) The name of the job. If provided, it must be a
non-empty string or Tensor. This argument makes it possible
for multiple datasets to share the same job. The default behavior is
that the dataset creates anonymous, exclusively owned jobs.
consumer_index: (Optional.) The index of the consumer in the range from
`0` to `num_consumers`. Must be specified alongside `num_consumers`.
When specified, consumers will read from the job in a strict round-robin
order, instead of the default first-come-first-served order.
num_consumers: (Optional.) The number of consumers which will consume from
the job. Must be specified alongside `consumer_index`. When specified,
consumers will read from the job in a strict round-robin order, instead
of the default first-come-first-served order. When `num_consumers` is
specified, the dataset must have infinite cardinality to prevent a
producer from running out of data early and causing consumers to go out
of sync.
max_outstanding_requests: (Optional.) A limit on how many elements may be
requested at the same time. You can use this option to control the
amount of memory used, since `distribute` won't use more than
`element_size` * `max_outstanding_requests` of memory.
task_refresh_interval_hint_ms: (Optional.) A hint for how often to query
the dispatcher for task changes.
target_workers: (Optional.) Which workers to read from. If `"AUTO"`,
tf.data runtime decides which workers to read from. If `"ANY"`, reads
from any tf.data service workers. If `"LOCAL"`, only reads from local
in-processs tf.data service workers. `"AUTO"` works well for most cases,
while users can specify other targets. For example, `"LOCAL"` helps
avoid RPCs and data copy if every TF worker colocates with a tf.data
service worker. Defaults to `"AUTO"`.
"""
processing_mode_def = data_service_pb2.ProcessingModeDef(
sharding_policy=_get_validated_sharding_policy(processing_mode))
if consumer_index is None != num_consumers is None:
raise ValueError(
"Must either set both consumer_index and num_consumers, or neither. ",
"consumer_index: ", consumer_index, ", num_consumers: ",
num_consumers)
if num_consumers is not None and job_name is None:
raise ValueError("job_name must be set when setting num_consumers")
if job_name is None:
job_name = ""
if max_outstanding_requests is None:
max_outstanding_requests = dataset_ops.AUTOTUNE
if task_refresh_interval_hint_ms is None:
task_refresh_interval_hint_ms = dataset_ops.AUTOTUNE
self._dataset_id = ops.convert_to_tensor(dataset_id,
dtype=dtypes.int64,
name="dataset_id")
self._processing_mode = ops.convert_to_tensor(
processing_mode_def.SerializeToString(),
dtype=dtypes.string,
name="processing_mode")
self._address = ops.convert_to_tensor(address,
dtype=dtypes.string,
name="address")
self._protocol = ops.convert_to_tensor(protocol,
dtype=dtypes.string,
name="protocol")
self._job_name = ops.convert_to_tensor(job_name,
dtype=dtypes.string,
name="job_name")
self._consumer_index = ops.convert_to_tensor(
-1 if consumer_index is None else consumer_index,
dtype=dtypes.int64,
name="consumer_index")
self._num_consumers = ops.convert_to_tensor(
-1 if num_consumers is None else num_consumers,
dtype=dtypes.int64,
name="num_consumers")
self._max_outstanding_requests = ops.convert_to_tensor(
max_outstanding_requests,
dtype=dtypes.int64,
name="max_outstanding_requests")
self._element_spec = element_spec
self._target_workers = target_workers
compat_kwargs = {}
if data_transfer_protocol is not None:
compat_kwargs["data_transfer_protocol"] = data_transfer_protocol
if compat.forward_compatible(2021, 7, 12) or target_workers != "AUTO":
compat_kwargs["target_workers"] = target_workers
variant_tensor = gen_experimental_dataset_ops.data_service_dataset_v2(
dataset_id=self._dataset_id,
processing_mode=self._processing_mode,
address=self._address,
protocol=self._protocol,
job_name=self._job_name,
consumer_index=self._consumer_index,
num_consumers=self._num_consumers,
max_outstanding_requests=self._max_outstanding_requests,
task_refresh_interval_hint_ms=task_refresh_interval_hint_ms,
iteration_counter=gen_experimental_dataset_ops.
dummy_iteration_counter(),
**compat_kwargs,
**self._flat_structure)
super(_DataServiceDatasetV2, self).__init__(variant_tensor)