Repository: skypilot-org/skypilot

Issue:
[UX] sky logs throws unnecessary errors when job id doesn't exist.
`sky logs` throws errors when the job id doesn't exist. We should catch the error instead and show `Job not found` or things like that.
```
(sky-tmp) weichiang@mbp sky % sky logs cpu 1
Tailing logs of job 1 on cluster 'cpu'...
Warning: Permanently added '35.202.33.57' (ED25519) to the list of known hosts.
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/opt/conda/lib/python3.7/posixpath.py", line 94, in join
    genericpath._check_arg_types('join', a, *p)
  File "/opt/conda/lib/python3.7/genericpath.py", line 153, in _check_arg_types
    (funcname, s.__class__.__name__)) from None
TypeError: join() argument must be str or bytes, not 'NoneType'
Shared connection to 35.202.33.57 closed.
```

Discussion:

Patch:
diff --git a/sky/cli.py b/sky/cli.py
index e26411c0..8d504a0c 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -2667,8 +2667,20 @@ def check():
               default=None,
               type=str,
               help='Cloud provider to query.')
+@click.option(
+    '--region',
+    required=False,
+    type=str,
+    help=
+    ('The region to use. If not specified, shows accelerators from all regions.'
+    ),
+)
 @usage_lib.entrypoint
-def show_gpus(gpu_name: Optional[str], all: bool, cloud: Optional[str]):  # pylint: disable=redefined-builtin
+def show_gpus(
+        gpu_name: Optional[str],
+        all: bool,  # pylint: disable=redefined-builtin
+        cloud: Optional[str],
+        region: Optional[str]):
     """Show supported GPU/TPU/accelerators.
 
     The names and counts shown can be set in the ``accelerators`` field in task
@@ -2682,9 +2694,14 @@ def show_gpus(gpu_name: Optional[str], all: bool, cloud: Optional[str]):  # pyli
     To show all accelerators, including less common ones and their detailed
     information, use ``sky show-gpus --all``.
 
-    NOTE: The price displayed for each instance type is the lowest across all
-    regions for both on-demand and spot instances.
+    NOTE: If region is not specified, the price displayed for each instance type
+    is the lowest across all regions for both on-demand and spot instances.
     """
+    # validation for the --region flag
+    if region is not None and cloud is None:
+        raise click.UsageError(
+            'The --region flag is only valid when the --cloud flag is set.')
+    service_catalog.validate_region_zone(region, None, clouds=cloud)
     show_all = all
     if show_all and gpu_name is not None:
         raise click.UsageError('--all is only allowed without a GPU name.')
@@ -2701,8 +2718,11 @@ def show_gpus(gpu_name: Optional[str], all: bool, cloud: Optional[str]):  # pyli
             ['OTHER_GPU', 'AVAILABLE_QUANTITIES'])
 
         if gpu_name is None:
-            result = service_catalog.list_accelerator_counts(gpus_only=True,
-                                                             clouds=cloud)
+            result = service_catalog.list_accelerator_counts(
+                gpus_only=True,
+                clouds=cloud,
+                region_filter=region,
+            )
             # NVIDIA GPUs
             for gpu in service_catalog.get_common_gpus():
                 if gpu in result:
@@ -2730,6 +2750,7 @@ def show_gpus(gpu_name: Optional[str], all: bool, cloud: Optional[str]):  # pyli
         # Show detailed accelerator information
         result = service_catalog.list_accelerators(gpus_only=True,
                                                    name_filter=gpu_name,
+                                                   region_filter=region,
                                                    clouds=cloud)
         if len(result) == 0:
             yield f'Resources \'{gpu_name}\' not found. '
@@ -2742,7 +2763,7 @@ def show_gpus(gpu_name: Optional[str], all: bool, cloud: Optional[str]):  # pyli
         yield 'the host VM\'s cost is not included.\n\n'
         import pandas as pd  # pylint: disable=import-outside-toplevel
         for i, (gpu, items) in enumerate(result.items()):
