anyscale 0.26.29__py3-none-any.whl → 0.26.31__py3-none-any.whl

anyscale/controllers/job_controller.py

@@ -1,13 +1,11 @@
 from __future__ import annotations
 
 import asyncio
-from collections import defaultdict
-from enum import Enum
 import os
 import random
 import string
 import time
-from typing import Any, Callable, cast, Dict, Iterable, List, Optional
+from typing import Any, cast, Dict, List, Optional
 
 import click
 import tabulate
@@ -27,21 +25,10 @@ from anyscale.client.openapi_client.models.create_internal_production_job import
 from anyscale.client.openapi_client.models.create_job_queue_config import (
     CreateJobQueueConfig,
 )
-from anyscale.client.openapi_client.models.decorated_job_queue import DecoratedJobQueue
 from anyscale.client.openapi_client.models.decorated_production_job import (
     DecoratedProductionJob,
 )
-from anyscale.client.openapi_client.models.decoratedjobqueue_response import (
-    DecoratedjobqueueResponse,
-)
 from anyscale.client.openapi_client.models.ha_job_states import HaJobStates
-from anyscale.client.openapi_client.models.job_queue_sort_directive import (
-    JobQueueSortDirective,
-)
-from anyscale.client.openapi_client.models.job_queues_query import JobQueuesQuery
-from anyscale.client.openapi_client.models.update_job_queue_request import (
-    UpdateJobQueueRequest,
-)
 from anyscale.controllers.base_controller import BaseController
 from anyscale.models.job_model import JobConfig
 from anyscale.project_utils import infer_project_id
@@ -60,7 +47,7 @@ from anyscale.util import (
     populate_unspecified_cluster_configs_from_current_workspace,
     validate_job_config_dict,
 )
-from anyscale.utils.connect_helpers import paginate, search_entities
+from anyscale.utils.connect_helpers import search_entities
 from anyscale.utils.runtime_env import override_runtime_env_config
 from anyscale.utils.workload_types import Workload
 
@@ -664,198 +651,6 @@ class JobController(BaseController):
         )
         return job_runs
 
