/*
* Copyright (C) 2022 The Android Open Source Project
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package android.adservices.ondevicepersonalization;
import android.adservices.ondevicepersonalization.aidl.IDataAccessService;
import android.adservices.ondevicepersonalization.aidl.IFederatedComputeCallback;
import android.adservices.ondevicepersonalization.aidl.IFederatedComputeService;
import android.annotation.CallbackExecutor;
import android.annotation.FlaggedApi;
import android.annotation.NonNull;
import android.annotation.WorkerThread;
import android.federatedcompute.common.TrainingOptions;
import android.os.OutcomeReceiver;
import android.os.RemoteException;
import com.android.adservices.ondevicepersonalization.flags.Flags;
import com.android.ondevicepersonalization.internal.util.LoggerFactory;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;
/**
* Handles scheduling federated compute jobs. See {@link
* IsolatedService#getFederatedComputeScheduler}.
*/
public class FederatedComputeScheduler {
private static final String TAG = FederatedComputeScheduler.class.getSimpleName();
private static final int FEDERATED_COMPUTE_SCHEDULE_TIMEOUT_SECONDS = 30;
private static final LoggerFactory.Logger sLogger = LoggerFactory.getLogger();
private final IFederatedComputeService mFcService;
private final IDataAccessService mDataAccessService;
/** @hide */
public FederatedComputeScheduler(
IFederatedComputeService binder, IDataAccessService dataService) {
mFcService = binder;
mDataAccessService = dataService;
}
/**
* Schedules a federated compute job. In {@link IsolatedService#onRequest}, the app can call
* {@link IsolatedService#getFederatedComputeScheduler} to pass the scheduler when constructing
* the {@link IsolatedWorker}.
*
* @param params parameters related to job scheduling.
* @param input the configuration of the federated computation. It should be consistent with the
* federated compute server setup as described in
* Federated Compute Server documentation. .
*/
@WorkerThread
public void schedule(@NonNull Params params, @NonNull FederatedComputeInput input) {
if (mFcService == null) {
logApiCallStats(
Constants.API_NAME_FEDERATED_COMPUTE_SCHEDULE,
0,
Constants.STATUS_INTERNAL_ERROR);
throw new IllegalStateException(
"FederatedComputeScheduler not available for this instance.");
}
final long startTimeMillis = System.currentTimeMillis();
TrainingOptions trainingOptions =
new TrainingOptions.Builder()
.setPopulationName(input.getPopulationName())
.setTrainingInterval(convertTrainingInterval(params.getTrainingInterval()))
.build();
CountDownLatch latch = new CountDownLatch(1);
final int[] err = {0};
int responseCode = Constants.STATUS_INTERNAL_ERROR;
try {
mFcService.schedule(
trainingOptions,
new IFederatedComputeCallback.Stub() {
@Override
public void onSuccess() {
latch.countDown();
}
@Override
public void onFailure(int i) {
err[0] = i;
latch.countDown();
}
});
boolean countedDown =
latch.await(FEDERATED_COMPUTE_SCHEDULE_TIMEOUT_SECONDS, TimeUnit.SECONDS);
if (err[0] != 0) {
sLogger.e(
TAG + " : Internal failure occurred while scheduling job, error code %d",
err[0]);
responseCode = err[0];
return;
} else if (!countedDown) {
sLogger.d(TAG + " : timed out waiting for schedule operation to complete.");
responseCode = Constants.STATUS_TIMEOUT;
return;
}
responseCode = Constants.STATUS_SUCCESS;
} catch (RemoteException | InterruptedException e) {
responseCode = Constants.STATUS_REMOTE_EXCEPTION;
sLogger.e(TAG + ": Failed to schedule federated compute job", e);
throw new IllegalStateException(e);
} finally {
logApiCallStats(
Constants.API_NAME_FEDERATED_COMPUTE_SCHEDULE,
System.currentTimeMillis() - startTimeMillis,
responseCode);
}
}
/**
* Schedules a federated compute job. In {@link IsolatedService#onRequest}, the app can call
* {@link IsolatedService#getFederatedComputeScheduler} to pass the scheduler when constructing
* the {@link IsolatedWorker}.
*
* @param federatedComputeScheduleRequest input parameters related to job scheduling.
* @param executor the {@link Executor} on which to invoke the callback.
* @param outcomeReceiver This either returns a {@link FederatedComputeScheduleResponse} on
* success, or {@link Exception} on failure. The exception type is {@link
* OnDevicePersonalizationException} with error code {@link
* OnDevicePersonalizationException#ERROR_INVALID_TRAINING_MANIFEST} if the manifest is
* missing the federated compute server URL or {@link
* OnDevicePersonalizationException#ERROR_SCHEDULE_TRAINING_FAILED} when scheduling fails
* for other reasons.
*/
@WorkerThread
@FlaggedApi(Flags.FLAG_FCP_SCHEDULE_WITH_OUTCOME_RECEIVER_ENABLED)
public void schedule(
@NonNull FederatedComputeScheduleRequest federatedComputeScheduleRequest,
@NonNull @CallbackExecutor Executor executor,
@NonNull OutcomeReceiver outcomeReceiver) {
if (mFcService == null) {
logApiCallStats(
Constants.API_NAME_FEDERATED_COMPUTE_SCHEDULE,
0,
Constants.STATUS_INTERNAL_ERROR);
executor.execute(
() -> {
outcomeReceiver.onError(new IllegalStateException(
"FederatedComputeScheduler not available for this instance."));
});
}
final long startTimeMillis = System.currentTimeMillis();
TrainingOptions trainingOptions =
new TrainingOptions.Builder()
.setPopulationName(federatedComputeScheduleRequest.getPopulationName())
.setTrainingInterval(
convertTrainingInterval(
federatedComputeScheduleRequest
.getParams()
.getTrainingInterval()))
.build();
try {
mFcService.schedule(
trainingOptions,
new IFederatedComputeCallback.Stub() {
@Override
public void onSuccess() {
logApiCallStats(
Constants.API_NAME_FEDERATED_COMPUTE_SCHEDULE,
System.currentTimeMillis() - startTimeMillis,
Constants.STATUS_SUCCESS);
executor.execute(
() -> {
outcomeReceiver.onResult(
new FederatedComputeScheduleResponse(
federatedComputeScheduleRequest));
});
}
@Override
public void onFailure(int errorCode) {
logApiCallStats(
Constants.API_NAME_FEDERATED_COMPUTE_SCHEDULE,
System.currentTimeMillis() - startTimeMillis,
errorCode);
executor.execute(
() -> {
outcomeReceiver.onError(
new OnDevicePersonalizationException(
translateErrorCode(errorCode)));
});
}
});
} catch (RemoteException e) {
sLogger.e(TAG + ": Failed to schedule federated compute job", e);
logApiCallStats(
Constants.API_NAME_FEDERATED_COMPUTE_SCHEDULE,
System.currentTimeMillis() - startTimeMillis,
Constants.STATUS_REMOTE_EXCEPTION);
executor.execute(
() -> {
outcomeReceiver.onError(e);
});
}
}
/**
* Translate the failed error code from the {@link IFederatedComputeService} to appropriate API
* surface error code.
*/
private static int translateErrorCode(int i) {
// Returns invalid/missing manifest or general error code to caller. The general error code
// includes personalization disable and all other errors populated from FCP service.
return i == Constants.STATUS_FCP_MANIFEST_INVALID
? OnDevicePersonalizationException.ERROR_INVALID_TRAINING_MANIFEST
: OnDevicePersonalizationException.ERROR_SCHEDULE_TRAINING_FAILED;
}
/**
* Cancels a federated compute job with input training params. In {@link
* IsolatedService#onRequest}, the app can call {@link
* IsolatedService#getFederatedComputeScheduler} to pass scheduler when constructing the {@link
* IsolatedWorker}.
*
* @param input the configuration of the federated compute. It should be consistent with the
* federated compute server setup.
*/
@WorkerThread
public void cancel(@NonNull FederatedComputeInput input) {
final long startTimeMillis = System.currentTimeMillis();
int responseCode = Constants.STATUS_INTERNAL_ERROR;
if (mFcService == null) {
logApiCallStats(
Constants.API_NAME_FEDERATED_COMPUTE_CANCEL,
System.currentTimeMillis() - startTimeMillis,
responseCode);
throw new IllegalStateException(
"FederatedComputeScheduler not available for this instance.");
}
CountDownLatch latch = new CountDownLatch(1);
final int[] err = {0};
try {
mFcService.cancel(
input.getPopulationName(),
new IFederatedComputeCallback.Stub() {
@Override
public void onSuccess() {
latch.countDown();
}
@Override
public void onFailure(int errorCode) {
err[0] = errorCode;
latch.countDown();
}
});
boolean countedDown =
latch.await(FEDERATED_COMPUTE_SCHEDULE_TIMEOUT_SECONDS, TimeUnit.SECONDS);
if (err[0] != 0) {
sLogger.e("Internal failure occurred while cancelling job, error code %d", err[0]);
responseCode = Constants.STATUS_INTERNAL_ERROR;
// Fail silently for now. TODO(b/346827691): update schedule/cancel API to return
// error status to caller.
return;
} else if (!countedDown) {
sLogger.d(TAG + " : timed out waiting for cancel operation to complete.");
responseCode = Constants.STATUS_INTERNAL_ERROR;
return;
}
responseCode = Constants.STATUS_SUCCESS;
} catch (RemoteException | InterruptedException e) {
sLogger.e(TAG + ": Failed to cancel federated compute job", e);
throw new IllegalStateException(e);
} finally {
logApiCallStats(
Constants.API_NAME_FEDERATED_COMPUTE_CANCEL,
System.currentTimeMillis() - startTimeMillis,
responseCode);
}
}
private static android.federatedcompute.common.TrainingInterval convertTrainingInterval(
TrainingInterval interval) {
return new android.federatedcompute.common.TrainingInterval.Builder()
.setMinimumIntervalMillis(interval.getMinimumInterval().toMillis())
.setSchedulingMode(convertSchedulingMode(interval))
.build();
}
private static @android.federatedcompute.common.TrainingInterval.SchedulingMode int
convertSchedulingMode(TrainingInterval interval) {
switch (interval.getSchedulingMode()) {
case TrainingInterval.SCHEDULING_MODE_ONE_TIME:
return android.federatedcompute.common.TrainingInterval.SCHEDULING_MODE_ONE_TIME;
case TrainingInterval.SCHEDULING_MODE_RECURRENT:
return android.federatedcompute.common.TrainingInterval.SCHEDULING_MODE_RECURRENT;
default:
throw new IllegalStateException(
"Unsupported scheduling mode " + interval.getSchedulingMode());
}
}
/** Helper method to log call stats based on response code. */
private void logApiCallStats(int apiName, long duration, int responseCode) {
try {
mDataAccessService.logApiCallStats(apiName, duration, responseCode);
} catch (Exception e) {
sLogger.d(e, TAG + ": failed to log metrics");
}
}
/** The parameters related to job scheduling. */
public static class Params {
/**
* If training interval is scheduled for recurrent tasks, the earliest time this task could
* start is after the minimum training interval expires. E.g. If the task is set to run
* maximum once per day, the first run of this task will be one day after this task is
* scheduled. When a one time job is scheduled, the earliest next runtime is calculated
* based on federated compute default interval.
*/
@NonNull private final TrainingInterval mTrainingInterval;
public Params(@NonNull TrainingInterval trainingInterval) {
mTrainingInterval = trainingInterval;
}
@NonNull
public TrainingInterval getTrainingInterval() {
return mTrainingInterval;
}
}
}