treebeardtech
diff --git a/‎docs/source/reference/local/setup.rst
+2-2 b/‎docs/source/reference/local/setup.rst
+2-2
diff --git a/‎examples/local/cluster-config.yaml
+1-1 b/‎examples/local/cluster-config.yaml
+1-1
diff --git a/‎sky/backends/backend_utils.py
+1-1 b/‎sky/backends/backend_utils.py
+1-1
diff --git a/‎sky/backends/cloud_vm_ray_backend.py
+8-6 b/‎sky/backends/cloud_vm_ray_backend.py
+8-6
diff --git a/‎sky/design_docs/onprem-design.md
+2-2 b/‎sky/design_docs/onprem-design.md
+2-2
diff --git a/‎sky/setup_files/setup.py
+5-3 b/‎sky/setup_files/setup.py
+5-3
diff --git a/‎sky/skylet/LICENCE
+6-6 b/‎sky/skylet/LICENCE
+6-6
diff --git a/‎sky/skylet/constants.py
+1-1 b/‎sky/skylet/constants.py
+1-1
diff --git a/‎sky/skylet/job_lib.py
+16-6 b/‎sky/skylet/job_lib.py
+16-6
diff --git a/‎sky/skylet/providers/aws/cloudwatch/cloudwatch_helper.py
+7-5 b/‎sky/skylet/providers/aws/cloudwatch/cloudwatch_helper.py
+7-5
diff --git a/‎sky/skylet/providers/aws/config.py
+24-23 b/‎sky/skylet/providers/aws/config.py
+24-23
@@ -14,13 +14,13 @@ For further reference, `here <https://docs.ray.io/en/latest/ray-core/configure.h
 Installing SkyPilot dependencies
 -----------------------------------
 
-SkyPilot On-prem requires :code:`python3`, :code:`ray==1.13.0`, and :code:`sky` to be setup on all local nodes and globally available to all users.
+SkyPilot On-prem requires :code:`python3`, :code:`ray==2.0.1`, and :code:`sky` to be setup on all local nodes and globally available to all users.
 
 To install Ray and SkyPilot for all users, run the following commands on all local nodes:
 
 .. code-block:: console
 
-   $ pip3 install ray[default]==1.13.0
+   $ pip3 install ray[default]==2.0.1
 
    $ # SkyPilot requires python >= 3.6.
    $ pip3 install skypilot
 
@@ -4,7 +4,7 @@
 # The system administrator must have `sudo` access to the local nodes.
 # Requirements:
 #   1) Python (> 3.6) on all nodes.
-#   2) Ray CLI (= 1.13.0) on all nodes.
+#   2) Ray CLI (= 2.0.1) on all nodes.
 #
 # Example usage:
 #   >> sky admin deploy cluster-config.yaml
 
@@ -1418,7 +1418,7 @@ def _ray_launch_hash(cluster_name: str, ray_config: Dict[str, Any]) -> Set[str]:
         return set(ray_launch_hashes)
     with suppress_output():
         ray_config = ray_commands._bootstrap_config(ray_config)  # pylint: disable=protected-access
-    # Adopted from https://github.com/ray-project/ray/blob/ray-1.13.0/python/ray/autoscaler/_private/node_launcher.py#L56-L64
+    # Adopted from https://github.com/ray-project/ray/blob/ray-2.0.1/python/ray/autoscaler/_private/node_launcher.py#L87-L97
     # TODO(zhwu): this logic is duplicated from the ray code above (keep in sync).
     launch_hashes = set()
     head_node_type = ray_config['head_node_type']
 
@@ -154,7 +154,7 @@ def add_prologue(self,
         # Should use 'auto' or 'ray://<internal_head_ip>:10001' rather than
         # 'ray://localhost:10001', or 'ray://127.0.0.1:10001', for public cloud.
         # Otherwise, it will a bug of ray job failed to get the placement group
-        # in ray <= 2.0.0.
+        # in ray <= 2.0.1.
         # TODO(mluo): Check why 'auto' not working with on-prem cluster and
         # whether the placement group issue also occurs in on-prem cluster.
         ray_address = 'ray://localhost:10001' if is_local else 'auto'
@@ -401,6 +401,8 @@ def add_epilogue(self) -> None:
                 # Need this to set the job status in ray job to be FAILED.
                 sys.exit(1)
             else:
+                sys.stdout.flush()
+                sys.stderr.flush()
                 job_lib.set_status({self.job_id!r}, job_lib.JobStatus.SUCCEEDED)
                 # This waits for all streaming logs to finish.
                 time.sleep(1)
@@ -1341,7 +1343,7 @@ def _ensure_cluster_ray_started(self,
         if isinstance(launched_resources.cloud, clouds.Local):
             raise RuntimeError(
                 'The command `ray status` errored out on the head node '
-                'of the local cluster. Check if ray[default]==1.13.0 '
+                'of the local cluster. Check if ray[default]==2.0.1 '
                 'is installed or running correctly.')
         backend.run_on_head(handle, 'ray stop', use_cached_head_ip=False)
 
@@ -2066,8 +2068,8 @@ def _exec_code_on_head(
         else:
             job_submit_cmd = (
                 f'{cd} && mkdir -p {remote_log_dir} && ray job submit '
-                f'--address=http://127.0.0.1:8265 --job-id {ray_job_id} '
-                '--no-wait -- '
+                f'--address=http://127.0.0.1:8265 --submission-id {ray_job_id} '
+                '--no-wait '
                 f'"{executable} -u {script_path} > {remote_log_path} 2>&1"')
 
         returncode, stdout, stderr = self.run_on_head(handle,
@@ -2151,8 +2153,8 @@ def _setup_and_create_job_cmd_on_local_head(
         switch_user_cmd = ' '.join(switch_user_cmd)
         job_submit_cmd = (
             'ray job submit '
-            f'--address=http://127.0.0.1:8265 --job-id {ray_job_id} --no-wait '
-            f'-- {switch_user_cmd}')
+            f'--address=http://127.0.0.1:8265 --submission-id {ray_job_id} '
+            f'--no-wait -- {switch_user_cmd}')
         return job_submit_cmd
 
     def _add_job(self, handle: ResourceHandle, job_name: str,
 
@@ -9,7 +9,7 @@
 - Does not support different types of accelerators within the same node (intranode).
 
 ## Installing Ray and SkyPilot 
-- Admin installs Ray==1.13.0 and SkyPilot globally on all machines. It is assumed that the admin regularly keeps SkyPilot updated on the cluster.
+- Admin installs Ray==2.0.1 and SkyPilot globally on all machines. It is assumed that the admin regularly keeps SkyPilot updated on the cluster.
 - Python >= 3.6 for all users.
 - When a regular user runs `sky launch`, a local version of SkyPilot will be installed on the machine for each user. The local installation of Ray is specified in `sky/templates/local-ray.yml.j2`.
 
@@ -36,7 +36,7 @@ ray.get(ray.remote(f).remote())
 ```
 
 - Therefore, SkyPilot On-prem transparently includes user-switching so that SkyPilot tasks are still run as the calling, unprivileged user. This user-switching (`sudo -H su --login [USER]` in appropriate places) works as follows:
-    - In `sky/backends/cloud_vm_ray_backend.py::_setup_and_create_job_cmd_on_local_head`, switching between users is called during Ray job submission. The command `ray job submit --address=http://127.0.0.1:8265 --job-id {ray_job_id} -- sudo -H su --login [SSH_USER] -c \"[JOB_COMMAND]\"` switches job submission execution from admin back to the original user `SSH_USER`. The `JOB_COMMAND` argument runs a bash script with the user's run commands.
+    - In `sky/backends/cloud_vm_ray_backend.py::_setup_and_create_job_cmd_on_local_head`, switching between users is called during Ray job submission. The command `ray job submit --address=http://127.0.0.1:8265 --submission-id {ray_job_id} -- sudo -H su --login [SSH_USER] -c \"[JOB_COMMAND]\"` switches job submission execution from admin back to the original user `SSH_USER`. The `JOB_COMMAND` argument runs a bash script with the user's run commands.
     - In `sky/skylet/log_lib.py::run_bash_command_with_log`, there is also another `sudo -H su` command to switch users. The function `run_bash_command_with_log` is part of the `RayCodeGen` job execution script uploaded to remote for job submission (located in `~/.sky/sky_app/sky_app_[JOB_ID].py`). This program initially runs under the calling user, but it executes the function `run_bash_command_with_log` from the context of the admin, as the function is executed within the Ray cluster as a Ray remote function (see above for why all Ray remote functions are run under admin).
 - SkyPilot ensures Ray-related environment variables (that are critical for execution) are preserved across switching users (check with `examples/env_check.yaml`).
 
 
