1 /*
2  * Copyright (C) 2024 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.training;
18 
19 import static com.android.federatedcompute.services.stats.FederatedComputeStatsLog.FEDERATED_COMPUTE_TRAINING_EVENT_REPORTED__KIND__TRAIN_ELIGIBILITY_EVAL_COMPUTATION_COMPLETED;
20 import static com.android.federatedcompute.services.stats.FederatedComputeStatsLog.FEDERATED_COMPUTE_TRAINING_EVENT_REPORTED__KIND__TRAIN_ELIGIBILITY_EVAL_COMPUTATION_ELIGIBLE;
21 import static com.android.federatedcompute.services.stats.FederatedComputeStatsLog.FEDERATED_COMPUTE_TRAINING_EVENT_REPORTED__KIND__TRAIN_ELIGIBILITY_EVAL_COMPUTATION_ERROR_EXAMPLE_ITERATOR;
22 import static com.android.federatedcompute.services.stats.FederatedComputeStatsLog.FEDERATED_COMPUTE_TRAINING_EVENT_REPORTED__KIND__TRAIN_ELIGIBILITY_EVAL_COMPUTATION_STARTED;
23 
24 import android.content.Context;
25 import android.federatedcompute.aidl.IExampleStoreIterator;
26 import android.federatedcompute.aidl.IExampleStoreService;
27 import android.os.SystemClock;
28 
29 import com.android.federatedcompute.internal.util.LogUtil;
30 import com.android.federatedcompute.services.common.ExampleStats;
31 import com.android.federatedcompute.services.common.TrainingEventLogger;
32 import com.android.federatedcompute.services.data.FederatedTrainingTask;
33 import com.android.federatedcompute.services.data.FederatedTrainingTaskDao;
34 import com.android.federatedcompute.services.data.TaskHistory;
35 import com.android.federatedcompute.services.examplestore.ExampleStoreServiceProvider;
36 import com.android.federatedcompute.services.training.util.EligibilityResult;
37 import com.android.internal.annotations.VisibleForTesting;
38 
39 import com.google.internal.federated.plan.ExampleSelector;
40 import com.google.ondevicepersonalization.federatedcompute.proto.DataAvailabilityPolicy;
41 import com.google.ondevicepersonalization.federatedcompute.proto.EligibilityPolicyEvalSpec;
42 import com.google.ondevicepersonalization.federatedcompute.proto.EligibilityTaskInfo;
43 import com.google.ondevicepersonalization.federatedcompute.proto.MinimumSeparationPolicy;
44 
45 /** Runs eligibility evaluation and decide if device is qualified for each task. */
46 public class EligibilityDecider {
47     private static final String TAG = EligibilityDecider.class.getSimpleName();
48     private final FederatedTrainingTaskDao mTaskDao;
49     private final ExampleStoreServiceProvider mExampleStoreServiceProvider;
50 
51     @VisibleForTesting
EligibilityDecider( FederatedTrainingTaskDao taskDao, ExampleStoreServiceProvider exampleStoreServiceProvider)52     EligibilityDecider(
53             FederatedTrainingTaskDao taskDao,
54             ExampleStoreServiceProvider exampleStoreServiceProvider) {
55         mExampleStoreServiceProvider = exampleStoreServiceProvider;
56         mTaskDao = taskDao;
57     }
58 
EligibilityDecider(FederatedTrainingTaskDao taskDao)59     public EligibilityDecider(FederatedTrainingTaskDao taskDao) {
60         this(taskDao, new ExampleStoreServiceProvider());
61     }
62 
63     /**
64      * Computes the eligibility of the client for the given tasks in the population eligibility
65      * spec. Returns true if device is eligible to execute this task.
66      */
computeEligibility( FederatedTrainingTask task, String taskId, EligibilityTaskInfo eligibilityTaskInfo, Context context, TrainingEventLogger trainingEventLogger, ExampleSelector exampleSelector)67     public EligibilityResult computeEligibility(
68             FederatedTrainingTask task,
69             String taskId,
70             EligibilityTaskInfo eligibilityTaskInfo,
71             Context context,
72             TrainingEventLogger trainingEventLogger,
73             ExampleSelector exampleSelector) {
74         boolean eligible = true;
75         ExampleStats exampleStats = new ExampleStats();
76         trainingEventLogger.logEventKind(
77                 FEDERATED_COMPUTE_TRAINING_EVENT_REPORTED__KIND__TRAIN_ELIGIBILITY_EVAL_COMPUTATION_STARTED);
78         EligibilityResult.Builder result = new EligibilityResult.Builder();
79         for (EligibilityPolicyEvalSpec policyEvalSpec :
80                 eligibilityTaskInfo.getEligibilityPoliciesList()) {
81             switch (policyEvalSpec.getPolicyTypeCase()) {
82                 case MIN_SEP_POLICY:
83                     eligible =
84                             computePerTaskMinSeparation(
85                                     policyEvalSpec.getMinSepPolicy(),
86                                     task.populationName(),
87                                     taskId,
88                                     task.jobId());
89                     break;
90                 case DATA_AVAILABILITY_POLICY:
91                     eligible =
92                             computePerTaskDataAvailability(
93                                     task,
94                                     policyEvalSpec.getDataAvailabilityPolicy(),
95                                     taskId,
96                                     context,
97                                     exampleStats,
98                                     trainingEventLogger,
99                                     result,
100                                     exampleSelector);
101                     break;
102                 default:
103                     throw new IllegalStateException(
104                             String.format("Unsupported policy %s", policyEvalSpec.getId()));
105             }
106             // Device has to meet all eligibility policies in order to execute task.
107             if (!eligible) {
108                 result.setEligible(false);
109                 break;
110             }
111         }
112         // Always record eligibility task complete event. To calculate not eligible tasks, it
113         // is (EVAL_COMPUTATION_COMPLETED - EVAL_COMPUTATION_ELIGIBLE).
114         trainingEventLogger.logEventWithExampleStats(
115                 FEDERATED_COMPUTE_TRAINING_EVENT_REPORTED__KIND__TRAIN_ELIGIBILITY_EVAL_COMPUTATION_COMPLETED,
116                 exampleStats);
117         EligibilityResult eligibilityResult = result.setEligible(eligible).build();
118 
119         if (eligible) {
120             trainingEventLogger.logEventKind(
121                     FEDERATED_COMPUTE_TRAINING_EVENT_REPORTED__KIND__TRAIN_ELIGIBILITY_EVAL_COMPUTATION_ELIGIBLE);
122             return eligibilityResult;
123         }
124 
125         // If device is not eligible, we should unbind from ExampleStore if needed.
126         if (eligibilityResult.getExampleStoreIterator() != null) {
127             mExampleStoreServiceProvider.unbindFromExampleStoreService();
128         }
129         return new EligibilityResult.Builder().setEligible(false).build();
130     }
131 
computePerTaskMinSeparation( MinimumSeparationPolicy minSepPolicy, String populationName, String taskId, int jobId)132     private boolean computePerTaskMinSeparation(
133             MinimumSeparationPolicy minSepPolicy, String populationName, String taskId, int jobId) {
134         TaskHistory taskHistory = mTaskDao.getLatestTaskHistory(jobId, populationName, taskId);
135         // Treat null as the task never run before, then device is qualified.
136         if (taskHistory == null) {
137             LogUtil.d(
138                     TAG,
139                     "population name %s task id %s job id %d doesn't have TaskHistory record.",
140                     populationName,
141                     taskId,
142                     jobId);
143             return true;
144         }
145         return minSepPolicy.getMinimumSeparation()
146                 <= minSepPolicy.getCurrentIndex() - taskHistory.getContributionRound();
147     }
148 
computePerTaskDataAvailability( FederatedTrainingTask task, DataAvailabilityPolicy dataAvailabilityPolicy, String taskId, Context context, ExampleStats exampleStats, TrainingEventLogger logger, EligibilityResult.Builder result, ExampleSelector exampleSelector)149     private boolean computePerTaskDataAvailability(
150             FederatedTrainingTask task,
151             DataAvailabilityPolicy dataAvailabilityPolicy,
152             String taskId,
153             Context context,
154             ExampleStats exampleStats,
155             TrainingEventLogger logger,
156             EligibilityResult.Builder result,
157             ExampleSelector exampleSelector) {
158         try {
159             long callStartTimeNanos = SystemClock.elapsedRealtimeNanos();
160             IExampleStoreService exampleStoreService =
161                     mExampleStoreServiceProvider.getExampleStoreService(
162                             task.appPackageName(), context);
163             if (exampleStoreService == null) {
164                 logger.logEventKind(
165                         FEDERATED_COMPUTE_TRAINING_EVENT_REPORTED__KIND__TRAIN_ELIGIBILITY_EVAL_COMPUTATION_ERROR_EXAMPLE_ITERATOR);
166                 LogUtil.e(
167                         TAG,
168                         "Failed to compute DataAvailabilityPolicy due to bind ExampleStore"
169                                 + " failure %s %s %s",
170                         task.appPackageName(),
171                         task.populationName(),
172                         taskId);
173                 return false;
174             }
175             exampleStats.mBindToExampleStoreLatencyNanos.addAndGet(
176                     SystemClock.elapsedRealtimeNanos() - callStartTimeNanos);
177             callStartTimeNanos = SystemClock.elapsedRealtimeNanos();
178             IExampleStoreIterator iterator =
179                     mExampleStoreServiceProvider.getExampleIterator(
180                             exampleStoreService,
181                             task,
182                             taskId,
183                             dataAvailabilityPolicy.getMinExampleCount(),
184                             exampleSelector);
185             if (iterator == null) {
186                 LogUtil.d(
187                         TAG,
188                         "Failed to compute DataAvailabilityPolicy due to iterator is null "
189                                 + "%s %s %s",
190                         task.appPackageName(),
191                         task.populationName(),
192                         taskId);
193                 logger.logEventKind(
194                         FEDERATED_COMPUTE_TRAINING_EVENT_REPORTED__KIND__TRAIN_ELIGIBILITY_EVAL_COMPUTATION_ERROR_EXAMPLE_ITERATOR);
195                 mExampleStoreServiceProvider.unbindFromExampleStoreService();
196                 return false;
197             }
198             result.setExampleStoreIterator(iterator);
199             exampleStats.mStartQueryLatencyNanos.addAndGet(
200                     SystemClock.elapsedRealtimeNanos() - callStartTimeNanos);
201             return true;
202         } catch (Exception e) {
203             mExampleStoreServiceProvider.unbindFromExampleStoreService();
204             LogUtil.e(TAG, e, "Failed to compute DataAvailabilityPolicy");
205             return false;
206         }
207     }
208 }
209