anyscale 0.26.19__py3-none-any.whl → 0.26.21__py3-none-any.whl

anyscale/controllers/job_controller.py

@@ -1,9 +1,13 @@
+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, cast, Dict, List, Optional
+from typing import Any, Callable, cast, Dict, Iterable, List, Optional
 
 import click
 import tabulate
@@ -23,10 +27,21 @@ 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
@@ -45,7 +60,7 @@ from anyscale.util import (
     populate_unspecified_cluster_configs_from_current_workspace,
     validate_job_config_dict,
 )
-from anyscale.utils.connect_helpers import search_entities
+from anyscale.utils.connect_helpers import paginate, search_entities
 from anyscale.utils.runtime_env import override_runtime_env_config
 from anyscale.utils.workload_types import Workload
 
@@ -294,7 +309,7 @@ class JobController(BaseController):
         project_id: Optional[str],
         include_archived: bool,
         max_items: int,
-        states: List[str],
+        states: List[HaJobStates],
     ) -> None:
         """
         This function will list jobs.
@@ -338,6 +353,7 @@ class JobController(BaseController):
                 )
             else:
                 creator_id = None
+
             resp = self.api_client.list_decorated_jobs_api_v2_decorated_ha_jobs_get(
                 project_id=project_id,
                 name=name,
@@ -647,3 +663,199 @@ 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]

anyscale/scripts.py

@@ -22,6 +22,7 @@ from anyscale.commands.experimental_integrations_commands import (
 )
 from anyscale.commands.image_commands import image_cli
 from anyscale.commands.job_commands import job_cli
+from anyscale.commands.job_queue_commands import job_queue_cli
 from anyscale.commands.list_commands import list_cli
 from anyscale.commands.llm.group import llm_cli
 from anyscale.commands.login_commands import anyscale_login, anyscale_logout
@@ -120,6 +121,7 @@ cli.add_command(version_cli)
 cli.add_command(list_cli)
 cli.add_command(cluster_env_cli)
 cli.add_command(job_cli)
+# cli.add_command(job_queue_cli) # TODO will be enabled later
 cli.add_command(schedule_cli)
 cli.add_command(service_cli)
 cli.add_command(cluster_cli)
@@ -153,6 +155,7 @@ ALIASES = {
     "h": anyscale_help,
     "schedules": schedule_cli,
     "jobs": job_cli,
+    "jq": job_queue_cli,
     "services": service_cli,
     "cluster-compute": compute_config_cli,
     "images": image_cli,

anyscale/sdk/anyscale_client/models/production_job_state_transition.py

@@ -108,7 +108,7 @@ class ProductionJobStateTransition(object):
     def state_transitioned_at(self):
         """Gets the state_transitioned_at of this ProductionJobStateTransition.  # noqa: E501
 
-        The last time the state of this job was updated. This includes updates to the state and to the goal state  # noqa: E501
+        The last time the state of this job was updated  # noqa: E501
 
         :return: The state_transitioned_at of this ProductionJobStateTransition.  # noqa: E501
         :rtype: datetime
