Donate to e Foundation | Murena handsets with /e/OS | Own a part of Murena! Learn more

Commit 2ff57111 authored by Sandeep Bandaru's avatar Sandeep Bandaru
Browse files

Implement registerInferenceServiceLifecycleListener for OnDeviceIntelligence 

This CL adds the `registerInferenceServiceLifecycleListener` API to `OnDeviceIntelligenceManager` and the corresponding implementation in the service side.

The `LifecycleListener` (now nested within `OnDeviceSandboxedInferenceService`) will be notified when a lifecycle event occurs, such as:
- `LIFECYCLE_EVENT_MODEL_LOADED`
- `LIFECYCLE_EVENT_MODEL_UNLOADED`
- `LIFECYCLE_EVENT_ATTEMPTING_MODEL_LOADING`

This allows clients to track the lifecycle status of models and the inference service.

Flag: android.app.ondeviceintelligence.flags.on_device_intelligence_25q4
Bug: 372658837
Test: cts
Change-Id: I3b03febe78871db81b1146462bec3dd4aadfbd67
parent ae2f4fc9

packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/ILifecycleListener.aidl

0 → 100644
+28 −0
Original line number Diff line number Diff line 
/*

 * Copyright (C) 2024 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.app.ondeviceintelligence;



import android.app.ondeviceintelligence.Feature;



/**

 * Callback for inference service lifecycle events.

 *

 * @hide

 */

oneway interface ILifecycleListener {

    void onLifecycleEvent(int eventType, in Feature feature);

}

packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl

+14 −5
Original line number Diff line number Diff line
@@ -34,6 +34,7 @@ 
import android.app.ondeviceintelligence.IStreamingResponseCallback;

import android.app.ondeviceintelligence.IProcessingSignal;

import android.app.ondeviceintelligence.ITokenInfoCallback;

import android.app.ondeviceintelligence.ILifecycleListener;





/**
@@ -41,7 +42,7 @@ 
  *

  * @hide

  */

// Next available field id: 12

// Next available field id: 14

interface IOnDeviceIntelligenceManager {

      @JavaPassthrough(annotation="@android.annotation.RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)")

      void getVersion(in RemoteCallback remoteCallback) = 1;
@@ -82,4 +83,12 @@ interface IOnDeviceIntelligenceManager { 
      String getRemoteServicePackageName() = 9;



      List<InferenceInfo> getLatestInferenceInfo(long startTimeEpochMillis) = 10;



      @JavaPassthrough(annotation="@android.annotation.RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)")

      @JavaPassthrough(annotation="@android.annotation.FlaggedApi(android.app.ondeviceintelligence.flags.Flags.FLAG_ON_DEVICE_INTELLIGENCE_25Q4)")

      void registerInferenceServiceLifecycleListener(in ILifecycleListener listener) = 12;



      @JavaPassthrough(annotation="@android.annotation.RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)")

      @JavaPassthrough(annotation="@android.annotation.FlaggedApi(android.app.ondeviceintelligence.flags.Flags.FLAG_ON_DEVICE_INTELLIGENCE_25Q4)")

      void unregisterInferenceServiceLifecycleListener(in ILifecycleListener listener) = 13;

 }

packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java

+66 −0
Original line number Diff line number Diff line
@@ -40,9 +40,16 @@ import android.os.OutcomeReceiver; 
import android.os.PersistableBundle;

import android.os.RemoteCallback;

import android.os.RemoteException;

import android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService;

import android.system.OsConstants;

import android.util.Log;



import java.util.List;

import java.util.Map;

import java.util.Objects;

import java.util.concurrent.ConcurrentHashMap;

import java.util.concurrent.Executor;



import com.android.internal.infra.AndroidFuture;



import java.lang.annotation.ElementType;
@@ -93,6 +100,9 @@ public final class OnDeviceIntelligenceManager { 
    private final Context mContext;

    private final IOnDeviceIntelligenceManager mService;



    private final Map<OnDeviceSandboxedInferenceService.LifecycleListener, ILifecycleListener.Stub>

            mLifecycleListeners = new ConcurrentHashMap<>();



    /**

     * @hide

     */
@@ -579,6 +589,62 @@ public final class OnDeviceIntelligenceManager { 
        }

    }



