Spaces:
Running
Running
File size: 4,117 Bytes
5f923cd | 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 | // Copyright 2026 The ODML Authors.
//
// 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.
#include "runtime/util/metrics_util.h"
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/time/clock.h" // from @com_google_absl
#include "absl/time/time.h" // from @com_google_absl
#include "runtime/engine/io_types.h"
#include "runtime/engine/litert_lm_lib.h"
#include "runtime/proto/engine.pb.h"
#include "runtime/proto/litert_lm_metrics.pb.h"
#include "runtime/util/test_utils.h" // IWYU pragma: keep
namespace litert::lm {
namespace {
using ::testing::Contains;
using ::testing::FloatEq;
using ::testing::Pair;
using ::testing::SizeIs;
TEST(MetricsUtilTest, EmptyMetrics) {
LitertLmMetrics metrics;
metrics.peak_mem_mb = 12.5f;
metrics.peak_private_mb = 10.0f;
ASSERT_OK_AND_ASSIGN(auto proto, ToProto(metrics));
EXPECT_THAT(proto.peak_mem_mb(), FloatEq(12.5f));
EXPECT_THAT(proto.peak_private_mb(), FloatEq(10.0f));
EXPECT_FALSE(proto.has_benchmark_params());
EXPECT_THAT(proto.init_phase_durations_us(), SizeIs(0));
EXPECT_THAT(proto.mark_durations_us(), SizeIs(0));
}
TEST(MetricsUtilTest, WithBenchmarkInfo) {
proto::BenchmarkParams params;
params.set_num_prefill_tokens(128);
params.set_num_decode_tokens(256);
BenchmarkInfo info(params);
// Simulate some init phases & turns
ASSERT_OK(info.InitPhaseRecord(BenchmarkInfo::InitPhase::kModelAssets,
absl::Microseconds(100)));
ASSERT_OK(info.TimePrefillTurnStart());
absl::SleepFor(absl::Milliseconds(2));
ASSERT_OK(info.TimePrefillTurnEnd(100));
ASSERT_OK(info.TimeDecodeTurnStart());
absl::SleepFor(absl::Milliseconds(2));
ASSERT_OK(info.TimeDecodeTurnEnd(1));
ASSERT_OK(info.TimeMarkDelta("test_mark"));
absl::SleepFor(absl::Milliseconds(2));
ASSERT_OK(info.TimeMarkDelta("test_mark"));
LitertLmMetrics metrics;
metrics.peak_mem_mb = 100.0f;
metrics.peak_private_mb = 80.0f;
metrics.benchmark_info = info;
ASSERT_OK_AND_ASSIGN(auto proto, ToProto(metrics));
EXPECT_THAT(proto.peak_mem_mb(), FloatEq(100.0f));
EXPECT_THAT(proto.peak_private_mb(), FloatEq(80.0f));
EXPECT_TRUE(proto.has_benchmark_params());
EXPECT_EQ(proto.benchmark_params().num_prefill_tokens(), 128);
EXPECT_EQ(proto.benchmark_params().num_decode_tokens(), 256);
EXPECT_THAT(proto.init_phase_durations_us(), SizeIs(1));
EXPECT_THAT(proto.init_phase_durations_us(),
Contains(Pair("Init Model assets", 100)));
EXPECT_THAT(proto.mark_durations_us().size(), 1);
EXPECT_THAT(proto.mark_durations_us(), Contains(testing::Key("test_mark")));
EXPECT_GT(proto.mark_durations_us().at("test_mark"), 0);
EXPECT_EQ(proto.prefill_turns().size(), 1);
EXPECT_EQ(proto.prefill_turns(0).num_tokens(), 100);
EXPECT_EQ(proto.prefill_turns(0).tokens_per_second(),
info.GetPrefillTokensPerSec(0));
EXPECT_EQ(proto.decode_turns().size(), 1);
EXPECT_EQ(proto.decode_turns(0).num_tokens(), 1);
EXPECT_EQ(proto.decode_turns(0).tokens_per_second(),
info.GetDecodeTokensPerSec(0));
EXPECT_EQ(proto.text_to_token_ids_turns().size(), 0);
EXPECT_EQ(proto.time_to_first_token_seconds(), info.GetTimeToFirstToken());
}
TEST(MetricsUtilTest, VectorToProtoList) {
LitertLmMetrics m1;
m1.peak_mem_mb = 1.0f;
LitertLmMetrics m2;
m2.peak_mem_mb = 2.0f;
ASSERT_OK_AND_ASSIGN(auto proto, ToProtoList({m1, m2}));
EXPECT_EQ(proto.metrics_size(), 2);
EXPECT_THAT(proto.metrics(0).peak_mem_mb(), FloatEq(1.0f));
EXPECT_THAT(proto.metrics(1).peak_mem_mb(), FloatEq(2.0f));
}
} // namespace
} // namespace litert::lm
|