Snap for 11130229 from 84855de1baa51f0a58f88f8a07275b22611b8e0a to 24Q1-release

Change-Id: Ib6517e5c731d9c8188230b8dfcd33a93aa6588ef

6 files changed, 55 insertions, 0 deletions

diff --git a/federatedcompute/src/com/android/federatedcompute/services/common/Constants.java b/federatedcompute/src/com/android/federatedcompute/services/common/Constants.java
index 33b64603..6c697525 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/common/Constants.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/common/Constants.java
@@ -36,5 +36,14 @@ public class Constants {
     public static final String ISOLATED_TRAINING_SERVICE_NAME =
             "com.android.federatedcompute.services.training.IsolatedTrainingService";
 
+    public static final String TRACE_HTTP_ISSUE_CHECKIN = "Http#issueCheckin";
+    public static final String TRACE_HTTP_REPORT_RESULT = "Http#reportResult";
+    public static final String TRACE_ISOLATED_PROCESS_RUN_FL_TRAINING =
+            "IsolatedProcess#runFlTraining";
+    public static final String TRACE_NATIVE_RUN_FEDERATED_COMPUTATION =
+            "Native#runFederatedComputation";
+    public static final String TRACE_WORKER_RUN_FL_COMPUTATION = "Worker#runFlComputation";
+    public static final String TRACE_WORKER_START_TRAINING_RUN = "Worker#startTrainingRun";
+
     private Constants() {}
 }
diff --git a/federatedcompute/src/com/android/federatedcompute/services/http/HttpFederatedProtocol.java b/federatedcompute/src/com/android/federatedcompute/services/http/HttpFederatedProtocol.java
index 0703d3f3..d7883542 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/http/HttpFederatedProtocol.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/http/HttpFederatedProtocol.java
@@ -16,6 +16,8 @@
 
 package com.android.federatedcompute.services.http;
 
+import static com.android.federatedcompute.services.common.Constants.TRACE_HTTP_ISSUE_CHECKIN;
+import static com.android.federatedcompute.services.common.Constants.TRACE_HTTP_REPORT_RESULT;
 import static com.android.federatedcompute.services.common.FederatedComputeExecutors.getBackgroundExecutor;
 import static com.android.federatedcompute.services.common.FederatedComputeExecutors.getLightweightExecutor;
 import static com.android.federatedcompute.services.common.FileUtils.createTempFile;
@@ -28,6 +30,8 @@ import static com.android.federatedcompute.services.http.HttpClientUtil.ODP_IDEM
 import static com.android.federatedcompute.services.http.HttpClientUtil.compressWithGzip;
 import static com.android.federatedcompute.services.http.HttpClientUtil.uncompressWithGzip;
 
+import android.os.Trace;
+
 import com.android.federatedcompute.internal.util.LogUtil;
 import com.android.federatedcompute.services.http.HttpClientUtil.HttpMethod;
 import com.android.federatedcompute.services.training.util.ComputationResult;
