1 /* 2 * Copyright (C) 2023 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package com.android.federatedcompute.services.examplestore; 18 19 import static android.federatedcompute.common.ClientConstants.EXTRA_EXAMPLE_ITERATOR_RESULT; 20 import static android.federatedcompute.common.ClientConstants.EXTRA_TASK_ID; 21 import static android.federatedcompute.common.ClientConstants.STATUS_INTERNAL_ERROR; 22 23 import android.annotation.NonNull; 24 import android.federatedcompute.ExampleStoreIterator; 25 import android.federatedcompute.ExampleStoreService; 26 import android.os.Bundle; 27 28 import com.google.common.collect.ImmutableList; 29 import com.google.protobuf.ByteString; 30 31 import org.tensorflow.example.BytesList; 32 import org.tensorflow.example.Example; 33 import org.tensorflow.example.Feature; 34 import org.tensorflow.example.Features; 35 36 import java.util.Iterator; 37 import java.util.List; 38 39 /** A sample ExampleStoreService implementation. */ 40 public class SampleExampleStoreService extends ExampleStoreService { 41 private static final String EXPECTED_TASK_NAME = "test_task"; 42 private static final Example EXAMPLE_PROTO_1 = 43 Example.newBuilder() 44 .setFeatures( 45 Features.newBuilder() 46 .putFeature( 47 "feature1", 48 Feature.newBuilder() 49 .setBytesList( 50 BytesList.newBuilder() 51 .addValue( 52 ByteString.copyFromUtf8( 53 "f1_value1"))) 54 .build())) 55 .build(); 56 57 @Override startQuery(@onNull Bundle params, @NonNull QueryCallback callback)58 public void startQuery(@NonNull Bundle params, @NonNull QueryCallback callback) { 59 String collection = params.getString(EXTRA_TASK_ID); 60 if (!collection.equals(EXPECTED_TASK_NAME)) { 61 callback.onStartQueryFailure(STATUS_INTERNAL_ERROR); 62 return; 63 } 64 callback.onStartQuerySuccess( 65 new ListExampleStoreIterator(ImmutableList.of(EXAMPLE_PROTO_1))); 66 } 67 68 @Override checkCallerPermission()69 protected boolean checkCallerPermission() { 70 return true; 71 } 72 73 /** 74 * A simple {@link ExampleStoreIterator} that returns the contents of the {@link List} it's 75 * constructed with. 76 */ 77 private static class ListExampleStoreIterator implements ExampleStoreIterator { 78 private final Iterator<Example> mExampleIterator; 79 ListExampleStoreIterator(List<Example> examples)80 ListExampleStoreIterator(List<Example> examples) { 81 mExampleIterator = examples.iterator(); 82 } 83 84 @Override next(IteratorCallback callback)85 public synchronized void next(IteratorCallback callback) { 86 if (mExampleIterator.hasNext()) { 87 Bundle bundle = new Bundle(); 88 bundle.putByteArray( 89 EXTRA_EXAMPLE_ITERATOR_RESULT, mExampleIterator.next().toByteArray()); 90 callback.onIteratorNextSuccess(bundle); 91 } else { 92 callback.onIteratorNextSuccess(null); 93 } 94 } 95 96 @Override close()97 public void close() {} 98 } 99 } 100