-    def update_job_queue(
-        self,
-        job_queue_id: str,
-        job_queue_name: str,
-        max_concurrency: Optional[int] = None,
-        idle_timeout_s: Optional[int] = None,
-    ):
-        job_queue: DecoratedJobQueue = _resolve_object(
-            fetch_by_id=cast(
-                Callable[[str], DecoratedjobqueueResponse],
-                self.api_client.get_job_queue_api_v2_job_queues_job_queue_id_get,
-            ),
-            fetch_by_id_param=job_queue_id,
-            fetch_by_name=cast(
-                Callable[[str], DecoratedjobqueueResponse],
-                self.api_client.list_job_queues_api_v2_job_queues_post,
-            ),
-            fetch_by_name_query={
-                "job_queues_query": {"name": {"equals": job_queue_name,}}
-            },
-            object_type_description="job queue",
-        )
-
-        queue: DecoratedJobQueue = self.api_client.update_job_queue_api_v2_job_queues_job_queue_id_put(
-            job_queue_id=job_queue.id,
-            update_job_queue_request=UpdateJobQueueRequest(
-                max_concurrency=max_concurrency, idle_timeout_sec=idle_timeout_s,
-            ),
-        ).result
-
-        _print_job_queue_vertical(queue, JobQueueView.ALL)
-
-    def get_job_queue(self, job_queue_id: str):
-        queue: DecoratedJobQueue = self.api_client.get_job_queue_api_v2_job_queues_job_queue_id_get(
-            job_queue_id=job_queue_id
-        ).result
-
-        _print_job_queue_vertical(queue, JobQueueView.ALL)
-
-    def list_job_queues(
-        self,
-        include_all_users: bool,
-        view: JobQueueView,
-        page_size: int,
-        max_items: Optional[int],
-        sorting_directives: List[JobQueueSortDirective],
-        interactive: bool,
-    ):
-        creator_id = (
-            None
-            if include_all_users
-            else self.api_client.get_user_info_api_v2_userinfo_get().result.id
-        )
-
-        def build_query(paging_token: Optional[str], count: int) -> Dict:
-            return {
-                "job_queues_query": JobQueuesQuery(
-                    creator_id=creator_id,
-                    paging=PageQuery(paging_token=paging_token, count=count),
-                    sorting_directives=sorting_directives,
-                )
-            }
-
-        for batch in paginate(
-            search_function=self.api_client.list_job_queues_api_v2_job_queues_post,
-            query_builder=build_query,
-            interactive=interactive,
-            page_size=page_size,
-            max_items=max_items,
-        ):
-            _render_job_queues(batch, view)
-
-
-def _render_jobs(jobs):
-    jobs_table = [
-        [
-            job.name,
-            job.id,
-            job.project.name,
-            job.last_job_run.cluster.name
-            if job.last_job_run and job.last_job_run.cluster
-            else None,
-            job.state.current_state,
-            job.creator.email,
-            job.config.entrypoint
-            if len(job.config.entrypoint) < 50
-            else job.config.entrypoint[:50] + " ...",
-        ]
-        for job in jobs
-    ]
-
-    table = tabulate.tabulate(
-        jobs_table,
-        headers=[
-            "NAME",
-            "ID",
-            "PROJECT NAME",
-            "CLUSTER NAME",
-            "CURRENT STATE",
-            "CREATOR",
-            "ENTRYPOINT",
-        ],
-        tablefmt="plain",
-    )
-    click.echo(f"{table}")
-
-
-class JobQueueView(Enum):
-    ALL = DecoratedJobQueue.attribute_map.keys()
-    STATS = [
-        "id",
-        "name",
-        "total_jobs",
-        "active_jobs",
-        "successful_jobs",
-        "failed_jobs",
-    ]
-    DEFAULT = [
-        "id",
-        "name",
-        "cluster_id",
-        "creator_id",
-        "max_concurrency",
-        "idle_timeout_sec",
-        "current_cluster_state",
-        "created_at",
-    ]
-
-
-def _format_job_queue(queue: DecoratedJobQueue, view: JobQueueView) -> List[str]:
-    formatters: Dict[str, Callable[[Any], Any]] = defaultdict(lambda: (lambda v: v))
-    formatters["created_at"] = lambda k: k.strftime("%Y-%m-%d %H:%M:%S")
-
-    return [formatters[field](getattr(queue, field, "")) or "" for field in view.value]
-
-
-def _render_job_queues(queues: Iterable[DecoratedJobQueue], view: JobQueueView):
-    if not queues:
-        click.echo("No job queues found!")
-        return
-    table = tabulate.tabulate(
-        [
-            _format_job_queue(queue, view)
-            for queue in cast(Iterable[DecoratedJobQueue], queues)
-        ],
-        headers=[field.replace("_", " ").upper() for field in view.value],
-        tablefmt="plain",
-        maxcolwidths=30,  # type: ignore
-        numalign="center",
-        stralign="center",
-    )
-    click.echo(table)
-
-
-def _print_job_queue_vertical(queue: DecoratedJobQueue, job_queue_view: JobQueueView):
-    """
-    Print single job queue with headers as a vertical table
-    """
-    for header, value in zip(
-        [field.replace("_", " ").upper() for field in job_queue_view.value],
-        _format_job_queue(queue, job_queue_view),
-    ):
-        print(f"{header:<{30}}: {value}")
-
-
-def _resolve_object(
-    fetch_by_id: Optional[Callable[[str], object]],
-    fetch_by_id_param: Optional[str],
-    fetch_by_name,
-    fetch_by_name_query,
-    object_type_description: str,
-) -> Any:
-    """Given job_id or job_name, retrieve decorated ha job spec"""
-    if fetch_by_id_param is None and fetch_by_name_query is None:
-        raise click.ClickException(
-            "Either `--id` or `--name` must be passed in for object."
-        )
-    if fetch_by_id_param:
-        try:
-            return fetch_by_id(fetch_by_id_param).result  # type: ignore
-        except Exception as e:  # noqa: BLE001
-            raise click.ClickException(
-                f"Could not fetch {object_type_description} by id: {e}"
-            )
-
-    object_list_resp: List[Any] = fetch_by_name(**fetch_by_name_query).results
-    if len(object_list_resp) == 0:
-        raise click.ClickException(
-            f"No {object_type_description} found with the provided name"
-        )
-    if len(object_list_resp) > 1:
-        raise click.ClickException(
-            f"Multiple {object_type_description}s found with the provided name"
-        )
-    return object_list_resp[0]
+    def get_authenticated_user_id(self) -> str:
+        user_info_response = self.api_client.get_user_info_api_v2_userinfo_get()
+        return user_info_response.result.id