    /**

     * Registers a listener for inference service lifecycle events of the

     * configured {@link OnDeviceSandboxedInferenceService}. This method is

     * thread-safe.

     *

     * @param executor The executor to run the listener callbacks on.

     * @param listener The listener to register.

     */

    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)

    @FlaggedApi(FLAG_ON_DEVICE_INTELLIGENCE_25Q4)

    public void registerInferenceServiceLifecycleListener(

            @NonNull @CallbackExecutor Executor executor,

            @NonNull OnDeviceSandboxedInferenceService.LifecycleListener listener) {

        Objects.requireNonNull(executor, "Executor must not be null");

        Objects.requireNonNull(listener, "Listener must not be null");

        ILifecycleListener.Stub stub = new ILifecycleListener.Stub() {

            @Override

            public void onLifecycleEvent(int event, @NonNull Feature feature) {

                Binder.withCleanCallingIdentity(

                        () -> executor.execute(() -> listener.onLifecycleEvent(event, feature)));

            }

        };



        mLifecycleListeners.put(listener, stub);

        try {

            mService.registerInferenceServiceLifecycleListener(stub);

        } catch (RemoteException e) {

            mLifecycleListeners.remove(listener);

            throw e.rethrowFromSystemServer();

        }

    }



    /**

     * Unregisters a {@link OnDeviceSandboxedInferenceService.LifecycleListener}

     * if previous registered with

     * {@link #registerInferenceServiceLifecycleListener}, otherwise no-op.

     * This method is thread-safe.

     *

     * @param listener The listener to unregister.

     */

    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)

    @FlaggedApi(FLAG_ON_DEVICE_INTELLIGENCE_25Q4)

    public void unregisterInferenceServiceLifecycleListener(

            @NonNull OnDeviceSandboxedInferenceService.LifecycleListener listener) {

        Objects.requireNonNull(listener, "Listener must not be null");

        ILifecycleListener.Stub stub = mLifecycleListeners.remove(listener);

        if (stub == null) {

            // Listener not registered or already unregistered.

            return;

        }

        try {

            mService.unregisterInferenceServiceLifecycleListener(stub);

        } catch (RemoteException e) {

            throw e.rethrowFromSystemServer();

        }

    }



    /** Request inference with provided Bundle and Params. */

    public static final int REQUEST_TYPE_INFERENCE = 0;

packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl

+3 −1
Original line number Diff line number Diff line
@@ -20,6 +20,7 @@ import android.app.ondeviceintelligence.IStreamingResponseCallback; 
import android.app.ondeviceintelligence.IResponseCallback;

import android.app.ondeviceintelligence.ITokenInfoCallback;

import android.app.ondeviceintelligence.IProcessingSignal;

import android.app.ondeviceintelligence.ILifecycleListener;

import android.app.ondeviceintelligence.Feature;

import android.os.IRemoteCallback;

import android.os.ICancellationSignal;
@@ -50,4 +51,5 @@ oneway interface IOnDeviceSandboxedInferenceService { 
                                in IStreamingResponseCallback callback) = 3;

    void updateProcessingState(in Bundle processingState,

                                     in IProcessingUpdateStatusCallback callback) = 4;

    void registerInferenceServiceLifecycleListener(in ILifecycleListener listener) = 5;

}

packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java

+80 −0
Original line number Diff line number Diff line
@@ -20,11 +20,13 @@ import static android.app.ondeviceintelligence.OnDeviceIntelligenceManager.AUGME 
import static android.app.ondeviceintelligence.flags.Flags.FLAG_ENABLE_ON_DEVICE_INTELLIGENCE;

import static android.app.ondeviceintelligence.flags.Flags.FLAG_ON_DEVICE_INTELLIGENCE_25Q4;



import static android.app.ondeviceintelligence.flags.Flags.FLAG_ON_DEVICE_INTELLIGENCE_25Q4;

import static com.android.internal.util.function.pooled.PooledLambda.obtainMessage;



import android.annotation.CallSuper;

import android.annotation.CallbackExecutor;

import android.annotation.FlaggedApi;

import android.annotation.IntDef;

import android.annotation.NonNull;

import android.annotation.Nullable;

import android.annotation.SdkConstant;
@@ -37,6 +39,7 @@ import android.app.ondeviceintelligence.IProcessingSignal; 
import android.app.ondeviceintelligence.IResponseCallback;

import android.app.ondeviceintelligence.IStreamingResponseCallback;

import android.app.ondeviceintelligence.ITokenInfoCallback;

import android.app.ondeviceintelligence.ILifecycleListener;

import android.app.ondeviceintelligence.OnDeviceIntelligenceException;

import android.app.ondeviceintelligence.OnDeviceIntelligenceManager;

