benchloop-client 0.2.23__tar.gz → 0.2.24__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.
Files changed (55) hide show
  1. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/PKG-INFO +10 -1
  2. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/README.md +9 -0
  3. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/pyproject.toml +1 -1
  4. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/_common.py +6 -0
  5. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/metrics.py +14 -3
  6. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/pull.py +8 -3
  7. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/trajectory.py +246 -56
  8. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_cmd_metrics.py +32 -2
  9. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_cmd_pull.py +80 -0
  10. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_cmd_trajectory.py +253 -42
  11. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/uv.lock +2 -2
  12. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/.gitignore +0 -0
  13. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/.python-version +0 -0
  14. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/RELEASE.md +0 -0
  15. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/benchloop-cli-design.md +0 -0
  16. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/bmcli_examples.sh +0 -0
  17. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/__init__.py +0 -0
  18. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/__main__.py +0 -0
  19. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/cli.py +0 -0
  20. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/client.py +0 -0
  21. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/__init__.py +0 -0
  22. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/ap_cmd.py +0 -0
  23. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/artifacts.py +0 -0
  24. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/brm.py +0 -0
  25. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/group_data.py +0 -0
  26. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/jobs.py +0 -0
  27. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/metadata.py +0 -0
  28. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/overview.py +0 -0
  29. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/report.py +0 -0
  30. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/commands/tasks.py +0 -0
  31. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/config.py +0 -0
  32. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/exit_codes.py +0 -0
  33. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/ids.py +0 -0
  34. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/output.py +0 -0
  35. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/runtime.py +0 -0
  36. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/src/aliyun/bmcli/version_check.py +0 -0
  37. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/__init__.py +0 -0
  38. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/conftest.py +0 -0
  39. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_cli_options.py +0 -0
  40. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_client.py +0 -0
  41. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_cmd_ap_cmd.py +0 -0
  42. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_cmd_artifacts.py +0 -0
  43. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_cmd_brm.py +0 -0
  44. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_cmd_metadata.py +0 -0
  45. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_cmd_overview.py +0 -0
  46. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_cmd_report.py +0 -0
  47. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_cmd_tasks_jobs.py +0 -0
  48. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_config.py +0 -0
  49. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_exit_code_behavior.py +0 -0
  50. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_exit_codes.py +0 -0
  51. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_group_data.py +0 -0
  52. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_ids.py +0 -0
  53. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_runtime.py +0 -0
  54. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_transfer_progress.py +0 -0
  55. {benchloop_client-0.2.23 → benchloop_client-0.2.24}/tests/test_version_check.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: benchloop_client
3
- Version: 0.2.23
3
+ Version: 0.2.24
4
4
  Summary: Benchloop CLI — command line tool for benchmark analysis data (metrics, trajectory, reports, artifacts).
5
5
  Requires-Python: >=3.10
6
6
  Requires-Dist: httpx>=0.25.2
@@ -99,6 +99,15 @@ bl brm status --task-id UID
99
99
  bl brm list -g G1
100
100
  ```
101
101
 
102
+ Metrics and Job trajectory requests always explicitly use the Trace source. The
103
+ CLI does not expose a `trace`/`artifacts` source selector.
104
+
105
+ Group trajectory retrieval and downloads use the backend Trace export API. The
106
+ CLI waits for the export URL, downloads the gzip JSON-lines payload, and keeps
107
+ the existing `group_id`/`trajectories` JSON layout. The export wait timeout is
108
+ 30 minutes by default; override it in seconds with
109
+ `--download-timeout SECONDS` on trajectory commands or `pull`.
110
+
102
111
  Group downloads always auto-paginate and include every Job in the Group. In
103
112
  particular, `trajectory`, detailed `metrics`, `report`, and `artifacts` enumerate
104
113
  the complete Job list instead of relying on the backend Group endpoints that
@@ -87,6 +87,15 @@ bl brm status --task-id UID
87
87
  bl brm list -g G1
88
88
  ```
89
89
 