@@ -56,7 +56,7 @@ def parse_footnote(readme: str) -> str:
 
 install_requires = [
     'wheel',
-    # NOTE: ray 1.13.0 requires click<=8.0.4,>=7.0; We disable the
+    # NOTE: ray 2.0.1 requires click<=8.0.4,>=7.0; We disable the
     # shell completion for click<8.0 for backward compatibility.
     'click<=8.0.4,>=7.0',
     'colorama',
@@ -70,14 +70,14 @@ def parse_footnote(readme: str) -> str:
     'PrettyTable',
     # Lower local ray version is not fully supported, due to the
     # autoscaler issues (also tracked in #537).
-    'ray[default]>=1.9.0,<=1.13.0',
+    'ray[default]>=1.9.0,<=2.0.1',
     'rich',
     'tabulate',
     'filelock',  # TODO(mraheja): Enforce >=3.6.0 when python version is >= 3.7
     # This is used by ray. The latest 1.44.0 will generate an error
     # `Fork support is only compatible with the epoll1 and poll
     # polling strategies`
-    'grpcio<=1.43.0',
+    'grpcio>=1.32.0,<=1.43.0',
     'packaging',
     # The latest 4.21.1 will break ray. Enforce < 4.0.0 until Ray releases the
     # fix.
@@ -98,6 +98,8 @@ def parse_footnote(readme: str) -> str:
     ],
     # TODO(zongheng): azure-cli is huge and takes a long time to install.
     # Tracked in: https://github.com/Azure/azure-cli/issues/7387
+    # azure-cli need to be pinned to 2.31.0 due to later versions
+    # do not have azure-identity (used in node_provider) installed
     'azure': ['azure-cli==2.31.0', 'azure-core'],
     'gcp': ['google-api-python-client', 'google-cloud-storage'],
     'docker': ['docker'],
 
@@ -203,16 +203,16 @@
 --------------------------------------------------------------------------------
 
 Code in providers/azure from
-https://github.com/ray-project/ray/tree/ray-1.13.0/python/ray/autoscaler/_private/_azure
-Git commit of the release 1.13.0: 4ce38d001dbbe09cd21c497fedd03d692b2be3e
+https://github.com/ray-project/ray/tree/ray-2.0.1/python/ray/autoscaler/_private/_azure
+Git commit of the release 2.0.1: 03b6bc7b5a305877501110ec04710a9c57011479
 
 Code in providers/gcp from
-https://github.com/ray-project/ray/tree/ray-1.13.0/python/ray/autoscaler/_private/gcp
-Git commit of the release 1.13.0: 4ce38d001dbbe09cd21c497fedd03d692b2be3e
+https://github.com/ray-project/ray/tree/ray-2.0.1/python/ray/autoscaler/_private/gcp
+Git commit of the release 2.0.1: 03b6bc7b5a305877501110ec04710a9c57011479
 
 Code in providers/aws from
-https://github.com/ray-project/ray/tree/ray-1.13.0/python/ray/autoscaler/_private/aws
-Git commit of the release 1.13.0: 4ce38d001dbbe09cd21c497fedd03d692b2be3e
+https://github.com/ray-project/ray/tree/ray-2.0.1/python/ray/autoscaler/_private/aws
+Git commit of the release 2.0.1: 03b6bc7b5a305877501110ec04710a9c57011479
 
 
 Copyright 2016-2022 Ray developers
 
@@ -2,7 +2,7 @@
 
 SKY_LOGS_DIRECTORY = '~/sky_logs'
 SKY_REMOTE_WORKDIR = '~/sky_workdir'
