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.training; 18 19 import static com.google.common.util.concurrent.MoreExecutors.directExecutor; 20 21 import android.app.job.JobParameters; 22 import android.app.job.JobService; 23 import android.util.Log; 24 25 import com.android.federatedcompute.services.common.FederatedComputeExecutors; 26 27 import com.google.common.util.concurrent.FutureCallback; 28 import com.google.common.util.concurrent.Futures; 29 import com.google.common.util.concurrent.ListenableFuture; 30 31 /** Main service for the scheduled federated computation jobs. */ 32 public class FederatedJobService extends JobService { 33 private static final String TAG = "FederatedJobService"; 34 private ListenableFuture<Boolean> mRunCompleteFuture; 35 36 @Override onStartJob(JobParameters params)37 public boolean onStartJob(JobParameters params) { 38 Log.d(TAG, "FederatedJobService.onStartJob"); 39 mRunCompleteFuture = 40 Futures.submit( 41 () -> 42 FederatedComputeWorker.getInstance(this) 43 .startTrainingRun(params.getJobId()), 44 FederatedComputeExecutors.getBackgroundExecutor()); 45 46 Futures.addCallback( 47 mRunCompleteFuture, 48 new FutureCallback<Boolean>() { 49 @Override 50 public void onSuccess(Boolean result) { 51 Log.d(TAG, "federated computation job is done!"); 52 jobFinished(params, /* wantsReschedule= */ false); 53 } 54 55 @Override 56 public void onFailure(Throwable t) { 57 Log.e(TAG, "Failed to handle computation job: " + params.getJobId()); 58 jobFinished(params, /* wantsReschedule= */ false); 59 } 60 }, 61 directExecutor()); 62 return true; 63 } 64 65 @Override onStopJob(JobParameters params)66 public boolean onStopJob(JobParameters params) { 67 if (mRunCompleteFuture != null) { 68 mRunCompleteFuture.cancel(true); 69 } 70 FederatedComputeWorker.getInstance(this).cancelActiveRun(); 71 // Reschedule the job since it's not done. TODO: we should implement specify reschedule 72 // logic instead. 73 return true; 74 } 75 } 76