skypilot-nightly 1.0.0.dev20250215__py3-none-any.whl → 1.0.0.dev20250217__py3-none-any.whl

sky/jobs/{core.py → server/core.py}

@@ -1,18 +1,21 @@
 """SDK functions for managed jobs."""
 import os
+import signal
+import subprocess
 import tempfile
+import time
 import typing
 from typing import Any, Dict, List, Optional, Tuple, Union
 import uuid
 
 import colorama
 
-import sky
 from sky import backends
+from sky import core
 from sky import exceptions
+from sky import execution
 from sky import provision as provision_lib
 from sky import sky_logging
-from sky import status_lib
 from sky import task as task_lib
 from sky.backends import backend_utils
 from sky.clouds.service_catalog import common as service_catalog_common
@@ -26,13 +29,17 @@ from sky.utils import common_utils
 from sky.utils import controller_utils
 from sky.utils import dag_utils
 from sky.utils import rich_utils
+from sky.utils import status_lib
 from sky.utils import subprocess_utils
 from sky.utils import timeline
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
+    import sky
     from sky.backends import cloud_vm_ray_backend
 
+logger = sky_logging.init_logger(__name__)
+
 
 @timeline.event
 @usage_lib.entrypoint
@@ -40,10 +47,9 @@ def launch(
     task: Union['sky.Task', 'sky.Dag'],
     name: Optional[str] = None,
     stream_logs: bool = True,
-    detach_run: bool = False,
 ) -> Tuple[Optional[int], Optional[backends.ResourceHandle]]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
-    """Launch a managed job.
+    """Launches a managed job.
 
     Please refer to sky.cli.job_launch for documentation.
 
