Add SmartDimMlAgent

SmartDimMlAgent is introduced to support loading model and preprocessor
from either rootfs or component update service. SmartDim logic can just
request decisions from SmartDimMlAgent without needing to know about it.

SmartDimMlAgent is a replacement for model_impl.* and
ml_serfvice_client.*, but they will coexist for a while, until the flag
kSmartDimNewMlAgent rolls out to 100% of the devices.

Bug: 1018065
Test: tast test power.SmartDim passes
Change-Id: I6bb16c642476775d14836de27ca3a56fa98a8e5a
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1903168Reviewed-by: Andrew Moylan <amoylan@chromium.org>
Reviewed-by: Jia Meng <jiameng@chromium.org>
Commit-Queue: Xinglong Luan <alanlxl@chromium.org>
Cr-Commit-Position: refs/heads/master@{#743877}

chrome/browser/chromeos/power/ml/smart_dim/BUILD.gn

@@ -4,16 +4,27 @@
 
 source_set("smart_dim") {
   public = [
+    "ml_agent.h",
     "model.h",
     "model_impl.h",
   ]
 
   sources = [
+    "builtin_worker.cc",
+    "builtin_worker.h",
+    "download_worker.cc",
+    "download_worker.h",
+    "ml_agent.cc",
+    "ml_agent.h",
+    "ml_agent_util.cc",
+    "ml_agent_util.h",
     "ml_service_client.cc",
     "ml_service_client.h",
     "model.h",
     "model_impl.cc",
     "model_impl.h",
+    "smart_dim_worker.cc",
+    "smart_dim_worker.h",
   ]
 
   deps = [

chrome/browser/chromeos/power/ml/smart_dim/builtin_worker.cc

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chrome/browser/chromeos/power/ml/smart_dim/builtin_worker.h"
+
+#include "base/bind.h"
+#include "base/memory/ref_counted_memory.h"
+#include "chrome/browser/chromeos/power/ml/smart_dim/ml_agent_util.h"
+#include "chrome/grit/browser_resources.h"
+#include "chromeos/constants/chromeos_features.h"
+#include "chromeos/services/machine_learning/public/cpp/service_connection.h"
+#include "components/assist_ranker/proto/example_preprocessor.pb.h"
+#include "ui/base/resource/resource_bundle.h"
+
+namespace chromeos {
+namespace power {
+namespace ml {
+
+namespace {
+
+using ::chromeos::machine_learning::mojom::BuiltinModelId;
+using ::chromeos::machine_learning::mojom::BuiltinModelSpec;
+using ::chromeos::machine_learning::mojom::BuiltinModelSpecPtr;
+
+constexpr size_t k20181115ModelInputVectorSize = 343;
+constexpr size_t k20190521ModelInputVectorSize = 592;
+
+constexpr double k20181115ModelDefaultDimThreshold = -1.0;
+constexpr double k20190521ModelDefaultDimThreshold = -0.6;
+
+}  // namespace
+
+BuiltinWorker::BuiltinWorker() : SmartDimWorker() {}
+
+BuiltinWorker::~BuiltinWorker() = default;
+
+const assist_ranker::ExamplePreprocessorConfig*
+BuiltinWorker::GetPreprocessorConfig() {
+  LazyInitialize();
+  return preprocessor_config_.get();
+}
+
+const mojo::Remote<::chromeos::machine_learning::mojom::GraphExecutor>&
+BuiltinWorker::GetExecutor() {
+  LazyInitialize();
+  return executor_;
+}
+
+void BuiltinWorker::LazyInitialize() {
+  const bool v3_enabled =
+      base::FeatureList::IsEnabled(features::kSmartDimModelV3);
+
+  // Initialize builtin meta info.
+  dim_threshold_ = v3_enabled ? k20190521ModelDefaultDimThreshold
+                              : k20181115ModelDefaultDimThreshold;
+  expected_feature_size_ = v3_enabled ? k20190521ModelInputVectorSize
+                                      : k20181115ModelInputVectorSize;
+
+  if (!preprocessor_config_) {
+    preprocessor_config_ =
+        std::make_unique<assist_ranker::ExamplePreprocessorConfig>();
+
+    const int resource_id =
+        v3_enabled ? IDR_SMART_DIM_20190521_EXAMPLE_PREPROCESSOR_CONFIG_PB
+                   : IDR_SMART_DIM_20181115_EXAMPLE_PREPROCESSOR_CONFIG_PB;
+
+    const scoped_refptr<base::RefCountedMemory> raw_config =
+        ui::ResourceBundle::GetSharedInstance().LoadDataResourceBytes(
+            resource_id);
+    if (!raw_config || !raw_config->front()) {
+      DLOG(FATAL)
+          << "Failed to load Builtin SmartDimModel example preprocessor "
+             "config.";
+      return;
+    }
+
+    if (!preprocessor_config_->ParseFromArray(raw_config->front(),
+                                              raw_config->size())) {
+      DLOG(FATAL) << "Failed to parse Builtin SmartDimModel example "
+                     "preprocessor config.";
+      preprocessor_config_.reset();
+      return;
+    }
+  }
+
+  if (!model_) {
+    // Load the model.
+    BuiltinModelSpecPtr spec =
+        BuiltinModelSpec::New(v3_enabled ? BuiltinModelId::SMART_DIM_20190521
+                                         : BuiltinModelId::SMART_DIM_20181115);
+    chromeos::machine_learning::ServiceConnection::GetInstance()
+        ->LoadBuiltinModel(std::move(spec), model_.BindNewPipeAndPassReceiver(),
+                           base::BindOnce(&LoadModelCallback));
+  }
+
+  if (!executor_) {
+    // Get the graph executor.
+    model_->CreateGraphExecutor(executor_.BindNewPipeAndPassReceiver(),
+                                base::BindOnce(&CreateGraphExecutorCallback));
+    executor_.set_disconnect_handler(base::BindOnce(
+        &BuiltinWorker::OnConnectionError, base::Unretained(this)));
+  }
+}
+
+}  // namespace ml
+}  // namespace power
+}  // namespace chromeos

chrome/browser/chromeos/power/ml/smart_dim/builtin_worker.h

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_BUILTIN_WORKER_H_
+#define CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_BUILTIN_WORKER_H_
+
+#include <memory>
+
+#include "base/containers/flat_map.h"
+#include "chrome/browser/chromeos/power/ml/smart_dim/smart_dim_worker.h"
+#include "chromeos/services/machine_learning/public/mojom/graph_executor.mojom.h"
+#include "chromeos/services/machine_learning/public/mojom/model.mojom.h"
+#include "mojo/public/cpp/bindings/remote.h"
+
+namespace chromeos {
+namespace power {
+namespace ml {
+
+// SmartDimWorker that uses a built-in preprocessor config and ML service model
+// file.
+class BuiltinWorker : public SmartDimWorker {
+ public:
+  BuiltinWorker();
+  ~BuiltinWorker() override;
+
+  // SmartDimWorker overrides:
+  const assist_ranker::ExamplePreprocessorConfig* GetPreprocessorConfig()
+      override;
+  const mojo::Remote<::chromeos::machine_learning::mojom::GraphExecutor>&
+  GetExecutor() override;
+
+ private:
+  // Loads the built-in preprocessor config if not loaded yet. Also
+  // initializes the model_ and executor_ with built-in model.
+  void LazyInitialize();
+
+  DISALLOW_COPY_AND_ASSIGN(BuiltinWorker);
+};
+
+}  // namespace ml
+}  // namespace power
+}  // namespace chromeos
+
+#endif  // CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_BUILTIN_WORKER_H_

chrome/browser/chromeos/power/ml/smart_dim/download_worker.cc

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chrome/browser/chromeos/power/ml/smart_dim/download_worker.h"
+
+#include "base/bind.h"
+#include "base/task/post_task.h"
+#include "base/task/task_traits.h"
+#include "base/threading/sequenced_task_runner_handle.h"
+#include "chrome/browser/chromeos/power/ml/smart_dim/ml_agent_util.h"
+#include "chromeos/services/machine_learning/public/cpp/service_connection.h"
+#include "components/assist_ranker/proto/example_preprocessor.pb.h"
+#include "content/public/browser/browser_task_traits.h"
+#include "content/public/browser/browser_thread.h"
+#include "ui/base/resource/resource_bundle.h"
+
+namespace chromeos {
+namespace power {
+namespace ml {
+
+namespace {
+using ::chromeos::machine_learning::mojom::FlatBufferModelSpec;
+}  // namespace
+
+DownloadWorker::DownloadWorker() : SmartDimWorker(), metrics_model_name_("") {}
+
+DownloadWorker::~DownloadWorker() = default;
+
+const assist_ranker::ExamplePreprocessorConfig*
+DownloadWorker::GetPreprocessorConfig() {
+  return preprocessor_config_.get();
+}
+
+const mojo::Remote<::chromeos::machine_learning::mojom::GraphExecutor>&
+DownloadWorker::GetExecutor() {
+  return executor_;
+}
+
+bool DownloadWorker::IsReady() {
+  return preprocessor_config_ && model_ && executor_ &&
+         expected_feature_size_ > 0 && metrics_model_name_ != "";
+}
+
+void DownloadWorker::InitializeFromComponent(
+    const std::string& metadata_json,
+    const std::string& preprocessor_proto,
+    const std::string& model_flatbuffer) {
+  // Meta data contains necessary info to construct model_spec_, and other
+  // optional info.
+  // TODO(crbug.com/1049886) move json parsing to the sandboxed separate parser.
+  // TODO(crbug.com/1049888) add new UMA metrics to log the json errors.
+  if (!ParseMetaInfoFromString(std::move(metadata_json), &metrics_model_name_,
+                               &dim_threshold_, &expected_feature_size_,
+                               &inputs_, &outputs_)) {
+    DVLOG(1) << "Failed to parse metadata_json.";
+    return;
+  }
+
+  preprocessor_config_ =
+      std::make_unique<assist_ranker::ExamplePreprocessorConfig>();
+  if (!preprocessor_config_->ParseFromString(preprocessor_proto)) {
+    DVLOG(1) << "Failed to load preprocessor_config.";
+    preprocessor_config_.reset();
+    return;
+  }
+
+  base::PostTask(
+      FROM_HERE, {content::BrowserThread::UI, base::TaskPriority::BEST_EFFORT},
+      base::BindOnce(&DownloadWorker::LoadModelAndCreateGraphExecutor,
+                     base::Unretained(this), std::move(model_flatbuffer)));
+}
+
+void DownloadWorker::LoadModelAndCreateGraphExecutor(
+    const std::string& model_flatbuffer) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+  DCHECK(!model_.is_bound() && !executor_.is_bound());
+
+  chromeos::machine_learning::ServiceConnection::GetInstance()
+      ->LoadFlatBufferModel(
+          FlatBufferModelSpec::New(std::move(model_flatbuffer), inputs_,
+                                   outputs_, metrics_model_name_),
+          model_.BindNewPipeAndPassReceiver(),
+          base::BindOnce(&LoadModelCallback));
+  model_->CreateGraphExecutor(executor_.BindNewPipeAndPassReceiver(),
+                              base::BindOnce(&CreateGraphExecutorCallback));
+  executor_.set_disconnect_handler(base::BindOnce(
+      &DownloadWorker::OnConnectionError, base::Unretained(this)));
+}
+
+}  // namespace ml
+}  // namespace power
+}  // namespace chromeos

chrome/browser/chromeos/power/ml/smart_dim/download_worker.h

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_DOWNLOAD_WORKER_H_
+#define CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_DOWNLOAD_WORKER_H_
+
+#include "base/containers/flat_map.h"
+#include "chrome/browser/chromeos/power/ml/smart_dim/smart_dim_worker.h"
+
+namespace chromeos {
+namespace power {
+namespace ml {
+
+// SmartDimWorker that loads meta info, preprocessor config and ML service model
+// files from smart dim components.
+class DownloadWorker : public SmartDimWorker {
+ public:
+  DownloadWorker();
+  ~DownloadWorker() override;
+
+  // SmartDimWorker overrides:
+  const assist_ranker::ExamplePreprocessorConfig* GetPreprocessorConfig()
+      override;
+  const mojo::Remote<::chromeos::machine_learning::mojom::GraphExecutor>&
+  GetExecutor() override;
+
+  // Returns true if it has loaded components successfully.
+  bool IsReady();
+
+  // Loads meta info, preprocessor config and ML service model from smart dim
+  // components.
+  // Called by component updater when it gets a verified smart dim component and
+  // DownloadWorker is not ready.
+  // If IsReady(), this function won't be called again.
+  void InitializeFromComponent(const std::string& metadata_json,
+                               const std::string& preprocessor_proto,
+                               const std::string& model_flatbuffer);
+
+ private:
+  base::flat_map<std::string, int> inputs_;
+  base::flat_map<std::string, int> outputs_;
+  std::string metrics_model_name_;
+
+  void LoadModelAndCreateGraphExecutor(const std::string& model_flatbuffer);
+
+  DISALLOW_COPY_AND_ASSIGN(DownloadWorker);
+};
+
+}  // namespace ml
+}  // namespace power
+}  // namespace chromeos
+
+#endif  // CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_DOWNLOAD_WORKER_H_

chrome/browser/chromeos/power/ml/smart_dim/ml_agent.h

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_ML_AGENT_H_
+#define CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_ML_AGENT_H_
+
+#include <string>
+
+#include "base/cancelable_callback.h"
+#include "base/no_destructor.h"
+#include "chrome/browser/chromeos/power/ml/smart_dim/builtin_worker.h"
+#include "chrome/browser/chromeos/power/ml/smart_dim/download_worker.h"
+#include "chrome/browser/chromeos/power/ml/smart_dim/model_impl.h"
+#include "chrome/browser/chromeos/power/ml/smart_dim/smart_dim_worker.h"
+#include "chrome/browser/chromeos/power/ml/user_activity_event.pb.h"
+
+namespace chromeos {
+namespace power {
+namespace ml {
+
+using DimDecisionCallback =
+    base::OnceCallback<void(UserActivityEvent::ModelPrediction)>;
+
+// SmartDimMlAgent is responsible for preprocessing the features and requesting
+// the inference from machine learning service.
+// Usage:
+//
+//     SmartDimMlAgent::GetInstance()->RequestDimDecision(
+//         features_, dim_decision_callback);
+class SmartDimMlAgent {
+ public:
+  static SmartDimMlAgent* GetInstance();
+
+  // Post a request to determine whether an upcoming dim should go ahead based
+  // on input |features|. When a decision is arrived at, runs the callback. If
+  // this method is called again before it calls the previous callback, the
+  // previous callback will be canceled.
+  void RequestDimDecision(const UserActivityEvent::Features& features,
+                          DimDecisionCallback callback);
+  void CancelPreviousRequest();
+
+  // Called by CUS(component update service). When new version of the component
+  // downloaded, CUS first uses IsDownloadWorkerReady to see if download worker
+  // is ready. If it's not, CUS then uses OnComponentReady to update the
+  // download metainfo, preprocessor and model.
+  bool IsDownloadWorkerReady();
+  void OnComponentReady(const std::string& metadata_json,
+                        const std::string& preprocessor_proto,
+                        const std::string& model_flatbuffer);
+
+ protected:
+  SmartDimMlAgent();
+  virtual ~SmartDimMlAgent();
+
+ private:
+  friend base::NoDestructor<SmartDimMlAgent>;
+
+  // Return download_worker_ if it's ready, otherwise builtin_worker_.
+  SmartDimWorker* GetWorker();
+
+  BuiltinWorker builtin_worker_;
+  DownloadWorker download_worker_;
+
+  base::CancelableOnceCallback<void(UserActivityEvent::ModelPrediction)>
+      dim_decision_callback_;
+
+  SEQUENCE_CHECKER(sequence_checker_);
+
+  DISALLOW_COPY_AND_ASSIGN(SmartDimMlAgent);
+};
+
+}  // namespace ml
+}  // namespace power
+}  // namespace chromeos
+
+#endif  // CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_ML_AGENT_H_

chrome/browser/chromeos/power/ml/smart_dim/ml_agent_util.cc

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chrome/browser/chromeos/power/ml/smart_dim/ml_agent_util.h"
+
+#include <string>
+
+#include "base/json/json_reader.h"
+#include "base/values.h"
+
+namespace chromeos {
+namespace power {
+namespace ml {
+
+namespace {
+
+// Extracts node names and id, then inserts into the name_2_node_map.
+bool PopulateMapFromNamesAndNodes(
+    const base::Value& names,
+    const base::Value& nodes,
+    base::flat_map<std::string, int>* name_2_node_map) {
+  if (names.GetList().size() != nodes.GetList().size()) {
+    DVLOG(1) << "names and nodes don't match";
+    return false;
+  }
+
+  for (size_t i = 0; i < names.GetList().size(); i++) {
+    if (!names.GetList()[i].is_string() || !nodes.GetList()[i].is_int()) {
+      DVLOG(1) << i << " names and nodes have unexpected type";
+      return false;
+    }
+    name_2_node_map->emplace(names.GetList()[i].GetString(),
+                             nodes.GetList()[i].GetInt());
+  }
+  return true;
+}
+
+}  // namespace
+
+bool ParseMetaInfoFromString(const std::string& metadata_json,
+                             std::string* metrics_model_name,
+                             double* dim_threshold,
+                             size_t* expected_feature_size,
+                             base::flat_map<std::string, int>* inputs,
+                             base::flat_map<std::string, int>* outputs) {
+  DCHECK(metrics_model_name && dim_threshold && expected_feature_size &&
+         inputs && outputs);
+  const base::Optional<base::Value> root =
+      base::JSONReader::Read(metadata_json);
+
+  if (!root || !root->is_dict()) {
+    DVLOG(1) << "Failed to read metadata_json as JSON dict.";
+    return false;
+  }
+
+  const std::string* metrics_model_name_value =
+      root->FindStringKey("metrics_model_name");
+  const base::Optional<double> dim_threshold_value =
+      root->FindDoubleKey("threshold");
+  const base::Optional<int> expected_feature_size_value =
+      root->FindIntKey("expected_feature_size");
+
+  if (!metrics_model_name_value || *metrics_model_name_value == "" ||
+      dim_threshold_value == base::nullopt ||
+      expected_feature_size_value == base::nullopt) {
+    DVLOG(1) << "metadata_json missing expected field(s).";
+    return false;
+  }
+
+  *metrics_model_name = *metrics_model_name_value;
+  *dim_threshold = dim_threshold_value.value();
+  *expected_feature_size =
+      static_cast<size_t>(expected_feature_size_value.value());
+
+  const base::Value* input_names = root->FindListKey("input_names");
+  const base::Value* input_nodes = root->FindListKey("input_nodes");
+  const base::Value* output_names = root->FindListKey("output_names");
+  const base::Value* output_nodes = root->FindListKey("output_nodes");
+
+  if (!input_names || !input_nodes || !output_names || !output_nodes ||
+      !PopulateMapFromNamesAndNodes(*input_names, *input_nodes, inputs) ||
+      !PopulateMapFromNamesAndNodes(*output_names, *output_nodes, outputs)) {
+    DVLOG(1) << "Failed to load inputs and outputs maps from metadata_json";
+    *metrics_model_name = "";
+    *dim_threshold = 0.0;
+    *expected_feature_size = 0;
+    inputs->clear();
+    outputs->clear();
+    return false;
+  }
+
+  return true;
+}
+
+void LoadModelCallback(LoadModelResult result) {
+  if (result != LoadModelResult::OK) {
+    DVLOG(1) << "Failed to load Smart Dim flatbuffer model.";
+  }
+}
+
+void CreateGraphExecutorCallback(CreateGraphExecutorResult result) {
+  if (result != CreateGraphExecutorResult::OK) {
+    DVLOG(1) << "Failed to create a Smart Dim graph executor.";
+  }
+}
+
+}  // namespace ml
+}  // namespace power
+}  // namespace chromeos

chrome/browser/chromeos/power/ml/smart_dim/ml_agent_util.h

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_ML_AGENT_UTIL_H_
+#define CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_ML_AGENT_UTIL_H_
+
+#include <string>
+
+#include "base/containers/flat_map.h"
+#include "chromeos/services/machine_learning/public/mojom/graph_executor.mojom.h"
+#include "chromeos/services/machine_learning/public/mojom/machine_learning_service.mojom.h"
+#include "chromeos/services/machine_learning/public/mojom/model.mojom.h"
+
+namespace chromeos {
+namespace power {
+namespace ml {
+
+using ::chromeos::machine_learning::mojom::CreateGraphExecutorResult;
+using ::chromeos::machine_learning::mojom::LoadModelResult;
+
+// Parse JSON dict and extract essential fields from metadata_json.
+bool ParseMetaInfoFromString(const std::string& metadata_json,
+                             std::string* metrics_model_name,
+                             double* dim_threshold,
+                             size_t* expected_feature_size,
+                             base::flat_map<std::string, int>* inputs,
+                             base::flat_map<std::string, int>* outputs);
+
+// Callbacks for LoadModel and CreateGraphExecutor used by BuiltinWorker and
+// DownloadWorker.
+void LoadModelCallback(LoadModelResult result);
+
+void CreateGraphExecutorCallback(CreateGraphExecutorResult result);
+
+}  // namespace ml
+}  // namespace power
+}  // namespace chromeos
+
+#endif  // CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_ML_AGENT_UTIL_H_

chrome/browser/chromeos/power/ml/smart_dim/smart_dim_worker.cc

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chrome/browser/chromeos/power/ml/smart_dim/smart_dim_worker.h"
+
+#include "components/assist_ranker/proto/example_preprocessor.pb.h"
+
+namespace chromeos {
+namespace power {
+namespace ml {
+
+namespace {
+
+using ::chromeos::machine_learning::mojom::FlatBufferModelSpec;
+
+}  // namespace
+
+SmartDimWorker::SmartDimWorker()
+    : dim_threshold_(0.0), expected_feature_size_(0) {}
+
+SmartDimWorker::~SmartDimWorker() = default;
+
+double SmartDimWorker::dim_threshold() const {
+  return dim_threshold_;
+}
+
+size_t SmartDimWorker::expected_feature_size() const {
+  return expected_feature_size_;
+}
+
+void SmartDimWorker::Reset() {
+  preprocessor_config_.reset();
+  executor_.reset();
+  model_.reset();
+}
+
+void SmartDimWorker::OnConnectionError() {
+  DVLOG(1) << "Mojo connection for ML service is closed!";
+  executor_.reset();
+  model_.reset();
+}
+
+}  // namespace ml
+}  // namespace power
+}  // namespace chromeos

chrome/browser/chromeos/power/ml/smart_dim/smart_dim_worker.h

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_SMART_DIM_WORKER_H_
+#define CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_SMART_DIM_WORKER_H_
+
+#include <memory>
+
+#include "chromeos/services/machine_learning/public/mojom/graph_executor.mojom.h"
+#include "chromeos/services/machine_learning/public/mojom/model.mojom.h"
+#include "mojo/public/cpp/bindings/remote.h"
+
+namespace assist_ranker {
+class ExamplePreprocessorConfig;
+}  // namespace assist_ranker
+
+namespace chromeos {
+namespace power {
+namespace ml {
+
+// This is the common interface & base class that holds the preprocessor, model
+// graph executor and other essential fields like dim_threshold,
+// expected_feature_size for SmartDimMlAgent to make decision.
+// It has two child classes, BuiltinWorker and DownloadWorker, both are members
+// of SmartDimMlAgent and are created when SmartDimWorker is created.
+// DownloadWorker is not ready until InitializeFromComponent is successfully
+// called by component update service, while BuiltinWorker can be lazily
+// initialized when GetPreprocessorConfig() or GetExecutor() are called.
+class SmartDimWorker {
+ public:
+  SmartDimWorker();
+  virtual ~SmartDimWorker();
+
+  // Gets model score threshold above which the screen dim is recommended.
+  double dim_threshold() const;
+  // Gets expected feature size.
+  size_t expected_feature_size() const;
+
+  // Returns a preprocessor_config for SmartDimMlAgent to convert feature proto
+  // to a vector.
+  virtual const assist_ranker::ExamplePreprocessorConfig*
+  GetPreprocessorConfig() = 0;
+  // Returns a mojo remote of ML service GraphExecutor to make inference.
+  virtual const mojo::Remote<
+      ::chromeos::machine_learning::mojom::GraphExecutor>&
+  GetExecutor() = 0;
+
+  // Release the members on connection error, or when download_worker_ is ready
+  // we can reset the builtin_worker_ to save memory.
+  void Reset();
+
+ protected:
+  void OnConnectionError();
+
+  double dim_threshold_;
+  size_t expected_feature_size_;
+  std::unique_ptr<assist_ranker::ExamplePreprocessorConfig>
+      preprocessor_config_;
+
+  // Remotes used to execute functions in the ML service side.
+  mojo::Remote<::chromeos::machine_learning::mojom::Model> model_;
+  mojo::Remote<::chromeos::machine_learning::mojom::GraphExecutor> executor_;
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(SmartDimWorker);
+};
+
+}  // namespace ml
+}  // namespace power
+}  // namespace chromeos
+
+#endif  // CHROME_BROWSER_CHROMEOS_POWER_ML_SMART_DIM_SMART_DIM_WORKER_H_

chrome/browser/chromeos/power/ml/user_activity_manager.cc

@@ -11,6 +11,7 @@
 #include "base/metrics/field_trial_params.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/strings/string_number_conversions.h"
+#include "chrome/browser/chromeos/power/ml/smart_dim/ml_agent.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/profiles/profile_manager.h"
 #include "chrome/browser/resource_coordinator/tab_metrics_logger.h"
@@ -263,13 +264,18 @@ void UserActivityManager::UpdateAndGetSmartDimDecision(
   }
   if (smart_dim_enabled &&
       base::FeatureList::IsEnabled(features::kUserActivityPrediction) &&
-      smart_dim_model_) {
+      SmartDimModelReady()) {
     waiting_for_model_decision_ = true;
     time_dim_decision_requested_ = base::TimeTicks::Now();
-    smart_dim_model_->RequestDimDecision(
-        features_,
+    auto request_callback =
         base::BindOnce(&UserActivityManager::HandleSmartDimDecision,
-                       weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
+                       weak_ptr_factory_.GetWeakPtr(), std::move(callback));
+    if (base::FeatureList::IsEnabled(features::kSmartDimNewMlAgent))
+      SmartDimMlAgent::GetInstance()->RequestDimDecision(
+          features_, std::move(request_callback));
+    else
+      smart_dim_model_->RequestDimDecision(features_,
+                                           std::move(request_callback));
   }
   waiting_for_final_action_ = true;
 }
@@ -555,7 +561,7 @@ void UserActivityManager::PopulatePreviousEventData(
   PreviousEventLoggingResult result = PreviousEventLoggingResult::kSuccess;
   if (!model_prediction_) {
     result = base::FeatureList::IsEnabled(features::kUserActivityPrediction) &&
-                     smart_dim_model_
+                     SmartDimModelReady()
                  ? PreviousEventLoggingResult::kErrorModelPredictionMissing
                  : PreviousEventLoggingResult::kErrorModelDisabled;
     LogPowerMLPreviousEventLoggingResult(result);
@@ -599,7 +605,11 @@ void UserActivityManager::ResetAfterLogging() {
 
 void UserActivityManager::CancelDimDecisionRequest() {
   LOG(WARNING) << "Cancelling pending Smart Dim decision request.";
-  smart_dim_model_->CancelPreviousRequest();
+  if (base::FeatureList::IsEnabled(features::kSmartDimNewMlAgent))
+    SmartDimMlAgent::GetInstance()->CancelPreviousRequest();
+  else
+    smart_dim_model_->CancelPreviousRequest();
+
   waiting_for_model_decision_ = false;
   const base::TimeDelta wait_time =
       base::TimeTicks::Now() - time_dim_decision_requested_;
@@ -607,6 +617,13 @@ void UserActivityManager::CancelDimDecisionRequest() {
   time_dim_decision_requested_ = base::TimeTicks();
 }
 
+bool UserActivityManager::SmartDimModelReady() {
+  // We assume that BuiltinWorker of SmartDimMlAgent can always load model and
+  // preprocessor config from rootfs, therefore SmartDimMlAgent is always ready.
+  return base::FeatureList::IsEnabled(features::kSmartDimNewMlAgent) ||
+         smart_dim_model_;
+}
+
 }  // namespace ml
 }  // namespace power
 }  // namespace chromeos

chrome/browser/chromeos/power/ml/user_activity_manager.h

@@ -158,6 +158,9 @@ class UserActivityManager : public ui::UserActivityObserver,
   // Cancel any pending request to |smart_dim_model_| to get a dim decision.
   void CancelDimDecisionRequest();
 
+  // If old smart dim model or new SmartDimMlAgent is ready, based on Finch.
+  bool SmartDimModelReady();
+
   // Time when an idle event is received and we start logging. Null if an idle
   // event hasn't been observed.
   base::Optional<base::TimeDelta> idle_event_start_since_boot_;

chromeos/constants/chromeos_features.cc

@@ -258,6 +258,11 @@ const base::Feature kShowBluetoothDeviceBattery{
 const base::Feature kShowPlayInDemoMode{"ShowPlayInDemoMode",
                                         base::FEATURE_DISABLED_BY_DEFAULT};
 
+// Uses the smart dim component updater to provide smart dim model and
+// preprocessor configuration.
+const base::Feature kSmartDimNewMlAgent{"SmartDimNewMlAgent",
+                                        base::FEATURE_DISABLED_BY_DEFAULT};
+
 // Uses the V3 (~2019-05 era) Smart Dim model instead of the default V2
 // (~2018-11) model.
 const base::Feature kSmartDimModelV3{"SmartDimModelV3",

chromeos/constants/chromeos_features.h

@@ -121,6 +121,8 @@ extern const base::Feature kShowBluetoothDeviceBattery;
 COMPONENT_EXPORT(CHROMEOS_CONSTANTS)
 extern const base::Feature kShowPlayInDemoMode;
 COMPONENT_EXPORT(CHROMEOS_CONSTANTS)
+extern const base::Feature kSmartDimNewMlAgent;
+COMPONENT_EXPORT(CHROMEOS_CONSTANTS)
 extern const base::Feature kSmartDimModelV3;
 COMPONENT_EXPORT(CHROMEOS_CONSTANTS)
 extern const base::Feature kSplitSettingsSync;