solpbc.org
personal memory agent
1# SPDX-License-Identifier: AGPL-3.0-only
2# Copyright (c) 2026 sol pbc
3
4from __future__ import annotations
5
6import pytest
7
8from solstone.think.providers import local_vulkan
9from solstone.think.providers.parakeet_placement import (
10 CO_FIT_MARGIN_MIB,
11 PARAKEET_WORST_CASE_MIB,
12 cpu_placement_suffix,
13 decide_parakeet_auto_placement,
14 discrete_hardware_gpu_count,
15 is_discrete,
16)
17
18PLACEMENT_LINE = (
19 "sol thinks on your GPU; transcription runs on your CPU on this machine"
20)
21
22
23def _device(
24 *,
25 index: int = 0,
26 name: str = "Test GPU",
27 device_type: int = local_vulkan.VK_TYPE_DISCRETE,
28 vram_mib: int = 6144,
29) -> local_vulkan.VulkanDevice:
30 return local_vulkan.VulkanDevice(
31 index=index,
32 name=name,
33 device_type=device_type,
34 vram_mib=vram_mib,
35 )
36
37
38def _decision(
39 vram_mib: int | None,
40 *,
41 selected_device_is_discrete: bool = True,
42 discrete_hardware_gpu_count: int = 1,
43 unified_memory: bool = False,
44 brain_lane_active: bool = True,
45):
46 return decide_parakeet_auto_placement(
47 vram_mib=vram_mib,
48 selected_device_is_discrete=selected_device_is_discrete,
49 discrete_hardware_gpu_count=discrete_hardware_gpu_count,
50 unified_memory=unified_memory,
51 brain_lane_active=brain_lane_active,
52 )
53
54
55def test_floor_tier_small_cards_force_cpu() -> None:
56 decision = _decision(6144)
57
58 assert decision.force_cpu is True
59 assert decision.reason_code == "co_location_requires_cpu"
60 assert decision.tier_name == "floor"
61 assert decision.tier_resident_mib == 4147
62 assert decision.parakeet_worst_case_mib == PARAKEET_WORST_CASE_MIB
63 assert decision.margin_mib == CO_FIT_MARGIN_MIB
64 assert decision.required_mib == 8118
65 assert decision.vram_mib == 6144
66
67
68@pytest.mark.parametrize(
69 ("vram_mib", "force_cpu"),
70 [
71 (6144, True),
72 (8117, True),
73 (8118, False),
74 (8192, False),
75 ],
76)
77def test_floor_tier_placement_boundary(vram_mib: int, force_cpu: bool) -> None:
78 decision = _decision(vram_mib)
79
80 assert decision.tier_name == "floor"
81 assert decision.tier_resident_mib == 4147
82 assert decision.parakeet_worst_case_mib == PARAKEET_WORST_CASE_MIB
83 assert decision.margin_mib == CO_FIT_MARGIN_MIB
84 assert decision.required_mib == 8118
85 assert decision.vram_mib == vram_mib
86 assert decision.force_cpu is force_cpu
87
88
89def test_capable_tier_unmeasured_residency_keeps_gpu() -> None:
90 decision = _decision(16000)
91
92 assert decision.force_cpu is False
93 assert decision.reason_code == "tier_residency_unmeasured"
94 assert decision.tier_name == "capable"
95 assert decision.tier_resident_mib is None
96 assert decision.required_mib is None
97
98
99@pytest.mark.parametrize(
100 ("kwargs", "reason_code"),
101 [
102 ({"vram_mib": None}, "vram_unknown"),
103 ({"vram_mib": 6144, "brain_lane_active": False}, "brain_lane_inactive"),
104 (
105 {"vram_mib": 6144, "discrete_hardware_gpu_count": 2},
106 "discrete_gpu_count_not_one",
107 ),
108 (
109 {"vram_mib": 6144, "selected_device_is_discrete": False},
110 "selected_device_not_discrete",
111 ),
112 ({"vram_mib": 6144, "unified_memory": True}, "unified_memory"),
113 ],
114)
115def test_non_matching_inputs_keep_today(kwargs: dict, reason_code: str) -> None:
116 decision = _decision(**kwargs)
117
118 assert decision.force_cpu is False
119 assert decision.reason_code == reason_code
120
121
122def test_is_discrete_centralizes_vulkan_classification() -> None:
123 assert is_discrete(_device(), local_vulkan) is True
124 assert (
125 is_discrete(
126 _device(device_type=local_vulkan.VK_TYPE_INTEGRATED),
127 local_vulkan,
128 )
129 is False
130 )
131
132
133def test_discrete_hardware_gpu_count_ignores_integrated_and_software() -> None:
134 devices = [
135 _device(index=0),
136 _device(index=1, device_type=local_vulkan.VK_TYPE_INTEGRATED),
137 _device(index=2, name="llvmpipe", device_type=local_vulkan.VK_TYPE_CPU),
138 ]
139
140 assert discrete_hardware_gpu_count(devices, local_vulkan) == 1
141
142
143def test_cpu_placement_suffix_owns_joiner_and_copy() -> None:
144 selected = _device(vram_mib=6144)
145
146 assert (
147 cpu_placement_suffix(
148 devices=[selected],
149 selected=selected,
150 local_vulkan=local_vulkan,
151 unified_memory=False,
152 brain_lane_active=True,
153 )
154 == f"; {PLACEMENT_LINE}"
155 )
156
157
158@pytest.mark.parametrize(
159 ("selected", "unified_memory", "brain_lane_active"),
160 [
161 (None, False, True),
162 (_device(vram_mib=6144), True, True),
163 (_device(vram_mib=6144), False, False),
164 (_device(vram_mib=12288), False, True),
165 ],
166)
167def test_cpu_placement_suffix_absent_outside_predicate(
168 selected: local_vulkan.VulkanDevice | None,
169 unified_memory: bool,
170 brain_lane_active: bool,
171) -> None:
172 devices = [selected] if selected is not None else []
173
174 assert (
175 cpu_placement_suffix(
176 devices=devices,
177 selected=selected,
178 local_vulkan=local_vulkan,
179 unified_memory=unified_memory,
180 brain_lane_active=brain_lane_active,
181 )
182 == ""
183 )