@@ -51,7 +57,6 @@ def launch(
         task: sky.Task, or sky.Dag (experimental; 1-task only) to launch as a
           managed job.
         name: Name of the managed job.
-        detach_run: Whether to detach the run.
 
     Raises:
         ValueError: cluster does not exist. Or, the entrypoint is not a valid
@@ -77,7 +82,7 @@ def launch(
         with ux_utils.print_exception_no_traceback():
             raise ValueError('Only single-task or chain DAG is '
                              f'allowed for job_launch. Dag: {dag}')
-
+    dag.validate()
     dag_utils.maybe_infer_and_fill_dag_and_task_names(dag)
 
     task_names = set()
@@ -120,6 +125,7 @@ def launch(
             'remote_user_config_path': remote_user_config_path,
             'modified_catalogs':
                 service_catalog_common.get_modified_catalog_file_mounts(),
+            'dashboard_setup_cmd': managed_job_constants.DASHBOARD_SETUP_CMD,
             **controller_utils.shared_controller_vars_to_fill(
                 controller_utils.Controllers.JOBS_CONTROLLER,
                 remote_user_config_path=remote_user_config_path,
@@ -143,15 +149,14 @@ def launch(
             f'{colorama.Fore.YELLOW}'
             f'Launching managed job {dag.name!r} from jobs controller...'
             f'{colorama.Style.RESET_ALL}')
-        return sky.launch(task=controller_task,
-                          stream_logs=stream_logs,
-                          cluster_name=controller_name,
-                          detach_run=detach_run,
-                          idle_minutes_to_autostop=skylet_constants.
-                          CONTROLLER_IDLE_MINUTES_TO_AUTOSTOP,
-                          retry_until_up=True,
-                          fast=True,
-                          _disable_controller_check=True)
+        return execution.launch(task=controller_task,
+                                cluster_name=controller_name,
+                                stream_logs=stream_logs,
+                                idle_minutes_to_autostop=skylet_constants.
+                                CONTROLLER_IDLE_MINUTES_TO_AUTOSTOP,
+                                retry_until_up=True,
+                                fast=True,
+                                _disable_controller_check=True)
 
 
 def queue_from_kubernetes_pod(
@@ -256,7 +261,22 @@ def _maybe_restart_controller(
     rich_utils.force_update_status(
         ux_utils.spinner_message(f'{spinner_message} - restarting '
                                  'controller'))
-    handle = sky.start(jobs_controller_type.value.cluster_name)
+    handle = core.start(cluster_name=jobs_controller_type.value.cluster_name)
+    # Make sure the dashboard is running when the controller is restarted.
+    # We should not directly use execution.launch() and have the dashboard cmd
+    # in the task setup because since we are using detached_setup, it will
+    # become a job on controller which messes up the job IDs (we assume the
+    # job ID in controller's job queue is consistent with managed job IDs).
+    with rich_utils.safe_status(
+            ux_utils.spinner_message('Starting dashboard...')):
+        runner = handle.get_command_runners()[0]
+        user_hash = common_utils.get_user_hash()
+        runner.run(
+            f'export '
+            f'{skylet_constants.USER_ID_ENV_VAR}={user_hash!r}; '
+            f'{managed_job_constants.DASHBOARD_SETUP_CMD}',
+            stream_logs=True,
+        )
     controller_status = status_lib.ClusterStatus.UP
     rich_utils.force_update_status(ux_utils.spinner_message(spinner_message))
 
@@ -267,7 +287,7 @@ def _maybe_restart_controller(
 @usage_lib.entrypoint
 def queue(refresh: bool, skip_finished: bool = False) -> List[Dict[str, Any]]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
-    """Get statuses of managed jobs.
+    """Gets statuses of managed jobs.
 
     Please refer to sky.cli.job_queue for documentation.
 
@@ -307,14 +327,10 @@ def queue(refresh: bool, skip_finished: bool = False) -> List[Dict[str, Any]]:
         stream_logs=False,
         separate_stderr=True)
 
-    try:
-        subprocess_utils.handle_returncode(returncode,
-                                           code,
-                                           'Failed to fetch managed jobs',
-                                           job_table_payload + stderr,
-                                           stream_logs=False)
-    except exceptions.CommandError as e:
-        raise RuntimeError(str(e)) from e
+    if returncode != 0:
+        logger.error(job_table_payload + stderr)
+        raise RuntimeError('Failed to fetch managed jobs with returncode: '
+                           f'{returncode}')
 
     jobs = managed_job_utils.load_managed_job_queue(job_table_payload)
     if skip_finished:
@@ -334,7 +350,7 @@ def cancel(name: Optional[str] = None,
            job_ids: Optional[List[int]] = None,
            all: bool = False) -> None:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
-    """Cancel managed jobs.
+    """Cancels managed jobs.
 
     Please refer to sky.cli.job_cancel for documentation.
 
@@ -428,22 +444,73 @@ def tail_logs(name: Optional[str], job_id: Optional[int], follow: bool,
                                   controller=controller)
 
 
+def start_dashboard_forwarding(refresh: bool = False) -> Tuple[int, int]:
+    """Opens a dashboard for managed jobs (needs controller to be UP)."""
+    # TODO(SKY-1212): ideally, the controller/dashboard server should expose the
+    # API perhaps via REST. Then here we would (1) not have to use SSH to try to
+    # see if the controller is UP first, which is slow; (2) not have to run SSH
+    # port forwarding first (we'd just launch a local dashboard which would make
+    # REST API calls to the controller dashboard server).
+    logger.info('Starting dashboard')
+    hint = ('Dashboard is not available if jobs controller is not up. Run '
+            'a managed job first or run: sky jobs queue --refresh')
+    handle = _maybe_restart_controller(
+        refresh=refresh,
+        stopped_message=hint,
+        spinner_message='Checking jobs controller')
+
+    # SSH forward a free local port to remote's dashboard port.
+    remote_port = skylet_constants.SPOT_DASHBOARD_REMOTE_PORT
+    free_port = common_utils.find_free_port(remote_port)
+    runner = handle.get_command_runners()[0]
+    port_forward_command = ' '.join(
+        runner.port_forward_command(port_forward=[(free_port, remote_port)],
+                                    connect_timeout=1))
+    port_forward_command = (
+        f'{port_forward_command} '
+        f'> ~/sky_logs/api_server/dashboard-{common_utils.get_user_hash()}.log '
+        '2>&1')
+    logger.info(f'Forwarding port: {colorama.Style.DIM}{port_forward_command}'
+                f'{colorama.Style.RESET_ALL}')
+
+    ssh_process = subprocess.Popen(port_forward_command,
+                                   shell=True,
+                                   start_new_session=True)
+    time.sleep(3)  # Added delay for ssh_command to initialize.
+    logger.info(f'{colorama.Fore.GREEN}Dashboard is now available at: '
+                f'http://127.0.0.1:{free_port}{colorama.Style.RESET_ALL}')
+
+    return free_port, ssh_process.pid
+
+
+def stop_dashboard_forwarding(pid: int) -> None:
+    # Exit the ssh command when the context manager is closed.
+    try:
+        os.killpg(os.getpgid(pid), signal.SIGTERM)
+    except ProcessLookupError:
+        # This happens if jobs controller is auto-stopped.
+        pass
+    logger.info('Forwarding port closed. Exiting.')
+
+
 @usage_lib.entrypoint
-def sync_down_logs(
+def download_logs(
         name: Optional[str],
         job_id: Optional[int],
         refresh: bool,
         controller: bool,
-        local_dir: str = skylet_constants.SKY_LOGS_DIRECTORY) -> None:
+        local_dir: str = skylet_constants.SKY_LOGS_DIRECTORY) -> Dict[str, str]:
     """Sync down logs of managed jobs.
 
     Please refer to sky.cli.job_logs for documentation.
 
+    Returns:
+        A dictionary mapping job ID to the local path.
+
     Raises:
         ValueError: invalid arguments.
         sky.exceptions.ClusterNotUpError: the jobs controller is not up.
     """
-    # TODO(zhwu): Automatically restart the jobs controller
     if name is not None and job_id is not None:
         with ux_utils.print_exception_no_traceback():
             raise ValueError('Cannot specify both name and job_id.')
@@ -467,8 +534,8 @@ def sync_down_logs(
     backend = backend_utils.get_backend_from_handle(handle)
     assert isinstance(backend, backends.CloudVmRayBackend), backend
 
-    backend.sync_down_managed_job_logs(handle,
-                                       job_id=job_id,
-                                       job_name=name,
-                                       controller=controller,
-                                       local_dir=local_dir)
+    return backend.sync_down_managed_job_logs(handle,
+                                              job_id=job_id,
+                                              job_name=name,
+                                              controller=controller,
+                                              local_dir=local_dir)

sky/jobs/server/dashboard_utils.py

@@ -0,0 +1,64 @@
+"""Persistent dashboard sessions."""
+import pathlib
+from typing import Tuple
+
+import filelock
+
+from sky.utils import db_utils
+
+
+def create_dashboard_table(cursor, conn):
+    cursor.execute("""\
+        CREATE TABLE IF NOT EXISTS dashboard_sessions (
+        user_hash TEXT PRIMARY KEY,
+        port INTEGER,
+        pid INTEGER)""")
+    conn.commit()
+
+
+def _get_db_path() -> str:
+    path = pathlib.Path('~/.sky/dashboard/sessions.db')
+    path = path.expanduser().absolute()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    return str(path)
+
+
+DB_PATH = _get_db_path()
+db_utils.SQLiteConn(DB_PATH, create_dashboard_table)
+LOCK_FILE_PATH = '~/.sky/dashboard/sessions-{user_hash}.lock'
+
+
+def get_dashboard_session(user_hash: str) -> Tuple[int, int]:
+    """Get the port and pid of the dashboard session for the user."""
+    with db_utils.safe_cursor(DB_PATH) as cursor:
+        cursor.execute(
+            'SELECT port, pid FROM dashboard_sessions WHERE user_hash=?',
+            (user_hash,))
+        result = cursor.fetchone()
+        if result is None:
+            return 0, 0
+        return result
+
+
+def add_dashboard_session(user_hash: str, port: int, pid: int) -> None:
+    """Add a dashboard session for the user."""
+    with db_utils.safe_cursor(DB_PATH) as cursor:
+        cursor.execute(
+            'INSERT OR REPLACE INTO dashboard_sessions (user_hash, port, pid) '
+            'VALUES (?, ?, ?)', (user_hash, port, pid))
+
+
+def remove_dashboard_session(user_hash: str) -> None:
+    """Remove the dashboard session for the user."""
+    with db_utils.safe_cursor(DB_PATH) as cursor:
+        cursor.execute('DELETE FROM dashboard_sessions WHERE user_hash=?',
+                       (user_hash,))
+    lock_path = pathlib.Path(LOCK_FILE_PATH.format(user_hash=user_hash))
+    lock_path.unlink(missing_ok=True)
+
+
+def get_dashboard_lock_for_user(user_hash: str) -> filelock.FileLock:
+    path = pathlib.Path(LOCK_FILE_PATH.format(user_hash=user_hash))
+    path = path.expanduser().absolute()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    return filelock.FileLock(path)

sky/jobs/server/server.py

@@ -0,0 +1,182 @@
+"""REST API for managed jobs."""
+import os
+
+import fastapi
+import httpx
+
+from sky import sky_logging
+from sky.jobs.server import core
+from sky.jobs.server import dashboard_utils
+from sky.server import common as server_common
+from sky.server import stream_utils
+from sky.server.requests import executor
+from sky.server.requests import payloads
+from sky.server.requests import requests as api_requests
+from sky.skylet import constants
+from sky.utils import common
+from sky.utils import common_utils
+
+logger = sky_logging.init_logger(__name__)
+
+router = fastapi.APIRouter()
+
+
+def _get_controller_name(request_body: payloads.RequestBody) -> str:
+    user_hash = request_body.user_hash
+    return common.get_controller_name(common.ControllerType.JOBS, user_hash)
+
+
+@router.post('/launch')
+async def launch(request: fastapi.Request,
+                 jobs_launch_body: payloads.JobsLaunchBody) -> None:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='jobs.launch',
+        request_body=jobs_launch_body,
+        func=core.launch,
+        schedule_type=api_requests.ScheduleType.LONG,
+        request_cluster_name=_get_controller_name(jobs_launch_body),
+    )
+
+
+@router.post('/queue')
+async def queue(request: fastapi.Request,
+                jobs_queue_body: payloads.JobsQueueBody) -> None:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='jobs.queue',
+        request_body=jobs_queue_body,
+        func=core.queue,
+        schedule_type=(api_requests.ScheduleType.LONG if jobs_queue_body.refresh
+                       else api_requests.ScheduleType.SHORT),
+        request_cluster_name=_get_controller_name(jobs_queue_body),
+    )
+
+
+@router.post('/cancel')
+async def cancel(request: fastapi.Request,
+                 jobs_cancel_body: payloads.JobsCancelBody) -> None:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='jobs.cancel',
+        request_body=jobs_cancel_body,
+        func=core.cancel,
+        schedule_type=api_requests.ScheduleType.SHORT,
+        request_cluster_name=_get_controller_name(jobs_cancel_body),
+    )
+
+
+@router.post('/logs')
+async def logs(
+    request: fastapi.Request, jobs_logs_body: payloads.JobsLogsBody,
+    background_tasks: fastapi.BackgroundTasks
+) -> fastapi.responses.StreamingResponse:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='jobs.logs',
+        request_body=jobs_logs_body,
+        func=core.tail_logs,
+        schedule_type=api_requests.ScheduleType.SHORT
+        if jobs_logs_body.refresh else api_requests.ScheduleType.LONG,
+        request_cluster_name=_get_controller_name(jobs_logs_body),
+    )
+    request_task = api_requests.get_request(request.state.request_id)
+
+    return stream_utils.stream_response(
+        request_id=request_task.request_id,
+        logs_path=request_task.log_path,
+        background_tasks=background_tasks,
+    )
+
+
+@router.post('/download_logs')
+async def download_logs(
+        request: fastapi.Request,
+        jobs_download_logs_body: payloads.JobsDownloadLogsBody) -> None:
+    user_hash = jobs_download_logs_body.env_vars[constants.USER_ID_ENV_VAR]
+    logs_dir_on_api_server = server_common.api_server_user_logs_dir_prefix(
+        user_hash)
+    logs_dir_on_api_server.expanduser().mkdir(parents=True, exist_ok=True)
+    # We should reuse the original request body, so that the env vars, such as
+    # user hash, are kept the same.
+    jobs_download_logs_body.local_dir = str(logs_dir_on_api_server)
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='jobs.download_logs',
+        request_body=jobs_download_logs_body,
+        func=core.download_logs,
+        schedule_type=api_requests.ScheduleType.LONG
+        if jobs_download_logs_body.refresh else api_requests.ScheduleType.SHORT,
+        request_cluster_name=_get_controller_name(jobs_download_logs_body),
+    )
+
+
+@router.get('/dashboard')
+async def dashboard(request: fastapi.Request,
+                    user_hash: str) -> fastapi.Response:
+    # Find the port for the dashboard of the user
+    os.environ[constants.USER_ID_ENV_VAR] = user_hash
+    server_common.reload_for_new_request(client_entrypoint=None,
+                                         client_command=None)
+    logger.info(f'Starting dashboard for user hash: {user_hash}')
+
+    with dashboard_utils.get_dashboard_lock_for_user(user_hash):
+        max_retries = 3
+        for attempt in range(max_retries):
+            port, pid = dashboard_utils.get_dashboard_session(user_hash)
+            if port == 0 or attempt > 0:
+                # Let the client know that we are waiting for starting the
+                # dashboard.
+                try:
+                    port, pid = core.start_dashboard_forwarding()
+                except Exception as e:  # pylint: disable=broad-except
+                    # We catch all exceptions to gracefully handle unknown
+                    # errors and raise an HTTPException to the client.
+                    msg = (
+                        'Dashboard failed to start: '
+                        f'{common_utils.format_exception(e, use_bracket=True)}')
+                    logger.error(msg)
+                    raise fastapi.HTTPException(status_code=503, detail=msg)
+                dashboard_utils.add_dashboard_session(user_hash, port, pid)
+
+            # Assuming the dashboard is forwarded to localhost on the API server
+            dashboard_url = f'http://localhost:{port}'
+            try:
+                # Ping the dashboard to check if it's still running
+                async with httpx.AsyncClient() as client:
+                    response = await client.request('GET',
+                                                    dashboard_url,
+                                                    timeout=1)
+                break  # Connection successful, proceed with the request
+            except Exception as e:  # pylint: disable=broad-except
+                # We catch all exceptions to gracefully handle unknown
+                # errors and retry or raise an HTTPException to the client.
+                msg = (
+                    f'Dashboard connection attempt {attempt + 1} failed with '
+                    f'{common_utils.format_exception(e, use_bracket=True)}')
+                logger.info(msg)
+                if attempt == max_retries - 1:
+                    raise fastapi.HTTPException(status_code=503, detail=msg)
+
+    # Create a client session to forward the request
+    try:
+        async with httpx.AsyncClient() as client:
+            # Make the request and get the response
+            response = await client.request(
+                method='GET',
+                url=f'{dashboard_url}',
+                headers=request.headers.raw,
+            )
+
+            # Create a new response with the content already read
+            content = await response.aread()
+            return fastapi.Response(
+                content=content,
+                status_code=response.status_code,
+                headers=dict(response.headers),
+                media_type=response.headers.get('content-type'))
+    except Exception as e:
+        msg = (f'Failed to forward request to dashboard: '
+               f'{common_utils.format_exception(e, use_bracket=True)}')
+        logger.error(msg)
+        raise fastapi.HTTPException(status_code=502, detail=msg)

sky/jobs/utils.py

@@ -20,7 +20,6 @@ import filelock
 import psutil
 from typing_extensions import Literal
 
-import sky
 from sky import backends
 from sky import exceptions
 from sky import global_user_state
@@ -35,21 +34,17 @@ from sky.skylet import log_lib
 from sky.usage import usage_lib
 from sky.utils import common_utils
 from sky.utils import log_utils
+from sky.utils import message_utils
 from sky.utils import rich_utils
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
+    import sky
     from sky import dag as dag_lib
 
 logger = sky_logging.init_logger(__name__)
 
-# Add user hash so that two users don't have the same controller VM on
-# shared-account clouds such as GCP.
-JOB_CONTROLLER_NAME: str = (
-    f'sky-jobs-controller-{common_utils.get_user_hash()}')
-LEGACY_JOB_CONTROLLER_NAME: str = (
-    f'sky-spot-controller-{common_utils.get_user_hash()}')
 SIGNAL_FILE_PREFIX = '/tmp/sky_jobs_controller_signal_{}'
 # Controller checks its job's status every this many seconds.
 JOB_STATUS_CHECK_GAP_SECONDS = 20
@@ -86,6 +81,7 @@ class UserSignal(enum.Enum):
 # ====== internal functions ======
 def terminate_cluster(cluster_name: str, max_retry: int = 6) -> None:
     """Terminate the cluster."""
+    from sky import core  # pylint: disable=import-outside-toplevel
     retry_cnt = 0
     # In some cases, e.g. botocore.exceptions.NoCredentialsError due to AWS
     # metadata service throttling, the failed sky.down attempt can take 10-11
@@ -102,7 +98,7 @@ def terminate_cluster(cluster_name: str, max_retry: int = 6) -> None:
     while True:
         try:
             usage_lib.messages.usage.set_internal()
-            sky.down(cluster_name)
+            core.down(cluster_name)
             return
         except exceptions.ClusterDoesNotExist:
             # The cluster is already down.
@@ -274,15 +270,23 @@ def update_managed_jobs_statuses(job_id: Optional[int] = None):
             failure_reason = ('Inconsistent internal job state. This is a bug.')
         elif pid is None:
             # Non-legacy job and controller process has not yet started.
-            if schedule_state in (
+            controller_status = job_lib.get_status(job_id)
+            if controller_status == job_lib.JobStatus.FAILED_SETUP:
+                # We should fail the case where the controller status is
+                # FAILED_SETUP, as it is due to the failure of dependency setup
+                # on the controller.
+                # TODO(cooperc): We should also handle the case where controller
+                # status is FAILED_DRIVER or FAILED.
+                logger.error('Failed to setup the cloud dependencies for '
+                             'the managed job.')
+            elif (schedule_state in [
                     managed_job_state.ManagedJobScheduleState.INACTIVE,
-                    managed_job_state.ManagedJobScheduleState.WAITING):
-                # For these states, the controller hasn't been started yet.
-                # This is expected.
+                    managed_job_state.ManagedJobScheduleState.WAITING,
+            ]):
+                # It is expected that the controller hasn't been started yet.
                 continue
-
-            if (schedule_state ==
-                    managed_job_state.ManagedJobScheduleState.LAUNCHING):
+            elif (schedule_state ==
+                  managed_job_state.ManagedJobScheduleState.LAUNCHING):
                 # This is unlikely but technically possible. There's a brief
                 # period between marking job as scheduled (LAUNCHING) and
                 # actually launching the controller process and writing the pid
@@ -368,7 +372,7 @@ def get_job_timestamp(backend: 'backends.CloudVmRayBackend', cluster_name: str,
     subprocess_utils.handle_returncode(returncode, code,
                                        'Failed to get job time.',
                                        stdout + stderr)
-    stdout = common_utils.decode_payload(stdout)
+    stdout = message_utils.decode_payload(stdout)
     return float(stdout)
 
 
@@ -531,7 +535,8 @@ def stream_logs_by_id(job_id: int, follow: bool = True) -> str:
                            f'{managed_job_state.get_failure_reason(job_id)}')
             log_file = managed_job_state.get_local_log_file(job_id, None)
             if log_file is not None:
-                with open(log_file, 'r', encoding='utf-8') as f:
+                with open(os.path.expanduser(log_file), 'r',
+                          encoding='utf-8') as f:
                     # Stream the logs to the console without reading the whole
                     # file into memory.
                     start_streaming = False
@@ -878,12 +883,12 @@ def dump_managed_job_queue() -> str:
             job['cluster_resources'] = '-'
             job['region'] = '-'
 
-    return common_utils.encode_payload(jobs)
+    return message_utils.encode_payload(jobs)
 
 
 def load_managed_job_queue(payload: str) -> List[Dict[str, Any]]:
     """Load job queue from json string."""
-    jobs = common_utils.decode_payload(payload)
+    jobs = message_utils.decode_payload(payload)
     for job in jobs:
         job['status'] = managed_job_state.ManagedJobStatus(job['status'])
     return jobs
@@ -1159,9 +1164,9 @@ class ManagedJobCodeGen:
     @classmethod
     def get_all_job_ids_by_name(cls, job_name: Optional[str]) -> str:
         code = textwrap.dedent(f"""\
-        from sky.utils import common_utils
+        from sky.utils import message_utils
         job_id = managed_job_state.get_all_job_ids_by_name({job_name!r})
-        print(common_utils.encode_payload(job_id), end="", flush=True)
+        print(message_utils.encode_payload(job_id), end="", flush=True)
         """)
         return cls._build(code)
 
@@ -1197,5 +1202,9 @@ class ManagedJobCodeGen:
     @classmethod
     def _build(cls, code: str) -> str:
         generated_code = cls._PREFIX + '\n' + code
-
-        return f'{constants.SKY_PYTHON_CMD} -u -c {shlex.quote(generated_code)}'
+        # Use the local user id to make sure the operation goes to the correct
+        # user.
+        return (
+            f'export {constants.USER_ID_ENV_VAR}='
+            f'"{common_utils.get_user_hash()}"; '
+            f'{constants.SKY_PYTHON_CMD} -u -c {shlex.quote(generated_code)}')

sky/models.py

@@ -0,0 +1,27 @@
+"""Data Models for SkyPilot."""
+
+import collections
+import dataclasses
+from typing import Dict, Optional
+
+
+@dataclasses.dataclass
+class User:
+    # User hash
+    id: str
+    # Display name of the user
+    name: Optional[str] = None
+
+
+RealtimeGpuAvailability = collections.namedtuple(
+    'RealtimeGpuAvailability', ['gpu', 'counts', 'capacity', 'available'])
+
+
+@dataclasses.dataclass
+class KubernetesNodeInfo:
+    """Dataclass to store Kubernetes node information."""
+    name: str
+    accelerator_type: Optional[str]
+    # Resources available on the node. E.g., {'nvidia.com/gpu': '2'}
+    total: Dict[str, int]
+    free: Dict[str, int]