skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/skypilot_config.py

@@ -1,7 +1,7 @@
 """Immutable user configurations (EXPERIMENTAL).
 
-On module import, we attempt to parse the config located at CONFIG_PATH. Caller
-can then use
+On module import, we attempt to parse the config located at CONFIG_PATH
+(default: ~/.sky/config.yaml). Caller can then use
 
   >> skypilot_config.loaded()
 
@@ -11,6 +11,13 @@ To read a nested-key config:
 
   >> skypilot_config.get_nested(('auth', 'some_auth_config'), default_value)
 
+The config can be overridden by the configs in task YAMLs. Callers are
+responsible to provide the override_configs. If the nested key is part of
+OVERRIDEABLE_CONFIG_KEYS, override_configs must be provided (can be empty):
+
+  >> skypilot_config.get_nested(('docker', 'run_options'), default_value
+                        override_configs={'docker': {'run_options': 'value'}})
+
 To set a value in the nested-key config:
 
   >> config_dict = skypilot_config.set_nested(('auth', 'some_key'), value)
@@ -41,18 +48,25 @@ then:
     skypilot_config.get_nested(('a', 'nonexist'), None)  # ==> None
     skypilot_config.get_nested(('a',), None)             # ==> None
 """
+import contextlib
 import copy
 import os
 import pprint
-from typing import Any, Dict, Iterable, Optional
+import tempfile
+from typing import Any, Dict, Iterator, Optional, Tuple
 
 import yaml
 
+from sky import exceptions
 from sky import sky_logging
+from sky.skylet import constants
 from sky.utils import common_utils
+from sky.utils import config_utils
 from sky.utils import schemas
 from sky.utils import ux_utils
 
+logger = sky_logging.init_logger(__name__)
+
 # The config path is discovered in this order:
 #
 # (1) (Used internally) If env var {ENV_VAR_SKYPILOT_CONFIG} exists, use its
@@ -65,68 +79,65 @@ from sky.utils import ux_utils
 # (Used internally) An env var holding the path to the local config file. This
 # is only used by jobs controller tasks to ensure recoveries of the same job
 # use the same config file.
-ENV_VAR_SKYPILOT_CONFIG = 'SKYPILOT_CONFIG'
+ENV_VAR_SKYPILOT_CONFIG = f'{constants.SKYPILOT_ENV_VAR_PREFIX}CONFIG'
 
 # Path to the local config file.
 CONFIG_PATH = '~/.sky/config.yaml'
 
-logger = sky_logging.init_logger(__name__)
-
 # The loaded config.
-_dict = None
-_loaded_config_path = None
+_dict = config_utils.Config()
+_loaded_config_path: Optional[str] = None
 
 
-def get_nested(keys: Iterable[str], default_value: Any) -> Any:
+def get_nested(keys: Tuple[str, ...],
+               default_value: Any,
+               override_configs: Optional[Dict[str, Any]] = None) -> Any:
     """Gets a nested key.
 
     If any key is not found, or any intermediate key does not point to a dict
     value, returns 'default_value'.
+
+    When 'keys' is within OVERRIDEABLE_CONFIG_KEYS, 'override_configs' must be
+    provided (can be empty). Otherwise, 'override_configs' must not be provided.
+
+    Args:
+        keys: A tuple of strings representing the nested keys.
+        default_value: The default value to return if the key is not found.
+        override_configs: A dict of override configs with the same schema as
+            the config file, but only containing the keys to override.
+
+    Returns:
+        The value of the nested key, or 'default_value' if not found.
     """
-    if _dict is None:
-        return default_value
-    curr = _dict
-    for key in keys:
-        if isinstance(curr, dict) and key in curr:
-            curr = curr[key]
-        else:
-            return default_value
-    logger.debug(f'User config: {".".join(keys)} -> {curr}')
-    return curr
+    return _dict.get_nested(
+        keys,
+        default_value,
+        override_configs,
+        allowed_override_keys=constants.OVERRIDEABLE_CONFIG_KEYS_IN_TASK,
+        disallowed_override_keys=None)
 
 
-def set_nested(keys: Iterable[str], value: Any) -> Dict[str, Any]:
+def set_nested(keys: Tuple[str, ...], value: Any) -> Dict[str, Any]:
     """Returns a deep-copied config with the nested key set to value.
 
     Like get_nested(), if any key is not found, this will not raise an error.
     """