anyscale/job_queue/__init__.py

@@ -0,0 +1,89 @@
+from typing import List, Optional
+
+from anyscale._private.anyscale_client import AnyscaleClient
+from anyscale._private.models.model_base import ResultIterator
+from anyscale._private.sdk import sdk_command, sdk_docs
+from anyscale._private.sdk.base_sdk import Timer
+from anyscale.cli_logger import BlockLogger
+from anyscale.client.openapi_client.models.job_queue_sort_directive import (
+    JobQueueSortDirective,
+)
+from anyscale.client.openapi_client.models.session_state import SessionState
+from anyscale.job_queue._private.job_queue_sdk import PrivateJobQueueSDK
+from anyscale.job_queue.commands import (
+    _JOB_QUEUE_SDK_SINGLETON_KEY,
+    _LIST_ARG_DOCSTRINGS,
+    _LIST_EXAMPLE,
+    _STATUS_ARG_DOCSTRINGS,
+    _STATUS_EXAMPLE,
+    _UPDATE_ARG_DOCSTRINGS,
+    _UPDATE_EXAMPLE,
+    list,
+    status,
+    update,
+)
+from anyscale.job_queue.models import JobQueueSortField, JobQueueState, JobQueueStatus
+
+
+class JobQueueSDK:
+    """Public SDK for interacting with Anyscale Job Queues."""
+
+    def __init__(
+        self,
+        *,
+        client: Optional[AnyscaleClient] = None,
+        logger: Optional[BlockLogger] = None,
+        timer: Optional[Timer] = None,
+    ):
+        self._private_sdk = PrivateJobQueueSDK(
+            client=client, logger=logger, timer=timer
+        )
+
+    @sdk_docs(doc_py_example=_LIST_EXAMPLE, arg_docstrings=_LIST_ARG_DOCSTRINGS)
+    def list(  # noqa: F811
+        self,
+        *,
+        job_queue_id: Optional[str] = None,
+        name: Optional[str] = None,
+        creator_id: Optional[str] = None,
+        cloud: Optional[str] = None,
+        project: Optional[str] = None,
+        cluster_status: Optional[SessionState] = None,
+        page_size: Optional[int] = None,
+        max_items: Optional[int] = None,
+        sorting_directives: Optional[List[JobQueueSortDirective]] = None,
+    ) -> ResultIterator[JobQueueStatus]:
+        """List job queues or fetch a single job queue by ID."""
+        return self._private_sdk.list(
+            job_queue_id=job_queue_id,
+            name=name,
+            creator_id=creator_id,
+            cloud=cloud,
+            project=project,
+            cluster_status=cluster_status,
+            page_size=page_size,
+            max_items=max_items,
+            sorting_directives=sorting_directives,
+        )
+
+    @sdk_docs(doc_py_example=_STATUS_EXAMPLE, arg_docstrings=_STATUS_ARG_DOCSTRINGS)
+    def status(self, job_queue_id: str) -> JobQueueStatus:  # noqa: F811
+        """Get the status and details for a specific job queue."""
+        return self._private_sdk.status(job_queue_id=job_queue_id)
+
+    @sdk_docs(doc_py_example=_UPDATE_EXAMPLE, arg_docstrings=_UPDATE_ARG_DOCSTRINGS)
+    def update(  # noqa: F811
+        self,
+        *,
+        job_queue_id: Optional[str] = None,
+        job_queue_name: Optional[str] = None,
+        max_concurrency: Optional[int] = None,
+        idle_timeout_s: Optional[int] = None,
+    ) -> JobQueueStatus:
+        """Update a job queue."""
+        return self._private_sdk.update(
+            job_queue_id=job_queue_id,
+            job_queue_name=job_queue_name,
+            max_concurrency=max_concurrency,
+            idle_timeout_s=idle_timeout_s,
+        )

