skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/jobs/client/sdk.py

@@ -0,0 +1,302 @@
+"""SDK functions for managed jobs."""
+import json
+import typing
+from typing import Dict, List, Optional, Union
+import webbrowser
+
+import click
+import requests
+
+from sky import sky_logging
+from sky.client import common as client_common
+from sky.client import sdk
+from sky.server import common as server_common
+from sky.server.requests import payloads
+from sky.skylet import constants
+from sky.usage import usage_lib
+from sky.utils import common_utils
+from sky.utils import dag_utils
+
+if typing.TYPE_CHECKING:
+    import io
+
+    import sky
+
+logger = sky_logging.init_logger(__name__)
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+def launch(
+    task: Union['sky.Task', 'sky.Dag'],
+    name: Optional[str] = None,
+    # Internal only:
+    # pylint: disable=invalid-name
+    _need_confirmation: bool = False,
+) -> server_common.RequestId:
+    """Launches a managed job.
+
+    Please refer to sky.cli.job_launch for documentation.
+
+    Args:
+        task: sky.Task, or sky.Dag (experimental; 1-task only) to launch as a
+            managed job.
+        name: Name of the managed job.
+        _need_confirmation: (Internal only) Whether to show a confirmation
+            prompt before launching the job.
+
+    Returns:
+        The request ID of the launch request.
+
+    Request Returns:
+        job_id (Optional[int]): Job ID for the managed job
+        controller_handle (Optional[ResourceHandle]): ResourceHandle of the
+          controller
+
+    Request Raises:
+        ValueError: cluster does not exist. Or, the entrypoint is not a valid
+          chain dag.
+        sky.exceptions.NotSupportedError: the feature is not supported.
+    """
+
+    dag = dag_utils.convert_entrypoint_to_dag(task)
+    sdk.validate(dag)
+    if _need_confirmation:
+        request_id = sdk.optimize(dag)
+        sdk.stream_and_get(request_id)
+        prompt = f'Launching a managed job {dag.name!r}. Proceed?'
+        if prompt is not None:
+            click.confirm(prompt, default=True, abort=True, show_default=True)
+
+    dag = client_common.upload_mounts_to_api_server(dag)
+    dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)
+    body = payloads.JobsLaunchBody(
+        task=dag_str,
+        name=name,
+    )
+    response = requests.post(
+        f'{server_common.get_server_url()}/jobs/launch',
+        json=json.loads(body.model_dump_json()),
+        timeout=(5, None),
+    )
+    return server_common.get_request_id(response)
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+def queue(refresh: bool,
+          skip_finished: bool = False,
+          all_users: bool = False) -> server_common.RequestId:
+    """Gets statuses of managed jobs.
+
+    Please refer to sky.cli.job_queue for documentation.
+
+    Args:
+        refresh: Whether to restart the jobs controller if it is stopped.
+        skip_finished: Whether to skip finished jobs.
+        all_users: Whether to show all users' jobs.
+
+    Returns:
+        The request ID of the queue request.
+
+    Request Returns:
+        job_records (List[Dict[str, Any]]): A list of dicts, with each dict
+          containing the information of a job.
+
+          .. code-block:: python
+
+            [
+              {
+                'job_id': (int) job id,
+                'job_name': (str) job name,
+                'resources': (str) resources of the job,
+                'submitted_at': (float) timestamp of submission,
+                'end_at': (float) timestamp of end,
+                'duration': (float) duration in seconds,
+                'recovery_count': (int) Number of retries,
+                'status': (sky.jobs.ManagedJobStatus) of the job,
+                'cluster_resources': (str) resources of the cluster,
+                'region': (str) region of the cluster,
+              }
+            ]
+
+    Request Raises:
+        sky.exceptions.ClusterNotUpError: the jobs controller is not up or
+          does not exist.
+        RuntimeError: if failed to get the managed jobs with ssh.
+    """
+    body = payloads.JobsQueueBody(
+        refresh=refresh,
+        skip_finished=skip_finished,
+        all_users=all_users,
+    )
+    response = requests.post(
+        f'{server_common.get_server_url()}/jobs/queue',
+        json=json.loads(body.model_dump_json()),
+        timeout=(5, None),
+    )
+    return server_common.get_request_id(response=response)
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+def cancel(
+    name: Optional[str] = None,
+    job_ids: Optional[List[int]] = None,
+    all: bool = False,  # pylint: disable=redefined-builtin
+    all_users: bool = False,
+) -> server_common.RequestId:
+    """Cancels managed jobs.
+
+    Please refer to sky.cli.job_cancel for documentation.
+
+    Args:
+        name: Name of the managed job to cancel.
+        job_ids: IDs of the managed jobs to cancel.
+        all: Whether to cancel all managed jobs.
+        all_users: Whether to cancel all managed jobs from all users.
+
+    Returns:
+        The request ID of the cancel request.
+
+    Request Raises:
+        sky.exceptions.ClusterNotUpError: the jobs controller is not up.
+        RuntimeError: failed to cancel the job.
+    """
+    body = payloads.JobsCancelBody(
+        name=name,
+        job_ids=job_ids,
+        all=all,
+        all_users=all_users,
+    )
+    response = requests.post(
+        f'{server_common.get_server_url()}/jobs/cancel',
+        json=json.loads(body.model_dump_json()),
+        timeout=(5, None),
+    )
+    return server_common.get_request_id(response=response)
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+def tail_logs(name: Optional[str] = None,
+              job_id: Optional[int] = None,
+              follow: bool = True,
+              controller: bool = False,
+              refresh: bool = False,
+              output_stream: Optional['io.TextIOBase'] = None) -> None:
+    """Tails logs of managed jobs.
+
+    You can provide either a job name or a job ID to tail logs. If both are not
+    provided, the logs of the latest job will be shown.
+
+    Args:
+        name: Name of the managed job to tail logs.
+        job_id: ID of the managed job to tail logs.
+        follow: Whether to follow the logs.
+        controller: Whether to tail logs from the jobs controller.
+        refresh: Whether to restart the jobs controller if it is stopped.
+        output_stream: The stream to write the logs to. If None, print to the
+            console.
+
+    Request Raises:
+        ValueError: invalid arguments.
+        sky.exceptions.ClusterNotUpError: the jobs controller is not up.
+    """
+    body = payloads.JobsLogsBody(
+        name=name,
+        job_id=job_id,
+        follow=follow,
+        controller=controller,
+        refresh=refresh,
+    )
+    response = requests.post(
+        f'{server_common.get_server_url()}/jobs/logs',
+        json=json.loads(body.model_dump_json()),
+        stream=True,
+        timeout=(5, None),
+    )
+    request_id = server_common.get_request_id(response)
+    sdk.stream_response(request_id, response, output_stream)
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+def download_logs(
+        name: Optional[str],
+        job_id: Optional[int],
+        refresh: bool,
+        controller: bool,
+        local_dir: str = constants.SKY_LOGS_DIRECTORY) -> Dict[int, str]:
+    """Sync down logs of managed jobs.
+
+    Please refer to sky.cli.job_logs for documentation.
+
+    Args:
+        name: Name of the managed job to sync down logs.
+        job_id: ID of the managed job to sync down logs.
+        refresh: Whether to restart the jobs controller if it is stopped.
+        controller: Whether to sync down logs from the jobs controller.
+        local_dir: Local directory to sync down logs.
+
+    Returns:
+        A dictionary mapping job ID to the local path.
+
+    Request Raises:
+        ValueError: invalid arguments.
+        sky.exceptions.ClusterNotUpError: the jobs controller is not up.
+    """
+
+    body = payloads.JobsDownloadLogsBody(
+        name=name,
+        job_id=job_id,
+        refresh=refresh,
+        controller=controller,
+        local_dir=local_dir,
+    )
+    response = requests.post(
+        f'{server_common.get_server_url()}/jobs/download_logs',
+        json=json.loads(body.model_dump_json()),
+        timeout=(5, None),
+    )
+    job_id_remote_path_dict = sdk.stream_and_get(
+        server_common.get_request_id(response))
+    remote2local_path_dict = client_common.download_logs_from_api_server(
+        job_id_remote_path_dict.values())
+    return {
+        job_id: remote2local_path_dict[remote_path]
+        for job_id, remote_path in job_id_remote_path_dict.items()
+    }
+
+
+spot_launch = common_utils.deprecated_function(
+    launch,
+    name='sky.jobs.launch',
+    deprecated_name='spot_launch',
+    removing_version='0.8.0',
+    override_argument={'use_spot': True})
+spot_queue = common_utils.deprecated_function(queue,
+                                              name='sky.jobs.queue',
+                                              deprecated_name='spot_queue',
+                                              removing_version='0.8.0')
+spot_cancel = common_utils.deprecated_function(cancel,
+                                               name='sky.jobs.cancel',
+                                               deprecated_name='spot_cancel',
+                                               removing_version='0.8.0')
+spot_tail_logs = common_utils.deprecated_function(
+    tail_logs,
+    name='sky.jobs.tail_logs',
+    deprecated_name='spot_tail_logs',
+    removing_version='0.8.0')
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+def dashboard() -> None:
+    """Starts a dashboard for managed jobs."""
+    user_hash = common_utils.get_user_hash()
+    api_server_url = server_common.get_server_url()
+    params = f'user_hash={user_hash}'
+    url = f'{api_server_url}/jobs/dashboard?{params}'
+    logger.info(f'Opening dashboard in browser: {url}')
+    webbrowser.open(url)

