anyscale 0.24.86__py3-none-any.whl

anyscale/controllers/job_controller.py

@@ -0,0 +1,647 @@
+import asyncio
+import os
+import random
+import string
+import time
+from typing import Any, cast, Dict, List, Optional
+
+import click
+import tabulate
+import yaml
+
+from anyscale.anyscale_pydantic import BaseModel
+from anyscale.api_utils.job_logs_util import _get_job_logs_from_storage_bucket_streaming
+from anyscale.cli_logger import LogsLogger
+from anyscale.client.openapi_client import (
+    ComputeTemplate,
+    ProductionJob,
+    ProductionJobConfig,
+)
+from anyscale.client.openapi_client.models.create_internal_production_job import (
+    CreateInternalProductionJob,
+)
+from anyscale.client.openapi_client.models.create_job_queue_config import (
+    CreateJobQueueConfig,
+)
+from anyscale.client.openapi_client.models.decorated_production_job import (
+    DecoratedProductionJob,
+)
+from anyscale.client.openapi_client.models.ha_job_states import HaJobStates
+from anyscale.controllers.base_controller import BaseController
+from anyscale.models.job_model import JobConfig
+from anyscale.project import infer_project_id
+from anyscale.sdk.anyscale_client.models.job import Job
+from anyscale.sdk.anyscale_client.models.jobs_query import JobsQuery
+from anyscale.sdk.anyscale_client.models.jobs_sort_field import JobsSortField
+from anyscale.sdk.anyscale_client.models.page_query import PageQuery
+from anyscale.sdk.anyscale_client.models.sort_by_clause_jobs_sort_field import (
+    SortByClauseJobsSortField,
+)
+from anyscale.sdk.anyscale_client.models.sort_order import SortOrder
+from anyscale.util import (
+    get_endpoint,
+    is_anyscale_workspace,
+    poll,
+    populate_unspecified_cluster_configs_from_current_workspace,
+    validate_job_config_dict,
+)
+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
+
+
+log = LogsLogger()
+
+_TERMINAL_STATES = {
+    HaJobStates.SUCCESS,
+    HaJobStates.TERMINATED,
+    HaJobStates.BROKEN,
+    HaJobStates.OUT_OF_RETRIES,
+}
+_PENDING_STATES = {
+    HaJobStates.PENDING,
+    HaJobStates.AWAITING_CLUSTER_START,
+    HaJobStates.ERRORED,
+    HaJobStates.RESTARTING,
+}
+
+
+class MiniJobRun(BaseModel):
+    last_job_run_id: str
+    job_state: str
+    error: Optional[str]
+
+
+class JobController(BaseController):
+    def __init__(
+        self,
+        log: Optional[LogsLogger] = None,
+        initialize_auth_api_client: bool = True,
+        raise_structured_exception: bool = False,
+        auth_token: str = None,
+    ):
+        if log is None:
+            log = LogsLogger()
+
+        super().__init__(
+            initialize_auth_api_client, raise_structured_exception, cli_token=auth_token
+        )
+        self.log = log
+        self.log.open_block("Output")
+        self._cluster_journal_events_start_line = 0
+
+    def submit(
+        self,
+        job_config_file: str,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        is_entrypoint_cmd: Optional[bool] = False,
+        entrypoint: Optional[List[str]] = None,
+    ) -> str:
+        entrypoint = entrypoint or []
+        workspace_id = os.environ.get("ANYSCALE_EXPERIMENTAL_WORKSPACE_ID", None)
+        if is_anyscale_workspace() and is_entrypoint_cmd:
+            entrypoint = [job_config_file, *entrypoint]
+            config = self.generate_config_from_entrypoint(
+                entrypoint, name, description, workspace_id
+            )
+            id = self.submit_from_config(config)  # noqa: A001
+        elif len(entrypoint) == 0:
+            # Assume that job_config_file is a file and submit it.
+            config = self.generate_config_from_file(
+                job_config_file,
+                name=name,
+                description=description,
+                workspace_id=workspace_id,
+            )
+            id = self.submit_from_config(config)  # noqa: A001
+        elif len(entrypoint) != 0:
+            msg = (
+                "Within an Anyscale Workspace, `anyscale job submit` takes either a file, or a command. To submit a command, use `anyscale job submit -- my command`."
+                if is_anyscale_workspace()
+                else "`anyscale job submit` takes one argument, a YAML file configuration. Please use `anyscale job submit my_file`."
+            )
+            raise click.ClickException(msg)
+        return id
+
+    def generate_config_from_entrypoint(
+        self,
+        entrypoint: List[str],
+        name: Optional[str],
+        description: Optional[str],
+        workspace_id: Optional[str] = None,
+    ) -> JobConfig:
+        config_dict = {
+            "entrypoint": " ".join(entrypoint),
+            "name": name,
+            "description": description,
+            "workspace_id": workspace_id,
+        }
+
+        config_dict = populate_unspecified_cluster_configs_from_current_workspace(
+            config_dict, self.anyscale_api_client,
+        )
+
+        return JobConfig.parse_obj(config_dict)
+
+    def generate_config_from_file(
+        self,
+        job_config_file: str,
+        name: Optional[str],
+        description: Optional[str],
+        workspace_id: Optional[str] = None,
+    ) -> JobConfig:
+        config_dict = self._load_config_dict_from_file(job_config_file)
+        config_dict["workspace_id"] = workspace_id
+        validate_job_config_dict(config_dict, self.api_client)
+        config_dict = populate_unspecified_cluster_configs_from_current_workspace(
+            config_dict, self.anyscale_api_client
+        )
+        job_config = JobConfig.parse_obj(config_dict)
+
+        if name:
+            job_config.name = name
+
+        if description:
+            job_config.description = description
+
+        return job_config
+
+    def submit_from_config(self, job_config: JobConfig):
+        # If project id is not specified, try to infer it
+        project_id = infer_project_id(
+            self.anyscale_api_client,
+            self.api_client,
+            self.log,
+            project_id=job_config.project_id,
+            cluster_compute_id=job_config.compute_config_id,
+            cluster_compute=job_config.compute_config,
+            cloud=job_config.cloud,
+        )
+
+        job_config.runtime_env = override_runtime_env_config(
+            runtime_env=job_config.runtime_env,
+            anyscale_api_client=self.anyscale_api_client,
+            api_client=self.api_client,
+            workload_type=Workload.JOBS,
+            compute_config_id=job_config.compute_config_id,
+            log=self.log,
+        )
+        config_object = ProductionJobConfig(
+            entrypoint=job_config.entrypoint,
+            runtime_env=job_config.runtime_env,
+            build_id=job_config.build_id,
+            compute_config_id=job_config.compute_config_id,
+            max_retries=job_config.max_retries,
+        )
+
+        job_queue_config = (
+            CreateJobQueueConfig(**job_config.job_queue_config)
+            if job_config.job_queue_config
+            else None
+        )
+
+        job = self.api_client.create_job_api_v2_decorated_ha_jobs_create_post(
+            CreateInternalProductionJob(
+                name=job_config.name or self._generate_random_job_name(),
+                description=job_config.description or "Job submitted from CLI",
+                project_id=project_id,
+                workspace_id=job_config.workspace_id,
+                config=config_object,
+                job_queue_config=job_queue_config,
+            )
+        ).result
+        self.log.info(
+            f"Maximum uptime is {self._get_maximum_uptime_output(job)} for clusters launched by this job."
+        )
+        self.log.info(
+            f"Job {job.id} has been successfully submitted. Current state of job: {job.state.current_state}."
+        )
+        self.log.info(
+            f"Query the status of the job with `anyscale job list --job-id {job.id}`."
+        )
+        self.log.info(
+            f"Get the logs for the job with `anyscale job logs --job-id {job.id} --follow`."
+        )
+        self.log.info(f'View the job in the UI at {get_endpoint(f"/jobs/{job.id}")}')
+        return job.id
+
+    def _get_maximum_uptime_output(self, job: ProductionJob) -> str:
+        compute_config: ComputeTemplate = self.api_client.get_compute_template_api_v2_compute_templates_template_id_get(
+            job.config.compute_config_id
+        ).result
+        maximum_uptime_minutes = compute_config.config.maximum_uptime_minutes
+        if maximum_uptime_minutes and maximum_uptime_minutes > 0:
+            return f"set to {maximum_uptime_minutes} minutes"
+        return "disabled"
+
+    def _load_config_dict_from_file(self, job_config_file: str) -> Dict[str, Any]:
+        if not os.path.exists(job_config_file):
+            raise click.ClickException(f"Config file {job_config_file} not found.")
+
+        with open(job_config_file) as f:
+            config_dict = yaml.safe_load(f)
+        return config_dict
+
+    def wait(
+        self,
+        job_name: Optional[str] = None,
+        job_id: Optional[str] = None,
+        target_state: str = HaJobStates.SUCCESS,
+        interval_secs: float = 10.0,
+        timeout_secs=None,
+    ):
+        if target_state not in HaJobStates.allowable_values:
+            raise click.ClickException(
+                f"`{target_state}` is not a valid Job state. Allowed states are {HaJobStates.allowable_values}"
+            )
+        job_id_or_name = job_name or job_id
+        job_id = cast(str, self._resolve_job_object(job_id, job_name).id,)
+        with self.log.spinner(f"Waiting for Job `{job_id_or_name}`...") as spinner:
+            state = None
+            for _ in poll(interval_secs=interval_secs, timeout_secs=timeout_secs):
+                state = self._get_job_state(job_id)
+                spinner.text = f"Job `{job_id_or_name}` is in state `{state}`. Waiting to reach state `{target_state}`."
+
+                if state == target_state:
+                    spinner.succeed(
+                        f"Job `{job_id_or_name}` reached state `{target_state}`."
+                    )
+                    return job_id
+
+                if state in _TERMINAL_STATES:
+                    msg = f"Job `{job_id_or_name}` reached terminal state `{state}`, and will not reach `{target_state}`."
+                    spinner.fail(msg)
+                    raise click.ClickException(msg)
+
+            msg = f"Timed out after waiting for {timeout_secs} seconds. The current state is {state}."
+            spinner.fail(msg)
+            raise click.ClickException(msg)
+
+    def _get_job_state(self, job_id: str):
+        job = self.api_client.get_job_api_v2_decorated_ha_jobs_production_job_id_get(
+            job_id
+        ).result
+        return job.state.current_state
+
+    def list(  # noqa: PLR0913
+        self,
+        include_all_users: bool,
+        name: Optional[str],
+        job_id: Optional[str],
+        project_id: Optional[str],
+        include_archived: bool,
+        max_items: int,
+    ) -> None:
+        """
+        This function will list jobs.
+        """
+        print(f'View your Jobs in the UI at {get_endpoint("/jobs")}')
+
+        jobs_list = []
+        if job_id:
+            job = self.api_client.get_job_api_v2_decorated_ha_jobs_production_job_id_get(
+                job_id
+            ).result
+            if job.is_service:
+                # `job_id` belongs to a service, but this function should list jobs.
+                raise click.ClickException(
+                    f"ID {job_id} belongs to a Anyscale service. Please get information about "
+                    f"this service with `anyscale service list --service-id {job_id}`."
+                )
+
+            output_map = {
+                "Name": job.name,
+                "Id": job.id,
+                "Cost (dollars)": job.cost_dollars,
+                "Project name": job.project.name,
+                "Cluster name": job.last_job_run.cluster.name
+                if job.last_job_run and job.last_job_run.cluster
+                else None,
+                "Current state": job.state.current_state,
+                "Creator": job.creator.email,
+                "Entrypoint": job.config.entrypoint
+                if len(job.config.entrypoint) < 100
+                else job.config.entrypoint[:100] + " ...",
+            }
+            output_str = "\n".join(
+                [f"\t{key}: {output_map[key]}" for key in output_map]
+            )
+            print(output_str)
+            return
+        else:
+            if not include_all_users:
+                creator_id = (
+                    self.api_client.get_user_info_api_v2_userinfo_get().result.id
+                )
+            else:
+                creator_id = None
+            resp = self.api_client.list_decorated_jobs_api_v2_decorated_ha_jobs_get(
+                project_id=project_id,
+                name=name,
+                creator_id=creator_id,
+                type_filter="BATCH_JOB",
+                archive_status="ALL" if include_archived else "NOT_ARCHIVED",
+                count=10,
+            )
+            jobs_list.extend(resp.results)
+            paging_token = resp.metadata.next_paging_token
+            has_more = (paging_token is not None) and (len(jobs_list) < max_items)
+            while has_more:
+                resp = self.api_client.list_decorated_jobs_api_v2_decorated_ha_jobs_get(
+                    project_id=project_id,
+                    name=name,
+                    creator_id=creator_id,
+                    type_filter="BATCH_JOB",
+                    archive_status="ALL" if include_archived else "NOT_ARCHIVED",
+                    count=10,
+                    paging_token=paging_token,
+                )
+                jobs_list.extend(resp.results)
+                paging_token = resp.metadata.next_paging_token
+                has_more = (paging_token is not None) and (len(jobs_list) < max_items)
+            jobs_list = jobs_list[:max_items]
+
+        jobs_table = [
+            [
+                job.name,
+                job.id,
+                job.cost_dollars,
+                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_list
+        ]
+
+        table = tabulate.tabulate(
+            jobs_table,
+            headers=[
+                "NAME",
+                "ID",
+                "COST",
+                "PROJECT NAME",
+                "CLUSTER NAME",
+                "CURRENT STATE",
+                "CREATOR",
+                "ENTRYPOINT",
+            ],
+            tablefmt="plain",
+        )
+        print(f"JOBS:\n{table}")
+
+    def archive(self, job_id: Optional[str], job_name: Optional[str]) -> None:
+        """
+        This function will archive jobs.
+        """
+        job_resp: DecoratedProductionJob = self._resolve_job_object(job_id, job_name)
+        self.api_client.archive_job_api_v2_decorated_ha_jobs_production_job_id_archive_post(
+            job_resp.id
+        )
+        self.log.info(f"Job {job_resp.id} is successfully archived.")
+
+    def terminate(self, job_id: Optional[str], job_name: Optional[str]) -> None:
+        """
+        This function will terminate jobs.
+        """
+        job_resp: DecoratedProductionJob = self._resolve_job_object(job_id, job_name)
+        job = self.api_client.terminate_job_api_v2_decorated_ha_jobs_production_job_id_terminate_post(
+            job_resp.id
+        ).result
+        self.log.info(f"Job {job.id} has begun terminating...")
+        self.log.info(
+            f" Current state of Job: {job.state.current_state}. Goal state of Job: {job.state.goal_state}"
+        )
+        self.log.info(
+            f"Query the status of the Job with `anyscale job list --job-id {job.id}`."
+        )
+
+    def _resolve_job_object(
+        self, job_id: Optional[str], job_name: Optional[str]
+    ) -> DecoratedProductionJob:
+        """Given job_id or job_name, retrieve decorated ha job spec"""
+        if job_id is None and job_name is None:
+            raise click.ClickException(
+                "Either `--id` or `--name` must be passed in for Job."
+            )
+        if job_id:
+            return self._get_job(job_id)
+
+        jobs_list_resp: List[
+            DecoratedProductionJob
+        ] = self.api_client.list_decorated_jobs_api_v2_decorated_ha_jobs_get(
+            name=job_name, type_filter="BATCH_JOB",
+        ).results
+        if len(jobs_list_resp) == 0:
+            raise click.ClickException(
+                f"No Job found with name {job_name}. Please either pass `--id` or list the "
+                f"available Jobs with `anyscale job list`."
+            )
+        if len(jobs_list_resp) > 1:
+            raise click.ClickException(
+                f"Multiple Jobs found with name {job_name}. Please specify the `--id` instead."
+            )
+        return jobs_list_resp[0]
+
+    def _get_job(self, job_id: str) -> DecoratedProductionJob:
+        job_object = self.api_client.get_job_api_v2_decorated_ha_jobs_production_job_id_get(
+            production_job_id=job_id
+        ).result
+        return job_object
+
+    def _get_formatted_latest_job_run(
+        self, job: DecoratedProductionJob
+    ) -> Optional[MiniJobRun]:
+        job_state = job.state.current_state
+        last_job_run_id = job.last_job_run_id
+        if job_state in _TERMINAL_STATES and last_job_run_id is None:
+            raise click.ClickException(
+                f"Can't find latest job run for {job_state} job."
+            )
+        if not last_job_run_id:
+            return None
+        return MiniJobRun(
+            last_job_run_id=last_job_run_id, job_state=job_state, error=job.state.error
+        )
+
+    def _wait_for_a_job_run(
+        self, job_id: Optional[str], job_name: Optional[str]
+    ) -> DecoratedProductionJob:
+        """Waits until the job has a job run and is in a non-pending state
+        (encountered when initially submitting).
+
+        Returns the job."""
+        start = time.monotonic()
+        job = self._resolve_job_object(job_id, job_name)
+        last_job_run_id: str = job.last_job_run_id
+        job_state = job.state.current_state
+        with self.log.spinner("Waiting for job run...") as spinner:
+            while last_job_run_id is None or job_state in _PENDING_STATES:
+                if job_state in _TERMINAL_STATES:
+                    raise click.ClickException(
+                        f"Can't find latest job run for {job_state} job."
+                    )
+                if self._cluster_journal_events_start_line == 0:
+                    spinner.text = (
+                        f"Waiting for a job run, current state is {job_state}..."
+                    )
+                time.sleep(max(0, 3 - (time.monotonic() - start)))
+                start = time.monotonic()
+                job = self._get_job(job.id)
+                last_job_run_id = job.last_job_run_id
+                if job.state.cluster_id:
+                    # Once we have a cluster id, we start printing the cluster journal event logs with each loop.
+                    cluster_journal_events = self._api_client.get_startup_logs_api_v2_sessions_session_id_startup_logs_get(
+                        job.state.cluster_id,
+                        start_line=self._cluster_journal_events_start_line,
+                        end_line=10000000,
+                    ).result
+                    lines = cluster_journal_events.lines.splitlines()
+                    if len(lines):
+                        spinner.succeed("Job run found. Cluster launching...")
+                    self._cluster_journal_events_start_line = (
+                        cluster_journal_events.start_line
+                        + cluster_journal_events.num_lines
+                    )
+
+                    for line in lines:
+                        self.log.info(line)
+
+                job_state = job.state.current_state
+        return job
+
+    def _print_logs_using_streaming_job_logs(
+        self,
+        job_run_id: str,
+        follow: bool = False,
+        no_cluster_journal_events: bool = False,
+    ) -> None:
+        """Retrieves logs from the streaming job logs folder in S3/GCS"""
+        logs = ""
+        next_page_token = None
+        # Keep track of if any logs have been printed yet,
+        # If not, we assume the cluster recently launched and we
+        # give more leeway for the logs to be available.
+        have_jobs_started_producing_logs = False
+        with self.log.spinner("Retrieving logs..."):
+            (
+                logs,
+                next_page_token,
+                cluster_journal_events_start_line,
+                job_finished,
+            ) = asyncio.run(
+                _get_job_logs_from_storage_bucket_streaming(
+                    self._api_client,
+                    self.log,
+                    job_run_id=job_run_id,
+                    remove_escape_chars=False,
+                    next_page_token=next_page_token,
+                    cluster_journal_events_start_line=self._cluster_journal_events_start_line,
+                    no_cluster_journal_events=no_cluster_journal_events,
+                )
+            )
+
+        if len(logs):
+            have_jobs_started_producing_logs = True
+
+        for line in logs.splitlines():
+            self.log.log(line)
+
+        # Continuously poll for logs.
+        terminal_state_count = 0
+        while follow:
+            start = time.monotonic()
+            if job_finished:
+                # Let's wait before terminating the loop to give the logs a chance to catch up.
+                terminal_state_count += 1
+                # Once jobs produce logs, we know we do not have to wait as long for those logs
+                # to be uploaded to S3 since the cluster is already running.
+                max_iterations = 4 if have_jobs_started_producing_logs else 6
+                if terminal_state_count >= max_iterations:
+                    self.log.info("Job run reached terminal state.")
+                    return
+
+            (
+                logs,
+                next_page_token,
+                cluster_journal_events_start_line,
+                job_finished,
+            ) = asyncio.run(
+                _get_job_logs_from_storage_bucket_streaming(
+                    self._api_client,
+                    self.log,
+                    job_run_id=job_run_id,
+                    remove_escape_chars=False,
+                    next_page_token=next_page_token,
+                    cluster_journal_events_start_line=cluster_journal_events_start_line,
+                    no_cluster_journal_events=no_cluster_journal_events,
+                )
+            )
+
+            for line in logs.splitlines():
+                self.log.log(line)
+
+            # Wait at least 5 seconds between iterations.
+            time.sleep(max(0, 5 - (time.monotonic() - start)))
+
+    def logs(
+        self,
+        job_id: Optional[str] = None,
+        job_name: Optional[str] = None,
+        should_follow: bool = False,
+        all_attempts: bool = False,
+        no_cluster_journal_events: bool = False,
+    ) -> None:
+        """
+        Fetches logs for a production job using the streaming log S3 bucket.
+        Params:
+        - no_cluster_journal_events: Controls whether we print out cluster journal logs. (Normally printed to stderr)
+        """
+        job = self._wait_for_a_job_run(job_id, job_name)
+        last_job_run_id: str = job.last_job_run_id
+        job_run_ids: List[str] = (
+            [last_job_run_id]
+            if not all_attempts
+            else [job_run.id for job_run in self._get_all_job_runs(job.id)]
+        )  # type: ignore
+        for job_run_id in job_run_ids:
+            self.log.open_block(
+                job_run_id, f"Job Run Id: {job_run_id}", auto_close=True,
+            )
+            self._print_logs_using_streaming_job_logs(
+                job_run_id,
+                follow=should_follow and job_run_id == last_job_run_id,
+                no_cluster_journal_events=no_cluster_journal_events,
+            )
+
+    def _generate_random_job_name(self) -> str:
+        """Generates a random job name
+        Format:
+            cli-job-{10 random characters}
+        """
+        random_chars = "".join(random.choices(string.ascii_letters, k=10)).lower()
+        return "cli-job-" + random_chars
+
+    def _get_all_job_runs(self, job_id: str) -> List[Job]:
+        """Returns all job runs for a given job id.
+        Returned in ascending order by creation time."""
+        job_runs: List[Job] = search_entities(
+            self.anyscale_api_client.search_jobs,
+            JobsQuery(
+                ha_job_id=job_id,
+                show_ray_client_runs_only=False,
+                sort_by_clauses=[
+                    SortByClauseJobsSortField(
+                        sort_field=JobsSortField.CREATED_AT, sort_order=SortOrder.ASC,
+                    )
+                ],
+                paging=PageQuery(),
+            ),
+        )
+        return job_runs