anyscale/job_queue/_private/job_queue_sdk.py

@@ -0,0 +1,158 @@
+from typing import List, Optional
+
+from anyscale._private.models.model_base import ResultIterator
+from anyscale._private.workload import WorkloadSDK
+from anyscale.client.openapi_client.models.decorated_job_queue import DecoratedJobQueue
+from anyscale.client.openapi_client.models.decoratedjobqueue_list_response import (
+    DecoratedjobqueueListResponse,
+)
+from anyscale.client.openapi_client.models.job_queue_sort_directive import (
+    JobQueueSortDirective,
+)
+from anyscale.client.openapi_client.models.list_response_metadata import (
+    ListResponseMetadata,
+)
+from anyscale.client.openapi_client.models.session_state import SessionState
+from anyscale.job_queue.models import JobQueueStatus
+
+
+class PrivateJobQueueSDK(WorkloadSDK):
+    """Internal SDK logic for Job Queue operations."""
+
+    def list(
+        self,
+        *,
+        job_queue_id: Optional[str] = None,
+        name: Optional[str] = None,
+        creator_id: Optional[str] = None,
+        cloud: Optional[str] = None,
+        project: Optional[str] = None,
+        cluster_status: Optional[SessionState] = None,
+        page_size: Optional[int] = None,
+        max_items: Optional[int] = None,
+        sorting_directives: Optional[List[JobQueueSortDirective]] = None,
+    ) -> ResultIterator[JobQueueStatus]:
+        """List job queues based on specified filters and pagination.
+
+        If job_queue_id is provided, fetches only that specific job queue.
+        """
+
+        if job_queue_id is not None:
+            raw = self._resolve_to_job_queue_model(job_queue_id=job_queue_id)
+
+            def _fetch_single_page(
+                _token: Optional[str],
+            ) -> DecoratedjobqueueListResponse:
+                # Only return data on the first call (token=None), simulate single-item page
+                if _token is None and raw is not None:
+                    results = [raw]
+                    metadata = ListResponseMetadata(total=1, next_paging_token=None)
+                else:
+                    results = []
+                    metadata = ListResponseMetadata(total=0, next_paging_token=None)
+
+                return DecoratedjobqueueListResponse(
+                    results=results, metadata=metadata,
+                )
+
+            return ResultIterator(
+                page_token=None,
+                max_items=0,
+                fetch_page=_fetch_single_page,
+                parse_fn=_parse_decorated_jq_to_status,
+            )
+
+        def _fetch_page(token: Optional[str]) -> DecoratedjobqueueListResponse:
+            return self.client.list_job_queues(
+                name=name,
+                creator_id=creator_id,
+                cloud=cloud,
+                project=project,
+                cluster_status=cluster_status,
+                count=page_size,
+                paging_token=token,
+                sorting_directives=sorting_directives,
+            )
+
+        return ResultIterator(
+            page_token=None,
+            max_items=max_items,
+            fetch_page=_fetch_page,
+            parse_fn=_parse_decorated_jq_to_status,
+        )
+
+    def status(self, job_queue_id: str) -> JobQueueStatus:
+        """Get the status and details for a specific job queue."""
+        raw = self._resolve_to_job_queue_model(job_queue_id=job_queue_id)
+        return _parse_decorated_jq_to_status(raw)
+
+    def update(
+        self,
+        *,
+        job_queue_id: Optional[str] = None,
+        job_queue_name: Optional[str] = None,
+        max_concurrency: Optional[int] = None,
+        idle_timeout_s: Optional[int] = None,
+    ) -> JobQueueStatus:
+        """Update a job queue."""
+
+        if max_concurrency is None and idle_timeout_s is None:
+            raise ValueError("No fields to update")
+
+        jq = self._resolve_to_job_queue_model(
+            job_queue_id=job_queue_id, name=job_queue_name
+        )
+
+        assert jq.id is not None
+        updated_jq = self.client.update_job_queue(
+            job_queue_id=jq.id,
+            max_concurrency=max_concurrency,
+            idle_timeout_s=idle_timeout_s,
+        )
+
+        return _parse_decorated_jq_to_status(updated_jq)
+
+    def _resolve_to_job_queue_model(
+        self, *, job_queue_id: Optional[str] = None, name: Optional[str] = None,
+    ) -> DecoratedJobQueue:
+        """Finds the specific Job Queue API model by ID or name.
+        """
+        if job_queue_id is None and name is None:
+            raise ValueError("Either 'job_queue_id' or 'name' must be provided.")
+
+        if job_queue_id:
+            job_queue = self.client.get_job_queue(job_queue_id)
+            if job_queue is None:
+                raise ValueError(f"Job Queue with ID '{job_queue_id}' not found.")
+            return job_queue
+        else:
+            job_queues = self.client.list_job_queues(name=name, count=1)
+            if len(job_queues) == 0:
+                raise ValueError(f"Job Queue with name '{name}' not found.")
+            return job_queues[0]
+
+
+def _parse_decorated_jq_to_status(decorated_jq: DecoratedJobQueue) -> JobQueueStatus:
+    """Helper to convert API model to SDK model."""
+
+    if decorated_jq.id is None or decorated_jq.current_job_queue_state is None:
+        raise ValueError("Job Queue ID or state is missing.")
+
+    return JobQueueStatus(
+        id=decorated_jq.id,
+        name=decorated_jq.name,
+        state=decorated_jq.current_job_queue_state,
+        creator_email=decorated_jq.creator_email,
+        project_id=decorated_jq.project_id,
+        created_at=decorated_jq.created_at,
+        max_concurrency=decorated_jq.max_concurrency,
+        idle_timeout_s=decorated_jq.idle_timeout_sec,
+        creator_id=decorated_jq.creator_id,
+        cloud_id=decorated_jq.cloud_id,
+        user_provided_id=decorated_jq.user_provided_id,
+        execution_mode=decorated_jq.execution_mode,
+        total_jobs=decorated_jq.total_jobs,
+        active_jobs=decorated_jq.active_jobs,
+        successful_jobs=decorated_jq.successful_jobs,
+        failed_jobs=decorated_jq.failed_jobs,
+    )