sky/jobs/constants.py

@@ -1,27 +1,65 @@
 """Constants used for Managed Jobs."""
+from typing import Dict, Union
+
+from sky.skylet import constants as skylet_constants
 
 JOBS_CONTROLLER_TEMPLATE = 'jobs-controller.yaml.j2'
 JOBS_CONTROLLER_YAML_PREFIX = '~/.sky/jobs_controller'
+JOBS_CONTROLLER_LOGS_DIR = '~/sky_logs/jobs_controller'
 
 JOBS_TASK_YAML_PREFIX = '~/.sky/managed_jobs'
 
 # Resources as a dict for the jobs controller.
-# Use default CPU instance type for jobs controller with >= 24GB, i.e.
-# m6i.2xlarge (8vCPUs, 32 GB) for AWS, Standard_D8s_v4 (8vCPUs, 32 GB)
-# for Azure, and n1-standard-8 (8 vCPUs, 32 GB) for GCP, etc.
-# Based on profiling, memory should be at least 3x (in GB) as num vCPUs to avoid
-# OOM (each vCPU can have 4 jobs controller processes as we set the CPU
-# requirement to 0.25, and 3 GB is barely enough for 4 job processes).
+# Use smaller CPU instance type for jobs controller, but with more memory, i.e.
+# r6i.xlarge (4vCPUs, 32 GB) for AWS, Standard_E4s_v5 (4vCPUs, 32 GB) for Azure,
+# and n2-highmem-4 (4 vCPUs, 32 GB) for GCP, etc.
+# Concurrently limits are set based on profiling. 4x num vCPUs is the launch
+# parallelism limit, and memory / 350MB is the limit to concurrently running
+# jobs. See _get_launch_parallelism and _get_job_parallelism in scheduler.py.
 # We use 50 GB disk size to reduce the cost.