import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.InferenceParams;
@@ -55,6 +58,7 @@ import android.os.IBinder; 
import android.os.ICancellationSignal;

import android.os.IRemoteCallback;

import android.os.Looper;

import android.os.Message;

import android.os.OutcomeReceiver;

import android.os.ParcelFileDescriptor;

import android.os.PersistableBundle;
@@ -67,6 +71,8 @@ import com.android.internal.infra.AndroidFuture; 


import java.io.FileInputStream;

import java.io.FileNotFoundException;

import java.lang.annotation.Retention;

import java.lang.annotation.RetentionPolicy;

import java.util.HashMap;

import java.util.Map;

import java.util.Objects;
@@ -147,6 +153,43 @@ public abstract class OnDeviceSandboxedInferenceService extends Service { 
     */

    public static final String DEVICE_CONFIG_UPDATE_BUNDLE_KEY = "device_config_update";



    /**

     * Listener for inference service lifecycle events.

     *

     * @hide

     */

    @SystemApi

    @FlaggedApi(FLAG_ON_DEVICE_INTELLIGENCE_25Q4)

    public interface LifecycleListener {

        /** A model for a feature was loaded. */

        int LIFECYCLE_EVENT_MODEL_LOADED = 1;

        /** A model for a feature was unloaded. */

        int LIFECYCLE_EVENT_MODEL_UNLOADED = 2;

        /** The service is attempting to load a model for a feature. */

        int LIFECYCLE_EVENT_ATTEMPTING_MODEL_LOADING = 3;





        /** @hide */

        @IntDef(prefix = {"LIFECYCLE_EVENT_"}, value = {

                LIFECYCLE_EVENT_MODEL_LOADED,

                LIFECYCLE_EVENT_MODEL_UNLOADED,

                LIFECYCLE_EVENT_ATTEMPTING_MODEL_LOADING

        })

        @Retention(RetentionPolicy.SOURCE)

        @interface LifecycleEvent {

        }



        /**

         * Called when a lifecycle event occurs.

         *

         * @param eventType The lifecycle event type

         *                  e.g., {@link #LIFECYCLE_EVENT_MODEL_LOADED}.

         * @param feature The feature associated with the event.

         */

        void onLifecycleEvent(@LifecycleEvent int eventType,

                @NonNull Feature feature);

    }



    private IRemoteStorageService mRemoteStorageService;

    private Handler mHandler;
@@ -268,6 +311,18 @@ public abstract class OnDeviceSandboxedInferenceService extends Service { 
                                    OnDeviceSandboxedInferenceService.this, processingState,

                                    wrapOutcomeReceiver(callback)));

                }



                @Override

                public void registerInferenceServiceLifecycleListener(

                        ILifecycleListener listener) {

                    Objects.requireNonNull(listener);

                    Message msg = obtainMessage(

                        OnDeviceSandboxedInferenceService

                                    ::onRegisterInferenceServiceLifecycleListener,

                        OnDeviceSandboxedInferenceService.this,

                        wrapLifecycleListener(listener));

                    mHandler.executeOrSendMessage(msg);

                }

            };

        }

        Slog.w(TAG, "Incorrect service interface, returning null.");
@@ -364,6 +419,15 @@ public abstract class OnDeviceSandboxedInferenceService extends Service { 
            @NonNull OutcomeReceiver<PersistableBundle,

                    OnDeviceIntelligenceException> callback);



    /**

     * Invoked when a caller wants to register a listener for inference service lifecycle events.

     *

     * @param listener The listener to be notified of model events.

     */

    @FlaggedApi(FLAG_ON_DEVICE_INTELLIGENCE_25Q4)

    public void onRegisterInferenceServiceLifecycleListener(

            @NonNull LifecycleListener listener) {

    }



    /**

     * Overrides {@link Context#openFileInput} to read files with the given file names under the
@@ -545,6 +609,22 @@ public abstract class OnDeviceSandboxedInferenceService extends Service { 
        };

    }



    private LifecycleListener wrapLifecycleListener(

            ILifecycleListener listener) {

        return new LifecycleListener() {

            @Override

            public void onLifecycleEvent(

                    @LifecycleListener.LifecycleEvent int event,

                    @NonNull Feature feature) {

                try {

                    listener.onLifecycleEvent(event, feature);

                } catch (RemoteException e) {

                    Slog.e(TAG, "Error sending onLifecycleEvent: ", e);

                }

            }

        };

    }



    private StreamingProcessingCallback wrapStreamingResponseCallback(

            IStreamingResponseCallback callback) {

        return new StreamingProcessingCallback() {