• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Tests for device placement."""
16
17from absl.testing import parameterized
18
19from tensorflow.python.eager import context
20from tensorflow.python.eager import def_function
21from tensorflow.python.eager import remote
22from tensorflow.python.eager import test
23from tensorflow.python.framework import config
24from tensorflow.python.framework import constant_op
25from tensorflow.python.framework import dtypes
26from tensorflow.python.framework import errors
27from tensorflow.python.framework import ops
28from tensorflow.python.framework import test_util
29from tensorflow.python.ops import array_ops
30from tensorflow.python.ops import math_ops
31from tensorflow.python.ops import random_ops
32
33
34class SoftDevicePlacementTest(test.TestCase, parameterized.TestCase):
35
36  def setUp(self):
37    super(SoftDevicePlacementTest, self).setUp()
38    context._reset_context()
39    context.ensure_initialized()
40    config.set_soft_device_placement(enabled=True)
41    context.context().log_device_placement = True
42
43  @test_util.run_gpu_only
44  def testDefaultPlacement(self):
45    a = constant_op.constant(1)
46    b = constant_op.constant(2)
47    c = a + b
48    with ops.device('CPU'):
49      d = a + b
50    self.assertIn('GPU', c.device)
51    self.assertIn('CPU', d.device)
52
53  @test_util.run_gpu_only
54  def testUnsupportedDevice(self):
55    a = constant_op.constant(1)
56    b = constant_op.constant(2)
57    s = constant_op.constant(list('hello world'))
58    with ops.device('GPU:0'):
59      c = a + b
60      t = s[a]
61    self.assertIn('GPU:0', c.device)
62    self.assertIn('CPU', t.device)
63
64  @test_util.run_gpu_only
65  def testUnknownDevice(self):
66    a = constant_op.constant(1)
67    b = constant_op.constant(2)
68    with ops.device('GPU:42'):
69      c = a + b
70    self.assertIn('GPU:0', c.device)
71
72  def testNoGpu(self):
73    if test_util.is_gpu_available():
74      # CPU only test.
75      return
76    a = constant_op.constant(1)
77    b = constant_op.constant(2)
78    c = a + b
79    with ops.device('GPU'):
80      d = a + b
81    self.assertIn('CPU', c.device)
82    self.assertIn('CPU', d.device)
83
84  @test_util.run_gpu_only
85  def testSoftPlacedGPU(self):
86    a = constant_op.constant(1)
87    b = constant_op.constant(2)
88    with ops.device('GPU:110'):
89      c = a + b
90    self.assertIn('GPU:0', c.device)
91
92  @test_util.run_gpu_only
93  def testNestedDeviceScope(self):
94    a = constant_op.constant(1)
95    b = constant_op.constant(2)
96    with ops.device('CPU:0'):
97      with ops.device('GPU:42'):
98        c = a + b
99    # We don't support nested device placement right now.
100    self.assertIn('GPU:0', c.device)
101
102  @parameterized.named_parameters(('float', 1.0, None),
103                                  ('int32', [1], dtypes.int32),
104                                  ('string', ['a'], None))
105  def testSoftPlacedCPUConstant(self, value, dtype):
106    if test_util.is_gpu_available():
107      self.skipTest('CPU only test')
108    with ops.device('GPU:0'):
109      a = constant_op.constant(value, dtype=dtype)
110    self.assertIn('CPU:0', a.device)
111    self.assertIn('CPU:0', a.backing_device)
112
113  def testPlacedToDeviceInFunction(self):
114
115    @def_function.function
116    def f():
117      a = random_ops.random_uniform([32, 32])
118      return math_ops.matmul(a, a)
119
120    gpus = config.list_physical_devices('GPU')
121    if not gpus:
122      self.assertIn('CPU:0', f().device)
123    else:
124      self.assertIn('GPU:0', f().device)
125
126  @test_util.disable_tfrt('b/173726713: Support properly inserting device at '
127                          'tf_to_corert lowering.')
128  def testUnknownDeviceInFunction(self):
129
130    @def_function.function
131    def f():
132      with ops.device('GPU:42'):
133        # With placer, the unknown GPU:42 will be replaced with GPU:0.
134        a = constant_op.constant(1) + constant_op.constant(2)
135      return a + constant_op.constant(2)
136
137    gpus = config.list_physical_devices('GPU')
138    if not gpus:
139      self.assertIn('CPU:0', f().device)
140    else:
141      self.assertIn('GPU:0', f().device)
142
143
144class HardDevicePlacementTest(test.TestCase, parameterized.TestCase):
145
146  def setUp(self):
147    super(HardDevicePlacementTest, self).setUp()
148    context._reset_context()
149    config.set_soft_device_placement(enabled=False)
150    context.context().log_device_placement = True
151    cpus = context.context().list_physical_devices('CPU')
152    # Set 2 virtual CPUs
153    context.context().set_logical_device_configuration(cpus[0], [
154        context.LogicalDeviceConfiguration(),
155        context.LogicalDeviceConfiguration()
156    ])
157    self.assertEqual(config.get_soft_device_placement(), False)
158    self.assertEqual(context.context().soft_device_placement, False)
159
160  @test_util.run_gpu_only
161  def testIdentityCanCopy(self):
162    config.set_device_policy('explicit')
163    with ops.device('CPU:0'):
164      x = constant_op.constant(1.0)
165      self.assertIn('CPU:0', x.device)
166      self.assertIn('CPU:0', x.backing_device)
167    with ops.device('GPU:0'):
168      y = array_ops.identity(x)
169      self.assertIn('GPU:0', y.device)
170      self.assertIn('GPU:0', y.backing_device)
171
172  @test_util.run_gpu_only
173  def testSimpleConstantsExplicitGPU(self):
174    config.set_device_policy('explicit')
175    with ops.device('GPU:0'):
176      self.assertAllClose(1., array_ops.ones([]))
177      self.assertAllClose(0., array_ops.zeros([]))
178      self.assertAllClose([1.], array_ops.fill([1], 1.))
179
180  def testSimpleConstantsExplicitCPU(self):
181    config.set_device_policy('explicit')
182    with ops.device('CPU:1'):
183      self.assertAllClose(1., array_ops.ones([]))
184      self.assertAllClose(0., array_ops.zeros([]))
185      self.assertAllClose([1.], array_ops.fill([1], 1.))
186      self.assertAllClose(2., constant_op.constant(1.) * 2.)
187
188  @parameterized.named_parameters(
189      ('float_cpu0', 'CPU:0', 1.0, None),
190      ('int32_cpu0', 'CPU:0', [1], dtypes.int32),
191      ('string_cpu0', 'CPU:0', ['a'], None),
192      ('float_cpu1', 'CPU:1', 1.0, None),
193      ('int32_cpu1', 'CPU:1', [1], dtypes.int32),
194      ('string_cpu1', 'CPU:1', ['a'], None),
195  )
196  def testHardPlacedCPUConstant(self, device, value, dtype):
197    with ops.device(device):
198      a = constant_op.constant(value, dtype=dtype)
199      self.assertIn(device, a.device)
200
201  @parameterized.named_parameters(
202      ('float', 'GPU:0', 1.0, None),
203      ('int32', 'GPU:0', [1], dtypes.int32),
204      ('string', 'GPU:0', ['a'], None),
205  )
206  def testHardPlacedGPUConstant(self, device, value, dtype):
207    if not test_util.is_gpu_available():
208      self.skipTest('Test requires a GPU')
209    with ops.device(device):
210      a = constant_op.constant(value, dtype=dtype)
211      self.assertIn(device, a.device)
212      if a.dtype == dtypes.float32:
213        self.assertIn(device, a.backing_device)
214
215
216class ClusterPlacementTest(test.TestCase):
217
218  def setUp(self):
219    super(ClusterPlacementTest, self).setUp()
220    context._reset_context()
221    config.set_soft_device_placement(enabled=True)
222    context.context().log_device_placement = True
223    workers, _ = test_util.create_local_cluster(2, 0)
224    remote.connect_to_remote_host([workers[0].target, workers[1].target])
225
226  @test_util.disable_tfrt('remote host not supported yet.')
227  def testNotFullySpecifiedTask(self):
228    a = constant_op.constant(1)
229    b = constant_op.constant(2)
230    with ops.device('/job:worker'):
231      c = a + b
232    self.assertIn('/job:worker/replica:0/task:0', c.device)
233
234  @test_util.disable_tfrt('remote host not supported yet.')
235  def testRemoteUnknownDevice(self):
236    a = constant_op.constant(1)
237    b = constant_op.constant(2)
238    # Right now we don't support soft device place on remote worker.
239    with self.assertRaises(errors.InvalidArgumentError) as cm:
240      with ops.device('/job:worker/replica:0/task:0/device:GPU:42'):
241        c = a + b
242        del c
243      self.assertIn('unknown device', cm.exception.message)
244
245  @test_util.disable_tfrt('remote host not supported yet.')
246  def testUnknownDeviceInFunctionReturnUnknowDevice(self):
247
248    @def_function.function
249    def f():
250      with ops.device('GPU:42'):
251        return constant_op.constant(1) + constant_op.constant(2)
252
253    gpus = config.list_physical_devices('GPU')
254    if not gpus:
255      self.assertIn('CPU:0', f().device)
256    else:
257      self.assertIn('GPU:0', f().device)
258
259  @test_util.disable_tfrt('remote host not supported yet.')
260  def testUnknownDeviceInFunction(self):
261
262    @def_function.function
263    def f():
264      with ops.device('GPU:42'):
265        a = constant_op.constant(1) + constant_op.constant(2)
266      return a + constant_op.constant(2)
267
268    gpus = config.list_physical_devices('GPU')
269    if not gpus:
270      self.assertIn('CPU:0', f().device)
271    else:
272      self.assertIn('GPU:0', f().device)
273
274
275if __name__ == '__main__':
276  test.main()
277