-CONTROLLER_RESOURCES = {'cpus': '8+', 'memory': '3x', 'disk_size': 50}
+CONTROLLER_RESOURCES: Dict[str, Union[str, int]] = {
+    'cpus': '4+',
+    'memory': '8x',
+    'disk_size': 50
+}
 
+# TODO(zhwu): This is no longer accurate, after #4592, which increases the
+# length of user hash appended to the cluster name from 4 to 8 chars. This makes
+# the cluster name on GCP being wrapped twice. However, we cannot directly
+# update this constant, because the job cluster cleanup and many other logic
+# in managed jobs depends on this constant, i.e., updating this constant will
+# break backward compatibility and existing jobs.
+#
 # Max length of the cluster name for GCP is 35, the user hash to be attached is
-# 4+1 chars, and we assume the maximum length of the job id is 4+1, so the max
-# length of the cluster name prefix is 25 to avoid the cluster name being too
-# long and truncated twice during the cluster creation.
+# 4(now 8)+1 chars, and we assume the maximum length of the job id is
+# 4(now 8)+1, so the max length of the cluster name prefix is 25(should be 21
+# now) to avoid the cluster name being too long and truncated twice during the
+# cluster creation.
 JOBS_CLUSTER_NAME_PREFIX_LENGTH = 25
 
 # The version of the lib files that jobs/utils use. Whenever there is an API
 # change for the jobs/utils, we need to bump this version and update
 # job.utils.ManagedJobCodeGen to handle the version update.
-MANAGED_JOBS_VERSION = 1
+MANAGED_JOBS_VERSION = 2
+
+# The command for setting up the jobs dashboard on the controller. It firstly
+# checks if the systemd services are available, and if not (e.g., Kubernetes
+# containers may not have systemd), it starts the dashboard manually.
+DASHBOARD_SETUP_CMD = (
+    'if command -v systemctl &>/dev/null && systemctl --user show &>/dev/null; '
+    'then '
+    '  systemctl --user daemon-reload; '
+    '  systemctl --user enable --now skypilot-dashboard; '
+    'else '
+    '  echo "Systemd services not found. Starting SkyPilot dashboard '
+    'manually."; '
+    # Kill any old dashboard processes;
+    '  ps aux | grep -v nohup | grep -v grep | '
+    '  grep -- \'-m sky.jobs.dashboard.dashboard\' | awk \'{print $2}\' | '
+    '  xargs kill > /dev/null 2>&1 || true;'
+    # Launch the dashboard in the background if not already running
+    '  (ps aux | grep -v nohup | grep -v grep | '
+    '  grep -q -- \'-m sky.jobs.dashboard.dashboard\') || '
+    f'(nohup {skylet_constants.SKY_PYTHON_CMD} -m sky.jobs.dashboard.dashboard '
+    '>> ~/.sky/job-dashboard.log 2>&1 &); '
+    'fi')