-SKY_REMOTE_RAY_VERSION = '1.13.0'
+SKY_REMOTE_RAY_VERSION = '2.0.1'
 
 # TODO(mluo): Make explicit `sky launch -c <name> ''` optional.
 UNINITIALIZED_ONPREM_CLUSTER_MESSAGE = (
 
@@ -7,6 +7,7 @@
 import pathlib
 import shlex
 import time
+import typing
 from typing import Any, Dict, List, Optional
 
 import filelock
@@ -17,6 +18,9 @@
 from sky.utils import db_utils
 from sky.utils import log_utils
 
+if typing.TYPE_CHECKING:
+    from ray.dashboard.modules.job import pydantic_models as ray_pydantic
+
 logger = sky_logging.init_logger(__name__)
 
 _JOB_STATUS_LOCK = '~/.sky/locks/.job_{}.lock'
@@ -331,7 +335,7 @@ def update_job_status(job_owner: str,
     we still need this to handle staleness problem, caused by instance
     restarting and other corner cases (if any).
 
-    This function should only be run on the remote instance with ray==1.13.0.
+    This function should only be run on the remote instance with ray==2.0.1.
     """
     if len(job_ids) == 0:
         return []
@@ -341,13 +345,19 @@ def update_job_status(job_owner: str,
 
     job_client = _create_ray_job_submission_client()
 
-    # In ray 1.13.0, job_client.list_jobs returns a dict of job_id to job_info,
-    # where job_info contains the job status (str).
-    ray_job_infos = job_client.list_jobs()
+    # In ray 2.0.1, job_client.list_jobs returns a list of JobDetails,
+    # which contains the job status (str) and submission_id (str).
+    job_detail_lists: List['ray_pydantic.JobDetails'] = job_client.list_jobs()
+
+    job_details = dict()
+    ray_job_ids_set = set(ray_job_ids)
+    for job_detail in job_detail_lists:
+        if job_detail.submission_id in ray_job_ids_set:
+            job_details[job_detail.submission_id] = job_detail
     job_statuses: List[JobStatus] = [None] * len(ray_job_ids)
     for i, ray_job_id in enumerate(ray_job_ids):
-        if ray_job_id in ray_job_infos:
-            ray_status = ray_job_infos[ray_job_id].status
+        if ray_job_id in job_details:
+            ray_status = job_details[ray_job_id].status
             job_statuses[i] = _RAY_TO_JOB_STATUS_MAP[ray_status]
 
     assert len(job_statuses) == len(job_ids), (job_statuses, job_ids)
 
@@ -1,13 +1,15 @@
-import botocore
 import copy
+import hashlib
 import json
-import os
 import logging
+import os
 import time
-import hashlib
-from typing import Any, Dict, List, Union, Tuple
+from typing import Any, Dict, List, Tuple, Union
+
+import botocore
+
 from sky.skylet.providers.aws.utils import client_cache, resource_cache
-from ray.autoscaler.tags import TAG_RAY_CLUSTER_NAME, NODE_KIND_HEAD, TAG_RAY_NODE_KIND
+from ray.autoscaler.tags import NODE_KIND_HEAD, TAG_RAY_CLUSTER_NAME, TAG_RAY_NODE_KIND
 
 logger = logging.getLogger(__name__)
 
 
@@ -1,30 +1,29 @@
-from distutils.version import StrictVersion
-from functools import lru_cache
-from functools import partial
 import copy
 import itertools
 import json
+import logging
 import os
 import time
+from distutils.version import StrictVersion
+from functools import lru_cache, partial
 from typing import Any, Dict, List, Optional, Set, Tuple
-import logging
 
 import boto3
 import botocore
 
-from ray.autoscaler._private.util import check_legacy_fields
-from ray.autoscaler.tags import NODE_TYPE_LEGACY_HEAD, NODE_TYPE_LEGACY_WORKER
-from ray.autoscaler._private.providers import _PROVIDER_PRETTY_NAMES
+from sky.skylet.providers.aws.cloudwatch.cloudwatch_helper import (
+    CloudwatchHelper as cwh,
+)
 from sky.skylet.providers.aws.utils import (
     LazyDefaultDict,
     handle_boto_error,
     resource_cache,
 )
-from ray.autoscaler._private.cli_logger import cli_logger, cf
+from ray.autoscaler._private.cli_logger import cf, cli_logger
 from ray.autoscaler._private.event_system import CreateClusterEvent, global_event_system
-from sky.skylet.providers.aws.cloudwatch.cloudwatch_helper import (
-    CloudwatchHelper as cwh,
-)
+from ray.autoscaler._private.providers import _PROVIDER_PRETTY_NAMES
+from ray.autoscaler._private.util import check_legacy_fields
+from ray.autoscaler.tags import NODE_TYPE_LEGACY_HEAD, NODE_TYPE_LEGACY_WORKER
 
 logger = logging.getLogger(__name__)
 
@@ -33,20 +32,22 @@
 DEFAULT_RAY_IAM_ROLE = RAY + "-v1"
 SECURITY_GROUP_TEMPLATE = RAY + "-{}"
 
-DEFAULT_AMI_NAME = "AWS Deep Learning AMI (Ubuntu 18.04) V30.0"
+# V61.0 has CUDA 11.2
+DEFAULT_AMI_NAME = "AWS Deep Learning AMI (Ubuntu 18.04) V61.0"
 
-# Obtained from https://aws.amazon.com/marketplace/pp/B07Y43P7X5 on 8/4/2020.
+# Obtained from https://aws.amazon.com/marketplace/pp/B07Y43P7X5 on 6/10/2022.
+# NOTE(skypilot): these are not used; skypilot instead uses the default AMIs in aws.py.
 DEFAULT_AMI = {
-    "us-east-1": "ami-029510cec6d69f121",  # US East (N. Virginia)
-    "us-east-2": "ami-08bf49c7b3a0c761e",  # US East (Ohio)
-    "us-west-1": "ami-0cc472544ce594a19",  # US West (N. California)
-    "us-west-2": "ami-0a2363a9cff180a64",  # US West (Oregon)
-    "ca-central-1": "ami-0a871851b2ab39f01",  # Canada (Central)
-    "eu-central-1": "ami-049fb1ea198d189d7",  # EU (Frankfurt)
-    "eu-west-1": "ami-0abcbc65f89fb220e",  # EU (Ireland)
-    "eu-west-2": "ami-0755b39fd4dab7cbe",  # EU (London)
-    "eu-west-3": "ami-020485d8df1d45530",  # EU (Paris)
-    "sa-east-1": "ami-058a6883cbdb4e599",  # SA (Sao Paulo)
+    "us-east-1": "ami-0dd6adfad4ad37eec",  # US East (N. Virginia)
+    "us-east-2": "ami-0c77cd5ca05bf1281",  # US East (Ohio)
+    "us-west-1": "ami-020ab1b368a5ed1db",  # US West (N. California)
+    "us-west-2": "ami-0387d929287ab193e",  # US West (Oregon)
+    "ca-central-1": "ami-07dbafdbd38f18d98",  # Canada (Central)
+    "eu-central-1": "ami-0383bd0c1fc4c63ec",  # EU (Frankfurt)
+    "eu-west-1": "ami-0a074b0a311a837ac",  # EU (Ireland)
+    "eu-west-2": "ami-094ba2b4651f761ca",  # EU (London)
+    "eu-west-3": "ami-031da10fbf225bf5f",  # EU (Paris)
+    "sa-east-1": "ami-0be7c1f1dd96d7337",  # SA (Sao Paulo)
 }
 
 # todo: cli_logger should handle this assert properly
Original file line number	Diff line number	Diff line change
`@@ -4,7 +4,7 @@`
`4`	`4`	# The system administrator must have `sudo` access to the local nodes.
`5`	`5`	`# Requirements:`
`6`	`6`	`# 1) Python (> 3.6) on all nodes.`
`7`		`-# 2) Ray CLI (= 1.13.0) on all nodes.`
	`7`	`+# 2) Ray CLI (= 2.0.1) on all nodes.`
`8`	`8`	`#`
`9`	`9`	`# Example usage:`
`10`	`10`	`# >> sky admin deploy cluster-config.yaml`