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.data; 18 19 import static com.android.federatedcompute.services.data.FederatedTraningTaskContract.FEDERATED_TRAINING_TASKS_TABLE; 20 21 import android.annotation.NonNull; 22 import android.annotation.Nullable; 23 import android.content.ContentValues; 24 import android.database.Cursor; 25 import android.database.sqlite.SQLiteDatabase; 26 27 import com.android.federatedcompute.services.data.FederatedTraningTaskContract.FederatedTrainingTaskColumns; 28 import com.android.federatedcompute.services.data.fbs.TrainingConstraints; 29 import com.android.federatedcompute.services.data.fbs.TrainingIntervalOptions; 30 31 import com.google.auto.value.AutoValue; 32 33 import java.nio.ByteBuffer; 34 import java.util.ArrayList; 35 import java.util.List; 36 37 /** Contains the details of a training task. */ 38 @AutoValue 39 public abstract class FederatedTrainingTask { 40 private static final String TAG = "FederatedTrainingTask"; 41 42 /** @return client app package name */ appPackageName()43 public abstract String appPackageName(); 44 45 /** 46 * @return the ID to use for the JobScheduler job that will run the training for this session. 47 */ jobId()48 public abstract int jobId(); 49 50 /** @return the population name to uniquely identify the training job by. */ populationName()51 public abstract String populationName(); 52 53 /** 54 * @return the byte array of training interval including scheduling mode and minimum latency. 55 * The byte array is constructed from TrainingConstraints flatbuffer. 56 */ 57 @Nullable 58 @SuppressWarnings("mutable") intervalOptions()59 public abstract byte[] intervalOptions(); 60 61 /** @return the training interval including scheduling mode and minimum latency. */ 62 @Nullable getTrainingIntervalOptions()63 public final TrainingIntervalOptions getTrainingIntervalOptions() { 64 if (intervalOptions() == null) { 65 return null; 66 } 67 return TrainingIntervalOptions.getRootAsTrainingIntervalOptions( 68 ByteBuffer.wrap(intervalOptions())); 69 } 70 71 /** @return the time the task was originally created. */ creationTime()72 public abstract Long creationTime(); 73 74 /** @return the time the task was last scheduled. */ lastScheduledTime()75 public abstract Long lastScheduledTime(); 76 77 /** @return the start time of the task's last run. */ 78 @Nullable lastRunStartTime()79 public abstract Long lastRunStartTime(); 80 81 /** @return the end time of the task's last run. */ 82 @Nullable lastRunEndTime()83 public abstract Long lastRunEndTime(); 84 85 /** @return the earliest time to run the task by. */ earliestNextRunTime()86 public abstract Long earliestNextRunTime(); 87 88 /** 89 * @return the byte array of training constraints that should apply to this task. The byte array 90 * is constructed from TrainingConstraints flatbuffer. 91 */ 92 @SuppressWarnings("mutable") constraints()93 public abstract byte[] constraints(); 94 95 /** @return the training constraints that should apply to this task. */ getTrainingConstraints()96 public final TrainingConstraints getTrainingConstraints() { 97 return TrainingConstraints.getRootAsTrainingConstraints(ByteBuffer.wrap(constraints())); 98 } 99 100 /** @return the reason to schedule the task. */ schedulingReason()101 public abstract int schedulingReason(); 102 103 /** Builder for {@link FederatedTrainingTask} */ 104 @AutoValue.Builder 105 public abstract static class Builder { 106 /** Set client application package name. */ appPackageName(String appPackageName)107 public abstract Builder appPackageName(String appPackageName); 108 109 /** Set job scheduler Id. */ jobId(int jobId)110 public abstract Builder jobId(int jobId); 111 112 /** Set population name which uniquely identify the job. */ populationName(String populationName)113 public abstract Builder populationName(String populationName); 114 115 /** Set the training interval including scheduling mode and minimum latency. */ 116 @SuppressWarnings("mutable") intervalOptions(@ullable byte[] intervalOptions)117 public abstract Builder intervalOptions(@Nullable byte[] intervalOptions); 118 119 /** Set the time the task was originally created. */ creationTime(Long creationTime)120 public abstract Builder creationTime(Long creationTime); 121 122 /** Set the time the task was last scheduled. */ lastScheduledTime(Long lastScheduledTime)123 public abstract Builder lastScheduledTime(Long lastScheduledTime); 124 125 /** Set the start time of the task's last run. */ lastRunStartTime(@ullable Long lastRunStartTime)126 public abstract Builder lastRunStartTime(@Nullable Long lastRunStartTime); 127 128 /** Set the end time of the task's last run. */ lastRunEndTime(@ullable Long lastRunEndTime)129 public abstract Builder lastRunEndTime(@Nullable Long lastRunEndTime); 130 131 /** Set the earliest time to run the task by. */ earliestNextRunTime(Long earliestNextRunTime)132 public abstract Builder earliestNextRunTime(Long earliestNextRunTime); 133 134 /** Set the training constraints that should apply to this task. */ 135 @SuppressWarnings("mutable") constraints(byte[] constraints)136 public abstract Builder constraints(byte[] constraints); 137 138 /** Set the reason to schedule the task. */ schedulingReason(int schedulingReason)139 public abstract Builder schedulingReason(int schedulingReason); 140 141 /** Build a federated training task instance. */ 142 @NonNull build()143 public abstract FederatedTrainingTask build(); 144 } 145 146 /** @return a builder of federated training task. */ toBuilder()147 public abstract Builder toBuilder(); 148 149 /** @return a generic builder. */ 150 @NonNull builder()151 public static Builder builder() { 152 return new AutoValue_FederatedTrainingTask.Builder(); 153 } 154 addToDatabase(SQLiteDatabase db)155 boolean addToDatabase(SQLiteDatabase db) { 156 ContentValues values = new ContentValues(); 157 values.put(FederatedTrainingTaskColumns.APP_PACKAGE_NAME, appPackageName()); 158 values.put(FederatedTrainingTaskColumns.JOB_SCHEDULER_JOB_ID, jobId()); 159 160 values.put(FederatedTrainingTaskColumns.POPULATION_NAME, populationName()); 161 if (intervalOptions() != null) { 162 values.put(FederatedTrainingTaskColumns.INTERVAL_OPTIONS, intervalOptions()); 163 } 164 165 values.put(FederatedTrainingTaskColumns.CREATION_TIME, creationTime()); 166 values.put(FederatedTrainingTaskColumns.LAST_SCHEDULED_TIME, lastScheduledTime()); 167 if (lastRunStartTime() != null) { 168 values.put(FederatedTrainingTaskColumns.LAST_RUN_START_TIME, lastRunStartTime()); 169 } 170 if (lastRunEndTime() != null) { 171 values.put(FederatedTrainingTaskColumns.LAST_RUN_END_TIME, lastRunEndTime()); 172 } 173 values.put(FederatedTrainingTaskColumns.EARLIEST_NEXT_RUN_TIME, earliestNextRunTime()); 174 values.put(FederatedTrainingTaskColumns.CONSTRAINTS, constraints()); 175 values.put(FederatedTrainingTaskColumns.SCHEDULING_REASON, schedulingReason()); 176 long jobId = 177 db.insertWithOnConflict( 178 FEDERATED_TRAINING_TASKS_TABLE, 179 "", 180 values, 181 SQLiteDatabase.CONFLICT_REPLACE); 182 return jobId != -1; 183 } 184 readFederatedTrainingTasksFromDatabase( SQLiteDatabase db, String selection, String[] selectionArgs)185 static List<FederatedTrainingTask> readFederatedTrainingTasksFromDatabase( 186 SQLiteDatabase db, String selection, String[] selectionArgs) { 187 List<FederatedTrainingTask> taskList = new ArrayList<>(); 188 String[] selectColumns = { 189 FederatedTrainingTaskColumns.APP_PACKAGE_NAME, 190 FederatedTrainingTaskColumns.JOB_SCHEDULER_JOB_ID, 191 FederatedTrainingTaskColumns.POPULATION_NAME, 192 FederatedTrainingTaskColumns.INTERVAL_OPTIONS, 193 FederatedTrainingTaskColumns.CREATION_TIME, 194 FederatedTrainingTaskColumns.LAST_SCHEDULED_TIME, 195 FederatedTrainingTaskColumns.LAST_RUN_START_TIME, 196 FederatedTrainingTaskColumns.LAST_RUN_END_TIME, 197 FederatedTrainingTaskColumns.EARLIEST_NEXT_RUN_TIME, 198 FederatedTrainingTaskColumns.CONSTRAINTS, 199 FederatedTrainingTaskColumns.SCHEDULING_REASON, 200 }; 201 Cursor cursor = null; 202 try { 203 cursor = 204 db.query( 205 FEDERATED_TRAINING_TASKS_TABLE, 206 selectColumns, 207 selection, 208 selectionArgs, 209 null, 210 null 211 /* groupBy= */ , 212 null 213 /* having= */ , 214 null 215 /* orderBy= */); 216 while (cursor.moveToNext()) { 217 FederatedTrainingTask.Builder trainingTaskBuilder = 218 FederatedTrainingTask.builder() 219 .appPackageName( 220 cursor.getString( 221 cursor.getColumnIndexOrThrow( 222 FederatedTrainingTaskColumns 223 .APP_PACKAGE_NAME))) 224 .jobId( 225 cursor.getInt( 226 cursor.getColumnIndexOrThrow( 227 FederatedTrainingTaskColumns 228 .JOB_SCHEDULER_JOB_ID))) 229 .populationName( 230 cursor.getString( 231 cursor.getColumnIndexOrThrow( 232 FederatedTrainingTaskColumns 233 .POPULATION_NAME))) 234 .creationTime( 235 cursor.getLong( 236 cursor.getColumnIndexOrThrow( 237 FederatedTrainingTaskColumns 238 .CREATION_TIME))) 239 .lastScheduledTime( 240 cursor.getLong( 241 cursor.getColumnIndexOrThrow( 242 FederatedTrainingTaskColumns 243 .LAST_SCHEDULED_TIME))) 244 .lastRunStartTime( 245 cursor.getLong( 246 cursor.getColumnIndexOrThrow( 247 FederatedTrainingTaskColumns 248 .LAST_RUN_START_TIME))) 249 .lastRunEndTime( 250 cursor.getLong( 251 cursor.getColumnIndexOrThrow( 252 FederatedTrainingTaskColumns 253 .LAST_RUN_END_TIME))) 254 .earliestNextRunTime( 255 cursor.getLong( 256 cursor.getColumnIndexOrThrow( 257 FederatedTrainingTaskColumns 258 .EARLIEST_NEXT_RUN_TIME))); 259 int schedulingReason = 260 cursor.getInt( 261 cursor.getColumnIndexOrThrow( 262 FederatedTrainingTaskColumns.SCHEDULING_REASON)); 263 if (!cursor.isNull(schedulingReason)) { 264 trainingTaskBuilder.schedulingReason(schedulingReason); 265 } 266 byte[] intervalOptions = 267 cursor.getBlob( 268 cursor.getColumnIndexOrThrow( 269 FederatedTrainingTaskColumns.INTERVAL_OPTIONS)); 270 if (intervalOptions != null) { 271 trainingTaskBuilder.intervalOptions(intervalOptions); 272 } 273 byte[] constraints = 274 cursor.getBlob( 275 cursor.getColumnIndexOrThrow( 276 FederatedTrainingTaskColumns.CONSTRAINTS)); 277 trainingTaskBuilder.constraints(constraints); 278 taskList.add(trainingTaskBuilder.build()); 279 } 280 } finally { 281 if (cursor != null) { 282 cursor.close(); 283 } 284 } 285 return taskList; 286 } 287 } 288