fix chmod

skypilot-org · Aug 7, 2024 · 6a2a644 · 6a2a644
1 parent 9de328f
commit 6a2a644
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 13 deletions.
diff --git a/sky/api/cli.py b/sky/api/cli.py
@@ -168,9 +168,9 @@ def _parse_env_var(env_var: str) -> Tuple[str, str]:
 
 
 def _async_call_or_wait(request_id: str, async_call: bool,
-                        request_name: str) -> None:
+                        request_name: str) -> Any:
     if not async_call:
-        sdk.stream_and_get(request_id)
+        return sdk.stream_and_get(request_id)
     else:
         click.secho(f'Submitted {request_name} request: {request_id}',
                     fg='green')
@@ -1091,10 +1091,11 @@ def launch(
         clone_disk_from=clone_disk_from,
         need_confirmation=not yes,
     )
-    _async_call_or_wait(request_id, async_call, 'Launch')
+    _, handle = _async_call_or_wait(request_id, async_call, 'Launch')
     if async_call:
         # Add ssh config for the cluster
-        _get_cluster_records_and_set_ssh_config(clusters=[cluster])
+        _get_cluster_records_and_set_ssh_config(
+            clusters=[handle.get_cluster_name()])
 
 
 @cli.command(cls=_DocumentedCodeCommand)

diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
@@ -53,7 +53,6 @@
 from sky.skylet import log_lib
 from sky.usage import usage_lib
 from sky.utils import accelerator_registry
-from sky.utils import cluster_utils
 from sky.utils import command_runner
 from sky.utils import common
 from sky.utils import common_utils
@@ -2844,8 +2843,7 @@ def _provision(
 
                 self._update_after_cluster_provisioned(
                     handle, to_provision_config.prev_handle, task,
-                    prev_cluster_status, handle.external_ips(),
-                    handle.external_ssh_ports(), lock_path)
+                    prev_cluster_status, lock_path)
                 return handle
 
             cluster_config_file = config_dict['ray']
@@ -2917,7 +2915,7 @@ def _get_zone(runner):
 
             self._update_after_cluster_provisioned(
                 handle, to_provision_config.prev_handle, task,
-                prev_cluster_status, ip_list, ssh_port_list, lock_path)
+                prev_cluster_status, lock_path)
             return handle
 
     def _open_ports(self, handle: CloudVmRayResourceHandle) -> None:
@@ -2935,7 +2933,6 @@ def _update_after_cluster_provisioned(
             prev_handle: Optional[CloudVmRayResourceHandle],
             task: task_lib.Task,
             prev_cluster_status: Optional[status_lib.ClusterStatus],
-            ip_list: List[str], ssh_port_list: List[int],
             lock_path: str) -> None:
         usage_lib.messages.usage.update_cluster_resources(
             handle.launched_nodes, handle.launched_resources)

diff --git a/sky/utils/cluster_utils.py b/sky/utils/cluster_utils.py
@@ -123,10 +123,13 @@ def add_cluster(
                 cluster_private_key_path)
             expanded_cluster_private_key_dir = os.path.dirname(
                 expanded_cluster_private_key_path)
-            os.makedirs(expanded_cluster_private_key_dir, exist_ok=True)
-            os.chmod(expanded_cluster_private_key_dir, 0o700)
-            with open(expanded_cluster_private_key_path, 'w',
-                      encoding='utf-8') as f:
+            os.makedirs(expanded_cluster_private_key_dir,
+                        exist_ok=True,
+                        mode=0o700)
+            with open(expanded_cluster_private_key_path,
+                      'w',
+                      encoding='utf-8',
+                      opener=functools.partial(os.open, mode=0o600)) as f:
                 f.write(key_content)
             auth_config['ssh_private_key'] = cluster_private_key_path
         key_path = os.path.expanduser(auth_config['ssh_private_key'])