Viewing files at commit 2735d7f5558eCommit detailsView latest

scripts/oakbench/remotes.py 332 lines · 11.5 KB · python
1
"""Shared remote plumbing for every lane that touches a remote.
2
3
One module owns: which disposable remote a lane talks to (env resolution with
4
structured skip reasons), local-bare git remotes (the core-equivalent track),
5
GitHub remotes (the agent-default track), collision-safe disposable branch
6
names, and the transport-honesty fields every remote-touching row must carry.
7
8
Transport honesty (AGENTS.md, ADR-0005): a network row and a local-file row
9
are different measurements. Every row built from this module records
10
`remote_transport` and `remote_server`, and reports must never aggregate or
11
delta across transports. This module makes the asymmetry explicit; it cannot
12
make it disappear.
13
14
Skips are work items, not absences (ADR-0002): an unresolved remote yields a
15
structured reason naming the exact env var to set, for the lane to emit as a
16
returncode-77 skip row.
17
"""
18
19
from __future__ import annotations
20
21
import itertools
22
import os
23
import secrets
24
import subprocess
25
import sys
26
from collections.abc import Mapping
27
from dataclasses import dataclass
28
from pathlib import Path
29
30
from .environment import base_env
31
32
# Purpose-scoped env vars keep lanes from polluting each other's disposable
33
# repos (fleet branch storms never land in the sync lane's history). Most
34
# purposes fall back to the generic OAK_BENCH_REMOTE; mirror intentionally
35
# does not because mount-vs-clone needs a large mirror, not a tiny disposable
36
# repo used by other lanes.
37
OAK_REMOTE_ENV = "OAK_BENCH_REMOTE"
38
_BRANCH_COUNTER = itertools.count()
39
OAK_PURPOSE_ENVS: dict[str, str] = {
40
    "fleet": "OAK_BENCH_FLEET_REPO",
41
    "sync": "OAK_BENCH_SYNC_REPO",
42
    "mirror": "OAK_BENCH_MIRROR_REPO",
43
}
44
OAK_PURPOSES_WITHOUT_GENERIC_FALLBACK = {"mirror"}
45
GIT_GITHUB_ENV = "GIT_BENCH_REMOTE"
46
GIT_MIRROR_ENV = "GIT_BENCH_MIRROR_URL"
47
48
TRANSPORT_NETWORK = "network"
49
TRANSPORT_LOCAL_FILE = "local_file"
50
# netshape lane (Linux netns + veth + tc/netem): both subjects served from
51
# local disk through the same shaped pipe. A shaped row is neither a real
52
# network row nor a local-file row — it is its own transport, and it is only
53
# comparable to another shaped row through the SAME pipe (same profile, same
54
# server identity). shaped_comparison_legal() is the one place that rule
55
# lives; existing transport rules are unchanged.
56
TRANSPORT_NETWORK_SHAPED = "network_shaped"
57
58
59
@dataclass(frozen=True)
60
class RemoteResolution:
61
    """Where a lane's remote lives, or exactly why it doesn't."""
62
63
    repo: str | None  # ORG/REPO, URL, or filesystem path
64
    transport: str | None  # TRANSPORT_NETWORK | TRANSPORT_LOCAL_FILE
65
    server: str | None  # host identity recorded on rows (e.g. "oak.space")
66
    env_var: str  # the env var that resolved (or would resolve) it
67
    skip_reason: str | None = None
68
69
    @property
70
    def resolved(self) -> bool:
71
        return self.repo is not None
72
73
    def row_fields(self) -> dict[str, str]:
74
        """Transport-honesty fields for every row touching this remote."""
75
        if not self.resolved:
76
            raise ValueError(f"unresolved remote ({self.skip_reason}); emit a skip row instead")
77
        return {
78
            "remote_repo": self.repo,
79
            "remote_transport": self.transport,
80
            "remote_server": self.server,
81
        }
82
83
84
@dataclass(frozen=True)
85
class RemoteRequirement:
86
    """One lane's need for an Oak remote before the run starts."""
87
88
    purpose: str
89
    lane: str
90
    operations: tuple[str, ...]
91
92
93
@dataclass(frozen=True)
94
class RemotePreflightWarning:
95
    """A remote skip that can be predicted before any fixture work starts."""
96
97
    purpose: str
98
    lane: str
99
    env_var: str
100
    operations: tuple[str, ...]
101
    skip_reason: str