-            accelerator_table = log_utils.create_table([
+            accelerator_table_headers = [
                 'GPU',
                 'QTY',
                 'CLOUD',
@@ -2751,7 +2772,11 @@ def show_gpus(gpu_name: Optional[str], all: bool, cloud: Optional[str]):  # pyli
                 'HOST_MEMORY',
                 'HOURLY_PRICE',
                 'HOURLY_SPOT_PRICE',
-            ])
+            ]
+            if not show_all:
+                accelerator_table_headers.append('REGION')
+            accelerator_table = log_utils.create_table(
+                accelerator_table_headers)
             for item in items:
                 instance_type_str = item.instance_type if not pd.isna(
                     item.instance_type) else '(attachable)'
@@ -2769,11 +2794,20 @@ def show_gpus(gpu_name: Optional[str], all: bool, cloud: Optional[str]):  # pyli
                     item.price) else '-'
                 spot_price_str = f'$ {item.spot_price:.3f}' if not pd.isna(
                     item.spot_price) else '-'
-                accelerator_table.add_row([
-                    item.accelerator_name, item.accelerator_count, item.cloud,
-                    instance_type_str, cpu_str, mem_str, price_str,
-                    spot_price_str
-                ])
+                region_str = item.region if not pd.isna(item.region) else '-'
+                accelerator_table_vals = [
+                    item.accelerator_name,
+                    item.accelerator_count,
+                    item.cloud,
+                    instance_type_str,
+                    cpu_str,
+                    mem_str,
+                    price_str,
+                    spot_price_str,
+                ]
+                if not show_all:
+                    accelerator_table_vals.append(region_str)
+                accelerator_table.add_row(accelerator_table_vals)
 
             if i != 0:
                 yield '\n\n'
diff --git a/sky/clouds/service_catalog/__init__.py b/sky/clouds/service_catalog/__init__.py
index 2acd61d5..a722701c 100644
--- a/sky/clouds/service_catalog/__init__.py
+++ b/sky/clouds/service_catalog/__init__.py
@@ -49,6 +49,7 @@ def _map_clouds_catalog(clouds: CloudFilter, method_name: str, *args, **kwargs):
 def list_accelerators(
     gpus_only: bool = True,
     name_filter: Optional[str] = None,
+    region_filter: Optional[str] = None,
     clouds: CloudFilter = None,
     case_sensitive: bool = True,
 ) -> 'Dict[str, List[common.InstanceTypeInfo]]':
