wehub-resource-sync
678 lines
26 KiB
Python
678 lines
26 KiB
Python
# 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.
|
||
# ==============================================================================
|
||
"""Tests for distributed_table."""
|
||
|
||
import copy
|
||
import os
|
||
|
||
from absl.testing import parameterized
|
||
|
||
from tensorflow.python.compat import v2_compat
|
||
from tensorflow.python.data.ops import dataset_ops
|
||
from tensorflow.python.distribute import combinations
|
||
from tensorflow.python.distribute import device_util
|
||
from tensorflow.python.distribute import multi_process_runner
|
||
from tensorflow.python.distribute import multi_worker_test_base
|
||
from tensorflow.python.distribute import parameter_server_strategy_v2
|
||
from tensorflow.python.distribute import ps_values
|
||
from tensorflow.python.distribute.coordinator import cluster_coordinator as coordinator_lib
|
||
from tensorflow.python.distribute.coordinator import coordinator_context
|
||
from tensorflow.python.eager import def_function
|
||
from tensorflow.python.eager import test
|
||
from tensorflow.python.framework import constant_op
|
||
from tensorflow.python.framework import dtypes
|
||
from tensorflow.python.framework import tensor_spec
|
||
from tensorflow.python.module import module
|
||
from tensorflow.python.ops import array_ops
|
||
from tensorflow.python.ops import lookup_ops
|
||
from tensorflow.python.ops import math_ops
|
||
from tensorflow.python.saved_model import load as tf_load
|
||
from tensorflow.python.saved_model import save as tf_save
|
||
|
||
|
||
source_combination = combinations.combine(source=["textfile", "keyvaluetensor"])
|
||
|
||
|
||
class DistributedTableTest(test.TestCase, parameterized.TestCase):
|
||
|
||
@classmethod
|
||
def setUpClass(cls):
|
||
super(DistributedTableTest, cls).setUpClass()
|
||
cls.cluster = multi_worker_test_base.create_multi_process_cluster(
|
||
num_workers=2, num_ps=3, rpc_layer="grpc")
|
||
cls.cluster_resolver = cls.cluster.cluster_resolver
|
||
|
||
@classmethod
|
||
def tearDownClass(cls):
|
||
super(DistributedTableTest, cls).tearDownClass()
|
||
cls.cluster.stop()
|
||
|
||
def make_initializer(self, init_source, vals):
|
||
if init_source == "textfile":
|
||
file = os.path.join(self.get_temp_dir(), "text_file_initializer")
|
||
with open(file, "w") as f:
|
||
f.write("\n".join(str(v) for v in vals) + "\n")
|
||
return lookup_ops.TextFileInitializer(
|
||
filename=file,
|
||
key_dtype=dtypes.int64,
|
||
key_index=lookup_ops.TextFileIndex.LINE_NUMBER,
|
||
value_dtype=dtypes.int64,
|
||
value_index=lookup_ops.TextFileIndex.WHOLE_LINE)
|
||
elif init_source == "keyvaluetensor":
|
||
keys_tensor = constant_op.constant(
|
||
list(range(len(vals))), dtype=dtypes.int64)
|
||
vals_tensor = constant_op.constant(vals, dtype=dtypes.int64)
|
||
return lookup_ops.KeyValueTensorInitializer(keys_tensor, vals_tensor)
|
||
else:
|
||
raise ValueError("Unrecognized init_source: " + init_source)
|
||
|
||
def createStaticHashTable(self,
|
||
init_source=None,
|
||
vals=None,
|
||
default_value=None,
|
||
initializer=None):
|
||
if not initializer:
|
||
initializer = self.make_initializer(init_source, vals)
|
||
return lookup_ops.StaticHashTable(
|
||
initializer=initializer, default_value=default_value)
|
||
|
||
def makeDatasetFromTensorWithoutUsingResource(self, input_context, tensor):
|
||
"""Returns a dataset made from `tensor`. To be called in a dataset_fn."""
|
||
global_batch_size = 24
|
||
batch_size = input_context.get_per_replica_batch_size(global_batch_size)
|
||
dataset = dataset_ops.DatasetV2.from_tensors(tensor).repeat().batch(
|
||
batch_size, drop_remainder=True)
|
||
dataset = dataset.shard(input_context.num_input_pipelines,
|
||
input_context.input_pipeline_id)
|
||
dataset = dataset.prefetch(2) # This prefetches 2 batches per device.
|
||
return dataset
|
||
|
||
@combinations.generate(source_combination)
|
||
def testCreateDistributedTableInScope(self, source):
|
||
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
|
||
self.cluster_resolver)
|
||
|
||
coordinator_lib.ClusterCoordinator(strategy=strategy)
|
||
|
||
with strategy.scope():
|
||
lookuptable = self.createStaticHashTable(
|
||
init_source=source, vals=[0, 1, 2], default_value=-2)
|
||
|
||
self.assertIsInstance(lookuptable, ps_values.DistributedTable)
|
||
self.assertEqual(self.evaluate(lookuptable.size()), 3)
|
||
|
||
# Lookup on the coordinator.
|
||
output = lookuptable.lookup(
|
||
constant_op.constant([0, 1, -1], dtype=dtypes.int64))
|
||
self.assertAllEqual([0, 1, -2], output)
|
||
self.assertEqual(lookuptable.size(), 3)
|
||
|
||
@combinations.generate(source_combination)
|
||
def testCopyDistributedTable(self, source):
|
||
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
|
||
self.cluster_resolver)
|
||
|
||
coordinator_lib.ClusterCoordinator(strategy=strategy)
|
||
|
||
with strategy.scope():
|
||
lookuptable = self.createStaticHashTable(
|
||
init_source=source, vals=[0, 1, 2], default_value=-2)
|
||
|
||
new_table = copy.copy(lookuptable)
|
||
# No new coordinator instance or distributed tables are created.
|
||
self.assertDictEqual(lookuptable.__dict__, new_table.__dict__)
|
||
|
||
@combinations.generate(source_combination)
|
||
def testCreateLookupInDatasetFnUnderScope(self, source):
|
||
# TODO(wxinyi): Warn the user of the inefficiency of this workflow (i.e.
|
||
# creating `StaticHashTable` inside a `@tf.function`-wrapped `dataset_fn` to
|
||
# be distributed with `distribute_datasets_from_function` and
|
||
# `create_per_worker_dataset`.
|
||
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
|
||
self.cluster_resolver)
|
||
|
||
coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
|
||
|
||
with strategy.scope():
|
||
|
||
def dataset_fn(input_context):
|
||
some_out_of_range_tensor = constant_op.constant(10, dtype=dtypes.int64)
|
||
lookuptable = self.createStaticHashTable(
|
||
init_source=source, vals=[0, 1, 2], default_value=-2)
|
||
|
||
self.assertNotIsInstance(lookuptable, ps_values.DistributedTable)
|
||
|
||
generation_tensor = lookuptable.lookup(some_out_of_range_tensor)
|
||
dataset = self.makeDatasetFromTensorWithoutUsingResource(
|
||
input_context, generation_tensor)
|
||
return dataset
|
||
|
||
@def_function.function
|
||
def per_worker_dataset_fn():
|
||
return strategy.distribute_datasets_from_function(dataset_fn)
|
||
|
||
per_worker_dataset = coordinator.create_per_worker_dataset(
|
||
per_worker_dataset_fn)
|
||
per_worker_iterator = iter(per_worker_dataset)
|
||
|
||
@def_function.function
|
||
def worker_fn(iterator):
|
||
return math_ops.reduce_sum(next(iterator))
|
||
|
||
result = []
|
||
for _ in range(10):
|
||
result.append(
|
||
coordinator.schedule(worker_fn, args=(per_worker_iterator,)))
|
||
|
||
for r in result:
|
||
returned_input = r.fetch()
|
||
self.assertAllClose(-48, returned_input)
|
||
|
||
@combinations.generate(source_combination)
|
||
def testAccessingResourceHandleInDatasetFnWithoutMap(self, source):
|
||
|
||
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
|
||
self.cluster_resolver)
|
||
|
||
coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
|
||
|
||
with strategy.scope():
|
||
lookuptable = self.createStaticHashTable(
|
||
init_source=source, vals=[0, 1, 2], default_value=-2)
|
||
|
||
def dataset_fn(input_context):
|
||
some_out_of_range_tensor = constant_op.constant(10, dtype=dtypes.int64)
|
||
|
||
self.assertIsInstance(lookuptable, ps_values.DistributedTable)
|
||
|
||
generation_tensor = lookuptable.lookup(some_out_of_range_tensor)
|
||
dataset = self.makeDatasetFromTensorWithoutUsingResource(
|
||
input_context, generation_tensor)
|
||
return dataset
|
||
|
||
@def_function.function
|
||
def per_worker_dataset_fn():
|
||
return strategy.distribute_datasets_from_function(dataset_fn)
|
||
|
||
per_worker_dataset = coordinator.create_per_worker_dataset(
|
||
per_worker_dataset_fn)
|
||
per_worker_iterator = iter(per_worker_dataset)
|
||
|
||
@def_function.function
|
||
def worker_fn(iterator):
|
||
return math_ops.reduce_sum(next(iterator))
|
||
|
||
result = []
|
||
for _ in range(10):
|
||
result.append(
|
||
coordinator.schedule(worker_fn, args=(per_worker_iterator,)))
|
||
|
||
for r in result:
|
||
returned_input = r.fetch()
|
||
self.assertAllClose(-48, returned_input)
|
||
|
||
@combinations.generate(
|
||
combinations.combine(
|
||
source=["textfile", "keyvaluetensor"],
|
||
create_datasets_under_scope=[True, False],
|
||
using_dataset_instance_not_function=[True, False],
|
||
create_per_worker_dataset_takes_instance=[True, False]))
|
||
def testCreateTableUnderScopeCombo(self, source,
|
||
create_datasets_under_scope,
|
||
using_dataset_instance_not_function,
|
||
create_per_worker_dataset_takes_instance):
|
||
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
|
||
self.cluster_resolver)
|
||
|
||
coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
|
||
|
||
with strategy.scope():
|
||
lookup_table = self.createStaticHashTable(
|
||
init_source=source, vals=[0, 1, 2], default_value=-2)
|
||
|
||
if using_dataset_instance_not_function:
|
||
|
||
def per_worker_dataset_fn():
|
||
dataset = dataset_ops.DatasetV2.from_tensors(
|
||
constant_op.constant([0, 1, 3], dtype=dtypes.int64))
|
||
dataset = dataset.repeat().batch(24, drop_remainder=True).prefetch(2)
|
||
dataset = dataset.map(lookup_table.lookup)
|
||
|
||
return strategy.experimental_distribute_dataset(dataset)
|
||
|
||
else:
|
||
|
||
def per_worker_dataset_fn():
|
||
def dataset_fn(input_context):
|
||
batch_size = input_context.get_per_replica_batch_size(24)
|
||
dataset = dataset_ops.DatasetV2.from_tensors(
|
||
constant_op.constant([0, 1, 3], dtype=dtypes.int64))
|
||
dataset = dataset.repeat().batch(batch_size, drop_remainder=True)
|
||
dataset = dataset.shard(input_context.num_input_pipelines,
|
||
input_context.input_pipeline_id)
|
||
dataset = dataset.prefetch(2) # This prefetches 2 batches per device.
|
||
dataset = dataset.map(lookup_table.lookup)
|
||
return dataset
|
||
return strategy.distribute_datasets_from_function(dataset_fn)
|
||
|
||
if create_datasets_under_scope:
|
||
with strategy.scope():
|
||
if create_per_worker_dataset_takes_instance:
|
||
per_worker_dataset = coordinator.create_per_worker_dataset(
|
||
per_worker_dataset_fn())
|
||
else:
|
||
per_worker_dataset = coordinator.create_per_worker_dataset(
|
||
per_worker_dataset_fn)
|
||
per_worker_iterator = iter(per_worker_dataset)
|
||
|
||
else:
|
||
if create_per_worker_dataset_takes_instance:
|
||
per_worker_dataset = coordinator.create_per_worker_dataset(
|
||
per_worker_dataset_fn())
|
||
else:
|
||
per_worker_dataset = coordinator.create_per_worker_dataset(
|
||
per_worker_dataset_fn)
|
||
per_worker_iterator = iter(per_worker_dataset)
|
||
|
||
@def_function.function
|
||
def worker_fn(iterator):
|
||
return math_ops.reduce_sum(next(iterator))
|
||
|
||
result = []
|
||
for _ in range(10):
|
||
result.append(
|
||
coordinator.schedule(worker_fn, args=(per_worker_iterator,)))
|
||
|
||
for r in result:
|
||
returned_input = r.fetch()
|
||
self.assertAllClose(-24, returned_input)
|
||
|
||
@combinations.generate(
|
||
combinations.combine(
|
||
source=["textfile", "keyvaluetensor"],
|
||
create_datasets_under_scope=[True, False],
|
||
using_dataset_instance_not_function=[True, False],
|
||
create_per_worker_dataset_takes_instance=[True, False]))
|
||
def testCreateTableInDatasetCombo(self, source, create_datasets_under_scope,
|
||
using_dataset_instance_not_function,
|
||
create_per_worker_dataset_takes_instance):
|
||
|
||
if using_dataset_instance_not_function and (
|
||
not create_per_worker_dataset_takes_instance):
|
||
# This is the case that uses the `experimental_distribute_dataset` API to
|
||
# distribute dataset (instead of the `distribute_datasets_from_function`
|
||
# API), and passes `create_per_worker_dataset` a function that returns
|
||
# the distributed dataset (instead of passing it the distributed dataset
|
||
# directly).
|
||
# TODO(b/201775366): evaluate whether we need to handle this case
|
||
self.skipTest("Failed to serialize the input pipeline graph")
|
||
|
||
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
|
||
self.cluster_resolver)
|
||
coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
|
||
|
||
if using_dataset_instance_not_function:
|
||
|
||
def per_worker_dataset_fn():
|
||
# If this line is being called under strategy.scope(), it becomes a
|
||
# DistributedTable. Interestingly, after
|
||
# `experimental_distribute_dataset` serializes the dataset on chief and
|
||
# deserializes it on workers, `lookup_table` becomes a
|
||
# RestoredDistributedTable instead of a DistributedTable. And when it’s
|
||
# `resource_handle` is being accessed on the worker, it does not detect
|
||
# a DispatchContext, so it returns the restored resource handle,
|
||
# which is also the one on the local worker. The LookupTableFindV2 ops
|
||
# is on the local worker, too.
|
||
lookup_table = self.createStaticHashTable(
|
||
init_source=source, vals=[0, 1, 2], default_value=-2)
|
||
|
||
if create_datasets_under_scope:
|
||
self.assertIsInstance(lookup_table, ps_values.DistributedTable)
|
||
|
||
dataset = dataset_ops.DatasetV2.from_tensors(
|
||
constant_op.constant([0, 1, 3], dtype=dtypes.int64))
|
||
dataset = dataset.repeat().batch(24, drop_remainder=True).prefetch(2)
|
||
dataset = dataset.map(lookup_table.lookup)
|
||
|
||
return strategy.experimental_distribute_dataset(dataset)
|
||
|
||
else:
|
||
|
||
def per_worker_dataset_fn():
|
||
|
||
def dataset_fn(input_context):
|
||
# When we're wrapping the initialization of a StaticHashTable inside a
|
||
# `dataset_fn` to be distributed with
|
||
# `distribute_datasets_from_function`, no matter it's called under
|
||
# strategy.scope() or not, this call creates a StaticHashTable on
|
||
# chief instead of a DistributedTable on chief and workers.
|
||
# And correspondingly, LookupTableFindV2 ops is on chief and there are
|
||
# send-recv communication for the lookup.
|
||
lookup_table = self.createStaticHashTable(
|
||
init_source=source, vals=[0, 1, 2], default_value=-2)
|
||
if create_datasets_under_scope:
|
||
self.assertIsInstance(lookup_table, lookup_ops.StaticHashTable)
|
||
self.assertNotIsInstance(lookup_table, ps_values.DistributedTable)
|
||
|
||
batch_size = input_context.get_per_replica_batch_size(24)
|
||
dataset = dataset_ops.DatasetV2.from_tensors(
|
||
constant_op.constant([0, 1, 3], dtype=dtypes.int64))
|
||
dataset = dataset.repeat().batch(batch_size, drop_remainder=True)
|
||
dataset = dataset.shard(input_context.num_input_pipelines,
|
||
input_context.input_pipeline_id)
|
||
dataset = dataset.prefetch(2) # This prefetches 2 batches per device.
|
||
dataset = dataset.map(lookup_table.lookup)
|
||
return dataset
|
||
|
||
return strategy.distribute_datasets_from_function(dataset_fn)
|
||
|
||
if create_datasets_under_scope:
|
||
with strategy.scope():
|
||
if create_per_worker_dataset_takes_instance:
|
||
per_worker_dataset = coordinator.create_per_worker_dataset(
|
||
per_worker_dataset_fn())
|
||
else:
|
||
per_worker_dataset = coordinator.create_per_worker_dataset(
|
||
per_worker_dataset_fn)
|
||
per_worker_iterator = iter(per_worker_dataset)
|
||
|
||
else:
|
||
if create_per_worker_dataset_takes_instance:
|
||
per_worker_dataset = coordinator.create_per_worker_dataset(
|
||
per_worker_dataset_fn())
|
||
else:
|
||
per_worker_dataset = coordinator.create_per_worker_dataset(
|
||
per_worker_dataset_fn)
|
||
per_worker_iterator = iter(per_worker_dataset)
|
||
|
||
@def_function.function
|
||
def worker_fn(iterator):
|
||
return math_ops.reduce_sum(next(iterator))
|
||
|
||
result = []
|
||
for _ in range(10):
|
||
result.append(
|
||
coordinator.schedule(worker_fn, args=(per_worker_iterator,)))
|
||
|
||
for r in result:
|
||
returned_input = r.fetch()
|
||
self.assertAllClose(-24, returned_input)
|
||
|
||
@combinations.generate(source_combination)
|
||
def testAccessingTableInStepFunction(self, source):
|
||
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
|
||
self.cluster_resolver)
|
||
coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
|
||
|
||
with strategy.scope():
|
||
lookup_table = self.createStaticHashTable(
|
||
init_source=source, vals=[0, 1, 2], default_value=-2)
|
||
|
||
dataset = (
|
||
dataset_ops.DatasetV2.from_tensors(
|
||
constant_op.constant([0, 1, 3], dtype=dtypes.int64)).repeat().batch(
|
||
24, drop_remainder=True).prefetch(2))
|
||
dataset = dataset.map(lookup_table.lookup)
|
||
|
||
distributed_dataset = strategy.experimental_distribute_dataset(dataset)
|
||
distributed_dataset = coordinator.create_per_worker_dataset(
|
||
distributed_dataset)
|
||
|
||
@def_function.function
|
||
def worker_fn(iterator):
|
||
|
||
def replica_fn(inputs):
|
||
return math_ops.reduce_sum(lookup_table.lookup(inputs))
|
||
|
||
all_results = strategy.run(replica_fn, args=(next(iterator),))
|
||
return all_results
|
||
|
||
steps_per_epoch = 10
|
||
distributed_iterator = iter(distributed_dataset)
|
||
result = []
|
||
for _ in range(steps_per_epoch):
|
||
|
||
result.append(
|
||
coordinator.schedule(worker_fn, args=(distributed_iterator,)))
|
||
|
||
coordinator.join()
|
||
|
||
for r in result:
|
||
returned_input = r.fetch()
|
||
self.assertAllClose(-24, returned_input)
|
||
|
||
@combinations.generate(source_combination)
|
||
def testAccessingResourceHandleInDatasetFnWithMapFnDefinedOutside(
|
||
self, source):
|
||
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
|
||
self.cluster_resolver)
|
||
|
||
coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
|
||
|
||
with strategy.scope():
|
||
lookuptable = self.createStaticHashTable(
|
||
init_source=source, vals=[0, 1, 2], default_value=-2)
|
||
|
||
def map_fn(vals):
|
||
return lookuptable.lookup(vals)
|
||
|
||
def dataset_fn(input_context):
|
||
generation_tensor = constant_op.constant([0, 1, 3], dtype=dtypes.int64)
|
||
dataset = self.makeDatasetFromTensorWithoutUsingResource(
|
||
input_context, generation_tensor)
|
||
dataset = dataset.map(map_fn)
|
||
return dataset
|
||
|
||
@def_function.function
|
||
def per_worker_dataset_fn():
|
||
return strategy.distribute_datasets_from_function(dataset_fn)
|
||
|
||
per_worker_dataset = coordinator.create_per_worker_dataset(
|
||
per_worker_dataset_fn)
|
||
per_worker_iterator = iter(per_worker_dataset)
|
||
|
||
@def_function.function
|
||
def worker_fn(iterator):
|
||
return math_ops.reduce_sum(next(iterator))
|
||
|
||
result = []
|
||
for _ in range(10):
|
||
# batch_size == 24 and each input is [0, 1, -2]
|
||
result.append(
|
||
coordinator.schedule(worker_fn, args=(per_worker_iterator,)))
|
||
|
||
for r in result:
|
||
returned_input = r.fetch()
|
||
self.assertAllClose(-24, returned_input)
|
||
|
||
class Model(module.Module):
|
||
|
||
def __init__(self, init_source, filepath):
|
||
vals = [0, 1, 2]
|
||
if init_source == "textfile":
|
||
|
||
with open(filepath, "w") as f:
|
||
f.write("\n".join(str(v) for v in vals) + "\n")
|
||
|
||
self.initializer = lookup_ops.TextFileInitializer(
|
||
filepath, dtypes.int64, lookup_ops.TextFileIndex.LINE_NUMBER,
|
||
dtypes.int64, lookup_ops.TextFileIndex.WHOLE_LINE)
|
||
else:
|
||
keys_tensor = constant_op.constant(
|
||
list(range(len(vals))), dtype=dtypes.int64)
|
||
vals_tensor = constant_op.constant(vals, dtype=dtypes.int64)
|
||
self.initializer = lookup_ops.KeyValueTensorInitializer(
|
||
keys_tensor, vals_tensor)
|
||
|
||
self.table = lookup_ops.StaticHashTable(
|
||
self.initializer, default_value=-2)
|
||
|
||
@def_function.function(
|
||
input_signature=[tensor_spec.TensorSpec(None, dtypes.int64)])
|
||
def use_table(self, x):
|
||
return self.table.lookup(x)
|
||
|
||
def verifyWorkerLocalInstance(self, coordinator, model):
|
||
# assert capturing a worker-local resource on each worker
|
||
for worker in coordinator._cluster.workers:
|
||
with coordinator_context.with_dispatch_context(worker):
|
||
captures = model.use_table.get_concrete_function().captured_inputs
|
||
resource_capture = [t for t in captures if t.dtype == dtypes.resource]
|
||
self.assertNotEmpty(resource_capture)
|
||
for capture in resource_capture:
|
||
self.assertEqual(
|
||
capture.device,
|
||
device_util.canonicalize("/CPU:0", default=worker.device_name))
|
||
|
||
@combinations.generate(source_combination)
|
||
def testInModelAndCapture(self, source):
|
||
|
||
file_path = os.path.join(self.get_temp_dir(), "text_file_initializer")
|
||
|
||
model = self.Model(source, file_path)
|
||
func_captures = model.use_table.get_concrete_function(
|
||
).graph.external_captures
|
||
self.assertLen(func_captures, 2)
|
||
self.assertTrue(
|
||
any(model.table.resource_handle is t for t in func_captures))
|
||
deferred_captures = model.use_table.get_concrete_function(
|
||
).graph.deferred_external_captures
|
||
self.assertEmpty(deferred_captures)
|
||
|
||
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
|
||
self.cluster_resolver)
|
||
coordinator = coordinator_lib.ClusterCoordinator(strategy)
|
||
with strategy.scope():
|
||
distributed_model = self.Model("value", file_path)
|
||
func_captures = distributed_model.use_table.get_concrete_function(
|
||
).graph.external_captures
|
||
# One less external_capture, since the table handle becomes a closure in the
|
||
# deferred_external_capture
|
||
self.assertLen(func_captures, 1)
|
||
self.assertFalse(
|
||
any(model.table.resource_handle is t for t in func_captures))
|
||
deferred_captures = distributed_model.use_table.get_concrete_function(
|
||
).graph.deferred_external_captures
|
||
self.assertNotEmpty(deferred_captures)
|
||
|
||
self.verifyWorkerLocalInstance(coordinator, distributed_model)
|
||
|
||
@combinations.generate(source_combination)
|
||
def testLookupInNestedTFWhileLoop(self, source):
|
||
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
|
||
self.cluster_resolver)
|
||
coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
|
||
|
||
file_path = os.path.join(self.get_temp_dir(), "text_file_initializer")
|
||
with strategy.scope():
|
||
model = self.Model(source, file_path)
|
||
|
||
@def_function.function
|
||
def replica_fn(batch_data):
|
||
replica_result = array_ops.zeros(shape=(), dtype=dtypes.int64)
|
||
for _ in math_ops.range(10):
|
||
replica_result += math_ops.reduce_sum(model.use_table(batch_data))
|
||
return replica_result
|
||
|
||
@def_function.function
|
||
def step_fn(iterator):
|
||
|
||
step_result = array_ops.zeros(shape=(), dtype=dtypes.int64)
|
||
for _ in math_ops.range(10):
|
||
step_result += strategy.run(replica_fn, args=(next(iterator),))
|
||
|
||
return step_result
|
||
|
||
dataset = (
|
||
dataset_ops.DatasetV2.from_tensors(
|
||
constant_op.constant([0, 1, 3], dtype=dtypes.int64)).repeat().batch(
|
||
24, drop_remainder=True).prefetch(2))
|
||
distributed_dataset = coordinator.create_per_worker_dataset(
|
||
strategy.experimental_distribute_dataset(dataset))
|
||
|
||
results = []
|
||
for _ in range(10):
|
||
results.append(
|
||
coordinator.schedule(step_fn, args=(iter(distributed_dataset),)))
|
||
|
||
coordinator.join()
|
||
|
||
for r in results:
|
||
self.assertAllClose(-2400, r.fetch())
|
||
|
||
@combinations.generate(source_combination)
|
||
def testDistributeTableSaveAndServe(self, source):
|
||
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
|
||
self.cluster_resolver)
|
||
file_path = os.path.join(self.get_temp_dir(), "text_file_initializer")
|
||
with strategy.scope():
|
||
model = self.Model(source, file_path)
|
||
|
||
model_dir = self.get_temp_dir()
|
||
tf_save.save(model, model_dir)
|
||
|
||
loaded_without_strategy = tf_load.load(model_dir)
|
||
loaded_func_captures_without_strategy = (
|
||
loaded_without_strategy.use_table.get_concrete_function().graph
|
||
.external_captures)
|
||
loaded_func_deferred_captures_without_strategy = (
|
||
loaded_without_strategy.use_table.get_concrete_function().graph
|
||
.deferred_external_captures)
|
||
self.assertLen(loaded_func_captures_without_strategy, 2)
|
||
self.assertEmpty(loaded_func_deferred_captures_without_strategy)
|
||
|
||
self.assertAllEqual(
|
||
loaded_without_strategy.use_table(
|
||
constant_op.constant([0, 1, 3], dtype=dtypes.int64)), [0, 1, -2])
|
||
|
||
@combinations.generate(source_combination)
|
||
def testDistributeTableSaveAndLoadUnderStrategy(self, source):
|
||
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
|
||
self.cluster_resolver)
|
||
coordinator = coordinator_lib.ClusterCoordinator(strategy)
|
||
file_path = os.path.join(self.get_temp_dir(), "text_file_initializer")
|
||
with strategy.scope():
|
||
model = self.Model(source, file_path)
|
||
model_dir = self.get_temp_dir()
|
||
tf_save.save(model, model_dir)
|
||
|
||
with strategy.scope():
|
||
loaded = tf_load.load(model_dir)
|
||
|
||
loaded_func_captures = (
|
||
loaded.use_table.get_concrete_function().graph.external_captures)
|
||
loaded_func_deferred_captures = (
|
||
loaded.use_table.get_concrete_function().graph
|
||
.deferred_external_captures)
|
||
# Compared with loading without strategy, there is one less
|
||
# external_capture, since the captured table handle has been swapped to a
|
||
# closure in the deferred_external_capture
|
||
self.assertLen(loaded_func_captures, 1)
|
||
self.assertNotEmpty(loaded_func_deferred_captures)
|
||
|
||
self.assertIsInstance(loaded.table, ps_values.DistributedTable)
|
||
|
||
self.assertLen([
|
||
t for t in loaded.use_table.get_concrete_function().captured_inputs
|
||
if t.dtype == dtypes.resource
|
||
], 1)
|
||
|
||
self.verifyWorkerLocalInstance(coordinator, loaded)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
v2_compat.enable_v2_behavior()
|
||
multi_process_runner.test_main()
|