-    _check_loaded_or_die()
-    assert _dict is not None
-    curr = copy.deepcopy(_dict)
-    to_return = curr
-    prev = None
-    for i, key in enumerate(keys):
-        if key not in curr:
-            curr[key] = {}
-        prev = curr
-        curr = curr[key]
-        if i == len(keys) - 1:
-            prev_value = prev[key]
-            prev[key] = value
-            logger.debug(f'Set the value of {keys} to {value} (previous: '
-                         f'{prev_value}). Returning conf: {to_return}')
-    return to_return
-
-
-def to_dict() -> Dict[str, Any]:
+    copied_dict = copy.deepcopy(_dict)
+    copied_dict.set_nested(keys, value)
+    return dict(**copied_dict)
+
+
+def to_dict() -> config_utils.Config:
     """Returns a deep-copied version of the current config."""
-    if _dict is not None:
-        return copy.deepcopy(_dict)
-    return {}
+    return copy.deepcopy(_dict)
 
 
-def _try_load_config() -> None:
+def _reload_config() -> None:
     global _dict, _loaded_config_path
+    # Reset the global variables, to avoid using stale values.
+    _dict = config_utils.Config()
+    _loaded_config_path = None
+
     config_path_via_env_var = os.environ.get(ENV_VAR_SKYPILOT_CONFIG)
     if config_path_via_env_var is not None:
         config_path = os.path.expanduser(config_path_via_env_var)
@@ -142,18 +153,19 @@ def _try_load_config() -> None:
     config_path = os.path.expanduser(config_path)
     if os.path.exists(config_path):
         logger.debug(f'Using config path: {config_path}')
-        _loaded_config_path = config_path
         try:
-            _dict = common_utils.read_yaml(config_path)
+            config = common_utils.read_yaml(config_path)
+            _dict = config_utils.Config.from_dict(config)
+            _loaded_config_path = config_path
             logger.debug(f'Config loaded:\n{pprint.pformat(_dict)}')
         except yaml.YAMLError as e:
             logger.error(f'Error in loading config file ({config_path}):', e)
-        if _dict is not None:
+        if _dict:
             common_utils.validate_schema(
                 _dict,
                 schemas.get_config_schema(),
                 f'Invalid config YAML ({config_path}). See: '
-                'https://skypilot.readthedocs.io/en/latest/reference/config.html. '  # pylint: disable=line-too-long
+                'https://docs.skypilot.co/en/latest/reference/config.html. '  # pylint: disable=line-too-long
                 'Error: ',
                 skip_none=False)
 
@@ -166,17 +178,67 @@ def loaded_config_path() -> Optional[str]:
 
 
 # Load on import.
-_try_load_config()
-
-
-def _check_loaded_or_die():
-    """Checks loaded() is true; otherwise raises RuntimeError."""
-    if _dict is None:
-        raise RuntimeError(
-            f'No user configs loaded. Check {CONFIG_PATH} exists and '
-            'can be loaded.')
+_reload_config()
 
 
 def loaded() -> bool:
     """Returns if the user configurations are loaded."""
-    return _dict is not None
+    return bool(_dict)
+
+
+@contextlib.contextmanager
+def override_skypilot_config(
+        override_configs: Optional[Dict[str, Any]]) -> Iterator[None]:
+    """Overrides the user configurations."""
+    # TODO(SKY-1215): allow admin user to extend the disallowed keys or specify
+    # allowed keys.
+    if not override_configs:
+        # If no override configs (None or empty dict), do nothing.
+        yield
+        return
+    original_env_config_path = _loaded_config_path
+    original_config = dict(_dict)
+    config = _dict.get_nested(
+        keys=tuple(),
+        default_value=None,
+        override_configs=override_configs,
+        allowed_override_keys=None,
+        disallowed_override_keys=constants.SKIPPED_CLIENT_OVERRIDE_KEYS)
+    with tempfile.NamedTemporaryFile(
+            mode='w',
+            prefix='skypilot_config',
+            # Have to avoid deleting the file as the underlying function needs
+            # to read the config file, and we need to close the file mode='w'
+            # to enable reading.
+            delete=False) as f:
+        common_utils.dump_yaml(f.name, dict(config))
+        os.environ[ENV_VAR_SKYPILOT_CONFIG] = f.name
+    try:
+        _reload_config()
+        yield
+    except exceptions.InvalidSkyPilotConfigError as e:
+        with ux_utils.print_exception_no_traceback():
+            raise exceptions.InvalidSkyPilotConfigError(
+                'Failed to override the SkyPilot config on API '
+                'server with your local SkyPilot config:\n'
+                '=== SkyPilot config on API server ===\n'
+                f'{common_utils.dump_yaml_str(original_config)}\n'
+                '=== Your local SkyPilot config ===\n'
+                f'{common_utils.dump_yaml_str(override_configs)}\n'
+                f'Details: {e}') from e
+
+    finally:
+        if original_env_config_path is not None:
+            os.environ[ENV_VAR_SKYPILOT_CONFIG] = original_env_config_path
+        else:
+            os.environ.pop(ENV_VAR_SKYPILOT_CONFIG, None)
+        # Reload the config to restore the original config to avoid the next
+        # request reusing the same process to use the config for the current
+        # request.
+        _reload_config()
+
+        try:
+            os.remove(f.name)
+        except Exception:  # pylint: disable=broad-except
+            # Failing to delete the file is not critical.
+            pass

