Refactor the file_mounts optimization and oslogin (skypilot-org#1108)

* Refactor the file_mounts optimization

* format

* format

* fix

* change optimization to be for all clouds

* Add ray yaml back

* address comments

Author: Michaelvll

sky/authentication.py

@@ -31,7 +31,9 @@
 PRIVATE_SSH_KEY_PATH = '~/.ssh/sky-key'
 
 GCP_CONFIGURE_PATH = '~/.config/gcloud/configurations/config_default'
-GCP_CONFIGURE_SKY_BACKUP_PATH = '~/.config/gcloud/configurations/.sky_config_default'  # pylint: disable=line-too-long
+# Do not place the backup under the gcloud config directory, as ray
+# autoscaler can overwrite that directory on the remote nodes.
+GCP_CONFIGURE_SKY_BACKUP_PATH = '~/.sky/.sky_gcp_config_default'
 
 
 def generate_rsa_key_pair():
@@ -202,17 +204,17 @@ def setup_gcp_authentication(config):
     project_oslogin = next(
         (item for item in project['commonInstanceMetadata'].get('items', [])
          if item['key'] == 'enable-oslogin'), {}).get('value', 'False')
+
     if project_oslogin.lower() == 'true':
         # project.
         logger.info(
             f'OS Login is enabled for GCP project {project_id}. Running '
             'additional authentication steps.')
+        # Read the account information from the credential file, since the user
+        # should be set according the account, when the oslogin is enabled.
         config_path = os.path.expanduser(GCP_CONFIGURE_PATH)
         sky_backup_config_path = os.path.expanduser(
             GCP_CONFIGURE_SKY_BACKUP_PATH)
-
-        # Read the account information from the credential file, since the user
-        # should be set according the account, when the oslogin is enabled.
         if not os.path.exists(sky_backup_config_path):
             if not os.path.exists(config_path):
                 with ux_utils.print_exception_no_traceback():
@@ -227,7 +229,10 @@ def setup_gcp_authentication(config):
             subprocess.run(f'cp {config_path} {sky_backup_config_path}',
                            shell=True,
                            check=True)
-
+        new_file_mounts = config.get('file_mounts', {})
+        new_file_mounts[
+            GCP_CONFIGURE_SKY_BACKUP_PATH] = GCP_CONFIGURE_SKY_BACKUP_PATH
+        config['file_mounts'] = new_file_mounts
         with open(sky_backup_config_path, 'r') as infile:
             for line in infile:
                 if line.startswith('account'):

sky/backends/backend_utils.py

@@ -31,7 +31,6 @@
 from ray.autoscaler._private import util as ray_util
 import rich.console as rich_console
 import rich.progress as rich_progress
-import yaml
 
 import sky
 from sky import authentication as auth
@@ -120,8 +119,7 @@ def is_ip(s: str) -> bool:
 def fill_template(template_name: str,
                   variables: Dict,
                   output_path: Optional[str] = None,
-                  output_prefix: str = SKY_USER_FILE_PATH,
-                  dryrun: bool = False) -> str:
+                  output_prefix: str = SKY_USER_FILE_PATH) -> str:
     """Create a file from a Jinja template and return the filename."""
     assert template_name.endswith('.j2'), template_name
     template_path = os.path.join(sky.__root_dir__, 'templates', template_name)
@@ -138,36 +136,44 @@ def fill_template(template_name: str,
         output_path = str(output_path)
     output_path = os.path.abspath(output_path)
 
-    # Runtime files handling
-    #
-    # List of runtime files to be uploaded to cluster:
-    #   - yaml config (for autostopping)
-    #   - wheel
-    #   - credentials
-    # Format is {dst: src}.
-    file_mounts = {SKY_RAY_YAML_REMOTE_PATH: output_path}
-
-    # fill_template() is also called to fill TPU/spot controller templates,
-    # which don't have all variables.
-    if 'sky_remote_path' in variables and 'sky_local_path' in variables:
-        file_mounts[variables['sky_remote_path']] = variables['sky_local_path']
-    if 'credentials' in variables:
-        file_mounts.update(variables['credentials'])
+    # Add yaml file path to the template variables.
+    variables['sky_ray_yaml_remote_path'] = SKY_RAY_YAML_REMOTE_PATH
+    variables['sky_ray_yaml_local_path'] = output_path
+    # Write out yaml config.
+    template = jinja2.Template(template)
+    content = template.render(**variables)
+    with open(output_path, 'w') as fout:
+        fout.write(content)
+    return output_path
+
+
+def _optimize_file_mounts(yaml_path: str) -> None:
+    """Optimize file mounts in the given ray yaml file.
+
+    Runtime files handling:
+    List of runtime files to be uploaded to cluster:
+      - yaml config (for autostopping)
+      - wheel
+      - credentials
+    Format is {dst: src}.
+    """
+    yaml_config = common_utils.read_yaml(yaml_path)
+
+    file_mounts = yaml_config.get('file_mounts', {})
+    # Remove the file mounts added by the newline.
+    if '' in file_mounts:
+        assert file_mounts[''] == '', file_mounts['']
+        file_mounts.pop('')
+
     # Putting these in file_mounts hurts provisioning speed, as each file
     # opens/closes an SSH connection.  Instead, we:
     #  - cp locally them into a directory
     #  - upload that directory as a file mount (1 connection)
     #  - use a remote command to move all runtime files to their right places.
 
-    # yaml config
-    variables['sky_ray_yaml_remote_path'] = SKY_RAY_YAML_REMOTE_PATH
-    variables['sky_ray_yaml_local_path'] = output_path
-
     # Local tmp dir holding runtime files.
     local_runtime_files_dir = tempfile.mkdtemp()
-    variables['local_runtime_files_dir'] = local_runtime_files_dir
-    # Remote dir.
-    variables['remote_runtime_files_dir'] = _REMOTE_RUNTIME_FILES_DIR
+    new_file_mounts = {_REMOTE_RUNTIME_FILES_DIR: local_runtime_files_dir}
 
     # (For remote) Build a command that copies runtime files to their right
     # destinations.
@@ -203,24 +209,28 @@ def fill_template(template_name: str,
               f'{dst_parent_dir}/{dst_basename}')
         fragment = f'({mkdir_parent} && {mv})'
         commands.append(fragment)
-    variables['postprocess_runtime_files_command'] = ' && '.join(commands)
+    postprocess_runtime_files_command = ' && '.join(commands)
 
-    # Write out yaml config.
-    template = jinja2.Template(template)
-    content = template.render(**variables)
-    with open(output_path, 'w') as fout:
-        fout.write(content)
+    setup_commands = yaml_config.get('setup_commands', [])
+    if setup_commands:
+        setup_commands[
+            0] = f'{postprocess_runtime_files_command}; {setup_commands[0]}'
+    else:
+        setup_commands = [postprocess_runtime_files_command]
+
+    yaml_config['file_mounts'] = new_file_mounts
+    yaml_config['setup_commands'] = setup_commands
 
     # (For local) Move all runtime files, including the just-written yaml, to
     # local_runtime_files_dir/.
-    if not dryrun:
-        all_local_sources = ' '.join(
-            local_src for local_src in file_mounts.values())
-        # Takes 10-20 ms on laptop incl. 3 clouds' credentials.
-        subprocess_utils.run(
-            f'cp -r {all_local_sources} {local_runtime_files_dir}/')
+    all_local_sources = ' '.join(
+        local_src for local_src in file_mounts.values())
+    # Takes 10-20 ms on laptop incl. 3 clouds' credentials.
+    subprocess.run(f'cp -r {all_local_sources} {local_runtime_files_dir}/',
+                   shell=True,
+                   check=True)
 
-    return output_path
+    common_utils.dump_yaml(yaml_path, yaml_config)
 
 
 def path_size_megabytes(path: str) -> int:
@@ -726,13 +736,18 @@ def write_cluster_config(to_provision: 'resources.Resources',
                 'ssh_private_key': (None if auth_config is None else
                                     auth_config['ssh_private_key']),
             }),
