def open(device_id=None): """Initiate and return a NPU device handle""" if device_id is None: device_id = int(os.getenv("ASCEND_DEVICE_ID", '0')) with _npu_ctx_lock: if not isinstance(context.context(), _ContextWithDefaultDevice): ctx = _ContextWithDefaultDevice() ctx.ensure_initialized() context._set_context(ctx) _npu_device_instances.clear( ) # Global context has changed since last init npu if device_id in _npu_device_instances.keys(): logging.info('Npu instance on device %s already created', str(device_id)) return _npu_device_instances.get(device_id) if len(_npu_device_instances): raise RuntimeError( 'Failed create npu instance on device {} as existed instance on {}' ''.format(device_id, list(_npu_device_instances.keys()))) global_kw_options = global_options().as_dict() workers_num = int(os.getenv('RANK_SIZE', '1')) if workers_num > 1: env_rank_table = os.getenv("RANK_TABLE_FILE") env_worker_id = os.getenv('RANK_ID') if not env_rank_table: raise RuntimeError( 'You must specify a rank table file by set env RANK_TABLE_FILE in distribution mode' ) if not env_worker_id: raise RuntimeError( 'You must specify rank id by set env RANK_ID in distribution mode' ) global_kw_options['_distribute.rank_table'] = env_rank_table global_kw_options['_distribute.rank_id'] = env_worker_id device_options = {} error_message = _npu_device_backends.Open(context.context()._handle, NPU, device_id, global_kw_options, device_options) if error_message: raise RuntimeError("Failed open npu device %s : %s" % (str(device_id), error_message)) if workers_num > 1: from hccl.manage.api import get_rank_id worker_id = get_rank_id() else: worker_id = 0 _npu_device_instances[device_id] = NpuDeviceHandle( context.context(), device_id, device_options, workers_num, worker_id) return _npu_device_instances[device_id]
def testSilentCopy(self): # Temporarily replace the context # pylint: disable=protected-access old_context = context.context() context._set_context(context.Context()) try: config.set_device_policy('silent') cpu_tensor = constant_op.constant(1.0) gpu_tensor = cpu_tensor.gpu() self.assertAllEqual(cpu_tensor + gpu_tensor, 2.0) finally: context._set_context(old_context)
def testCrossContextTensorCache(self): old_context = context.context() old_x = constant_op.constant(9.5) context._set_context(context.Context()) try: new_x = constant_op.constant(9.5) self.assertEqual(new_x.numpy(), 9.5) finally: context._set_context(old_context) self.assertEqual(old_x.numpy(), 9.5)
def testSoftPlacement(self): # Temporarily replace the context # pylint: disable=protected-access old_context = context.context() context._set_context(context.Context()) try: config.set_device_policy('silent') config.set_soft_device_placement(True) cpu_tensor = constant_op.constant(1.0) result = cpu_tensor + cpu_tensor self.assertEqual(result.device, '/job:localhost/replica:0/task:0/device:GPU:0') finally: context._set_context(old_context)
def testSoftPlacement(self): # Temporarily replace the context # pylint: disable=protected-access old_context = context.context() context._set_context(context.Context()) try: config.set_device_policy('silent') config.set_soft_device_placement(True) # Avoid the TensorHandle cache hit. # TODO(b/169790439): include Context to the TensorHandle cache. cpu_tensor = constant_op.constant(1.1) result = cpu_tensor + cpu_tensor self.assertEqual(result.device, '/job:localhost/replica:0/task:0/device:GPU:0') finally: context._set_context(old_context)
def testContextIsDestroyedAfterTensors(self): # Create a new context new_context = context.Context() weak_c = weakref.ref(new_context) new_context.ensure_initialized() # Create a tensor with the new context as default. # Make sure to restore the original context. original_context = context.context() try: context._set_context(new_context) # Use a 2D tensor so that it is not cached. tensor1 = constant_op.constant([[3.]]) # Produce a tensor as an operation output. This uses a different code path # from tensors created from Python. tensor2 = tensor1 * tensor1 context._set_context(original_context) except: context._set_context(original_context) raise # Deleting our context reference should not delete the underlying object. del new_context self.assertIsNot(weak_c(), None) # Deleting the first tensor should not delete the context since there is # another tensor. del tensor1 self.assertIsNot(weak_c(), None) # Deleting the last tensor should result in deleting its context. del tensor2 self.assertIs(weak_c(), None)