sky/task.py

@@ -1,4 +1,5 @@
 """Task: a coarse-grained stage in an application."""
+import collections
 import inspect
 import json
 import os
@@ -121,6 +122,9 @@ def _check_docker_login_config(task_envs: Dict[str, str]) -> bool:
 
     If any of the docker login env vars is set, all of them must be set.
 
+    Returns:
+        True if there is a valid docker login config in task_envs.
+        False otherwise.
     Raises:
         ValueError: if any of the docker login env vars is set, but not all of
             them are set.
@@ -168,6 +172,23 @@ def _with_docker_login_config(
     return type(resources)(new_resources)
 
 
+def _with_docker_username_for_runpod(
+    resources: Union[Set['resources_lib.Resources'],
+                     List['resources_lib.Resources']],
+    task_envs: Dict[str, str],
+) -> Union[Set['resources_lib.Resources'], List['resources_lib.Resources']]:
+    docker_username_for_runpod = task_envs.get(
+        constants.RUNPOD_DOCKER_USERNAME_ENV_VAR)
+
+    # We should not call r.copy() if docker_username_for_runpod is None,
+    # to prevent `DummyResources` instance becoming a `Resources` instance.
+    if docker_username_for_runpod is None:
+        return resources
+    return (type(resources)(
+        r.copy(_docker_username_for_runpod=docker_username_for_runpod)
+        for r in resources))
+
+
 class Task:
     """Task: a computation to be run on the cloud."""
 
@@ -184,6 +205,8 @@ class Task:
         docker_image: Optional[str] = None,
         event_callback: Optional[str] = None,
         blocked_resources: Optional[Iterable['resources_lib.Resources']] = None,
+        # Internal use only.
+        file_mounts_mapping: Optional[Dict[str, str]] = None,
     ):
         """Initializes a Task.
 
@@ -280,21 +303,33 @@ class Task:
 
         # Filled in by the optimizer.  If None, this Task is not planned.
         self.best_resources = None
-        # Check if the task is legal.
-        self._validate()
+
+        # For internal use only.
+        self.file_mounts_mapping = file_mounts_mapping
 
         dag = sky.dag.get_current_dag()
         if dag is not None:
             dag.add(self)
 
-    def _validate(self):
-        """Checks if the Task fields are valid."""
+    def validate(self, workdir_only: bool = False):
+        """Validate all fields of the task."""
+        self.validate_name()
+        self.validate_run()
+        self.expand_and_validate_workdir()
+        if not workdir_only:
+            self.expand_and_validate_file_mounts()
+        for r in self.resources:
+            r.validate()
+
+    def validate_name(self):
+        """Validates if the task name is valid."""
         if not _is_valid_name(self.name):
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(f'Invalid task name {self.name}. Valid name: '
                                  f'{_VALID_NAME_DESCR}')
 
-        # Check self.run
+    def validate_run(self):
+        """Validates if the run command is valid."""
         if callable(self.run):
             run_sig = inspect.signature(self.run)
             # Check that run is a function with 2 arguments.
@@ -333,15 +368,65 @@ class Task:
                                  f'a command generator ({CommandGen}). '
                                  f'Got {type(self.run)}')
 
-        # Workdir.
-        if self.workdir is not None:
-            full_workdir = os.path.abspath(os.path.expanduser(self.workdir))
-            if not os.path.isdir(full_workdir):
-                # Symlink to a dir is legal (isdir() follows symlinks).
+    def expand_and_validate_file_mounts(self):
+        """Expand file_mounts paths to absolute paths and validate them.
+
+        Note: if this function is called on a remote SkyPilot API server,
+        it must be after the client side has sync-ed all files to the
+        remote server.
+        """
+        if self.file_mounts is None:
+            return
+        for target, source in self.file_mounts.items():
+            if target.endswith('/') or source.endswith('/'):
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        'File mount paths cannot end with a slash '
+                        '(try "/mydir: /mydir" or "/myfile: /myfile"). '
+                        f'Found: target={target} source={source}')
+            if data_utils.is_cloud_store_url(target):
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        'File mount destination paths cannot be cloud storage')
+            if not data_utils.is_cloud_store_url(source):
+                self.file_mounts[target] = os.path.abspath(
+                    os.path.expanduser(source))
+                if not os.path.exists(self.file_mounts[target]
+                                     ) and not source.startswith('skypilot:'):
+                    with ux_utils.print_exception_no_traceback():
+                        raise ValueError(
+                            f'File mount source {source!r} does not exist '
+                            'locally. To fix: check if it exists, and correct '
+                            'the path.')
+            # TODO(zhwu): /home/username/sky_workdir as the target path need
+            # to be filtered out as well.
+            if (target == constants.SKY_REMOTE_WORKDIR and
+                    self.workdir is not None):
                 with ux_utils.print_exception_no_traceback():
                     raise ValueError(
-                        'Workdir must exist and must be a directory (or '
-                        f'a symlink to a directory). {self.workdir} not found.')
+                        f'Cannot use {constants.SKY_REMOTE_WORKDIR!r} as a '
+                        'destination path of a file mount, as it will be used '
+                        'by the workdir. If uploading a file/folder to the '
+                        'workdir is needed, please specify the full path to '
+                        'the file/folder.')
+
+    def expand_and_validate_workdir(self):
+        """Expand workdir to absolute path and validate it.
+
+        Note: if this function is called on a remote SkyPilot API server,
+        it must be after the client side has sync-ed all files to the
+        remote server.
+        """
+        if self.workdir is None:
+            return
+        user_workdir = self.workdir
+        self.workdir = os.path.abspath(os.path.expanduser(user_workdir))
+        if not os.path.isdir(self.workdir):
+            # Symlink to a dir is legal (isdir() follows symlinks).
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    'Workdir must be a valid directory (or '
+                    f'a symlink to a directory). {user_workdir} not found.')
 
     @staticmethod
     def from_yaml_config(
@@ -393,6 +478,11 @@ class Task:
             config['service'] = _fill_in_env_vars(config['service'],
                                                   config.get('envs', {}))
 
+        # Fill in any Task.envs into workdir
+        if config.get('workdir') is not None:
+            config['workdir'] = _fill_in_env_vars(config['workdir'],
+                                                  config.get('envs', {}))
+
         task = Task(
             config.pop('name', None),
             run=config.pop('run', None),
@@ -401,6 +491,7 @@ class Task:
             num_nodes=config.pop('num_nodes', None),
             envs=config.pop('envs', None),
             event_callback=config.pop('event_callback', None),
+            file_mounts_mapping=config.pop('file_mounts_mapping', None),
         )
 
         # Create lists to store storage objects inlined in file_mounts.
@@ -456,8 +547,25 @@ class Task:
             task.set_outputs(outputs=outputs,
                              estimated_size_gigabytes=estimated_size_gigabytes)
 
+        # Experimental configs.
+        experimnetal_configs = config.pop('experimental', None)
+        cluster_config_override = None
+        if experimnetal_configs is not None:
+            cluster_config_override = experimnetal_configs.pop(
+                'config_overrides', None)
+            logger.debug('Overriding skypilot config with task-level config: '
+                         f'{cluster_config_override}')
+        assert not experimnetal_configs, ('Invalid task args: '
+                                          f'{experimnetal_configs.keys()}')
+
         # Parse resources field.
-        resources_config = config.pop('resources', None)
+        resources_config = config.pop('resources', {})
+        if cluster_config_override is not None:
+            assert resources_config.get('_cluster_config_overrides') is None, (
+                'Cannot set _cluster_config_overrides in both resources and '
+                'experimental.config_overrides')
+            resources_config[
+                '_cluster_config_overrides'] = cluster_config_override
         task.set_resources(sky.Resources.from_yaml_config(resources_config))
 
         service = config.pop('service', None)
@@ -560,6 +668,8 @@ class Task:
         if _check_docker_login_config(self._envs):
             self.resources = _with_docker_login_config(self.resources,
                                                        self._envs)
+        self.resources = _with_docker_username_for_runpod(
+            self.resources, self._envs)
         return self
 
     @property
@@ -625,6 +735,9 @@ class Task:
             resources = {resources}
         # TODO(woosuk): Check if the resources are None.
         self.resources = _with_docker_login_config(resources, self.envs)
+        # Only have effect on RunPod.
+        self.resources = _with_docker_username_for_runpod(
+            self.resources, self.envs)
 
         # Evaluate if the task requires FUSE and set the requires_fuse flag
         for _, storage_obj in self.storage_mounts.items():
@@ -710,45 +823,7 @@ class Task:
 
         Returns:
           self: the current task, with file mounts set.
-
-        Raises:
-          ValueError: if input paths are invalid.
         """
