def _clear_executor(self): request = executor_pb2.ClearExecutorRequest() try: _request(self._stub.ClearExecutor, request) except (grpc.RpcError, execution_context.RetryableError): logging.debug('RPC error caught during attempt to clear state on the ' 'server; this likely indicates a broken connection, and ' 'therefore there is no state to clear.') return
def test_executor_service_raises_after_cleanup_without_configuration(self): ex_factory = executor_stacks.ResourceManagingExecutorFactory( lambda _: eager_tf_executor.EagerTFExecutor()) env = TestEnv(ex_factory) env.stub.ClearExecutor(executor_pb2.ClearExecutorRequest()) value_proto, _ = executor_serialization.serialize_value( tf.constant(10.0).numpy(), tf.float32) with self.assertRaises(grpc.RpcError): env.stub.CreateValue( executor_pb2.CreateValueRequest(value=value_proto))
def test_clear_executor_calls_cleanup(self, mock_cleanup): ex_factory = executor_stacks.ResourceManagingExecutorFactory( lambda _: eager_tf_executor.EagerTFExecutor()) env = TestEnv(ex_factory) env.stub.ClearExecutor(executor_pb2.ClearExecutorRequest()) mock_cleanup.assert_called_once()