102
103
104
def resolve_oak_remote(
105
    purpose: str = "default",
106
    env: Mapping[str, str] | None = None,
107
) -> RemoteResolution:
108
    """Resolve the disposable oak ORG/REPO for a lane purpose.
109
110
    Checks the purpose-scoped env var first, then OAK_BENCH_REMOTE, except
111
    for mirror which must be explicitly configured. The server identity is the
112
    configured oak host (oak.space unless OAK_SERVER overrides it for a test
113
    rig).
114
    """
115
    source_env = os.environ if env is None else env
116
    candidates = []
117
    scoped = OAK_PURPOSE_ENVS.get(purpose)
118
    if scoped:
119
        candidates.append(scoped)
120
    if purpose not in OAK_PURPOSES_WITHOUT_GENERIC_FALLBACK:
121
        candidates.append(OAK_REMOTE_ENV)
122
    for env_var in candidates:
123
        repo = source_env.get(env_var, "").strip()
124
        if repo:
125
            return RemoteResolution(
126
                repo=repo,
127
                transport=TRANSPORT_NETWORK,
128
                server=source_env.get("OAK_SERVER", "oak.space"),
129
                env_var=env_var,
130
            )
131
    primary = candidates[0]
132
    return RemoteResolution(
133
        repo=None,
134
        transport=None,
135
        server=None,
136
        env_var=primary,
137
        skip_reason=(
138
            f"oak_remote_not_configured: set {primary} to a DISPOSABLE ORG/REPO "
139
            f"(e.g. oak/bench-{purpose}-tmp) to measure real-server operations"
140
        ),
141
    )
142
143
144
def oak_remote_preflight_warnings(
145
    requirements: list[RemoteRequirement],
146
    env: Mapping[str, str] | None = None,
147
) -> list[RemotePreflightWarning]:
148
    """Resolve remote requirements without mutating state.
149
150
    Skip rows are still emitted by the lanes themselves. This is an early,
151
    human-visible forecast so a long benchmark run does not wait minutes before
152
    revealing that an Oak remote env var was missing.
153
    """
154
    grouped: dict[tuple[str, str, str, str], list[str]] = {}
155
    for requirement in requirements:
156
        resolution = resolve_oak_remote(requirement.purpose, env=env)
157
        if resolution.resolved:
158
            continue
159
        key = (
160
            requirement.purpose,
161
            requirement.lane,
162
            resolution.env_var,
163
            resolution.skip_reason or "",
164
        )
165
        grouped.setdefault(key, []).extend(requirement.operations)
166
    return [
167
        RemotePreflightWarning(
168
            purpose=purpose,
169
            lane=lane,
170
            env_var=env_var,
171
            operations=tuple(dict.fromkeys(operations)),
172
            skip_reason=skip_reason,
173
        )
174
        for (purpose, lane, env_var, skip_reason), operations in sorted(grouped.items())
175
    ]
176
177
178
def format_remote_preflight_warnings(warnings: list[RemotePreflightWarning]) -> list[str]:
179
    if not warnings:
180
        return []
181
    lines = [
182
        "[preflight] Oak remote configuration missing; these lanes will emit skip rows:",
183
    ]
184
    for warning in warnings:
185
        operations = ", ".join(warning.operations[:6])
186
        if len(warning.operations) > 6:
187
            operations += f", +{len(warning.operations) - 6} more"
188
        lines.append(
189
            f"[preflight] - {warning.lane}/{warning.purpose}: set {warning.env_var}; "
190
            f"affected: {operations}; skip_reason: {warning.skip_reason}"
191
        )
192
    return lines
193
194
195
def print_remote_preflight_warnings(warnings: list[RemotePreflightWarning]) -> None:
196
    for line in format_remote_preflight_warnings(warnings):
197
        print(line, file=sys.stderr, flush=True)
198
199
200
def resolve_git_github_remote(env_var: str = GIT_GITHUB_ENV) -> RemoteResolution:
201
    """Resolve git's agent-default network remote (a GitHub URL or org/repo).
202
203
    This is the transport-comparable counterpart to oak's network remote: git
204
    pays a real `git push`/`git clone` over the wire like oak pays a real
205
    `oak push`. Core-equivalent runs use make_git_bare_remote() instead.
206
    """
207
    url = os.environ.get(env_var, "").strip()
208
    if not url:
209
        return RemoteResolution(
210
            repo=None,
211
            transport=None,
212
            server=None,
213
            env_var=env_var,
214
            skip_reason=(
215
                f"git_network_remote_not_configured: set {env_var} to a DISPOSABLE "
216
                "GitHub repo (URL or owner/repo) for the agent-default track"
217
            ),
218
        )
219
    server = "github.com"
220
    if "://" in url:
221
        server = url.split("://", 1)[1].split("/", 1)[0]