anyscale/job_queue/commands.py

@@ -0,0 +1,130 @@
+from typing import List, Optional
+
+from anyscale._private.models.model_base import ResultIterator
+from anyscale._private.sdk import sdk_command
+from anyscale.client.openapi_client.models.job_queue_sort_directive import (
+    JobQueueSortDirective,
+)
+from anyscale.client.openapi_client.models.session_state import SessionState
+from anyscale.job_queue._private.job_queue_sdk import PrivateJobQueueSDK
+from anyscale.job_queue.models import JobQueueStatus
+
+
+_JOB_QUEUE_SDK_SINGLETON_KEY = "job_queue_sdk"
+
+
+_LIST_EXAMPLE = """
+import anyscale
+
+# Example: List the first 50 job queues
+for jq in anyscale.job_queue.list(max_items=50):
+    print(jq.id, jq.name, jq.state)
+"""
+
+_LIST_ARG_DOCSTRINGS = {
+    "job_queue_id": "If provided, fetches only the job queue with this ID.",
+    "name": "Filter by job queue name.",
+    "creator_id": "Filter by the user ID of the creator.",
+    "cloud": "Filter by cloud name.",
+    "project": "Filter by project name.",
+    "cluster_status": "Filter by the state of the associated cluster.",
+    "page_size": "Number of items per API request page.",
+    "max_items": "Maximum total number of items to return.",
+    "sorting_directives": "List of directives to sort the results.",
+}
+
+_STATUS_EXAMPLE = """
+import anyscale
+
+status = anyscale.job_queue.status(job_queue_id=\"jobq_abc123\")
+print(status)
+"""
+
+_STATUS_ARG_DOCSTRINGS = {
+    "job_queue_id": "The unique ID of the job queue.",
+}
+
+_UPDATE_EXAMPLE = """
+import anyscale
+
+updated_jq = anyscale.job_queue.update(job_queue_id=\"jobq_abc123\", max_concurrency=5)
+print(updated_jq)
+"""
+
+_UPDATE_ARG_DOCSTRINGS = {
+    "job_queue_id": "ID of the job queue to update.",
+    "job_queue_name": "Name of the job queue to update (alternative to ID).",
+    "max_concurrency": "New maximum concurrency value.",
+    "idle_timeout_s": "New idle timeout in seconds.",
+}
+
+
+@sdk_command(
+    _JOB_QUEUE_SDK_SINGLETON_KEY,
+    PrivateJobQueueSDK,
+    doc_py_example=_LIST_EXAMPLE,
+    arg_docstrings=_LIST_ARG_DOCSTRINGS,
+)
+def list(  # noqa: A001
+    *,
+    job_queue_id: Optional[str] = None,
+    name: Optional[str] = None,
+    creator_id: Optional[str] = None,
+    cloud: Optional[str] = None,
+    project: Optional[str] = None,
+    cluster_status: Optional[SessionState] = None,
+    page_size: Optional[int] = None,
+    max_items: Optional[int] = None,
+    sorting_directives: Optional[List[JobQueueSortDirective]] = None,
+    _private_sdk: Optional[PrivateJobQueueSDK] = None,
+) -> ResultIterator[JobQueueStatus]:
+    """List job queues or fetch a single job queue by ID."""
+    return _private_sdk.list(  # type: ignore
+        job_queue_id=job_queue_id,
+        name=name,
+        creator_id=creator_id,
+        cloud=cloud,
+        project=project,
+        cluster_status=cluster_status,
+        page_size=page_size,
+        max_items=max_items,
+        sorting_directives=sorting_directives,
+    )
+
+
+@sdk_command(
+    _JOB_QUEUE_SDK_SINGLETON_KEY,
+    PrivateJobQueueSDK,
+    doc_py_example=_STATUS_EXAMPLE,
+    arg_docstrings=_STATUS_ARG_DOCSTRINGS,
+)
+def status(
+    job_queue_id: str, _private_sdk: Optional[PrivateJobQueueSDK] = None
+) -> JobQueueStatus:
+    """Get the status and details for a specific job queue."""
+    return _private_sdk.status(  # type: ignore
+        job_queue_id=job_queue_id
+    )
+
+
+@sdk_command(
+    _JOB_QUEUE_SDK_SINGLETON_KEY,
+    PrivateJobQueueSDK,
+    doc_py_example=_UPDATE_EXAMPLE,
+    arg_docstrings=_UPDATE_ARG_DOCSTRINGS,
+)
+def update(
+    *,
+    job_queue_id: Optional[str] = None,
+    job_queue_name: Optional[str] = None,
+    max_concurrency: Optional[int] = None,
+    idle_timeout_s: Optional[int] = None,
+    _private_sdk: Optional[PrivateJobQueueSDK] = None,
+) -> JobQueueStatus:
+    """Update a job queue."""
+    return _private_sdk.update(  # type: ignore
+        job_queue_id=job_queue_id,
+        job_queue_name=job_queue_name,
+        max_concurrency=max_concurrency,
+        idle_timeout_s=idle_timeout_s,
+    )