-        if file_mounts is None:
-            self.file_mounts = None
-            return self
-        for target, source in file_mounts.items():
-            if target.endswith('/') or source.endswith('/'):
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        'File mount paths cannot end with a slash '
-                        '(try "/mydir: /mydir" or "/myfile: /myfile"). '
-                        f'Found: target={target} source={source}')
-            if data_utils.is_cloud_store_url(target):
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        'File mount destination paths cannot be cloud storage')
-            if not data_utils.is_cloud_store_url(source):
-                if (not os.path.exists(
-                        os.path.abspath(os.path.expanduser(source))) and
-                        not source.startswith('skypilot:')):
-                    with ux_utils.print_exception_no_traceback():
-                        raise ValueError(
-                            f'File mount source {source!r} does not exist '
-                            'locally. To fix: check if it exists, and correct '
-                            'the path.')
-            # TODO(zhwu): /home/username/sky_workdir as the target path need
-            # to be filtered out as well.
-            if (target == constants.SKY_REMOTE_WORKDIR and
-                    self.workdir is not None):
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        f'Cannot use {constants.SKY_REMOTE_WORKDIR!r} as a '
-                        'destination path of a file mount, as it will be used '
-                        'by the workdir. If uploading a file/folder to the '
-                        'workdir is needed, please specify the full path to '
-                        'the file/folder.')
-
         self.file_mounts = file_mounts
         return self
 
