diff --git a/docs/cli/reference.mdx b/docs/cli/reference.mdx index d7c47002de..1a6849e8e8 100644 --- a/docs/cli/reference.mdx +++ b/docs/cli/reference.mdx @@ -2019,15 +2019,16 @@ nemo inference deployments create [OPTIONS] [NAME] * `--input-file`: Path to JSON file (use '-' for stdin) * `--input-data`: Input data for the request (JSON or YAML) -**Output Options:** +**Lifecycle Options:** -* `--output-format, -f `: Output format for an entity. [possible values: json, yaml, raw, code] +* `--wait`: Wait for the created deployment to be up and running +* `--watch`: Watch the created deployment until it is stable, streaming status updates +* `--timeout `: Maximum time to wait or watch in seconds [default: 1200] +* `--poll-interval `: Seconds between status checks [default: 3] -**Wait Options:** +**Output Options:** -* `--wait`: Wait for the created deployment to reach a terminal state -* `--timeout `: Maximum time to wait in seconds [default: 1200] -* `--poll-interval `: Seconds between status checks [default: 3] +* `--output-format, -f `: Output format for an entity. [possible values: json, yaml, raw, code] ##### nemo inference deployments delete @@ -3814,6 +3815,7 @@ nemo jobs [OPTIONS] COMMAND [ARGS]... * `resume`: Resume a paused platform job. * `get`: Get a platform job by name. * `update-status-details`: Update the status details of a platform job. +* `watch`: Watch a platform job until it reaches a terminal status. * `results`: Manage results * `steps`: Manage steps * `tasks`: Manage tasks @@ -3890,6 +3892,13 @@ nemo jobs create [OPTIONS] [NAME] * `--input-file`: Path to JSON file (use '-' for stdin) * `--input-data`: Input data for the request (JSON or YAML) +**Lifecycle Options:** + +* `--wait`: Wait for the created job to reach a terminal state without streaming logs +* `--watch`: Watch the created job to a terminal state +* `--timeout `: Maximum time to wait or watch in seconds +* `--poll-interval `: Seconds between status checks [default: 3] + **Output Options:** * `--output-format, -f `: Output format for an entity. [possible values: json, yaml, raw, code] @@ -4159,6 +4168,35 @@ nemo jobs update-status-details [OPTIONS] NAME * `--output-format, -f `: Output format for an entity. [possible values: json, yaml, raw, code] +#### nemo jobs watch + +Watch a platform job until it reaches a terminal status. + +**Usage:** + +```shell +nemo jobs watch [OPTIONS] NAME +``` + +**Arguments:** + +* ``: Name of the platform job to watch + +**Options:** + +* `--workspace`: Workspace containing the job +* `--attempt-id `: Filter logs to an attempt ID +* `--step-id`: Filter logs to a step ID +* `--task-id`: Filter logs to a task ID +* `--limit `: Maximum logs to fetch per page +* `--timeout `: Maximum watch time in seconds +* `--poll-interval `: Seconds between status checks [default: 3] +* `--history, --no-history`: Include logs already present before watching + +**Help:** + +* `--help, -h`: Show this message and exit. + #### nemo jobs results Manage results diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/inference/deployments/__init__.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/inference/deployments/__init__.py index e1120e9c8b..894a7d6ab0 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/inference/deployments/__init__.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/inference/deployments/__init__.py @@ -76,15 +76,28 @@ def create_deployments( wait: Annotated[ bool, typer.Option( - "--wait", help="Wait for the created deployment to reach a terminal state", rich_help_panel="Wait Options" + "--wait", help="Wait for the created deployment to be up and running", rich_help_panel="Lifecycle Options" + ), + ] = False, + watch: Annotated[ + bool, + typer.Option( + "--watch", + help="Watch the created deployment until it is stable, streaming status updates", + rich_help_panel="Lifecycle Options", ), ] = False, timeout: Annotated[ - int, typer.Option("--timeout", min=1, help="Maximum time to wait in seconds", rich_help_panel="Wait Options") + int, + typer.Option( + "--timeout", min=1, help="Maximum time to wait or watch in seconds", rich_help_panel="Lifecycle Options" + ), ] = 1200, poll_interval: Annotated[ int, - typer.Option("--poll-interval", min=1, help="Seconds between status checks", rich_help_panel="Wait Options"), + typer.Option( + "--poll-interval", min=1, help="Seconds between status checks", rich_help_panel="Lifecycle Options" + ), ] = 3, ) -> None: """Create a new ModelDeployment (version 1). @@ -131,12 +144,17 @@ def create_deployments( state: CLIContext = ctx.obj output_format = state.get_output_format(output_format) + if wait and watch: + raise typer.BadParameter("Cannot combine --wait and --watch.") + if handle_code_generation( ["inference", "deployments"], "create", all_kwargs, output_format, state, + watch_config={"type": "inference_deployment", "resource_label": "deployment"} if watch else None, + watch_options={"timeout": timeout, "poll_interval": poll_interval} if watch else None, wait_config={"type": "inference_deployment", "resource_label": "deployment"} if wait else None, wait_options={"timeout": timeout, "poll_interval": poll_interval} if wait else None, ): @@ -152,11 +170,10 @@ def create_deployments( no_truncate=state.get_no_truncate(), timestamp_format=state.get_timestamp_format(), ) - - if wait: + if wait or watch: wait_name = getattr(result, "name", None) or all_kwargs.get("name") if not wait_name: - raise RuntimeError("Unable to determine created resource name for --wait") + raise RuntimeError("Unable to determine created resource name for --wait/--watch") wait_workspace = all_kwargs.get("workspace") if not wait_for_inference_deployment( client, @@ -166,6 +183,7 @@ def create_deployments( poll_interval=poll_interval, ): raise typer.Exit(1) + return @app.command("delete") diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/jobs/__init__.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/jobs/__init__.py index a389bdf0f0..e0e3aef60a 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/jobs/__init__.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/api/jobs/__init__.py @@ -8,6 +8,8 @@ from typing import Annotated, Literal import typer +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.jobs.client import JobsClient from nemo_platform_ext.cli.core.api import build_kwargs, merge_filter_dict from nemo_platform_ext.cli.core.code_generator import handle_code_generation @@ -15,6 +17,7 @@ from nemo_platform_ext.cli.core.errors import handle_errors from nemo_platform_ext.cli.core.formatters import Column, check_output_columns_with_format, format_output from nemo_platform_ext.cli.core.help_formatter import collect_warnings, create_typer_app +from nemo_platform_ext.cli.core.job_watch_renderer import render_job_watch_events from nemo_platform_ext.cli.core.pagination import PaginationType, fetch_all_pages, warn_if_more_pages from nemo_platform_ext.cli.core.stdin_utils import read_data_input_with_flags, read_payload, validate_required_fields from nemo_platform_ext.cli.core.types import ( @@ -23,6 +26,7 @@ NoTruncateOption, OutputColumnsOption, ) +from nemo_platform_ext.cli.core.waiters import wait_for_platform_job _cli_child_results = _importlib_import_module("nemo_platform_ext.cli.commands.api.jobs.results") _cli_child_steps = _importlib_import_module("nemo_platform_ext.cli.commands.api.jobs.steps") @@ -96,6 +100,30 @@ def create_jobs( typer.Option("--input-data", help="Input data for the request (JSON or YAML)", rich_help_panel="Input Options"), ] = None, output_format: EntityOutputFormatOption = None, + wait: Annotated[ + bool, + typer.Option( + "--wait", + help="Wait for the created job to reach a terminal state without streaming logs", + rich_help_panel="Lifecycle Options", + ), + ] = False, + watch: Annotated[ + bool, + typer.Option("--watch", help="Watch the created job to a terminal state", rich_help_panel="Lifecycle Options"), + ] = False, + timeout: Annotated[ + int | None, + typer.Option( + "--timeout", min=1, help="Maximum time to wait or watch in seconds", rich_help_panel="Lifecycle Options" + ), + ] = None, + poll_interval: Annotated[ + int, + typer.Option( + "--poll-interval", min=1, help="Seconds between status checks", rich_help_panel="Lifecycle Options" + ), + ] = 3, ) -> None: """Create a new platform job. @@ -150,12 +178,26 @@ def create_jobs( state: CLIContext = ctx.obj output_format = state.get_output_format(output_format) - if handle_code_generation(["jobs"], "create", all_kwargs, output_format, state): + if wait and watch: + raise typer.BadParameter("Cannot combine --wait and --watch.") + + if handle_code_generation( + ["jobs"], + "create", + all_kwargs, + output_format, + state, + watch_config={"type": "platform_job", "resource_label": "job"} if watch else None, + watch_options={"timeout": timeout, "poll_interval": poll_interval} if watch else None, + wait_config={"type": "platform_job", "resource_label": "job"} if wait else None, + wait_options={"timeout": timeout if timeout is not None else 1200, "poll_interval": poll_interval} + if wait + else None, + ): return client = state.get_client() result = client.jobs.create(**all_kwargs) - format_output( result, is_list=False, @@ -163,6 +205,34 @@ def create_jobs( no_truncate=state.get_no_truncate(), timestamp_format=state.get_timestamp_format(), ) + if wait or watch: + wait_name = getattr(result, "name", None) or all_kwargs.get("name") + if not wait_name: + raise RuntimeError("Unable to determine created resource name for --wait/--watch") + wait_workspace = getattr(result, "workspace", None) or all_kwargs.get("workspace") + if wait_workspace is None: + wait_workspace = client._get_workspace_path_param() + jobs_client = client_from_platform(client, JobsClient) + if wait: + if not wait_for_platform_job( + jobs_client, + wait_name, + workspace=wait_workspace, + resource_label="job", + timeout=timeout if timeout is not None else 1200, + poll_interval=poll_interval, + ): + raise typer.Exit(1) + return + events = jobs_client.watch_job( + wait_name, + workspace=wait_workspace, + timeout=timeout, + poll_interval=poll_interval, + ) + if not render_job_watch_events(events, resource_label="job"): + raise typer.Exit(1) + return @app.command("delete") @@ -581,3 +651,46 @@ def update_status_details_jobs( no_truncate=state.get_no_truncate(), timestamp_format=state.get_timestamp_format(), ) + + +@app.command("watch") +@collect_warnings +@handle_errors +def watch_platform_job( + ctx: typer.Context, + name: Annotated[str, typer.Argument(help="Name of the platform job to watch")], + workspace: Annotated[str | None, typer.Option("--workspace", help="Workspace containing the job")] = None, + attempt_id: Annotated[int | None, typer.Option("--attempt-id", help="Filter logs to an attempt ID")] = None, + step_id: Annotated[str | None, typer.Option("--step-id", help="Filter logs to a step ID")] = None, + task_id: Annotated[str | None, typer.Option("--task-id", help="Filter logs to a task ID")] = None, + limit: Annotated[int | None, typer.Option("--limit", min=1, help="Maximum logs to fetch per page")] = None, + timeout: Annotated[int | None, typer.Option("--timeout", min=1, help="Maximum watch time in seconds")] = None, + poll_interval: Annotated[ + int, + typer.Option("--poll-interval", min=1, help="Seconds between status checks"), + ] = 3, + include_history: Annotated[ + bool, + typer.Option("--history/--no-history", help="Include logs already present before watching"), + ] = True, +) -> None: + """Watch a platform job until it reaches a terminal status.""" + state: CLIContext = ctx.obj + client = state.get_client() + jobs_client = client_from_platform(client, JobsClient) + if workspace is None: + workspace = client._get_workspace_path_param() + + events = jobs_client.watch_job( + name, + workspace=workspace, + attempt_id=attempt_id, + step_id=step_id, + task_id=task_id, + limit=limit, + timeout=timeout, + poll_interval=poll_interval, + include_history=include_history, + ) + if not render_job_watch_events(events, resource_label="job"): + raise typer.Exit(1) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/code_generator.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/code_generator.py index 3996810491..443cc883aa 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/code_generator.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/code_generator.py @@ -11,6 +11,11 @@ from nemo_platform_ext.cli.core.context import CLIContext +_INFERENCE_DEPLOYMENT_LIFECYCLE = "inference_deployment" +_PLATFORM_JOB_LIFECYCLE = "platform_job" +_LIFECYCLE_TYPES_WITH_DEADLINES = {_INFERENCE_DEPLOYMENT_LIFECYCLE} +_LIFECYCLE_TYPES_WITH_STATUS_ERROR_HANDLING = {_INFERENCE_DEPLOYMENT_LIFECYCLE} + def handle_code_generation( resource_path: list[str], @@ -20,6 +25,8 @@ def handle_code_generation( context: CLIContext, wait_config: dict[str, Any] | None = None, wait_options: dict[str, Any] | None = None, + watch_config: dict[str, Any] | None = None, + watch_options: dict[str, Any] | None = None, ) -> bool: """ Check if in code generation mode and generate code if needed. @@ -42,6 +49,8 @@ def handle_code_generation( base_url=base_url, wait_config=wait_config, wait_options=wait_options, + watch_config=watch_config, + watch_options=watch_options, ) formatted_code = format_code_output(code, language="python") print(formatted_code) @@ -57,6 +66,8 @@ def generate_python_code( base_url: str | None = None, wait_config: dict[str, Any] | None = None, wait_options: dict[str, Any] | None = None, + watch_config: dict[str, Any] | None = None, + watch_options: dict[str, Any] | None = None, ) -> str: """ Generate Python SDK code equivalent to a CLI command. @@ -72,16 +83,28 @@ def generate_python_code( """ lines = [] - wait_type = wait_config.get("type") if wait_config else None + if wait_config and watch_config: + raise ValueError("Only one of wait_config or watch_config may be provided") + + lifecycle_config = watch_config or wait_config + lifecycle_options = watch_options if watch_config else wait_options + lifecycle_type = lifecycle_config.get("type") if lifecycle_config else None + + lifecycle_mode = "watch" if watch_config else "wait" if wait_config else None - if wait_config: + if _lifecycle_uses_deadline(lifecycle_type, lifecycle_mode): lines.append("import time") - if wait_type == "inference_deployment": + if _lifecycle_uses_status_error_handling(lifecycle_type, lifecycle_mode): lines.append( "from nemo_platform import APIConnectionError, APIStatusError, APITimeoutError, NeMoPlatform, NotFoundError" ) else: lines.append("from nemo_platform import NeMoPlatform") + if lifecycle_type == _PLATFORM_JOB_LIFECYCLE: + lines.append("from nemo_platform_plugin.client.adapter import client_from_platform") + lines.append("from nemo_platform_plugin.jobs.client import JobsClient") + if lifecycle_mode == "wait": + lines.append("from nemo_platform_plugin.jobs.watch_types import JobStatusEvent, JobWatchTimeoutError") lines.append("") if base_url: @@ -89,15 +112,30 @@ def generate_python_code( else: lines.append("client = NeMoPlatform()") lines.append("") + if lifecycle_type == _PLATFORM_JOB_LIFECYCLE: + lines.append("jobs_client = client_from_platform(client, JobsClient)") + lines.append("") resource_chain = "client." + ".".join(resource_path) _append_method_call(lines, resource_chain, method, _format_method_args(args)) - if wait_config: - lines.extend(["", _render_wait_code(resource_path, args, wait_config, wait_options or {})]) + if lifecycle_config: + lines.extend( + [ + "", + _render_lifecycle_code( + resource_path, + args, + lifecycle_config, + lifecycle_options or {}, + mode=lifecycle_mode, + ), + ] + ) - lines.append("") - lines.append("print(response)") + if lifecycle_type != _PLATFORM_JOB_LIFECYCLE: + lines.append("") + lines.append("print(response)") return "\n".join(lines) @@ -138,16 +176,36 @@ def _format_python_literal(value: Any) -> str: return repr(value) -def _render_wait_code( +def _lifecycle_uses_deadline(lifecycle_type: object, mode: str | None) -> bool: + return lifecycle_type in _LIFECYCLE_TYPES_WITH_DEADLINES and not ( + lifecycle_type == _PLATFORM_JOB_LIFECYCLE and mode == "watch" + ) + + +def _lifecycle_uses_status_error_handling(lifecycle_type: object, mode: str | None) -> bool: + return lifecycle_type in _LIFECYCLE_TYPES_WITH_STATUS_ERROR_HANDLING and not ( + lifecycle_type == _PLATFORM_JOB_LIFECYCLE and mode == "watch" + ) + + +def _require_timeout(timeout: Any, lifecycle_type: object, mode: str | None) -> Any: + if timeout is not None: + return timeout + mode_label = f"{mode} " if mode else "" + raise ValueError(f"{mode_label}{lifecycle_type!r} lifecycle code generation requires timeout") + + +def _render_lifecycle_code( resource_path: list[str], args: dict[str, Any], - wait_config: dict[str, Any], - wait_options: dict[str, Any], + lifecycle_config: dict[str, Any], + lifecycle_options: dict[str, Any], + *, + mode: str | None, ) -> str: - wait_type = wait_config.get("type") - resource_label = str(wait_config.get("resource_label") or "resource") - timeout = wait_options.get("timeout", 1200) - poll_interval = wait_options.get("poll_interval", 3) + lifecycle_type = lifecycle_config.get("type") + timeout = lifecycle_options.get("timeout") + poll_interval = lifecycle_options.get("poll_interval", 3) resource_chain = "client." + ".".join(resource_path) status_kwargs = _format_keyword_args(args, ["workspace"]) resource_name = 'getattr(response, "name", None)' @@ -159,11 +217,14 @@ def _render_wait_code( resource_name = {resource_name} if not resource_name: raise RuntimeError("Unable to determine created resource name for --wait") - deadline = time.monotonic() + {timeout} """ ).strip() + if mode == "watch": + prelude = prelude.replace("--wait", "--watch") - if wait_type == "inference_deployment": + if lifecycle_type == _INFERENCE_DEPLOYMENT_LIFECYCLE: + timeout = _require_timeout(timeout, lifecycle_type, mode) + prelude = "\n".join([prelude, f"deadline = time.monotonic() + {timeout}"]) workspace_literal = _format_python_literal(args["workspace"]) if args.get("workspace") is not None else "None" return "\n\n".join( [ @@ -177,15 +238,23 @@ def _render_wait_code( ] ) - if wait_type == "platform_job": + if lifecycle_type == _PLATFORM_JOB_LIFECYCLE and mode == "watch": return "\n\n".join( [ prelude, - _render_platform_job_wait_code(resource_chain, status_kwargs, resource_label, poll_interval), + _render_platform_job_watch_code(args, timeout, poll_interval), + ] + ) + if lifecycle_type == _PLATFORM_JOB_LIFECYCLE and mode == "wait": + timeout = _require_timeout(timeout, lifecycle_type, mode) + return "\n\n".join( + [ + prelude, + _render_platform_job_wait_code(args, timeout, poll_interval), ] ) - raise ValueError(f"Unsupported wait config type: {wait_type!r}") + raise ValueError(f"Unsupported lifecycle config type: {lifecycle_type!r}") def _render_inference_deployment_wait_code( @@ -240,34 +309,52 @@ def _render_inference_deployment_wait_code( ).strip() +def _render_platform_job_watch_code( + args: dict[str, Any], + timeout: int | None, + poll_interval: int, +) -> str: + workspace = _format_python_literal(args["workspace"]) if args.get("workspace") is not None else "None" + + return dedent( + f""" + for event in jobs_client.watch_job( + resource_name, + workspace={workspace}, + timeout={timeout}, + poll_interval={poll_interval}, + ): + print(event) + """ + ).strip() + + def _render_platform_job_wait_code( - resource_chain: str, - status_kwargs: str, - resource_label: str, + args: dict[str, Any], + timeout: int, poll_interval: int, ) -> str: - resource_label_literal = _format_python_literal(resource_label) + workspace = _format_python_literal(args["workspace"]) if args.get("workspace") is not None else "None" return dedent( f""" - while True: - status_response = {resource_chain}.get_status(resource_name{status_kwargs}) - status = str(status_response.status or "").lower() - if status == "completed": - response = status_response - break - if status in {{"cancelled", "error"}}: - raise RuntimeError( - {resource_label_literal} + f" {{resource_name!r}} ended with status {{status!r}}" - ) - remaining = deadline - time.monotonic() - if remaining <= 0: - raise TimeoutError( - "Timed out waiting for " - + {resource_label_literal} - + f" {{resource_name!r}} to complete" - ) - time.sleep(min({poll_interval}, remaining)) + try: + for event in jobs_client.watch_job( + resource_name, + workspace={workspace}, + timeout={timeout}, + poll_interval={poll_interval}, + include_logs=False, + ): + if not isinstance(event, JobStatusEvent): + continue + if not event.terminal: + continue + if event.successful: + break + raise RuntimeError(f"Job {{resource_name!r}} ended with status {{event.status!r}}") + except JobWatchTimeoutError as exc: + raise TimeoutError(f"Timed out waiting for job {{resource_name!r}} to complete") from exc """ ).strip() diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/job_watch_renderer.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/job_watch_renderer.py new file mode 100644 index 0000000000..64eb01dec7 --- /dev/null +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/job_watch_renderer.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import time +from collections.abc import Iterable, Mapping +from datetime import datetime + +from nemo_platform_plugin.jobs.watch_types import ( + JobLogEvent, + JobStatusEvent, + JobWarningEvent, + JobWatchEvent, + JobWatchTimeoutError, +) +from rich.console import Console +from rich.text import Text + + +def render_job_watch_events( + events: Iterable[JobWatchEvent], + *, + console: Console | None = None, + error_console: Console | None = None, + resource_label: str | None = None, + start_time: float | None = None, +) -> bool: + """Render job watch events and return whether the terminal status succeeded.""" + output = console or Console() + errors = error_console or Console(stderr=True) + started_at = time.time() if start_time is None else start_time + terminal_event: JobStatusEvent | None = None + + try: + for event in events: + if isinstance(event, JobStatusEvent): + _render_status(output, event) + if event.terminal: + terminal_event = event + elif isinstance(event, JobLogEvent): + _render_log(output, event) + elif isinstance(event, JobWarningEvent): + _render_warning(output, event) + except JobWatchTimeoutError as exc: + errors.print(str(exc), style="red") + return False + + if terminal_event is None: + return False + + _emit_terminal_job_run_event(terminal_event, resource_label=resource_label, start_time=started_at) + + if terminal_event.successful: + output.print(f"Job {terminal_event.job_name!r} completed", style="green") + return True + + message = f"Job {terminal_event.job_name!r} ended with status {terminal_event.status!r}" + error_details = _status_details(terminal_event.error_details or {}) + if error_details: + message = f"{message}: {error_details}" + output.print(message, style="red") + return False + + +def _emit_terminal_job_run_event( + event: JobStatusEvent, + *, + resource_label: str | None, + start_time: float, +) -> None: + if resource_label is None: + return + from .waiters import _emit_job_run_event + + _emit_job_run_event(event, resource_label=resource_label, status=event.status, start_time=start_time) + + +def _render_status(console: Console, event: JobStatusEvent) -> None: + line = Text() + line.append(f"[{_time_label()}] ", style="dim") + line.append("Status: ") + line.append(event.status, style=_status_style(event)) + details = _status_details(event.status_details) + if details: + line.append(f" {details}", style="dim") + console.print(line) + + +def _render_log(console: Console, event: JobLogEvent) -> None: + line = Text() + line.append(f"[{_time_label(event.timestamp)}] ", style="dim") + scope = _scope(event) + if scope: + line.append(f"{scope} | ", style="dim") + line.append(event.message) + console.print(line) + + +def _render_warning(console: Console, event: JobWarningEvent) -> None: + line = Text() + line.append(f"[{_time_label()}] ", style="dim") + line.append(event.message, style="yellow") + console.print(line) + + +def _time_label(timestamp: datetime | None = None) -> str: + value = timestamp.astimezone() if timestamp is not None else datetime.now() + return value.strftime("%H:%M:%S") + + +def _status_style(event: JobStatusEvent) -> str: + if event.successful: + return "green bold" + if event.successful is False: + return "red bold" + return "cyan bold" + + +def _status_details(details: Mapping[str, object]) -> str: + parts = [] + for key, value in details.items(): + if value in (None, "", [], {}): + continue + if isinstance(value, str | int | float | bool): + parts.append(f"{key}={value}") + if len(parts) >= 6: + break + return " ".join(parts) + + +def _scope(event: JobLogEvent) -> str: + if event.step_id and event.task_id: + return f"{event.step_id}/{event.task_id}" + if event.step_id: + return event.step_id + if event.task_id: + return event.task_id + return "" diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/waiters.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/waiters.py index bba00723e3..552cf7c34f 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/waiters.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/waiters.py @@ -19,10 +19,22 @@ import logging import time +from collections.abc import Iterator from datetime import datetime, timezone from typing import Any from nemo_platform import APIConnectionError, APIStatusError, APITimeoutError, NotFoundError +from nemo_platform_plugin.client.response import NemoPaginatedResponse, NemoResponse +from nemo_platform_plugin.client.types import CursorPagination +from nemo_platform_plugin.jobs.client import JobsWatchClient +from nemo_platform_plugin.jobs.schemas import PlatformJobLog, PlatformJobStatusResponse +from nemo_platform_plugin.jobs.types import JobLogsQueryParams +from nemo_platform_plugin.jobs.watch_types import ( + JobStatusEvent, + JobWarningEvent, + JobWatchEvent, + JobWatchTimeoutError, +) from rich.console import Console from rich.live import Live from rich.text import Text @@ -51,6 +63,63 @@ def _pause(seconds: float) -> None: time.sleep(seconds) +class _WatchedJobsClient: + def __init__(self, jobs_client: JobsWatchClient) -> None: + self._jobs_client = jobs_client + self.last_status: PlatformJobStatusResponse | None = None + + def get_job_status( + self, + *, + workspace: str | None = None, + name: str, + ) -> NemoResponse[PlatformJobStatusResponse]: + response = self._jobs_client.get_job_status(workspace=workspace, name=name) + self.last_status = response.data() + return response + + def list_job_logs( + self, + *, + workspace: str | None = None, + name: str, + query_params: JobLogsQueryParams | None = None, + ) -> NemoPaginatedResponse[PlatformJobLog, CursorPagination]: + return self._jobs_client.list_job_logs(workspace=workspace, name=name, query_params=query_params) + + def watch_job( + self, + name: str, + *, + workspace: str | None = None, + poll_interval: float = 3, + timeout: float | None = None, + include_history: bool = True, + include_logs: bool = True, + attempt_id: int | None = None, + step_id: str | None = None, + task_id: str | None = None, + limit: int | None = None, + page_cursor: str | None = None, + ) -> Iterator[JobWatchEvent]: + from nemo_platform_plugin.jobs.watch import watch_job + + return watch_job( + self, + name, + workspace=workspace, + poll_interval=poll_interval, + timeout=timeout, + include_history=include_history, + include_logs=include_logs, + attempt_id=attempt_id, + step_id=step_id, + task_id=task_id, + limit=limit, + page_cursor=page_cursor, + ) + + def _seconds_since_creation(entry_timestamp: datetime | str | None, created_at: datetime | None) -> int | None: if created_at is None or entry_timestamp is None: return None @@ -187,6 +256,20 @@ def _make_live_display( return text +class _PlatformJobWaitLiveDisplay: + def __init__(self, *, start_time: float, timeout: int, poll_interval: int) -> None: + self.start_time = start_time + self.timeout = timeout + self.poll_interval = poll_interval + + def snapshot(self) -> tuple[str, int]: + return datetime.now().strftime("%H:%M:%S"), int(time.time() - self.start_time) + + def __rich__(self) -> Text: + polling_time, wait_elapsed = self.snapshot() + return _make_live_display(polling_time, self.timeout, self.poll_interval, wait_elapsed) + + def _sleep_until_next_poll(start_time: float, timeout: float, poll_interval: int) -> bool: if poll_interval <= 0: raise ValueError(f"_sleep_until_next_poll poll_interval must be greater than 0, got {poll_interval}") @@ -221,6 +304,7 @@ def wait_for_inference_deployment( timeout: int = 1200, poll_interval: int = 3, check_gateway: bool = True, + verbose: bool = True, ) -> bool: """Wait for an inference deployment to reach the requested status.""" if workspace is None: @@ -231,13 +315,15 @@ def wait_for_inference_deployment( last_status = "" last_message = "" - console.print(f"[bold]Waiting for deployment '{name}' to reach status: {status}[/bold]\n") + if verbose: + console.print(f"[bold]Waiting for deployment '{name}' to reach status: {status}[/bold]\n") with Live(console=console, refresh_per_second=4, transient=True) as live: while time.time() - start_time < timeout: wait_elapsed = int(time.time() - start_time) polling_time = datetime.now().strftime("%H:%M:%S") - live.update(_make_live_display(polling_time, timeout, poll_interval, wait_elapsed)) + if verbose: + live.update(_make_live_display(polling_time, timeout, poll_interval, wait_elapsed)) try: deployment = client.inference.deployments.retrieve(name, workspace=workspace) @@ -253,7 +339,7 @@ def wait_for_inference_deployment( last_status = current_status last_message = current_message - if history and len(history) > last_history_len: + if verbose and history and len(history) > last_history_len: live.stop() if last_history_len == 0: console.print() @@ -269,11 +355,13 @@ def wait_for_inference_deployment( console.print() live.start() - live.update(_make_live_display(polling_time, timeout, poll_interval, wait_elapsed)) + if verbose: + live.update(_make_live_display(polling_time, timeout, poll_interval, wait_elapsed)) if current_status == status and status != "DELETED": live.stop() - console.print(f"\n[green]✓ Deployment reached {status} status![/green]") + if verbose: + console.print(f"\n[green]✓ Deployment reached {status} status![/green]") if status == "READY" and check_gateway: remaining_timeout = timeout - (time.time() - start_time) if remaining_timeout <= 0: @@ -286,6 +374,7 @@ def wait_for_inference_deployment( provider_workspace, timeout=remaining_timeout, poll_interval=poll_interval, + verbose=verbose, ) return True @@ -303,11 +392,13 @@ def wait_for_inference_deployment( console.print("\n[red]✗ Deployment not found[/red]") return False except (APIConnectionError, APITimeoutError) as exc: - _print_transient_wait_error(live, "deployment status", exc) + if verbose: + _print_transient_wait_error(live, "deployment status", exc) except APIStatusError as exc: if exc.status_code not in _TRANSIENT_GATEWAY_STATUS_CODES: raise - _print_transient_wait_error(live, "deployment status", exc) + if verbose: + _print_transient_wait_error(live, "deployment status", exc) if not _sleep_until_next_poll(start_time, timeout, poll_interval): break @@ -321,7 +412,7 @@ def wait_for_inference_deployment( def wait_for_platform_job( - jobs_resource: Any, + jobs_client: JobsWatchClient, name: str, *, workspace: str | None = None, @@ -332,60 +423,56 @@ def wait_for_platform_job( """Wait for a platform job resource to complete.""" start_time = time.time() last_status = "" + jobs = _WatchedJobsClient(jobs_client) console.print(f"[bold]Waiting for {resource_label} '{name}' to complete[/bold]\n") - with Live(console=console, refresh_per_second=4, transient=True) as live: - while time.time() - start_time < timeout: - wait_elapsed = int(time.time() - start_time) - polling_time = datetime.now().strftime("%H:%M:%S") - live.update(_make_live_display(polling_time, timeout, poll_interval, wait_elapsed)) - - try: - job_status = jobs_resource.get_status(name, workspace=workspace) - except NotFoundError: - live.stop() - console.print(f"\n[red]✗ {resource_label.title()} not found[/red]") - return False - except (APIConnectionError, APITimeoutError) as exc: - _print_transient_wait_error(live, f"{resource_label} status", exc) - if not _sleep_until_next_poll(start_time, timeout, poll_interval): - break - continue - except APIStatusError as exc: - if exc.status_code not in _TRANSIENT_GATEWAY_STATUS_CODES: - raise - _print_transient_wait_error(live, f"{resource_label} status", exc) - if not _sleep_until_next_poll(start_time, timeout, poll_interval): - break - continue - - current_status = _status_text(getattr(job_status, "status", "")).lower() - if current_status != last_status: - live.stop() - console.print(_make_history_line(polling_time, wait_elapsed, current_status)) - last_status = current_status - console.print() - live.start() + live_display = _PlatformJobWaitLiveDisplay(start_time=start_time, timeout=timeout, poll_interval=poll_interval) + with Live(live_display, console=console, refresh_per_second=4, transient=True) as live: + try: + for event in jobs.watch_job( + name, + workspace=workspace, + poll_interval=poll_interval, + timeout=timeout, + include_logs=False, + ): + polling_time, wait_elapsed = live_display.snapshot() + live.update(live_display) + + if isinstance(event, JobWarningEvent): + live.stop() + console.print(f"\n[yellow]{event.message}[/yellow]") + live.start() + continue - if current_status == "completed": - live.stop() - _emit_job_run_event( - job_status, resource_label=resource_label, status=current_status, start_time=start_time - ) - console.print(f"\n[green]✓ {resource_label.title()} completed![/green]") - return True + if not isinstance(event, JobStatusEvent): + continue - if current_status in {"cancelled", "error"}: - live.stop() - _emit_job_run_event( - job_status, resource_label=resource_label, status=current_status, start_time=start_time - ) - console.print(f"\n[red]✗ {resource_label.title()} entered {current_status} state[/red]") - return False + current_status = event.status + if current_status != last_status: + live.stop() + console.print(_make_history_line(polling_time, wait_elapsed, current_status)) + last_status = current_status + console.print() + live.start() - if not _sleep_until_next_poll(start_time, timeout, poll_interval): - break + if event.terminal: + live.stop() + _emit_job_run_event( + jobs.last_status, resource_label=resource_label, status=current_status, start_time=start_time + ) + if event.successful: + console.print(f"\n[green]✓ {resource_label.title()} completed![/green]") + return True + console.print(f"\n[red]✗ {resource_label.title()} entered {current_status} state[/red]") + return False + except NotFoundError: + live.stop() + console.print(f"\n[red]✗ {resource_label.title()} not found[/red]") + return False + except JobWatchTimeoutError: + pass wait_elapsed = int(time.time() - start_time) detail = f"Last status: {last_status}" if last_status else "No status returned" @@ -399,12 +486,14 @@ def wait_for_gateway( workspace: str, timeout: float = 60, poll_interval: int = 1, + verbose: bool = True, ) -> bool: """Wait for the inference gateway to be able to route to a provider.""" start_time = time.time() start_timestamp = datetime.now().strftime("%H:%M:%S") - console.print(f"[bold]Waiting for gateway to be ready for provider '{provider_name}'[/bold]\n") + if verbose: + console.print(f"[bold]Waiting for gateway to be ready for provider '{provider_name}'[/bold]\n") def _make_gateway_display(polling_time: str, elapsed: int, status: str) -> Text: text = Text() @@ -418,12 +507,14 @@ def _make_gateway_display(polling_time: str, elapsed: int, status: str) -> Text: while time.time() - start_time < timeout: elapsed = int(time.time() - start_time) polling_time = datetime.now().strftime("%H:%M:%S") - live.update(_make_gateway_display(polling_time, elapsed, "Checking gateway...")) + if verbose: + live.update(_make_gateway_display(polling_time, elapsed, "Checking gateway...")) try: client.inference.gateway.provider.ready(provider_name, workspace=workspace) live.stop() - console.print(f" [{polling_time}] ({elapsed}s) [green]Gateway is ready![/green]") + if verbose: + console.print(f" [{polling_time}] ({elapsed}s) [green]Gateway is ready![/green]") return True except NotFoundError: pass diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_create_wait.py b/packages/nemo_platform_ext/tests/cli/commands/test_create_wait.py index b8892b742e..060212eb4e 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_create_wait.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_create_wait.py @@ -8,9 +8,18 @@ import pytest import typer - -# from nemo_platform_ext.cli.commands.api.customization.jobs import create_jobs as create_customization_job from nemo_platform_ext.cli.commands.api.inference.deployments import create_deployments +from nemo_platform_ext.cli.commands.api.jobs import create_jobs, watch_platform_job + + +class _CreatedJob: + def __init__(self, *, name: str = "created-job", workspace: str = "result-workspace") -> None: + self.name = name + self.workspace = workspace + + def model_dump(self, *, mode: str = "json") -> dict[str, str]: + assert mode == "json" + return {"name": self.name, "workspace": self.workspace} def _ctx(client: object) -> SimpleNamespace: @@ -23,55 +32,310 @@ def _ctx(client: object) -> SimpleNamespace: return SimpleNamespace(obj=state) -# def test_customization_job_create_waits_for_created_job() -> None: -# jobs = MagicMock() -# jobs.create.return_value = SimpleNamespace(name="created-job") -# client = SimpleNamespace(customization=SimpleNamespace(jobs=jobs)) -# ctx = _ctx(client) - -# with ( -# patch( -# "nemo_platform_ext.cli.commands.api.customization.jobs.handle_code_generation", -# return_value=False, -# ) as handle_code_generation, -# patch("nemo_platform_ext.cli.commands.api.customization.jobs.format_output"), -# patch( -# "nemo_platform_ext.cli.commands.api.customization.jobs.wait_for_platform_job", -# return_value=True, -# ) as wait_for_platform_job, -# ): -# create_customization_job( -# ctx, -# name="input-job", -# workspace="test-workspace", -# spec='{"training_type": "sft"}', -# wait=True, -# timeout=42, -# poll_interval=7, -# ) - -# handle_code_generation.assert_called_once_with( -# ["customization", "jobs"], -# "create", -# {"workspace": "test-workspace", "spec": {"training_type": "sft"}, "name": "input-job"}, -# None, -# ctx.obj, -# wait_config={"type": "platform_job", "resource_label": "customization job"}, -# wait_options={"timeout": 42, "poll_interval": 7}, -# ) -# jobs.create.assert_called_once_with( -# workspace="test-workspace", -# spec={"training_type": "sft"}, -# name="input-job", -# ) -# wait_for_platform_job.assert_called_once_with( -# jobs, -# "created-job", -# workspace="test-workspace", -# resource_label="customization job", -# timeout=42, -# poll_interval=7, -# ) +def test_jobs_create_watch_uses_sdk_watcher_and_outputs_created_job() -> None: + jobs = MagicMock() + jobs.create.return_value = _CreatedJob() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + ctx = _ctx(client) + events = object() + jobs_client = MagicMock() + jobs_client.watch_job.return_value = events + + with ( + patch( + "nemo_platform_ext.cli.commands.api.jobs.handle_code_generation", + return_value=False, + ) as handle_code_generation, + patch("nemo_platform_ext.cli.commands.api.jobs.client_from_platform", return_value=jobs_client), + patch("nemo_platform_ext.cli.commands.api.jobs.format_output") as format_output, + patch("nemo_platform_ext.cli.commands.api.jobs.render_job_watch_events", return_value=True) as render_events, + ): + create_jobs( + ctx, + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=False, + watch=True, + timeout=42, + poll_interval=7, + ) + + expected_kwargs = { + "workspace": "test-workspace", + "platform_spec": {}, + "source": "test-source", + "spec": {}, + "name": "input-job", + } + handle_code_generation.assert_called_once_with( + ["jobs"], + "create", + expected_kwargs, + None, + ctx.obj, + watch_config={"type": "platform_job", "resource_label": "job"}, + watch_options={"timeout": 42, "poll_interval": 7}, + wait_config=None, + wait_options=None, + ) + jobs.create.assert_called_once_with(**expected_kwargs) + format_output.assert_called_once_with( + jobs.create.return_value, + is_list=False, + output_format=None, + no_truncate=False, + timestamp_format=None, + ) + jobs_client.watch_job.assert_called_once_with( + "created-job", + workspace="result-workspace", + timeout=42, + poll_interval=7, + ) + render_events.assert_called_once_with(events, resource_label="job") + + +def test_jobs_create_watch_exits_when_renderer_reports_failure() -> None: + jobs = MagicMock() + jobs.create.return_value = _CreatedJob() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + jobs_client = MagicMock() + jobs_client.watch_job.return_value = object() + + with ( + patch("nemo_platform_ext.cli.commands.api.jobs.handle_code_generation", return_value=False), + patch("nemo_platform_ext.cli.commands.api.jobs.client_from_platform", return_value=jobs_client), + patch("nemo_platform_ext.cli.commands.api.jobs.render_job_watch_events", return_value=False), + pytest.raises(typer.Exit) as exc_info, + ): + create_jobs( + _ctx(client), + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=False, + watch=True, + ) + + assert exc_info.value.exit_code == 1 + + +def test_jobs_create_watch_has_no_default_timeout() -> None: + jobs = MagicMock() + jobs.create.return_value = _CreatedJob() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + ctx = _ctx(client) + events = object() + jobs_client = MagicMock() + jobs_client.watch_job.return_value = events + + with ( + patch( + "nemo_platform_ext.cli.commands.api.jobs.handle_code_generation", + return_value=False, + ) as handle_code_generation, + patch("nemo_platform_ext.cli.commands.api.jobs.client_from_platform", return_value=jobs_client), + patch("nemo_platform_ext.cli.commands.api.jobs.render_job_watch_events", return_value=True), + ): + create_jobs( + ctx, + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=False, + watch=True, + ) + + handle_code_generation.assert_called_once() + assert handle_code_generation.call_args.kwargs["watch_options"] == {"timeout": None, "poll_interval": 3} + jobs_client.watch_job.assert_called_once() + assert jobs_client.watch_job.call_args.kwargs["timeout"] is None + + +def test_jobs_create_wait_uses_quiet_waiter_and_outputs_created_job() -> None: + jobs = MagicMock() + jobs.create.return_value = _CreatedJob() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + ctx = _ctx(client) + jobs_client = MagicMock() + + with ( + patch( + "nemo_platform_ext.cli.commands.api.jobs.handle_code_generation", + return_value=False, + ) as handle_code_generation, + patch("nemo_platform_ext.cli.commands.api.jobs.client_from_platform", return_value=jobs_client), + patch("nemo_platform_ext.cli.commands.api.jobs.format_output") as format_output, + patch("nemo_platform_ext.cli.commands.api.jobs.wait_for_platform_job", return_value=True) as wait_for_job, + patch("nemo_platform_ext.cli.commands.api.jobs.render_job_watch_events") as render_events, + ): + create_jobs( + ctx, + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=True, + watch=False, + timeout=42, + poll_interval=7, + ) + + expected_kwargs = { + "workspace": "test-workspace", + "platform_spec": {}, + "source": "test-source", + "spec": {}, + "name": "input-job", + } + handle_code_generation.assert_called_once_with( + ["jobs"], + "create", + expected_kwargs, + None, + ctx.obj, + watch_config=None, + watch_options=None, + wait_config={"type": "platform_job", "resource_label": "job"}, + wait_options={"timeout": 42, "poll_interval": 7}, + ) + jobs.create.assert_called_once_with(**expected_kwargs) + wait_for_job.assert_called_once_with( + jobs_client, + "created-job", + workspace="result-workspace", + resource_label="job", + timeout=42, + poll_interval=7, + ) + format_output.assert_called_once_with( + jobs.create.return_value, + is_list=False, + output_format=None, + no_truncate=False, + timestamp_format=None, + ) + jobs_client.watch_job.assert_not_called() + render_events.assert_not_called() + + +def test_jobs_create_wait_uses_waiter_default_timeout() -> None: + jobs = MagicMock() + jobs.create.return_value = _CreatedJob() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + jobs_client = MagicMock() + + with ( + patch( + "nemo_platform_ext.cli.commands.api.jobs.handle_code_generation", + return_value=False, + ) as handle_code_generation, + patch("nemo_platform_ext.cli.commands.api.jobs.client_from_platform", return_value=jobs_client), + patch("nemo_platform_ext.cli.commands.api.jobs.wait_for_platform_job", return_value=True) as wait_for_job, + ): + create_jobs( + _ctx(client), + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=True, + watch=False, + ) + + assert handle_code_generation.call_args.kwargs["wait_options"] == {"timeout": 1200, "poll_interval": 3} + wait_for_job.assert_called_once() + assert wait_for_job.call_args.kwargs["timeout"] == 1200 + + +def test_jobs_create_wait_exits_when_waiter_reports_failure() -> None: + jobs = MagicMock() + jobs.create.return_value = _CreatedJob() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + + with ( + patch("nemo_platform_ext.cli.commands.api.jobs.handle_code_generation", return_value=False), + patch("nemo_platform_ext.cli.commands.api.jobs.wait_for_platform_job", return_value=False), + pytest.raises(typer.Exit) as exc_info, + ): + create_jobs( + _ctx(client), + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=True, + watch=False, + ) + + assert exc_info.value.exit_code == 1 + + +def test_jobs_create_rejects_wait_and_watch_together() -> None: + jobs = MagicMock() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + + with pytest.raises(SystemExit) as exc_info: + create_jobs( + _ctx(client), + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=True, + watch=True, + ) + + assert exc_info.value.code == 2 + jobs.create.assert_not_called() + + +def test_jobs_watch_command_uses_sdk_watcher() -> None: + client = SimpleNamespace(_get_workspace_path_param=MagicMock(return_value="default")) + events = object() + jobs_client = MagicMock() + jobs_client.watch_job.return_value = events + + with ( + patch("nemo_platform_ext.cli.commands.api.jobs.client_from_platform", return_value=jobs_client), + patch("nemo_platform_ext.cli.commands.api.jobs.render_job_watch_events", return_value=True) as render_events, + ): + watch_platform_job( + _ctx(client), + name="job-a", + workspace=None, + attempt_id=1, + step_id="step-1", + task_id="task-1", + limit=25, + timeout=42, + poll_interval=7, + include_history=False, + ) + + jobs_client.watch_job.assert_called_once_with( + "job-a", + workspace="default", + attempt_id=1, + step_id="step-1", + task_id="task-1", + limit=25, + timeout=42, + poll_interval=7, + include_history=False, + ) + render_events.assert_called_once_with(events, resource_label="job") def test_inference_deployment_create_exits_when_wait_fails() -> None: @@ -94,6 +358,7 @@ def test_inference_deployment_create_exits_when_wait_fails() -> None: workspace="test-workspace", config="deployment-config", wait=True, + watch=False, timeout=90, poll_interval=10, ) @@ -106,3 +371,76 @@ def test_inference_deployment_create_exits_when_wait_fails() -> None: timeout=90, poll_interval=10, ) + + +def test_inference_deployment_create_watch_uses_waiter() -> None: + deployments = MagicMock() + deployments.create.return_value = SimpleNamespace(name="deployment-a") + client = SimpleNamespace(inference=SimpleNamespace(deployments=deployments)) + ctx = _ctx(client) + + with ( + patch( + "nemo_platform_ext.cli.commands.api.inference.deployments.handle_code_generation", + return_value=False, + ) as handle_code_generation, + patch("nemo_platform_ext.cli.commands.api.inference.deployments.format_output") as format_output, + patch( + "nemo_platform_ext.cli.commands.api.inference.deployments.wait_for_inference_deployment", + return_value=True, + ) as wait_for_inference_deployment, + ): + create_deployments( + ctx, + name="deployment-a", + workspace="test-workspace", + config="deployment-config", + wait=False, + watch=True, + timeout=90, + poll_interval=10, + ) + + expected_kwargs = { + "workspace": "test-workspace", + "config": "deployment-config", + "name": "deployment-a", + } + handle_code_generation.assert_called_once_with( + ["inference", "deployments"], + "create", + expected_kwargs, + None, + ctx.obj, + watch_config={"type": "inference_deployment", "resource_label": "deployment"}, + watch_options={"timeout": 90, "poll_interval": 10}, + wait_config=None, + wait_options=None, + ) + deployments.create.assert_called_once_with(**expected_kwargs) + format_output.assert_called_once() + wait_for_inference_deployment.assert_called_once_with( + client, + "deployment-a", + workspace="test-workspace", + timeout=90, + poll_interval=10, + ) + + +def test_inference_deployment_create_rejects_wait_and_watch_together() -> None: + deployments = MagicMock() + client = SimpleNamespace(inference=SimpleNamespace(deployments=deployments)) + + with pytest.raises(SystemExit) as exc_info: + create_deployments( + _ctx(client), + name="deployment-a", + workspace="test-workspace", + config="deployment-config", + wait=True, + watch=True, + ) + + assert exc_info.value.code == 2 + deployments.create.assert_not_called() diff --git a/packages/nemo_platform_ext/tests/cli/core/test_code_generator.py b/packages/nemo_platform_ext/tests/cli/core/test_code_generator.py index 4b7de0c39d..672764e7b7 100644 --- a/packages/nemo_platform_ext/tests/cli/core/test_code_generator.py +++ b/packages/nemo_platform_ext/tests/cli/core/test_code_generator.py @@ -98,28 +98,90 @@ def test_generate_python_code_multiline_format(): assert any("namespace=" in line and line.strip().startswith("namespace=") for line in lines) -def test_generate_python_code_with_platform_job_wait(): +def test_generate_python_code_with_platform_job_watch(): code = generate_python_code( resource_path=["customization", "jobs"], method="create", args={"workspace": "default", "name": "job-a", "spec": {"training_type": "sft"}}, - wait_config={"type": "platform_job", "resource_label": "customization job"}, - wait_options={"timeout": 42, "poll_interval": 7}, + watch_config={"type": "platform_job", "resource_label": "customization job"}, + watch_options={"timeout": 42, "poll_interval": 7}, ) - assert "import time" in code + assert "import time" not in code + assert "from nemo_platform.jobs.watch import watch_job" not in code + assert "from nemo_platform_plugin.client.adapter import client_from_platform" in code + assert "from nemo_platform_plugin.jobs.client import JobsClient" in code + assert "jobs_client = client_from_platform(client, JobsClient)" in code assert "response = client.customization.jobs.create" in code assert 'resource_name = getattr(response, "name", None) or "job-a"' in code + assert 'raise RuntimeError("Unable to determine created resource name for --watch")' in code + assert "jobs_client.watch_job(" in code + assert 'workspace="default"' in code + assert "timeout=42" in code + assert "poll_interval=7" in code + assert "print(event)" in code + assert "get_status" not in code + assert "time.sleep" not in code + assert "print(response)" not in code + compile(code, "", "exec") + + +def test_generate_python_code_with_platform_job_watch_has_no_default_timeout(): + code = generate_python_code( + resource_path=["customization", "jobs"], + method="create", + args={"workspace": "default", "name": "job-a", "spec": {"training_type": "sft"}}, + watch_config={"type": "platform_job", "resource_label": "customization job"}, + watch_options={"poll_interval": 7}, + ) + + assert "timeout=None" in code + assert "deadline = time.monotonic()" not in code + assert "poll_interval=7" in code + compile(code, "", "exec") + + +def test_generate_python_code_with_platform_job_wait(): + code = generate_python_code( + resource_path=["jobs"], + method="create", + args={"workspace": "default", "name": "job-a", "spec": {"training_type": "sft"}}, + wait_config={"type": "platform_job", "resource_label": "job"}, + wait_options={"timeout": 42, "poll_interval": 7}, + ) + + assert "import time" not in code + for symbol in ("JobStatusEvent", "JobWatchTimeoutError", "JobsClient", "NeMoPlatform"): + assert symbol in code + assert "from nemo_platform_plugin.jobs.watch_types import JobStatusEvent, JobWatchTimeoutError" in code + assert "jobs_client = client_from_platform(client, JobsClient)" in code + assert "APIConnectionError" not in code + assert "APIStatusError" not in code + assert "APITimeoutError" not in code + assert "NotFoundError" not in code assert 'raise RuntimeError("Unable to determine created resource name for --wait")' in code - assert "deadline = time.monotonic() + 42" in code - assert 'client.customization.jobs.get_status(resource_name, workspace="default")' in code - assert 'status = str(status_response.status or "").lower()' in code - assert "response = status_response" in code - assert code.rindex("print(response)") > code.index("response = status_response") - assert "time.sleep(min(7, remaining))" in code + assert "deadline = time.monotonic()" not in code + assert "get_status" not in code + assert "jobs_client.watch_job(" in code + assert "include_logs=False" in code + assert "isinstance(event, JobStatusEvent)" in code + assert "except JobWatchTimeoutError as exc:" in code + assert "time.sleep" not in code + assert "print(response)" not in code compile(code, "", "exec") +def test_generate_python_code_with_platform_job_wait_requires_timeout(): + with pytest.raises(ValueError, match=r"wait 'platform_job' lifecycle code generation requires timeout"): + generate_python_code( + resource_path=["jobs"], + method="create", + args={"workspace": "default", "name": "job-a", "spec": {"training_type": "sft"}}, + wait_config={"type": "platform_job", "resource_label": "job"}, + wait_options={"poll_interval": 7}, + ) + + def test_generate_python_code_with_inference_deployment_wait(): code = generate_python_code( resource_path=["inference", "deployments"], @@ -149,21 +211,55 @@ def test_generate_python_code_with_inference_deployment_wait(): compile(code, "", "exec") -def test_generate_python_code_escapes_platform_job_wait_label(): +def test_generate_python_code_with_inference_deployment_watch(): + code = generate_python_code( + resource_path=["inference", "deployments"], + method="create", + args={"workspace": "default", "name": "deployment-a", "config": "deployment-config"}, + watch_config={"type": "inference_deployment", "resource_label": "deployment"}, + watch_options={"timeout": 90, "poll_interval": 10}, + ) + + assert "import time" in code + for symbol in ("APIConnectionError", "APIStatusError", "APITimeoutError", "NeMoPlatform", "NotFoundError"): + assert symbol in code + assert "deadline = time.monotonic() + 90" in code + assert 'resource_name = getattr(response, "name", None) or "deployment-a"' in code + assert 'raise RuntimeError("Unable to determine created resource name for --watch")' in code + assert 'client.inference.deployments.retrieve(resource_name, workspace="default")' in code + assert "client.inference.gateway.provider.ready(provider_name, workspace=provider_workspace)" in code + assert "response = deployment" in code + assert code.rindex("print(response)") > code.index("response = deployment") + assert "time.sleep(min(10, remaining))" in code + compile(code, "", "exec") + + +def test_generate_python_code_with_inference_deployment_wait_requires_timeout(): + with pytest.raises(ValueError, match=r"wait 'inference_deployment' lifecycle code generation requires timeout"): + generate_python_code( + resource_path=["inference", "deployments"], + method="create", + args={"workspace": "default", "name": "deployment-a", "config": "deployment-config"}, + wait_config={"type": "inference_deployment", "resource_label": "deployment"}, + wait_options={"poll_interval": 10}, + ) + + +def test_generate_python_code_with_platform_job_watch_ignores_label_formatting(): code = generate_python_code( resource_path=["customization", "jobs"], method="create", args={"workspace": "default", "name": "job-a"}, - wait_config={"type": "platform_job", "resource_label": 'customization "job" {label}'}, - wait_options={"timeout": 42, "poll_interval": 7}, + watch_config={"type": "platform_job", "resource_label": 'customization "job" {label}'}, + watch_options={"timeout": 42, "poll_interval": 7}, ) compile(code, "", "exec") - assert '"customization \\"job\\" {label}" + f" {resource_name!r}' in code + assert 'raise RuntimeError("Unable to determine created resource name for --watch")' in code -def test_generate_python_code_rejects_unknown_wait_type(): - with pytest.raises(ValueError, match="Unsupported wait config type: 'unknown'"): +def test_generate_python_code_rejects_unknown_lifecycle_type(): + with pytest.raises(ValueError, match="Unsupported lifecycle config type: 'unknown'"): generate_python_code( resource_path=["customization", "jobs"], method="create", diff --git a/packages/nemo_platform_ext/tests/cli/core/test_job_watch_renderer.py b/packages/nemo_platform_ext/tests/cli/core/test_job_watch_renderer.py new file mode 100644 index 0000000000..aaa1ada36c --- /dev/null +++ b/packages/nemo_platform_ext/tests/cli/core/test_job_watch_renderer.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Iterator +from datetime import datetime, timezone +from io import StringIO +from unittest.mock import patch + +from nemo_platform_ext.cli.core.job_watch_renderer import render_job_watch_events +from nemo_platform_ext.cli.telemetry.events import TaskStatusEnum +from nemo_platform_plugin.jobs.watch_types import JobLogEvent, JobStatusEvent, JobWatchEvent, JobWatchTimeoutError +from rich.console import Console + + +def _console_pair() -> tuple[Console, StringIO]: + output = StringIO() + return Console(file=output, force_terminal=False, color_system=None, width=120), output + + +def test_render_job_watch_events_returns_true_for_completed_status() -> None: + console, output = _console_pair() + error_console, error_output = _console_pair() + events: list[JobWatchEvent] = [ + JobStatusEvent( + kind="status", + job_name="job-a", + status="active", + status_details={"phase": "training", "progress_pct": 41}, + terminal=False, + successful=None, + ), + log_event := JobLogEvent( + kind="log", + job_name="job-a", + timestamp=datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc), + step_id="step-a", + task_id="task-a", + message="started", + ), + JobStatusEvent( + kind="status", + job_name="job-a", + status="completed", + status_details={}, + terminal=True, + successful=True, + ), + ] + + assert render_job_watch_events(events, console=console, error_console=error_console) is True + + rendered = output.getvalue() + expected_log_time = log_event.timestamp.astimezone().strftime("%H:%M:%S") + assert "Status: active phase=training progress_pct=41" in rendered + assert f"[{expected_log_time}] step-a/task-a | started" in rendered + assert "step-a/task-a | started" in rendered + assert "Job 'job-a' completed" in rendered + assert error_output.getvalue() == "" + + +def test_render_job_watch_events_emits_job_run_event_for_terminal_status() -> None: + console, _ = _console_pair() + event = JobStatusEvent( + kind="status", + job_name="job-a", + status="completed", + status_details={"model": "nemotron"}, + terminal=True, + successful=True, + ) + + with patch("nemo_platform_ext.cli.telemetry.emit.emit_event") as emit_event: + assert render_job_watch_events([event], console=console, resource_label="job", start_time=123.0) is True + + emit_event.assert_called_once() + telemetry_event = emit_event.call_args.args[0] + assert telemetry_event.job_type == "job" + assert telemetry_event.task_status is TaskStatusEnum.COMPLETED + assert telemetry_event.model == "defined" + + +def test_render_job_watch_events_returns_false_for_failed_terminal_status() -> None: + console, output = _console_pair() + event = JobStatusEvent( + kind="status", + job_name="job-a", + status="error", + status_details={}, + error_details={"reason": "container exited", "exit_code": 137, "empty": {}}, + terminal=True, + successful=False, + ) + + assert render_job_watch_events([event], console=console) is False + + assert "Job 'job-a' ended with status 'error': reason=container exited exit_code=137" in output.getvalue() + + +def test_render_job_watch_events_catches_timeout() -> None: + console, output = _console_pair() + error_console, error_output = _console_pair() + + def events() -> Iterator[JobWatchEvent]: + yield JobStatusEvent( + kind="status", + job_name="job-a", + status="active", + status_details={}, + terminal=False, + successful=None, + ) + raise JobWatchTimeoutError("Timed out watching job 'job-a'") + + assert render_job_watch_events(events(), console=console, error_console=error_console) is False + + assert "Status: active" in output.getvalue() + assert "Timed out watching job 'job-a'" in error_output.getvalue() diff --git a/packages/nemo_platform_ext/tests/cli/core/test_waiters.py b/packages/nemo_platform_ext/tests/cli/core/test_waiters.py index b819eb55aa..5cbb96b5f1 100644 --- a/packages/nemo_platform_ext/tests/cli/core/test_waiters.py +++ b/packages/nemo_platform_ext/tests/cli/core/test_waiters.py @@ -4,6 +4,7 @@ from __future__ import annotations from collections.abc import Iterator +from datetime import datetime, timezone from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -11,8 +12,11 @@ import pytest from nemo_platform import APIConnectionError, APIStatusError, AuthenticationError from nemo_platform_ext.cli.core import waiters +from nemo_platform_plugin.jobs.schemas import PlatformJobStatus, PlatformJobStatusResponse WAITERS_MODULE = "nemo_platform_ext.cli.core.waiters" +WATCH_MODULE = "nemo_platform_plugin.jobs.watch" +JOB_TIMESTAMP = datetime(2026, 1, 1, tzinfo=timezone.utc) class _DummyLive: @@ -35,6 +39,41 @@ def start(self) -> None: pass +class _RecordingLive(_DummyLive): + instances: list[_RecordingLive] = [] + + def __init__(self, renderable: object | None = None, *_args: object, **_kwargs: object) -> None: + self.renderable = renderable + self.updates: list[object] = [] + self.instances.append(self) + + def update(self, renderable: object, *_args: object, **_kwargs: object) -> None: + self.updates.append(renderable) + + +class _StatusResponse: + def __init__(self, status: PlatformJobStatusResponse) -> None: + self._status = status + + def data(self) -> PlatformJobStatusResponse: + return self._status + + +def _status_response(status: str | PlatformJobStatus) -> _StatusResponse: + return _StatusResponse( + PlatformJobStatusResponse( + id="job-a", + name="job-a", + status=PlatformJobStatus(status), + status_details={}, + error_details=None, + steps=[], + created_at=JOB_TIMESTAMP, + updated_at=JOB_TIMESTAMP, + ) + ) + + @pytest.fixture(autouse=True) def _quiet_rich_output() -> Iterator[None]: with ( @@ -71,22 +110,56 @@ def gateway_wait() -> Iterator[MagicMock]: def test_wait_for_platform_job_returns_true_on_completed(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="completed") + jobs.get_job_status.return_value = _status_response("completed") assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is True - jobs.get_status.assert_called_once_with("job-a", workspace="default") + jobs.get_job_status.assert_called_once_with(workspace="default", name="job-a") frozen_time.assert_called() def test_wait_for_platform_job_returns_false_on_error(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="error") + jobs.get_job_status.return_value = _status_response("error") assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is False frozen_time.assert_called() +def test_platform_job_wait_live_display_recomputes_elapsed() -> None: + display = waiters._PlatformJobWaitLiveDisplay(start_time=100.0, timeout=1200, poll_interval=3) + + with ( + patch(f"{WAITERS_MODULE}.datetime") as datetime_mock, + patch(f"{WAITERS_MODULE}.time.time", side_effect=[101.0, 109.0]), + ): + datetime_mock.now.return_value.strftime.return_value = "12:34:56" + + assert "Wait: 1s" in display.__rich__().plain + assert "Wait: 9s" in display.__rich__().plain + + +def test_wait_for_platform_job_uses_dynamic_live_display_for_unchanged_status_polls() -> None: + jobs = MagicMock() + jobs.get_job_status.side_effect = [ + _status_response("active"), + _status_response("active"), + _status_response("completed"), + ] + _RecordingLive.instances = [] + + with ( + patch(f"{WAITERS_MODULE}.Live", _RecordingLive), + patch(f"{WATCH_MODULE}.time.sleep"), + ): + assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is True + + live = _RecordingLive.instances[0] + assert isinstance(live.renderable, waiters._PlatformJobWaitLiveDisplay) + assert live.updates == [live.renderable, live.renderable] + assert jobs.get_job_status.call_count == 3 + + def test_wait_for_inference_deployment_uses_remaining_timeout_for_gateway(gateway_wait: MagicMock) -> None: client = MagicMock() client.inference.deployments.retrieve.return_value = SimpleNamespace( @@ -131,6 +204,27 @@ def test_wait_for_inference_deployment_uses_model_provider_id_for_gateway(gatewa assert gateway_wait.call_args.args[:3] == (client, "generated-provider", "provider-workspace") +def test_wait_for_inference_deployment_quiet_mode_uses_quiet_gateway(gateway_wait: MagicMock) -> None: + client = MagicMock() + client.inference.deployments.retrieve.return_value = SimpleNamespace( + status="READY", + status_message="", + status_history=[], + ) + + with patch(f"{WAITERS_MODULE}.time.time", side_effect=[100.0, 104.0, 104.0, 104.0]): + assert waiters.wait_for_inference_deployment( + client, + "deployment-a", + workspace="default", + timeout=10, + poll_interval=2, + verbose=False, + ) + + assert gateway_wait.call_args.kwargs["verbose"] is False + + def test_wait_for_inference_deployment_retries_transient_status_error( frozen_time: MagicMock, waiter_pause: MagicMock, gateway_wait: MagicMock ) -> None: @@ -204,29 +298,34 @@ def test_wait_for_inference_deployment_does_not_sleep_past_timeout(waiter_pause: waiter_pause.assert_called_once_with(1.0) -def test_wait_for_platform_job_does_not_sleep_past_timeout(waiter_pause: MagicMock) -> None: +def test_wait_for_platform_job_does_not_sleep_past_timeout() -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="running") + jobs.get_job_status.return_value = _status_response("active") - with patch(f"{WAITERS_MODULE}.time.time", side_effect=[0.0, 0.0, 0.0, 4.0, 5.0, 5.0]): + with ( + patch(f"{WAITERS_MODULE}.time.time", return_value=0.0), + patch(f"{WATCH_MODULE}.time.monotonic", side_effect=[0.0, 0.0, 4.0, 5.0]), + patch(f"{WATCH_MODULE}.time.sleep") as watch_sleep, + ): assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", timeout=5, poll_interval=10) is False - waiter_pause.assert_called_once_with(1.0) + watch_sleep.assert_called_once_with(1.0) -def test_wait_for_platform_job_retries_transient_status_error(frozen_time: MagicMock, waiter_pause: MagicMock) -> None: +def test_wait_for_platform_job_retries_transient_status_error(frozen_time: MagicMock) -> None: jobs = MagicMock() request = httpx.Request("GET", "http://test") response = httpx.Response(503, request=request) - jobs.get_status.side_effect = [ + jobs.get_job_status.side_effect = [ APIStatusError("service unavailable", response=response, body=None), - SimpleNamespace(status="completed"), + _status_response("completed"), ] - assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", timeout=10, poll_interval=1) is True + with patch(f"{WATCH_MODULE}.time.sleep") as watch_sleep: + assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", timeout=10, poll_interval=1) is True frozen_time.assert_called() - waiter_pause.assert_called_once_with(1) + watch_sleep.assert_called_once_with(1) def test_wait_for_gateway_does_not_sleep_past_timeout(waiter_pause: MagicMock) -> None: diff --git a/packages/nemo_platform_ext/tests/cli/telemetry/test_job_events.py b/packages/nemo_platform_ext/tests/cli/telemetry/test_job_events.py index e00cc180af..7279034fe6 100644 --- a/packages/nemo_platform_ext/tests/cli/telemetry/test_job_events.py +++ b/packages/nemo_platform_ext/tests/cli/telemetry/test_job_events.py @@ -5,15 +5,21 @@ from collections.abc import Iterator from datetime import datetime, timezone -from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from nemo_platform_ext.cli.core import waiters from nemo_platform_ext.cli.telemetry.events import TaskStatusEnum +from nemo_platform_plugin.jobs.schemas import ( + PlatformJobStatus, + PlatformJobStatusResponse, + PlatformJobStepStatusResponse, +) WAITERS_MODULE = "nemo_platform_ext.cli.core.waiters" +WATCH_MODULE = "nemo_platform_plugin.jobs.watch" EMIT_TARGET = "nemo_platform_ext.cli.telemetry.emit.emit_event" +JOB_TIMESTAMP = datetime(2026, 1, 1, tzinfo=timezone.utc) class _DummyLive: @@ -36,6 +42,48 @@ def start(self) -> None: pass +class _StatusResponse: + def __init__(self, status: PlatformJobStatusResponse) -> None: + self._status = status + + def data(self) -> PlatformJobStatusResponse: + return self._status + + +def _step(name: str) -> PlatformJobStepStatusResponse: + return PlatformJobStepStatusResponse( + id=name, + name=name, + status=PlatformJobStatus.COMPLETED, + status_details={}, + error_details=None, + tasks=[], + created_at=JOB_TIMESTAMP, + updated_at=JOB_TIMESTAMP, + ) + + +def _status_response( + status: str | PlatformJobStatus, + *, + steps: list[str] | None = None, + status_details: dict[str, object] | None = None, + created_at: datetime = JOB_TIMESTAMP, +) -> _StatusResponse: + return _StatusResponse( + PlatformJobStatusResponse( + id="job-a", + name="job-a", + status=PlatformJobStatus(status), + status_details=status_details or {}, + error_details=None, + steps=[_step(step) for step in steps or []], + created_at=created_at, + updated_at=created_at, + ) + ) + + @pytest.fixture(autouse=True) def _quiet_rich_output() -> Iterator[None]: with ( @@ -57,14 +105,14 @@ def waiter_pause() -> Iterator[MagicMock]: yield pause -def _completed_status() -> SimpleNamespace: - return SimpleNamespace( - status="completed", +def _completed_status() -> _StatusResponse: + return _status_response( + "completed", steps=[ - SimpleNamespace(name="audit-job"), - SimpleNamespace(name="evaluate"), - SimpleNamespace(name="evaluate-suite"), - SimpleNamespace(name="customer-project-step"), + "audit-job", + "evaluate", + "evaluate-suite", + "customer-project-step", ], status_details={"input_tokens": 512, "output_tokens": 2048, "model": "nemotron-super-49b"}, ) @@ -72,7 +120,7 @@ def _completed_status() -> SimpleNamespace: def test_completed_emits_single_job_run_event(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = _completed_status() + jobs.get_job_status.return_value = _completed_status() with patch(EMIT_TARGET) as emit_event: assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", resource_label="customization") is True @@ -90,9 +138,7 @@ def test_completed_emits_single_job_run_event(frozen_time: MagicMock) -> None: def test_static_step_name_is_not_emitted_as_job_type(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace( - status="completed", steps=[SimpleNamespace(name="audit-job")], status_details={} - ) + jobs.get_job_status.return_value = _status_response("completed", steps=["audit-job"]) with patch(EMIT_TARGET) as emit_event: assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", resource_label="audit") is True @@ -104,7 +150,7 @@ def test_static_step_name_is_not_emitted_as_job_type(frozen_time: MagicMock) -> def test_job_type_falls_back_to_resource_label_without_steps(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="completed", steps=[], status_details={}) + jobs.get_job_status.return_value = _status_response("completed") with patch(EMIT_TARGET) as emit_event: assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", resource_label="customization") is True @@ -116,11 +162,7 @@ def test_job_type_falls_back_to_resource_label_without_steps(frozen_time: MagicM def test_unsafe_resource_label_falls_back_to_custom(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace( - status="completed", - steps=[SimpleNamespace(name="private-customer-step")], - status_details={}, - ) + jobs.get_job_status.return_value = _status_response("completed", steps=["private-customer-step"]) with patch(EMIT_TARGET) as emit_event: assert ( @@ -135,7 +177,7 @@ def test_unsafe_resource_label_falls_back_to_custom(frozen_time: MagicMock) -> N def test_status_details_defaults_when_absent(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="completed", steps=[], status_details={}) + jobs.get_job_status.return_value = _status_response("completed") with patch(EMIT_TARGET) as emit_event: assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is True @@ -150,9 +192,8 @@ def test_status_details_defaults_when_absent(frozen_time: MagicMock) -> None: def test_null_status_details_still_emits(frozen_time: MagicMock) -> None: """Explicit nulls must not drop the event; a real 0 token count must survive.""" jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace( - status="completed", - steps=[], + jobs.get_job_status.return_value = _status_response( + "completed", status_details={"model": None, "input_tokens": 0, "output_tokens": None}, ) @@ -168,9 +209,8 @@ def test_null_status_details_still_emits(frozen_time: MagicMock) -> None: def test_non_string_model_details_still_emit_safe_bucket(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace( - status="completed", - steps=[], + jobs.get_job_status.return_value = _status_response( + "completed", status_details={"model": {"name": "private-model"}, "input_tokens": 7, "output_tokens": 9}, ) @@ -186,15 +226,13 @@ def test_non_string_model_details_still_emit_safe_bucket(frozen_time: MagicMock) def test_duration_uses_job_created_at_when_available() -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace( - status="completed", - steps=[], - status_details={}, + jobs.get_job_status.return_value = _status_response( + "completed", created_at=datetime.fromtimestamp(90.0, tz=timezone.utc), ) with ( - patch(f"{WAITERS_MODULE}.time.time", side_effect=[100.0, 100.0, 100.0, 130.0]), + patch(f"{WAITERS_MODULE}.time.time", side_effect=[100.0, 100.0, 130.0]), patch(EMIT_TARGET) as emit_event, ): assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is True @@ -205,7 +243,7 @@ def test_duration_uses_job_created_at_when_available() -> None: def test_error_status_maps_to_error(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="error", steps=[], status_details={}) + jobs.get_job_status.return_value = _status_response("error") with patch(EMIT_TARGET) as emit_event: assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is False @@ -216,7 +254,7 @@ def test_error_status_maps_to_error(frozen_time: MagicMock) -> None: def test_cancelled_status_maps_to_canceled(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="cancelled", steps=[], status_details={}) + jobs.get_job_status.return_value = _status_response("cancelled") with patch(EMIT_TARGET) as emit_event: assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is False @@ -225,11 +263,16 @@ def test_cancelled_status_maps_to_canceled(frozen_time: MagicMock) -> None: assert emit_event.call_args.args[0].task_status is TaskStatusEnum.CANCELED -def test_timeout_emits_nothing_and_does_not_crash(waiter_pause: MagicMock) -> None: +def test_timeout_emits_nothing_and_does_not_crash() -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="running", steps=[], status_details={}) + jobs.get_job_status.return_value = _status_response("active") - with patch(EMIT_TARGET) as emit_event: + with ( + patch(f"{WAITERS_MODULE}.time.time", return_value=0.0), + patch(f"{WATCH_MODULE}.time.monotonic", side_effect=[0.0, 0.0, 4.0, 5.0]), + patch(f"{WATCH_MODULE}.time.sleep"), + patch(EMIT_TARGET) as emit_event, + ): assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", timeout=5, poll_interval=10) is False emit_event.assert_not_called() diff --git a/packages/nemo_platform_ext/tests/cli/test_app.py b/packages/nemo_platform_ext/tests/cli/test_app.py index 91194e7cb3..d4760b6e5e 100644 --- a/packages/nemo_platform_ext/tests/cli/test_app.py +++ b/packages/nemo_platform_ext/tests/cli/test_app.py @@ -267,6 +267,36 @@ def test_members_api_command_is_not_registered_at_top_level(): assert "nemo_platform_ext.cli.commands.api.members" not in sys.modules +def test_jobs_watch_command_is_registered(): + runner = CliRunner() + result = runner.invoke(app, ["jobs", "watch", "--help"]) + + assert result.exit_code == 0 + assert "Watch a platform job until it reaches a terminal status." in result.stdout + assert "--history" in result.stdout + assert "--no-history" in result.stdout + + +def test_jobs_create_exposes_wait_and_watch_flags(): + runner = CliRunner() + result = runner.invoke(app, ["jobs", "create", "--help"]) + + assert result.exit_code == 0 + assert "--watch" in result.stdout + assert "--wait" in result.stdout + + +def test_inference_deployments_create_exposes_wait_and_watch_flags(): + runner = CliRunner() + result = runner.invoke(app, ["inference", "deployments", "create", "--help"]) + + assert result.exit_code == 0 + assert "--watch" in result.stdout + assert "--wait" in result.stdout + assert "up and running" in result.stdout + assert "until it is stable" in result.stdout + + def test_root_help_excludes_hidden_commands_and_context_option(): runner = CliRunner() result = runner.invoke(app, ["--help"]) diff --git a/packages/nemo_platform_ext/tests/jobs/test_watch.py b/packages/nemo_platform_ext/tests/jobs/test_watch.py new file mode 100644 index 0000000000..0df79a2cdd --- /dev/null +++ b/packages/nemo_platform_ext/tests/jobs/test_watch.py @@ -0,0 +1,974 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections import deque +from collections.abc import AsyncIterator, Iterable +from datetime import datetime, timezone +from typing import TypedDict, TypeVar + +import httpx +import pytest +from nemo_platform import APIStatusError +from nemo_platform_plugin.client.errors import NemoHTTPError, NemoTransportError +from nemo_platform_plugin.client.response import AsyncNemoPaginatedResponse, NemoPaginatedResponse, NemoResponse +from nemo_platform_plugin.client.types import CursorPagination, PreparedRequest +from nemo_platform_plugin.jobs import watch as watch_module +from nemo_platform_plugin.jobs.client import AsyncJobsClient, JobsClient +from nemo_platform_plugin.jobs.schemas import PlatformJobLog, PlatformJobStatus, PlatformJobStatusResponse +from nemo_platform_plugin.jobs.types import JobLogsQueryParams +from nemo_platform_plugin.jobs.watch import async_watch_job, watch_job +from nemo_platform_plugin.jobs.watch_types import ( + JobLogEvent, + JobStatusEvent, + JobWarningEvent, + JobWatchTimeoutError, +) + +ResponseT = TypeVar("ResponseT") + + +class _StatusCall(TypedDict): + workspace: str | None + name: str + + +class _LogCall(TypedDict): + workspace: str | None + name: str + query_params: JobLogsQueryParams | None + + +def _prepared_request(response_type: type[ResponseT] | None = None) -> PreparedRequest[ResponseT]: + return PreparedRequest( + path_template="/test", + path_params={}, + method="GET", + content=None, + content_type=None, + response_type=response_type, + ) + + +def _status_response(body: PlatformJobStatusResponse) -> NemoResponse[PlatformJobStatusResponse]: + return NemoResponse( + http_response=httpx.Response(200), + body=body, + request=_prepared_request(PlatformJobStatusResponse), + ) + + +def _warning_message(event: JobStatusEvent | JobLogEvent | JobWarningEvent) -> str: + assert isinstance(event, JobWarningEvent) + return event.message + + +def _page_cursor(call: _LogCall) -> str | None: + query_params = call["query_params"] + if query_params is None: + return None + return query_params.get("page_cursor") + + +class _PageResponse( + NemoPaginatedResponse[PlatformJobLog, CursorPagination], + AsyncNemoPaginatedResponse[PlatformJobLog, CursorPagination], +): + def __init__(self, items: list[PlatformJobLog], next_page: str | None = None) -> None: + response = httpx.Response( + 200, + json={ + "data": [item.model_dump(mode="json") for item in items], + "total": len(items), + "next_page": next_page, + "prev_page": None, + }, + ) + super().__init__( + first_http_response=response, + model_type=PlatformJobLog, + request=_prepared_request(), + fetch_page=_unexpected_page_fetch, + strategy=CursorPagination, + ) + + +def _unexpected_page_fetch(_request: PreparedRequest[object], _page: object) -> httpx.Response: + raise AssertionError("Unexpected paginated fetch") + + +class _JobsClientState: + def __init__( + self, + *, + statuses: Iterable[PlatformJobStatusResponse | Exception], + log_results: Iterable[_PageResponse | Exception], + ) -> None: + self._statuses = deque(statuses) + self._last_status: PlatformJobStatusResponse | None = None + self._log_results = deque(log_results) + self.status_calls: list[_StatusCall] = [] + self.log_calls: list[_LogCall] = [] + + def _next_status(self, *, workspace: str | None, name: str) -> PlatformJobStatusResponse: + self.status_calls.append({"workspace": workspace, "name": name}) + if self._statuses: + result = self._statuses.popleft() + if isinstance(result, Exception): + raise result + self._last_status = result + if self._last_status is None: + raise AssertionError("No status result configured") + return self._last_status + + def _next_logs( + self, + *, + workspace: str | None, + name: str, + query_params: JobLogsQueryParams | None, + ) -> _PageResponse: + self.log_calls.append({"workspace": workspace, "name": name, "query_params": query_params}) + if not self._log_results: + return _PageResponse([]) + result = self._log_results.popleft() + if isinstance(result, Exception): + raise result + return result + + +class _SyncJobsClient(_JobsClientState): + def get_job_status(self, *, workspace: str | None = None, name: str) -> NemoResponse[PlatformJobStatusResponse]: + return _status_response(self._next_status(workspace=workspace, name=name)) + + def list_job_logs( + self, + *, + workspace: str | None = None, + name: str, + query_params: JobLogsQueryParams | None = None, + ) -> NemoPaginatedResponse[PlatformJobLog, CursorPagination]: + return self._next_logs(workspace=workspace, name=name, query_params=query_params) + + +class _AsyncJobsClient(_JobsClientState): + async def get_job_status( + self, + *, + workspace: str | None = None, + name: str, + ) -> NemoResponse[PlatformJobStatusResponse]: + return _status_response(self._next_status(workspace=workspace, name=name)) + + async def list_job_logs( + self, + *, + workspace: str | None = None, + name: str, + query_params: JobLogsQueryParams | None = None, + ) -> AsyncNemoPaginatedResponse[PlatformJobLog, CursorPagination]: + return self._next_logs(workspace=workspace, name=name, query_params=query_params) + + +def _status(status: str, status_details: dict[str, object] | None = None) -> PlatformJobStatusResponse: + timestamp = datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc) + return PlatformJobStatusResponse( + id="job-a", + name="job-a", + status=PlatformJobStatus(status), + status_details=status_details or {}, + error_details=None, + steps=[], + created_at=timestamp, + updated_at=timestamp, + ) + + +def _log( + message: str, + *, + timestamp: datetime | None = None, + job_step: str = "step-a", + job_task: str = "task-a", +) -> PlatformJobLog: + return PlatformJobLog( + job="job-a", + timestamp=timestamp or datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc), + job_step=job_step, + job_task=job_task, + message=message, + ) + + +def _record_completed_log_drain( + state: watch_module._WatchState, + logs: list[PlatformJobLog], +) -> list[JobLogEvent]: + state.start_log_drain() + events = list( + watch_module._new_log_events( + logs, + state=state, + occurrence_counts={}, + name="job-a", + emit=True, + ) + ) + state.complete_log_drain() + return events + + +def _invalid_page_cursor_error() -> NemoHTTPError: + request = httpx.Request("GET", "http://test") + response = httpx.Response(422, request=request, json={"detail": "Invalid page cursor"}) + return NemoHTTPError(response) + + +def _http_error_body(status_code: int, body: object) -> NemoHTTPError: + request = httpx.Request("GET", "http://test") + response = httpx.Response(status_code, request=request, json=body) + return NemoHTTPError(response) + + +def _http_text_error(status_code: int, text: str) -> NemoHTTPError: + request = httpx.Request("GET", "http://test") + response = httpx.Response(status_code, request=request, text=text) + return NemoHTTPError(response) + + +def _http_error(status_code: int, detail: str) -> NemoHTTPError: + request = httpx.Request("GET", "http://test") + response = httpx.Response(status_code, request=request, json={"detail": detail}) + return NemoHTTPError(response) + + +def test_can_retry_log_scan_from_start_accepts_invalid_cursor_code() -> None: + exc = _http_error_body( + 422, + {"detail": {"code": "invalid_page_cursor", "message": "The saved cursor is no longer valid"}}, + ) + + assert watch_module._can_retry_log_scan_from_start(exc, "cursor-0") is True + + +def test_can_retry_log_scan_from_start_keeps_page_cursor_decode_compatibility() -> None: + assert watch_module._can_retry_log_scan_from_start(_invalid_page_cursor_error(), "cursor-0") is True + + +def test_can_retry_log_scan_from_start_requires_saved_cursor() -> None: + assert watch_module._can_retry_log_scan_from_start(_invalid_page_cursor_error(), None) is False + + +def test_can_retry_log_scan_from_start_ignores_plain_text_detail_fallback() -> None: + exc = _http_text_error(422, "Invalid page cursor") + + assert watch_module._can_retry_log_scan_from_start(exc, "cursor-0") is False + + +def test_can_retry_log_scan_from_start_rejects_other_422_errors() -> None: + exc = _http_error(422, "Invalid page size") + + assert watch_module._can_retry_log_scan_from_start(exc, "cursor-0") is False + + +def test_watch_job_yields_status_logs_terminal_and_passes_log_query_params() -> None: + client = _SyncJobsClient( + statuses=[ + _status("active", {"phase": "training"}), + _status("completed", {"phase": "done"}), + ], + log_results=[ + _PageResponse([_log("starting")], next_page="cursor-1"), + _PageResponse([_log("still running")]), + _PageResponse([_log("starting"), _log("still running"), _log("done")]), + ], + ) + + events = list( + watch_job( + client, + "job-a", + workspace="default", + poll_interval=0, + attempt_id=1, + step_id="step-1", + task_id="task-1", + limit=2, + page_cursor="cursor-0", + ) + ) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("log", None, "starting"), + ("log", None, "still running"), + ("status", "completed", None), + ("log", None, "done"), + ] + assert isinstance(events[0], JobStatusEvent) + assert events[0].terminal is False + assert events[0].successful is None + assert events[0].status_details == {"phase": "training"} + assert isinstance(events[3], JobStatusEvent) + assert events[3].terminal is True + assert events[3].successful is True + + assert client.status_calls == [ + {"workspace": "default", "name": "job-a"}, + {"workspace": "default", "name": "job-a"}, + ] + assert client.log_calls == [ + { + "workspace": "default", + "name": "job-a", + "query_params": { + "attempt_id": 1, + "step_id": "step-1", + "task_id": "task-1", + "limit": 2, + "page_cursor": "cursor-0", + }, + }, + { + "workspace": "default", + "name": "job-a", + "query_params": { + "attempt_id": 1, + "step_id": "step-1", + "task_id": "task-1", + "limit": 2, + "page_cursor": "cursor-1", + }, + }, + { + "workspace": "default", + "name": "job-a", + "query_params": { + "attempt_id": 1, + "step_id": "step-1", + "task_id": "task-1", + "limit": 2, + "page_cursor": "cursor-1", + }, + }, + ] + + +def test_watch_job_rejects_negative_poll_interval_eagerly() -> None: + client = _SyncJobsClient(statuses=[], log_results=[]) + + with pytest.raises(ValueError, match="poll_interval"): + watch_job(client, "job-a", poll_interval=-1) + + +def test_jobs_client_watch_job_delegates_to_source_owned_watcher(monkeypatch: pytest.MonkeyPatch) -> None: + calls: dict[str, object] = {} + + def fake_watch_job(client: JobsClient, name: str, **kwargs: object) -> Iterable[JobWarningEvent]: + calls["client"] = client + calls["name"] = name + calls["kwargs"] = kwargs + return iter([JobWarningEvent(kind="warning", job_name=name, message="delegated")]) + + monkeypatch.setattr("nemo_platform_plugin.jobs.watch.watch_job", fake_watch_job) + http_client = httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request))) + client = JobsClient(base_url="http://test", http_client=http_client) + + try: + events = list( + client.watch_job( + "job-a", + workspace="default", + poll_interval=0, + timeout=5, + include_history=False, + include_logs=False, + attempt_id=1, + step_id="step-1", + task_id="task-1", + limit=2, + page_cursor="cursor-0", + ) + ) + finally: + http_client.close() + + assert calls == { + "client": client, + "name": "job-a", + "kwargs": { + "workspace": "default", + "poll_interval": 0, + "timeout": 5, + "include_history": False, + "include_logs": False, + "attempt_id": 1, + "step_id": "step-1", + "task_id": "task-1", + "limit": 2, + "page_cursor": "cursor-0", + }, + } + assert events == [JobWarningEvent(kind="warning", job_name="job-a", message="delegated")] + + +def test_watch_job_can_skip_existing_log_history() -> None: + client = _SyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _PageResponse([_log("old")]), + _PageResponse([_log("old"), _log("new")]), + ], + ) + + events = list(watch_job(client, "job-a", include_history=False, poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("status", "completed", None), + ("log", None, "new"), + ] + + +def test_watch_state_replaces_retained_log_window_after_successful_drain() -> None: + state = watch_module._WatchState(history_seen=True, log_cursor=None) + first_log = _log("first") + second_log = _log("second") + + assert [event.message for event in _record_completed_log_drain(state, [first_log])] == ["first"] + assert state.previous_drain_seen_logs == {watch_module._log_key(first_log): 1} + + assert [event.message for event in _record_completed_log_drain(state, [second_log])] == ["second"] + assert state.previous_drain_seen_logs == {watch_module._log_key(second_log): 1} + assert state.current_drain_seen_logs == {} + + +def test_new_log_events_suppresses_seen_occurrences_after_partial_drain_failure() -> None: + state = watch_module._WatchState(history_seen=True, log_cursor=None) + duplicate_log = _log("duplicate") + + state.start_log_drain() + first_events = list( + watch_module._new_log_events( + [duplicate_log, duplicate_log], + state=state, + occurrence_counts={}, + name="job-a", + emit=True, + ) + ) + + state.start_log_drain() + retry_events = list( + watch_module._new_log_events( + [duplicate_log, duplicate_log, duplicate_log], + state=state, + occurrence_counts={}, + name="job-a", + emit=True, + ) + ) + state.complete_log_drain() + + assert [event.message for event in first_events] == ["duplicate", "duplicate"] + assert [event.message for event in retry_events] == ["duplicate"] + assert state.previous_drain_seen_logs == {watch_module._log_key(duplicate_log): 3} + + +def test_watch_job_can_poll_status_without_logs() -> None: + client = _SyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[AssertionError("logs should not be fetched")], + ) + + events = list(watch_job(client, "job-a", include_logs=False, poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None)) for event in events] == [ + ("status", "active"), + ("status", "completed"), + ] + assert client.log_calls == [] + + +def test_watch_job_stops_when_status_is_paused() -> None: + client = _SyncJobsClient( + statuses=[_status("paused")], + log_results=[_PageResponse([_log("paused")])], + ) + + events = list(watch_job(client, "job-a", poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "paused", None), + ("log", None, "paused"), + ] + assert isinstance(events[0], JobStatusEvent) + assert events[0].terminal is True + assert events[0].successful is False + assert len(client.status_calls) == 1 + assert len(client.log_calls) == 1 + + +def test_watch_job_retries_sdk_transient_status_error() -> None: + request = httpx.Request("GET", "http://test") + response = httpx.Response(503, request=request) + client = _SyncJobsClient( + statuses=[ + APIStatusError("service unavailable", response=response, body=None), + _status("completed"), + ], + log_results=[AssertionError("logs should not be fetched")], + ) + + events = list(watch_job(client, "job-a", include_logs=False, poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("warning", None, "Transient status check failed: service unavailable"), + ("status", "completed", None), + ] + assert client.log_calls == [] + + +def test_watch_job_backs_off_and_deduplicates_consecutive_transient_status_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + request = httpx.Request("GET", "http://test") + response = httpx.Response(503, request=request) + client = _SyncJobsClient( + statuses=[ + APIStatusError("service unavailable", response=response, body=None), + APIStatusError("service unavailable", response=response, body=None), + APIStatusError("service unavailable", response=response, body=None), + APIStatusError("service unavailable", response=response, body=None), + APIStatusError("service unavailable", response=response, body=None), + APIStatusError("service unavailable", response=response, body=None), + APIStatusError("service unavailable", response=response, body=None), + _status("completed"), + ], + log_results=[AssertionError("logs should not be fetched")], + ) + sleeps: list[float] = [] + monkeypatch.setattr(watch_module.time, "sleep", sleeps.append) + + events = list(watch_job(client, "job-a", include_logs=False, poll_interval=1)) + + assert sleeps == [1, 2, 4, 8, 16, 30.0, 30.0] + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("warning", None, "Transient status check failed: service unavailable"), + ("status", "completed", None), + ] + assert client.log_calls == [] + + +def test_watch_job_suppresses_unread_history_after_partial_history_drain_failure() -> None: + client = _SyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _PageResponse([_log("old")], next_page="cursor-1"), + NemoTransportError(httpx.TransportError("temporary log failure")), + _PageResponse([_log("old"), _log("also-old")]), + ], + ) + + events = list(watch_job(client, "job-a", include_history=False, poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("warning", None, "Transient log check failed: temporary log failure"), + ("status", "completed", None), + ] + + +def test_watch_job_continues_status_polling_after_non_retryable_log_http_error() -> None: + client = _SyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _http_error(500, "log store unavailable"), + _PageResponse([_log("done")]), + ], + ) + + events = list(watch_job(client, "job-a", poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("warning", None, "Log check failed: HTTP 500: log store unavailable"), + ("status", "completed", None), + ("log", None, "done"), + ] + + +def test_watch_job_retries_terminal_status_until_logs_drain_after_transient_failure() -> None: + client = _SyncJobsClient( + statuses=[_status("completed"), _status("completed")], + log_results=[ + NemoTransportError(httpx.TransportError("temporary log failure")), + _PageResponse([_log("done")]), + ], + ) + + events = list(watch_job(client, "job-a", poll_interval=0)) + + assert [event.kind for event in events] == ["status", "warning", "log"] + assert _warning_message(events[1]) == "Transient log check failed: temporary log failure" + assert isinstance(events[2], JobLogEvent) + assert events[2].message == "done" + + +def test_watch_job_stops_after_terminal_log_drain_retry_cap() -> None: + retry_cap = watch_module._TerminalLogDrainRetryBudget.RETRY_CAP + client = _SyncJobsClient( + statuses=[_status("completed")], + log_results=[ + NemoTransportError(httpx.TransportError(f"temporary log failure {attempt}")) for attempt in range(retry_cap) + ], + ) + + events = list(watch_job(client, "job-a", poll_interval=0)) + + assert [event.kind for event in events] == ["status"] + ["warning"] * (retry_cap + 1) + assert [getattr(event, "message", None) for event in events[1:-1]] == [ + f"Transient log check failed: temporary log failure {attempt}" for attempt in range(retry_cap) + ] + assert _warning_message(events[-1]) == ( + f"Terminal log drain retry cap reached ({retry_cap}); stopping watch for job 'job-a'" + ) + assert len(client.status_calls) == retry_cap + assert len(client.log_calls) == retry_cap + + +def test_watch_job_falls_back_to_full_rescan_when_saved_cursor_is_invalid() -> None: + client = _SyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _PageResponse([_log("old")], next_page="cursor-1"), + _PageResponse([_log("new")]), + _invalid_page_cursor_error(), + _PageResponse([_log("old"), _log("new"), _log("done")]), + ], + ) + + events = list(watch_job(client, "job-a", poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("log", None, "old"), + ("log", None, "new"), + ("status", "completed", None), + ("log", None, "done"), + ] + assert [_page_cursor(call) for call in client.log_calls] == [None, "cursor-1", "cursor-1", None] + + +def test_watch_job_stops_log_pagination_when_cursor_does_not_advance() -> None: + client = _SyncJobsClient( + statuses=[_status("completed")], + log_results=[_PageResponse([_log("done")], next_page="cursor-0")], + ) + + events = list(watch_job(client, "job-a", poll_interval=0, page_cursor="cursor-0")) + + assert [(event.kind, getattr(event, "message", None)) for event in events] == [ + ("status", None), + ("log", "done"), + ] + assert [_page_cursor(call) for call in client.log_calls] == ["cursor-0"] + + +def test_watch_job_enforces_timeout_between_log_pages(monkeypatch: pytest.MonkeyPatch) -> None: + client = _SyncJobsClient( + statuses=[_status("completed")], + log_results=[ + _PageResponse([], next_page="cursor-1"), + _PageResponse([]), + ], + ) + monotonic_values = iter([0.0, 0.0, 1.0, 10.0]) + monkeypatch.setattr(watch_module.time, "monotonic", lambda: next(monotonic_values, 10.0)) + + with pytest.raises(JobWatchTimeoutError, match="job-a"): + list(watch_job(client, "job-a", timeout=5, poll_interval=0)) + + assert len(client.log_calls) == 1 + + +async def test_async_watch_job_suppresses_unread_history_after_partial_history_drain_failure() -> None: + client = _AsyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _PageResponse([_log("old")], next_page="cursor-1"), + NemoTransportError(httpx.TransportError("temporary log failure")), + _PageResponse([_log("old"), _log("also-old")]), + ], + ) + + events = [event async for event in async_watch_job(client, "job-a", include_history=False, poll_interval=0)] + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("warning", None, "Transient log check failed: temporary log failure"), + ("status", "completed", None), + ] + + +async def test_async_watch_job_backs_off_and_deduplicates_consecutive_transient_log_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = _AsyncJobsClient( + statuses=[_status("active"), _status("active"), _status("completed")], + log_results=[ + NemoTransportError(httpx.TransportError("temporary log failure")), + NemoTransportError(httpx.TransportError("temporary log failure")), + _PageResponse([_log("done")]), + ], + ) + sleeps: list[float] = [] + + async def fake_sleep(sleep_for: float) -> None: + sleeps.append(sleep_for) + + monkeypatch.setattr(watch_module.asyncio, "sleep", fake_sleep) + + events = [event async for event in async_watch_job(client, "job-a", poll_interval=1)] + + assert sleeps == [1, 2] + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("warning", None, "Transient log check failed: temporary log failure"), + ("status", "completed", None), + ("log", None, "done"), + ] + + +async def test_async_watch_job_continues_status_polling_after_non_retryable_log_http_error() -> None: + client = _AsyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _http_error(500, "log store unavailable"), + _PageResponse([_log("done")]), + ], + ) + + events = [event async for event in async_watch_job(client, "job-a", poll_interval=0)] + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("warning", None, "Log check failed: HTTP 500: log store unavailable"), + ("status", "completed", None), + ("log", None, "done"), + ] + + +async def test_async_watch_job_can_poll_status_without_logs() -> None: + client = _AsyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[AssertionError("logs should not be fetched")], + ) + + events = [event async for event in async_watch_job(client, "job-a", include_logs=False, poll_interval=0)] + + assert [(event.kind, getattr(event, "status", None)) for event in events] == [ + ("status", "active"), + ("status", "completed"), + ] + assert client.log_calls == [] + + +async def test_async_watch_job_stops_when_status_is_paused() -> None: + client = _AsyncJobsClient( + statuses=[_status("paused")], + log_results=[_PageResponse([_log("paused")])], + ) + + events = [event async for event in async_watch_job(client, "job-a", poll_interval=0)] + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "paused", None), + ("log", None, "paused"), + ] + assert isinstance(events[0], JobStatusEvent) + assert events[0].terminal is True + assert events[0].successful is False + assert len(client.status_calls) == 1 + assert len(client.log_calls) == 1 + + +async def test_async_watch_job_falls_back_to_full_rescan_when_saved_cursor_is_invalid() -> None: + client = _AsyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _PageResponse([_log("old")], next_page="cursor-1"), + _PageResponse([_log("new")]), + _invalid_page_cursor_error(), + _PageResponse([_log("old"), _log("new"), _log("done")]), + ], + ) + + events = [event async for event in async_watch_job(client, "job-a", poll_interval=0)] + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("log", None, "old"), + ("log", None, "new"), + ("status", "completed", None), + ("log", None, "done"), + ] + assert [_page_cursor(call) for call in client.log_calls] == [None, "cursor-1", "cursor-1", None] + + +async def test_async_watch_job_stops_log_pagination_when_cursor_repeats() -> None: + client = _AsyncJobsClient( + statuses=[_status("completed")], + log_results=[ + _PageResponse([_log("one")], next_page="cursor-1"), + _PageResponse([_log("two")], next_page="cursor-0"), + ], + ) + + events = [event async for event in async_watch_job(client, "job-a", poll_interval=0, page_cursor="cursor-0")] + + assert [(event.kind, getattr(event, "message", None)) for event in events] == [ + ("status", None), + ("log", "one"), + ("log", "two"), + ] + assert [_page_cursor(call) for call in client.log_calls] == ["cursor-0", "cursor-1"] + + +async def test_async_watch_job_enforces_timeout_between_log_pages(monkeypatch: pytest.MonkeyPatch) -> None: + client = _AsyncJobsClient( + statuses=[_status("completed")], + log_results=[ + _PageResponse([], next_page="cursor-1"), + _PageResponse([]), + ], + ) + monotonic_values = iter([0.0, 0.0, 1.0, 10.0]) + monkeypatch.setattr(watch_module.time, "monotonic", lambda: next(monotonic_values, 10.0)) + + with pytest.raises(JobWatchTimeoutError, match="job-a"): + [event async for event in async_watch_job(client, "job-a", timeout=5, poll_interval=0)] + + assert len(client.log_calls) == 1 + + +async def test_async_watch_job_stops_after_terminal_log_drain_retry_cap() -> None: + retry_cap = watch_module._TerminalLogDrainRetryBudget.RETRY_CAP + client = _AsyncJobsClient( + statuses=[_status("completed")], + log_results=[ + NemoTransportError(httpx.TransportError(f"temporary log failure {attempt}")) for attempt in range(retry_cap) + ], + ) + + events = [event async for event in async_watch_job(client, "job-a", poll_interval=0)] + + assert [event.kind for event in events] == ["status"] + ["warning"] * (retry_cap + 1) + assert [getattr(event, "message", None) for event in events[1:-1]] == [ + f"Transient log check failed: temporary log failure {attempt}" for attempt in range(retry_cap) + ] + assert _warning_message(events[-1]) == ( + f"Terminal log drain retry cap reached ({retry_cap}); stopping watch for job 'job-a'" + ) + assert len(client.status_calls) == retry_cap + assert len(client.log_calls) == retry_cap + + +def test_watch_job_raises_timeout_with_job_name(monkeypatch: pytest.MonkeyPatch) -> None: + client = _SyncJobsClient(statuses=[_status("active")], log_results=[_PageResponse([])]) + monotonic_values = iter([0.0, 0.0, 10.0]) + monkeypatch.setattr(watch_module.time, "monotonic", lambda: next(monotonic_values, 10.0)) + + with pytest.raises(JobWatchTimeoutError, match="job-a"): + list(watch_job(client, "job-a", timeout=5, poll_interval=1)) + + +async def test_async_watch_job_is_async_iterator_and_uses_async_jobs_client() -> None: + client = _AsyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _PageResponse([_log("starting")]), + _PageResponse([_log("starting"), _log("done")]), + ], + ) + + iterator = async_watch_job(client, "job-a", workspace="default", poll_interval=0) + assert hasattr(iterator, "__aiter__") + + events = [event async for event in iterator] + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("log", None, "starting"), + ("status", "completed", None), + ("log", None, "done"), + ] + assert client.status_calls == [ + {"workspace": "default", "name": "job-a"}, + {"workspace": "default", "name": "job-a"}, + ] + + +def test_async_watch_job_rejects_negative_poll_interval_eagerly() -> None: + client = _AsyncJobsClient(statuses=[], log_results=[]) + + with pytest.raises(ValueError, match="poll_interval"): + async_watch_job(client, "job-a", poll_interval=-1) + + +async def test_async_jobs_client_watch_job_delegates_to_source_owned_watcher( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: dict[str, object] = {} + + async def fake_events(job_name: str) -> AsyncIterator[JobWarningEvent]: + yield JobWarningEvent(kind="warning", job_name=job_name, message="delegated") + + def fake_async_watch_job(client: AsyncJobsClient, name: str, **kwargs: object) -> AsyncIterator[JobWarningEvent]: + calls["client"] = client + calls["name"] = name + calls["kwargs"] = kwargs + return fake_events(name) + + monkeypatch.setattr("nemo_platform_plugin.jobs.watch.async_watch_job", fake_async_watch_job) + http_client = httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request))) + client = AsyncJobsClient(base_url="http://test", http_client=http_client) + + try: + events = [ + event + async for event in client.watch_job( + "job-a", + workspace="default", + poll_interval=0, + timeout=5, + include_history=False, + include_logs=False, + attempt_id=1, + step_id="step-1", + task_id="task-1", + limit=2, + page_cursor="cursor-0", + ) + ] + finally: + await http_client.aclose() + + assert calls == { + "client": client, + "name": "job-a", + "kwargs": { + "workspace": "default", + "poll_interval": 0, + "timeout": 5, + "include_history": False, + "include_logs": False, + "attempt_id": 1, + "step_id": "step-1", + "task_id": "task-1", + "limit": 2, + "page_cursor": "cursor-0", + }, + } + assert events == [JobWarningEvent(kind="warning", job_name="job-a", message="delegated")] diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/client.py index f1862c2f2b..8d2affeb36 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/client.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/client.py @@ -23,9 +23,63 @@ ... """ +from collections.abc import AsyncIterator, Awaitable, Iterator +from typing import Protocol + from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient from nemo_platform_plugin.client.method import method +from nemo_platform_plugin.client.response import AsyncNemoPaginatedResponse, NemoPaginatedResponse, NemoResponse +from nemo_platform_plugin.client.types import CursorPagination from nemo_platform_plugin.jobs import endpoints +from nemo_platform_plugin.jobs.schemas import PlatformJobLog, PlatformJobStatusResponse +from nemo_platform_plugin.jobs.types import JobLogsQueryParams +from nemo_platform_plugin.jobs.watch_types import JobWatchEvent + + +class JobStatusClient(Protocol): + def get_job_status( + self, + *, + workspace: str | None = None, + name: str, + ) -> NemoResponse[PlatformJobStatusResponse]: ... + + +class AsyncJobStatusClient(Protocol): + def get_job_status( + self, + *, + workspace: str | None = None, + name: str, + ) -> Awaitable[NemoResponse[PlatformJobStatusResponse]]: ... + + +class JobLogsClient(Protocol): + def list_job_logs( + self, + *, + workspace: str | None = None, + name: str, + query_params: JobLogsQueryParams | None = None, + ) -> NemoPaginatedResponse[PlatformJobLog, CursorPagination]: ... + + +class AsyncJobLogsClient(Protocol): + def list_job_logs( + self, + *, + workspace: str | None = None, + name: str, + query_params: JobLogsQueryParams | None = None, + ) -> Awaitable[AsyncNemoPaginatedResponse[PlatformJobLog, CursorPagination]]: ... + + +class JobsWatchClient(JobStatusClient, JobLogsClient, Protocol): + """Structural sync Jobs client accepted by the watcher implementation.""" + + +class AsyncJobsWatchClient(AsyncJobStatusClient, AsyncJobLogsClient, Protocol): + """Structural async Jobs client accepted by the watcher implementation.""" class _JobsMethods: @@ -68,6 +122,80 @@ class _JobsMethods: class JobsClient(_JobsMethods, NemoClient): """Sync client for the Jobs service API.""" + def watch_job( + self, + name: str, + *, + workspace: str | None = None, + poll_interval: float = 3, + timeout: float | None = None, + include_history: bool = True, + include_logs: bool = True, + attempt_id: int | None = None, + step_id: str | None = None, + task_id: str | None = None, + limit: int | None = None, + page_cursor: str | None = None, + ) -> Iterator[JobWatchEvent]: + """Watch a platform job and yield status, log, and warning events. + + Poll-based log pagination can miss delayed log entries that sort before + the current cursor. + """ + from nemo_platform_plugin.jobs.watch import watch_job + + return watch_job( + self, + name, + workspace=workspace, + poll_interval=poll_interval, + timeout=timeout, + include_history=include_history, + include_logs=include_logs, + attempt_id=attempt_id, + step_id=step_id, + task_id=task_id, + limit=limit, + page_cursor=page_cursor, + ) + class AsyncJobsClient(_JobsMethods, AsyncNemoClient): """Async client for the Jobs service API.""" + + def watch_job( + self, + name: str, + *, + workspace: str | None = None, + poll_interval: float = 3, + timeout: float | None = None, + include_history: bool = True, + include_logs: bool = True, + attempt_id: int | None = None, + step_id: str | None = None, + task_id: str | None = None, + limit: int | None = None, + page_cursor: str | None = None, + ) -> AsyncIterator[JobWatchEvent]: + """Watch a platform job asynchronously and yield status, log, and warning events. + + Poll-based log pagination can miss delayed log entries that sort before + the current cursor. + """ + from nemo_platform_plugin.jobs.watch import async_watch_job + + return async_watch_job( + self, + name, + workspace=workspace, + poll_interval=poll_interval, + timeout=timeout, + include_history=include_history, + include_logs=include_logs, + attempt_id=attempt_id, + step_id=step_id, + task_id=task_id, + limit=limit, + page_cursor=page_cursor, + ) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/watch.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/watch.py new file mode 100644 index 0000000000..402e776379 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/watch.py @@ -0,0 +1,689 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +import time +from collections.abc import AsyncIterator, Iterable, Iterator +from dataclasses import dataclass, field +from enum import Enum +from typing import ClassVar, Literal, TypeAlias + +from nemo_platform._exceptions import APIConnectionError, APIStatusError, APITimeoutError +from nemo_platform_plugin.client.errors import NemoHTTPError, NemoTransportError +from nemo_platform_plugin.jobs.client import ( + AsyncJobLogsClient, + AsyncJobStatusClient, + AsyncJobsWatchClient, + JobLogsClient, + JobStatusClient, + JobsWatchClient, +) +from nemo_platform_plugin.jobs.schemas import PlatformJobLog, PlatformJobStatusResponse +from nemo_platform_plugin.jobs.types import JobLogsQueryParams +from nemo_platform_plugin.jobs.watch_types import ( + JobLogEvent, + JobStatusEvent, + JobWarningEvent, + JobWatchEvent, + JobWatchTimeoutError, +) + +_SUCCESSFUL_TERMINAL_STATUSES = {"completed"} +# A paused job is resumable, but it is no longer making progress. Treat it as +# failure-terminal so wait/watch callers do not report incomplete work as done. +_FAILED_TERMINAL_STATUSES = {"cancelled", "error", "paused"} +_TERMINAL_STATUSES = _SUCCESSFUL_TERMINAL_STATUSES | _FAILED_TERMINAL_STATUSES +_TRANSIENT_STATUS_CODES = {429, 502, 503, 504} +_TRANSPORT_ERRORS = (NemoTransportError, APIConnectionError, APITimeoutError) +_HTTP_STATUS_ERRORS = (NemoHTTPError, APIStatusError) +_INVALID_PAGE_CURSOR_STATUS_CODE = 422 +_INVALID_PAGE_CURSOR_ERROR_CODE = "invalid_page_cursor" +_PAGE_CURSOR_DECODE_ERROR_DETAIL = "Invalid page cursor" +_LogKey: TypeAlias = tuple[str, str, str, str, str] + + +@dataclass(frozen=True, slots=True) +class _SyncWatchClients: + status: JobStatusClient + logs: JobLogsClient | None + + +@dataclass(frozen=True, slots=True) +class _AsyncWatchClients: + status: AsyncJobStatusClient + logs: AsyncJobLogsClient | None + + +@dataclass(frozen=True, slots=True) +class _WatchOptions: + workspace: str | None + attempt_id: int | None + step_id: str | None + task_id: str | None + limit: int | None + + +@dataclass(slots=True) +class _WatchState: + history_seen: bool + log_cursor: str | None + last_status: str | None = None + previous_drain_seen_logs: dict[_LogKey, int] = field(default_factory=dict) + current_drain_seen_logs: dict[_LogKey, int] = field(default_factory=dict) + + def start_log_drain(self) -> None: + if self.current_drain_seen_logs: + self.previous_drain_seen_logs = self.current_drain_seen_logs + self.current_drain_seen_logs = {} + + def complete_log_drain(self) -> None: + self.previous_drain_seen_logs = self.current_drain_seen_logs + self.current_drain_seen_logs = {} + + def seen_log_occurrences(self, key: _LogKey) -> int: + return max( + self.previous_drain_seen_logs.get(key, 0), + self.current_drain_seen_logs.get(key, 0), + ) + + def record_log_occurrence(self, key: _LogKey, occurrences: int) -> None: + self.current_drain_seen_logs[key] = occurrences + + +@dataclass(slots=True) +class _TerminalLogDrainRetryBudget: + RETRY_CAP: ClassVar[int] = 3 + + _failures_before_warning: Iterator[int] = field(init=False) + + def __post_init__(self) -> None: + self.reset() + + def reset(self) -> None: + self._failures_before_warning = iter(range(1, self.RETRY_CAP)) + + def can_retry_after_failure(self) -> bool: + try: + next(self._failures_before_warning) + except StopIteration: + return False + return True + + def warning(self, job_name: str) -> JobWarningEvent: + message = f"Terminal log drain retry cap reached ({self.RETRY_CAP}); stopping watch for job {job_name!r}" + return JobWarningEvent(kind="warning", job_name=job_name, message=message) + + +@dataclass(slots=True) +class _TransientRetryBackoff: + MAX_SLEEP_SECONDS: ClassVar[float] = 30.0 + + _next_sleep_interval: float | None = None + _last_warning_key: tuple[Literal["status", "log"], str] | None = None + + def reset(self) -> None: + self._next_sleep_interval = None + self._last_warning_key = None + + def next_sleep_interval(self, poll_interval: float) -> float: + if poll_interval <= 0: + self._next_sleep_interval = 0 + return 0 + + max_sleep = max(poll_interval, self.MAX_SLEEP_SECONDS) + sleep_interval = poll_interval if self._next_sleep_interval is None else self._next_sleep_interval + sleep_interval = min(sleep_interval, max_sleep) + self._next_sleep_interval = min(sleep_interval * 2, max_sleep) + return sleep_interval + + def warning( + self, + operation: Literal["status", "log"], + job_name: str, + exc: Exception, + *, + transient: bool = True, + ) -> JobWarningEvent | None: + warning = ( + _transient_failure_warning(operation, job_name, exc) if transient else _log_failure_warning(job_name, exc) + ) + warning_key = (operation, warning.message) + if warning_key == self._last_warning_key: + return None + self._last_warning_key = warning_key + return warning + + +@dataclass(frozen=True, slots=True) +class _TransientRetry: + warning: JobWarningEvent | None + sleep_interval: float + + +@dataclass(frozen=True, slots=True) +class _WatchDeadline: + job_name: str + expires_at: float | None + + @classmethod + def from_timeout(cls, job_name: str, timeout: float | None) -> _WatchDeadline: + expires_at = None if timeout is None else time.monotonic() + timeout + return cls(job_name=job_name, expires_at=expires_at) + + def raise_if_expired(self) -> None: + remaining = self.remaining() + if remaining is not None and remaining <= 0: + raise JobWatchTimeoutError(f"Timed out watching job {self.job_name!r}") + + def remaining(self) -> float | None: + if self.expires_at is None: + return None + return self.expires_at - time.monotonic() + + def sleep_seconds(self, poll_interval: float) -> float: + remaining = self.remaining() + if remaining is None: + return poll_interval + if remaining <= 0: + raise JobWatchTimeoutError(f"Timed out watching job {self.job_name!r}") + return min(poll_interval, remaining) + + +def watch_job( + client: JobsWatchClient, + name: str, + *, + workspace: str | None = None, + poll_interval: float = 3, + timeout: float | None = None, + include_history: bool = True, + include_logs: bool = True, + attempt_id: int | None = None, + step_id: str | None = None, + task_id: str | None = None, + limit: int | None = None, + page_cursor: str | None = None, +) -> Iterator[JobWatchEvent]: + """Watch a platform job until completion. + + The iterator yields typed status, log, and warning events. It accepts a + ``JobsClient``-compatible client from the source-owned Jobs service. + """ + if poll_interval < 0: + raise ValueError("poll_interval must be greater than or equal to 0") + + jobs = _SyncWatchClients(status=client, logs=client if include_logs else None) + options = _WatchOptions( + workspace=workspace, + attempt_id=attempt_id, + step_id=step_id, + task_id=task_id, + limit=limit, + ) + state = _WatchState(history_seen=include_history, log_cursor=page_cursor) + return _watch_job( + jobs, + name, + options=options, + state=state, + deadline=_WatchDeadline.from_timeout(name, timeout), + poll_interval=poll_interval, + ) + + +def async_watch_job( + client: AsyncJobsWatchClient, + name: str, + *, + workspace: str | None = None, + poll_interval: float = 3, + timeout: float | None = None, + include_history: bool = True, + include_logs: bool = True, + attempt_id: int | None = None, + step_id: str | None = None, + task_id: str | None = None, + limit: int | None = None, + page_cursor: str | None = None, +) -> AsyncIterator[JobWatchEvent]: + """Async variant of :func:`watch_job`. + + This has the same poll-based log pagination limitation as + :func:`watch_job`. + """ + if poll_interval < 0: + raise ValueError("poll_interval must be greater than or equal to 0") + + jobs = _AsyncWatchClients(status=client, logs=client if include_logs else None) + options = _WatchOptions( + workspace=workspace, + attempt_id=attempt_id, + step_id=step_id, + task_id=task_id, + limit=limit, + ) + state = _WatchState(history_seen=include_history, log_cursor=page_cursor) + return _async_watch_job( + jobs, + name, + options=options, + state=state, + deadline=_WatchDeadline.from_timeout(name, timeout), + poll_interval=poll_interval, + ) + + +def _watch_job( + jobs: _SyncWatchClients, + name: str, + *, + options: _WatchOptions, + state: _WatchState, + deadline: _WatchDeadline, + poll_interval: float, +) -> Iterator[JobWatchEvent]: + terminal_log_drain_budget = _TerminalLogDrainRetryBudget() + status_retry = _TransientRetryBackoff() + log_retry = _TransientRetryBackoff() + while True: + deadline.raise_if_expired() + try: + status_response = jobs.status.get_job_status(workspace=options.workspace, name=name) + status_event = _status_event(status_response.data(), name) + except _TRANSPORT_ERRORS as exc: + retry = _transient_retry(status_retry, "status", name, exc, poll_interval) + if retry.warning is not None: + yield retry.warning + _sleep(retry.sleep_interval, deadline) + continue + except _HTTP_STATUS_ERRORS as exc: + if not _is_retryable_http_error(exc): + raise + retry = _transient_retry(status_retry, "status", name, exc, poll_interval) + if retry.warning is not None: + yield retry.warning + _sleep(retry.sleep_interval, deadline) + continue + status_retry.reset() + + if status_event.status != state.last_status: + yield status_event + state.last_status = status_event.status + + logs_drained = True + sleep_interval = poll_interval + if jobs.logs is not None: + logs_drained = False + try: + yield from _drain_logs_with_cursor_recovery( + jobs.logs, + name, + options=options, + state=state, + deadline=deadline, + ) + logs_drained = True + log_retry.reset() + except _TRANSPORT_ERRORS as exc: + retry = _transient_retry(log_retry, "log", name, exc, poll_interval) + if retry.warning is not None: + yield retry.warning + sleep_interval = retry.sleep_interval + except _HTTP_STATUS_ERRORS as exc: + retry = ( + _transient_retry(log_retry, "log", name, exc, poll_interval) + if _is_retryable_http_error(exc) + else _log_failure_retry(log_retry, name, exc, poll_interval) + ) + if retry.warning is not None: + yield retry.warning + sleep_interval = retry.sleep_interval + + if status_event.terminal: + if logs_drained: + return + if not terminal_log_drain_budget.can_retry_after_failure(): + yield terminal_log_drain_budget.warning(name) + return + else: + terminal_log_drain_budget.reset() + + _sleep(sleep_interval, deadline) + + +async def _async_watch_job( + jobs: _AsyncWatchClients, + name: str, + *, + options: _WatchOptions, + state: _WatchState, + deadline: _WatchDeadline, + poll_interval: float, +) -> AsyncIterator[JobWatchEvent]: + terminal_log_drain_budget = _TerminalLogDrainRetryBudget() + status_retry = _TransientRetryBackoff() + log_retry = _TransientRetryBackoff() + while True: + deadline.raise_if_expired() + try: + status_response = await jobs.status.get_job_status(workspace=options.workspace, name=name) + status_event = _status_event(status_response.data(), name) + except _TRANSPORT_ERRORS as exc: + retry = _transient_retry(status_retry, "status", name, exc, poll_interval) + if retry.warning is not None: + yield retry.warning + await _async_sleep(retry.sleep_interval, deadline) + continue + except _HTTP_STATUS_ERRORS as exc: + if not _is_retryable_http_error(exc): + raise + retry = _transient_retry(status_retry, "status", name, exc, poll_interval) + if retry.warning is not None: + yield retry.warning + await _async_sleep(retry.sleep_interval, deadline) + continue + status_retry.reset() + + if status_event.status != state.last_status: + yield status_event + state.last_status = status_event.status + + logs_drained = True + sleep_interval = poll_interval + if jobs.logs is not None: + logs_drained = False + try: + async for event in _async_drain_logs_with_cursor_recovery( + jobs.logs, + name, + options=options, + state=state, + deadline=deadline, + ): + yield event + logs_drained = True + log_retry.reset() + except _TRANSPORT_ERRORS as exc: + retry = _transient_retry(log_retry, "log", name, exc, poll_interval) + if retry.warning is not None: + yield retry.warning + sleep_interval = retry.sleep_interval + except _HTTP_STATUS_ERRORS as exc: + retry = ( + _transient_retry(log_retry, "log", name, exc, poll_interval) + if _is_retryable_http_error(exc) + else _log_failure_retry(log_retry, name, exc, poll_interval) + ) + if retry.warning is not None: + yield retry.warning + sleep_interval = retry.sleep_interval + + if status_event.terminal: + if logs_drained: + return + if not terminal_log_drain_budget.can_retry_after_failure(): + yield terminal_log_drain_budget.warning(name) + return + else: + terminal_log_drain_budget.reset() + + await _async_sleep(sleep_interval, deadline) + + +def _sleep(poll_interval: float, deadline: _WatchDeadline) -> None: + sleep_for = deadline.sleep_seconds(poll_interval) + if sleep_for > 0: + time.sleep(sleep_for) + + +async def _async_sleep(poll_interval: float, deadline: _WatchDeadline) -> None: + sleep_for = deadline.sleep_seconds(poll_interval) + if sleep_for > 0: + await asyncio.sleep(sleep_for) + + +def _status_event(status_response: PlatformJobStatusResponse, job_name: str) -> JobStatusEvent: + status = _normalized_status(status_response.status) + terminal = status in _TERMINAL_STATUSES + successful = ( + True if status in _SUCCESSFUL_TERMINAL_STATUSES else False if status in _FAILED_TERMINAL_STATUSES else None + ) + error_details = status_response.error_details + return JobStatusEvent( + kind="status", + job_name=job_name, + status=status, + status_details=dict(status_response.status_details), + terminal=terminal, + successful=successful, + error_details=dict(error_details) if error_details is not None else None, + ) + + +def _normalized_status(value: str | Enum) -> str: + if isinstance(value, Enum): + return str(value.value).lower() + return value.lower() + + +def _drain_logs_with_cursor_recovery( + jobs: JobLogsClient, + name: str, + *, + options: _WatchOptions, + state: _WatchState, + deadline: _WatchDeadline, +) -> Iterator[JobLogEvent]: + state.start_log_drain() + try: + yield from _drain_logs(jobs, name, options=options, state=state, deadline=deadline) + except NemoHTTPError as exc: + if not _can_retry_log_scan_from_start(exc, state.log_cursor): + raise + state.log_cursor = None + yield from _drain_logs(jobs, name, options=options, state=state, deadline=deadline) + state.complete_log_drain() + + +async def _async_drain_logs_with_cursor_recovery( + jobs: AsyncJobLogsClient, + name: str, + *, + options: _WatchOptions, + state: _WatchState, + deadline: _WatchDeadline, +) -> AsyncIterator[JobLogEvent]: + state.start_log_drain() + try: + async for event in _async_drain_logs(jobs, name, options=options, state=state, deadline=deadline): + yield event + except NemoHTTPError as exc: + if not _can_retry_log_scan_from_start(exc, state.log_cursor): + raise + state.log_cursor = None + async for event in _async_drain_logs(jobs, name, options=options, state=state, deadline=deadline): + yield event + state.complete_log_drain() + + +def _drain_logs( + jobs: JobLogsClient, + name: str, + *, + options: _WatchOptions, + state: _WatchState, + deadline: _WatchDeadline, +) -> Iterator[JobLogEvent]: + current_cursor = state.log_cursor + emit_logs = state.history_seen + visited_cursors = {current_cursor} + occurrence_counts: dict[_LogKey, int] = {} + while True: + deadline.raise_if_expired() + page = jobs.list_job_logs( + workspace=options.workspace, + name=name, + query_params=_log_query_params(options, page_cursor=current_cursor), + ).page() + yield from _new_log_events( + page.items, + state=state, + occurrence_counts=occurrence_counts, + name=name, + emit=emit_logs, + ) + + next_cursor = page.metadata["next_page"] + if next_cursor is None or next_cursor in visited_cursors: + state.history_seen = True + return + visited_cursors.add(next_cursor) + state.log_cursor = next_cursor + current_cursor = next_cursor + + +async def _async_drain_logs( + jobs: AsyncJobLogsClient, + name: str, + *, + options: _WatchOptions, + state: _WatchState, + deadline: _WatchDeadline, +) -> AsyncIterator[JobLogEvent]: + current_cursor = state.log_cursor + emit_logs = state.history_seen + visited_cursors = {current_cursor} + occurrence_counts: dict[_LogKey, int] = {} + while True: + deadline.raise_if_expired() + page_response = await jobs.list_job_logs( + workspace=options.workspace, + name=name, + query_params=_log_query_params(options, page_cursor=current_cursor), + ) + page = page_response.page() + for event in _new_log_events( + page.items, + state=state, + occurrence_counts=occurrence_counts, + name=name, + emit=emit_logs, + ): + yield event + + next_cursor = page.metadata["next_page"] + if next_cursor is None or next_cursor in visited_cursors: + state.history_seen = True + return + visited_cursors.add(next_cursor) + state.log_cursor = next_cursor + current_cursor = next_cursor + + +def _new_log_events( + logs: Iterable[PlatformJobLog], + *, + state: _WatchState, + occurrence_counts: dict[_LogKey, int], + name: str, + emit: bool, +) -> Iterator[JobLogEvent]: + for log in logs: + key = _log_key(log) + occurrence_counts[key] = occurrence_counts.get(key, 0) + 1 + seen_occurrences = state.seen_log_occurrences(key) + state.record_log_occurrence(key, occurrence_counts[key]) + if occurrence_counts[key] <= seen_occurrences: + continue + if emit: + yield _log_event(log, name) + + +def _log_query_params(options: _WatchOptions, *, page_cursor: str | None) -> JobLogsQueryParams | None: + params: JobLogsQueryParams = {} + if options.attempt_id is not None: + params["attempt_id"] = options.attempt_id + if options.step_id is not None: + params["step_id"] = options.step_id + if options.task_id is not None: + params["task_id"] = options.task_id + if options.limit is not None: + params["limit"] = options.limit + if page_cursor is not None: + params["page_cursor"] = page_cursor + return params or None + + +def _log_event(log: PlatformJobLog, job_name: str) -> JobLogEvent: + return JobLogEvent( + kind="log", + job_name=job_name, + timestamp=log.timestamp, + step_id=log.job_step, + task_id=log.job_task, + message=log.message, + ) + + +def _log_key(log: PlatformJobLog) -> _LogKey: + return (log.job, log.timestamp.isoformat(), log.job_step, log.job_task, log.message) + + +def _transient_retry( + retry_backoff: _TransientRetryBackoff, + operation: Literal["status", "log"], + job_name: str, + exc: Exception, + poll_interval: float, +) -> _TransientRetry: + return _TransientRetry( + warning=retry_backoff.warning(operation, job_name, exc), + sleep_interval=retry_backoff.next_sleep_interval(poll_interval), + ) + + +def _log_failure_retry( + retry_backoff: _TransientRetryBackoff, + job_name: str, + exc: Exception, + poll_interval: float, +) -> _TransientRetry: + return _TransientRetry( + warning=retry_backoff.warning("log", job_name, exc, transient=False), + sleep_interval=retry_backoff.next_sleep_interval(poll_interval), + ) + + +def _transient_failure_warning( + operation: Literal["status", "log"], + job_name: str, + exc: Exception, +) -> JobWarningEvent: + return JobWarningEvent(kind="warning", job_name=job_name, message=f"Transient {operation} check failed: {exc}") + + +def _log_failure_warning(job_name: str, exc: Exception) -> JobWarningEvent: + return JobWarningEvent(kind="warning", job_name=job_name, message=f"Log check failed: {exc}") + + +def _is_retryable_http_error(exc: NemoHTTPError | APIStatusError) -> bool: + return exc.status_code in _TRANSIENT_STATUS_CODES + + +def _can_retry_log_scan_from_start(exc: NemoHTTPError, page_cursor: str | None) -> bool: + return page_cursor is not None and _is_invalid_page_cursor_error(exc) + + +def _is_invalid_page_cursor_error(exc: NemoHTTPError) -> bool: + if exc.status_code != _INVALID_PAGE_CURSOR_STATUS_CODE: + return False + + match exc.body: + case {"code": code} | {"detail": {"code": code}} | {"error": {"code": code}}: + return code == _INVALID_PAGE_CURSOR_ERROR_CODE + case {"detail": detail}: + return detail == _PAGE_CURSOR_DECODE_ERROR_DETAIL + case _: + return False diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/watch_types.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/watch_types.py new file mode 100644 index 0000000000..3d6b0eda74 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/watch_types.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime +from typing import Literal + + +@dataclass(frozen=True, slots=True) +class JobStatusEvent: + kind: Literal["status"] + job_name: str + status: str + status_details: Mapping[str, object] + terminal: bool + successful: bool | None + error_details: Mapping[str, object] | None = None + + +@dataclass(frozen=True, slots=True) +class JobLogEvent: + kind: Literal["log"] + job_name: str + timestamp: datetime | None + step_id: str | None + task_id: str | None + message: str + + +@dataclass(frozen=True, slots=True) +class JobWarningEvent: + kind: Literal["warning"] + job_name: str + message: str + + +JobWatchEvent = JobStatusEvent | JobLogEvent | JobWarningEvent + + +class JobWatchTimeoutError(TimeoutError): + """Raised when a job watch exceeds its timeout.""" diff --git a/packages/nemo_platform_plugin/tests/jobs/test_watch.py b/packages/nemo_platform_plugin/tests/jobs/test_watch.py new file mode 100644 index 0000000000..c27daaa406 --- /dev/null +++ b/packages/nemo_platform_plugin/tests/jobs/test_watch.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime, timezone + +from nemo_platform_plugin.jobs.schemas import PlatformJobStatus, PlatformJobStatusResponse +from nemo_platform_plugin.jobs.watch import _status_event + + +def test_status_event_preserves_error_details_for_failed_job() -> None: + timestamp = datetime(2026, 1, 1, tzinfo=timezone.utc) + status_response = PlatformJobStatusResponse( + id="job-id", + name="job-name", + status=PlatformJobStatus.ERROR, + status_details={"phase": "failed"}, + error_details={"reason": "container exited"}, + steps=[], + created_at=timestamp, + updated_at=timestamp, + ) + + event = _status_event(status_response, "job-name") + + assert event.status == "error" + assert event.status_details == {"phase": "failed"} + assert event.error_details == {"reason": "container exited"} + assert event.terminal is True + assert event.successful is False diff --git a/pytest.ini b/pytest.ini index 0568982260..bbc7ddb37a 100644 --- a/pytest.ini +++ b/pytest.ini @@ -99,6 +99,8 @@ asyncio_mode = auto addopts = --verbose --strict-markers + # Pytest defaults to prepend; importlib avoids duplicate test module basename collisions. + --import-mode=importlib --tb=short -W ignore::UserWarning:pytest_only.version # Coverage options (can be overridden by command line) diff --git a/sdk/python/nemo-platform/src/nemo_platform/_client.py b/sdk/python/nemo-platform/src/nemo_platform/_client.py index 9c5761f202..e62bc44c9b 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/_client.py +++ b/sdk/python/nemo-platform/src/nemo_platform/_client.py @@ -19,9 +19,14 @@ import os from typing import TYPE_CHECKING, Any, Mapping +from pathlib import Path from typing_extensions import Self, override import httpx +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR + +from nemo_platform.client.tls import client_verify_from_env +from nemo_platform._base_client import DefaultHttpxClient, DefaultAsyncHttpxClient from . import _exceptions from ._qs import Querystring @@ -48,12 +53,9 @@ SyncAPIClient, AsyncAPIClient, ) -from nemo_platform._base_client import DefaultAsyncHttpxClient, DefaultHttpxClient -from nemo_platform.client.tls import client_verify_from_env -from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR -from pathlib import Path if TYPE_CHECKING: + from .models import ModelsResource, AsyncModelsResource from .resources import ( iam, auth, @@ -73,11 +75,10 @@ experiments, ) from .resources.iam.iam import IamResource, AsyncIamResource + from .filesets.resources import FilesResource, AsyncFilesResource from .resources.auth.auth import AuthResource, AsyncAuthResource from .resources.jobs.jobs import JobsResource, AsyncJobsResource - from .filesets.resources import FilesResource, AsyncFilesResource from .resources.intake.intake import IntakeResource, AsyncIntakeResource - from .models import ModelsResource, AsyncModelsResource from .resources.secrets.secrets import SecretsResource, AsyncSecretsResource from .resources.adapters.adapters import AdaptersResource, AsyncAdaptersResource from .resources.entities.entities import EntitiesResource, AsyncEntitiesResource diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/inference/deployments/__init__.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/inference/deployments/__init__.py index ef3ba6c879..87318c63a3 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/inference/deployments/__init__.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/inference/deployments/__init__.py @@ -76,15 +76,28 @@ def create_deployments( wait: Annotated[ bool, typer.Option( - "--wait", help="Wait for the created deployment to reach a terminal state", rich_help_panel="Wait Options" + "--wait", help="Wait for the created deployment to be up and running", rich_help_panel="Lifecycle Options" + ), + ] = False, + watch: Annotated[ + bool, + typer.Option( + "--watch", + help="Watch the created deployment until it is stable, streaming status updates", + rich_help_panel="Lifecycle Options", ), ] = False, timeout: Annotated[ - int, typer.Option("--timeout", min=1, help="Maximum time to wait in seconds", rich_help_panel="Wait Options") + int, + typer.Option( + "--timeout", min=1, help="Maximum time to wait or watch in seconds", rich_help_panel="Lifecycle Options" + ), ] = 1200, poll_interval: Annotated[ int, - typer.Option("--poll-interval", min=1, help="Seconds between status checks", rich_help_panel="Wait Options"), + typer.Option( + "--poll-interval", min=1, help="Seconds between status checks", rich_help_panel="Lifecycle Options" + ), ] = 3, ) -> None: """Create a new ModelDeployment (version 1). @@ -131,12 +144,17 @@ def create_deployments( state: CLIContext = ctx.obj output_format = state.get_output_format(output_format) + if wait and watch: + raise typer.BadParameter("Cannot combine --wait and --watch.") + if handle_code_generation( ["inference", "deployments"], "create", all_kwargs, output_format, state, + watch_config={"type": "inference_deployment", "resource_label": "deployment"} if watch else None, + watch_options={"timeout": timeout, "poll_interval": poll_interval} if watch else None, wait_config={"type": "inference_deployment", "resource_label": "deployment"} if wait else None, wait_options={"timeout": timeout, "poll_interval": poll_interval} if wait else None, ): @@ -152,11 +170,10 @@ def create_deployments( no_truncate=state.get_no_truncate(), timestamp_format=state.get_timestamp_format(), ) - - if wait: + if wait or watch: wait_name = getattr(result, "name", None) or all_kwargs.get("name") if not wait_name: - raise RuntimeError("Unable to determine created resource name for --wait") + raise RuntimeError("Unable to determine created resource name for --wait/--watch") wait_workspace = all_kwargs.get("workspace") if not wait_for_inference_deployment( client, @@ -166,6 +183,7 @@ def create_deployments( poll_interval=poll_interval, ): raise typer.Exit(1) + return @app.command("delete") diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/jobs/__init__.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/jobs/__init__.py index 83e0c9fba3..9df6af0128 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/jobs/__init__.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/api/jobs/__init__.py @@ -8,6 +8,8 @@ from typing import Annotated, Literal import typer +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.jobs.client import JobsClient from nemo_platform.cli.core.api import build_kwargs, merge_filter_dict from nemo_platform.cli.core.code_generator import handle_code_generation @@ -15,6 +17,7 @@ from nemo_platform.cli.core.errors import handle_errors from nemo_platform.cli.core.formatters import Column, check_output_columns_with_format, format_output from nemo_platform.cli.core.help_formatter import collect_warnings, create_typer_app +from nemo_platform.cli.core.job_watch_renderer import render_job_watch_events from nemo_platform.cli.core.pagination import PaginationType, fetch_all_pages, warn_if_more_pages from nemo_platform.cli.core.stdin_utils import read_data_input_with_flags, read_payload, validate_required_fields from nemo_platform.cli.core.types import ( @@ -23,6 +26,7 @@ NoTruncateOption, OutputColumnsOption, ) +from nemo_platform.cli.core.waiters import wait_for_platform_job _cli_child_results = _importlib_import_module("nemo_platform.cli.commands.api.jobs.results") _cli_child_steps = _importlib_import_module("nemo_platform.cli.commands.api.jobs.steps") @@ -96,6 +100,30 @@ def create_jobs( typer.Option("--input-data", help="Input data for the request (JSON or YAML)", rich_help_panel="Input Options"), ] = None, output_format: EntityOutputFormatOption = None, + wait: Annotated[ + bool, + typer.Option( + "--wait", + help="Wait for the created job to reach a terminal state without streaming logs", + rich_help_panel="Lifecycle Options", + ), + ] = False, + watch: Annotated[ + bool, + typer.Option("--watch", help="Watch the created job to a terminal state", rich_help_panel="Lifecycle Options"), + ] = False, + timeout: Annotated[ + int | None, + typer.Option( + "--timeout", min=1, help="Maximum time to wait or watch in seconds", rich_help_panel="Lifecycle Options" + ), + ] = None, + poll_interval: Annotated[ + int, + typer.Option( + "--poll-interval", min=1, help="Seconds between status checks", rich_help_panel="Lifecycle Options" + ), + ] = 3, ) -> None: """Create a new platform job. @@ -150,12 +178,26 @@ def create_jobs( state: CLIContext = ctx.obj output_format = state.get_output_format(output_format) - if handle_code_generation(["jobs"], "create", all_kwargs, output_format, state): + if wait and watch: + raise typer.BadParameter("Cannot combine --wait and --watch.") + + if handle_code_generation( + ["jobs"], + "create", + all_kwargs, + output_format, + state, + watch_config={"type": "platform_job", "resource_label": "job"} if watch else None, + watch_options={"timeout": timeout, "poll_interval": poll_interval} if watch else None, + wait_config={"type": "platform_job", "resource_label": "job"} if wait else None, + wait_options={"timeout": timeout if timeout is not None else 1200, "poll_interval": poll_interval} + if wait + else None, + ): return client = state.get_client() result = client.jobs.create(**all_kwargs) - format_output( result, is_list=False, @@ -163,6 +205,34 @@ def create_jobs( no_truncate=state.get_no_truncate(), timestamp_format=state.get_timestamp_format(), ) + if wait or watch: + wait_name = getattr(result, "name", None) or all_kwargs.get("name") + if not wait_name: + raise RuntimeError("Unable to determine created resource name for --wait/--watch") + wait_workspace = getattr(result, "workspace", None) or all_kwargs.get("workspace") + if wait_workspace is None: + wait_workspace = client._get_workspace_path_param() + jobs_client = client_from_platform(client, JobsClient) + if wait: + if not wait_for_platform_job( + jobs_client, + wait_name, + workspace=wait_workspace, + resource_label="job", + timeout=timeout if timeout is not None else 1200, + poll_interval=poll_interval, + ): + raise typer.Exit(1) + return + events = jobs_client.watch_job( + wait_name, + workspace=wait_workspace, + timeout=timeout, + poll_interval=poll_interval, + ) + if not render_job_watch_events(events, resource_label="job"): + raise typer.Exit(1) + return @app.command("delete") @@ -581,3 +651,46 @@ def update_status_details_jobs( no_truncate=state.get_no_truncate(), timestamp_format=state.get_timestamp_format(), ) + + +@app.command("watch") +@collect_warnings +@handle_errors +def watch_platform_job( + ctx: typer.Context, + name: Annotated[str, typer.Argument(help="Name of the platform job to watch")], + workspace: Annotated[str | None, typer.Option("--workspace", help="Workspace containing the job")] = None, + attempt_id: Annotated[int | None, typer.Option("--attempt-id", help="Filter logs to an attempt ID")] = None, + step_id: Annotated[str | None, typer.Option("--step-id", help="Filter logs to a step ID")] = None, + task_id: Annotated[str | None, typer.Option("--task-id", help="Filter logs to a task ID")] = None, + limit: Annotated[int | None, typer.Option("--limit", min=1, help="Maximum logs to fetch per page")] = None, + timeout: Annotated[int | None, typer.Option("--timeout", min=1, help="Maximum watch time in seconds")] = None, + poll_interval: Annotated[ + int, + typer.Option("--poll-interval", min=1, help="Seconds between status checks"), + ] = 3, + include_history: Annotated[ + bool, + typer.Option("--history/--no-history", help="Include logs already present before watching"), + ] = True, +) -> None: + """Watch a platform job until it reaches a terminal status.""" + state: CLIContext = ctx.obj + client = state.get_client() + jobs_client = client_from_platform(client, JobsClient) + if workspace is None: + workspace = client._get_workspace_path_param() + + events = jobs_client.watch_job( + name, + workspace=workspace, + attempt_id=attempt_id, + step_id=step_id, + task_id=task_id, + limit=limit, + timeout=timeout, + poll_interval=poll_interval, + include_history=include_history, + ) + if not render_job_watch_events(events, resource_label="job"): + raise typer.Exit(1) diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/core/code_generator.py b/sdk/python/nemo-platform/src/nemo_platform/cli/core/code_generator.py index f69cc028f0..0586e692bc 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/core/code_generator.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/core/code_generator.py @@ -11,6 +11,11 @@ from nemo_platform.cli.core.context import CLIContext +_INFERENCE_DEPLOYMENT_LIFECYCLE = "inference_deployment" +_PLATFORM_JOB_LIFECYCLE = "platform_job" +_LIFECYCLE_TYPES_WITH_DEADLINES = {_INFERENCE_DEPLOYMENT_LIFECYCLE} +_LIFECYCLE_TYPES_WITH_STATUS_ERROR_HANDLING = {_INFERENCE_DEPLOYMENT_LIFECYCLE} + def handle_code_generation( resource_path: list[str], @@ -20,6 +25,8 @@ def handle_code_generation( context: CLIContext, wait_config: dict[str, Any] | None = None, wait_options: dict[str, Any] | None = None, + watch_config: dict[str, Any] | None = None, + watch_options: dict[str, Any] | None = None, ) -> bool: """ Check if in code generation mode and generate code if needed. @@ -42,6 +49,8 @@ def handle_code_generation( base_url=base_url, wait_config=wait_config, wait_options=wait_options, + watch_config=watch_config, + watch_options=watch_options, ) formatted_code = format_code_output(code, language="python") print(formatted_code) @@ -57,6 +66,8 @@ def generate_python_code( base_url: str | None = None, wait_config: dict[str, Any] | None = None, wait_options: dict[str, Any] | None = None, + watch_config: dict[str, Any] | None = None, + watch_options: dict[str, Any] | None = None, ) -> str: """ Generate Python SDK code equivalent to a CLI command. @@ -72,16 +83,28 @@ def generate_python_code( """ lines = [] - wait_type = wait_config.get("type") if wait_config else None + if wait_config and watch_config: + raise ValueError("Only one of wait_config or watch_config may be provided") + + lifecycle_config = watch_config or wait_config + lifecycle_options = watch_options if watch_config else wait_options + lifecycle_type = lifecycle_config.get("type") if lifecycle_config else None + + lifecycle_mode = "watch" if watch_config else "wait" if wait_config else None - if wait_config: + if _lifecycle_uses_deadline(lifecycle_type, lifecycle_mode): lines.append("import time") - if wait_type == "inference_deployment": + if _lifecycle_uses_status_error_handling(lifecycle_type, lifecycle_mode): lines.append( "from nemo_platform import APIConnectionError, APIStatusError, APITimeoutError, NeMoPlatform, NotFoundError" ) else: lines.append("from nemo_platform import NeMoPlatform") + if lifecycle_type == _PLATFORM_JOB_LIFECYCLE: + lines.append("from nemo_platform_plugin.client.adapter import client_from_platform") + lines.append("from nemo_platform_plugin.jobs.client import JobsClient") + if lifecycle_mode == "wait": + lines.append("from nemo_platform_plugin.jobs.watch_types import JobStatusEvent, JobWatchTimeoutError") lines.append("") if base_url: @@ -89,15 +112,30 @@ def generate_python_code( else: lines.append("client = NeMoPlatform()") lines.append("") + if lifecycle_type == _PLATFORM_JOB_LIFECYCLE: + lines.append("jobs_client = client_from_platform(client, JobsClient)") + lines.append("") resource_chain = "client." + ".".join(resource_path) _append_method_call(lines, resource_chain, method, _format_method_args(args)) - if wait_config: - lines.extend(["", _render_wait_code(resource_path, args, wait_config, wait_options or {})]) + if lifecycle_config: + lines.extend( + [ + "", + _render_lifecycle_code( + resource_path, + args, + lifecycle_config, + lifecycle_options or {}, + mode=lifecycle_mode, + ), + ] + ) - lines.append("") - lines.append("print(response)") + if lifecycle_type != _PLATFORM_JOB_LIFECYCLE: + lines.append("") + lines.append("print(response)") return "\n".join(lines) @@ -138,16 +176,36 @@ def _format_python_literal(value: Any) -> str: return repr(value) -def _render_wait_code( +def _lifecycle_uses_deadline(lifecycle_type: object, mode: str | None) -> bool: + return lifecycle_type in _LIFECYCLE_TYPES_WITH_DEADLINES and not ( + lifecycle_type == _PLATFORM_JOB_LIFECYCLE and mode == "watch" + ) + + +def _lifecycle_uses_status_error_handling(lifecycle_type: object, mode: str | None) -> bool: + return lifecycle_type in _LIFECYCLE_TYPES_WITH_STATUS_ERROR_HANDLING and not ( + lifecycle_type == _PLATFORM_JOB_LIFECYCLE and mode == "watch" + ) + + +def _require_timeout(timeout: Any, lifecycle_type: object, mode: str | None) -> Any: + if timeout is not None: + return timeout + mode_label = f"{mode} " if mode else "" + raise ValueError(f"{mode_label}{lifecycle_type!r} lifecycle code generation requires timeout") + + +def _render_lifecycle_code( resource_path: list[str], args: dict[str, Any], - wait_config: dict[str, Any], - wait_options: dict[str, Any], + lifecycle_config: dict[str, Any], + lifecycle_options: dict[str, Any], + *, + mode: str | None, ) -> str: - wait_type = wait_config.get("type") - resource_label = str(wait_config.get("resource_label") or "resource") - timeout = wait_options.get("timeout", 1200) - poll_interval = wait_options.get("poll_interval", 3) + lifecycle_type = lifecycle_config.get("type") + timeout = lifecycle_options.get("timeout") + poll_interval = lifecycle_options.get("poll_interval", 3) resource_chain = "client." + ".".join(resource_path) status_kwargs = _format_keyword_args(args, ["workspace"]) resource_name = 'getattr(response, "name", None)' @@ -159,11 +217,14 @@ def _render_wait_code( resource_name = {resource_name} if not resource_name: raise RuntimeError("Unable to determine created resource name for --wait") - deadline = time.monotonic() + {timeout} """ ).strip() + if mode == "watch": + prelude = prelude.replace("--wait", "--watch") - if wait_type == "inference_deployment": + if lifecycle_type == _INFERENCE_DEPLOYMENT_LIFECYCLE: + timeout = _require_timeout(timeout, lifecycle_type, mode) + prelude = "\n".join([prelude, f"deadline = time.monotonic() + {timeout}"]) workspace_literal = _format_python_literal(args["workspace"]) if args.get("workspace") is not None else "None" return "\n\n".join( [ @@ -177,15 +238,23 @@ def _render_wait_code( ] ) - if wait_type == "platform_job": + if lifecycle_type == _PLATFORM_JOB_LIFECYCLE and mode == "watch": return "\n\n".join( [ prelude, - _render_platform_job_wait_code(resource_chain, status_kwargs, resource_label, poll_interval), + _render_platform_job_watch_code(args, timeout, poll_interval), + ] + ) + if lifecycle_type == _PLATFORM_JOB_LIFECYCLE and mode == "wait": + timeout = _require_timeout(timeout, lifecycle_type, mode) + return "\n\n".join( + [ + prelude, + _render_platform_job_wait_code(args, timeout, poll_interval), ] ) - raise ValueError(f"Unsupported wait config type: {wait_type!r}") + raise ValueError(f"Unsupported lifecycle config type: {lifecycle_type!r}") def _render_inference_deployment_wait_code( @@ -240,34 +309,52 @@ def _render_inference_deployment_wait_code( ).strip() +def _render_platform_job_watch_code( + args: dict[str, Any], + timeout: int | None, + poll_interval: int, +) -> str: + workspace = _format_python_literal(args["workspace"]) if args.get("workspace") is not None else "None" + + return dedent( + f""" + for event in jobs_client.watch_job( + resource_name, + workspace={workspace}, + timeout={timeout}, + poll_interval={poll_interval}, + ): + print(event) + """ + ).strip() + + def _render_platform_job_wait_code( - resource_chain: str, - status_kwargs: str, - resource_label: str, + args: dict[str, Any], + timeout: int, poll_interval: int, ) -> str: - resource_label_literal = _format_python_literal(resource_label) + workspace = _format_python_literal(args["workspace"]) if args.get("workspace") is not None else "None" return dedent( f""" - while True: - status_response = {resource_chain}.get_status(resource_name{status_kwargs}) - status = str(status_response.status or "").lower() - if status == "completed": - response = status_response - break - if status in {{"cancelled", "error"}}: - raise RuntimeError( - {resource_label_literal} + f" {{resource_name!r}} ended with status {{status!r}}" - ) - remaining = deadline - time.monotonic() - if remaining <= 0: - raise TimeoutError( - "Timed out waiting for " - + {resource_label_literal} - + f" {{resource_name!r}} to complete" - ) - time.sleep(min({poll_interval}, remaining)) + try: + for event in jobs_client.watch_job( + resource_name, + workspace={workspace}, + timeout={timeout}, + poll_interval={poll_interval}, + include_logs=False, + ): + if not isinstance(event, JobStatusEvent): + continue + if not event.terminal: + continue + if event.successful: + break + raise RuntimeError(f"Job {{resource_name!r}} ended with status {{event.status!r}}") + except JobWatchTimeoutError as exc: + raise TimeoutError(f"Timed out waiting for job {{resource_name!r}} to complete") from exc """ ).strip() diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/core/job_watch_renderer.py b/sdk/python/nemo-platform/src/nemo_platform/cli/core/job_watch_renderer.py new file mode 100644 index 0000000000..64eb01dec7 --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/core/job_watch_renderer.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import time +from collections.abc import Iterable, Mapping +from datetime import datetime + +from nemo_platform_plugin.jobs.watch_types import ( + JobLogEvent, + JobStatusEvent, + JobWarningEvent, + JobWatchEvent, + JobWatchTimeoutError, +) +from rich.console import Console +from rich.text import Text + + +def render_job_watch_events( + events: Iterable[JobWatchEvent], + *, + console: Console | None = None, + error_console: Console | None = None, + resource_label: str | None = None, + start_time: float | None = None, +) -> bool: + """Render job watch events and return whether the terminal status succeeded.""" + output = console or Console() + errors = error_console or Console(stderr=True) + started_at = time.time() if start_time is None else start_time + terminal_event: JobStatusEvent | None = None + + try: + for event in events: + if isinstance(event, JobStatusEvent): + _render_status(output, event) + if event.terminal: + terminal_event = event + elif isinstance(event, JobLogEvent): + _render_log(output, event) + elif isinstance(event, JobWarningEvent): + _render_warning(output, event) + except JobWatchTimeoutError as exc: + errors.print(str(exc), style="red") + return False + + if terminal_event is None: + return False + + _emit_terminal_job_run_event(terminal_event, resource_label=resource_label, start_time=started_at) + + if terminal_event.successful: + output.print(f"Job {terminal_event.job_name!r} completed", style="green") + return True + + message = f"Job {terminal_event.job_name!r} ended with status {terminal_event.status!r}" + error_details = _status_details(terminal_event.error_details or {}) + if error_details: + message = f"{message}: {error_details}" + output.print(message, style="red") + return False + + +def _emit_terminal_job_run_event( + event: JobStatusEvent, + *, + resource_label: str | None, + start_time: float, +) -> None: + if resource_label is None: + return + from .waiters import _emit_job_run_event + + _emit_job_run_event(event, resource_label=resource_label, status=event.status, start_time=start_time) + + +def _render_status(console: Console, event: JobStatusEvent) -> None: + line = Text() + line.append(f"[{_time_label()}] ", style="dim") + line.append("Status: ") + line.append(event.status, style=_status_style(event)) + details = _status_details(event.status_details) + if details: + line.append(f" {details}", style="dim") + console.print(line) + + +def _render_log(console: Console, event: JobLogEvent) -> None: + line = Text() + line.append(f"[{_time_label(event.timestamp)}] ", style="dim") + scope = _scope(event) + if scope: + line.append(f"{scope} | ", style="dim") + line.append(event.message) + console.print(line) + + +def _render_warning(console: Console, event: JobWarningEvent) -> None: + line = Text() + line.append(f"[{_time_label()}] ", style="dim") + line.append(event.message, style="yellow") + console.print(line) + + +def _time_label(timestamp: datetime | None = None) -> str: + value = timestamp.astimezone() if timestamp is not None else datetime.now() + return value.strftime("%H:%M:%S") + + +def _status_style(event: JobStatusEvent) -> str: + if event.successful: + return "green bold" + if event.successful is False: + return "red bold" + return "cyan bold" + + +def _status_details(details: Mapping[str, object]) -> str: + parts = [] + for key, value in details.items(): + if value in (None, "", [], {}): + continue + if isinstance(value, str | int | float | bool): + parts.append(f"{key}={value}") + if len(parts) >= 6: + break + return " ".join(parts) + + +def _scope(event: JobLogEvent) -> str: + if event.step_id and event.task_id: + return f"{event.step_id}/{event.task_id}" + if event.step_id: + return event.step_id + if event.task_id: + return event.task_id + return "" diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/core/waiters.py b/sdk/python/nemo-platform/src/nemo_platform/cli/core/waiters.py index e96af0253c..767746a73e 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/core/waiters.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/core/waiters.py @@ -19,10 +19,22 @@ import logging import time +from collections.abc import Iterator from datetime import datetime, timezone from typing import Any from nemo_platform import APIConnectionError, APIStatusError, APITimeoutError, NotFoundError +from nemo_platform_plugin.client.response import NemoPaginatedResponse, NemoResponse +from nemo_platform_plugin.client.types import CursorPagination +from nemo_platform_plugin.jobs.client import JobsWatchClient +from nemo_platform_plugin.jobs.schemas import PlatformJobLog, PlatformJobStatusResponse +from nemo_platform_plugin.jobs.types import JobLogsQueryParams +from nemo_platform_plugin.jobs.watch_types import ( + JobStatusEvent, + JobWarningEvent, + JobWatchEvent, + JobWatchTimeoutError, +) from rich.console import Console from rich.live import Live from rich.text import Text @@ -51,6 +63,63 @@ def _pause(seconds: float) -> None: time.sleep(seconds) +class _WatchedJobsClient: + def __init__(self, jobs_client: JobsWatchClient) -> None: + self._jobs_client = jobs_client + self.last_status: PlatformJobStatusResponse | None = None + + def get_job_status( + self, + *, + workspace: str | None = None, + name: str, + ) -> NemoResponse[PlatformJobStatusResponse]: + response = self._jobs_client.get_job_status(workspace=workspace, name=name) + self.last_status = response.data() + return response + + def list_job_logs( + self, + *, + workspace: str | None = None, + name: str, + query_params: JobLogsQueryParams | None = None, + ) -> NemoPaginatedResponse[PlatformJobLog, CursorPagination]: + return self._jobs_client.list_job_logs(workspace=workspace, name=name, query_params=query_params) + + def watch_job( + self, + name: str, + *, + workspace: str | None = None, + poll_interval: float = 3, + timeout: float | None = None, + include_history: bool = True, + include_logs: bool = True, + attempt_id: int | None = None, + step_id: str | None = None, + task_id: str | None = None, + limit: int | None = None, + page_cursor: str | None = None, + ) -> Iterator[JobWatchEvent]: + from nemo_platform_plugin.jobs.watch import watch_job + + return watch_job( + self, + name, + workspace=workspace, + poll_interval=poll_interval, + timeout=timeout, + include_history=include_history, + include_logs=include_logs, + attempt_id=attempt_id, + step_id=step_id, + task_id=task_id, + limit=limit, + page_cursor=page_cursor, + ) + + def _seconds_since_creation(entry_timestamp: datetime | str | None, created_at: datetime | None) -> int | None: if created_at is None or entry_timestamp is None: return None @@ -187,6 +256,20 @@ def _make_live_display( return text +class _PlatformJobWaitLiveDisplay: + def __init__(self, *, start_time: float, timeout: int, poll_interval: int) -> None: + self.start_time = start_time + self.timeout = timeout + self.poll_interval = poll_interval + + def snapshot(self) -> tuple[str, int]: + return datetime.now().strftime("%H:%M:%S"), int(time.time() - self.start_time) + + def __rich__(self) -> Text: + polling_time, wait_elapsed = self.snapshot() + return _make_live_display(polling_time, self.timeout, self.poll_interval, wait_elapsed) + + def _sleep_until_next_poll(start_time: float, timeout: float, poll_interval: int) -> bool: if poll_interval <= 0: raise ValueError(f"_sleep_until_next_poll poll_interval must be greater than 0, got {poll_interval}") @@ -221,6 +304,7 @@ def wait_for_inference_deployment( timeout: int = 1200, poll_interval: int = 3, check_gateway: bool = True, + verbose: bool = True, ) -> bool: """Wait for an inference deployment to reach the requested status.""" if workspace is None: @@ -231,13 +315,15 @@ def wait_for_inference_deployment( last_status = "" last_message = "" - console.print(f"[bold]Waiting for deployment '{name}' to reach status: {status}[/bold]\n") + if verbose: + console.print(f"[bold]Waiting for deployment '{name}' to reach status: {status}[/bold]\n") with Live(console=console, refresh_per_second=4, transient=True) as live: while time.time() - start_time < timeout: wait_elapsed = int(time.time() - start_time) polling_time = datetime.now().strftime("%H:%M:%S") - live.update(_make_live_display(polling_time, timeout, poll_interval, wait_elapsed)) + if verbose: + live.update(_make_live_display(polling_time, timeout, poll_interval, wait_elapsed)) try: deployment = client.inference.deployments.retrieve(name, workspace=workspace) @@ -253,7 +339,7 @@ def wait_for_inference_deployment( last_status = current_status last_message = current_message - if history and len(history) > last_history_len: + if verbose and history and len(history) > last_history_len: live.stop() if last_history_len == 0: console.print() @@ -269,11 +355,13 @@ def wait_for_inference_deployment( console.print() live.start() - live.update(_make_live_display(polling_time, timeout, poll_interval, wait_elapsed)) + if verbose: + live.update(_make_live_display(polling_time, timeout, poll_interval, wait_elapsed)) if current_status == status and status != "DELETED": live.stop() - console.print(f"\n[green]✓ Deployment reached {status} status![/green]") + if verbose: + console.print(f"\n[green]✓ Deployment reached {status} status![/green]") if status == "READY" and check_gateway: remaining_timeout = timeout - (time.time() - start_time) if remaining_timeout <= 0: @@ -286,6 +374,7 @@ def wait_for_inference_deployment( provider_workspace, timeout=remaining_timeout, poll_interval=poll_interval, + verbose=verbose, ) return True @@ -303,11 +392,13 @@ def wait_for_inference_deployment( console.print("\n[red]✗ Deployment not found[/red]") return False except (APIConnectionError, APITimeoutError) as exc: - _print_transient_wait_error(live, "deployment status", exc) + if verbose: + _print_transient_wait_error(live, "deployment status", exc) except APIStatusError as exc: if exc.status_code not in _TRANSIENT_GATEWAY_STATUS_CODES: raise - _print_transient_wait_error(live, "deployment status", exc) + if verbose: + _print_transient_wait_error(live, "deployment status", exc) if not _sleep_until_next_poll(start_time, timeout, poll_interval): break @@ -321,7 +412,7 @@ def wait_for_inference_deployment( def wait_for_platform_job( - jobs_resource: Any, + jobs_client: JobsWatchClient, name: str, *, workspace: str | None = None, @@ -332,60 +423,56 @@ def wait_for_platform_job( """Wait for a platform job resource to complete.""" start_time = time.time() last_status = "" + jobs = _WatchedJobsClient(jobs_client) console.print(f"[bold]Waiting for {resource_label} '{name}' to complete[/bold]\n") - with Live(console=console, refresh_per_second=4, transient=True) as live: - while time.time() - start_time < timeout: - wait_elapsed = int(time.time() - start_time) - polling_time = datetime.now().strftime("%H:%M:%S") - live.update(_make_live_display(polling_time, timeout, poll_interval, wait_elapsed)) - - try: - job_status = jobs_resource.get_status(name, workspace=workspace) - except NotFoundError: - live.stop() - console.print(f"\n[red]✗ {resource_label.title()} not found[/red]") - return False - except (APIConnectionError, APITimeoutError) as exc: - _print_transient_wait_error(live, f"{resource_label} status", exc) - if not _sleep_until_next_poll(start_time, timeout, poll_interval): - break - continue - except APIStatusError as exc: - if exc.status_code not in _TRANSIENT_GATEWAY_STATUS_CODES: - raise - _print_transient_wait_error(live, f"{resource_label} status", exc) - if not _sleep_until_next_poll(start_time, timeout, poll_interval): - break - continue - - current_status = _status_text(getattr(job_status, "status", "")).lower() - if current_status != last_status: - live.stop() - console.print(_make_history_line(polling_time, wait_elapsed, current_status)) - last_status = current_status - console.print() - live.start() + live_display = _PlatformJobWaitLiveDisplay(start_time=start_time, timeout=timeout, poll_interval=poll_interval) + with Live(live_display, console=console, refresh_per_second=4, transient=True) as live: + try: + for event in jobs.watch_job( + name, + workspace=workspace, + poll_interval=poll_interval, + timeout=timeout, + include_logs=False, + ): + polling_time, wait_elapsed = live_display.snapshot() + live.update(live_display) + + if isinstance(event, JobWarningEvent): + live.stop() + console.print(f"\n[yellow]{event.message}[/yellow]") + live.start() + continue - if current_status == "completed": - live.stop() - _emit_job_run_event( - job_status, resource_label=resource_label, status=current_status, start_time=start_time - ) - console.print(f"\n[green]✓ {resource_label.title()} completed![/green]") - return True + if not isinstance(event, JobStatusEvent): + continue - if current_status in {"cancelled", "error"}: - live.stop() - _emit_job_run_event( - job_status, resource_label=resource_label, status=current_status, start_time=start_time - ) - console.print(f"\n[red]✗ {resource_label.title()} entered {current_status} state[/red]") - return False + current_status = event.status + if current_status != last_status: + live.stop() + console.print(_make_history_line(polling_time, wait_elapsed, current_status)) + last_status = current_status + console.print() + live.start() - if not _sleep_until_next_poll(start_time, timeout, poll_interval): - break + if event.terminal: + live.stop() + _emit_job_run_event( + jobs.last_status, resource_label=resource_label, status=current_status, start_time=start_time + ) + if event.successful: + console.print(f"\n[green]✓ {resource_label.title()} completed![/green]") + return True + console.print(f"\n[red]✗ {resource_label.title()} entered {current_status} state[/red]") + return False + except NotFoundError: + live.stop() + console.print(f"\n[red]✗ {resource_label.title()} not found[/red]") + return False + except JobWatchTimeoutError: + pass wait_elapsed = int(time.time() - start_time) detail = f"Last status: {last_status}" if last_status else "No status returned" @@ -399,12 +486,14 @@ def wait_for_gateway( workspace: str, timeout: float = 60, poll_interval: int = 1, + verbose: bool = True, ) -> bool: """Wait for the inference gateway to be able to route to a provider.""" start_time = time.time() start_timestamp = datetime.now().strftime("%H:%M:%S") - console.print(f"[bold]Waiting for gateway to be ready for provider '{provider_name}'[/bold]\n") + if verbose: + console.print(f"[bold]Waiting for gateway to be ready for provider '{provider_name}'[/bold]\n") def _make_gateway_display(polling_time: str, elapsed: int, status: str) -> Text: text = Text() @@ -418,12 +507,14 @@ def _make_gateway_display(polling_time: str, elapsed: int, status: str) -> Text: while time.time() - start_time < timeout: elapsed = int(time.time() - start_time) polling_time = datetime.now().strftime("%H:%M:%S") - live.update(_make_gateway_display(polling_time, elapsed, "Checking gateway...")) + if verbose: + live.update(_make_gateway_display(polling_time, elapsed, "Checking gateway...")) try: client.inference.gateway.provider.ready(provider_name, workspace=workspace) live.stop() - console.print(f" [{polling_time}] ({elapsed}s) [green]Gateway is ready![/green]") + if verbose: + console.print(f" [{polling_time}] ({elapsed}s) [green]Gateway is ready![/green]") return True except NotFoundError: pass diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_create_wait.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_create_wait.py index 1251a44f97..c01ad9594e 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_create_wait.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_create_wait.py @@ -8,9 +8,18 @@ import pytest import typer - -# from nemo_platform_ext.cli.commands.api.customization.jobs import create_jobs as create_customization_job from nemo_platform.cli.commands.api.inference.deployments import create_deployments +from nemo_platform.cli.commands.api.jobs import create_jobs, watch_platform_job + + +class _CreatedJob: + def __init__(self, *, name: str = "created-job", workspace: str = "result-workspace") -> None: + self.name = name + self.workspace = workspace + + def model_dump(self, *, mode: str = "json") -> dict[str, str]: + assert mode == "json" + return {"name": self.name, "workspace": self.workspace} def _ctx(client: object) -> SimpleNamespace: @@ -23,55 +32,310 @@ def _ctx(client: object) -> SimpleNamespace: return SimpleNamespace(obj=state) -# def test_customization_job_create_waits_for_created_job() -> None: -# jobs = MagicMock() -# jobs.create.return_value = SimpleNamespace(name="created-job") -# client = SimpleNamespace(customization=SimpleNamespace(jobs=jobs)) -# ctx = _ctx(client) - -# with ( -# patch( -# "nemo_platform_ext.cli.commands.api.customization.jobs.handle_code_generation", -# return_value=False, -# ) as handle_code_generation, -# patch("nemo_platform_ext.cli.commands.api.customization.jobs.format_output"), -# patch( -# "nemo_platform_ext.cli.commands.api.customization.jobs.wait_for_platform_job", -# return_value=True, -# ) as wait_for_platform_job, -# ): -# create_customization_job( -# ctx, -# name="input-job", -# workspace="test-workspace", -# spec='{"training_type": "sft"}', -# wait=True, -# timeout=42, -# poll_interval=7, -# ) - -# handle_code_generation.assert_called_once_with( -# ["customization", "jobs"], -# "create", -# {"workspace": "test-workspace", "spec": {"training_type": "sft"}, "name": "input-job"}, -# None, -# ctx.obj, -# wait_config={"type": "platform_job", "resource_label": "customization job"}, -# wait_options={"timeout": 42, "poll_interval": 7}, -# ) -# jobs.create.assert_called_once_with( -# workspace="test-workspace", -# spec={"training_type": "sft"}, -# name="input-job", -# ) -# wait_for_platform_job.assert_called_once_with( -# jobs, -# "created-job", -# workspace="test-workspace", -# resource_label="customization job", -# timeout=42, -# poll_interval=7, -# ) +def test_jobs_create_watch_uses_sdk_watcher_and_outputs_created_job() -> None: + jobs = MagicMock() + jobs.create.return_value = _CreatedJob() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + ctx = _ctx(client) + events = object() + jobs_client = MagicMock() + jobs_client.watch_job.return_value = events + + with ( + patch( + "nemo_platform.cli.commands.api.jobs.handle_code_generation", + return_value=False, + ) as handle_code_generation, + patch("nemo_platform.cli.commands.api.jobs.client_from_platform", return_value=jobs_client), + patch("nemo_platform.cli.commands.api.jobs.format_output") as format_output, + patch("nemo_platform.cli.commands.api.jobs.render_job_watch_events", return_value=True) as render_events, + ): + create_jobs( + ctx, + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=False, + watch=True, + timeout=42, + poll_interval=7, + ) + + expected_kwargs = { + "workspace": "test-workspace", + "platform_spec": {}, + "source": "test-source", + "spec": {}, + "name": "input-job", + } + handle_code_generation.assert_called_once_with( + ["jobs"], + "create", + expected_kwargs, + None, + ctx.obj, + watch_config={"type": "platform_job", "resource_label": "job"}, + watch_options={"timeout": 42, "poll_interval": 7}, + wait_config=None, + wait_options=None, + ) + jobs.create.assert_called_once_with(**expected_kwargs) + format_output.assert_called_once_with( + jobs.create.return_value, + is_list=False, + output_format=None, + no_truncate=False, + timestamp_format=None, + ) + jobs_client.watch_job.assert_called_once_with( + "created-job", + workspace="result-workspace", + timeout=42, + poll_interval=7, + ) + render_events.assert_called_once_with(events, resource_label="job") + + +def test_jobs_create_watch_exits_when_renderer_reports_failure() -> None: + jobs = MagicMock() + jobs.create.return_value = _CreatedJob() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + jobs_client = MagicMock() + jobs_client.watch_job.return_value = object() + + with ( + patch("nemo_platform.cli.commands.api.jobs.handle_code_generation", return_value=False), + patch("nemo_platform.cli.commands.api.jobs.client_from_platform", return_value=jobs_client), + patch("nemo_platform.cli.commands.api.jobs.render_job_watch_events", return_value=False), + pytest.raises(typer.Exit) as exc_info, + ): + create_jobs( + _ctx(client), + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=False, + watch=True, + ) + + assert exc_info.value.exit_code == 1 + + +def test_jobs_create_watch_has_no_default_timeout() -> None: + jobs = MagicMock() + jobs.create.return_value = _CreatedJob() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + ctx = _ctx(client) + events = object() + jobs_client = MagicMock() + jobs_client.watch_job.return_value = events + + with ( + patch( + "nemo_platform.cli.commands.api.jobs.handle_code_generation", + return_value=False, + ) as handle_code_generation, + patch("nemo_platform.cli.commands.api.jobs.client_from_platform", return_value=jobs_client), + patch("nemo_platform.cli.commands.api.jobs.render_job_watch_events", return_value=True), + ): + create_jobs( + ctx, + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=False, + watch=True, + ) + + handle_code_generation.assert_called_once() + assert handle_code_generation.call_args.kwargs["watch_options"] == {"timeout": None, "poll_interval": 3} + jobs_client.watch_job.assert_called_once() + assert jobs_client.watch_job.call_args.kwargs["timeout"] is None + + +def test_jobs_create_wait_uses_quiet_waiter_and_outputs_created_job() -> None: + jobs = MagicMock() + jobs.create.return_value = _CreatedJob() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + ctx = _ctx(client) + jobs_client = MagicMock() + + with ( + patch( + "nemo_platform.cli.commands.api.jobs.handle_code_generation", + return_value=False, + ) as handle_code_generation, + patch("nemo_platform.cli.commands.api.jobs.client_from_platform", return_value=jobs_client), + patch("nemo_platform.cli.commands.api.jobs.format_output") as format_output, + patch("nemo_platform.cli.commands.api.jobs.wait_for_platform_job", return_value=True) as wait_for_job, + patch("nemo_platform.cli.commands.api.jobs.render_job_watch_events") as render_events, + ): + create_jobs( + ctx, + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=True, + watch=False, + timeout=42, + poll_interval=7, + ) + + expected_kwargs = { + "workspace": "test-workspace", + "platform_spec": {}, + "source": "test-source", + "spec": {}, + "name": "input-job", + } + handle_code_generation.assert_called_once_with( + ["jobs"], + "create", + expected_kwargs, + None, + ctx.obj, + watch_config=None, + watch_options=None, + wait_config={"type": "platform_job", "resource_label": "job"}, + wait_options={"timeout": 42, "poll_interval": 7}, + ) + jobs.create.assert_called_once_with(**expected_kwargs) + wait_for_job.assert_called_once_with( + jobs_client, + "created-job", + workspace="result-workspace", + resource_label="job", + timeout=42, + poll_interval=7, + ) + format_output.assert_called_once_with( + jobs.create.return_value, + is_list=False, + output_format=None, + no_truncate=False, + timestamp_format=None, + ) + jobs_client.watch_job.assert_not_called() + render_events.assert_not_called() + + +def test_jobs_create_wait_uses_waiter_default_timeout() -> None: + jobs = MagicMock() + jobs.create.return_value = _CreatedJob() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + jobs_client = MagicMock() + + with ( + patch( + "nemo_platform.cli.commands.api.jobs.handle_code_generation", + return_value=False, + ) as handle_code_generation, + patch("nemo_platform.cli.commands.api.jobs.client_from_platform", return_value=jobs_client), + patch("nemo_platform.cli.commands.api.jobs.wait_for_platform_job", return_value=True) as wait_for_job, + ): + create_jobs( + _ctx(client), + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=True, + watch=False, + ) + + assert handle_code_generation.call_args.kwargs["wait_options"] == {"timeout": 1200, "poll_interval": 3} + wait_for_job.assert_called_once() + assert wait_for_job.call_args.kwargs["timeout"] == 1200 + + +def test_jobs_create_wait_exits_when_waiter_reports_failure() -> None: + jobs = MagicMock() + jobs.create.return_value = _CreatedJob() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + + with ( + patch("nemo_platform.cli.commands.api.jobs.handle_code_generation", return_value=False), + patch("nemo_platform.cli.commands.api.jobs.wait_for_platform_job", return_value=False), + pytest.raises(typer.Exit) as exc_info, + ): + create_jobs( + _ctx(client), + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=True, + watch=False, + ) + + assert exc_info.value.exit_code == 1 + + +def test_jobs_create_rejects_wait_and_watch_together() -> None: + jobs = MagicMock() + client = SimpleNamespace(jobs=jobs, _get_workspace_path_param=MagicMock(return_value="default")) + + with pytest.raises(SystemExit) as exc_info: + create_jobs( + _ctx(client), + name="input-job", + workspace="test-workspace", + platform_spec="{}", + source="test-source", + spec="{}", + wait=True, + watch=True, + ) + + assert exc_info.value.code == 2 + jobs.create.assert_not_called() + + +def test_jobs_watch_command_uses_sdk_watcher() -> None: + client = SimpleNamespace(_get_workspace_path_param=MagicMock(return_value="default")) + events = object() + jobs_client = MagicMock() + jobs_client.watch_job.return_value = events + + with ( + patch("nemo_platform.cli.commands.api.jobs.client_from_platform", return_value=jobs_client), + patch("nemo_platform.cli.commands.api.jobs.render_job_watch_events", return_value=True) as render_events, + ): + watch_platform_job( + _ctx(client), + name="job-a", + workspace=None, + attempt_id=1, + step_id="step-1", + task_id="task-1", + limit=25, + timeout=42, + poll_interval=7, + include_history=False, + ) + + jobs_client.watch_job.assert_called_once_with( + "job-a", + workspace="default", + attempt_id=1, + step_id="step-1", + task_id="task-1", + limit=25, + timeout=42, + poll_interval=7, + include_history=False, + ) + render_events.assert_called_once_with(events, resource_label="job") def test_inference_deployment_create_exits_when_wait_fails() -> None: @@ -94,6 +358,7 @@ def test_inference_deployment_create_exits_when_wait_fails() -> None: workspace="test-workspace", config="deployment-config", wait=True, + watch=False, timeout=90, poll_interval=10, ) @@ -106,3 +371,76 @@ def test_inference_deployment_create_exits_when_wait_fails() -> None: timeout=90, poll_interval=10, ) + + +def test_inference_deployment_create_watch_uses_waiter() -> None: + deployments = MagicMock() + deployments.create.return_value = SimpleNamespace(name="deployment-a") + client = SimpleNamespace(inference=SimpleNamespace(deployments=deployments)) + ctx = _ctx(client) + + with ( + patch( + "nemo_platform.cli.commands.api.inference.deployments.handle_code_generation", + return_value=False, + ) as handle_code_generation, + patch("nemo_platform.cli.commands.api.inference.deployments.format_output") as format_output, + patch( + "nemo_platform.cli.commands.api.inference.deployments.wait_for_inference_deployment", + return_value=True, + ) as wait_for_inference_deployment, + ): + create_deployments( + ctx, + name="deployment-a", + workspace="test-workspace", + config="deployment-config", + wait=False, + watch=True, + timeout=90, + poll_interval=10, + ) + + expected_kwargs = { + "workspace": "test-workspace", + "config": "deployment-config", + "name": "deployment-a", + } + handle_code_generation.assert_called_once_with( + ["inference", "deployments"], + "create", + expected_kwargs, + None, + ctx.obj, + watch_config={"type": "inference_deployment", "resource_label": "deployment"}, + watch_options={"timeout": 90, "poll_interval": 10}, + wait_config=None, + wait_options=None, + ) + deployments.create.assert_called_once_with(**expected_kwargs) + format_output.assert_called_once() + wait_for_inference_deployment.assert_called_once_with( + client, + "deployment-a", + workspace="test-workspace", + timeout=90, + poll_interval=10, + ) + + +def test_inference_deployment_create_rejects_wait_and_watch_together() -> None: + deployments = MagicMock() + client = SimpleNamespace(inference=SimpleNamespace(deployments=deployments)) + + with pytest.raises(SystemExit) as exc_info: + create_deployments( + _ctx(client), + name="deployment-a", + workspace="test-workspace", + config="deployment-config", + wait=True, + watch=True, + ) + + assert exc_info.value.code == 2 + deployments.create.assert_not_called() diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_code_generator.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_code_generator.py index f7ff03c87f..2a722ad442 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_code_generator.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_code_generator.py @@ -98,28 +98,90 @@ def test_generate_python_code_multiline_format(): assert any("namespace=" in line and line.strip().startswith("namespace=") for line in lines) -def test_generate_python_code_with_platform_job_wait(): +def test_generate_python_code_with_platform_job_watch(): code = generate_python_code( resource_path=["customization", "jobs"], method="create", args={"workspace": "default", "name": "job-a", "spec": {"training_type": "sft"}}, - wait_config={"type": "platform_job", "resource_label": "customization job"}, - wait_options={"timeout": 42, "poll_interval": 7}, + watch_config={"type": "platform_job", "resource_label": "customization job"}, + watch_options={"timeout": 42, "poll_interval": 7}, ) - assert "import time" in code + assert "import time" not in code + assert "from nemo_platform.jobs.watch import watch_job" not in code + assert "from nemo_platform_plugin.client.adapter import client_from_platform" in code + assert "from nemo_platform_plugin.jobs.client import JobsClient" in code + assert "jobs_client = client_from_platform(client, JobsClient)" in code assert "response = client.customization.jobs.create" in code assert 'resource_name = getattr(response, "name", None) or "job-a"' in code + assert 'raise RuntimeError("Unable to determine created resource name for --watch")' in code + assert "jobs_client.watch_job(" in code + assert 'workspace="default"' in code + assert "timeout=42" in code + assert "poll_interval=7" in code + assert "print(event)" in code + assert "get_status" not in code + assert "time.sleep" not in code + assert "print(response)" not in code + compile(code, "", "exec") + + +def test_generate_python_code_with_platform_job_watch_has_no_default_timeout(): + code = generate_python_code( + resource_path=["customization", "jobs"], + method="create", + args={"workspace": "default", "name": "job-a", "spec": {"training_type": "sft"}}, + watch_config={"type": "platform_job", "resource_label": "customization job"}, + watch_options={"poll_interval": 7}, + ) + + assert "timeout=None" in code + assert "deadline = time.monotonic()" not in code + assert "poll_interval=7" in code + compile(code, "", "exec") + + +def test_generate_python_code_with_platform_job_wait(): + code = generate_python_code( + resource_path=["jobs"], + method="create", + args={"workspace": "default", "name": "job-a", "spec": {"training_type": "sft"}}, + wait_config={"type": "platform_job", "resource_label": "job"}, + wait_options={"timeout": 42, "poll_interval": 7}, + ) + + assert "import time" not in code + for symbol in ("JobStatusEvent", "JobWatchTimeoutError", "JobsClient", "NeMoPlatform"): + assert symbol in code + assert "from nemo_platform_plugin.jobs.watch_types import JobStatusEvent, JobWatchTimeoutError" in code + assert "jobs_client = client_from_platform(client, JobsClient)" in code + assert "APIConnectionError" not in code + assert "APIStatusError" not in code + assert "APITimeoutError" not in code + assert "NotFoundError" not in code assert 'raise RuntimeError("Unable to determine created resource name for --wait")' in code - assert "deadline = time.monotonic() + 42" in code - assert 'client.customization.jobs.get_status(resource_name, workspace="default")' in code - assert 'status = str(status_response.status or "").lower()' in code - assert "response = status_response" in code - assert code.rindex("print(response)") > code.index("response = status_response") - assert "time.sleep(min(7, remaining))" in code + assert "deadline = time.monotonic()" not in code + assert "get_status" not in code + assert "jobs_client.watch_job(" in code + assert "include_logs=False" in code + assert "isinstance(event, JobStatusEvent)" in code + assert "except JobWatchTimeoutError as exc:" in code + assert "time.sleep" not in code + assert "print(response)" not in code compile(code, "", "exec") +def test_generate_python_code_with_platform_job_wait_requires_timeout(): + with pytest.raises(ValueError, match=r"wait 'platform_job' lifecycle code generation requires timeout"): + generate_python_code( + resource_path=["jobs"], + method="create", + args={"workspace": "default", "name": "job-a", "spec": {"training_type": "sft"}}, + wait_config={"type": "platform_job", "resource_label": "job"}, + wait_options={"poll_interval": 7}, + ) + + def test_generate_python_code_with_inference_deployment_wait(): code = generate_python_code( resource_path=["inference", "deployments"], @@ -149,21 +211,55 @@ def test_generate_python_code_with_inference_deployment_wait(): compile(code, "", "exec") -def test_generate_python_code_escapes_platform_job_wait_label(): +def test_generate_python_code_with_inference_deployment_watch(): + code = generate_python_code( + resource_path=["inference", "deployments"], + method="create", + args={"workspace": "default", "name": "deployment-a", "config": "deployment-config"}, + watch_config={"type": "inference_deployment", "resource_label": "deployment"}, + watch_options={"timeout": 90, "poll_interval": 10}, + ) + + assert "import time" in code + for symbol in ("APIConnectionError", "APIStatusError", "APITimeoutError", "NeMoPlatform", "NotFoundError"): + assert symbol in code + assert "deadline = time.monotonic() + 90" in code + assert 'resource_name = getattr(response, "name", None) or "deployment-a"' in code + assert 'raise RuntimeError("Unable to determine created resource name for --watch")' in code + assert 'client.inference.deployments.retrieve(resource_name, workspace="default")' in code + assert "client.inference.gateway.provider.ready(provider_name, workspace=provider_workspace)" in code + assert "response = deployment" in code + assert code.rindex("print(response)") > code.index("response = deployment") + assert "time.sleep(min(10, remaining))" in code + compile(code, "", "exec") + + +def test_generate_python_code_with_inference_deployment_wait_requires_timeout(): + with pytest.raises(ValueError, match=r"wait 'inference_deployment' lifecycle code generation requires timeout"): + generate_python_code( + resource_path=["inference", "deployments"], + method="create", + args={"workspace": "default", "name": "deployment-a", "config": "deployment-config"}, + wait_config={"type": "inference_deployment", "resource_label": "deployment"}, + wait_options={"poll_interval": 10}, + ) + + +def test_generate_python_code_with_platform_job_watch_ignores_label_formatting(): code = generate_python_code( resource_path=["customization", "jobs"], method="create", args={"workspace": "default", "name": "job-a"}, - wait_config={"type": "platform_job", "resource_label": 'customization "job" {label}'}, - wait_options={"timeout": 42, "poll_interval": 7}, + watch_config={"type": "platform_job", "resource_label": 'customization "job" {label}'}, + watch_options={"timeout": 42, "poll_interval": 7}, ) compile(code, "", "exec") - assert '"customization \\"job\\" {label}" + f" {resource_name!r}' in code + assert 'raise RuntimeError("Unable to determine created resource name for --watch")' in code -def test_generate_python_code_rejects_unknown_wait_type(): - with pytest.raises(ValueError, match="Unsupported wait config type: 'unknown'"): +def test_generate_python_code_rejects_unknown_lifecycle_type(): + with pytest.raises(ValueError, match="Unsupported lifecycle config type: 'unknown'"): generate_python_code( resource_path=["customization", "jobs"], method="create", diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_job_watch_renderer.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_job_watch_renderer.py new file mode 100644 index 0000000000..03386c1aae --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_job_watch_renderer.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Iterator +from datetime import datetime, timezone +from io import StringIO +from unittest.mock import patch + +from nemo_platform.cli.core.job_watch_renderer import render_job_watch_events +from nemo_platform.cli.telemetry.events import TaskStatusEnum +from nemo_platform_plugin.jobs.watch_types import JobLogEvent, JobStatusEvent, JobWatchEvent, JobWatchTimeoutError +from rich.console import Console + + +def _console_pair() -> tuple[Console, StringIO]: + output = StringIO() + return Console(file=output, force_terminal=False, color_system=None, width=120), output + + +def test_render_job_watch_events_returns_true_for_completed_status() -> None: + console, output = _console_pair() + error_console, error_output = _console_pair() + events: list[JobWatchEvent] = [ + JobStatusEvent( + kind="status", + job_name="job-a", + status="active", + status_details={"phase": "training", "progress_pct": 41}, + terminal=False, + successful=None, + ), + log_event := JobLogEvent( + kind="log", + job_name="job-a", + timestamp=datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc), + step_id="step-a", + task_id="task-a", + message="started", + ), + JobStatusEvent( + kind="status", + job_name="job-a", + status="completed", + status_details={}, + terminal=True, + successful=True, + ), + ] + + assert render_job_watch_events(events, console=console, error_console=error_console) is True + + rendered = output.getvalue() + expected_log_time = log_event.timestamp.astimezone().strftime("%H:%M:%S") + assert "Status: active phase=training progress_pct=41" in rendered + assert f"[{expected_log_time}] step-a/task-a | started" in rendered + assert "step-a/task-a | started" in rendered + assert "Job 'job-a' completed" in rendered + assert error_output.getvalue() == "" + + +def test_render_job_watch_events_emits_job_run_event_for_terminal_status() -> None: + console, _ = _console_pair() + event = JobStatusEvent( + kind="status", + job_name="job-a", + status="completed", + status_details={"model": "nemotron"}, + terminal=True, + successful=True, + ) + + with patch("nemo_platform.cli.telemetry.emit.emit_event") as emit_event: + assert render_job_watch_events([event], console=console, resource_label="job", start_time=123.0) is True + + emit_event.assert_called_once() + telemetry_event = emit_event.call_args.args[0] + assert telemetry_event.job_type == "job" + assert telemetry_event.task_status is TaskStatusEnum.COMPLETED + assert telemetry_event.model == "defined" + + +def test_render_job_watch_events_returns_false_for_failed_terminal_status() -> None: + console, output = _console_pair() + event = JobStatusEvent( + kind="status", + job_name="job-a", + status="error", + status_details={}, + error_details={"reason": "container exited", "exit_code": 137, "empty": {}}, + terminal=True, + successful=False, + ) + + assert render_job_watch_events([event], console=console) is False + + assert "Job 'job-a' ended with status 'error': reason=container exited exit_code=137" in output.getvalue() + + +def test_render_job_watch_events_catches_timeout() -> None: + console, output = _console_pair() + error_console, error_output = _console_pair() + + def events() -> Iterator[JobWatchEvent]: + yield JobStatusEvent( + kind="status", + job_name="job-a", + status="active", + status_details={}, + terminal=False, + successful=None, + ) + raise JobWatchTimeoutError("Timed out watching job 'job-a'") + + assert render_job_watch_events(events(), console=console, error_console=error_console) is False + + assert "Status: active" in output.getvalue() + assert "Timed out watching job 'job-a'" in error_output.getvalue() diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_waiters.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_waiters.py index 23ab7f7646..7ce04f55ed 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_waiters.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_waiters.py @@ -4,6 +4,7 @@ from __future__ import annotations from collections.abc import Iterator +from datetime import datetime, timezone from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -11,8 +12,11 @@ import pytest from nemo_platform import APIConnectionError, APIStatusError, AuthenticationError from nemo_platform.cli.core import waiters +from nemo_platform_plugin.jobs.schemas import PlatformJobStatus, PlatformJobStatusResponse WAITERS_MODULE = "nemo_platform.cli.core.waiters" +WATCH_MODULE = "nemo_platform_plugin.jobs.watch" +JOB_TIMESTAMP = datetime(2026, 1, 1, tzinfo=timezone.utc) class _DummyLive: @@ -35,6 +39,41 @@ def start(self) -> None: pass +class _RecordingLive(_DummyLive): + instances: list[_RecordingLive] = [] + + def __init__(self, renderable: object | None = None, *_args: object, **_kwargs: object) -> None: + self.renderable = renderable + self.updates: list[object] = [] + self.instances.append(self) + + def update(self, renderable: object, *_args: object, **_kwargs: object) -> None: + self.updates.append(renderable) + + +class _StatusResponse: + def __init__(self, status: PlatformJobStatusResponse) -> None: + self._status = status + + def data(self) -> PlatformJobStatusResponse: + return self._status + + +def _status_response(status: str | PlatformJobStatus) -> _StatusResponse: + return _StatusResponse( + PlatformJobStatusResponse( + id="job-a", + name="job-a", + status=PlatformJobStatus(status), + status_details={}, + error_details=None, + steps=[], + created_at=JOB_TIMESTAMP, + updated_at=JOB_TIMESTAMP, + ) + ) + + @pytest.fixture(autouse=True) def _quiet_rich_output() -> Iterator[None]: with ( @@ -71,22 +110,56 @@ def gateway_wait() -> Iterator[MagicMock]: def test_wait_for_platform_job_returns_true_on_completed(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="completed") + jobs.get_job_status.return_value = _status_response("completed") assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is True - jobs.get_status.assert_called_once_with("job-a", workspace="default") + jobs.get_job_status.assert_called_once_with(workspace="default", name="job-a") frozen_time.assert_called() def test_wait_for_platform_job_returns_false_on_error(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="error") + jobs.get_job_status.return_value = _status_response("error") assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is False frozen_time.assert_called() +def test_platform_job_wait_live_display_recomputes_elapsed() -> None: + display = waiters._PlatformJobWaitLiveDisplay(start_time=100.0, timeout=1200, poll_interval=3) + + with ( + patch(f"{WAITERS_MODULE}.datetime") as datetime_mock, + patch(f"{WAITERS_MODULE}.time.time", side_effect=[101.0, 109.0]), + ): + datetime_mock.now.return_value.strftime.return_value = "12:34:56" + + assert "Wait: 1s" in display.__rich__().plain + assert "Wait: 9s" in display.__rich__().plain + + +def test_wait_for_platform_job_uses_dynamic_live_display_for_unchanged_status_polls() -> None: + jobs = MagicMock() + jobs.get_job_status.side_effect = [ + _status_response("active"), + _status_response("active"), + _status_response("completed"), + ] + _RecordingLive.instances = [] + + with ( + patch(f"{WAITERS_MODULE}.Live", _RecordingLive), + patch(f"{WATCH_MODULE}.time.sleep"), + ): + assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is True + + live = _RecordingLive.instances[0] + assert isinstance(live.renderable, waiters._PlatformJobWaitLiveDisplay) + assert live.updates == [live.renderable, live.renderable] + assert jobs.get_job_status.call_count == 3 + + def test_wait_for_inference_deployment_uses_remaining_timeout_for_gateway(gateway_wait: MagicMock) -> None: client = MagicMock() client.inference.deployments.retrieve.return_value = SimpleNamespace( @@ -131,6 +204,27 @@ def test_wait_for_inference_deployment_uses_model_provider_id_for_gateway(gatewa assert gateway_wait.call_args.args[:3] == (client, "generated-provider", "provider-workspace") +def test_wait_for_inference_deployment_quiet_mode_uses_quiet_gateway(gateway_wait: MagicMock) -> None: + client = MagicMock() + client.inference.deployments.retrieve.return_value = SimpleNamespace( + status="READY", + status_message="", + status_history=[], + ) + + with patch(f"{WAITERS_MODULE}.time.time", side_effect=[100.0, 104.0, 104.0, 104.0]): + assert waiters.wait_for_inference_deployment( + client, + "deployment-a", + workspace="default", + timeout=10, + poll_interval=2, + verbose=False, + ) + + assert gateway_wait.call_args.kwargs["verbose"] is False + + def test_wait_for_inference_deployment_retries_transient_status_error( frozen_time: MagicMock, waiter_pause: MagicMock, gateway_wait: MagicMock ) -> None: @@ -204,29 +298,34 @@ def test_wait_for_inference_deployment_does_not_sleep_past_timeout(waiter_pause: waiter_pause.assert_called_once_with(1.0) -def test_wait_for_platform_job_does_not_sleep_past_timeout(waiter_pause: MagicMock) -> None: +def test_wait_for_platform_job_does_not_sleep_past_timeout() -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="running") + jobs.get_job_status.return_value = _status_response("active") - with patch(f"{WAITERS_MODULE}.time.time", side_effect=[0.0, 0.0, 0.0, 4.0, 5.0, 5.0]): + with ( + patch(f"{WAITERS_MODULE}.time.time", return_value=0.0), + patch(f"{WATCH_MODULE}.time.monotonic", side_effect=[0.0, 0.0, 4.0, 5.0]), + patch(f"{WATCH_MODULE}.time.sleep") as watch_sleep, + ): assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", timeout=5, poll_interval=10) is False - waiter_pause.assert_called_once_with(1.0) + watch_sleep.assert_called_once_with(1.0) -def test_wait_for_platform_job_retries_transient_status_error(frozen_time: MagicMock, waiter_pause: MagicMock) -> None: +def test_wait_for_platform_job_retries_transient_status_error(frozen_time: MagicMock) -> None: jobs = MagicMock() request = httpx.Request("GET", "http://test") response = httpx.Response(503, request=request) - jobs.get_status.side_effect = [ + jobs.get_job_status.side_effect = [ APIStatusError("service unavailable", response=response, body=None), - SimpleNamespace(status="completed"), + _status_response("completed"), ] - assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", timeout=10, poll_interval=1) is True + with patch(f"{WATCH_MODULE}.time.sleep") as watch_sleep: + assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", timeout=10, poll_interval=1) is True frozen_time.assert_called() - waiter_pause.assert_called_once_with(1) + watch_sleep.assert_called_once_with(1) def test_wait_for_gateway_does_not_sleep_past_timeout(waiter_pause: MagicMock) -> None: diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/telemetry/test_job_events.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/telemetry/test_job_events.py index ba350b6036..1b92c96a67 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/telemetry/test_job_events.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/telemetry/test_job_events.py @@ -5,15 +5,21 @@ from collections.abc import Iterator from datetime import datetime, timezone -from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from nemo_platform.cli.core import waiters from nemo_platform.cli.telemetry.events import TaskStatusEnum +from nemo_platform_plugin.jobs.schemas import ( + PlatformJobStatus, + PlatformJobStatusResponse, + PlatformJobStepStatusResponse, +) WAITERS_MODULE = "nemo_platform.cli.core.waiters" +WATCH_MODULE = "nemo_platform_plugin.jobs.watch" EMIT_TARGET = "nemo_platform.cli.telemetry.emit.emit_event" +JOB_TIMESTAMP = datetime(2026, 1, 1, tzinfo=timezone.utc) class _DummyLive: @@ -36,6 +42,48 @@ def start(self) -> None: pass +class _StatusResponse: + def __init__(self, status: PlatformJobStatusResponse) -> None: + self._status = status + + def data(self) -> PlatformJobStatusResponse: + return self._status + + +def _step(name: str) -> PlatformJobStepStatusResponse: + return PlatformJobStepStatusResponse( + id=name, + name=name, + status=PlatformJobStatus.COMPLETED, + status_details={}, + error_details=None, + tasks=[], + created_at=JOB_TIMESTAMP, + updated_at=JOB_TIMESTAMP, + ) + + +def _status_response( + status: str | PlatformJobStatus, + *, + steps: list[str] | None = None, + status_details: dict[str, object] | None = None, + created_at: datetime = JOB_TIMESTAMP, +) -> _StatusResponse: + return _StatusResponse( + PlatformJobStatusResponse( + id="job-a", + name="job-a", + status=PlatformJobStatus(status), + status_details=status_details or {}, + error_details=None, + steps=[_step(step) for step in steps or []], + created_at=created_at, + updated_at=created_at, + ) + ) + + @pytest.fixture(autouse=True) def _quiet_rich_output() -> Iterator[None]: with ( @@ -57,14 +105,14 @@ def waiter_pause() -> Iterator[MagicMock]: yield pause -def _completed_status() -> SimpleNamespace: - return SimpleNamespace( - status="completed", +def _completed_status() -> _StatusResponse: + return _status_response( + "completed", steps=[ - SimpleNamespace(name="audit-job"), - SimpleNamespace(name="evaluate"), - SimpleNamespace(name="evaluate-suite"), - SimpleNamespace(name="customer-project-step"), + "audit-job", + "evaluate", + "evaluate-suite", + "customer-project-step", ], status_details={"input_tokens": 512, "output_tokens": 2048, "model": "nemotron-super-49b"}, ) @@ -72,7 +120,7 @@ def _completed_status() -> SimpleNamespace: def test_completed_emits_single_job_run_event(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = _completed_status() + jobs.get_job_status.return_value = _completed_status() with patch(EMIT_TARGET) as emit_event: assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", resource_label="customization") is True @@ -90,9 +138,7 @@ def test_completed_emits_single_job_run_event(frozen_time: MagicMock) -> None: def test_static_step_name_is_not_emitted_as_job_type(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace( - status="completed", steps=[SimpleNamespace(name="audit-job")], status_details={} - ) + jobs.get_job_status.return_value = _status_response("completed", steps=["audit-job"]) with patch(EMIT_TARGET) as emit_event: assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", resource_label="audit") is True @@ -104,7 +150,7 @@ def test_static_step_name_is_not_emitted_as_job_type(frozen_time: MagicMock) -> def test_job_type_falls_back_to_resource_label_without_steps(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="completed", steps=[], status_details={}) + jobs.get_job_status.return_value = _status_response("completed") with patch(EMIT_TARGET) as emit_event: assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", resource_label="customization") is True @@ -116,11 +162,7 @@ def test_job_type_falls_back_to_resource_label_without_steps(frozen_time: MagicM def test_unsafe_resource_label_falls_back_to_custom(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace( - status="completed", - steps=[SimpleNamespace(name="private-customer-step")], - status_details={}, - ) + jobs.get_job_status.return_value = _status_response("completed", steps=["private-customer-step"]) with patch(EMIT_TARGET) as emit_event: assert ( @@ -135,7 +177,7 @@ def test_unsafe_resource_label_falls_back_to_custom(frozen_time: MagicMock) -> N def test_status_details_defaults_when_absent(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="completed", steps=[], status_details={}) + jobs.get_job_status.return_value = _status_response("completed") with patch(EMIT_TARGET) as emit_event: assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is True @@ -150,9 +192,8 @@ def test_status_details_defaults_when_absent(frozen_time: MagicMock) -> None: def test_null_status_details_still_emits(frozen_time: MagicMock) -> None: """Explicit nulls must not drop the event; a real 0 token count must survive.""" jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace( - status="completed", - steps=[], + jobs.get_job_status.return_value = _status_response( + "completed", status_details={"model": None, "input_tokens": 0, "output_tokens": None}, ) @@ -168,9 +209,8 @@ def test_null_status_details_still_emits(frozen_time: MagicMock) -> None: def test_non_string_model_details_still_emit_safe_bucket(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace( - status="completed", - steps=[], + jobs.get_job_status.return_value = _status_response( + "completed", status_details={"model": {"name": "private-model"}, "input_tokens": 7, "output_tokens": 9}, ) @@ -186,15 +226,13 @@ def test_non_string_model_details_still_emit_safe_bucket(frozen_time: MagicMock) def test_duration_uses_job_created_at_when_available() -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace( - status="completed", - steps=[], - status_details={}, + jobs.get_job_status.return_value = _status_response( + "completed", created_at=datetime.fromtimestamp(90.0, tz=timezone.utc), ) with ( - patch(f"{WAITERS_MODULE}.time.time", side_effect=[100.0, 100.0, 100.0, 130.0]), + patch(f"{WAITERS_MODULE}.time.time", side_effect=[100.0, 100.0, 130.0]), patch(EMIT_TARGET) as emit_event, ): assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is True @@ -205,7 +243,7 @@ def test_duration_uses_job_created_at_when_available() -> None: def test_error_status_maps_to_error(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="error", steps=[], status_details={}) + jobs.get_job_status.return_value = _status_response("error") with patch(EMIT_TARGET) as emit_event: assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is False @@ -216,7 +254,7 @@ def test_error_status_maps_to_error(frozen_time: MagicMock) -> None: def test_cancelled_status_maps_to_canceled(frozen_time: MagicMock) -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="cancelled", steps=[], status_details={}) + jobs.get_job_status.return_value = _status_response("cancelled") with patch(EMIT_TARGET) as emit_event: assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default") is False @@ -225,11 +263,16 @@ def test_cancelled_status_maps_to_canceled(frozen_time: MagicMock) -> None: assert emit_event.call_args.args[0].task_status is TaskStatusEnum.CANCELED -def test_timeout_emits_nothing_and_does_not_crash(waiter_pause: MagicMock) -> None: +def test_timeout_emits_nothing_and_does_not_crash() -> None: jobs = MagicMock() - jobs.get_status.return_value = SimpleNamespace(status="running", steps=[], status_details={}) + jobs.get_job_status.return_value = _status_response("active") - with patch(EMIT_TARGET) as emit_event: + with ( + patch(f"{WAITERS_MODULE}.time.time", return_value=0.0), + patch(f"{WATCH_MODULE}.time.monotonic", side_effect=[0.0, 0.0, 4.0, 5.0]), + patch(f"{WATCH_MODULE}.time.sleep"), + patch(EMIT_TARGET) as emit_event, + ): assert waiters.wait_for_platform_job(jobs, "job-a", workspace="default", timeout=5, poll_interval=10) is False emit_event.assert_not_called() diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/test_app.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/test_app.py index d9cab1c087..f2502241ac 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/test_app.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/test_app.py @@ -267,6 +267,36 @@ def test_members_api_command_is_not_registered_at_top_level(): assert "nemo_platform.cli.commands.api.members" not in sys.modules +def test_jobs_watch_command_is_registered(): + runner = CliRunner() + result = runner.invoke(app, ["jobs", "watch", "--help"]) + + assert result.exit_code == 0 + assert "Watch a platform job until it reaches a terminal status." in result.stdout + assert "--history" in result.stdout + assert "--no-history" in result.stdout + + +def test_jobs_create_exposes_wait_and_watch_flags(): + runner = CliRunner() + result = runner.invoke(app, ["jobs", "create", "--help"]) + + assert result.exit_code == 0 + assert "--watch" in result.stdout + assert "--wait" in result.stdout + + +def test_inference_deployments_create_exposes_wait_and_watch_flags(): + runner = CliRunner() + result = runner.invoke(app, ["inference", "deployments", "create", "--help"]) + + assert result.exit_code == 0 + assert "--watch" in result.stdout + assert "--wait" in result.stdout + assert "up and running" in result.stdout + assert "until it is stable" in result.stdout + + def test_root_help_excludes_hidden_commands_and_context_option(): runner = CliRunner() result = runner.invoke(app, ["--help"]) diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/jobs/__init__.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/jobs/__init__.py new file mode 100644 index 0000000000..1275d78dff --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/jobs/__init__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/jobs/test_watch.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/jobs/test_watch.py new file mode 100644 index 0000000000..0df79a2cdd --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/jobs/test_watch.py @@ -0,0 +1,974 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections import deque +from collections.abc import AsyncIterator, Iterable +from datetime import datetime, timezone +from typing import TypedDict, TypeVar + +import httpx +import pytest +from nemo_platform import APIStatusError +from nemo_platform_plugin.client.errors import NemoHTTPError, NemoTransportError +from nemo_platform_plugin.client.response import AsyncNemoPaginatedResponse, NemoPaginatedResponse, NemoResponse +from nemo_platform_plugin.client.types import CursorPagination, PreparedRequest +from nemo_platform_plugin.jobs import watch as watch_module +from nemo_platform_plugin.jobs.client import AsyncJobsClient, JobsClient +from nemo_platform_plugin.jobs.schemas import PlatformJobLog, PlatformJobStatus, PlatformJobStatusResponse +from nemo_platform_plugin.jobs.types import JobLogsQueryParams +from nemo_platform_plugin.jobs.watch import async_watch_job, watch_job +from nemo_platform_plugin.jobs.watch_types import ( + JobLogEvent, + JobStatusEvent, + JobWarningEvent, + JobWatchTimeoutError, +) + +ResponseT = TypeVar("ResponseT") + + +class _StatusCall(TypedDict): + workspace: str | None + name: str + + +class _LogCall(TypedDict): + workspace: str | None + name: str + query_params: JobLogsQueryParams | None + + +def _prepared_request(response_type: type[ResponseT] | None = None) -> PreparedRequest[ResponseT]: + return PreparedRequest( + path_template="/test", + path_params={}, + method="GET", + content=None, + content_type=None, + response_type=response_type, + ) + + +def _status_response(body: PlatformJobStatusResponse) -> NemoResponse[PlatformJobStatusResponse]: + return NemoResponse( + http_response=httpx.Response(200), + body=body, + request=_prepared_request(PlatformJobStatusResponse), + ) + + +def _warning_message(event: JobStatusEvent | JobLogEvent | JobWarningEvent) -> str: + assert isinstance(event, JobWarningEvent) + return event.message + + +def _page_cursor(call: _LogCall) -> str | None: + query_params = call["query_params"] + if query_params is None: + return None + return query_params.get("page_cursor") + + +class _PageResponse( + NemoPaginatedResponse[PlatformJobLog, CursorPagination], + AsyncNemoPaginatedResponse[PlatformJobLog, CursorPagination], +): + def __init__(self, items: list[PlatformJobLog], next_page: str | None = None) -> None: + response = httpx.Response( + 200, + json={ + "data": [item.model_dump(mode="json") for item in items], + "total": len(items), + "next_page": next_page, + "prev_page": None, + }, + ) + super().__init__( + first_http_response=response, + model_type=PlatformJobLog, + request=_prepared_request(), + fetch_page=_unexpected_page_fetch, + strategy=CursorPagination, + ) + + +def _unexpected_page_fetch(_request: PreparedRequest[object], _page: object) -> httpx.Response: + raise AssertionError("Unexpected paginated fetch") + + +class _JobsClientState: + def __init__( + self, + *, + statuses: Iterable[PlatformJobStatusResponse | Exception], + log_results: Iterable[_PageResponse | Exception], + ) -> None: + self._statuses = deque(statuses) + self._last_status: PlatformJobStatusResponse | None = None + self._log_results = deque(log_results) + self.status_calls: list[_StatusCall] = [] + self.log_calls: list[_LogCall] = [] + + def _next_status(self, *, workspace: str | None, name: str) -> PlatformJobStatusResponse: + self.status_calls.append({"workspace": workspace, "name": name}) + if self._statuses: + result = self._statuses.popleft() + if isinstance(result, Exception): + raise result + self._last_status = result + if self._last_status is None: + raise AssertionError("No status result configured") + return self._last_status + + def _next_logs( + self, + *, + workspace: str | None, + name: str, + query_params: JobLogsQueryParams | None, + ) -> _PageResponse: + self.log_calls.append({"workspace": workspace, "name": name, "query_params": query_params}) + if not self._log_results: + return _PageResponse([]) + result = self._log_results.popleft() + if isinstance(result, Exception): + raise result + return result + + +class _SyncJobsClient(_JobsClientState): + def get_job_status(self, *, workspace: str | None = None, name: str) -> NemoResponse[PlatformJobStatusResponse]: + return _status_response(self._next_status(workspace=workspace, name=name)) + + def list_job_logs( + self, + *, + workspace: str | None = None, + name: str, + query_params: JobLogsQueryParams | None = None, + ) -> NemoPaginatedResponse[PlatformJobLog, CursorPagination]: + return self._next_logs(workspace=workspace, name=name, query_params=query_params) + + +class _AsyncJobsClient(_JobsClientState): + async def get_job_status( + self, + *, + workspace: str | None = None, + name: str, + ) -> NemoResponse[PlatformJobStatusResponse]: + return _status_response(self._next_status(workspace=workspace, name=name)) + + async def list_job_logs( + self, + *, + workspace: str | None = None, + name: str, + query_params: JobLogsQueryParams | None = None, + ) -> AsyncNemoPaginatedResponse[PlatformJobLog, CursorPagination]: + return self._next_logs(workspace=workspace, name=name, query_params=query_params) + + +def _status(status: str, status_details: dict[str, object] | None = None) -> PlatformJobStatusResponse: + timestamp = datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc) + return PlatformJobStatusResponse( + id="job-a", + name="job-a", + status=PlatformJobStatus(status), + status_details=status_details or {}, + error_details=None, + steps=[], + created_at=timestamp, + updated_at=timestamp, + ) + + +def _log( + message: str, + *, + timestamp: datetime | None = None, + job_step: str = "step-a", + job_task: str = "task-a", +) -> PlatformJobLog: + return PlatformJobLog( + job="job-a", + timestamp=timestamp or datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc), + job_step=job_step, + job_task=job_task, + message=message, + ) + + +def _record_completed_log_drain( + state: watch_module._WatchState, + logs: list[PlatformJobLog], +) -> list[JobLogEvent]: + state.start_log_drain() + events = list( + watch_module._new_log_events( + logs, + state=state, + occurrence_counts={}, + name="job-a", + emit=True, + ) + ) + state.complete_log_drain() + return events + + +def _invalid_page_cursor_error() -> NemoHTTPError: + request = httpx.Request("GET", "http://test") + response = httpx.Response(422, request=request, json={"detail": "Invalid page cursor"}) + return NemoHTTPError(response) + + +def _http_error_body(status_code: int, body: object) -> NemoHTTPError: + request = httpx.Request("GET", "http://test") + response = httpx.Response(status_code, request=request, json=body) + return NemoHTTPError(response) + + +def _http_text_error(status_code: int, text: str) -> NemoHTTPError: + request = httpx.Request("GET", "http://test") + response = httpx.Response(status_code, request=request, text=text) + return NemoHTTPError(response) + + +def _http_error(status_code: int, detail: str) -> NemoHTTPError: + request = httpx.Request("GET", "http://test") + response = httpx.Response(status_code, request=request, json={"detail": detail}) + return NemoHTTPError(response) + + +def test_can_retry_log_scan_from_start_accepts_invalid_cursor_code() -> None: + exc = _http_error_body( + 422, + {"detail": {"code": "invalid_page_cursor", "message": "The saved cursor is no longer valid"}}, + ) + + assert watch_module._can_retry_log_scan_from_start(exc, "cursor-0") is True + + +def test_can_retry_log_scan_from_start_keeps_page_cursor_decode_compatibility() -> None: + assert watch_module._can_retry_log_scan_from_start(_invalid_page_cursor_error(), "cursor-0") is True + + +def test_can_retry_log_scan_from_start_requires_saved_cursor() -> None: + assert watch_module._can_retry_log_scan_from_start(_invalid_page_cursor_error(), None) is False + + +def test_can_retry_log_scan_from_start_ignores_plain_text_detail_fallback() -> None: + exc = _http_text_error(422, "Invalid page cursor") + + assert watch_module._can_retry_log_scan_from_start(exc, "cursor-0") is False + + +def test_can_retry_log_scan_from_start_rejects_other_422_errors() -> None: + exc = _http_error(422, "Invalid page size") + + assert watch_module._can_retry_log_scan_from_start(exc, "cursor-0") is False + + +def test_watch_job_yields_status_logs_terminal_and_passes_log_query_params() -> None: + client = _SyncJobsClient( + statuses=[ + _status("active", {"phase": "training"}), + _status("completed", {"phase": "done"}), + ], + log_results=[ + _PageResponse([_log("starting")], next_page="cursor-1"), + _PageResponse([_log("still running")]), + _PageResponse([_log("starting"), _log("still running"), _log("done")]), + ], + ) + + events = list( + watch_job( + client, + "job-a", + workspace="default", + poll_interval=0, + attempt_id=1, + step_id="step-1", + task_id="task-1", + limit=2, + page_cursor="cursor-0", + ) + ) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("log", None, "starting"), + ("log", None, "still running"), + ("status", "completed", None), + ("log", None, "done"), + ] + assert isinstance(events[0], JobStatusEvent) + assert events[0].terminal is False + assert events[0].successful is None + assert events[0].status_details == {"phase": "training"} + assert isinstance(events[3], JobStatusEvent) + assert events[3].terminal is True + assert events[3].successful is True + + assert client.status_calls == [ + {"workspace": "default", "name": "job-a"}, + {"workspace": "default", "name": "job-a"}, + ] + assert client.log_calls == [ + { + "workspace": "default", + "name": "job-a", + "query_params": { + "attempt_id": 1, + "step_id": "step-1", + "task_id": "task-1", + "limit": 2, + "page_cursor": "cursor-0", + }, + }, + { + "workspace": "default", + "name": "job-a", + "query_params": { + "attempt_id": 1, + "step_id": "step-1", + "task_id": "task-1", + "limit": 2, + "page_cursor": "cursor-1", + }, + }, + { + "workspace": "default", + "name": "job-a", + "query_params": { + "attempt_id": 1, + "step_id": "step-1", + "task_id": "task-1", + "limit": 2, + "page_cursor": "cursor-1", + }, + }, + ] + + +def test_watch_job_rejects_negative_poll_interval_eagerly() -> None: + client = _SyncJobsClient(statuses=[], log_results=[]) + + with pytest.raises(ValueError, match="poll_interval"): + watch_job(client, "job-a", poll_interval=-1) + + +def test_jobs_client_watch_job_delegates_to_source_owned_watcher(monkeypatch: pytest.MonkeyPatch) -> None: + calls: dict[str, object] = {} + + def fake_watch_job(client: JobsClient, name: str, **kwargs: object) -> Iterable[JobWarningEvent]: + calls["client"] = client + calls["name"] = name + calls["kwargs"] = kwargs + return iter([JobWarningEvent(kind="warning", job_name=name, message="delegated")]) + + monkeypatch.setattr("nemo_platform_plugin.jobs.watch.watch_job", fake_watch_job) + http_client = httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request))) + client = JobsClient(base_url="http://test", http_client=http_client) + + try: + events = list( + client.watch_job( + "job-a", + workspace="default", + poll_interval=0, + timeout=5, + include_history=False, + include_logs=False, + attempt_id=1, + step_id="step-1", + task_id="task-1", + limit=2, + page_cursor="cursor-0", + ) + ) + finally: + http_client.close() + + assert calls == { + "client": client, + "name": "job-a", + "kwargs": { + "workspace": "default", + "poll_interval": 0, + "timeout": 5, + "include_history": False, + "include_logs": False, + "attempt_id": 1, + "step_id": "step-1", + "task_id": "task-1", + "limit": 2, + "page_cursor": "cursor-0", + }, + } + assert events == [JobWarningEvent(kind="warning", job_name="job-a", message="delegated")] + + +def test_watch_job_can_skip_existing_log_history() -> None: + client = _SyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _PageResponse([_log("old")]), + _PageResponse([_log("old"), _log("new")]), + ], + ) + + events = list(watch_job(client, "job-a", include_history=False, poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("status", "completed", None), + ("log", None, "new"), + ] + + +def test_watch_state_replaces_retained_log_window_after_successful_drain() -> None: + state = watch_module._WatchState(history_seen=True, log_cursor=None) + first_log = _log("first") + second_log = _log("second") + + assert [event.message for event in _record_completed_log_drain(state, [first_log])] == ["first"] + assert state.previous_drain_seen_logs == {watch_module._log_key(first_log): 1} + + assert [event.message for event in _record_completed_log_drain(state, [second_log])] == ["second"] + assert state.previous_drain_seen_logs == {watch_module._log_key(second_log): 1} + assert state.current_drain_seen_logs == {} + + +def test_new_log_events_suppresses_seen_occurrences_after_partial_drain_failure() -> None: + state = watch_module._WatchState(history_seen=True, log_cursor=None) + duplicate_log = _log("duplicate") + + state.start_log_drain() + first_events = list( + watch_module._new_log_events( + [duplicate_log, duplicate_log], + state=state, + occurrence_counts={}, + name="job-a", + emit=True, + ) + ) + + state.start_log_drain() + retry_events = list( + watch_module._new_log_events( + [duplicate_log, duplicate_log, duplicate_log], + state=state, + occurrence_counts={}, + name="job-a", + emit=True, + ) + ) + state.complete_log_drain() + + assert [event.message for event in first_events] == ["duplicate", "duplicate"] + assert [event.message for event in retry_events] == ["duplicate"] + assert state.previous_drain_seen_logs == {watch_module._log_key(duplicate_log): 3} + + +def test_watch_job_can_poll_status_without_logs() -> None: + client = _SyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[AssertionError("logs should not be fetched")], + ) + + events = list(watch_job(client, "job-a", include_logs=False, poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None)) for event in events] == [ + ("status", "active"), + ("status", "completed"), + ] + assert client.log_calls == [] + + +def test_watch_job_stops_when_status_is_paused() -> None: + client = _SyncJobsClient( + statuses=[_status("paused")], + log_results=[_PageResponse([_log("paused")])], + ) + + events = list(watch_job(client, "job-a", poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "paused", None), + ("log", None, "paused"), + ] + assert isinstance(events[0], JobStatusEvent) + assert events[0].terminal is True + assert events[0].successful is False + assert len(client.status_calls) == 1 + assert len(client.log_calls) == 1 + + +def test_watch_job_retries_sdk_transient_status_error() -> None: + request = httpx.Request("GET", "http://test") + response = httpx.Response(503, request=request) + client = _SyncJobsClient( + statuses=[ + APIStatusError("service unavailable", response=response, body=None), + _status("completed"), + ], + log_results=[AssertionError("logs should not be fetched")], + ) + + events = list(watch_job(client, "job-a", include_logs=False, poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("warning", None, "Transient status check failed: service unavailable"), + ("status", "completed", None), + ] + assert client.log_calls == [] + + +def test_watch_job_backs_off_and_deduplicates_consecutive_transient_status_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + request = httpx.Request("GET", "http://test") + response = httpx.Response(503, request=request) + client = _SyncJobsClient( + statuses=[ + APIStatusError("service unavailable", response=response, body=None), + APIStatusError("service unavailable", response=response, body=None), + APIStatusError("service unavailable", response=response, body=None), + APIStatusError("service unavailable", response=response, body=None), + APIStatusError("service unavailable", response=response, body=None), + APIStatusError("service unavailable", response=response, body=None), + APIStatusError("service unavailable", response=response, body=None), + _status("completed"), + ], + log_results=[AssertionError("logs should not be fetched")], + ) + sleeps: list[float] = [] + monkeypatch.setattr(watch_module.time, "sleep", sleeps.append) + + events = list(watch_job(client, "job-a", include_logs=False, poll_interval=1)) + + assert sleeps == [1, 2, 4, 8, 16, 30.0, 30.0] + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("warning", None, "Transient status check failed: service unavailable"), + ("status", "completed", None), + ] + assert client.log_calls == [] + + +def test_watch_job_suppresses_unread_history_after_partial_history_drain_failure() -> None: + client = _SyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _PageResponse([_log("old")], next_page="cursor-1"), + NemoTransportError(httpx.TransportError("temporary log failure")), + _PageResponse([_log("old"), _log("also-old")]), + ], + ) + + events = list(watch_job(client, "job-a", include_history=False, poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("warning", None, "Transient log check failed: temporary log failure"), + ("status", "completed", None), + ] + + +def test_watch_job_continues_status_polling_after_non_retryable_log_http_error() -> None: + client = _SyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _http_error(500, "log store unavailable"), + _PageResponse([_log("done")]), + ], + ) + + events = list(watch_job(client, "job-a", poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("warning", None, "Log check failed: HTTP 500: log store unavailable"), + ("status", "completed", None), + ("log", None, "done"), + ] + + +def test_watch_job_retries_terminal_status_until_logs_drain_after_transient_failure() -> None: + client = _SyncJobsClient( + statuses=[_status("completed"), _status("completed")], + log_results=[ + NemoTransportError(httpx.TransportError("temporary log failure")), + _PageResponse([_log("done")]), + ], + ) + + events = list(watch_job(client, "job-a", poll_interval=0)) + + assert [event.kind for event in events] == ["status", "warning", "log"] + assert _warning_message(events[1]) == "Transient log check failed: temporary log failure" + assert isinstance(events[2], JobLogEvent) + assert events[2].message == "done" + + +def test_watch_job_stops_after_terminal_log_drain_retry_cap() -> None: + retry_cap = watch_module._TerminalLogDrainRetryBudget.RETRY_CAP + client = _SyncJobsClient( + statuses=[_status("completed")], + log_results=[ + NemoTransportError(httpx.TransportError(f"temporary log failure {attempt}")) for attempt in range(retry_cap) + ], + ) + + events = list(watch_job(client, "job-a", poll_interval=0)) + + assert [event.kind for event in events] == ["status"] + ["warning"] * (retry_cap + 1) + assert [getattr(event, "message", None) for event in events[1:-1]] == [ + f"Transient log check failed: temporary log failure {attempt}" for attempt in range(retry_cap) + ] + assert _warning_message(events[-1]) == ( + f"Terminal log drain retry cap reached ({retry_cap}); stopping watch for job 'job-a'" + ) + assert len(client.status_calls) == retry_cap + assert len(client.log_calls) == retry_cap + + +def test_watch_job_falls_back_to_full_rescan_when_saved_cursor_is_invalid() -> None: + client = _SyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _PageResponse([_log("old")], next_page="cursor-1"), + _PageResponse([_log("new")]), + _invalid_page_cursor_error(), + _PageResponse([_log("old"), _log("new"), _log("done")]), + ], + ) + + events = list(watch_job(client, "job-a", poll_interval=0)) + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("log", None, "old"), + ("log", None, "new"), + ("status", "completed", None), + ("log", None, "done"), + ] + assert [_page_cursor(call) for call in client.log_calls] == [None, "cursor-1", "cursor-1", None] + + +def test_watch_job_stops_log_pagination_when_cursor_does_not_advance() -> None: + client = _SyncJobsClient( + statuses=[_status("completed")], + log_results=[_PageResponse([_log("done")], next_page="cursor-0")], + ) + + events = list(watch_job(client, "job-a", poll_interval=0, page_cursor="cursor-0")) + + assert [(event.kind, getattr(event, "message", None)) for event in events] == [ + ("status", None), + ("log", "done"), + ] + assert [_page_cursor(call) for call in client.log_calls] == ["cursor-0"] + + +def test_watch_job_enforces_timeout_between_log_pages(monkeypatch: pytest.MonkeyPatch) -> None: + client = _SyncJobsClient( + statuses=[_status("completed")], + log_results=[ + _PageResponse([], next_page="cursor-1"), + _PageResponse([]), + ], + ) + monotonic_values = iter([0.0, 0.0, 1.0, 10.0]) + monkeypatch.setattr(watch_module.time, "monotonic", lambda: next(monotonic_values, 10.0)) + + with pytest.raises(JobWatchTimeoutError, match="job-a"): + list(watch_job(client, "job-a", timeout=5, poll_interval=0)) + + assert len(client.log_calls) == 1 + + +async def test_async_watch_job_suppresses_unread_history_after_partial_history_drain_failure() -> None: + client = _AsyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _PageResponse([_log("old")], next_page="cursor-1"), + NemoTransportError(httpx.TransportError("temporary log failure")), + _PageResponse([_log("old"), _log("also-old")]), + ], + ) + + events = [event async for event in async_watch_job(client, "job-a", include_history=False, poll_interval=0)] + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("warning", None, "Transient log check failed: temporary log failure"), + ("status", "completed", None), + ] + + +async def test_async_watch_job_backs_off_and_deduplicates_consecutive_transient_log_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = _AsyncJobsClient( + statuses=[_status("active"), _status("active"), _status("completed")], + log_results=[ + NemoTransportError(httpx.TransportError("temporary log failure")), + NemoTransportError(httpx.TransportError("temporary log failure")), + _PageResponse([_log("done")]), + ], + ) + sleeps: list[float] = [] + + async def fake_sleep(sleep_for: float) -> None: + sleeps.append(sleep_for) + + monkeypatch.setattr(watch_module.asyncio, "sleep", fake_sleep) + + events = [event async for event in async_watch_job(client, "job-a", poll_interval=1)] + + assert sleeps == [1, 2] + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("warning", None, "Transient log check failed: temporary log failure"), + ("status", "completed", None), + ("log", None, "done"), + ] + + +async def test_async_watch_job_continues_status_polling_after_non_retryable_log_http_error() -> None: + client = _AsyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _http_error(500, "log store unavailable"), + _PageResponse([_log("done")]), + ], + ) + + events = [event async for event in async_watch_job(client, "job-a", poll_interval=0)] + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("warning", None, "Log check failed: HTTP 500: log store unavailable"), + ("status", "completed", None), + ("log", None, "done"), + ] + + +async def test_async_watch_job_can_poll_status_without_logs() -> None: + client = _AsyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[AssertionError("logs should not be fetched")], + ) + + events = [event async for event in async_watch_job(client, "job-a", include_logs=False, poll_interval=0)] + + assert [(event.kind, getattr(event, "status", None)) for event in events] == [ + ("status", "active"), + ("status", "completed"), + ] + assert client.log_calls == [] + + +async def test_async_watch_job_stops_when_status_is_paused() -> None: + client = _AsyncJobsClient( + statuses=[_status("paused")], + log_results=[_PageResponse([_log("paused")])], + ) + + events = [event async for event in async_watch_job(client, "job-a", poll_interval=0)] + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "paused", None), + ("log", None, "paused"), + ] + assert isinstance(events[0], JobStatusEvent) + assert events[0].terminal is True + assert events[0].successful is False + assert len(client.status_calls) == 1 + assert len(client.log_calls) == 1 + + +async def test_async_watch_job_falls_back_to_full_rescan_when_saved_cursor_is_invalid() -> None: + client = _AsyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _PageResponse([_log("old")], next_page="cursor-1"), + _PageResponse([_log("new")]), + _invalid_page_cursor_error(), + _PageResponse([_log("old"), _log("new"), _log("done")]), + ], + ) + + events = [event async for event in async_watch_job(client, "job-a", poll_interval=0)] + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("log", None, "old"), + ("log", None, "new"), + ("status", "completed", None), + ("log", None, "done"), + ] + assert [_page_cursor(call) for call in client.log_calls] == [None, "cursor-1", "cursor-1", None] + + +async def test_async_watch_job_stops_log_pagination_when_cursor_repeats() -> None: + client = _AsyncJobsClient( + statuses=[_status("completed")], + log_results=[ + _PageResponse([_log("one")], next_page="cursor-1"), + _PageResponse([_log("two")], next_page="cursor-0"), + ], + ) + + events = [event async for event in async_watch_job(client, "job-a", poll_interval=0, page_cursor="cursor-0")] + + assert [(event.kind, getattr(event, "message", None)) for event in events] == [ + ("status", None), + ("log", "one"), + ("log", "two"), + ] + assert [_page_cursor(call) for call in client.log_calls] == ["cursor-0", "cursor-1"] + + +async def test_async_watch_job_enforces_timeout_between_log_pages(monkeypatch: pytest.MonkeyPatch) -> None: + client = _AsyncJobsClient( + statuses=[_status("completed")], + log_results=[ + _PageResponse([], next_page="cursor-1"), + _PageResponse([]), + ], + ) + monotonic_values = iter([0.0, 0.0, 1.0, 10.0]) + monkeypatch.setattr(watch_module.time, "monotonic", lambda: next(monotonic_values, 10.0)) + + with pytest.raises(JobWatchTimeoutError, match="job-a"): + [event async for event in async_watch_job(client, "job-a", timeout=5, poll_interval=0)] + + assert len(client.log_calls) == 1 + + +async def test_async_watch_job_stops_after_terminal_log_drain_retry_cap() -> None: + retry_cap = watch_module._TerminalLogDrainRetryBudget.RETRY_CAP + client = _AsyncJobsClient( + statuses=[_status("completed")], + log_results=[ + NemoTransportError(httpx.TransportError(f"temporary log failure {attempt}")) for attempt in range(retry_cap) + ], + ) + + events = [event async for event in async_watch_job(client, "job-a", poll_interval=0)] + + assert [event.kind for event in events] == ["status"] + ["warning"] * (retry_cap + 1) + assert [getattr(event, "message", None) for event in events[1:-1]] == [ + f"Transient log check failed: temporary log failure {attempt}" for attempt in range(retry_cap) + ] + assert _warning_message(events[-1]) == ( + f"Terminal log drain retry cap reached ({retry_cap}); stopping watch for job 'job-a'" + ) + assert len(client.status_calls) == retry_cap + assert len(client.log_calls) == retry_cap + + +def test_watch_job_raises_timeout_with_job_name(monkeypatch: pytest.MonkeyPatch) -> None: + client = _SyncJobsClient(statuses=[_status("active")], log_results=[_PageResponse([])]) + monotonic_values = iter([0.0, 0.0, 10.0]) + monkeypatch.setattr(watch_module.time, "monotonic", lambda: next(monotonic_values, 10.0)) + + with pytest.raises(JobWatchTimeoutError, match="job-a"): + list(watch_job(client, "job-a", timeout=5, poll_interval=1)) + + +async def test_async_watch_job_is_async_iterator_and_uses_async_jobs_client() -> None: + client = _AsyncJobsClient( + statuses=[_status("active"), _status("completed")], + log_results=[ + _PageResponse([_log("starting")]), + _PageResponse([_log("starting"), _log("done")]), + ], + ) + + iterator = async_watch_job(client, "job-a", workspace="default", poll_interval=0) + assert hasattr(iterator, "__aiter__") + + events = [event async for event in iterator] + + assert [(event.kind, getattr(event, "status", None), getattr(event, "message", None)) for event in events] == [ + ("status", "active", None), + ("log", None, "starting"), + ("status", "completed", None), + ("log", None, "done"), + ] + assert client.status_calls == [ + {"workspace": "default", "name": "job-a"}, + {"workspace": "default", "name": "job-a"}, + ] + + +def test_async_watch_job_rejects_negative_poll_interval_eagerly() -> None: + client = _AsyncJobsClient(statuses=[], log_results=[]) + + with pytest.raises(ValueError, match="poll_interval"): + async_watch_job(client, "job-a", poll_interval=-1) + + +async def test_async_jobs_client_watch_job_delegates_to_source_owned_watcher( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: dict[str, object] = {} + + async def fake_events(job_name: str) -> AsyncIterator[JobWarningEvent]: + yield JobWarningEvent(kind="warning", job_name=job_name, message="delegated") + + def fake_async_watch_job(client: AsyncJobsClient, name: str, **kwargs: object) -> AsyncIterator[JobWarningEvent]: + calls["client"] = client + calls["name"] = name + calls["kwargs"] = kwargs + return fake_events(name) + + monkeypatch.setattr("nemo_platform_plugin.jobs.watch.async_watch_job", fake_async_watch_job) + http_client = httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request))) + client = AsyncJobsClient(base_url="http://test", http_client=http_client) + + try: + events = [ + event + async for event in client.watch_job( + "job-a", + workspace="default", + poll_interval=0, + timeout=5, + include_history=False, + include_logs=False, + attempt_id=1, + step_id="step-1", + task_id="task-1", + limit=2, + page_cursor="cursor-0", + ) + ] + finally: + await http_client.aclose() + + assert calls == { + "client": client, + "name": "job-a", + "kwargs": { + "workspace": "default", + "poll_interval": 0, + "timeout": 5, + "include_history": False, + "include_logs": False, + "attempt_id": 1, + "step_id": "step-1", + "task_id": "task-1", + "limit": 2, + "page_cursor": "cursor-0", + }, + } + assert events == [JobWarningEvent(kind="warning", job_name="job-a", message="delegated")] diff --git a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/cli_config.yaml b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/cli_config.yaml index c473774a62..4d95f61dfd 100644 --- a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/cli_config.yaml +++ b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/cli_config.yaml @@ -166,7 +166,17 @@ config: - description - created_at - resource: [jobs] + additional_methods: + watch: + override: jobs/watch.py methods: + create: + wait: + type: platform_job + resource_label: job + watch: + type: platform_job + resource_label: job list: columns: - name @@ -218,6 +228,9 @@ config: wait: type: inference_deployment resource_label: deployment + watch: + type: inference_deployment + resource_label: deployment list: columns: - name diff --git a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/config.py b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/config.py index 12bc338184..9dad5c90f5 100644 --- a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/config.py +++ b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/config.py @@ -13,6 +13,7 @@ from nemo_platform_sdk_tools.sdk.core.common import get_project_dir VALID_WAIT_CONFIG_TYPES = {"inference_deployment", "platform_job"} +VALID_WATCH_CONFIG_TYPES = {"inference_deployment", "platform_job"} def get_cli_generator_root() -> Path: @@ -210,39 +211,63 @@ def should_skip_method(self, resource_path: list[str], method_name: str) -> bool def get_wait_config(self, resource_path: list[str], method_name: str) -> dict[str, object] | None: """Get inline wait configuration for a generated command.""" + return self._get_lifecycle_config( + resource_path, + method_name, + config_key="wait", + valid_types=VALID_WAIT_CONFIG_TYPES, + ) + + def get_watch_config(self, resource_path: list[str], method_name: str) -> dict[str, object] | None: + """Get inline watch configuration for a generated command.""" + return self._get_lifecycle_config( + resource_path, + method_name, + config_key="watch", + valid_types=VALID_WATCH_CONFIG_TYPES, + ) + + def _get_lifecycle_config( + self, + resource_path: list[str], + method_name: str, + *, + config_key: str, + valid_types: set[str], + ) -> dict[str, object] | None: if method_config := self.get_method_config(resource_path, method_name): - wait_config = method_config.get("wait") - if wait_config is None: + lifecycle_config = method_config.get(config_key) + if lifecycle_config is None: return None resource = ".".join(resource_path) - if not isinstance(wait_config, dict): + if not isinstance(lifecycle_config, dict): raise ValueError( - f"Invalid wait config for {resource}.{method_name}. Expected a mapping, got " - f"{type(wait_config).__name__}." + f"Invalid {config_key} config for {resource}.{method_name}. Expected a mapping, got " + f"{type(lifecycle_config).__name__}." ) - wait_type = wait_config.get("type") - valid_types = ", ".join(sorted(VALID_WAIT_CONFIG_TYPES)) + lifecycle_type = lifecycle_config.get("type") + valid_types_list = ", ".join(sorted(valid_types)) try: - wait_type_is_valid = wait_type in VALID_WAIT_CONFIG_TYPES + lifecycle_type_is_valid = lifecycle_type in valid_types except TypeError as exc: raise ValueError( - f"Invalid wait config wait_type={wait_type!r} for resource_path={resource_path!r}, " - f"method_name={method_name!r}. Expected one of VALID_WAIT_CONFIG_TYPES: {valid_types}" + f"Invalid {config_key} config type={lifecycle_type!r} for resource_path={resource_path!r}, " + f"method_name={method_name!r}. Expected one of: {valid_types_list}" ) from exc - if not wait_type_is_valid: + if not lifecycle_type_is_valid: raise ValueError( - f"Invalid wait config type {wait_type!r} for {resource}.{method_name}. " - f"Expected one of: {valid_types}" + f"Invalid {config_key} config type {lifecycle_type!r} for {resource}.{method_name}. " + f"Expected one of: {valid_types_list}" ) - resource_label = wait_config.get("resource_label") + resource_label = lifecycle_config.get("resource_label") if not isinstance(resource_label, str) or not resource_label.strip(): raise ValueError( - f"Invalid wait config resource_label {resource_label!r} for {resource}.{method_name}. " + f"Invalid {config_key} config resource_label {resource_label!r} for {resource}.{method_name}. " "Expected a non-empty string." ) - return wait_config + return lifecycle_config return None def get_additional_methods(self, resource_path: list[str]) -> dict[str, Any]: diff --git a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/context_collectors/create_collector.py b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/context_collectors/create_collector.py index 63e19dd00f..ed1c4e29ff 100644 --- a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/context_collectors/create_collector.py +++ b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/context_collectors/create_collector.py @@ -92,4 +92,5 @@ def collect( "required_fields": required_fields, "required_fields_example": build_required_fields_example(parameters, required_fields), "wait_config": self._cli_config.get_wait_config(resource_path, method_name), + "watch_config": self._cli_config.get_watch_config(resource_path, method_name), } diff --git a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/overrides/jobs/watch.py b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/overrides/jobs/watch.py new file mode 100644 index 0000000000..c5a0b2c7fc --- /dev/null +++ b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/overrides/jobs/watch.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Annotated, Any, cast + +import typer +from nemo_platform_ext.cli.core.context import CLIContext +from nemo_platform_ext.cli.core.errors import handle_errors +from nemo_platform_ext.cli.core.help_formatter import collect_warnings +from nemo_platform_ext.cli.core.job_watch_renderer import render_job_watch_events +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.jobs.client import JobsClient + +app = cast(Any, None) # override-skip: provided by generated file + + +@app.command("watch") +@collect_warnings +@handle_errors +def watch_platform_job( + ctx: typer.Context, + name: Annotated[str, typer.Argument(help="Name of the platform job to watch")], + workspace: Annotated[str | None, typer.Option("--workspace", help="Workspace containing the job")] = None, + attempt_id: Annotated[int | None, typer.Option("--attempt-id", help="Filter logs to an attempt ID")] = None, + step_id: Annotated[str | None, typer.Option("--step-id", help="Filter logs to a step ID")] = None, + task_id: Annotated[str | None, typer.Option("--task-id", help="Filter logs to a task ID")] = None, + limit: Annotated[int | None, typer.Option("--limit", min=1, help="Maximum logs to fetch per page")] = None, + timeout: Annotated[int | None, typer.Option("--timeout", min=1, help="Maximum watch time in seconds")] = None, + poll_interval: Annotated[ + int, + typer.Option("--poll-interval", min=1, help="Seconds between status checks"), + ] = 3, + include_history: Annotated[ + bool, + typer.Option("--history/--no-history", help="Include logs already present before watching"), + ] = True, +) -> None: + """Watch a platform job until it reaches a terminal status.""" + state: CLIContext = ctx.obj + client = state.get_client() + jobs_client = client_from_platform(client, JobsClient) + if workspace is None: + workspace = client._get_workspace_path_param() + + events = jobs_client.watch_job( + name, + workspace=workspace, + attempt_id=attempt_id, + step_id=step_id, + task_id=task_id, + limit=limit, + timeout=timeout, + poll_interval=poll_interval, + include_history=include_history, + ) + if not render_job_watch_events(events, resource_label="job"): + raise typer.Exit(1) diff --git a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/templates/create_command.py.j2 b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/templates/create_command.py.j2 index 0ca6673817..61cba64eb6 100644 --- a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/templates/create_command.py.j2 +++ b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/cli_generator/templates/create_command.py.j2 @@ -2,6 +2,15 @@ from typing import Annotated, Literal import typer +{% set has_wait_config = wait_config is defined and wait_config is not none %} +{% set has_watch_config = watch_config is defined and watch_config is not none %} +{% set wait_type = wait_config.type if has_wait_config else none %} +{% set watch_type = watch_config.type if has_watch_config else none %} +{% set has_lifecycle_config = has_wait_config or has_watch_config %} +{% set has_platform_job_wait = wait_type == "platform_job" %} +{% set has_platform_job_watch = watch_type == "platform_job" %} +{% set has_inference_deployment_wait = wait_type == "inference_deployment" %} +{% set has_inference_deployment_watch = watch_type == "inference_deployment" %} from nemo_platform_ext.cli.core.code_generator import handle_code_generation from nemo_platform_ext.cli.core.context import CLIContext from nemo_platform_ext.cli.core.errors import handle_errors @@ -9,11 +18,19 @@ from nemo_platform_ext.cli.core.formatters import format_output from nemo_platform_ext.cli.core.help_formatter import collect_warnings from nemo_platform_ext.cli.core.stdin_utils import read_data_input_with_flags, read_payload, validate_required_fields from nemo_platform_ext.cli.core.types import EntityOutputFormatOption -{% if wait_config and wait_config.type == "inference_deployment" %} +{% if has_inference_deployment_wait or has_inference_deployment_watch %} from nemo_platform_ext.cli.core.waiters import wait_for_inference_deployment -{% elif wait_config and wait_config.type == "platform_job" %} +{% endif %} +{% if has_platform_job_watch %} +from nemo_platform_ext.cli.core.job_watch_renderer import render_job_watch_events +{% endif %} +{% if has_platform_job_wait %} from nemo_platform_ext.cli.core.waiters import wait_for_platform_job {% endif %} +{% if has_platform_job_wait or has_platform_job_watch %} +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.jobs.client import JobsClient +{% endif %} @app.command("{{ cli_command_name }}") @@ -27,10 +44,29 @@ def {{ function_name }}( {% endfor %}input_file: Annotated[str | None, typer.Option("--input-file", help="Path to JSON file (use '-' for stdin)", rich_help_panel="Input Options")] = None, input_data: Annotated[str | None, typer.Option("--input-data", help="Input data for the request (JSON or YAML)", rich_help_panel="Input Options")] = None, output_format: EntityOutputFormatOption = None, -{% if wait_config %} - wait: Annotated[bool, typer.Option("--wait", help={{ ("Wait for the created " ~ wait_config.resource_label ~ " to reach a terminal state")|repr }}, rich_help_panel="Wait Options")] = False, - timeout: Annotated[int, typer.Option("--timeout", min=1, help="Maximum time to wait in seconds", rich_help_panel="Wait Options")] = 1200, - poll_interval: Annotated[int, typer.Option("--poll-interval", min=1, help="Seconds between status checks", rich_help_panel="Wait Options")] = 3, +{% if has_lifecycle_config %} +{% if has_wait_config %} +{% if has_platform_job_wait %} + wait: Annotated[bool, typer.Option("--wait", help={{ ("Wait for the created " ~ wait_config.resource_label ~ " to reach a terminal state without streaming logs")|repr }}, rich_help_panel="Lifecycle Options")] = False, +{% elif has_inference_deployment_wait %} + wait: Annotated[bool, typer.Option("--wait", help={{ ("Wait for the created " ~ wait_config.resource_label ~ " to be up and running")|repr }}, rich_help_panel="Lifecycle Options")] = False, +{% else %} + wait: Annotated[bool, typer.Option("--wait", help={{ ("Wait for the created " ~ wait_config.resource_label ~ " to reach a terminal state")|repr }}, rich_help_panel="Lifecycle Options")] = False, +{% endif %} +{% endif %} +{% if has_watch_config %} +{% if has_inference_deployment_watch %} + watch: Annotated[bool, typer.Option("--watch", help={{ ("Watch the created " ~ watch_config.resource_label ~ " until it is stable, streaming status updates")|repr }}, rich_help_panel="Lifecycle Options")] = False, +{% else %} + watch: Annotated[bool, typer.Option("--watch", help={{ ("Watch the created " ~ watch_config.resource_label ~ " to a terminal state")|repr }}, rich_help_panel="Lifecycle Options")] = False, +{% endif %} +{% endif %} +{% if has_platform_job_watch %} + timeout: Annotated[int | None, typer.Option("--timeout", min=1, help="Maximum time to wait or watch in seconds", rich_help_panel="Lifecycle Options")] = None, +{% else %} + timeout: Annotated[int, typer.Option("--timeout", min=1, help="Maximum time to wait or watch in seconds", rich_help_panel="Lifecycle Options")] = 1200, +{% endif %} + poll_interval: Annotated[int, typer.Option("--poll-interval", min=1, help="Seconds between status checks", rich_help_panel="Lifecycle Options")] = 3, {% endif %} ) -> None: """{{ help_text }} @@ -84,15 +120,30 @@ def {{ function_name }}( state: CLIContext = ctx.obj output_format = state.get_output_format(output_format) -{% if wait_config %} +{% if has_wait_config and has_watch_config %} + if wait and watch: + raise typer.BadParameter("Cannot combine --wait and --watch.") + +{% endif %} +{% if has_lifecycle_config %} if handle_code_generation( [{{ resource_path_quoted }}], "{{ method_name }}", all_kwargs, output_format, state, + {% if has_watch_config %} + watch_config={"type": {{ watch_config.type|repr }}, "resource_label": {{ watch_config.resource_label|repr }}} if watch else None, + watch_options={"timeout": timeout, "poll_interval": poll_interval} if watch else None, + {% endif %} + {% if has_wait_config %} wait_config={"type": {{ wait_config.type|repr }}, "resource_label": {{ wait_config.resource_label|repr }}} if wait else None, + {% if has_platform_job_watch %} + wait_options={"timeout": timeout if timeout is not None else 1200, "poll_interval": poll_interval} if wait else None, + {% else %} wait_options={"timeout": timeout, "poll_interval": poll_interval} if wait else None, + {% endif %} + {% endif %} ): return {% else %} @@ -102,7 +153,7 @@ def {{ function_name }}( client = state.get_client() result = client.{{ sdk_accessor }}.{{ method_name }}(**all_kwargs) - +{% if has_platform_job_wait or has_platform_job_watch %} format_output( result, is_list=False, @@ -110,31 +161,100 @@ def {{ function_name }}( no_truncate=state.get_no_truncate(), timestamp_format=state.get_timestamp_format(), ) -{% if wait_config %} - - if wait: +{% endif %} +{% if has_platform_job_wait or has_platform_job_watch %} +{% if has_platform_job_wait and has_platform_job_watch %} + if wait or watch: wait_name = getattr(result, "name", None) or all_kwargs.get("name") if not wait_name: - raise RuntimeError("Unable to determine created resource name for --wait") - wait_workspace = all_kwargs.get("workspace") - {% if wait_config.type == "inference_deployment" %} - if not wait_for_inference_deployment( - client, + raise RuntimeError("Unable to determine created resource name for --wait/--watch") + wait_workspace = getattr(result, "workspace", None) or all_kwargs.get("workspace") + if wait_workspace is None: + wait_workspace = client._get_workspace_path_param() + jobs_client = client_from_platform(client, JobsClient) + if wait: + if not wait_for_platform_job( + jobs_client, + wait_name, + workspace=wait_workspace, + resource_label={{ wait_config.resource_label|repr }}, + timeout=timeout if timeout is not None else 1200, + poll_interval=poll_interval, + ): + raise typer.Exit(1) + return + events = jobs_client.watch_job( wait_name, workspace=wait_workspace, timeout=timeout, poll_interval=poll_interval, - ): + ) + if not render_job_watch_events(events, resource_label={{ watch_config.resource_label|repr }}): raise typer.Exit(1) - {% elif wait_config.type == "platform_job" %} + return +{% elif has_platform_job_wait %} + if wait: + wait_name = getattr(result, "name", None) or all_kwargs.get("name") + if not wait_name: + raise RuntimeError("Unable to determine created resource name for --wait") + wait_workspace = getattr(result, "workspace", None) or all_kwargs.get("workspace") + if wait_workspace is None: + wait_workspace = client._get_workspace_path_param() + jobs_client = client_from_platform(client, JobsClient) if not wait_for_platform_job( - client.{{ sdk_accessor }}, + jobs_client, wait_name, workspace=wait_workspace, resource_label={{ wait_config.resource_label|repr }}, + timeout=timeout if timeout is not None else 1200, + poll_interval=poll_interval, + ): + raise typer.Exit(1) + return +{% elif has_platform_job_watch %} + if watch: + wait_name = getattr(result, "name", None) or all_kwargs.get("name") + if not wait_name: + raise RuntimeError("Unable to determine created resource name for --watch") + wait_workspace = getattr(result, "workspace", None) or all_kwargs.get("workspace") + if wait_workspace is None: + wait_workspace = client._get_workspace_path_param() + jobs_client = client_from_platform(client, JobsClient) + events = jobs_client.watch_job( + wait_name, + workspace=wait_workspace, + timeout=timeout, + poll_interval=poll_interval, + ) + if not render_job_watch_events(events, resource_label={{ watch_config.resource_label|repr }}): + raise typer.Exit(1) + return +{% endif %} + +{% endif %} +{% if not (has_platform_job_wait or has_platform_job_watch) %} + + format_output( + result, + is_list=False, + output_format=output_format, + no_truncate=state.get_no_truncate(), + timestamp_format=state.get_timestamp_format(), + ) +{% endif %} +{% if has_inference_deployment_wait or has_inference_deployment_watch %} + if {% if has_inference_deployment_wait and has_inference_deployment_watch %}wait or watch{% elif has_inference_deployment_wait %}wait{% else %}watch{% endif %}: + wait_name = getattr(result, "name", None) or all_kwargs.get("name") + if not wait_name: + raise RuntimeError("Unable to determine created resource name for {% if has_inference_deployment_wait and has_inference_deployment_watch %}--wait/--watch{% elif has_inference_deployment_wait %}--wait{% else %}--watch{% endif %}") + wait_workspace = all_kwargs.get("workspace") + if not wait_for_inference_deployment( + client, + wait_name, + workspace=wait_workspace, timeout=timeout, poll_interval=poll_interval, ): raise typer.Exit(1) - {% endif %} + return {% endif %} diff --git a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/vendor/vendor_package.py b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/vendor/vendor_package.py index 64eb931554..99be461c6f 100644 --- a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/vendor/vendor_package.py +++ b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/vendor/vendor_package.py @@ -116,7 +116,10 @@ def __init__(self, methods: dict[str, dict[str, cst.FunctionDef]]) -> None: def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef: class_name = original_node.name.value - replacements = self._methods.get(class_name) + if class_name not in _CLIENT_CLASS_NAMES: + return updated_node + + replacements = self._methods.get(class_name, {}) if not replacements: return updated_node diff --git a/tools/nemo-platform-sdk-tools/tests/sdk/cli_generator/test_config.py b/tools/nemo-platform-sdk-tools/tests/sdk/cli_generator/test_config.py index cbe7d4ebe2..0e55838e3c 100644 --- a/tools/nemo-platform-sdk-tools/tests/sdk/cli_generator/test_config.py +++ b/tools/nemo-platform-sdk-tools/tests/sdk/cli_generator/test_config.py @@ -350,16 +350,16 @@ def _create_config(self, config_yaml: str) -> CLIConfig: def test_returns_configured_wait(self): config = self._create_config(""" config: - - resource: [customization, jobs] + - resource: [inference, deployments] methods: create: wait: - type: platform_job - resource_label: customization job + type: inference_deployment + resource_label: deployment """) - assert config.get_wait_config(["customization", "jobs"], "create") == { - "type": "platform_job", - "resource_label": "customization job", + assert config.get_wait_config(["inference", "deployments"], "create") == { + "type": "inference_deployment", + "resource_label": "deployment", } def test_returns_none_when_not_configured(self): @@ -395,6 +395,22 @@ def test_rejects_unknown_wait_config_type(self): with pytest.raises(ValueError, match="Invalid wait config type 'unknown'"): config.get_wait_config(["customization", "jobs"], "create") + def test_returns_platform_job_wait_config(self): + config = self._create_config(""" +config: + - resource: [customization, jobs] + methods: + create: + wait: + type: platform_job + resource_label: customization job +""") + + assert config.get_wait_config(["customization", "jobs"], "create") == { + "type": "platform_job", + "resource_label": "customization job", + } + def test_rejects_unhashable_wait_config_type(self): config = self._create_config(""" config: @@ -408,7 +424,7 @@ def test_rejects_unhashable_wait_config_type(self): with pytest.raises( ValueError, - match=r"wait_type=\['platform_job'\].*resource_path=\['customization', 'jobs'\].*method_name='create'.*VALID_WAIT_CONFIG_TYPES", + match=r"type=\['platform_job'\].*resource_path=\['customization', 'jobs'\].*method_name='create'.*Expected one of:", ): config.get_wait_config(["customization", "jobs"], "create") @@ -419,7 +435,7 @@ def test_rejects_missing_wait_resource_label(self): methods: create: wait: - type: platform_job + type: inference_deployment """) with pytest.raises(ValueError, match="Invalid wait config resource_label None"): @@ -432,7 +448,7 @@ def test_rejects_empty_wait_resource_label(self): methods: create: wait: - type: platform_job + type: inference_deployment resource_label: " " """) @@ -440,6 +456,47 @@ def test_rejects_empty_wait_resource_label(self): config.get_wait_config(["customization", "jobs"], "create") +class TestGetWatchConfig: + """Tests for CLIConfig.get_watch_config method.""" + + def _create_config(self, config_yaml: str) -> CLIConfig: + with NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write(config_yaml) + f.flush() + return CLIConfig(Path(f.name)) + + def test_returns_configured_watch(self): + config = self._create_config(""" +config: + - resource: [customization, jobs] + methods: + create: + watch: + type: platform_job + resource_label: customization job +""") + assert config.get_watch_config(["customization", "jobs"], "create") == { + "type": "platform_job", + "resource_label": "customization job", + } + + def test_returns_inference_deployment_watch_config(self): + config = self._create_config(""" +config: + - resource: [inference, deployments] + methods: + create: + watch: + type: inference_deployment + resource_label: deployment +""") + + assert config.get_watch_config(["inference", "deployments"], "create") == { + "type": "inference_deployment", + "resource_label": "deployment", + } + + class TestDiscoverEntityTypes: """Tests for discover_entity_types function.""" diff --git a/tools/nemo-platform-sdk-tools/tests/sdk/cli_generator/test_generator.py b/tools/nemo-platform-sdk-tools/tests/sdk/cli_generator/test_generator.py index 1b0c6b36d8..71c02ffdc4 100644 --- a/tools/nemo-platform-sdk-tools/tests/sdk/cli_generator/test_generator.py +++ b/tools/nemo-platform-sdk-tools/tests/sdk/cli_generator/test_generator.py @@ -10,8 +10,9 @@ from typing import Any import pytest +from jinja2 import Environment, FileSystemLoader from nemo_platform._types import Omit -from nemo_platform_sdk_tools.sdk.cli_generator.config import CLIConfig +from nemo_platform_sdk_tools.sdk.cli_generator.config import CLIConfig, get_templates_dir from nemo_platform_sdk_tools.sdk.cli_generator.context_collectors.base import ( build_path_params, promote_name_to_positional, @@ -696,6 +697,18 @@ def _make_body_param(name: str, help_text: str | None = None) -> Parameter: ) +def _render_create_command(context: dict[str, Any]) -> str: + env = Environment( # noqa: S701 # nosec B701 + loader=FileSystemLoader(get_templates_dir()), + trim_blocks=True, + lstrip_blocks=True, + autoescape=False, + ) + env.filters["repr"] = repr + env.filters["to_kebab"] = lambda value: str(value).replace("_", "-") + return env.get_template("create_command.py.j2").render(**context) + + class TestCLIConfig: def test_top_level_command_config_ignores_non_mapping_top_level(self): config = _make_config(""" @@ -907,7 +920,30 @@ def test_create_no_name_param_does_not_crash(self): ctx = CreateContextCollector(self._config()).collect(["things"], sdk_method, "create") assert all(p.var_name != "name" for p in ctx["parameters"]) - def test_create_includes_wait_config(self): + def test_create_includes_watch_config(self): + config = self._config(""" +config: + - resource: [things] + methods: + create: + watch: + type: platform_job + resource_label: thing job +""") + sdk_method = _make_sdk_method_with_body([("name", None), ("spec", None)]) + ctx = CreateContextCollector(config).collect(["things"], sdk_method, "create") + + assert ctx["watch_config"] == {"type": "platform_job", "resource_label": "thing job"} + assert ctx["wait_config"] is None + rendered = _render_create_command(ctx) + assert "from nemo_platform_ext.jobs.watch import watch_job" not in rendered + assert "from nemo_platform_ext.cli.core.job_watch_renderer import render_job_watch_events" in rendered + assert "from nemo_platform_ext.cli.core.waiters import wait_for_platform_job" not in rendered + assert "from nemo_platform_plugin.client.adapter import client_from_platform" in rendered + assert "from nemo_platform_plugin.jobs.client import JobsClient" in rendered + assert "jobs_client.watch_job(" in rendered + + def test_create_platform_job_watch_config_renders_wait_and_watch(self): config = self._config(""" config: - resource: [things] @@ -916,11 +952,100 @@ def test_create_includes_wait_config(self): wait: type: platform_job resource_label: thing job + watch: + type: platform_job + resource_label: thing job """) sdk_method = _make_sdk_method_with_body([("name", None), ("spec", None)]) ctx = CreateContextCollector(config).collect(["things"], sdk_method, "create") - assert ctx["wait_config"] == {"type": "platform_job", "resource_label": "thing job"} + rendered = _render_create_command(ctx) + + assert 'typer.Option("--watch"' in rendered + assert 'typer.Option("--wait"' in rendered + assert "timeout: Annotated[int | None" in rendered + assert "] = None" in rendered + assert 'watch_options={"timeout": timeout, "poll_interval": poll_interval} if watch else None' in rendered + assert ( + 'wait_options={"timeout": timeout if timeout is not None else 1200, "poll_interval": poll_interval} ' + "if wait else None" + ) in rendered + assert "wait_config={\"type\": 'platform_job', \"resource_label\": 'thing job'} if wait else None" in rendered + assert "from nemo_platform_ext.jobs.watch import watch_job" not in rendered + assert "from nemo_platform_ext.cli.core.waiters import wait_for_platform_job" in rendered + assert "from nemo_platform_ext.cli.core.job_watch_renderer import render_job_watch_events" in rendered + assert "from nemo_platform_plugin.client.adapter import client_from_platform" in rendered + assert "from nemo_platform_plugin.jobs.client import JobsClient" in rendered + assert "jobs_client = client_from_platform(client, JobsClient)" in rendered + assert "client.things.get_status" not in rendered + assert "wait_for_platform_job(" in rendered + assert "jobs_client.watch_job(" in rendered + assert "resource_label='thing job'" in rendered + assert "render_job_watch_events(events, resource_label='thing job')" in rendered + assert rendered.index("format_output(") < rendered.index("if wait or watch:") + + def test_create_inference_deployment_config_renders_wait_and_watch(self): + config = self._config(""" +config: + - resource: [inference, deployments] + methods: + create: + wait: + type: inference_deployment + resource_label: deployment + watch: + type: inference_deployment + resource_label: deployment +""") + sdk_method = _make_sdk_method_with_body([("name", None), ("workspace", None), ("config", None)]) + ctx = CreateContextCollector(config).collect(["inference", "deployments"], sdk_method, "create") + + rendered = _render_create_command(ctx) + + assert 'typer.Option("--wait"' in rendered + assert 'typer.Option("--watch"' in rendered + assert "Wait for the created deployment to be up and running" in rendered + assert "Watch the created deployment until it is stable, streaming status updates" in rendered + assert "created deployment to reach a terminal state" not in rendered + assert "timeout if timeout is not None else 1200" not in rendered + assert "timeout: Annotated[int, typer.Option" in rendered + assert "] = 1200" in rendered + assert "from nemo_platform_ext.cli.core.waiters import wait_for_inference_deployment" in rendered + assert "from nemo_platform_ext.jobs.watch import watch_job" not in rendered + assert ( + "watch_config={\"type\": 'inference_deployment', \"resource_label\": 'deployment'} if watch else None" + in rendered + ) + assert 'watch_options={"timeout": timeout, "poll_interval": poll_interval} if watch else None' in rendered + assert ( + "wait_config={\"type\": 'inference_deployment', \"resource_label\": 'deployment'} if wait else None" + in rendered + ) + assert "if wait or watch:" in rendered + assert "wait_for_inference_deployment(" in rendered + assert "verbose=" not in rendered + assert "Unable to determine created resource name for --wait/--watch" in rendered + + def test_create_inference_deployment_wait_only_does_not_reference_watch(self): + config = self._config(""" +config: + - resource: [inference, deployments] + methods: + create: + wait: + type: inference_deployment + resource_label: deployment +""") + sdk_method = _make_sdk_method_with_body([("name", None), ("workspace", None), ("config", None)]) + ctx = CreateContextCollector(config).collect(["inference", "deployments"], sdk_method, "create") + + rendered = _render_create_command(ctx) + + assert 'typer.Option("--wait"' in rendered + assert 'typer.Option("--watch"' not in rendered + assert "if wait:" in rendered + assert "verbose=" not in rendered + assert "verbose=watch" not in rendered def test_delete_suppressed_by_config(self): sdk_method = _make_sdk_method_with_body([("name", None)])