222
    elif "@" in url and ":" in url:
223
        server = url.split("@", 1)[1].split(":", 1)[0]
224
    return RemoteResolution(
225
        repo=url,
226
        transport=TRANSPORT_NETWORK,
227
        server=server,
228
        env_var=env_var,
229
    )
230
231
232
def make_git_bare_remote(git_bin: str | Path, dest_dir: Path, name: str = "remote.git") -> RemoteResolution:
233
    """Create an untimed local bare repo: git's core-equivalent remote.
234
235
    Local-file transport isolates object-transfer/checkout cost from network
236
    jitter. Setup is a precondition, never the measurement — callers must not
237
    time this call.
238
    """
239
    remote_dir = dest_dir / name
240
    subprocess.run(
241
        [str(git_bin), "init", "--bare", str(remote_dir)],
242
        env=base_env(),
243
        stdout=subprocess.DEVNULL,
244
        stderr=subprocess.DEVNULL,
245
        check=False,
246
    )
247
    return RemoteResolution(
248
        repo=str(remote_dir),
249
        transport=TRANSPORT_LOCAL_FILE,
250
        server="local",
251
        env_var="",
252
    )
253
254
255
def attach_git_origin(git_bin: str | Path, repo: Path, remote: RemoteResolution) -> None:
256
    """Untimed: point `origin` at the resolved remote (precondition, not measurement)."""
257
    subprocess.run(
258
        [str(git_bin), "remote", "add", "origin", str(remote.repo)],
259
        cwd=repo,
260
        env=base_env(),
261
        stdout=subprocess.DEVNULL,
262
        stderr=subprocess.DEVNULL,
263
        check=False,
264
    )
265
266
267
def current_git_branch(git_bin: str | Path, repo: Path) -> str:
268
    """Untimed probe of the checked-out branch name (defaults to main)."""
269
    probe = subprocess.run(
270
        [str(git_bin), "rev-parse", "--abbrev-ref", "HEAD"],
271
        cwd=repo,
272
        env=base_env(),
273
        capture_output=True,
274
        text=True,
275
        check=False,
276
    )
277
    return probe.stdout.strip() or "main"
278
279
280
def point_git_bare_head(git_bin: str | Path, remote: RemoteResolution, branch: str) -> None:
281
    """Untimed: point the bare remote's HEAD at a branch so clones check out
282
    a working tree (precondition, not measurement)."""
283
    subprocess.run(
284
        [str(git_bin), "--git-dir", str(remote.repo), "symbolic-ref", "HEAD", f"refs/heads/{branch}"],
285
        env=base_env(),
286
        stdout=subprocess.DEVNULL,
287
        stderr=subprocess.DEVNULL,
288
        check=False,
289
    )
290
291
292
def shaped_comparison_legal(row_a: dict, row_b: dict) -> bool:
293
    """Pure rule: two shaped rows may be compared only same-pipe.
294
295
    Legal iff BOTH rows carry remote_transport == TRANSPORT_NETWORK_SHAPED,
296
    the same non-empty netshape_profile, and the same non-empty remote_server.
297
    Anything else — one side unshaped, profile mismatch, missing fields — is
298
    illegal here; comparisons at other transports are governed by the existing
299
    transport-honesty rules, which this helper does not relax.
300
    """
301
    if row_a.get("remote_transport") != TRANSPORT_NETWORK_SHAPED:
302
        return False
303
    if row_b.get("remote_transport") != TRANSPORT_NETWORK_SHAPED:
304
        return False
305
    profile_a = row_a.get("netshape_profile")
306
    profile_b = row_b.get("netshape_profile")
307
    if not profile_a or profile_a != profile_b:
308
        return False
309
    server_a = row_a.get("remote_server")
310
    server_b = row_b.get("remote_server")
311
    return bool(server_a) and server_a == server_b
312
313
314
def disposable_branch(
315
    bench_id: str,
316
    scenario_name: str,
317
    subject_name: str,
318
    run_index: int,
319
    *,
320
    pid: int | None = None,
321
    token: str | None = None,
322
) -> str:
323
    """Disposable remote branch name: readable prefix plus per-run uniquifier."""
324
    if pid is None:
325
        pid = os.getpid()
326
    if token is None:
327
        token = f"{secrets.token_hex(4)}-{next(_BRANCH_COUNTER):x}"
328
    # Same-second parallel runs must never collide on the server.
329
    return (
330
        f"bench-{bench_id}-{scenario_name}-{subject_name}-r{run_index}"
331
        f"-{pid:x}-{token}"
332
    )