@@ -784,8 +859,8 @@ class Task:
             self.file_mounts = {}
         assert self.file_mounts is not None
         self.file_mounts.update(file_mounts)
-        # For validation logic:
-        return self.set_file_mounts(self.file_mounts)
+        self.expand_and_validate_file_mounts()
+        return self
 
     def set_storage_mounts(
         self,
@@ -933,18 +1008,46 @@ class Task:
         file_mounts of the form ``{ /remote/path: {s3,gs,..}://<bucket path>
         }``.
         """
+        # The same storage can be used multiple times, and we should construct
+        # the storage with stores first, so that the storage will be created on
+        # the correct cloud.
+        name_to_storage = collections.defaultdict(list)
         for storage in self.storage_mounts.values():
-            if len(storage.stores) == 0:
-                store_type, store_region = self._get_preferred_store()
-                self.storage_plans[storage] = store_type
-                storage.add_store(store_type, store_region)
-            else:
-                # We will download the first store that is added to remote.
-                self.storage_plans[storage] = list(storage.stores.keys())[0]
-
+            name_to_storage[storage.name].append(storage)
+        for storages in name_to_storage.values():
+            # Place the storage with most stores first, so that the storage will
+            # be created on the correct cloud.
+            storage_to_construct = sorted(storages,
+                                          key=lambda x: len(x.stores),
+                                          reverse=True)
+            for storage in storage_to_construct:
+                storage.construct()
+                assert storage.name is not None, storage
+                if not storage.stores:
+                    store_type, store_region = self._get_preferred_store()
+                    self.storage_plans[storage] = store_type
+                    storage.add_store(store_type, store_region)
+                else:
+                    # We don't need to sync the storage here as if the stores
+                    # are not empty, it measn the storage has been synced during
+                    # construct() above.
+                    # We will download the first store that is added to remote.
+                    assert all(store is not None
+                               for store in storage.stores.values()), storage
+                    self.storage_plans[storage] = list(storage.stores.keys())[0]
+
+        # The following logic converts the storage mounts with COPY mode into
+        # inline file mounts with cloud URIs, so that the _execute_file_mounts()
+        # in cloud_vm_ray_backend.py can correctly download from the specific
+        # cloud storage on the remote cluster.
+        # Note that this will cause duplicate destination paths in file_mounts,
+        # and storage_mounts, which should be fine as our to_yaml_config() will
+        # only dump the storage mount version, i.e. what user specified.
         storage_mounts = self.storage_mounts
         storage_plans = self.storage_plans
         for mnt_path, storage in storage_mounts.items():
+            assert storage.name is not None, storage
+
             if storage.mode == storage_lib.StorageMode.COPY:
                 store_type = storage_plans[storage]
                 if store_type is storage_lib.StoreType.S3:
@@ -955,6 +1058,7 @@ class Task:
                     else:
                         assert storage.name is not None, storage
                         blob_path = 's3://' + storage.name
+                    blob_path = storage.get_bucket_sub_path_prefix(blob_path)
                     self.update_file_mounts({
                         mnt_path: blob_path,
                     })
@@ -965,6 +1069,26 @@ class Task:
                     else:
                         assert storage.name is not None, storage
                         blob_path = 'gs://' + storage.name
+                    blob_path = storage.get_bucket_sub_path_prefix(blob_path)
+                    self.update_file_mounts({
+                        mnt_path: blob_path,
+                    })
+                elif store_type is storage_lib.StoreType.AZURE:
+                    if (isinstance(storage.source, str) and
+                            data_utils.is_az_container_endpoint(
+                                storage.source)):
+                        blob_path = storage.source
+                    else:
+                        assert storage.name is not None, storage
+                        store_object = storage.stores[
+                            storage_lib.StoreType.AZURE]
+                        assert isinstance(store_object,
+                                          storage_lib.AzureBlobStore)
+                        storage_account_name = store_object.storage_account_name
+                        blob_path = data_utils.AZURE_CONTAINER_URL.format(
+                            storage_account_name=storage_account_name,
+                            container_name=storage.name)
+                    blob_path = storage.get_bucket_sub_path_prefix(blob_path)
                     self.update_file_mounts({
                         mnt_path: blob_path,
                     })
@@ -975,6 +1099,7 @@ class Task:
                         blob_path = storage.source
                     else:
                         blob_path = 'r2://' + storage.name
+                    blob_path = storage.get_bucket_sub_path_prefix(blob_path)
                     self.update_file_mounts({
                         mnt_path: blob_path,
                     })
@@ -990,15 +1115,29 @@ class Task:
                         cos_region = data_utils.Rclone.get_region_from_rclone(
                             storage.name, data_utils.Rclone.RcloneClouds.IBM)
                         blob_path = f'cos://{cos_region}/{storage.name}'
+                    blob_path = storage.get_bucket_sub_path_prefix(blob_path)
                     self.update_file_mounts({mnt_path: blob_path})
-                elif store_type is storage_lib.StoreType.AZURE:
-                    # TODO when Azure Blob is done: sync ~/.azure
-                    raise NotImplementedError('Azure Blob not mountable yet')
+                elif store_type is storage_lib.StoreType.OCI:
+                    if storage.source is not None and not isinstance(
+                            storage.source,
+                            list) and storage.source.startswith('oci://'):
+                        blob_path = storage.source
+                    else:
+                        blob_path = 'oci://' + storage.name
+                    self.update_file_mounts({
+                        mnt_path: blob_path,
+                    })
                 else:
                     with ux_utils.print_exception_no_traceback():
                         raise ValueError(f'Storage Type {store_type} '
                                          'does not exist!')
 
+                # TODO: Delete from storage_mounts, now that the storage is
+                # translated into file_mounts. Note: as is, this will break
+                # controller_utils.
+                # _maybe_translate_local_file_mounts_and_sync_up(), which still
+                # needs the storage, but not the file_mounts.
+
     def get_local_to_remote_file_mounts(self) -> Optional[Dict[str, str]]:
         """Returns file mounts of the form (dst=VM path, src=local path).
 
@@ -1093,6 +1232,8 @@ class Task:
                 mount_path: storage.to_yaml_config()
                 for mount_path, storage in self.storage_mounts.items()
             })
+
+        add_if_not_none('file_mounts_mapping', self.file_mounts_mapping)
         return config
 
     def get_required_cloud_features(