-        dryrun=dryrun,
     )
     config_dict['cluster_name'] = cluster_name
     config_dict['ray'] = yaml_path
     if dryrun:
         return config_dict
     _add_auth_to_cluster_config(cloud, yaml_path)
+    # Delay the optimization of the config until the authentication files is added.
+    if not isinstance(cloud, clouds.Local):
+        # Only optimize the file mounts for public clouds now, as local has not
+        # been fully tested yet.
+        _optimize_file_mounts(yaml_path)
+
     usage_lib.messages.usage.update_ray_yaml(yaml_path)
     # For TPU nodes. TPU VMs do not need TPU_NAME.
     if (resources_vars.get('tpu_type') is not None and
@@ -768,8 +783,7 @@ def _add_auth_to_cluster_config(cloud: clouds.Cloud, cluster_config_file: str):
 
     This function's output removes comments included in the jinja2 template.
     """
-    with open(cluster_config_file, 'r') as f:
-        config = yaml.safe_load(f)
+    config = common_utils.read_yaml(cluster_config_file)
     # Check the availability of the cloud type.
     if isinstance(cloud, clouds.AWS):
         config = auth.setup_aws_authentication(config)

sky/templates/gcp-ray.yml.j2

@@ -104,7 +104,11 @@ head_node_type: ray_head_default
 
 # Format: `REMOTE_PATH : LOCAL_PATH`
 file_mounts: {
-  "{{remote_runtime_files_dir}}": "{{local_runtime_files_dir}}",
+  "{{sky_ray_yaml_remote_path}}": "{{sky_ray_yaml_local_path}}",
+  "{{sky_remote_path}}": "{{sky_local_path}}",
+{%- for remote_path, local_path in credentials.items() %}
+  "{{remote_path}}": "{{local_path}}",
+{%- endfor %}
 }
 
 rsync_exclude: []
@@ -125,8 +129,7 @@ setup_commands:
   # default. 'source ~/.bashrc' is needed so conda takes effect for the next
   # commands.
   # Line 'python3 -c ..': patch the buggy ray files and enable `-o allow_other` option for `goofys`
-  - {{postprocess_runtime_files_command}};
-    mkdir -p ~/.ssh; touch ~/.ssh/config;
+  - mkdir -p ~/.ssh; touch ~/.ssh/config;
     pip3 --version > /dev/null 2>&1 || (curl -sSL https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py && echo "PATH=$HOME/.local/bin:$PATH" >> ~/.bashrc); (type -a python | grep -q python3) || echo 'alias python=python3' >> ~/.bashrc; (type -a pip | grep -q pip3) || echo 'alias pip=pip3' >> ~/.bashrc;
     which conda > /dev/null 2>&1 || (wget -nc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.sh -b && eval "$(/home/gcpuser/miniconda3/bin/conda shell.bash hook)" && conda init && conda config --set auto_activate_base true);
     source ~/.bashrc;

tests/test_smoke.py

@@ -82,6 +82,8 @@ def run_one_test(test: Test) -> Tuple[int, str, str]:
             log_file.flush()
             test.echo(f'Timeout after {test.timeout} seconds.')
             test.echo(e)
+            log_file.write(f'Timeout after {test.timeout} seconds.\n')
+            log_file.flush()
             # Kill the current process.
             proc.terminate()
             proc.returncode = 1  # None if we don't set it.