@@ -119,7 +119,7 @@ class ProductionJobStateTransition(object):
     def state_transitioned_at(self, state_transitioned_at):
         """Sets the state_transitioned_at of this ProductionJobStateTransition.
 
-        The last time the state of this job was updated. This includes updates to the state and to the goal state  # noqa: E501
+        The last time the state of this job was updated  # noqa: E501
 
         :param state_transitioned_at: The state_transitioned_at of this ProductionJobStateTransition.  # noqa: E501
         :type: datetime

anyscale/util.py

@@ -987,7 +987,9 @@ def validate_job_config_dict(
         )
 
 
-def validate_list_jobs_state_filter(ctx, param, value) -> List[str]:  # noqa: ARG001
+def validate_list_jobs_state_filter(
+    _, param, value
+) -> List[HaJobStates]:  # noqa: ARG001
     """
     Validate the job state filter for list jobs CLI method
     """

anyscale/utils/connect_helpers.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass
 import inspect
+import sys
 from typing import (
     Any,
     Callable,
@@ -166,3 +167,36 @@ def search_entities(
         entities = entities[:max_to_return]
 
     return entities
+
+
+def paginate(
+    search_function: Callable[..., ListResponse[T]],
+    query_builder: Callable[[Optional[str], int], Dict],
+    page_size: int,
+    max_items: Optional[int] = None,
+    interactive: bool = True,
+):
+    max_items = max_items or sys.maxsize
+    queues = []
+    has_more, token = True, None
+
+    while has_more:
+        query_kwargs = query_builder(token, min(page_size, max_items))
+        resp = search_function(**query_kwargs)
+        token = resp.metadata.next_paging_token
+        batch, has_more = resp.results, token is not None
+        queues.extend(batch)
+        yield batch
+
+        if (
+            not has_more
+            or (not interactive and len(queues) >= max_items)
+            or (
+                interactive
+                and input("Press Enter to load more, or 'q' to quit: ").strip().lower()
+                == "q"
+            )
+        ):
+            break
+
+        max_items -= page_size

anyscale/utils/gcp_utils.py

@@ -248,10 +248,26 @@ def get_gcp_filestore_config(
     filestore_instance_id: str,
     logger: CloudSetupLogger,
 ):
-    client = factory.filestore_v1.CloudFilestoreManagerClient()
     instance_name = "projects/{}/locations/{}/instances/{}".format(
         project_id, filestore_location, filestore_instance_id
     )
+    return get_gcp_filestore_config_from_full_name(
+        factory=factory, vpc_name=vpc_name, instance_name=instance_name, logger=logger,
+    )
+
+
+def get_gcp_filestore_config_from_full_name(
+    factory: GoogleCloudClientFactory,
+    vpc_name: str,
+    instance_name: str,
+    logger: CloudSetupLogger,
+):
+    if not re.search("projects/.+/locations/.+/instances/.+", instance_name):
+        raise ValueError(
+            "Please provide the full filestore instance name. Example: projects/<project number>/locations/<location>/instances/<instance id>"
+        )
+
+    client = factory.filestore_v1.CloudFilestoreManagerClient()
     try:
         file_store = client.get_instance(name=instance_name)
     except NotFound as e:
@@ -260,7 +276,7 @@ def get_gcp_filestore_config(
             CloudSetupError.RESOURCE_NOT_FOUND,
         )
         raise ClickException(
-            f"Could not find Filestore with id {filestore_instance_id} in project {project_id}. Please validate that you're using the correct GCP project and that the resource values are correct. Error details: {e}"
+            f"Could not find Filestore with id {instance_name}. Please validate that you're using the correct GCP project and that the resource values are correct. Error details: {e}"
         )
     root_dir = file_store.file_shares[0].name
     for v in file_store.networks:
@@ -271,7 +287,7 @@ def get_gcp_filestore_config(
             break
     else:
         logger.error(
-            f"Filestore {filestore_instance_id} is not connected to {vpc_name}, but to {[v.network for v in file_store.networks]}. "
+            f"Filestore {instance_name} is not connected to {vpc_name}, but to {[v.network for v in file_store.networks]}. "
             f"This cannot be edited on an existing Filestore instance. Please recreate the filestore and connect it to {vpc_name}."
         )
         logger.log_resource_error(
@@ -279,7 +295,7 @@ def get_gcp_filestore_config(
             CloudSetupError.FILESTORE_NOT_CONNECTED_TO_VPC,
         )
         raise ClickException(
-            f"Filestore {filestore_instance_id} is not connected to {vpc_name}."
+            f"Filestore {instance_name} is not connected to {vpc_name}."
         )
     return GCPFileStoreConfig(
         instance_name=instance_name, root_dir=root_dir, mount_target_ip=mount_target_ip,

anyscale/version.py

@@ -1 +1 @@
-__version__ = "0.26.19"
+__version__ = "0.26.21"

anyscale/workspace/_private/workspace_sdk.py

@@ -422,11 +422,18 @@ class PrivateWorkspaceSDK(WorkloadSDK):
             should_warn_delete = True
             dry_run_options.append("--delete-excluded")
 
-        result = subprocess.run(
-            rsync_command + dry_run_options, capture_output=True, text=True, check=True
-        )
-        if result.returncode != 0:
-            raise RuntimeError(f"Rsync failed with return code {result.returncode}")
+        try:
+            result = subprocess.run(
+                rsync_command + dry_run_options,
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+        except subprocess.CalledProcessError as e:
+            self._logger.error(f"Error running rsync command: {e}")
+            self._logger.error(f">>> stdout: {e.stdout}")
+            self._logger.error(f">>> stderr: {e.stderr}")
+            raise RuntimeError(f"Rsync failed with return code {e.returncode}")
 
         _, deleting_files = self._parse_rsync_dry_run_output(result.stdout)
 
@@ -514,7 +521,13 @@ class PrivateWorkspaceSDK(WorkloadSDK):
             # Add -v / --verbose to the rsync command to be explicit about what is being transferred
             args += ["-v"]
 
-            subprocess.run(args, check=True)
+            try:
+                subprocess.run(args, check=True, capture_output=True, text=True)
+            except subprocess.CalledProcessError as e:
+                self._logger.error(f">>> Error running rsync command: {e}")
+                self._logger.error(f">>> stdout: {e.stdout}")
+                self._logger.error(f">>> stderr: {e.stderr}")
+                raise RuntimeError(f"Rsync failed with return code {e.returncode}")
 
     def pull(
         self,

{anyscale-0.26.19.dist-info → anyscale-0.26.21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anyscale
-Version: 0.26.19
+Version: 0.26.21
 Summary: Command Line Interface for Anyscale
 Author: Anyscale Inc.
 License: AS License