@@ -58,7 +59,7 @@ def list_accelerators(
     of instance type offerings. See usage in cli.py.
     """
     results = _map_clouds_catalog(clouds, 'list_accelerators', gpus_only,
-                                  name_filter, case_sensitive)
+                                  name_filter, region_filter, case_sensitive)
     if not isinstance(results, list):
         results = [results]
     ret: Dict[str,
@@ -72,6 +73,7 @@ def list_accelerators(
 def list_accelerator_counts(
     gpus_only: bool = True,
     name_filter: Optional[str] = None,
+    region_filter: Optional[str] = None,
     clouds: CloudFilter = None,
 ) -> Dict[str, List[int]]:
     """List all accelerators offered by Sky and available counts.
@@ -80,7 +82,7 @@ def list_accelerator_counts(
     of available counts. See usage in cli.py.
     """
     results = _map_clouds_catalog(clouds, 'list_accelerators', gpus_only,
-                                  name_filter)
+                                  name_filter, region_filter, False)
     if not isinstance(results, list):
         results = [results]
     accelerator_counts: Dict[str, Set[int]] = collections.defaultdict(set)
diff --git a/sky/clouds/service_catalog/aws_catalog.py b/sky/clouds/service_catalog/aws_catalog.py
index 9fccd495..150e45b3 100644
--- a/sky/clouds/service_catalog/aws_catalog.py
+++ b/sky/clouds/service_catalog/aws_catalog.py
@@ -71,7 +71,7 @@ def instance_type_exists(instance_type: str) -> bool:
 def validate_region_zone(
         region: Optional[str],
         zone: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
-    return common.validate_region_zone_impl(_df, region, zone)
+    return common.validate_region_zone_impl('aws', _df, region, zone)
 
 
 def accelerator_in_region_or_zone(acc_name: str,
@@ -134,13 +134,15 @@ def get_region_zones_for_instance_type(instance_type: str,
     return us_region_list + other_region_list
 
 
-def list_accelerators(gpus_only: bool,
-                      name_filter: Optional[str],
-                      case_sensitive: bool = True
-                     ) -> Dict[str, List[common.InstanceTypeInfo]]:
+def list_accelerators(
+        gpus_only: bool,
+        name_filter: Optional[str],
+        region_filter: Optional[str],
+        case_sensitive: bool = True
+) -> Dict[str, List[common.InstanceTypeInfo]]:
     """Returns all instance types in AWS offering accelerators."""
     return common.list_accelerators_impl('AWS', _df, gpus_only, name_filter,
-                                         case_sensitive)
+                                         region_filter, case_sensitive)
 
 
 def get_image_id_from_tag(tag: str, region: Optional[str]) -> Optional[str]:
diff --git a/sky/clouds/service_catalog/azure_catalog.py b/sky/clouds/service_catalog/azure_catalog.py
index fc74fc79..c07f8e8a 100644
--- a/sky/clouds/service_catalog/azure_catalog.py
+++ b/sky/clouds/service_catalog/azure_catalog.py
@@ -22,7 +22,7 @@ def validate_region_zone(
     if zone is not None:
         with ux_utils.print_exception_no_traceback():
             raise ValueError('Azure does not support zones.')
-    return common.validate_region_zone_impl(_df, region, zone)
+    return common.validate_region_zone_impl('azure', _df, region, zone)
 
 
 def accelerator_in_region_or_zone(acc_name: str,
@@ -89,10 +89,12 @@ def get_gen_version_from_instance_type(instance_type: str) -> Optional[int]:
     return _df[_df['InstanceType'] == instance_type]['Generation'].iloc[0]
 
 
-def list_accelerators(gpus_only: bool,
-                      name_filter: Optional[str],
-                      case_sensitive: bool = True
-                     ) -> Dict[str, List[common.InstanceTypeInfo]]:
+def list_accelerators(
+        gpus_only: bool,
+        name_filter: Optional[str],
+        region_filter: Optional[str],
+        case_sensitive: bool = True
+) -> Dict[str, List[common.InstanceTypeInfo]]:
     """Returns all instance types in Azure offering GPUs."""
     return common.list_accelerators_impl('Azure', _df, gpus_only, name_filter,
-                                         case_sensitive)
+                                         region_filter, case_sensitive)
diff --git a/sky/clouds/service_catalog/common.py b/sky/clouds/service_catalog/common.py
index d6466731..84b2c67b 100644
--- a/sky/clouds/service_catalog/common.py
+++ b/sky/clouds/service_catalog/common.py
@@ -34,6 +34,7 @@ class InstanceTypeInfo(NamedTuple):
     - memory: Instance memory in GiB.
     - price: Regular instance price per hour (cheapest across all regions).
     - spot_price: Spot instance price per hour (cheapest across all regions).
+    - region: Region where this instance type belongs to.
     """
     cloud: str
     instance_type: Optional[str]
@@ -43,6 +44,7 @@ class InstanceTypeInfo(NamedTuple):
     memory: Optional[float]
     price: float
     spot_price: float
+    region: str
 
 
 def get_catalog_path(filename: str) -> str:
@@ -161,7 +163,7 @@ def instance_type_exists_impl(df: pd.DataFrame, instance_type: str) -> bool:
 
 
 def validate_region_zone_impl(
-        df: pd.DataFrame, region: Optional[str],
+        cloud_name: str, df: pd.DataFrame, region: Optional[str],
         zone: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
     """Validates whether region and zone exist in the catalog."""
 
@@ -174,6 +176,11 @@ def validate_region_zone_impl(
             candidate_strs = f'\nDid you mean one of these: {candidate_strs!r}?'
         return candidate_strs
 
+    def _get_all_supported_regions_str() -> str:
+        all_regions: List[str] = sorted(df['Region'].unique().tolist())
+        return \
+        f'\nList of supported {cloud_name} regions: {", ".join(all_regions)!r}'
+
     validated_region, validated_zone = region, zone
 
     filter_df = df
@@ -182,7 +189,12 @@ def validate_region_zone_impl(
         if len(filter_df) == 0:
             with ux_utils.print_exception_no_traceback():
                 error_msg = (f'Invalid region {region!r}')
-                error_msg += _get_candidate_str(region, df['Region'].unique())
+                candidate_strs = _get_candidate_str(region,
+                                                    df['Region'].unique())
+                if not candidate_strs:
+                    error_msg += _get_all_supported_regions_str()
+                    raise ValueError(error_msg)
+                error_msg += candidate_strs
                 raise ValueError(error_msg)
 
     if zone is not None:
@@ -321,6 +333,7 @@ def list_accelerators_impl(
     df: pd.DataFrame,
     gpus_only: bool,
     name_filter: Optional[str],
+    region_filter: Optional[str],
     case_sensitive: bool = True,
 ) -> Dict[str, List[InstanceTypeInfo]]:
     """Lists accelerators offered in a cloud service catalog.
@@ -341,11 +354,16 @@ def list_accelerators_impl(
         'MemoryGiB',
         'Price',
         'SpotPrice',
+        'Region',
     ]].dropna(subset=['AcceleratorName']).drop_duplicates()
     if name_filter is not None:
         df = df[df['AcceleratorName'].str.contains(name_filter,
                                                    case=case_sensitive,
                                                    regex=True)]
+    if region_filter is not None:
+        df = df[df['Region'].str.contains(region_filter,
+                                          case=case_sensitive,
+                                          regex=True)]
     df['AcceleratorCount'] = df['AcceleratorCount'].astype(int)
     grouped = df.groupby('AcceleratorName')
 
@@ -366,6 +384,7 @@ def list_accelerators_impl(
                 row['MemoryGiB'],
                 row['Price'],
                 row['SpotPrice'],
+                row['Region'],
             ),
             axis='columns',
         ).tolist()
diff --git a/sky/clouds/service_catalog/gcp_catalog.py b/sky/clouds/service_catalog/gcp_catalog.py
index 39e3eecb..b2cb8f4b 100644
--- a/sky/clouds/service_catalog/gcp_catalog.py
+++ b/sky/clouds/service_catalog/gcp_catalog.py
@@ -210,7 +210,7 @@ def get_instance_type_for_accelerator(
 def validate_region_zone(
         region: Optional[str],
         zone: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
-    return common.validate_region_zone_impl(_df, region, zone)
+    return common.validate_region_zone_impl('gcp', _df, region, zone)
 
 
 def accelerator_in_region_or_zone(acc_name: str,
@@ -273,11 +273,12 @@ def get_accelerator_hourly_cost(accelerator: str,
 def list_accelerators(
     gpus_only: bool,
     name_filter: Optional[str] = None,
+    region_filter: Optional[str] = None,
     case_sensitive: bool = True,
 ) -> Dict[str, List[common.InstanceTypeInfo]]:
     """Returns all instance types in GCP offering GPUs."""
     results = common.list_accelerators_impl('GCP', _df, gpus_only, name_filter,
-                                            case_sensitive)
+                                            region_filter, case_sensitive)
 
     a100_infos = results.get('A100', []) + results.get('A100-80GB', [])
     if not a100_infos:
diff --git a/sky/skylet/job_lib.py b/sky/skylet/job_lib.py
index 62909440..2946bcc8 100644
--- a/sky/skylet/job_lib.py
+++ b/sky/skylet/job_lib.py
@@ -362,17 +362,14 @@ def _get_jobs_by_ids(job_ids: List[int]) -> List[Dict[str, Any]]:
 def update_job_status(job_owner: str,
                       job_ids: List[int],
                       silent: bool = False) -> List[JobStatus]:
-    """Updates and returns the job statuses matching our `JobStatus` semantics
+    """Updates and returns the job statuses matching our `JobStatus` semantics.
 
-    "True" statuses: this function queries `ray job status` and processes
-    those results to match our semantics.
+    This function queries `ray job status` and processes those results to match
+    our semantics.
 
-    This function queries `ray job status` and processes those results to
-    match our semantics.
-
-    Though we update job status actively in ray program and job cancelling,
-    we still need this to handle staleness problem, caused by instance
-    restarting and other corner cases (if any).
+    Though we update job status actively in the generated ray program and
+    during job cancelling, we still need this to handle the staleness problem,
+    caused by instance restarting and other corner cases (if any).
 
     This function should only be run on the remote instance with ray==2.0.1.
     """
@@ -413,7 +410,8 @@ def update_job_status(job_owner: str,
             if status is None:
                 original_status = get_status_no_lock(job_id)
                 status = original_status
-                if not original_status.is_terminal():
+                if (original_status is not None and
+                        not original_status.is_terminal()):
                     # The job may be stale, when the instance is restarted
                     # (the ray redis is volatile). We need to reset the
                     # status of the task to FAILED if its original status
@@ -665,14 +663,12 @@ class JobLibCodeGen:
                   job_id: Optional[int],
                   spot_job_id: Optional[int],
                   follow: bool = True) -> str:
+        # pylint: disable=line-too-long
         code = [
-            f'job_id = {job_id} if {job_id} is not None '
-            'else job_lib.get_latest_job_id()',
+            f'job_id = {job_id} if {job_id} is not None else job_lib.get_latest_job_id()',
             'run_timestamp = job_lib.get_run_timestamp(job_id)',
-            (f'log_dir = os.path.join({constants.SKY_LOGS_DIRECTORY!r}, '
-             'run_timestamp)'),
-            (f'log_lib.tail_logs({job_owner!r},'
-             f'job_id, log_dir, {spot_job_id!r}, follow={follow})'),
+            f'log_dir = None if run_timestamp is None else os.path.join({constants.SKY_LOGS_DIRECTORY!r}, run_timestamp)',
+            f'log_lib.tail_logs({job_owner!r}, job_id, log_dir, {spot_job_id!r}, follow={follow})',
         ]
         return cls._build(code)
 
diff --git a/sky/skylet/log_lib.py b/sky/skylet/log_lib.py
index 6ab16dcf..afc82209 100644
--- a/sky/skylet/log_lib.py
+++ b/sky/skylet/log_lib.py
@@ -383,14 +383,14 @@ def tail_logs(job_owner: str,
     job_str = f'job {job_id}'
     if spot_job_id is not None:
         job_str = f'spot job {spot_job_id}'
-    logger.debug(f'Tailing logs for job, real job_id {job_id}, spot_job_id '
-                 f'{spot_job_id}.')
-    logger.info(f'{colorama.Fore.YELLOW}Start streaming logs for {job_str}.'
-                f'{colorama.Style.RESET_ALL}')
     if log_dir is None:
         print(f'{job_str.capitalize()} not found (see `sky queue`).',
               file=sys.stderr)
         return
+    logger.debug(f'Tailing logs for job, real job_id {job_id}, spot_job_id '
+                 f'{spot_job_id}.')
+    logger.info(f'{colorama.Fore.YELLOW}Start streaming logs for {job_str}.'
+                f'{colorama.Style.RESET_ALL}')
     log_path = os.path.join(log_dir, 'run.log')
     log_path = os.path.expanduser(log_path)
 


Test patch:
diff --git a/tests/test_list_accelerators.py b/tests/test_list_accelerators.py
index 56bde6b4..c8e0579b 100644
--- a/tests/test_list_accelerators.py
+++ b/tests/test_list_accelerators.py
@@ -17,6 +17,17 @@ def test_list_ccelerators_all():
     assert 'A100-80GB' in result, result
 
 
-def test_list_accelerators_filters():
+def test_list_accelerators_name_filter():
     result = sky.list_accelerators(gpus_only=False, name_filter='V100')
     assert sorted(result.keys()) == ['V100', 'V100-32GB'], result
+
+
+def test_list_accelerators_region_filter():
+    result = sky.list_accelerators(gpus_only=False,
+                                   clouds="aws",
+                                   region_filter='us-west-1')
+    all_regions = []
+    for res in result.values():
+        for instance in res:
+            all_regions.append(instance.region)
+    assert all([region == 'us-west-1' for region in all_regions])