ap-client 0.3.0.dev0__tar.gz → 0.3.1__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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: ap-client
3
- Version: 0.3.0.dev0
3
+ Version: 0.3.1
4
4
  Summary: Agent Platform API Client & CLI
5
5
  Requires-Python: >=3.10
6
6
  Requires-Dist: pyyaml>=6.0
@@ -15,6 +15,7 @@ from urllib.parse import parse_qsl, quote, urlencode, urlsplit, urlunsplit
15
15
  import requests
16
16
 
17
17
  from .config import Config, ConfigurationError, get_config, has_cluster_header
18
+ from .managed_ak import require_managed_ak_ack, validate_ak_id, validate_ak_selection
18
19
 
19
20
  _DEFAULT_CONNECT_TIMEOUT = 300
20
21
  _DEFAULT_READ_TIMEOUT = 86400
@@ -459,7 +460,12 @@ class APIClient:
459
460
 
460
461
  def _post(self, path: str, data: Any, timeout: TimeoutType = None) -> Any:
461
462
  """Send a POST request."""
462
- return self._request("POST", path, json_body=data, timeout=timeout)
463
+ from .managed_ak import parameter_ak_ids, require_parameter_ak_ack
464
+
465
+ references = parameter_ak_ids(data) if path == "/jobs" else set()
466
+ result = self._request("POST", path, json_body=data, timeout=timeout)
467
+ require_parameter_ak_ack(result, references)
468
+ return result
463
469
 
464
470
  def _delete(self, path: str, *, params: Optional[dict] = None) -> Any:
465
471
  """Send a DELETE request."""
@@ -911,9 +917,19 @@ class APIClient:
911
917
  checkpoint: Optional[dict] = None,
912
918
  benchmark: Optional[str] = None,
913
919
  experiment: Optional[str] = None,
920
+ ak_ids: Optional[list[str]] = None,
914
921
  ) -> dict:
915
922
  """Create a group."""
923
+ if ak_ids is not None and not isinstance(ak_ids, list):
924
+ raise ValueError("ak_ids must be a list")
925
+ references = list(ak_ids or [])
926
+ for reference in references:
927
+ if validate_ak_id(reference) is None:
928
+ raise ValueError("ak_ids entries must be nonempty strings")
929
+ references = list(dict.fromkeys(references))
916
930
  body: dict = {}
931
+ if references:
932
+ body["ak_ids"] = references
917
933
  if name:
918
934
  body["name"] = name
919
935
  effective_tags = _merge_meta_tags(
@@ -965,11 +981,14 @@ class APIClient:
965
981
  body["benchmark"] = benchmark
966
982
  if experiment is not None:
967
983
  body["experiment"] = experiment
968
- return self._post("/groups", body)
984
+ result = self._post("/groups", body)
985
+ for reference in references:
986
+ require_managed_ak_ack(result, reference)
987
+ return result
969
988
 
970
989
  def create_job(
971
990
  self,
972
- template: str,
991
+ template: Optional[str] = None,
973
992
  params: Optional[dict] = None,
974
993
  params_list: Optional[list] = None,
975
994
  suite_name: Optional[str] = None,
@@ -1022,6 +1041,7 @@ class APIClient:
1022
1041
  checkpoint: Optional[dict] = None,
1023
1042
  benchmark: Optional[str] = None,
1024
1043
  experiment: Optional[str] = None,
1044
+ ak_id: Optional[str] = None,
1025
1045
  ) -> dict:
1026
1046
  """Create a job."""
1027
1047
  body = self.build_create_job_body(
@@ -1068,13 +1088,16 @@ class APIClient:
1068
1088
  checkpoint=checkpoint,
1069
1089
  benchmark=benchmark,
1070
1090
  experiment=experiment,
1091
+ ak_id=ak_id,
1071
1092
  )
1072
1093
 
1073
- return self._post("/jobs", body, timeout=timeout)
1094
+ result = self._post("/jobs", body, timeout=timeout)
1095
+ require_managed_ak_ack(result, ak_id)
1096
+ return result
1074
1097
 
1075
1098
  def build_create_job_body(
1076
1099
  self,
1077
- template: str,
1100
+ template: Optional[str] = None,
1078
1101
  params: Optional[dict] = None,
1079
1102
  params_list: Optional[list] = None,
1080
1103
  suite_name: Optional[str] = None,
@@ -1126,6 +1149,7 @@ class APIClient:
1126
1149
  checkpoint: Optional[dict] = None,
1127
1150
  benchmark: Optional[str] = None,
1128
1151
  experiment: Optional[str] = None,
1152
+ ak_id: Optional[str] = None,
1129
1153
  ) -> dict:
1130
1154
  """Build the /jobs request body for job creation."""
1131
1155
  _validate_credential_source_selection(
@@ -1134,7 +1158,14 @@ class APIClient:
1134
1158
  params=params,
1135
1159
  params_list=params_list,
1136
1160
  )
1137
- body: dict = {"template": template}
1161
+ if not template and not profile_id:
1162
+ raise ValueError("Either template or profile_id is required")
1163
+ validate_ak_selection(ak_id, params, params_list, overrides)
1164
+ body: dict = {}
1165
+ if template is not None:
1166
+ body["template"] = template
1167
+ if ak_id is not None:
1168
+ body["ak_id"] = ak_id
1138
1169
 
1139
1170
  if params is not None:
1140
1171
  body["params"] = params
@@ -1289,9 +1320,13 @@ class APIClient:
1289
1320
  next_token: Optional[str] = None,
1290
1321
  pagination: Optional[str] = None,
1291
1322
  include_total: bool = True,
1323
+ ak_id: Optional[str] = None,
1292
1324
  ) -> dict:
1293
1325
  """List jobs."""
1326
+ validate_ak_id(ak_id)
1294
1327
  params = {"skip": skip, "limit": limit}
1328
+ if ak_id is not None:
1329
+ params["ak_id"] = ak_id
1295
1330
  if template:
1296
1331
  params["template"] = template
1297
1332
  if group_id:
@@ -1485,9 +1520,13 @@ class APIClient:
1485
1520
  meta_job_type: Optional[str] = None,
1486
1521
  upstream_platform: Optional[str] = None,
1487
1522
  upstream_job_id: Optional[str] = None,
1523
+ ak_id: Optional[str] = None,
1488
1524
  ) -> dict:
1489
1525
  """List groups."""
1526
+ validate_ak_id(ak_id)
1490
1527
  params: dict = {"skip": skip, "limit": limit}
1528
+ if ak_id is not None:
1529
+ params["ak_id"] = ak_id
1491
1530
  if name is not None:
1492
1531
  params["name"] = name
1493
1532
  if group_id is not None:
@@ -17,13 +17,22 @@ from ap_client.api import (
17
17
  _OTHERS_MODEL,
18
18
  APIError,
19
19
  PaginationProtocolError,
20
+ _is_sensitive_key,
20
21
  _merge_meta_tags,
22
+ _redact_sensitive_data,
21
23
  set_cluster_override,
22
24
  set_verbose_override,
23
25
  )
24
26
  from ap_client.config import ENV_VAR_SPECS, ConfigurationError, _parse_bool, normalize_output_format
25
27
  from ap_client.exporter import export_group, export_job
26
28
  from ap_client.irepo_commands import register as _register_dataset_repo
29
+ from ap_client.managed_ak import (
30
+ ManagedAKProtocolError,
31
+ require_managed_ak_ack,
32
+ response_ak_ids,
33
+ validate_ak_id,
34
+ validate_ak_selection,
35
+ )
27
36
  from ap_client.profile_commands import profile_app
28
37
  from ap_client.waiter import (
29
38
  WaitTimeoutError,
@@ -188,9 +197,28 @@ def _shell_command(args: list[str]) -> str:
188
197
  return " ".join(shlex.quote(str(arg)) for arg in args)
189
198
 
190
199
 
200
+ def _retry_argument_contains_secret(raw: str) -> bool:
201
+ """Inspect display arguments without changing the submission payload."""
202
+ try:
203
+ value = json.loads(raw)
204
+ except (json.JSONDecodeError, TypeError):
205
+ value = raw
206
+
207
+ def contains_secret(item: Any) -> bool:
208
+ if isinstance(item, dict):
209
+ return any(
210
+ _is_sensitive_key(str(key)) or contains_secret(field) for key, field in item.items()
211
+ )
212
+ if isinstance(item, list):
213
+ return any(contains_secret(field) for field in item)
214
+ return _redact_sensitive_data(item) != item
215
+
216
+ return contains_secret(value)
217
+
218
+
191
219
  def _build_job_create_retry_command(
192
220
  *,
193
- template: str,
221
+ template: Optional[str],
194
222
  agenthub_ref: Optional[str],
195
223
  instance_id: Optional[str],
196
224
  dataset: Optional[str],
@@ -223,8 +251,13 @@ def _build_job_create_retry_command(
223
251
  profile: Optional[str] = None,
224
252
  resource_profile_id: Optional[str] = None,
225
253
  timeout: Optional[float] = None,
254
+ ak_id: Optional[str] = None,
226
255
  ) -> str:
227
- args = ["ap", "job", "create", template]
256
+ args = ["ap", "job", "create"]
257
+ if template is not None:
258
+ args.append(template)
259
+ if ak_id is not None:
260
+ args.extend(["--ak-id", ak_id])
228
261
  if profile:
229
262
  args.extend(["--profile", profile])
230
263
  if resource_profile_id:
@@ -287,6 +320,14 @@ def _build_job_create_retry_command(
287
320
  args.extend(["--group-post-process", group_post_process])
288
321
  elif enable_post_process:
289
322
  args.append("--enable-post-process")
323
+ if any(_retry_argument_contains_secret(str(arg)) for arg in args):
324
+ # A redacted command would silently retry with different credentials.
325
+ # Keep the original invocation private and print only the retry key.
326
+ return (
327
+ "Rerun your original command with "
328
+ f"--idempotency-key {shlex.quote(idempotency_key)}. "
329
+ "The command is omitted because it contains inline credentials."
330
+ )
290
331
  args.extend(["--idempotency-key", idempotency_key])
291
332
  return _shell_command(args)
292
333
 
@@ -1085,6 +1126,7 @@ def _idempotency_item_key(
1085
1126
  account_pool: Optional[str],
1086
1127
  resource_profile_id: Optional[str] = None,
1087
1128
  max_failure_retries: int = 0,
1129
+ ak_id: Optional[str] = None,
1088
1130
  ) -> str:
1089
1131
  # Client-supplied item keys are authoritative; server hashing is only a
1090
1132
  # fallback for non-CLI callers, so this payload intentionally need not
@@ -1102,6 +1144,8 @@ def _idempotency_item_key(
1102
1144
  "queue": queue,
1103
1145
  "account_pool": account_pool,
1104
1146
  }
1147
+ if ak_id is not None:
1148
+ payload_data["ak_id"] = ak_id
1105
1149
  if resource_profile_id is not None:
1106
1150
  payload_data["resource_profile_id"] = resource_profile_id
1107
1151
  if max_failure_retries:
@@ -2595,6 +2639,7 @@ def meta_job_type_list(
2595
2639
  @job_app.command("list")
2596
2640
  def job_list(
2597
2641
  template: Optional[str] = typer.Option(None, "--template", help="Template name"),
2642
+ ak_id: Optional[str] = typer.Option(None, "--ak-id", help="AP-managed AK ID"),
2598
2643
  group_id: Optional[str] = typer.Option(None, "--group-id", help="Group ID"),
2599
2644
  status: Optional[str] = typer.Option(None, "--status", help=_JOB_STATUS_HELP),
2600
2645
  job_id: Optional[str] = typer.Option(None, "--job-id", help="Job ID"),
@@ -2670,6 +2715,10 @@ def job_list(
2670
2715
  ),
2671
2716
  ):
2672
2717
  """List jobs."""
2718
+ try:
2719
+ validate_ak_id(ak_id)
2720
+ except ValueError as exc:
2721
+ raise typer.BadParameter(str(exc), param_hint="--ak-id") from None
2673
2722
  output_format = _normalize_output_format(output_format, keep_table=True)
2674
2723
  selected_columns = (
2675
2724
  _parse_columns_option(columns) if output_format in ("plain", "table") else None
@@ -2696,6 +2745,7 @@ def job_list(
2696
2745
  client = get_client()
2697
2746
  result = client.list_jobs(
2698
2747
  template=template,
2748
+ **({"ak_id": ak_id} if ak_id is not None else {}),
2699
2749
  group_id=group_id,
2700
2750
  status=status,
2701
2751
  job_id=job_id,
@@ -2735,7 +2785,10 @@ def job_list(
2735
2785
 
2736
2786
  @job_app.command("create")
2737
2787
  def job_create(
2738
- template: str = typer.Argument(..., help="Template name"),
2788
+ template: Optional[str] = typer.Argument(None, help="Template name; optional with --profile"),
2789
+ ak_id: Optional[str] = typer.Option(
2790
+ None, "--ak-id", help="AP-managed AK ID; requires server support"
2791
+ ),
2739
2792
  agenthub_ref: Optional[str] = typer.Option(
2740
2793
  None, "--agenthub-ref", help="Agent-Hub branch/commit; overrides AP_AGENTHUB_REF"
2741
2794
  ),
@@ -3016,7 +3069,7 @@ def job_create(
3016
3069
  """Submit a job.
3017
3070
 
3018
3071
  Model configuration:
3019
- All jobs must provide model, model_base_url, and model_api_key through --params/-p.
3072
+ Supply template parameters through --params/-p; select a managed credential with --ak-id.
3020
3073
 
3021
3074
  Examples:
3022
3075
  # Single job - full parameters
@@ -3057,7 +3110,10 @@ def job_create(
3057
3110
  different requested group_id (409 idempotency_conflict).
3058
3111
  """
3059
3112
  output_format = _normalize_output_format(output_format)
3113
+ if not template and not profile:
3114
+ raise typer.BadParameter("TEMPLATE is required unless --profile is provided")
3060
3115
  try:
3116
+ validate_ak_id(ak_id)
3061
3117
  resource_profile_id = _normalize_resource_profile_id(resource_profile_id)
3062
3118
  except ValueError as exc:
3063
3119
  _emit_error(str(exc))
@@ -3076,6 +3132,7 @@ def job_create(
3076
3132
  or idempotency_key is not None
3077
3133
  or profile is not None
3078
3134
  or implicit_resource_profile_idempotency
3135
+ or ak_id is not None
3079
3136
  )
3080
3137
  submission_idempotency_key = (
3081
3138
  idempotency_key or str(uuid.uuid4()) if idempotency_enabled else None
@@ -3110,6 +3167,11 @@ def job_create(
3110
3167
  _emit_error(str(exc))
3111
3168
  raise typer.Exit(1)
3112
3169
  overrides_dict = json.loads(overrides) if overrides else None
3170
+ try:
3171
+ validate_ak_selection(ak_id, params_dict, None, overrides_dict)
3172
+ except ValueError as exc:
3173
+ _emit_error(str(exc))
3174
+ raise typer.Exit(1) from None
3113
3175
  if isinstance(overrides_dict, dict) and overrides_dict.get("sidecar_resources"):
3114
3176
  typer.secho(
3115
3177
  "warning: 'overrides.sidecar_resources' is deprecated; "
@@ -3215,7 +3277,7 @@ def job_create(
3215
3277
  if not suite_name and not group_id:
3216
3278
  from datetime import datetime as _dt
3217
3279
 
3218
- suite_name = f"{template}-sp-{profile}-{_dt.now().strftime('%m%d-%H%M')}"
3280
+ suite_name = f"{template or 'profile'}-sp-{profile}-{_dt.now().strftime('%m%d-%H%M')}"
3219
3281
 
3220
3282
  # Parse instances from -i (simple comma split, no range/exclude expansion)
3221
3283
  instances_list: Optional[list[str]] = None
@@ -3255,6 +3317,7 @@ def job_create(
3255
3317
  eval_config=eval_config_dict,
3256
3318
  profile_id=profile,
3257
3319
  resource_profile_id=resource_profile_id,
3320
+ **({"ak_id": ak_id} if ak_id is not None else {}),
3258
3321
  profile_version=profile_version_req,
3259
3322
  instances=instances_list,
3260
3323
  dataset=dataset,
@@ -3276,6 +3339,11 @@ def job_create(
3276
3339
  _emit_dry_run_warnings(dry_run_warnings)
3277
3340
  body["dry_run"] = True
3278
3341
  result = client._post("/jobs", body, timeout=timeout)
3342
+ try:
3343
+ require_managed_ak_ack(result, ak_id)
3344
+ except ManagedAKProtocolError as exc:
3345
+ _emit_error(str(exc))
3346
+ raise typer.Exit(1) from None
3279
3347
  _print_formatted(result, "json" if output_format == "plain" else output_format)
3280
3348
  return
3281
3349
 
@@ -3314,6 +3382,7 @@ def job_create(
3314
3382
  idempotency_key=submission_idempotency_key,
3315
3383
  profile=profile,
3316
3384
  resource_profile_id=resource_profile_id,
3385
+ **({"ak_id": ak_id} if ak_id is not None else {}),
3317
3386
  timeout=timeout,
3318
3387
  )
3319
3388
  if idempotency_enabled and submission_idempotency_key is not None
@@ -3357,6 +3426,7 @@ def job_create(
3357
3426
  eval_config=eval_config_dict,
3358
3427
  profile_id=profile,
3359
3428
  resource_profile_id=resource_profile_id,
3429
+ **({"ak_id": ak_id} if ak_id is not None else {}),
3360
3430
  profile_version=profile_version_req,
3361
3431
  instances=instances_list,
3362
3432
  dataset=dataset,
@@ -3370,7 +3440,11 @@ def job_create(
3370
3440
  timeout=timeout,
3371
3441
  **post_process_kwargs,
3372
3442
  )
3443
+ require_managed_ak_ack(result, ak_id)
3373
3444
  _require_resource_profile_ack(result, resource_profile_id)
3445
+ except ManagedAKProtocolError as exc:
3446
+ _emit_error(str(exc))
3447
+ raise typer.Exit(1) from None
3374
3448
  except Exception:
3375
3449
  emit_profile_retry_hint_once()
3376
3450
  raise
@@ -3385,6 +3459,8 @@ def job_create(
3385
3459
  else:
3386
3460
  print("[green]Profile submission completed:[/]")
3387
3461
  print(f" profile: {profile}")
3462
+ if ak_id is not None:
3463
+ print(f" ak_ids: {response_ak_ids(result)}")
3388
3464
  if resource_profile_id is not None:
3389
3465
  print(f" resource_profile_id: {resource_profile_id}")
3390
3466
  print(f" group_id: {result.get('group_id')}")
@@ -3496,6 +3572,12 @@ def job_create(
3496
3572
  *(item.get("credential_type") for item in parsed_params_list if isinstance(item, dict)),
3497
3573
  )
3498
3574
 
3575
+ try:
3576
+ validate_ak_selection(ak_id, params_dict, parsed_params_list, overrides_dict)
3577
+ except ValueError as exc:
3578
+ _emit_error(str(exc))
3579
+ raise typer.Exit(1) from None
3580
+
3499
3581
  is_batch = False
3500
3582
  if parsed_params_list:
3501
3583
  is_batch = True
@@ -3555,6 +3637,7 @@ def job_create(
3555
3637
  group_post_process=group_post_process,
3556
3638
  idempotency_key=submission_idempotency_key,
3557
3639
  resource_profile_id=resource_profile_id,
3640
+ **({"ak_id": ak_id} if ak_id is not None else {}),
3558
3641
  timeout=timeout,
3559
3642
  )
3560
3643
  if not dry_run
@@ -3627,6 +3710,7 @@ def job_create(
3627
3710
  total_submitted = 0
3628
3711
  total_failed = 0
3629
3712
  batch_results: list[dict] = []
3713
+ observed_ak_ids: set[str] = set()
3630
3714
  observed_group_id = target_group_id
3631
3715
  observed_queue_id: Optional[str] = None
3632
3716
 
@@ -3649,6 +3733,7 @@ def job_create(
3649
3733
  queue=queue,
3650
3734
  account_pool=account_pool,
3651
3735
  resource_profile_id=resource_profile_id,
3736
+ **({"ak_id": ak_id} if ak_id is not None else {}),
3652
3737
  max_failure_retries=max_failure_retries,
3653
3738
  )
3654
3739
  batch.append(item)
@@ -3681,6 +3766,7 @@ def job_create(
3681
3766
  credential_type=credential_type,
3682
3767
  account_pool=account_pool,
3683
3768
  resource_profile_id=resource_profile_id,
3769
+ **({"ak_id": ak_id} if ak_id is not None else {}),
3684
3770
  checkpoint_id=checkpoint_id,
3685
3771
  checkpoint_path=checkpoint_path,
3686
3772
  checkpoint=checkpoint_dict,
@@ -3723,6 +3809,7 @@ def job_create(
3723
3809
  credential_type=credential_type,
3724
3810
  account_pool=account_pool,
3725
3811
  resource_profile_id=resource_profile_id,
3812
+ **({"ak_id": ak_id} if ak_id is not None else {}),
3726
3813
  checkpoint_id=checkpoint_id,
3727
3814
  checkpoint_path=checkpoint_path,
3728
3815
  checkpoint=checkpoint_dict,
@@ -3779,6 +3866,7 @@ def job_create(
3779
3866
  credential_type=credential_type,
3780
3867
  account_pool=account_pool,
3781
3868
  resource_profile_id=resource_profile_id,
3869
+ **({"ak_id": ak_id} if ak_id is not None else {}),
3782
3870
  checkpoint_id=checkpoint_id,
3783
3871
  checkpoint_path=checkpoint_path,
3784
3872
  checkpoint=checkpoint_dict,
@@ -3787,6 +3875,10 @@ def job_create(
3787
3875
  timeout=timeout,
3788
3876
  **post_process_kwargs,
3789
3877
  )
3878
+ require_managed_ak_ack(batch_result, ak_id)
3879
+ observed_ak_ids.update(response_ak_ids(batch_result))
3880
+ for batch_job in batch_result.get("jobs", []):
3881
+ observed_ak_ids.update(response_ak_ids(batch_job))
3790
3882
  _require_resource_profile_ack(batch_result, resource_profile_id)
3791
3883
  batch_group_id = batch_result.get("group_id")
3792
3884
  if not observed_group_id and batch_group_id:
@@ -3820,6 +3912,9 @@ def job_create(
3820
3912
  current_job=f"group={observed_group_id}" if observed_group_id else "",
3821
3913
  current_stage=f"submitted={total_submitted} failed={total_failed}",
3822
3914
  )
3915
+ except ManagedAKProtocolError as exc:
3916
+ _emit_error(str(exc))
3917
+ raise typer.Exit(1) from None
3823
3918
  except Exception:
3824
3919
  emit_retry_hint_once()
3825
3920
  raise
@@ -3832,6 +3927,7 @@ def job_create(
3832
3927
  "failed": total_failed,
3833
3928
  "batches": batch_results,
3834
3929
  }
3930
+ result["ak_ids"] = sorted(observed_ak_ids)
3835
3931
  if submission_idempotency_key is not None:
3836
3932
  result["idempotency_key"] = submission_idempotency_key
3837
3933
  if resource_profile_id is not None:
@@ -3841,6 +3937,8 @@ def job_create(
3841
3937
  print(f" group_id: {result.get('group_id')}")
3842
3938
  if result.get("queue_id") is not None:
3843
3939
  print(f" queue_id: {result.get('queue_id')}")
3940
+ if ak_id is not None:
3941
+ print(f" ak_ids: {response_ak_ids(result)}")
3844
3942
  if resource_profile_id is not None:
3845
3943
  print(f" resource_profile_id: {resource_profile_id}")
3846
3944
  if submission_idempotency_key is not None:
@@ -3874,6 +3972,7 @@ def job_create(
3874
3972
  credential_type=credential_type,
3875
3973
  account_pool=account_pool,
3876
3974
  resource_profile_id=resource_profile_id,
3975
+ **({"ak_id": ak_id} if ak_id is not None else {}),
3877
3976
  checkpoint_id=checkpoint_id,
3878
3977
  checkpoint_path=checkpoint_path,
3879
3978
  checkpoint=checkpoint_dict,
@@ -3910,6 +4009,7 @@ def job_create(
3910
4009
  credential_type=credential_type,
3911
4010
  account_pool=account_pool,
3912
4011
  resource_profile_id=resource_profile_id,
4012
+ **({"ak_id": ak_id} if ak_id is not None else {}),
3913
4013
  checkpoint_id=checkpoint_id,
3914
4014
  checkpoint_path=checkpoint_path,
3915
4015
  checkpoint=checkpoint_dict,
@@ -3956,6 +4056,7 @@ def job_create(
3956
4056
  credential_type=credential_type,
3957
4057
  account_pool=account_pool,
3958
4058
  resource_profile_id=resource_profile_id,
4059
+ **({"ak_id": ak_id} if ak_id is not None else {}),
3959
4060
  checkpoint_id=checkpoint_id,
3960
4061
  checkpoint_path=checkpoint_path,
3961
4062
  checkpoint=checkpoint_dict,
@@ -3964,6 +4065,7 @@ def job_create(
3964
4065
  timeout=timeout,
3965
4066
  **post_process_kwargs,
3966
4067
  )
4068
+ require_managed_ak_ack(result, ak_id)
3967
4069
  _require_resource_profile_ack(result, resource_profile_id)
3968
4070
  jobs = result.get("jobs", [])
3969
4071
  if not jobs:
@@ -3978,6 +4080,9 @@ def job_create(
3978
4080
  current_job=f"job={job.get('job_id')}" if job.get("job_id") else "",
3979
4081
  current_stage=f"submitted={submitted} failed={failed}",
3980
4082
  )
4083
+ except ManagedAKProtocolError as exc:
4084
+ _emit_error(str(exc))
4085
+ raise typer.Exit(1) from None
3981
4086
  except Exception:
3982
4087
  emit_retry_hint_once()
3983
4088
  raise
@@ -4009,6 +4114,8 @@ def job_create(
4009
4114
  )
4010
4115
  if result.get("group_id"):
4011
4116
  print(f" group_id: {result.get('group_id')}")
4117
+ if ak_id is not None:
4118
+ print(f" ak_ids: {response_ak_ids(result)}")
4012
4119
  if resource_profile_id is not None:
4013
4120
  print(f" resource_profile_id: {resource_profile_id}")
4014
4121
 
@@ -4970,6 +5077,9 @@ def checkpoint_get(
4970
5077
 
4971
5078
  @group_app.command("create")
4972
5079
  def group_create(
5080
+ ak_id: Optional[list[str]] = typer.Option(
5081
+ None, "--ak-id", help="Managed AK reference (repeatable) to associate with the Group"
5082
+ ),
4973
5083
  name: Optional[str] = typer.Option(None, "--name", "-n", help="Group name"),
4974
5084
  template: Optional[str] = typer.Option(
4975
5085
  None,
@@ -5108,6 +5218,11 @@ def group_create(
5108
5218
  ):
5109
5219
  """Create a Group."""
5110
5220
  output_format = _normalize_output_format(output_format)
5221
+ try:
5222
+ for reference in ak_id or []:
5223
+ validate_ak_id(reference)
5224
+ except ValueError as exc:
5225
+ raise typer.BadParameter(str(exc), param_hint="--ak-id") from None
5111
5226
  client = get_client()
5112
5227
  eval_config_dict = json.loads(eval_config) if eval_config else None
5113
5228
  group_post_process_dict = _parse_json_object(group_post_process, "--group-post-process")
@@ -5163,30 +5278,37 @@ def group_create(
5163
5278
  _emit_error(str(exc))
5164
5279
  raise typer.Exit(1)
5165
5280
 
5166
- result = client.create_group(
5167
- name=name,
5168
- tags=group_tags,
5169
- template=template,
5170
- agenthub_revision=agenthub_revision,
5171
- queue_id=queue, # CLI --queue / -q 走标准 queue_id 字段
5172
- max_concurrency=max_concurrency,
5173
- eval_config=eval_config_dict,
5174
- model_base_url_collection=model_base_url_collection_list,
5175
- group_post_process=group_post_process_dict,
5176
- dataset=dataset,
5177
- model=model_name,
5178
- scaffold=scaffold,
5179
- agent=agent_dict,
5180
- checkpoint_id=checkpoint_id,
5181
- checkpoint_path=checkpoint_path,
5182
- checkpoint=checkpoint_dict,
5183
- benchmark=benchmark,
5184
- experiment=experiment,
5185
- )
5281
+ try:
5282
+ result = client.create_group(
5283
+ ak_ids=ak_id,
5284
+ name=name,
5285
+ tags=group_tags,
5286
+ template=template,
5287
+ agenthub_revision=agenthub_revision,
5288
+ queue_id=queue, # CLI --queue / -q 走标准 queue_id 字段
5289
+ max_concurrency=max_concurrency,
5290
+ eval_config=eval_config_dict,
5291
+ model_base_url_collection=model_base_url_collection_list,
5292
+ group_post_process=group_post_process_dict,
5293
+ dataset=dataset,
5294
+ model=model_name,
5295
+ scaffold=scaffold,
5296
+ agent=agent_dict,
5297
+ checkpoint_id=checkpoint_id,
5298
+ checkpoint_path=checkpoint_path,
5299
+ checkpoint=checkpoint_dict,
5300
+ benchmark=benchmark,
5301
+ experiment=experiment,
5302
+ )
5303
+ except ManagedAKProtocolError as exc:
5304
+ _emit_error(str(exc))
5305
+ raise typer.Exit(1) from None
5186
5306
  if output_format == "plain":
5187
5307
  print("[green]Group created:[/]")
5188
5308
  print(f" group_id: {result.get('group_id')}")
5189
5309
  print(f" name: {result.get('name')}")
5310
+ if response_ak_ids(result):
5311
+ print(f" ak_ids: {response_ak_ids(result)}")
5190
5312
  if result.get("queue_id") is not None:
5191
5313
  print(f" queue_id: {result.get('queue_id')}")
5192
5314
  if result.get("template") is not None:
@@ -5226,6 +5348,9 @@ def group_create(
5226
5348
  def group_list(
5227
5349
  name: Optional[str] = typer.Option(None, "--name", help="Group name"),
5228
5350
  group_id: Optional[str] = typer.Option(None, "--group-id", help="Group ID"),
5351
+ ak_id: Optional[str] = typer.Option(
5352
+ None, "--ak-id", help="Filter groups containing this managed AK"
5353
+ ),
5229
5354
  user_id: Optional[str] = typer.Option(None, "--user-id", help="User ID"),
5230
5355
  user_name: Optional[str] = typer.Option(None, "--user-name", help="User name"),
5231
5356
  skip: int = typer.Option(0, "--skip", help="Skip the first N entries"),
@@ -5280,6 +5405,7 @@ def group_list(
5280
5405
  profile=profile,
5281
5406
  name=name,
5282
5407
  group_id=group_id,
5408
+ **({"ak_id": ak_id} if ak_id is not None else {}),
5283
5409
  user_id=user_id,
5284
5410
  user_name=user_name,
5285
5411
  dataset=dataset,
@@ -7299,19 +7425,21 @@ async def _non_interactive_exec(client, job_id: str, container: str, command: li
7299
7425
  else:
7300
7426
  try:
7301
7427
  data = json.loads(message)
7302
- if "error" in data:
7428
+ if isinstance(data, dict) and "error" in data:
7303
7429
  sys.stderr.write(f"[错误] {data['error']}\n")
7304
7430
  sys.stderr.flush()
7305
- break
7431
+ raise typer.Exit(1)
7432
+ sys.stdout.write(message)
7433
+ sys.stdout.flush()
7306
7434
  except (json.JSONDecodeError, TypeError):
7307
7435
  sys.stdout.write(message)
7308
7436
  sys.stdout.flush()
7437
+ except typer.Exit:
7438
+ raise
7309
7439
  except websockets.exceptions.ConnectionClosed:
7310
7440
  # 命令执行完毕后服务端关闭连接,属于正常结束
7311
7441
  pass
7312
7442
  except Exception as e:
7313
- import typer
7314
-
7315
7443
  typer.echo(f"连接失败: {e}", err=True)
7316
7444
  raise typer.Exit(1)
7317
7445
 
@@ -0,0 +1,86 @@
1
+ """Client-side validation and acknowledgement of managed AK references."""
2
+
3
+ import re
4
+ from typing import Any
5
+
6
+
7
+ def validate_ak_id(value: Any) -> str | None:
8
+ if value is None:
9
+ return None
10
+ if not isinstance(value, str) or not 1 <= len(value) <= 128 or value != value.strip():
11
+ raise ValueError(
12
+ "ak_id must be a nonempty string of at most 128 characters without surrounding whitespace"
13
+ )
14
+ return value
15
+
16
+
17
+ def validate_ak_selection(ak_id: Any, params: Any, params_list: Any, overrides: Any) -> None:
18
+ validate_ak_id(ak_id)
19
+ candidates = [params, overrides]
20
+ if isinstance(params_list, list):
21
+ candidates.extend(params_list)
22
+ for value in candidates:
23
+ parameter_ak_ids(value)
24
+ if (
25
+ isinstance(value, dict)
26
+ and {"ak_id", "ak_selection_source", "user_ak_id", "key_id"} & value.keys()
27
+ ):
28
+ raise ValueError(
29
+ "AK selection belongs at the request top level; ak_selection_source is server-generated"
30
+ )
31
+
32
+
33
+ def parameter_ak_ids(value: Any) -> set[str]:
34
+ """Find explicit parameter references without resolving credential plaintext."""
35
+ if isinstance(value, dict):
36
+ return set().union(*(parameter_ak_ids(v) for v in value.values())) if value else set()
37
+ if isinstance(value, list):
38
+ return set().union(*(parameter_ak_ids(v) for v in value)) if value else set()
39
+ if isinstance(value, str) and "$akpool" in value:
40
+ match = re.fullmatch(r"\$akpool:([A-Za-z0-9][A-Za-z0-9_.-]{0,127})", value)
41
+ if not match:
42
+ raise ValueError("Use a complete $akpool:<ak-id> parameter value")
43
+ return {match[1]}
44
+ return set()
45
+
46
+
47
+ def require_parameter_ak_ack(response: Any, ak_ids: set[str]) -> None:
48
+ if ak_ids and (not isinstance(response, dict) or response.get("ak_placeholder_version") != 1):
49
+ raise ManagedAKProtocolError(",".join(sorted(ak_ids)), response)
50
+
51
+
52
+ class ManagedAKProtocolError(RuntimeError):
53
+ """A response cannot confirm that the requested AK reference was honored."""
54
+
55
+ def __init__(self, ak_id: str, response: Any):
56
+ self.response = response
57
+ self.job_ids = []
58
+ if isinstance(response, dict):
59
+ jobs = response.get("jobs") or [response]
60
+ self.job_ids = [
61
+ job["job_id"] for job in jobs if isinstance(job, dict) and job.get("job_id")
62
+ ]
63
+ self.group_id = response.get("group_id") if isinstance(response, dict) else None
64
+ super().__init__(
65
+ f"Managed AK protocol mismatch: server did not acknowledge ak_id={ak_id}. "
66
+ f"Returned Job IDs: {self.job_ids}. Verify these jobs before resubmitting; "
67
+ "the target AP server may need an upgrade."
68
+ + (f" Returned Group ID: {self.group_id}." if self.group_id else "")
69
+ )
70
+
71
+
72
+ def response_ak_ids(response: Any) -> list[str]:
73
+ if not isinstance(response, dict):
74
+ return []
75
+ value = response.get("ak_ids", [])
76
+ return value if isinstance(value, list) and all(isinstance(v, str) for v in value) else []
77
+
78
+
79
+ def require_managed_ak_ack(response: Any, ak_id: str | None) -> None:
80
+ if ak_id is None:
81
+ return
82
+ if ak_id not in response_ak_ids(response):
83
+ raise ManagedAKProtocolError(ak_id, response)
84
+ jobs = response.get("jobs", [])
85
+ if not isinstance(jobs, list) or any(ak_id not in response_ak_ids(job) for job in jobs):
86
+ raise ManagedAKProtocolError(ak_id, response)
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "ap-client"
7
- version = "0.3.0.dev0"
7
+ version = "0.3.1"
8
8
  description = "Agent Platform API Client & CLI"
9
9
  readme = { text = "A lightweight Python SDK and command line interface for Agent Platform. It provides helpers for configuring API access and managing templates, datasets, jobs, and groups.", content-type = "text/markdown" }
10
10
  requires-python = ">=3.10"