deltav 0.0.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- deltav-0.0.0/.gitignore +12 -0
- deltav-0.0.0/PKG-INFO +47 -0
- deltav-0.0.0/README.md +33 -0
- deltav-0.0.0/pyproject.toml +36 -0
- deltav-0.0.0/src/deltav/__init__.py +53 -0
- deltav-0.0.0/src/deltav/core.py +178 -0
- deltav-0.0.0/src/deltav/store_convex.py +276 -0
- deltav-0.0.0/src/deltav/store_local.py +63 -0
- deltav-0.0.0/src/deltav/store_postgres.py +205 -0
- deltav-0.0.0/src/deltav/store_remote.py +140 -0
- deltav-0.0.0/tests/test_budget.py +33 -0
- deltav-0.0.0/tests/test_store_convex.py +42 -0
- deltav-0.0.0/tests/test_store_postgres.py +75 -0
- deltav-0.0.0/tests/test_store_remote.py +55 -0
- deltav-0.0.0/uv.lock +272 -0
deltav-0.0.0/.gitignore
ADDED
deltav-0.0.0/PKG-INFO
ADDED
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: deltav
|
|
3
|
+
Version: 0.0.0
|
|
4
|
+
Summary: DeltaV reservation-based budget enforcement for agent and tool runs.
|
|
5
|
+
Project-URL: Homepage, https://github.com/cachetronaut/deltav
|
|
6
|
+
Project-URL: Repository, https://github.com/cachetronaut/deltav
|
|
7
|
+
Project-URL: Issues, https://github.com/cachetronaut/deltav/issues
|
|
8
|
+
License: MIT
|
|
9
|
+
Requires-Python: >=3.11
|
|
10
|
+
Requires-Dist: dockbay>=0.0.0
|
|
11
|
+
Requires-Dist: psycopg-pool>=3.2
|
|
12
|
+
Requires-Dist: psycopg[binary]>=3.2
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
|
|
15
|
+
# deltav
|
|
16
|
+
|
|
17
|
+
Python implementation of DeltaV.
|
|
18
|
+
|
|
19
|
+
For product-level context, shared contracts, and cross-language repository information, see the public repository: https://github.com/cachetronaut/deltav.
|
|
20
|
+
|
|
21
|
+
## Install
|
|
22
|
+
|
|
23
|
+
```sh
|
|
24
|
+
pip install deltav
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
## Import
|
|
28
|
+
|
|
29
|
+
```python
|
|
30
|
+
import deltav
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
## Development
|
|
34
|
+
|
|
35
|
+
Run from `py/`:
|
|
36
|
+
|
|
37
|
+
```sh
|
|
38
|
+
uv sync --dev
|
|
39
|
+
uv run --with ruff ruff check .
|
|
40
|
+
uv run --with ruff ruff format --check .
|
|
41
|
+
uv run --with ty ty check
|
|
42
|
+
uv run --with pytest --with pytest-asyncio python -m pytest
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
## License
|
|
46
|
+
|
|
47
|
+
MIT
|
deltav-0.0.0/README.md
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
# deltav
|
|
2
|
+
|
|
3
|
+
Python implementation of DeltaV.
|
|
4
|
+
|
|
5
|
+
For product-level context, shared contracts, and cross-language repository information, see the public repository: https://github.com/cachetronaut/deltav.
|
|
6
|
+
|
|
7
|
+
## Install
|
|
8
|
+
|
|
9
|
+
```sh
|
|
10
|
+
pip install deltav
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
## Import
|
|
14
|
+
|
|
15
|
+
```python
|
|
16
|
+
import deltav
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
## Development
|
|
20
|
+
|
|
21
|
+
Run from `py/`:
|
|
22
|
+
|
|
23
|
+
```sh
|
|
24
|
+
uv sync --dev
|
|
25
|
+
uv run --with ruff ruff check .
|
|
26
|
+
uv run --with ruff ruff format --check .
|
|
27
|
+
uv run --with ty ty check
|
|
28
|
+
uv run --with pytest --with pytest-asyncio python -m pytest
|
|
29
|
+
```
|
|
30
|
+
|
|
31
|
+
## License
|
|
32
|
+
|
|
33
|
+
MIT
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "deltav"
|
|
3
|
+
version = "0.0.0"
|
|
4
|
+
description = "DeltaV reservation-based budget enforcement for agent and tool runs."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.11"
|
|
7
|
+
license = { text = "MIT" }
|
|
8
|
+
dependencies = ["dockbay>=0.0.0", "psycopg[binary]>=3.2", "psycopg-pool>=3.2"]
|
|
9
|
+
|
|
10
|
+
[project.urls]
|
|
11
|
+
Homepage = "https://github.com/cachetronaut/deltav"
|
|
12
|
+
Repository = "https://github.com/cachetronaut/deltav"
|
|
13
|
+
Issues = "https://github.com/cachetronaut/deltav/issues"
|
|
14
|
+
|
|
15
|
+
[build-system]
|
|
16
|
+
requires = ["hatchling"]
|
|
17
|
+
build-backend = "hatchling.build"
|
|
18
|
+
|
|
19
|
+
[tool.hatch.build.targets.wheel]
|
|
20
|
+
packages = ["src/deltav"]
|
|
21
|
+
|
|
22
|
+
[dependency-groups]
|
|
23
|
+
dev = ["pytest>=8", "pytest-asyncio>=0.23", "ruff>=0.6", "ty>=0.0.1a8"]
|
|
24
|
+
|
|
25
|
+
[tool.ruff]
|
|
26
|
+
line-length = 100
|
|
27
|
+
target-version = "py311"
|
|
28
|
+
|
|
29
|
+
[tool.ruff.lint]
|
|
30
|
+
select = ["E", "F", "I", "UP", "B", "SIM"]
|
|
31
|
+
|
|
32
|
+
[tool.pytest.ini_options]
|
|
33
|
+
pythonpath = ["src"]
|
|
34
|
+
|
|
35
|
+
[tool.ty.environment]
|
|
36
|
+
extra-paths = ["src"]
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
from .core import (
|
|
2
|
+
Breach,
|
|
3
|
+
Budget,
|
|
4
|
+
BudgetStore,
|
|
5
|
+
Decision,
|
|
6
|
+
Reservation,
|
|
7
|
+
Spend,
|
|
8
|
+
SpendRequest,
|
|
9
|
+
Usage,
|
|
10
|
+
UsageMap,
|
|
11
|
+
canonicalize,
|
|
12
|
+
check,
|
|
13
|
+
check_stack,
|
|
14
|
+
empty_usage,
|
|
15
|
+
release,
|
|
16
|
+
release_stack,
|
|
17
|
+
reserve,
|
|
18
|
+
reserve_stack,
|
|
19
|
+
settle,
|
|
20
|
+
settle_stack,
|
|
21
|
+
)
|
|
22
|
+
from .store_convex import ConvexBudgetStore, create_budget_operations
|
|
23
|
+
from .store_local import InMemoryBudgetStore
|
|
24
|
+
from .store_postgres import DriverBudgetStore, PostgresBudgetStore
|
|
25
|
+
from .store_remote import RemoteBudgetStore
|
|
26
|
+
|
|
27
|
+
__all__ = [
|
|
28
|
+
"Breach",
|
|
29
|
+
"Budget",
|
|
30
|
+
"BudgetStore",
|
|
31
|
+
"ConvexBudgetStore",
|
|
32
|
+
"Decision",
|
|
33
|
+
"DriverBudgetStore",
|
|
34
|
+
"InMemoryBudgetStore",
|
|
35
|
+
"PostgresBudgetStore",
|
|
36
|
+
"RemoteBudgetStore",
|
|
37
|
+
"Reservation",
|
|
38
|
+
"Spend",
|
|
39
|
+
"SpendRequest",
|
|
40
|
+
"Usage",
|
|
41
|
+
"UsageMap",
|
|
42
|
+
"canonicalize",
|
|
43
|
+
"check",
|
|
44
|
+
"check_stack",
|
|
45
|
+
"create_budget_operations",
|
|
46
|
+
"empty_usage",
|
|
47
|
+
"release",
|
|
48
|
+
"release_stack",
|
|
49
|
+
"reserve",
|
|
50
|
+
"reserve_stack",
|
|
51
|
+
"settle",
|
|
52
|
+
"settle_stack",
|
|
53
|
+
]
|
|
@@ -0,0 +1,178 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
from dataclasses import dataclass, field
|
|
5
|
+
from typing import Protocol
|
|
6
|
+
|
|
7
|
+
Spend = dict[str, float]
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
@dataclass(frozen=True)
|
|
11
|
+
class Budget:
|
|
12
|
+
id: str
|
|
13
|
+
layer: str
|
|
14
|
+
limits: dict[str, float]
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
@dataclass(frozen=True)
|
|
18
|
+
class Usage:
|
|
19
|
+
cumulative: Spend = field(default_factory=dict)
|
|
20
|
+
concurrent: Spend = field(default_factory=dict)
|
|
21
|
+
peak: Spend = field(default_factory=dict)
|
|
22
|
+
reserved: Spend = field(default_factory=dict)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
@dataclass(frozen=True)
|
|
26
|
+
class SpendRequest:
|
|
27
|
+
estimate: Spend
|
|
28
|
+
run_id: str | None = None
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
@dataclass(frozen=True)
|
|
32
|
+
class Reservation:
|
|
33
|
+
id: str
|
|
34
|
+
estimate: Spend
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
@dataclass(frozen=True)
|
|
38
|
+
class Breach:
|
|
39
|
+
budget_id: str
|
|
40
|
+
layer: str
|
|
41
|
+
dimension: str
|
|
42
|
+
limit: float
|
|
43
|
+
would_be: float
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
@dataclass(frozen=True)
|
|
47
|
+
class Decision:
|
|
48
|
+
ok: bool
|
|
49
|
+
reservation: Reservation | None = None
|
|
50
|
+
breach: Breach | None = None
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
UsageMap = dict[str, Usage]
|
|
54
|
+
_next_reservation = 1
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
class BudgetStore(Protocol):
|
|
58
|
+
async def load(self, stack_id: str) -> tuple[list[Budget], UsageMap]: ...
|
|
59
|
+
async def try_reserve(self, stack_id: str, req: SpendRequest) -> Decision: ...
|
|
60
|
+
async def settle(self, stack_id: str, reservation: Reservation, actual: Spend) -> None: ...
|
|
61
|
+
async def release(self, stack_id: str, reservation: Reservation) -> None: ...
|
|
62
|
+
async def snapshot(self, stack_id: str) -> UsageMap: ...
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def empty_usage() -> Usage:
|
|
66
|
+
return Usage()
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def check(budget: Budget, usage: Usage, req: SpendRequest) -> Decision:
|
|
70
|
+
for dimension, amount in req.estimate.items():
|
|
71
|
+
limit = budget.limits.get(dimension)
|
|
72
|
+
if limit is None:
|
|
73
|
+
continue
|
|
74
|
+
if amount < 0:
|
|
75
|
+
return Decision(
|
|
76
|
+
ok=False,
|
|
77
|
+
breach=Breach(budget.id, budget.layer, dimension, 0, amount),
|
|
78
|
+
)
|
|
79
|
+
would_be = usage.cumulative.get(dimension, 0) + usage.reserved.get(dimension, 0) + amount
|
|
80
|
+
if would_be > limit:
|
|
81
|
+
return Decision(
|
|
82
|
+
ok=False,
|
|
83
|
+
breach=Breach(budget.id, budget.layer, dimension, limit, would_be),
|
|
84
|
+
)
|
|
85
|
+
return Decision(ok=True, reservation=_reservation(req.estimate))
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
def check_stack(stack: list[Budget], usages: UsageMap, req: SpendRequest) -> Decision:
|
|
89
|
+
reservation: Reservation | None = None
|
|
90
|
+
for budget in stack:
|
|
91
|
+
decision = check(budget, usages.get(budget.id, empty_usage()), req)
|
|
92
|
+
if not decision.ok:
|
|
93
|
+
return decision
|
|
94
|
+
reservation = decision.reservation
|
|
95
|
+
return Decision(ok=True, reservation=reservation or _reservation(req.estimate))
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def reserve(usage: Usage, req: SpendRequest) -> Usage:
|
|
99
|
+
return Usage(
|
|
100
|
+
cumulative=dict(usage.cumulative),
|
|
101
|
+
concurrent=dict(usage.concurrent),
|
|
102
|
+
peak=dict(usage.peak),
|
|
103
|
+
reserved=_add_delta(usage.reserved, req.estimate),
|
|
104
|
+
)
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def settle(usage: Usage, reservation: Reservation, actual: Spend) -> Usage:
|
|
108
|
+
cumulative = dict(usage.cumulative)
|
|
109
|
+
for dimension, amount in actual.items():
|
|
110
|
+
cumulative[dimension] = cumulative.get(dimension, 0) + amount
|
|
111
|
+
return Usage(
|
|
112
|
+
cumulative=cumulative,
|
|
113
|
+
concurrent=dict(usage.concurrent),
|
|
114
|
+
peak=dict(usage.peak),
|
|
115
|
+
reserved=_add_delta(usage.reserved, {k: -v for k, v in reservation.estimate.items()}),
|
|
116
|
+
)
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
def release(usage: Usage, reservation: Reservation) -> Usage:
|
|
120
|
+
return Usage(
|
|
121
|
+
cumulative=dict(usage.cumulative),
|
|
122
|
+
concurrent=dict(usage.concurrent),
|
|
123
|
+
peak=dict(usage.peak),
|
|
124
|
+
reserved=_add_delta(usage.reserved, {k: -v for k, v in reservation.estimate.items()}),
|
|
125
|
+
)
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
def reserve_stack(stack: list[Budget], usages: UsageMap, req: SpendRequest) -> UsageMap:
|
|
129
|
+
next_usages = dict(usages)
|
|
130
|
+
for budget in stack:
|
|
131
|
+
next_usages[budget.id] = reserve(next_usages.get(budget.id, empty_usage()), req)
|
|
132
|
+
return next_usages
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
def settle_stack(
|
|
136
|
+
stack: list[Budget], usages: UsageMap, reservation: Reservation, actual: Spend
|
|
137
|
+
) -> UsageMap:
|
|
138
|
+
next_usages = dict(usages)
|
|
139
|
+
for budget in stack:
|
|
140
|
+
next_usages[budget.id] = settle(
|
|
141
|
+
next_usages.get(budget.id, empty_usage()), reservation, actual
|
|
142
|
+
)
|
|
143
|
+
return next_usages
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
def release_stack(stack: list[Budget], usages: UsageMap, reservation: Reservation) -> UsageMap:
|
|
147
|
+
next_usages = dict(usages)
|
|
148
|
+
for budget in stack:
|
|
149
|
+
next_usages[budget.id] = release(next_usages.get(budget.id, empty_usage()), reservation)
|
|
150
|
+
return next_usages
|
|
151
|
+
|
|
152
|
+
|
|
153
|
+
def canonicalize(value: object) -> str:
|
|
154
|
+
return json.dumps(value, sort_keys=True, separators=(",", ":"), default=_json_default)
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
def _reservation(estimate: Spend) -> Reservation:
|
|
158
|
+
global _next_reservation
|
|
159
|
+
reservation = Reservation(id=f"reservation_{_next_reservation}", estimate=dict(estimate))
|
|
160
|
+
_next_reservation += 1
|
|
161
|
+
return reservation
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
def _add_delta(values: Spend, delta: Spend) -> Spend:
|
|
165
|
+
out = dict(values)
|
|
166
|
+
for dimension, amount in delta.items():
|
|
167
|
+
value = out.get(dimension, 0) + amount
|
|
168
|
+
if value == 0:
|
|
169
|
+
out.pop(dimension, None)
|
|
170
|
+
else:
|
|
171
|
+
out[dimension] = value
|
|
172
|
+
return out
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
def _json_default(value: object) -> object:
|
|
176
|
+
if hasattr(value, "__dict__"):
|
|
177
|
+
return value.__dict__
|
|
178
|
+
raise TypeError(f"Cannot serialize {type(value)!r}")
|
|
@@ -0,0 +1,276 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from copy import deepcopy
|
|
4
|
+
from typing import Any
|
|
5
|
+
|
|
6
|
+
from dockbay import (
|
|
7
|
+
ConvexOperationContext,
|
|
8
|
+
ConvexOperationDriver,
|
|
9
|
+
ConvexStoreOperation,
|
|
10
|
+
JsonValue,
|
|
11
|
+
)
|
|
12
|
+
|
|
13
|
+
from .core import (
|
|
14
|
+
Budget,
|
|
15
|
+
Decision,
|
|
16
|
+
Reservation,
|
|
17
|
+
Spend,
|
|
18
|
+
SpendRequest,
|
|
19
|
+
Usage,
|
|
20
|
+
UsageMap,
|
|
21
|
+
check_stack,
|
|
22
|
+
release_stack,
|
|
23
|
+
reserve_stack,
|
|
24
|
+
settle_stack,
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
LOAD = "budget.load"
|
|
28
|
+
TRY_RESERVE = "budget.tryReserve"
|
|
29
|
+
SETTLE = "budget.settle"
|
|
30
|
+
RELEASE = "budget.release"
|
|
31
|
+
SNAPSHOT = "budget.snapshot"
|
|
32
|
+
|
|
33
|
+
InitialState = dict[str, tuple[list[Budget], UsageMap]]
|
|
34
|
+
_StackState = dict[str, Any]
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
class ConvexBudgetStore:
|
|
38
|
+
def __init__(
|
|
39
|
+
self,
|
|
40
|
+
driver: ConvexOperationDriver,
|
|
41
|
+
*,
|
|
42
|
+
operations: dict[str, str] | None = None,
|
|
43
|
+
) -> None:
|
|
44
|
+
self._driver = driver
|
|
45
|
+
self._operations = {
|
|
46
|
+
"load": LOAD,
|
|
47
|
+
"try_reserve": TRY_RESERVE,
|
|
48
|
+
"settle": SETTLE,
|
|
49
|
+
"release": RELEASE,
|
|
50
|
+
"snapshot": SNAPSHOT,
|
|
51
|
+
**(operations or {}),
|
|
52
|
+
}
|
|
53
|
+
|
|
54
|
+
async def load(self, stack_id: str) -> tuple[list[Budget], UsageMap]:
|
|
55
|
+
result = await self._driver.call(self._operations["load"], {"stackId": stack_id})
|
|
56
|
+
assert isinstance(result, dict)
|
|
57
|
+
return _row_to_stack(result), _row_to_usage_map(result["usages"])
|
|
58
|
+
|
|
59
|
+
async def try_reserve(self, stack_id: str, req: SpendRequest) -> Decision:
|
|
60
|
+
result = await self._driver.call(
|
|
61
|
+
self._operations["try_reserve"],
|
|
62
|
+
{"stackId": stack_id, "req": _spend_request_to_row(req)},
|
|
63
|
+
)
|
|
64
|
+
assert isinstance(result, dict)
|
|
65
|
+
return _row_to_decision(result)
|
|
66
|
+
|
|
67
|
+
async def settle(self, stack_id: str, reservation: Reservation, actual: Spend) -> None:
|
|
68
|
+
await self._driver.call(
|
|
69
|
+
self._operations["settle"],
|
|
70
|
+
{
|
|
71
|
+
"stackId": stack_id,
|
|
72
|
+
"reservation": _reservation_to_row(reservation),
|
|
73
|
+
"actual": actual,
|
|
74
|
+
},
|
|
75
|
+
)
|
|
76
|
+
|
|
77
|
+
async def release(self, stack_id: str, reservation: Reservation) -> None:
|
|
78
|
+
await self._driver.call(
|
|
79
|
+
self._operations["release"],
|
|
80
|
+
{"stackId": stack_id, "reservation": _reservation_to_row(reservation)},
|
|
81
|
+
)
|
|
82
|
+
|
|
83
|
+
async def snapshot(self, stack_id: str) -> UsageMap:
|
|
84
|
+
result = await self._driver.call(self._operations["snapshot"], {"stackId": stack_id})
|
|
85
|
+
return _row_to_usage_map(result)
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
def create_budget_operations(initial: InitialState) -> list[ConvexStoreOperation]:
|
|
89
|
+
state: dict[str, _StackState] = {
|
|
90
|
+
stack_id: {
|
|
91
|
+
"stack": [_budget_to_row(budget) for budget in stack],
|
|
92
|
+
"usages": _usage_map_to_row(usages),
|
|
93
|
+
"openReservations": [],
|
|
94
|
+
}
|
|
95
|
+
for stack_id, (stack, usages) in initial.items()
|
|
96
|
+
}
|
|
97
|
+
|
|
98
|
+
async def load(_ctx: ConvexOperationContext, input_value: JsonValue) -> JsonValue:
|
|
99
|
+
stack_id = _stack_id_from(input_value)
|
|
100
|
+
current = _require_stack(state, stack_id)
|
|
101
|
+
return {
|
|
102
|
+
"stack": deepcopy(current["stack"]),
|
|
103
|
+
"usages": deepcopy(current["usages"]),
|
|
104
|
+
}
|
|
105
|
+
|
|
106
|
+
async def try_reserve(_ctx: ConvexOperationContext, input_value: JsonValue) -> JsonValue:
|
|
107
|
+
assert isinstance(input_value, dict)
|
|
108
|
+
stack_id = _stack_id_from(input_value)
|
|
109
|
+
current = _require_stack(state, stack_id)
|
|
110
|
+
stack = _row_to_stack(current)
|
|
111
|
+
usages = _row_to_usage_map(current["usages"])
|
|
112
|
+
decision = check_stack(stack, usages, _row_to_spend_request(input_value["req"]))
|
|
113
|
+
if not decision.ok or decision.reservation is None:
|
|
114
|
+
return _decision_to_row(decision)
|
|
115
|
+
state[stack_id] = {
|
|
116
|
+
**current,
|
|
117
|
+
"usages": _usage_map_to_row(
|
|
118
|
+
reserve_stack(stack, usages, _row_to_spend_request(input_value["req"]))
|
|
119
|
+
),
|
|
120
|
+
"openReservations": [*current["openReservations"], decision.reservation.id],
|
|
121
|
+
}
|
|
122
|
+
return _decision_to_row(decision)
|
|
123
|
+
|
|
124
|
+
async def settle(_ctx: ConvexOperationContext, input_value: JsonValue) -> JsonValue:
|
|
125
|
+
assert isinstance(input_value, dict)
|
|
126
|
+
stack_id = _stack_id_from(input_value)
|
|
127
|
+
current = _require_stack(state, stack_id)
|
|
128
|
+
reservation = _row_to_reservation(input_value["reservation"])
|
|
129
|
+
if reservation.id not in current["openReservations"]:
|
|
130
|
+
return None
|
|
131
|
+
stack = _row_to_stack(current)
|
|
132
|
+
usages = _row_to_usage_map(current["usages"])
|
|
133
|
+
state[stack_id] = {
|
|
134
|
+
**current,
|
|
135
|
+
"usages": _usage_map_to_row(
|
|
136
|
+
settle_stack(stack, usages, reservation, input_value["actual"])
|
|
137
|
+
),
|
|
138
|
+
"openReservations": [
|
|
139
|
+
item for item in current["openReservations"] if item != reservation.id
|
|
140
|
+
],
|
|
141
|
+
}
|
|
142
|
+
return None
|
|
143
|
+
|
|
144
|
+
async def release(_ctx: ConvexOperationContext, input_value: JsonValue) -> JsonValue:
|
|
145
|
+
assert isinstance(input_value, dict)
|
|
146
|
+
stack_id = _stack_id_from(input_value)
|
|
147
|
+
current = _require_stack(state, stack_id)
|
|
148
|
+
reservation = _row_to_reservation(input_value["reservation"])
|
|
149
|
+
if reservation.id not in current["openReservations"]:
|
|
150
|
+
return None
|
|
151
|
+
stack = _row_to_stack(current)
|
|
152
|
+
usages = _row_to_usage_map(current["usages"])
|
|
153
|
+
state[stack_id] = {
|
|
154
|
+
**current,
|
|
155
|
+
"usages": _usage_map_to_row(release_stack(stack, usages, reservation)),
|
|
156
|
+
"openReservations": [
|
|
157
|
+
item for item in current["openReservations"] if item != reservation.id
|
|
158
|
+
],
|
|
159
|
+
}
|
|
160
|
+
return None
|
|
161
|
+
|
|
162
|
+
async def snapshot(_ctx: ConvexOperationContext, input_value: JsonValue) -> JsonValue:
|
|
163
|
+
stack_id = _stack_id_from(input_value)
|
|
164
|
+
return deepcopy(_require_stack(state, stack_id)["usages"])
|
|
165
|
+
|
|
166
|
+
return [
|
|
167
|
+
ConvexStoreOperation(name=LOAD, kind="query", run=load),
|
|
168
|
+
ConvexStoreOperation(name=TRY_RESERVE, kind="mutation", run=try_reserve),
|
|
169
|
+
ConvexStoreOperation(name=SETTLE, kind="mutation", run=settle),
|
|
170
|
+
ConvexStoreOperation(name=RELEASE, kind="mutation", run=release),
|
|
171
|
+
ConvexStoreOperation(name=SNAPSHOT, kind="query", run=snapshot),
|
|
172
|
+
]
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
def _require_stack(state: dict[str, _StackState], stack_id: str) -> _StackState:
|
|
176
|
+
try:
|
|
177
|
+
return state[stack_id]
|
|
178
|
+
except KeyError as exc:
|
|
179
|
+
raise KeyError(f"Unknown budget stack: {stack_id}") from exc
|
|
180
|
+
|
|
181
|
+
|
|
182
|
+
def _stack_id_from(input_value: JsonValue) -> str:
|
|
183
|
+
assert isinstance(input_value, dict)
|
|
184
|
+
stack_id = input_value.get("stackId")
|
|
185
|
+
if not isinstance(stack_id, str):
|
|
186
|
+
raise ValueError("Convex budget operation requires stackId")
|
|
187
|
+
return stack_id
|
|
188
|
+
|
|
189
|
+
|
|
190
|
+
def _budget_to_row(budget: Budget) -> dict[str, Any]:
|
|
191
|
+
return {"id": budget.id, "layer": budget.layer, "limits": dict(budget.limits)}
|
|
192
|
+
|
|
193
|
+
|
|
194
|
+
def _row_to_budget(row: Any) -> Budget:
|
|
195
|
+
return Budget(id=row["id"], layer=row["layer"], limits=dict(row["limits"]))
|
|
196
|
+
|
|
197
|
+
|
|
198
|
+
def _usage_to_row(usage: Usage) -> dict[str, Any]:
|
|
199
|
+
return {
|
|
200
|
+
"cumulative": dict(usage.cumulative),
|
|
201
|
+
"concurrent": dict(usage.concurrent),
|
|
202
|
+
"peak": dict(usage.peak),
|
|
203
|
+
"reserved": dict(usage.reserved),
|
|
204
|
+
}
|
|
205
|
+
|
|
206
|
+
|
|
207
|
+
def _row_to_usage(row: Any) -> Usage:
|
|
208
|
+
return Usage(
|
|
209
|
+
cumulative=dict(row["cumulative"]),
|
|
210
|
+
concurrent=dict(row["concurrent"]),
|
|
211
|
+
peak=dict(row["peak"]),
|
|
212
|
+
reserved=dict(row["reserved"]),
|
|
213
|
+
)
|
|
214
|
+
|
|
215
|
+
|
|
216
|
+
def _usage_map_to_row(usages: UsageMap) -> dict[str, Any]:
|
|
217
|
+
return {budget_id: _usage_to_row(usage) for budget_id, usage in usages.items()}
|
|
218
|
+
|
|
219
|
+
|
|
220
|
+
def _row_to_usage_map(row: Any) -> UsageMap:
|
|
221
|
+
return {budget_id: _row_to_usage(usage) for budget_id, usage in deepcopy(row).items()}
|
|
222
|
+
|
|
223
|
+
|
|
224
|
+
def _row_to_stack(row: Any) -> list[Budget]:
|
|
225
|
+
return [_row_to_budget(budget) for budget in deepcopy(row["stack"])]
|
|
226
|
+
|
|
227
|
+
|
|
228
|
+
def _spend_request_to_row(req: SpendRequest) -> dict[str, Any]:
|
|
229
|
+
return {"estimate": dict(req.estimate), "run_id": req.run_id}
|
|
230
|
+
|
|
231
|
+
|
|
232
|
+
def _row_to_spend_request(row: Any) -> SpendRequest:
|
|
233
|
+
return SpendRequest(estimate=dict(row["estimate"]), run_id=row.get("run_id"))
|
|
234
|
+
|
|
235
|
+
|
|
236
|
+
def _reservation_to_row(reservation: Reservation) -> dict[str, Any]:
|
|
237
|
+
return {"id": reservation.id, "estimate": dict(reservation.estimate)}
|
|
238
|
+
|
|
239
|
+
|
|
240
|
+
def _row_to_reservation(row: Any) -> Reservation:
|
|
241
|
+
return Reservation(id=row["id"], estimate=dict(row["estimate"]))
|
|
242
|
+
|
|
243
|
+
|
|
244
|
+
def _decision_to_row(decision: Decision) -> dict[str, Any]:
|
|
245
|
+
if decision.ok:
|
|
246
|
+
assert decision.reservation is not None
|
|
247
|
+
return {"ok": True, "reservation": _reservation_to_row(decision.reservation)}
|
|
248
|
+
assert decision.breach is not None
|
|
249
|
+
return {
|
|
250
|
+
"ok": False,
|
|
251
|
+
"breach": {
|
|
252
|
+
"budget_id": decision.breach.budget_id,
|
|
253
|
+
"layer": decision.breach.layer,
|
|
254
|
+
"dimension": decision.breach.dimension,
|
|
255
|
+
"limit": decision.breach.limit,
|
|
256
|
+
"would_be": decision.breach.would_be,
|
|
257
|
+
},
|
|
258
|
+
}
|
|
259
|
+
|
|
260
|
+
|
|
261
|
+
def _row_to_decision(row: Any) -> Decision:
|
|
262
|
+
if row["ok"]:
|
|
263
|
+
return Decision(ok=True, reservation=_row_to_reservation(row["reservation"]))
|
|
264
|
+
breach = row["breach"]
|
|
265
|
+
from .core import Breach
|
|
266
|
+
|
|
267
|
+
return Decision(
|
|
268
|
+
ok=False,
|
|
269
|
+
breach=Breach(
|
|
270
|
+
budget_id=breach["budget_id"],
|
|
271
|
+
layer=breach["layer"],
|
|
272
|
+
dimension=breach["dimension"],
|
|
273
|
+
limit=breach["limit"],
|
|
274
|
+
would_be=breach["would_be"],
|
|
275
|
+
),
|
|
276
|
+
)
|
|
@@ -0,0 +1,63 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
from copy import deepcopy
|
|
5
|
+
|
|
6
|
+
from .core import (
|
|
7
|
+
Budget,
|
|
8
|
+
Decision,
|
|
9
|
+
Reservation,
|
|
10
|
+
Spend,
|
|
11
|
+
SpendRequest,
|
|
12
|
+
UsageMap,
|
|
13
|
+
check_stack,
|
|
14
|
+
release_stack,
|
|
15
|
+
reserve_stack,
|
|
16
|
+
settle_stack,
|
|
17
|
+
)
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class InMemoryBudgetStore:
|
|
21
|
+
def __init__(self, initial: dict[str, tuple[list[Budget], UsageMap]] | None = None) -> None:
|
|
22
|
+
self._stacks: dict[str, tuple[list[Budget], UsageMap]] = initial or {}
|
|
23
|
+
self._open: dict[str, set[str]] = {stack_id: set() for stack_id in self._stacks}
|
|
24
|
+
self._lock = asyncio.Lock()
|
|
25
|
+
|
|
26
|
+
async def load(self, stack_id: str) -> tuple[list[Budget], UsageMap]:
|
|
27
|
+
async with self._lock:
|
|
28
|
+
stack, usages = self._require(stack_id)
|
|
29
|
+
return list(stack), deepcopy(usages)
|
|
30
|
+
|
|
31
|
+
async def try_reserve(self, stack_id: str, req: SpendRequest) -> Decision:
|
|
32
|
+
async with self._lock:
|
|
33
|
+
stack, usages = self._require(stack_id)
|
|
34
|
+
decision = check_stack(stack, usages, req)
|
|
35
|
+
if decision.ok and decision.reservation is not None:
|
|
36
|
+
self._stacks[stack_id] = (stack, reserve_stack(stack, usages, req))
|
|
37
|
+
self._open.setdefault(stack_id, set()).add(decision.reservation.id)
|
|
38
|
+
return decision
|
|
39
|
+
|
|
40
|
+
async def settle(self, stack_id: str, reservation: Reservation, actual: Spend) -> None:
|
|
41
|
+
async with self._lock:
|
|
42
|
+
if reservation.id not in self._open.setdefault(stack_id, set()):
|
|
43
|
+
return
|
|
44
|
+
stack, usages = self._require(stack_id)
|
|
45
|
+
self._stacks[stack_id] = (stack, settle_stack(stack, usages, reservation, actual))
|
|
46
|
+
self._open[stack_id].remove(reservation.id)
|
|
47
|
+
|
|
48
|
+
async def release(self, stack_id: str, reservation: Reservation) -> None:
|
|
49
|
+
async with self._lock:
|
|
50
|
+
if reservation.id not in self._open.setdefault(stack_id, set()):
|
|
51
|
+
return
|
|
52
|
+
stack, usages = self._require(stack_id)
|
|
53
|
+
self._stacks[stack_id] = (stack, release_stack(stack, usages, reservation))
|
|
54
|
+
self._open[stack_id].remove(reservation.id)
|
|
55
|
+
|
|
56
|
+
async def snapshot(self, stack_id: str) -> UsageMap:
|
|
57
|
+
async with self._lock:
|
|
58
|
+
return deepcopy(self._require(stack_id)[1])
|
|
59
|
+
|
|
60
|
+
def _require(self, stack_id: str) -> tuple[list[Budget], UsageMap]:
|
|
61
|
+
if stack_id not in self._stacks:
|
|
62
|
+
raise KeyError(f"Unknown budget stack: {stack_id}")
|
|
63
|
+
return self._stacks[stack_id]
|