90
+ Metrics and Job trajectory requests always explicitly use the Trace source. The
91
+ CLI does not expose a `trace`/`artifacts` source selector.
92
+
93
+ Group trajectory retrieval and downloads use the backend Trace export API. The
94
+ CLI waits for the export URL, downloads the gzip JSON-lines payload, and keeps
95
+ the existing `group_id`/`trajectories` JSON layout. The export wait timeout is
96
+ 30 minutes by default; override it in seconds with
97
+ `--download-timeout SECONDS` on trajectory commands or `pull`.
98
+
90
99
  Group downloads always auto-paginate and include every Job in the Group. In
91
100
  particular, `trajectory`, detailed `metrics`, `report`, and `artifacts` enumerate
92
101
  the complete Job list instead of relying on the backend Group endpoints that
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "benchloop_client"
7
- version = "0.2.23"
7
+ version = "0.2.24"
8
8
  description = "Benchloop CLI — command line tool for benchmark analysis data (metrics, trajectory, reports, artifacts)."
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.10"
@@ -31,6 +31,12 @@ OutputDirOption = typer.Option(
31
31
  SkipExistingOption = typer.Option(
32
32
  False, "--skip-existing", help="Skip files that already exist on disk."
33
33
  )
34
+ DownloadTimeoutOption = typer.Option(
35
+ 1800,
36
+ "--download-timeout",
37
+ min=1,
38
+ help="Maximum seconds to wait for a Group Trace trajectory export.",
39
+ )
34
40
 
35
41
 
36
42
  def expand(values: Optional[list[str]]) -> list[str]:
@@ -25,6 +25,7 @@ from aliyun.bmcli.runtime import get_runtime, report_and_exit, safe_run_concurre
25
25
  app = typer.Typer(help="Run metrics (Group overall/detail, Job).", no_args_is_help=True)
26
26
 
27
27
  _SCOPES = ("overall", "detail", "all")
28
+ _TRACE_SOURCE = "trace"
28
29
 
29
30
  _GROUP_RUN_METRIC_FIELDS = (
30
31
  "job_count",
@@ -112,7 +113,9 @@ def _validate_scope(scope: str) -> str:
112
113
 
113
114
 
114
115
  def _fetch_job(client, job_id: str):
115
- data = client.get_json(f"/openapi/metrics/{job_id}")
116
+ data = client.get_json(
117
+ f"/openapi/metrics/{job_id}", params={"source": _TRACE_SOURCE}
118
+ )
116
119
  return _format_job_payload(data, job_id)
117
120
 
118
121
 
@@ -125,7 +128,12 @@ def fetch_group_detail(
125
128
  nonlocal tokenized_metrics
126
129
  data = client.get_json(
127
130
  "/openapi/metrics",
128
- params={"group_id": group_id, "skip": skip, "limit": limit},
131
+ params={
132
+ "group_id": group_id,
133
+ "skip": skip,
134
+ "limit": limit,
135
+ "source": _TRACE_SOURCE,
136
+ },
129
137
  )
130
138
  if isinstance(data, dict) and tokenized_metrics is None:
131
139
  tokenized_metrics = data.get("tokenized_metrics")
@@ -154,7 +162,10 @@ def fetch_group_detail(
154
162
 
155
163
 
156
164
  def _fetch_group_overall(client, group_id: str):
157
- data = client.get_json(f"/openapi/groups/{group_id}/metrics/overall")
165
+ data = client.get_json(
166
+ f"/openapi/groups/{group_id}/metrics/overall",
167
+ params={"source": _TRACE_SOURCE},
168
+ )
158
169
  return _format_group_overall(data, group_id)
159
170
 
160
171
 
@@ -14,6 +14,7 @@ from aliyun.bmcli.commands._common import (
14
14
  JobOption,
15
15
  OutputDirOption,
16
16
  SkipExistingOption,
17
+ DownloadTimeoutOption,
17
18
  expand,
18
19
  resolve_output_dir,
19
20
  write_empty_tgz,
@@ -29,7 +30,10 @@ from aliyun.bmcli.commands.metrics import (
29
30
  fetch_group_detail,
30
31
  )
31
32
  from aliyun.bmcli.commands.tasks import _fetch_all as fetch_all_tasks
32
- from aliyun.bmcli.commands.trajectory import write_group_trajectories
33
+ from aliyun.bmcli.commands.trajectory import (
34
+ fetch_job_trajectory,
35
+ write_group_trajectories,
36
+ )
33
37
  from aliyun.bmcli.client import ArtifactMissingError, AuthError
34
38
  from aliyun.bmcli.ids import enforce_mutually_exclusive, require_any
35
39
  from aliyun.bmcli.runtime import get_runtime
@@ -61,6 +65,7 @@ def pull(
61
65
  ),
62
66
  output_dir: Optional[str] = OutputDirOption,
63
67
  skip_existing: bool = SkipExistingOption,
68
+ download_timeout: int = DownloadTimeoutOption,
64
69
  ) -> None:
65
70
  """Aggregate downloads for the given groups/jobs into the layout."""
66
71
  rt = get_runtime(ctx)
@@ -107,7 +112,7 @@ def pull(
107
112
  f"{jid}:trajectory",
108
113
  lambda jid=jid: _fetch_and_write_json(
109
114
  root / f"{jid}_trajectory.json",
110
- lambda: client.get_json(f"/openapi/jobs/{jid}/trajectory"),
115
+ lambda: fetch_job_trajectory(client, jid),
111
116
  ),
112
117
  )
113
118
  if "report" in selected:
@@ -160,7 +165,7 @@ def pull(
160
165
  rt.concurrency,
161
166
  trajectory_dest,
162
167
  pretty=rt.pretty,
163
- ids=job_ids(_group_jobs()),
168
+ download_timeout=download_timeout,
164
169
  ),
165
170
  )
166
171
  if "report" in selected:
@@ -2,13 +2,15 @@
2
2
 
3
3
  from __future__ import annotations
4
4
 
5
+ import gzip
5
6
  import json
6
7
  import os
7
8
  import shutil
8
9
  import tempfile
10
+ import time
9
11
  import zipfile
10
12
  from pathlib import Path
11
- from typing import Optional
13
+ from typing import Any, Iterator, Optional
12
14
  from urllib.parse import quote
13
15
 
14
16
  import typer
@@ -20,23 +22,230 @@ from aliyun.bmcli.commands._common import (
20
22
  OutputDirOption,
21
23
  SkipExistingOption,
22
24
  TaskIdOption,
25
+ DownloadTimeoutOption,
23
26
  expand,
24
27
  resolve_output_dir,
25
28
  resolve_task_id,
26
29
  task_job_ids,
27
30
  )
28
- from aliyun.bmcli.commands.group_data import (
29
- fetch_all_jobs,
30
- fetch_job_resources,
31
- job_ids,
32
- raise_for_failures,
33
- )
34
31
  from aliyun.bmcli.exit_codes import OK, PARTIAL
35
32
  from aliyun.bmcli.ids import enforce_mutually_exclusive, require_any
36
33
  from aliyun.bmcli.runtime import get_runtime, report_and_exit, safe_run_concurrent
37
34
 
38
35
  app = typer.Typer(help="Trajectory data (Group, Job).", no_args_is_help=True)
39
36
 
37
+ _TRACE_SOURCE = "trace"
38
+ _TRACE_DOWNLOAD_POLL_INTERVAL_SECONDS = 15.0
39
+ _TRACE_DOWNLOAD_TOTAL_TIMEOUT_SECONDS = 1800
40
+
41
+
42
+ def fetch_job_trajectory(client, job_id: str):
43
+ """Fetch one Job trajectory from the Trace-backed OpenAPI path."""
44
+ return client.get_json(
45
+ f"/openapi/jobs/{job_id}/trajectory",
46
+ params={"source": _TRACE_SOURCE},
47
+ )
48
+
49
+
50
+ def _trace_field(data: dict[str, Any], snake_name: str, camel_name: str) -> Any:
51
+ if snake_name in data:
52
+ return data[snake_name]
53
+ return data.get(camel_name)
54
+
55
+
56
+ def _trace_remaining(deadline: float) -> float:
57
+ remaining = deadline - time.monotonic()
58
+ if remaining <= 0:
59
+ raise RuntimeError("trace trajectory download total timeout exceeded")
60
+ return remaining
61
+
62
+
63
+ def _trace_sleep(seconds: float, deadline: float) -> None:
64
+ time.sleep(min(seconds, _trace_remaining(deadline)))
65
+
66
+
67
+ def _trace_retry_delay(data: dict[str, Any]) -> float:
68
+ raw_delay = _trace_field(data, "retry_after_seconds", "retryAfterSeconds")
69
+ try:
70
+ return max(1.0, min(float(raw_delay), _TRACE_DOWNLOAD_POLL_INTERVAL_SECONDS))
71
+ except (TypeError, ValueError):
72
+ return _TRACE_DOWNLOAD_POLL_INTERVAL_SECONDS
73
+
74
+
75
+ def _trace_download_url(data: dict[str, Any]) -> str:
76
+ value = _trace_field(data, "download_url", "downloadUrl")
77
+ if not isinstance(value, str) or not value.strip():
78
+ raise RuntimeError("trace trajectory download response has no download_url")
79
+ return value.strip()
80
+
81
+
82
+ def _wait_for_group_trace_download_url(
83
+ client, group_id: str, timeout_seconds: int = _TRACE_DOWNLOAD_TOTAL_TIMEOUT_SECONDS
84
+ ) -> str:
85
+ """Create and poll the backend Group Trace download job."""
86
+ if timeout_seconds <= 0:
87
+ raise ValueError("trace download timeout must be greater than 0 seconds")
88
+ encoded_group_id = quote(group_id, safe="")
89
+ create_path = f"/openapi/groups/{encoded_group_id}/trace-trajectory/download"
90
+ deadline = time.monotonic() + timeout_seconds
91
+
92
+ while True:
93
+ payload = client.get_json(create_path)
94
+ if not isinstance(payload, dict):
95
+ raise RuntimeError("trace trajectory download response must be a JSON object")
96
+ status = str(payload.get("status", "")).strip().lower()
97
+
98
+ if status == "succeeded":
99
+ return _trace_download_url(payload)
100
+ if status == "waiting_for_data":
101
+ delay = _trace_retry_delay(payload)
102
+ output.log(
103
+ f"trace trajectory not ready for group_id={group_id}; "
104
+ f"retrying in {delay:.1f}s"
105
+ )
106
+ _trace_sleep(delay, deadline)
107
+ continue
108
+ if status != "running":
109
+ raise RuntimeError(
110
+ f"unknown trace trajectory download status: {status or '<missing>'}"
111
+ )
112
+
113
+ download_job_name = _trace_field(payload, "download_job_name", "downloadJobName")
114
+ if not isinstance(download_job_name, str) or not download_job_name.strip():
115
+ raise RuntimeError(
116
+ "trace trajectory download response has no download_job_name"
117
+ )
118
+ encoded_name = quote(download_job_name.strip(), safe="")
119
+ status_path = (
120
+ f"/openapi/groups/{encoded_group_id}/trace-trajectory/download/{encoded_name}"
121
+ )
122
+ while True:
123
+ status_payload = client.get_json(status_path)
124
+ if not isinstance(status_payload, dict):
125
+ raise RuntimeError(
126
+ "trace trajectory download status must be a JSON object"
127
+ )
128
+ status = str(status_payload.get("status", "")).strip().lower()
129
+ if status == "succeeded":
130
+ return _trace_download_url(status_payload)
131
+ if status != "running":
132
+ raise RuntimeError(
133
+ f"unknown trace trajectory download status: "
134
+ f"{status or '<missing>'}"
135
+ )
136
+ _trace_sleep(_TRACE_DOWNLOAD_POLL_INTERVAL_SECONDS, deadline)
137
+
138
+
139
+ def _download_group_trace_archive(client, download_url: str, parent: Path) -> Path:
140
+ raw_path: Path | None = None
141
+ with tempfile.NamedTemporaryFile(
142
+ prefix=".trace-trajectory-",
143
+ suffix=".tmp",
144
+ dir=parent,
145
+ delete=False,
146
+ ) as raw_file:
147
+ raw_path = Path(raw_file.name)
148
+ try:
149
+ client.download_file(download_url, raw_path)
150
+ return raw_path
151
+ except Exception:
152
+ raw_path.unlink(missing_ok=True)
153
+ raise
154
+
155
+
156
+ def _iter_trace_records(download_path: Path) -> Iterator[dict[str, Any]]:
157
+ with download_path.open("rb") as probe:
158
+ is_gzip = probe.read(2) == b"\x1f\x8b"
159
+ opener = gzip.open if is_gzip else open
160
+ with opener(download_path, "rt", encoding="utf-8-sig") as source:
161
+ for line_number, line in enumerate(source, start=1):
162
+ if not line.strip():
163
+ continue
164
+ try:
165
+ record = json.loads(line)
166
+ except json.JSONDecodeError as exc:
167
+ raise RuntimeError(
168
+ f"invalid trace trajectory JSON at line {line_number}"
169
+ ) from exc
170
+ if not isinstance(record, dict):
171
+ raise RuntimeError(
172
+ f"trace trajectory row at line {line_number} is not an object"
173
+ )
174
+ yield record
175
+
176
+
177
+ def _trace_record_job_id(record: dict[str, Any]) -> str | None:
178
+ value = _trace_field(record, "job_id", "jobId")
179
+ if value is None:
180
+ value = record.get("JobId")
181
+ if not isinstance(value, str) or not value.strip():
182
+ return None
183
+ return value.strip()
184
+
185
+
186
+ def _fetch_group_trace_job_ids(client, group_id: str) -> list[str]:
187
+ path = f"/openapi/groups/{quote(group_id, safe='')}/job-ids"
188
+ payload = client.get_json(path)
189
+ if not isinstance(payload, dict):
190
+ raise RuntimeError("group job-ids response must be a JSON object")
191
+ raw_job_ids = payload.get("job_ids", [])
192
+ if not isinstance(raw_job_ids, list):
193
+ raise RuntimeError("job-ids response field 'job_ids' must be a list")
194
+
195
+ result: list[str] = []
196
+ seen: set[str] = set()
197
+ for raw_job_id in raw_job_ids:
198
+ if not isinstance(raw_job_id, str) or not raw_job_id.strip():
199
+ raise RuntimeError("job-ids response contains an invalid job_id")
200
+ job_id = raw_job_id.strip()
201
+ if job_id not in seen:
202
+ seen.add(job_id)
203
+ result.append(job_id)
204
+ return result
205
+
206
+
207
+ def fetch_group_trace_trajectory(
208
+ client,
209
+ group_id: str,
210
+ *,
211
+ download_timeout: int = _TRACE_DOWNLOAD_TOTAL_TIMEOUT_SECONDS,
212
+ ) -> dict[str, Any]:
213
+ """Fetch a Group trajectory through the backend Trace download API.
214
+
215
+ The backend performs the SLS export. The CLI only converts the downloaded
216
+ JSON-lines file into its existing Group response shape; it never requests
217
+ trajectory data one Job at a time.
218
+ """
219
+ with tempfile.TemporaryDirectory(prefix=".trace-trajectory-") as temp_dir:
220
+ temp_root = Path(temp_dir)
221
+ download_url = _wait_for_group_trace_download_url(
222
+ client, group_id, download_timeout
223
+ )
224
+ raw_path = _download_group_trace_archive(client, download_url, temp_root)
225
+ try:
226
+ group_job_ids = _fetch_group_trace_job_ids(client, group_id)
227
+ rows_by_job: dict[str, list[dict[str, Any]]] = {}
228
+ for record in _iter_trace_records(raw_path):
229
+ job_id = _trace_record_job_id(record)
230
+ if job_id is None:
231
+ raise RuntimeError("trace trajectory row is missing job_id")
232
+ rows_by_job.setdefault(job_id, []).append(record)
233
+
234
+ ordered_job_ids = list(group_job_ids)
235
+ known_job_ids = set(ordered_job_ids)
236
+ ordered_job_ids.extend(
237
+ job_id for job_id in rows_by_job if job_id not in known_job_ids
238
+ )
239
+ return {
240
+ "group_id": group_id,
241
+ "trajectories": [
242
+ {"job_id": job_id, "trajectory": rows_by_job.get(job_id, [])}
243
+ for job_id in ordered_job_ids
244
+ ],
245
+ }
246
+ finally:
247
+ raw_path.unlink(missing_ok=True)
248
+
40
249
 
41
250
  def write_group_trajectories(
42
251
  client,
@@ -46,52 +255,30 @@ def write_group_trajectories(
46
255
  *,
47
256
  pretty: bool,
48
257
  ids: list[str] | None = None,
258
+ download_timeout: int = _TRACE_DOWNLOAD_TOTAL_TIMEOUT_SECONDS,
49
259
  ) -> Path:
50
- if ids is None:
51
- ids = job_ids(fetch_all_jobs(client, group_id))
260
+ # ``concurrency`` and ``ids`` remain in the helper signature for callers
261
+ # from older CLI code, but Group Trace downloads are now backend-owned.
262
+ del concurrency, ids
263
+ data = fetch_group_trace_trajectory(
264
+ client, group_id, download_timeout=download_timeout
265
+ )
52
266
  output.ensure_dir(dest.parent)
53
267
 
54
- with tempfile.TemporaryDirectory(
55
- prefix=f".{group_id}-trajectories-", dir=dest.parent
56
- ) as temp_dir:
57
- temp_root = Path(temp_dir)
58
-
59
- def fetch_one(api, job_id: str) -> Path:
60
- data = api.get_json(f"/openapi/jobs/{job_id}/trajectory")
61
- path = temp_root / f"{quote(job_id, safe='-_.')}.json"
62
- output.write_json_file(path, data, pretty=pretty)
63
- return path
64
-
65
- result = fetch_job_resources(client, ids, fetch_one, concurrency)
66
- raise_for_failures("trajectories", result.failures)
67
-
68
- temp_output = None
69
- try:
70
- with tempfile.NamedTemporaryFile(
71
- mode="w",
72
- encoding="utf-8",
73
- prefix=f".{dest.name}.",
74
- suffix=".tmp",
75
- dir=dest.parent,
76
- delete=False,
77
- ) as target:
78
- temp_output = Path(target.name)
79
- separator = ",\n" if pretty else ","
80
- target.write('{\n "group_id": ' if pretty else '{"group_id":')
81
- json.dump(group_id, target, ensure_ascii=False)
82
- target.write(
83
- ',\n "trajectories": [\n' if pretty else ',"trajectories":['
84
- )
85
- for index, path in enumerate(result.items):
86
- if index:
87
- target.write(separator)
88
- with path.open("r", encoding="utf-8") as source:
89
- shutil.copyfileobj(source, target)
90
- target.write(" ]\n}\n" if pretty else "]}\n")
91
- os.replace(temp_output, dest)
92
- finally:
93
- if temp_output is not None:
94
- temp_output.unlink(missing_ok=True)
268
+ temp_output: Path | None = None
269
+ try:
270
+ with tempfile.NamedTemporaryFile(
271
+ prefix=f".{dest.name}.",
272
+ suffix=".tmp",
273
+ dir=dest.parent,
274
+ delete=False,
275
+ ) as target:
276
+ temp_output = Path(target.name)
277
+ output.write_json_file(temp_output, data, pretty=pretty)
278
+ os.replace(temp_output, dest)
279
+ finally:
280
+ if temp_output is not None:
281
+ temp_output.unlink(missing_ok=True)
95
282
  return dest
96
283
 
97
284
 
@@ -174,6 +361,7 @@ def get(
174
361
  ctx: typer.Context,
175
362
  group: Optional[list[str]] = GroupOption,
176
363
  job: Optional[list[str]] = JobOption,
364
+ download_timeout: int = DownloadTimeoutOption,
177
365
  ) -> None:
178
366
  """Query trajectory; prints a JSON array to stdout."""
179
367
  rt = get_runtime(ctx)
@@ -185,13 +373,15 @@ def get(
185
373
  if jobs:
186
374
  res = safe_run_concurrent(
187
375
  jobs,
188
- lambda jid: rt.client.get_json(f"/openapi/jobs/{jid}/trajectory"),
376
+ lambda jid: fetch_job_trajectory(rt.client, jid),
189
377
  rt.concurrency,
190
378
  )
191
379
  else:
192
380
  res = safe_run_concurrent(
193
381
  groups,
194
- lambda gid: rt.client.get_json(f"/openapi/groups/{gid}/trajectory"),
382
+ lambda gid: fetch_group_trace_trajectory(
383
+ rt.client, gid, download_timeout=download_timeout
384
+ ),
195
385
  rt.concurrency,
196
386
  )
197
387
  report_and_exit(res, emit=True, pretty=rt.pretty)
@@ -205,6 +395,7 @@ def download(
205
395
  task_id: Optional[str] = TaskIdOption,
206
396
  output_dir: Optional[str] = OutputDirOption,
207
397
  skip_existing: bool = SkipExistingOption,
398
+ download_timeout: int = DownloadTimeoutOption,
208
399
  ) -> None:
209
400
  """Download complete trajectory data to the standard file layout."""
210
401
  rt = get_runtime(ctx)
@@ -258,9 +449,7 @@ def download(
258
449
  job_id: str,
259
450
  ) -> tuple[str, Path, Optional[Path], Optional[str]]:
260
451
  try:
261
- data = rt.client.get_json(
262
- f"/openapi/jobs/{quote(job_id, safe='')}/trajectory"
263
- )
452
+ data = fetch_job_trajectory(rt.client, quote(job_id, safe=""))
264
453
  # job_id comes from the server-side task snapshot. Encode it
265
454
  # before using it as a local filename so an unexpected slash
266
455
  # or traversal segment can never escape the temporary folder.
@@ -325,7 +514,7 @@ def download(
325
514
  output.log(f"skip existing {dest}")
326
515
  job_progress.advance(str(dest))
327
516
  return str(dest)
328
- data = rt.client.get_json(f"/openapi/jobs/{job_id}/trajectory")
517
+ data = fetch_job_trajectory(rt.client, job_id)
329
518
  output.write_json_file(dest, data, pretty=rt.pretty)
330
519
  output.log(f"wrote {dest}")
331
520
  job_progress.advance(str(dest))
@@ -343,6 +532,7 @@ def download(
343
532
  rt.concurrency,
344
533
  dest,
345
534
  pretty=rt.pretty,
535
+ download_timeout=download_timeout,
346
536
  )
347
537
  output.log(f"wrote {dest}")
348
538
  group_progress.advance(str(dest))
@@ -17,7 +17,9 @@ class TestMetricsGetJobUrl:
17
17
  mock_client.get_json.return_value = {"score": 0.9}
18
18
  result = invoke(runner, ["--api-key", "k", "metrics", "get", "-j", "J1"])
19
19
  assert result.exit_code == 0
20
- mock_client.get_json.assert_called_once_with("/openapi/metrics/J1")
20
+ mock_client.get_json.assert_called_once_with(
21
+ "/openapi/metrics/J1", params={"source": "trace"}
22
+ )
21
23
 
22
24
  def test_get_job_stdout_is_json_array(self, runner, mock_client):
23
25
  mock_client.get_json.return_value = {"score": 0.9}
@@ -114,6 +116,13 @@ class TestMetricsGetGroupUrl:
114
116
  urls = [c.args[0] for c in mock_client.get_json.call_args_list]
115
117
  assert "/openapi/groups/G1/metrics/overall" in urls
116
118
  assert "/openapi/metrics" in urls
119
+ metric_calls = [
120
+ call
121
+ for call in mock_client.get_json.call_args_list
122
+ if call.args[0]
123
+ in {"/openapi/groups/G1/metrics/overall", "/openapi/metrics"}
124
+ ]
125
+ assert {call.kwargs["params"]["source"] for call in metric_calls} == {"trace"}
117
126
 
118
127
  def test_get_group_scope_overall_calls_one_endpoint(self, runner, mock_client):
119
128
  mock_client.get_json.return_value = {}
@@ -125,6 +134,7 @@ class TestMetricsGetGroupUrl:
125
134
  urls = [c.args[0] for c in mock_client.get_json.call_args_list]
126
135
  assert "/openapi/groups/G1/metrics/overall" in urls
127
136
  assert "/openapi/metrics" not in urls
137
+ assert mock_client.get_json.call_args.kwargs["params"] == {"source": "trace"}
128
138
 
129
139
  def test_get_group_scope_detail_calls_one_endpoint(self, runner, mock_client):
130
140
  mock_client.get_json.return_value = {}
@@ -136,6 +146,12 @@ class TestMetricsGetGroupUrl:
136
146
  urls = [c.args[0] for c in mock_client.get_json.call_args_list]
137
147
  assert "/openapi/metrics" in urls
138
148
  assert "/openapi/groups/G1/metrics/overall" not in urls
149
+ metric_calls = [
150
+ call
151
+ for call in mock_client.get_json.call_args_list
152
+ if call.args[0] == "/openapi/metrics"
153
+ ]
154
+ assert all(call.kwargs["params"]["source"] == "trace" for call in metric_calls)
139
155
 
140
156
  def test_get_group_stdout_is_json_array(self, runner, mock_client):
141
157
  mock_client.get_json.return_value = {}
@@ -349,7 +365,9 @@ class TestMetricsDownloadJob:
349
365
  runner,
350
366
  ["--api-key", "k", "metrics", "download", "-j", "J1", "-o", str(tmp_path)],
351
367
  )
352
- mock_client.get_json.assert_called_once_with("/openapi/metrics/J1")
368
+ mock_client.get_json.assert_called_once_with(
369
+ "/openapi/metrics/J1", params={"source": "trace"}
370
+ )
353
371
 
354
372
 
355
373
  class TestMetricsDownloadGroup:
@@ -445,6 +463,13 @@ class TestMetricsDownloadGroup:
445
463
  urls = [c.args[0] for c in mock_client.get_json.call_args_list]
446
464
  assert "/openapi/groups/G1/metrics/overall" in urls
447
465
  assert "/openapi/metrics" in urls
466
+ metric_calls = [
467
+ call
468
+ for call in mock_client.get_json.call_args_list
469
+ if call.args[0]
470
+ in {"/openapi/groups/G1/metrics/overall", "/openapi/metrics"}
471
+ ]
472
+ assert {call.kwargs["params"]["source"] for call in metric_calls} == {"trace"}
448
473
 
449
474
  def test_download_group_detail_auto_paginates_all_metrics(
450
475
  self, runner, mock_client, tmp_path
@@ -491,12 +516,17 @@ class TestMetricsDownloadGroup:
491
516
  call.kwargs["params"]["skip"]
492
517
  for call in mock_client.get_json.call_args_list
493
518
  ] == [0, 2]
519
+ assert all(
520
+ call.kwargs["params"]["source"] == "trace"
521
+ for call in mock_client.get_json.call_args_list
522
+ )
494
523
 
495
524
  def test_download_group_detail_recovers_from_duplicate_page_boundaries(
496
525
  self, runner, mock_client, tmp_path
497
526
  ):
498
527
  def response(path, params=None):
499
528
  assert path == "/openapi/metrics"
529
+ assert params["source"] == "trace"
500
530
  skip = params["skip"]
501
531
  limit = params["limit"]
502
532
  if limit == 1000: