ap-client 0.5.0.dev0__tar.gz → 0.5.0.dev1__tar.gz

{ap_client-0.5.0.dev0 → ap_client-0.5.0.dev1}/PKG-INFO

@@ -1,8 +1,9 @@
 Metadata-Version: 2.5
 Name: ap-client
-Version: 0.5.0.dev0
+Version: 0.5.0.dev1
 Summary: Agent Platform API Client & CLI
 Requires-Python: >=3.10
+Requires-Dist: irepo-artifact==0.2.0
 Requires-Dist: pyyaml>=6.0
 Requires-Dist: requests>=2.28.0
 Requires-Dist: rich>=13.0.0

{ap_client-0.5.0.dev0 → ap_client-0.5.0.dev1}/ap_client/api.py

@@ -1283,6 +1283,7 @@ class APIClient:
         benchmark: Optional[str] = None,
         experiment: Optional[str] = None,
         ak_id: Optional[str] = None,
+        ak_mode: str = "auto",
         instance_range: Optional[str] = None,
         profile: Optional[str] = None,
     ) -> dict:
@@ -1334,6 +1335,7 @@ class APIClient:
             benchmark=benchmark,
             experiment=experiment,
             ak_id=ak_id,
+            ak_mode=ak_mode,
         )
 
         result = self._post("/jobs", body, timeout=timeout)
@@ -1395,6 +1397,7 @@ class APIClient:
         benchmark: Optional[str] = None,
         experiment: Optional[str] = None,
         ak_id: Optional[str] = None,
+        ak_mode: str = "auto",
         instance_range: Optional[str] = None,
         profile: Optional[str] = None,
     ) -> dict:
@@ -1410,11 +1413,15 @@ class APIClient:
         if not template and not profile_id and profile is None:
             raise ValueError("Either template, profile or profile_id is required")
         validate_ak_selection(ak_id, params, params_list, overrides)
+        if ak_mode not in {"auto", "manual", "synced"}:
+            raise ValueError("ak_mode must be auto, manual or synced")
         body: dict = {}
         if template is not None:
             body["template"] = template
         if ak_id is not None:
             body["ak_id"] = ak_id
+        if ak_mode != "auto":
+            body["ak_mode"] = ak_mode
 
         if params is not None:
             body["params"] = params

{ap_client-0.5.0.dev0 → ap_client-0.5.0.dev1}/ap_client/cli.py

@@ -26,7 +26,7 @@ from ap_client.api import (
 )
 from ap_client.config import ENV_VAR_SPECS, ConfigurationError, _parse_bool, normalize_output_format
 from ap_client.dataset_commands import register as _register_dataset_commands
-from ap_client.exporter import export_group, export_job
+from ap_client.exporter import export_group, export_job, export_jobs
 from ap_client.fs_commands import register as _register_fs_commands
 from ap_client.instance_commands import register as _register_instance_commands
 from ap_client.managed_ak import (
@@ -267,12 +267,15 @@ def _build_job_create_retry_command(
     resource_profile_id: Optional[str] = None,
     timeout: Optional[float] = None,
     ak_id: Optional[str] = None,
+    ak_mode: str = "auto",
 ) -> str:
     args = ["ap", "job", "create"]
     if template is not None:
         args.append(template)
     if ak_id is not None:
         args.extend(["--ak-id", ak_id])
+    if ak_mode != "auto":
+        args.extend(["--ak-mode", ak_mode])
     if profile is not None:
         args.extend(["--profile", profile])
     if resource_profile_id:
@@ -1240,6 +1243,7 @@ def _idempotency_item_key(
     resource_profile_id: Optional[str] = None,
     max_failure_retries: int = 0,
     ak_id: Optional[str] = None,
+    ak_mode: str = "auto",
 ) -> str:
     # Client-supplied item keys are authoritative; server hashing is only a
     # fallback for non-CLI callers, so this payload intentionally need not
@@ -1259,6 +1263,8 @@ def _idempotency_item_key(
     }
     if ak_id is not None:
         payload_data["ak_id"] = ak_id
+    if ak_mode != "auto":
+        payload_data["ak_mode"] = ak_mode
     if resource_profile_id is not None:
         payload_data["resource_profile_id"] = resource_profile_id
     if max_failure_retries:
@@ -2990,6 +2996,9 @@ def job_create(
     ak_id: Optional[str] = typer.Option(
         None, "--ak-id", help="AP-managed AK ID; requires server support"
     ),
+    ak_mode: str = typer.Option(
+        "auto", "--ak-mode", help="AK credential source: auto, manual or synced"
+    ),
     agenthub_ref: Optional[str] = typer.Option(
         None, "--agenthub-ref", help="Agent-Hub branch/commit; overrides AP_AGENTHUB_REF"
     ),
@@ -3332,6 +3341,8 @@ def job_create(
     output_format = _normalize_output_format(output_format)
     if not template and not profile:
         raise typer.BadParameter("TEMPLATE is required unless --profile is provided")
+    if ak_mode not in {"auto", "manual", "synced"}:
+        raise typer.BadParameter("Choose auto, manual or synced", param_hint="--ak-mode")
     try:
         validate_ak_id(ak_id)
         resource_profile_id = _normalize_resource_profile_id(resource_profile_id)
@@ -3528,6 +3539,7 @@ def job_create(
             profile=profile,
             resource_profile_id=resource_profile_id,
             **({"ak_id": ak_id} if ak_id is not None else {}),
+            **({"ak_mode": ak_mode} if ak_mode != "auto" else {}),
             instances=instances_list,
             instance_range=instance_range,
             dataset=dataset,
@@ -3601,6 +3613,7 @@ def job_create(
                 profile=profile,
                 resource_profile_id=resource_profile_id,
                 **({"ak_id": ak_id} if ak_id is not None else {}),
+                **({"ak_mode": ak_mode} if ak_mode != "auto" else {}),
                 timeout=timeout,
             )
             if idempotency_enabled and submission_idempotency_key is not None
@@ -3645,6 +3658,7 @@ def job_create(
                     profile=profile,
                     resource_profile_id=resource_profile_id,
                     **({"ak_id": ak_id} if ak_id is not None else {}),
+                    **({"ak_mode": ak_mode} if ak_mode != "auto" else {}),
                     instances=instances_list,
                     instance_range=instance_range,
                     dataset=dataset,
@@ -3860,6 +3874,7 @@ def job_create(
             idempotency_key=submission_idempotency_key,
             resource_profile_id=resource_profile_id,
             **({"ak_id": ak_id} if ak_id is not None else {}),
+            **({"ak_mode": ak_mode} if ak_mode != "auto" else {}),
             timeout=timeout,
         )
         if not dry_run
@@ -3956,6 +3971,7 @@ def job_create(
                         account_pool=account_pool,
                         resource_profile_id=resource_profile_id,
                         **({"ak_id": ak_id} if ak_id is not None else {}),
+                        **({"ak_mode": ak_mode} if ak_mode != "auto" else {}),
                         max_failure_retries=max_failure_retries,
                     )
                 batch.append(item)
@@ -3989,6 +4005,7 @@ def job_create(
                 account_pool=account_pool,
                 resource_profile_id=resource_profile_id,
                 **({"ak_id": ak_id} if ak_id is not None else {}),
+                **({"ak_mode": ak_mode} if ak_mode != "auto" else {}),
                 checkpoint_id=checkpoint_id,
                 checkpoint_path=checkpoint_path,
                 checkpoint=checkpoint_dict,
@@ -4032,6 +4049,7 @@ def job_create(
                         account_pool=account_pool,
                         resource_profile_id=resource_profile_id,
                         **({"ak_id": ak_id} if ak_id is not None else {}),
+                        **({"ak_mode": ak_mode} if ak_mode != "auto" else {}),
                         checkpoint_id=checkpoint_id,
                         checkpoint_path=checkpoint_path,
                         checkpoint=checkpoint_dict,
@@ -4089,6 +4107,7 @@ def job_create(
                         account_pool=account_pool,
                         resource_profile_id=resource_profile_id,
                         **({"ak_id": ak_id} if ak_id is not None else {}),
+                        **({"ak_mode": ak_mode} if ak_mode != "auto" else {}),
                         checkpoint_id=checkpoint_id,
                         checkpoint_path=checkpoint_path,
                         checkpoint=checkpoint_dict,
@@ -4195,6 +4214,7 @@ def job_create(
                 account_pool=account_pool,
                 resource_profile_id=resource_profile_id,
                 **({"ak_id": ak_id} if ak_id is not None else {}),
+                **({"ak_mode": ak_mode} if ak_mode != "auto" else {}),
                 checkpoint_id=checkpoint_id,
                 checkpoint_path=checkpoint_path,
                 checkpoint=checkpoint_dict,
@@ -4232,6 +4252,7 @@ def job_create(
                 account_pool=account_pool,
                 resource_profile_id=resource_profile_id,
                 **({"ak_id": ak_id} if ak_id is not None else {}),
+                **({"ak_mode": ak_mode} if ak_mode != "auto" else {}),
                 checkpoint_id=checkpoint_id,
                 checkpoint_path=checkpoint_path,
                 checkpoint=checkpoint_dict,
@@ -4279,6 +4300,7 @@ def job_create(
                     account_pool=account_pool,
                     resource_profile_id=resource_profile_id,
                     **({"ak_id": ak_id} if ak_id is not None else {}),
+                    **({"ak_mode": ak_mode} if ak_mode != "auto" else {}),
                     checkpoint_id=checkpoint_id,
                     checkpoint_path=checkpoint_path,
                     checkpoint=checkpoint_dict,
@@ -4857,10 +4879,47 @@ def job_artifacts(
         _print_formatted(results, output_format)
 
 
+def _validate_export_timeout(value: float) -> float:
+    import math
+
+    if not math.isfinite(value) or value <= 0:
+        raise typer.BadParameter("Timeout must be a finite number greater than zero")
+    return value
+
+
+def _export_channel_reporter(emit):
+    """Report completed Jobs, including callbacks from Group worker threads."""
+    from threading import Lock
+
+    from rich.text import Text
+
+    counts = {"oss": 0, "proxy": 0}
+    lock = Lock()
+
+    def report(job_id, status, channel):
+        with lock:
+            suffix = ""
+            if status == "downloaded" and channel in counts:
+                counts[channel] += 1
+                suffix = f"  channel={channel}"
+            emit(Text(f"{job_id}  {status}{suffix}"))
+
+    return report, counts
+
+
 @job_app.command("export")
 def job_export(
-    job_id: str = typer.Argument(..., help="Job ID"),
+    job_id: str = typer.Argument(..., help="Job ID or comma-separated Job IDs"),
     output: Optional[Path] = typer.Option(None, "--output", "-o", help="Export directory"),
+    workers: int = typer.Option(
+        1, "--workers", "-w", min=1, help="Number of concurrent export workers"
+    ),
+    timeout: float = typer.Option(
+        300,
+        "--timeout",
+        callback=_validate_export_timeout,
+        help="SDK request timeout in seconds, not the total export duration",
+    ),
     logs: bool = typer.Option(False, "--logs", help="Include container logs"),
     events: bool = typer.Option(False, "--events", help="Include events"),
     no_extract_artifacts: bool = typer.Option(
@@ -4870,17 +4929,50 @@ def job_export(
     ),
 ):
     """Export job artifacts to a local directory. Use --logs/--events to include logs and events."""
+    ids = list(dict.fromkeys(value.strip() for value in job_id.split(",")))
+    if not all(ids):
+        raise typer.BadParameter("Job IDs must not be empty", param_hint="job_id")
     client = get_client()
+    if len(ids) > 1:
+        report, channels = _export_channel_reporter(print)
+        dest, summary = export_jobs(
+            client,
+            ids,
+            output,
+            workers=workers,
+            result_callback=report,
+            timeout=timeout,
+            include_logs=logs,
+            include_events=events,
+            extract_artifacts=not no_extract_artifacts,
+        )
+        print(
+            f"Jobs exported: downloaded={summary.downloaded} "
+            f"skipped={summary.skipped} failed={len(summary.failures)}"
+        )
+        print(f"channels: oss={channels['oss']} proxy={channels['proxy']}")
+        print(f"  path: {dest}")
+        for failed_id, error in summary.failures.items():
+            print(f"[red]Failed job {failed_id}:[/] {error}")
+        if summary.failures:
+            raise typer.Exit(code=1)
+        return
+    job_id = ids[0]
+    result = {}
     dest = export_job(
         client,
         job_id,
         output,
+        result_callback=lambda jid, status, channel: result.update(status=status, channel=channel),
+        timeout=timeout,
         include_logs=logs,
         include_events=events,
         extract_artifacts=not no_extract_artifacts,
     )
     print(f"[green]Job exported:[/] {job_id}")
     print(f"  path: {dest}")
+    if result.get("channel") in {"oss", "proxy"} and result.get("status") == "downloaded":
+        print(f"  channel: {result['channel']}")
 
 
 @job_app.command("wait")
@@ -6355,6 +6447,12 @@ def group_export(
     workers: int = typer.Option(
         4, "--workers", "-w", min=1, help="Number of concurrent export workers"
     ),
+    timeout: float = typer.Option(
+        300,
+        "--timeout",
+        callback=_validate_export_timeout,
+        help="SDK request timeout in seconds, not the total export duration",
+    ),
     logs: bool = typer.Option(False, "--logs", help="Include container logs"),
     events: bool = typer.Option(False, "--events", help="Include events"),
     no_extract_artifacts: bool = typer.Option(
@@ -6415,12 +6513,15 @@ def group_export(
                 current_stage=current_stage,
             )
 
+        report, channels = _export_channel_reporter(progress.console.print)
         dest, summary = export_group(
             client,
             group_id,
             output,
             progress_callback=_on_progress,
             workers=workers,
+            result_callback=report,
+            timeout=timeout,
             include_logs=logs,
             include_events=events,
             extract_artifacts=not no_extract_artifacts,
@@ -6433,6 +6534,7 @@ def group_export(
         f"  [green]{summary.downloaded}[/] jobs downloaded, "
         f"[yellow]{summary.skipped}[/] jobs skipped (no artifact)"
     )
+    print(f"channels: oss={channels['oss']} proxy={channels['proxy']}")
     if summary.skipped_job_names:
         print(f"  Skipped: {', '.join(summary.skipped_job_names)}")