@@ -84,6 +88,7 @@ public final class HttpFederatedProtocol {
 
     /** Helper function to perform check in and download federated task from remote servers. */
     public ListenableFuture<CheckinResult> issueCheckin() {
+        Trace.beginAsyncSection(TRACE_HTTP_ISSUE_CHECKIN, 0);
         return FluentFuture.from(createTaskAssignment())
                 .transformAsync(
                         federatedComputeHttpResponse -> {
@@ -129,6 +134,7 @@ public final class HttpFederatedProtocol {
 
     /** Helper functions to reporting result and upload result. */
     public FluentFuture<RejectionInfo> reportResult(ComputationResult computationResult) {
+        Trace.beginAsyncSection(TRACE_HTTP_REPORT_RESULT, 0);
         if (computationResult != null && computationResult.isResultSuccess()) {
             return FluentFuture.from(performReportResult(computationResult))
                     .transformAsync(
@@ -146,6 +152,8 @@ public final class HttpFederatedProtocol {
                                                 resp -> {
                                                     validateHttpResponseStatus(
                                                             "Upload result", resp);
+                                                    Trace.endAsyncSection(
+                                                            TRACE_HTTP_REPORT_RESULT, 0);
                                                     return null;
                                                 },
                                                 getLightweightExecutor());
@@ -260,6 +268,7 @@ public final class HttpFederatedProtocol {
             checkpointData = uncompressWithGzip(checkpointData);
         }
         writeToFile(inputCheckpointFile, checkpointData);
+        Trace.endAsyncSection(TRACE_HTTP_ISSUE_CHECKIN, 0);
         return new CheckinResult(inputCheckpointFile, clientOnlyPlan, taskAssignment);
     }
 
@@ -366,6 +375,7 @@ public final class HttpFederatedProtocol {
                     // https://cloud.google.com/storage/docs/transcoding#decompressive_transcoding
                     headerList.put(ACCEPT_ENCODING_HDR, GZIP_ENCODING_HDR);
                 }
+                LogUtil.d(TAG, "start fetch task resources");
                 FederatedComputeHttpRequest httpRequest =
                         FederatedComputeHttpRequest.create(
                                 resource.getUri(),
diff --git a/federatedcompute/src/com/android/federatedcompute/services/training/FederatedComputeWorker.java b/federatedcompute/src/com/android/federatedcompute/services/training/FederatedComputeWorker.java
index 58eb6038..ac54dde5 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/training/FederatedComputeWorker.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/training/FederatedComputeWorker.java
@@ -18,6 +18,8 @@ package com.android.federatedcompute.services.training;
 
 import static com.android.federatedcompute.services.common.Constants.CLIENT_ONLY_PLAN_FILE_NAME;
 import static com.android.federatedcompute.services.common.Constants.ISOLATED_TRAINING_SERVICE_NAME;
+import static com.android.federatedcompute.services.common.Constants.TRACE_WORKER_RUN_FL_COMPUTATION;
+import static com.android.federatedcompute.services.common.Constants.TRACE_WORKER_START_TRAINING_RUN;
 import static com.android.federatedcompute.services.common.FederatedComputeExecutors.getBackgroundExecutor;
 import static com.android.federatedcompute.services.common.FederatedComputeExecutors.getLightweightExecutor;
 import static com.android.federatedcompute.services.common.FileUtils.createTempFile;
@@ -33,6 +35,7 @@ import android.federatedcompute.common.ClientConstants;
 import android.federatedcompute.common.ExampleConsumption;
 import android.os.Bundle;
 import android.os.ParcelFileDescriptor;
+import android.os.Trace;
 
 import androidx.concurrent.futures.CallbackToFutureAdapter;
 
@@ -160,6 +163,7 @@ public class FederatedComputeWorker {
             int jobId, FederatedTrainingTask trainingTask) {
         synchronized (mLock) {
             // Only allows one concurrent job running.
+            Trace.beginAsyncSection(TRACE_WORKER_START_TRAINING_RUN, jobId);
             TrainingRun run = new TrainingRun(jobId, trainingTask);
             mActiveRun = run;
             ListenableFuture<FLRunnerResult> runCompletedFuture = doTraining(run);
@@ -168,6 +172,8 @@ public class FederatedComputeWorker {
                             .call(
                                     () -> {
                                         unBindServicesIfNecessary(run);
+                                        Trace.endAsyncSection(
+                                                TRACE_WORKER_START_TRAINING_RUN, jobId);
                                         return null;
                                     },
                                     mInjector.getBgExecutor());
@@ -445,6 +451,7 @@ public class FederatedComputeWorker {
             CheckinResult checkinResult,
             String outputCheckpointFile,
             IExampleStoreIterator iterator) {
+        Trace.beginAsyncSection(TRACE_WORKER_RUN_FL_COMPUTATION, 0);
         ParcelFileDescriptor outputCheckpointFd =
                 createTempFileDescriptor(
                         outputCheckpointFile, ParcelFileDescriptor.MODE_READ_WRITE);
@@ -510,6 +517,7 @@ public class FederatedComputeWorker {
                                     unbindFromIsolatedTrainingService();
                                     run.mIsolatedTrainingService = null;
                                 }
+                                Trace.endAsyncSection(TRACE_WORKER_RUN_FL_COMPUTATION, 0);
                                 return computationResult;
                             },
                             getLightweightExecutor());
diff --git a/federatedcompute/src/com/android/federatedcompute/services/training/IsolatedTrainingServiceImpl.java b/federatedcompute/src/com/android/federatedcompute/services/training/IsolatedTrainingServiceImpl.java
index 7d23f926..79966d99 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/training/IsolatedTrainingServiceImpl.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/training/IsolatedTrainingServiceImpl.java
@@ -16,6 +16,8 @@
 
 package com.android.federatedcompute.services.training;
 
+import static com.android.federatedcompute.services.common.Constants.TRACE_ISOLATED_PROCESS_RUN_FL_TRAINING;
+
 import android.annotation.NonNull;
 import android.federatedcompute.aidl.IExampleStoreIterator;
 import android.federatedcompute.common.ClientConstants;
@@ -23,6 +25,7 @@ import android.federatedcompute.common.ExampleConsumption;
 import android.os.Bundle;
 import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
+import android.os.Trace;
 
 import com.android.federatedcompute.internal.util.LogUtil;
 import com.android.federatedcompute.services.common.Constants;
@@ -70,6 +73,7 @@ public class IsolatedTrainingServiceImpl extends IIsolatedTrainingService.Stub {
     public void runFlTraining(@NonNull Bundle params, @NonNull ITrainingResultCallback callback) {
         Objects.requireNonNull(params);
         Objects.requireNonNull(callback);
+        Trace.beginAsyncSection(TRACE_ISOLATED_PROCESS_RUN_FL_TRAINING, 0);
 
         IExampleStoreIterator exampleStoreIteratorBinder =
                 IExampleStoreIterator.Stub.asInterface(
@@ -149,6 +153,7 @@ public class IsolatedTrainingServiceImpl extends IIsolatedTrainingService.Stub {
                                 ClientConstants.EXTRA_EXAMPLE_CONSUMPTION_LIST,
                                 exampleConsumptionArrayList);
                         sendResult(bundle, callback);
+                        Trace.endAsyncSection(TRACE_ISOLATED_PROCESS_RUN_FL_TRAINING, 0);
                     }
 
                     @Override
diff --git a/federatedcompute/src/com/android/federatedcompute/services/training/jni/FlRunnerWrapper.java b/federatedcompute/src/com/android/federatedcompute/services/training/jni/FlRunnerWrapper.java
index 33ec656f..30fa4ce1 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/training/jni/FlRunnerWrapper.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/training/jni/FlRunnerWrapper.java
@@ -16,6 +16,10 @@
 
 package com.android.federatedcompute.services.training.jni;
 
+import static com.android.federatedcompute.services.common.Constants.TRACE_NATIVE_RUN_FEDERATED_COMPUTATION;
+
+import android.os.Trace;
+
 import com.android.federatedcompute.internal.util.LogUtil;
 import com.android.federatedcompute.services.examplestore.ExampleIterator;
 import com.android.federatedcompute.services.training.util.ListenableSupplier;
@@ -57,6 +61,7 @@ public final class FlRunnerWrapper implements Closeable {
             ClientOnlyPlan clientOnlyPlan,
             String checkpointInputFileName,
             String checkpointOutputFileName) {
+        Trace.beginAsyncSection(TRACE_NATIVE_RUN_FEDERATED_COMPUTATION, 0);
         SimpleTaskEnvironmentImpl simpleTaskEnv =
                 new SimpleTaskEnvironmentImpl(mInterruptionFlag, mExampleIterator);
         byte[] flRunnerResultSerialized =
@@ -77,6 +82,8 @@ public final class FlRunnerWrapper implements Closeable {
             // recover from it.
             LogUtil.e(TAG, "Cannot parse FLRunnerResult", e);
             throw new IllegalArgumentException(e);
+        } finally {
+            Trace.endAsyncSection(TRACE_NATIVE_RUN_FEDERATED_COMPUTATION, 0);
         }
     }
 
diff --git a/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/TestHelper.java b/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/TestHelper.java
index 4a5f2186..c8ff1f88 100644
--- a/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/TestHelper.java
+++ b/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/TestHelper.java
@@ -70,6 +70,22 @@ public class TestHelper {
         executeShellCommand(
                 "am broadcast -a android.intent.action.BOOT_COMPLETED -p "
                     + "com.google.android.federatedcompute");
+        executeShellCommand(
+                "cmd jobscheduler run -f "
+                        + "com.google.android.ondevicepersonalization.services 1000");
+        SystemClock.sleep(5000);
+        executeShellCommand(
+                "cmd jobscheduler run -f "
+                        + "com.google.android.ondevicepersonalization.services 1006");
+        SystemClock.sleep(5000);
+        executeShellCommand(
+                "cmd jobscheduler run -f "
+                        + "com.google.android.ondevicepersonalization.services 1003");
+        SystemClock.sleep(5000);
+        executeShellCommand(
+                "cmd jobscheduler run -f "
+                        + "com.google.android.ondevicepersonalization.services 1004");
+        SystemClock.sleep(5000);
     }
 
     /** Kill running processes to get performance measurement under cold start */