# Copyright 2019 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 that the system configuration methods work properly.""" from absl.testing import parameterized from tensorflow.core.protobuf import cluster_pb2 from tensorflow.core.protobuf import config_pb2 from tensorflow.core.protobuf import rewriter_config_pb2 from tensorflow.core.protobuf import tensorflow_server_pb2 from tensorflow.python.eager import context from tensorflow.python.eager import def_function from tensorflow.python.framework import config from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.framework import test_ops from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.platform import test from tensorflow.python.util import compat def reset_eager(fn): def wrapper(*args, **kwargs): try: return fn(*args, **kwargs) finally: # Reset the context. context._reset_jit_compiler_flags() context._reset_context() ops.enable_eager_execution_internal() assert context._context is not None return wrapper @test_util.with_eager_op_as_function class ConfigTest(test.TestCase, parameterized.TestCase): @test_util.disable_eager_op_as_function('b/204320409') @test_util.run_gpu_only @reset_eager def testDevicePolicy(self): self.assertEqual(context.DEVICE_PLACEMENT_SILENT, context.context().device_policy) # If no op has been executed we should be able to set the device policy as # well as any init-time configs. config.set_intra_op_parallelism_threads(1) config.set_device_policy('silent') config.set_intra_op_parallelism_threads(2) context.ensure_initialized() def copy_tensor(dtype=dtypes.int32): with ops.device('CPU:0'): cpu_tensor = constant_op.constant(1, dtype=dtype) gpu_tensor = cpu_tensor.gpu() self.assertAllEqual(cpu_tensor + gpu_tensor, 2.0) config.set_device_policy('silent') self.assertEqual(config.get_device_policy(), 'silent') self.assertEqual(context.DEVICE_PLACEMENT_SILENT, context.context().device_policy) copy_tensor() config.set_device_policy('silent_for_int32') self.assertEqual(config.get_device_policy(), 'silent_for_int32') self.assertEqual(context.DEVICE_PLACEMENT_SILENT_FOR_INT32, context.context().device_policy) with self.assertRaisesRegex(errors.InvalidArgumentError, 'Tensors on conflicting devices'): copy_tensor(dtypes.float32) copy_tensor() config.set_device_policy('warn') self.assertEqual(config.get_device_policy(), 'warn') self.assertEqual(context.DEVICE_PLACEMENT_WARN, context.context().device_policy) copy_tensor() config.set_device_policy('explicit') self.assertEqual(config.get_device_policy(), 'explicit') self.assertEqual(context.DEVICE_PLACEMENT_EXPLICIT, context.context().device_policy) with self.assertRaisesRegex(errors.InvalidArgumentError, 'Tensors on conflicting devices'): copy_tensor() config.set_device_policy(None) self.assertEqual(config.get_device_policy(), 'silent') @reset_eager def testExecutionMode(self): self.assertTrue(config.get_synchronous_execution()) self.assertEqual(context.SYNC, context.context().execution_mode) # If no op has been executed we should be able to set the execution mode as # well as any init-time configs. config.set_intra_op_parallelism_threads(1) config.set_synchronous_execution(False) config.set_intra_op_parallelism_threads(2) config.set_synchronous_execution(True) self.assertTrue(config.get_synchronous_execution()) self.assertEqual(context.SYNC, context.context().execution_mode) config.set_synchronous_execution(False) self.assertFalse(config.get_synchronous_execution()) self.assertEqual(context.ASYNC, context.context().execution_mode) @reset_eager def testIntraOpParallelismThreads(self): config.set_intra_op_parallelism_threads(10) self.assertEqual(config.get_intra_op_parallelism_threads(), context.context().intra_op_parallelism_threads) context.ensure_initialized() with self.assertRaises(RuntimeError): config.set_intra_op_parallelism_threads(1) config.set_intra_op_parallelism_threads(10) @reset_eager def testInterOpParallelismThreads(self): config.set_inter_op_parallelism_threads(10) self.assertEqual(config.get_inter_op_parallelism_threads(), context.context().inter_op_parallelism_threads) context.ensure_initialized() with self.assertRaises(RuntimeError): config.set_inter_op_parallelism_threads(1) config.set_inter_op_parallelism_threads(10) @test_util.run_gpu_only @reset_eager def testSoftPlacement(self): if context.executing_eagerly(): self.assertTrue(config.get_soft_device_placement()) else: self.assertFalse(config.get_soft_device_placement()) def test_attr(): with ops.device('/device:GPU:0'): return test_ops.test_attr(T=dtypes.float32, name='test_attr') config.set_soft_device_placement(True) self.assertEqual(config.get_soft_device_placement(), True) self.assertEqual(config.get_soft_device_placement(), context.context().soft_device_placement) # Since soft placement is enabled, the test_attr operation should fallback # to CPU with pure eager execution as well as functions test_attr() def_function.function(test_attr)() config.set_soft_device_placement(False) self.assertEqual(config.get_soft_device_placement(), False) self.assertEqual(config.get_soft_device_placement(), context.context().soft_device_placement) # Since soft placement is disabled, the test_attr operation should fail on # GPU with pure eager execution as well as functions with self.assertRaises(errors.InvalidArgumentError): test_attr() with self.assertRaises(errors.InvalidArgumentError): def_function.function(test_attr)() @reset_eager def testLogDevicePlacement(self): self.assertFalse(context.get_log_device_placement()) context.set_log_device_placement(True) self.assertEqual(context.get_log_device_placement(), True) self.assertEqual(context.get_log_device_placement(), context.context().log_device_placement) context.set_log_device_placement(False) self.assertEqual(context.get_log_device_placement(), False) self.assertEqual(context.get_log_device_placement(), context.context().log_device_placement) context.ensure_initialized() # Changing the device placement should not throw an exception context.set_log_device_placement(True) @reset_eager def testEnableMlirBridge(self): # Default value of enable_mlir_bridge is false. self.assertFalse(context.context().config.experimental.enable_mlir_bridge) self.assertEqual( context.context().config.experimental.mlir_bridge_rollout, config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_UNSPECIFIED) # Tests enabling mlir bridge. config.enable_mlir_bridge() self.assertTrue(context.context().config.experimental.enable_mlir_bridge) self.assertEqual( context.context().config.experimental.mlir_bridge_rollout, config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_ENABLED) # Tests disabling mlir bridge. config.disable_mlir_bridge() self.assertFalse(context.context().config.experimental.enable_mlir_bridge) self.assertEqual( context.context().config.experimental.mlir_bridge_rollout, config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_DISABLED) @reset_eager def testResetMlirFlags(self): # Default value of enable_mlir_bridge is false. self.assertFalse(context.context().config.experimental.enable_mlir_bridge) self.assertEqual( context.context().config.experimental.mlir_bridge_rollout, config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_UNSPECIFIED) @test_util.run_gpu_only @reset_eager def testJit(self): self.assertEqual(config.get_optimizer_jit(), '') # the following function should cause Op fusion to occur. However, there is # unfortunately no straightforward way to ensure this. We will just have to # settle for creating a test that can trigger JIT. @def_function.function def fun(a, b): c = a * b d = c + a return d a = constant_op.constant([2., 2.]) b = constant_op.constant([2., 2.]) self.evaluate(fun(a, b)) config.set_optimizer_jit('autoclustering') self.assertEqual(config.get_optimizer_jit(), 'autoclustering') self.evaluate(fun(a, b)) config.set_optimizer_jit('') self.assertEqual(config.get_optimizer_jit(), '') self.evaluate(fun(a, b)) @parameterized.named_parameters( ('LayoutOptimizer', 'layout_optimizer'), ('ConstantFolding', 'constant_folding'), ('ShapeOptimization', 'shape_optimization'), ('Remapping', 'remapping'), ('ArithmeticOptimization', 'arithmetic_optimization'), ('DependencyOptimization', 'dependency_optimization'), ('LoopOptimization', 'loop_optimization'), ('FunctionOptimization', 'function_optimization'), ('DebugStripper', 'debug_stripper'), ('ScopedAllocatorOptimization', 'scoped_allocator_optimization'), ('ImplementationSelector', 'implementation_selector'), ('AutoMixedPrecision', 'auto_mixed_precision')) @reset_eager def testOptimizerToggleOption(self, field): # TODO(b/128531235): Improve testing of option options = config.get_optimizer_experimental_options() self.assertIsNone(options.get(field)) config.set_optimizer_experimental_options({field: True}) options[field] = True self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual(context.context().get_optimizer_experimental_options(), options) config.set_optimizer_experimental_options({field: False}) options[field] = False self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual(context.context().get_optimizer_experimental_options(), options) @parameterized.named_parameters( ('DisableModelPruning', 'disable_model_pruning'), ('DisableMetaOptimizer', 'disable_meta_optimizer')) @reset_eager def testOptimizerBoolOption(self, field): # TODO(b/128531235): Improve testing of option options = config.get_optimizer_experimental_options() self.assertFalse(options.get(field)) config.set_optimizer_experimental_options({field: True}) options[field] = True self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual(context.context().get_optimizer_experimental_options(), options) config.set_optimizer_experimental_options({field: False}) options[field] = False self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual(context.context().get_optimizer_experimental_options(), options) @test_util.run_gpu_only @reset_eager def testOptimizerToggleOptionPinToHost(self): options = config.get_optimizer_experimental_options() self.assertIsNone(options.get('pin_to_host_optimization')) @def_function.function def fun(): op = test_ops.device_placement_op() return op # Force optimizer to run for all graphs config.set_optimizer_experimental_options({'min_graph_nodes': -1}) options['min_graph_nodes'] = -1 # Since pin to host is disabled, the operation should go on GPU gpu = self.evaluate(fun()) self.assertIn(compat.as_bytes('GPU'), gpu) config.set_optimizer_experimental_options( {'pin_to_host_optimization': True}) options['pin_to_host_optimization'] = True self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual(context.context().get_optimizer_experimental_options(), options) # Since pin to host is enabled, the operation should go on CPU cpu = self.evaluate(fun()) self.assertIn(compat.as_bytes('CPU'), cpu) config.set_optimizer_experimental_options( {'pin_to_host_optimization': False}) options['pin_to_host_optimization'] = False self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual(context.context().get_optimizer_experimental_options(), options) # Since pin to host is disabled again, the operation should go on GPU gpu2 = self.evaluate(fun()) self.assertIn(compat.as_bytes('GPU'), gpu2) class DeviceTest(test.TestCase): @reset_eager def testPhysicalDevices(self): cpus = config.list_physical_devices('CPU') self.assertGreater(len(cpus), 0) if test_util.is_gpu_available(): gpus = config.list_physical_devices('GPU') self.assertGreater(len(gpus), 0) @reset_eager def testCpuMultiple(self): cpus = config.list_physical_devices('CPU') self.assertEqual(len(cpus), 1) config.set_logical_device_configuration(cpus[0], [ context.LogicalDeviceConfiguration(), context.LogicalDeviceConfiguration() ]) context.ensure_initialized() vcpus = config.list_logical_devices('CPU') self.assertEqual(len(vcpus), 2) with ops.device('/device:CPU:0'): a = constant_op.constant(1.0) self.evaluate(a) with ops.device('/device:CPU:1'): b = constant_op.constant(1.0) self.evaluate(b) with ops.device('/device:CPU:2'): c = constant_op.constant(1.0) self.evaluate(c) if test_util.is_gpu_available(): self.assertIn('GPU:0', c.device) else: self.assertIn('CPU:0', c.device) # Ensure we can place ops on each of the device names for vcpu in vcpus: with ops.device(vcpu.name): d = constant_op.constant(1.0) self.evaluate(d) # Modifying the CPU configuration is not supported with self.assertRaisesRegex(RuntimeError, 'cannot be modified'): config.set_logical_device_configuration(cpus[0], [ context.LogicalDeviceConfiguration(), context.LogicalDeviceConfiguration(), context.LogicalDeviceConfiguration() ]) # Setting the same CPU configuration is fine config.set_logical_device_configuration(cpus[0], [ context.LogicalDeviceConfiguration(), context.LogicalDeviceConfiguration() ]) @test_util.run_gpu_only @reset_eager def testGpuNone(self): config.set_soft_device_placement(False) gpus = config.list_physical_devices('GPU') self.assertGreater(len(gpus), 0) cpus = config.list_physical_devices('CPU') self.assertEqual(len(cpus), 1) self.assertEqual(len(config.get_visible_devices('CPU')), 1) self.assertGreater(len(config.get_visible_devices('GPU')), 0) self.assertEqual(len(config.get_visible_devices('XLA_GPU')), 0) config.set_visible_devices(cpus[0]) self.assertEqual(len(config.get_visible_devices('CPU')), 1) self.assertEqual(len(config.get_visible_devices('GPU')), 0) self.assertEqual(len(config.list_logical_devices('XLA_GPU')), 0) with self.assertRaisesRegex(errors.InvalidArgumentError, 'Could not satisfy'): with ops.device('/device:GPU:0'): a = array_ops.identity(1.0) self.evaluate(a) # Modifying the visible devices is not supported with self.assertRaisesRegex(RuntimeError, 'cannot be modified'): config.set_visible_devices(gpus) # Setting the same visible devices is fine config.set_visible_devices(cpus[0]) @reset_eager def testGpuMultiple(self): config.set_soft_device_placement(False) gpus = config.list_physical_devices('GPU') if len(gpus) < 2: self.skipTest('Need at least 2 GPUs') context.ensure_initialized() for i in range(0, len(gpus)): with ops.device('/device:GPU:' + str(i)): a = constant_op.constant(1.0) self.evaluate(a) with self.assertRaisesRegex(errors.InvalidArgumentError, 'Could not satisfy device specification'): with ops.device('/device:GPU:' + str(len(gpus))): a = constant_op.constant(1.0) self.evaluate(a) @reset_eager def testDeviceDetails(self): (cpu,) = config.list_physical_devices('CPU') details = config.get_device_details(cpu) self.assertEqual(details, {}) if not test_util.is_gpu_available(): return gpus = config.list_physical_devices('GPU') details = config.get_device_details(gpus[0]) self.assertIsInstance(details['device_name'], str) self.assertNotEmpty(details['device_name']) if test.is_built_with_rocm(): # AMD GPUs do not have a compute capability self.assertNotIn('compute_capability', details) else: cc = details['compute_capability'] self.assertIsInstance(cc, tuple) major, minor = cc self.assertGreater(major, 0) self.assertGreaterEqual(minor, 0) # Test GPU returned from get_visible_devices if len(gpus) > 2: config.set_visible_devices(gpus[1], 'GPU') (visible_gpu,) = config.get_visible_devices('GPU') details = config.get_device_details(visible_gpu) self.assertIsInstance(details['device_name'], str) @reset_eager def testDeviceDetailsErrors(self): logical_devices = config.list_logical_devices() with self.assertRaisesRegex(ValueError, 'must be a tf.config.PhysicalDevice'): config.get_device_details(logical_devices[0]) phys_dev = context.PhysicalDevice('/physical_device:CPU:100', 'CPU') with self.assertRaisesRegex( ValueError, 'The PhysicalDevice must be one obtained from ' 'calling `tf.config.list_physical_devices`'): config.get_device_details(phys_dev) @test_util.run_gpu_only @reset_eager def testVirtualGpu(self): config.set_soft_device_placement(False) gpus = config.list_physical_devices('GPU') self.assertNotEqual(len(gpus), 0) self.assertIsNone(config.get_logical_device_configuration(gpus[-1])) config.set_logical_device_configuration(gpus[-1], [ context.LogicalDeviceConfiguration(memory_limit=10), context.LogicalDeviceConfiguration(memory_limit=10) ]) self.assertEqual(len(config.get_logical_device_configuration(gpus[-1])), 2) logical_gpus = config.list_logical_devices('GPU') self.assertTrue(len(logical_gpus), len(gpus) + 1) for i in range(0, len(logical_gpus)): with ops.device('/device:GPU:' + str(i)): a = array_ops.identity(1.0) self.evaluate(a) with self.assertRaisesRegex(errors.InvalidArgumentError, 'Could not satisfy'): with ops.device('/device:GPU:' + str(len(logical_gpus))): a = array_ops.identity(1.0) self.evaluate(a) # Modifying the GPU configuration is not supported with self.assertRaisesRegex(RuntimeError, 'cannot be modified'): config.set_logical_device_configuration(gpus[-1], [ context.LogicalDeviceConfiguration(memory_limit=20), context.LogicalDeviceConfiguration(memory_limit=20) ]) with self.assertRaisesRegex(RuntimeError, 'cannot be modified'): config.set_logical_device_configuration(gpus[-1], [ context.LogicalDeviceConfiguration(memory_limit=10), context.LogicalDeviceConfiguration(memory_limit=10), context.LogicalDeviceConfiguration(memory_limit=10) ]) # Setting the same GPU configuration is fine config.set_logical_device_configuration(gpus[-1], [ context.LogicalDeviceConfiguration(memory_limit=10), context.LogicalDeviceConfiguration(memory_limit=10) ]) @test_util.run_gpu_only @reset_eager def testGpuGrowth(self): gpus = config.list_physical_devices('GPU') self.assertNotEqual(len(gpus), 0) self.assertIsNone(config.get_memory_growth(gpus[-1])) for gpu in gpus: config.set_memory_growth(gpu, True) c = context.context().config self.assertTrue(c.gpu_options.allow_growth) logical_gpus = config.list_logical_devices('GPU') self.assertTrue(len(logical_gpus), len(gpus)) # Modifying the GPU configuration is not supported with self.assertRaisesRegex(RuntimeError, 'cannot be modified'): for gpu in gpus: config.set_memory_growth(gpu, False) # Setting the same GPU configuration is fine for gpu in gpus: config.set_memory_growth(gpu, True) @test_util.run_gpu_or_tpu @reset_eager def testGetMemoryInfoBasic(self, device_type): with ops.device(f'{device_type}:0'): device = array_ops.zeros([]).backing_device info = config.get_memory_info(device) self.assertGreater(info['current'], 0) self.assertGreater(info['peak'], 0) self.assertEqual(info.keys(), {'current', 'peak'}) self.assertEqual(config.get_memory_usage(device), info['current']) @test_util.run_gpu_or_tpu @reset_eager def testGetMemoryUsageSubstring(self, device_type): info = config.get_memory_info(f'{device_type}:0') self.assertGreater(info['current'], 0) @reset_eager def testGetMemoryInfoCPU(self): if test_util.IsMklEnabled(): # TODO(gzmkl) work with Google team to address design issue in allocator.h self.skipTest('MklCPUAllocator does not throw exception. So skip test.') with self.assertRaisesRegex(ValueError, 'Allocator stats not available'): config.get_memory_info('CPU:0') with self.assertRaisesRegex(ValueError, 'Allocator stats not available'): config.get_memory_usage('CPU:0') @reset_eager def testGetMemoryInfoUnknownDevice(self): with self.assertRaisesRegex(ValueError, 'No matching devices found'): config.get_memory_info('unknown_device:0') with self.assertRaisesRegex(ValueError, 'No matching devices found'): config.get_memory_usage('unknown_device:0') @reset_eager def testGetMemoryInfoInvalidDeviceString(self): with self.assertRaisesRegex(ValueError, 'Failed parsing device name'): context.context().get_memory_info('GPU') with self.assertRaisesRegex(ValueError, 'Failed parsing device name'): context.context().get_memory_info('GPU:') with self.assertRaisesRegex(ValueError, 'Failed parsing device name'): context.context().get_memory_info('GPU:CPU') @test_util.run_gpu_or_tpu @reset_eager def testPeakMemoryUsage(self, device_type): device = f'{device_type}:0' with ops.device(device): x1 = array_ops.zeros((1000, 1000)) peak1 = config.get_memory_info(device)['peak'] self.assertGreaterEqual(peak1, 4 * 1000 * 1000) with ops.device(device): x2 = array_ops.ones((1000, 1000)) peak2 = config.get_memory_info(device)['peak'] self.assertGreaterEqual(peak2, peak1 + 4 * 1000 * 1000) del x1, x2 # With CPython, causes tensor memory to be immediately freed peak3 = config.get_memory_info(device)['peak'] self.assertGreaterEqual(peak3, peak2) self.assertGreaterEqual(peak3, config.get_memory_info(device)['current']) @test_util.run_gpu_or_tpu @reset_eager def testResetMemoryStats(self, device_type): device = f'{device_type}:0' with ops.device(device): x = array_ops.zeros((1000, 1000), dtype=dtypes.float32) config.reset_memory_stats(device) info1 = config.get_memory_info(device) self.assertGreaterEqual(info1['peak'], 4 * 1000 * 1000) self.assertGreaterEqual(info1['peak'], info1['current']) self.assertGreater(info1['current'], 0) del x # With CPython, causes tensor memory to be immediately freed config.reset_memory_stats(device) info2 = config.get_memory_info(device) self.assertLess(info2['peak'], info1['peak']) @reset_eager def testResetMemoryStatsCPU(self): if test_util.IsMklEnabled(): # TODO(gzmkl) work with Google team to address design issue in allocator.h self.skipTest('MklCPUAllocator does not throw exception. So skip test.') with self.assertRaisesRegex(ValueError, 'Cannot reset memory stats'): config.reset_memory_stats('CPU:0') @reset_eager def testResetMemoryStatsUnknownDevice(self): with self.assertRaisesRegex(ValueError, 'No matching devices found'): config.reset_memory_stats('unknown_device:0') @reset_eager def testResetMemoryStatsInvalidDeviceString(self): with self.assertRaisesRegex(ValueError, 'Failed parsing device name'): context.context().reset_memory_stats('GPU') with self.assertRaisesRegex(ValueError, 'Failed parsing device name'): context.context().reset_memory_stats('GPU:') with self.assertRaisesRegex(ValueError, 'Failed parsing device name'): context.context().reset_memory_stats('GPU:CPU') @test_util.run_gpu_only @reset_eager def testGpuInvalidConfig(self): gpus = config.list_physical_devices('GPU') self.assertNotEqual(len(gpus), 0) if len(gpus) > 1: # Assert if other GPUs were not configured config.set_memory_growth(gpus[0], True) with self.assertRaisesRegex(ValueError, 'cannot differ'): c = context.context().config # If we limit visibility to GPU 0, growth is fine config.set_visible_devices(gpus[0], 'GPU') c = context.context().config self.assertTrue(c.gpu_options.allow_growth) # Default setting for second GPU is False and works if we set visibility config.set_visible_devices(gpus[1], 'GPU') c = context.context().config self.assertFalse(c.gpu_options.allow_growth) # Growth now fails because all the GPUs are visible and not the same config.set_visible_devices(gpus, 'GPU') with self.assertRaisesRegex(ValueError, 'cannot differ'): c = context.context().config for gpu in gpus: config.set_memory_growth(gpu, True) c = context.context().config self.assertTrue(c.gpu_options.allow_growth) with self.assertRaisesRegex(ValueError, 'memory limit'): config.set_logical_device_configuration(gpus[-1], [ context.LogicalDeviceConfiguration(), context.LogicalDeviceConfiguration() ]) self.assertIsNone(config.get_logical_device_configuration(gpus[-1])) config.set_logical_device_configuration(gpus[-1], [ context.LogicalDeviceConfiguration(memory_limit=10), context.LogicalDeviceConfiguration(memory_limit=10) ]) c = context.context().config self.assertFalse(c.gpu_options.allow_growth) with self.assertRaisesRegex(ValueError, 'virtual devices'): config.set_memory_growth(gpus[-1], False) @test_util.run_gpu_only @reset_eager def testRemote(self): gpus = config.list_logical_devices('GPU') self.assertNotEqual(len(gpus), 0) context.ensure_initialized() gpus = config.list_logical_devices('GPU') self.assertNotEqual(len(gpus), 0) for gpu in gpus: self.assertIsNotNone(gpu.name) context.ensure_initialized() job_name = 'test' cluster_def = cluster_pb2.ClusterDef() job_def = cluster_def.job.add() job_def.name = job_name job_def.tasks[0] = 'localhost:0' server_def = tensorflow_server_pb2.ServerDef( cluster=cluster_def, job_name=job_name, task_index=0, protocol='grpc') context.set_server_def(server_def) gpus = config.list_logical_devices('GPU') for gpu in gpus: self.assertIsNotNone(gpu.name) @reset_eager def testV1CompatibilityDummyInvisibleDeviceList(self): gpus = config.list_physical_devices('GPU') if gpus: self.skipTest('Test requires no GPUs') # Ensure GPU options left untouched on CPU only environments context.context()._physical_devices = None context.context()._config = config_pb2.ConfigProto( gpu_options=config_pb2.GPUOptions(visible_device_list='0')) new_config = context.context().config self.assertEqual(new_config.gpu_options.visible_device_list, '0') @test_util.run_gpu_only @reset_eager def testV1Compatibility(self): # Ensure we set 1 CPU by default context.context()._config = config_pb2.ConfigProto() new_config = context.context().config self.assertEqual(new_config.device_count['CPU'], 1) context.context()._physical_devices = None # Ensure CPU is split context.context()._config = config_pb2.ConfigProto(device_count={'CPU': 2}) new_config = context.context().config self.assertEqual(new_config.device_count['CPU'], 2) context.context()._physical_devices = None # Handle empty visible device list context.context()._config = config_pb2.ConfigProto( gpu_options=config_pb2.GPUOptions(visible_device_list='')) gpus = config.list_physical_devices('GPU') gpu_count = len(gpus) new_config = context.context().config self.assertEqual(new_config.gpu_options.visible_device_list, ','.join(str(i) for i in range(len(gpus)))) context.context()._physical_devices = None # Handle invalid visible device list context.context()._config = config_pb2.ConfigProto( gpu_options=config_pb2.GPUOptions(visible_device_list=str(gpu_count))) with self.assertRaisesRegex(ValueError, 'Invalid visible device index'): gpus = config.list_physical_devices('GPU') new_config = context.context().config context.context()._physical_devices = None # Handle single visible device list context.context()._config = config_pb2.ConfigProto( gpu_options=config_pb2.GPUOptions( visible_device_list=str(gpu_count - 1))) gpus = config.list_physical_devices('GPU') new_config = context.context().config self.assertEqual(new_config.gpu_options.visible_device_list, str(gpu_count - 1)) context.context()._physical_devices = None def testConfigureCollectiveOps(self): context.context().configure_collective_ops( collective_leader='/job:worker/replica:0/task:0', scoped_allocator_enabled_ops=('CollectiveReduce',), use_nccl_communication=False, device_filters=['/job:worker/task:1']) new_config = context.context().config # Verify group leader self.assertEqual('/job:worker/replica:0/task:0', new_config.experimental.collective_group_leader) # Verify device filters. self.assertEqual(['/job:worker/task:1'], new_config.device_filters) # Verify rewrite options. new_rewrite_options = new_config.graph_options.rewrite_options self.assertEqual(rewriter_config_pb2.RewriterConfig.ON, new_rewrite_options.scoped_allocator_optimization) self.assertEqual(['CollectiveReduce'], new_rewrite_options.scoped_allocator_opts.enable_op) def testDeterminism(self): # This does not test any ops are deterministic, because that is tested by # many kernel tests. try: config.disable_op_determinism() self.assertFalse(config.is_op_determinism_enabled()) config.enable_op_determinism() self.assertTrue(config.is_op_determinism_enabled()) finally: config.disable_op_determinism() class TensorFloat32Test(test.TestCase): def tearDown(self): super(TensorFloat32Test, self).tearDown() config.enable_tensor_float_32_execution(True) def test_tensor_float_32_global_variable(self): self.assertTrue(config.tensor_float_32_execution_enabled()) self.assertTrue(test_ops.is_tensor_float32_enabled()) config.enable_tensor_float_32_execution(False) self.assertFalse(config.tensor_float_32_execution_enabled()) self.assertFalse(test_ops.is_tensor_float32_enabled()) config.enable_tensor_float_32_execution(True) self.assertTrue(config.tensor_float_32_execution_enabled()) self.assertTrue(test_ops.is_tensor_float32_enabled()) def _skip_if_tensor_float_32_unsupported(self): if not test_util.is_gpu_available( cuda_only=True, min_cuda_compute_capability=(8, 0)): self.skipTest('TensorFloat-32 requires an NVIDIA GPU with compute ' 'capability of at least 8.0') # Size of each dimension of matrices to test. cuBLAS does not use TF32 for # small matrices, so we must choose a large enough size to cause TF32 to be # used. DIM = 2 ** 10 def test_tensor_float_32_enabled(self): self._skip_if_tensor_float_32_unsupported() self.assertTrue(config.tensor_float_32_execution_enabled()) x = array_ops.fill((self.DIM, self.DIM), 1 + 2**-12) y = array_ops.ones((self.DIM, self.DIM)) out = math_ops.matmul(x, y) # In TensorFloat-32, each element of x is rounded to 1, so each output # element should be self.DIM. expected = array_ops.fill((self.DIM, self.DIM), float(self.DIM)) self.assertAllEqual(out, expected) def test_tensor_float_32_disabled(self): self._skip_if_tensor_float_32_unsupported() self.assertTrue(config.tensor_float_32_execution_enabled()) config.enable_tensor_float_32_execution(False) self.assertFalse(config.tensor_float_32_execution_enabled()) x = array_ops.fill((self.DIM, self.DIM), 1 + 2**-12) y = array_ops.ones((self.DIM, self.DIM)) out = math_ops.matmul(x, y) expected = array_ops.fill((self.DIM, self.DIM), self.DIM * (1 + 2**-12)) self.assertAllClose(out, expected, rtol=2**-13, atol=0) if __name__ == '__main__': ops.enable_eager_execution() test.main()