1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
|
// Copyright 2023 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#ifndef COMPONENTS_SEGMENTATION_PLATFORM_INTERNAL_EXECUTION_MODEL_MANAGER_IMPL_H_
#define COMPONENTS_SEGMENTATION_PLATFORM_INTERNAL_EXECUTION_MODEL_MANAGER_IMPL_H_
#include <map>
#include <memory>
#include <optional>
#include <vector>
#include "base/containers/flat_set.h"
#include "base/functional/callback.h"
#include "base/memory/raw_ptr.h"
#include "base/memory/weak_ptr.h"
#include "components/segmentation_platform/internal/database/segment_info_database.h"
#include "components/segmentation_platform/internal/execution/model_manager.h"
#include "components/segmentation_platform/public/proto/segmentation_platform.pb.h"
namespace base {
class Clock;
} // namespace base
namespace segmentation_platform {
class ModelProvider;
class ModelProviderFactory;
namespace proto {
class SegmentInfo;
} // namespace proto
// ModelManager that owns and manages non-default (OptimizationGuide
// based) ModelProvider(s) and uses SegmentInfoDatabase (metadata) for storing
// the latest model from ModelProvider. The vector of OptimizationTargets need
// to be passed in at construction time so the SegmentationModelHandler
// instances can be created early.
class ModelManagerImpl : public ModelManager {
public:
ModelManagerImpl(
const base::flat_set<SegmentId>& segment_ids,
ModelProviderFactory* model_provider_factory,
base::Clock* clock,
SegmentInfoDatabase* segment_database,
const SegmentationModelUpdatedCallback& model_updated_callback);
~ModelManagerImpl() override;
// Disallow copy/assign.
ModelManagerImpl(const ModelManagerImpl&) = delete;
ModelManagerImpl& operator=(const ModelManagerImpl&) = delete;
void Initialize() override;
// ModelManager override:
ModelProvider* GetModelProvider(proto::SegmentId segment_id,
proto::ModelSource model_source) override;
void SetSegmentationModelUpdatedCallbackForTesting(
ModelManager::SegmentationModelUpdatedCallback model_updated_callback)
override;
private:
friend class SegmentationPlatformServiceImplTest;
friend class TestServicesForPlatform;
// Callback for whenever a SegmentationModelHandler is informed that the
// underlying ML model file has been updated. If there is an available
// model, this will be called at least once per session.
void OnSegmentationModelUpdated(
proto::ModelSource model_source,
proto::SegmentId segment_id,
std::optional<proto::SegmentationModelMetadata> metadata,
int64_t model_version);
// Callback after fetching the current SegmentInfo from the
// SegmentInfoDatabase. This is part of the flow for informing the
// SegmentationModelUpdatedCallback about a changed model.
// Merges the PredictionResult from the previously stored SegmentInfo with
// the newly updated one, and stores the new version in the DB.
void OnSegmentInfoFetchedForModelUpdate(
proto::SegmentId segment_id,
proto::ModelSource model_source,
proto::SegmentationModelMetadata metadata,
int64_t model_version,
const proto::SegmentInfo* segment_info);
// Callback after storing the updated version of the SegmentInfo.
// Responsible for invoking the SegmentationModelUpdatedCallback.
void OnUpdatedSegmentInfoStored(proto::SegmentInfo segment_info,
std::optional<int64_t> old_model_version,
bool success);
// Callback after deleting the previous version of the SegmentInfo.
void OnSegmentInfoDeleted(SegmentId segment_id,
proto::ModelSource model_source,
int64_t deleted_version,
bool success);
const base::flat_set<SegmentId>& segment_ids_;
// All the relevant handlers for each of the segments.
std::map<std::pair<SegmentId, proto::ModelSource>,
std::unique_ptr<ModelProvider>>
model_providers_;
// Creates model provider.
const raw_ptr<ModelProviderFactory> model_provider_factory_;
// Used to access the current time.
raw_ptr<base::Clock> clock_;
// Database for segment information and metadata.
raw_ptr<SegmentInfoDatabase> segment_database_;
// Invoked whenever there is an update to any of the relevant ML models.
SegmentationModelUpdatedCallback model_updated_callback_;
base::WeakPtrFactory<ModelManagerImpl> weak_ptr_factory_{this};
};
} // namespace segmentation_platform
#endif // COMPONENTS_SEGMENTATION_PLATFORM_INTERNAL_EXECUTION_MODEL_MANAGER_IMPL_H_
|
