skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/task.py

@@ -1,43 +1,33 @@
 """Task: a coarse-grained stage in an application."""
 import collections
-import inspect
 import json
 import os
 import re
-import typing
 from typing import (Any, Callable, Dict, Iterable, List, Optional, Set, Tuple,
                     Union)
 
 import colorama
 
-import sky
 from sky import clouds
+from sky import dag as dag_lib
 from sky import exceptions
+from sky import resources as resources_lib
 from sky import sky_logging
-from sky.adaptors import common as adaptors_common
-import sky.dag
 from sky.data import data_utils
 from sky.data import storage as storage_lib
 from sky.provision import docker_utils
 from sky.serve import service_spec
 from sky.skylet import constants
 from sky.utils import common_utils
+from sky.utils import git
+from sky.utils import registry
 from sky.utils import schemas
 from sky.utils import ux_utils
-
-if typing.TYPE_CHECKING:
-    import yaml
-
-    from sky import resources as resources_lib
-else:
-    yaml = adaptors_common.LazyImport('yaml')
+from sky.utils import volume as volume_lib
+from sky.utils import yaml_utils
 
 logger = sky_logging.init_logger(__name__)
 
-# A lambda generating commands (node rank_i, node addrs -> cmd_i).
-CommandGen = Callable[[int, List[str]], Optional[str]]
-CommandOrCommandGen = Union[str, CommandGen]
-
 _VALID_NAME_REGEX = '[a-zA-Z0-9]+(?:[._-]{1,2}[a-zA-Z0-9]+)*'
 _VALID_NAME_DESCR = ('ASCII characters and may contain lowercase and'
                      ' uppercase letters, digits, underscores, periods,'
@@ -121,27 +111,61 @@ def _fill_in_env_vars(
     return json.loads(yaml_field_str)
 
 
-def _check_docker_login_config(task_envs: Dict[str, str]) -> bool:
-    """Checks if there is a valid docker login config in task_envs.
+def _check_docker_login_config(task_envs: Dict[str, str],
+                               task_secrets: Dict[str, str]) -> bool:
+    """Validates a valid docker login config in task_envs and task_secrets.
 
-    If any of the docker login env vars is set, all of them must be set.
+    Docker login variables must be specified together either in envs OR secrets,
+    not split across both. If any of the docker login env vars is set, all of
+    them must be set in the same location.
+
+    Args:
+        task_envs: Environment variables
+        task_secrets: Secret variables (optional, defaults to empty dict)
 
     Returns:
-        True if there is a valid docker login config in task_envs.
+        True if there is a valid docker login config.
         False otherwise.
     Raises:
-        ValueError: if any of the docker login env vars is set, but not all of
-            them are set.
+        ValueError: if docker login configuration is invalid.
     """
+    if task_secrets is None:
+        task_secrets = {}
+
     all_keys = constants.DOCKER_LOGIN_ENV_VARS
-    existing_keys = all_keys & set(task_envs.keys())
-    if not existing_keys:
+    envs_keys = all_keys & set(task_envs.keys())
+    secrets_keys = all_keys & set(task_secrets.keys())
+
+    # Check if any docker variables exist
+    if not envs_keys and not secrets_keys:
         return False
-    if len(existing_keys) != len(all_keys):
+
+    # Check if variables are split across envs and secrets
+    if envs_keys and secrets_keys:
         with ux_utils.print_exception_no_traceback():
             raise ValueError(
-                f'If any of {", ".join(all_keys)} is set, all of them must '
-                f'be set. Missing envs: {all_keys - existing_keys}')
+                'Docker login variables must be specified together either '
+                'in envs OR secrets, not split across both. '
+                f'Found in envs: {sorted(envs_keys)}, '
+                f'Found in secrets: {sorted(secrets_keys)}')
+
+    # Check if all variables are present in the chosen location
+    if envs_keys:
+        if len(envs_keys) != len(all_keys):
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    'Docker login variables must be specified together '
+                    'in envs. '
+                    f'Missing from envs: {sorted(all_keys - envs_keys)}')
+
+    if secrets_keys:
+        if len(secrets_keys) != len(all_keys):
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    'Docker login variables must be specified together '
+                    'in secrets. '
+                    f'Missing from secrets: {sorted(all_keys - secrets_keys)}')
+
     return True
 
 
@@ -149,11 +173,13 @@ def _with_docker_login_config(
     resources: Union[Set['resources_lib.Resources'],
                      List['resources_lib.Resources']],
     task_envs: Dict[str, str],
+    task_secrets: Dict[str, str],
 ) -> Union[Set['resources_lib.Resources'], List['resources_lib.Resources']]:
-    if not _check_docker_login_config(task_envs):
+    if not _check_docker_login_config(task_envs, task_secrets):
         return resources
-    docker_login_config = docker_utils.DockerLoginConfig.from_env_vars(
-        task_envs)
+    envs = task_envs.copy()
+    envs.update(task_secrets)
+    docker_login_config = docker_utils.DockerLoginConfig.from_env_vars(envs)
 
     def _add_docker_login_config(resources: 'resources_lib.Resources'):
         docker_image = resources.extract_docker_image()
@@ -165,7 +191,8 @@ def _with_docker_login_config(
                            f'ignored.{colorama.Style.RESET_ALL}')
             return resources
         # Already checked in extract_docker_image
-        assert len(resources.image_id) == 1, resources.image_id
+        assert resources.image_id is not None and len(
+            resources.image_id) == 1, resources.image_id
         region = list(resources.image_id.keys())[0]
         return resources.copy(image_id={region: 'docker:' + docker_image},
                               _docker_login_config=docker_login_config)
@@ -180,8 +207,11 @@ def _with_docker_username_for_runpod(
     resources: Union[Set['resources_lib.Resources'],
                      List['resources_lib.Resources']],
     task_envs: Dict[str, str],
+    task_secrets: Dict[str, str],
 ) -> Union[Set['resources_lib.Resources'], List['resources_lib.Resources']]:
-    docker_username_for_runpod = task_envs.get(
+    envs = task_envs.copy()
+    envs.update(task_secrets)
+    docker_username_for_runpod = envs.get(
         constants.RUNPOD_DOCKER_USERNAME_ENV_VAR)
 
     # We should not call r.copy() if docker_username_for_runpod is None,
@@ -200,17 +230,27 @@ class Task:
         self,
         name: Optional[str] = None,
         *,
-        setup: Optional[str] = None,
-        run: Optional[CommandOrCommandGen] = None,
+        setup: Optional[Union[str, List[str]]] = None,
+        run: Optional[Union[str, List[str]]] = None,
         envs: Optional[Dict[str, str]] = None,
-        workdir: Optional[str] = None,
+        secrets: Optional[Dict[str, str]] = None,
+        workdir: Optional[Union[str, Dict[str, Any]]] = None,
         num_nodes: Optional[int] = None,
+        file_mounts: Optional[Dict[str, str]] = None,
+        storage_mounts: Optional[Dict[str, storage_lib.Storage]] = None,
+        volumes: Optional[Dict[str, str]] = None,
+        resources: Optional[Union['resources_lib.Resources',
+                                  List['resources_lib.Resources'],
+                                  Set['resources_lib.Resources']]] = None,
         # Advanced:
         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,
+        _file_mounts_mapping: Optional[Dict[str, str]] = None,
+        _volume_mounts: Optional[List[volume_lib.VolumeMount]] = None,
+        _metadata: Optional[Dict[str, Any]] = None,
+        _user_specified_yaml: Optional[str] = None,
     ):
         """Initializes a Task.
 
@@ -242,38 +282,81 @@ class Task:
 
         Args:
           name: A string name for the Task for display purposes.
-          setup: A setup command, which will be run before executing the run
+          setup: A setup command(s), which will be run before executing the run
             commands ``run``, and executed under ``workdir``.
           run: The actual command for the task. If not None, either a shell
-            command (str) or a command generator (callable).  If latter, it
-            must take a node rank and a list of node addresses as input and
-            return a shell command (str) (valid to return None for some nodes,
-            in which case no commands are run on them).  Run commands will be
-            run under ``workdir``. Note the command generator should be a
-            self-contained lambda.
+            command(s) (str, list(str)) or a command generator (callable). If
+            latter, it must take a node rank and a list of node addresses as
+            input and return a shell command (str) (valid to return None for
+            some nodes, in which case no commands are run on them). Run
+            commands will be run under ``workdir``. Note the command generator
+            should be a self-contained lambda.
           envs: A dictionary of environment variables to set before running the
             setup and run commands.
-          workdir: The local working directory.  This directory will be synced
+          secrets: A dictionary of secret environment variables to set before
+            running the setup and run commands. These will be redacted in logs
+            and YAML output.
+          workdir: The local working directory or a git repository.
+            For a local working directory, this directory will be synced
             to a location on the remote VM(s), and ``setup`` and ``run``
             commands will be run under that location (thus, they can rely on
             relative paths when invoking binaries).
+            If a git repository is provided, the repository will be cloned to
+            the working directory and the ``setup`` and ``run`` commands will
+            be run under the cloned repository.
           num_nodes: The number of nodes to provision for this Task.  If None,
             treated as 1 node.  If > 1, each node will execute its own
             setup/run command, where ``run`` can either be a str, meaning all
             nodes get the same command, or a lambda, with the semantics
             documented above.
+          file_mounts: An optional dict of ``{remote_path: (local_path|cloud
+            URI)}``, where remote means the VM(s) on which this Task will
+            eventually run on, and local means the node from which the task is
+            launched.
+          storage_mounts: an optional dict of ``{mount_path: sky.Storage
+            object}``, where mount_path is the path inside the remote VM(s)
+            where the Storage object will be mounted on.
+          volumes: A dict of volumes to be mounted for the task. The dict has
+            the form of ``{mount_path: volume_name}``.
+          resources: either a sky.Resources, a set of them, or a list of them.
+            A set or a list of resources asks the optimizer to "pick the
+            best of these resources" to run this task.
           docker_image: (EXPERIMENTAL: Only in effect when LocalDockerBackend
             is used.) The base docker image that this Task will be built on.
             Defaults to 'gpuci/miniforge-cuda:11.4-devel-ubuntu18.04'.
+          event_callback: A bash script that will be executed when the task
+            changes state.
           blocked_resources: A set of resources that this task cannot run on.
+          _file_mounts_mapping: (Internal use only) A dictionary of file mounts
+            mapping.
+          _volume_mounts: (Internal use only) A list of volume mounts.
+          _metadata: (Internal use only) A dictionary of metadata to be added to
+            the task.
+          _user_specified_yaml: (Internal use only) A string of user-specified
+            YAML config.
         """
         self.name = name
-        self.run = run
         self.storage_mounts: Dict[str, storage_lib.Storage] = {}
         self.storage_plans: Dict[storage_lib.Storage,
                                  storage_lib.StoreType] = {}
-        self.setup = setup
         self._envs = envs or {}
+        self._secrets = secrets or {}
+        self._volumes = volumes or {}
+
+        # concatenate commands if given as list
+        def _concat(commands: Optional[Union[str, List[str]]]) -> Optional[str]:
+            if isinstance(commands, list):
+                return '\n'.join(commands)
+            return commands
+
+        self.run = _concat(run)
+        self.setup = _concat(setup)
+
+        # Validate Docker login configuration early if both envs and secrets
+        # contain Docker variables
+        if self._envs or self._secrets:
+            _check_docker_login_config(self._envs, self._secrets)
+
         self.workdir = workdir
         self.docker_image = (docker_image if docker_image else
                              'gpuci/miniforge-cuda:11.4-devel-ubuntu18.04')
@@ -288,33 +371,50 @@ class Task:
         self.estimated_inputs_size_gigabytes: Optional[float] = None
         self.estimated_outputs_size_gigabytes: Optional[float] = None
         # Default to CPU VM
-        self.resources: Union[List[sky.Resources],
-                              Set[sky.Resources]] = {sky.Resources()}
+        self.resources: Union[List['resources_lib.Resources'],
+                              Set['resources_lib.Resources']] = {
+                                  resources_lib.Resources()
+                              }
         self._service: Optional[service_spec.SkyServiceSpec] = None
+
         # Resources that this task cannot run on.
         self.blocked_resources = blocked_resources
 
-        self.time_estimator_func: Optional[Callable[['sky.Resources'],
+        self.time_estimator_func: Optional[Callable[['resources_lib.Resources'],
                                                     int]] = None
         self.file_mounts: Optional[Dict[str, str]] = None
 
         # Only set when 'self' is a jobs controller task: 'self.managed_job_dag'
         # is the underlying managed job dag (sky.Dag object).
-        self.managed_job_dag: Optional['sky.Dag'] = None
+        self.managed_job_dag: Optional['dag_lib.Dag'] = None
 
         # Only set when 'self' is a sky serve controller task.
         self.service_name: Optional[str] = None
 
         # Filled in by the optimizer.  If None, this Task is not planned.
-        self.best_resources: Optional[sky.Resources] = None
+        self.best_resources: Optional['resources_lib.Resources'] = None
 
         # For internal use only.
-        self.file_mounts_mapping = file_mounts_mapping
+        self.file_mounts_mapping: Optional[Dict[str,
+                                                str]] = _file_mounts_mapping
+        self.volume_mounts: Optional[List[volume_lib.VolumeMount]] = (
+            _volume_mounts)
+
+        self._metadata = _metadata if _metadata is not None else {}
 
-        dag = sky.dag.get_current_dag()
+        if resources is not None:
+            self.set_resources(resources)
+        if storage_mounts is not None:
+            self.set_storage_mounts(storage_mounts)
+        if file_mounts is not None:
+            self.set_file_mounts(file_mounts)
+
+        dag = dag_lib.get_current_dag()
         if dag is not None:
             dag.add(self)
 
+        self._user_specified_yaml = _user_specified_yaml
+
     def validate(self,
                  skip_file_mounts: bool = False,
                  skip_workdir: bool = False):
@@ -342,42 +442,9 @@ class Task:
 
     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.
-            if len(run_sig.parameters) != 2:
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(_RUN_FN_CHECK_FAIL_MSG.format(run_sig))
-
-            type_list = [int, List[str]]
-            # Check annotations, if exists
-            for i, param in enumerate(run_sig.parameters.values()):
-                if param.annotation != inspect.Parameter.empty:
-                    if param.annotation != type_list[i]:
-                        with ux_utils.print_exception_no_traceback():
-                            raise ValueError(
-                                _RUN_FN_CHECK_FAIL_MSG.format(run_sig))
-
-            # Check self containedness.
-            run_closure = inspect.getclosurevars(self.run)
-            if run_closure.nonlocals:
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        'run command generator must be self contained. '
-                        f'Found nonlocals: {run_closure.nonlocals}')
-            if run_closure.globals:
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        'run command generator must be self contained. '
-                        f'Found globals: {run_closure.globals}')
-            if run_closure.unbound:
-                # Do not raise an error here. Import statements, which are
-                # allowed, will be considered as unbounded.
-                pass
-        elif self.run is not None and not isinstance(self.run, str):
+        if self.run is not None and not isinstance(self.run, str):
             with ux_utils.print_exception_no_traceback():
-                raise ValueError('run must be either a shell script (str) or '
-                                 f'a command generator ({CommandGen}). '
+                raise ValueError('run must be a shell script (str). '
                                  f'Got {type(self.run)}')
 
     def expand_and_validate_file_mounts(self):
@@ -390,12 +457,9 @@ class Task:
         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}')
+            location = f'file_mounts.{target}: {source}'
+            self._validate_mount_path(target, location)
+            self._validate_path(source, location)
             if data_utils.is_cloud_store_url(target):
                 with ux_utils.print_exception_no_traceback():
                     raise ValueError(
@@ -410,17 +474,25 @@ class Task:
                             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.')
+
+    def _validate_mount_path(self, path: str, location: str):
+        self._validate_path(path, location)
+        # TODO(zhwu): /home/username/sky_workdir as the target path need
+        # to be filtered out as well.
+        if (path == 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.')
+
+    def _validate_path(self, path: str, location: str):
+        if path.endswith('/'):
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError('Mount paths cannot end with a slash '
+                                 f'Found: {path} in {location}')
 
     def expand_and_validate_workdir(self):
         """Expand workdir to absolute path and validate it.
@@ -431,6 +503,12 @@ class Task:
         """
         if self.workdir is None:
             return
+        # Only expand the workdir if it is a string
+        if isinstance(self.workdir, dict):
+            git_ref = self.workdir.get('ref')
+            if git_ref is not None:
+                self._metadata['git_commit'] = git_ref
+            return
         user_workdir = self.workdir
         self.workdir = os.path.abspath(os.path.expanduser(user_workdir))
         if not os.path.isdir(self.workdir):
@@ -440,11 +518,16 @@ class Task:
                     'Workdir must be a valid directory (or '
                     f'a symlink to a directory). {user_workdir} not found.')
 
+        self._metadata['git_commit'] = common_utils.get_git_commit(self.workdir)
+
     @staticmethod
     def from_yaml_config(
         config: Dict[str, Any],
         env_overrides: Optional[List[Tuple[str, str]]] = None,
+        secrets_overrides: Optional[List[Tuple[str, str]]] = None,
     ) -> 'Task':
+        user_specified_yaml = config.pop('_user_specified_yaml',
+                                         yaml_utils.dump_yaml_str(config))
         # More robust handling for 'envs': explicitly convert keys and values to
         # str, since users may pass '123' as keys/values which will get parsed
         # as int causing validate_schema() to fail.
@@ -457,6 +540,20 @@ class Task:
                 else:
                     new_envs[str(k)] = None
             config['envs'] = new_envs
+
+        # More robust handling for 'secrets': explicitly convert keys and values
+        # to str, since users may pass '123' as keys/values which will get
+        # parsed as int causing validate_schema() to fail.
+        secrets = config.get('secrets')
+        if secrets is not None and isinstance(secrets, dict):
+            new_secrets: Dict[str, Optional[str]] = {}
+            for k, v in secrets.items():
+                if v is not None:
+                    new_secrets[str(k)] = str(v)
+                else:
+                    new_secrets[str(k)] = None
+            config['secrets'] = new_secrets
+
         common_utils.validate_schema(config, schemas.get_task_schema(),
                                      'Invalid task YAML: ')
         if env_overrides is not None:
@@ -470,6 +567,12 @@ class Task:
             new_envs.update(env_overrides)
             config['envs'] = new_envs
 
+        if secrets_overrides is not None:
+            # Override secrets vars from CLI.
+            new_secrets = config.get('secrets', {})
+            new_secrets.update(secrets_overrides)
+            config['secrets'] = new_secrets
+
         for k, v in config.get('envs', {}).items():
             if v is None:
                 with ux_utils.print_exception_no_traceback():
@@ -479,21 +582,38 @@ class Task:
                         f'To set it to be empty, use an empty string ({k}: "" '
                         f'in task YAML or --env {k}="" in CLI).')
 
+        for k, v in config.get('secrets', {}).items():
+            if v is None:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        f'Secret variable {k!r} is None. Please set a '
+                        'value for it in task YAML or with --secret flag. '
+                        f'To set it to be empty, use an empty string ({k}: "" '
+                        f'in task YAML or --secret {k}="" in CLI).')
+
         # Fill in any Task.envs into file_mounts (src/dst paths, storage
         # name/source).
+        env_vars = config.get('envs', {})
+        secrets = config.get('secrets', {})
+        env_and_secrets = env_vars.copy()
+        env_and_secrets.update(secrets)
         if config.get('file_mounts') is not None:
             config['file_mounts'] = _fill_in_env_vars(config['file_mounts'],
-                                                      config.get('envs', {}))
+                                                      env_and_secrets)
 
         # Fill in any Task.envs into service (e.g. MODEL_NAME).
         if config.get('service') is not None:
             config['service'] = _fill_in_env_vars(config['service'],
-                                                  config.get('envs', {}))
+                                                  env_and_secrets)
 
         # 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', {}))
+                                                  env_and_secrets)
+
+        if config.get('volumes') is not None:
+            config['volumes'] = _fill_in_env_vars(config['volumes'],
+                                                  env_and_secrets)
 
         task = Task(
             config.pop('name', None),
@@ -502,8 +622,12 @@ class Task:
             setup=config.pop('setup', None),
             num_nodes=config.pop('num_nodes', None),
             envs=config.pop('envs', None),
+            secrets=config.pop('secrets', None),
+            volumes=config.pop('volumes', None),
             event_callback=config.pop('event_callback', None),
-            file_mounts_mapping=config.pop('file_mounts_mapping', None),
+            _file_mounts_mapping=config.pop('file_mounts_mapping', None),
+            _metadata=config.pop('_metadata', None),
+            _user_specified_yaml=user_specified_yaml,
         )
 
         # Create lists to store storage objects inlined in file_mounts.
@@ -511,6 +635,7 @@ class Task:
         # storage objects with the storage/storage_mount objects.
         fm_storages = []
         file_mounts = config.pop('file_mounts', None)
+        volumes = []
         if file_mounts is not None:
             copy_mounts = {}
             for dst_path, src in file_mounts.items():
@@ -520,7 +645,27 @@ class Task:
                 # If the src is not a str path, it is likely a dict. Try to
                 # parse storage object.
                 elif isinstance(src, dict):
-                    fm_storages.append((dst_path, src))
+                    if (src.get('store') ==
+                            storage_lib.StoreType.VOLUME.value.lower()):
+                        # Build the volumes config for resources.
+                        volume_config = {
+                            'path': dst_path,
+                        }
+                        if src.get('name'):
+                            volume_config['name'] = src.get('name')
+                        persistent = src.get('persistent', False)
+                        volume_config['auto_delete'] = not persistent
+                        volume_config_detail = src.get('config', {})
+                        volume_config.update(volume_config_detail)
+                        volumes.append(volume_config)
+                        source_path = src.get('source')
+                        if source_path:
+                            # For volume, copy the source path to the
+                            # data directory of the volume mount point.
+                            copy_mounts[
+                                f'{dst_path.rstrip("/")}/data'] = source_path
+                    else:
+                        fm_storages.append((dst_path, src))
                 else:
                     with ux_utils.print_exception_no_traceback():
                         raise ValueError(f'Unable to parse file_mount '
@@ -598,12 +743,35 @@ class Task:
                 'experimental.config_overrides')
             resources_config[
                 '_cluster_config_overrides'] = cluster_config_override
-        task.set_resources(sky.Resources.from_yaml_config(resources_config))
+        if volumes:
+            resources_config['volumes'] = volumes
+        task.set_resources(
+            resources_lib.Resources.from_yaml_config(resources_config))
 
         service = config.pop('service', None)
+        pool = config.pop('pool', None)
+        if service is not None and pool is not None:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    'Cannot set both service and pool in the same task.')
+
         if service is not None:
             service = service_spec.SkyServiceSpec.from_yaml_config(service)
-        task.set_service(service)
+            task.set_service(service)
+        elif pool is not None:
+            pool['pool'] = True
+            pool = service_spec.SkyServiceSpec.from_yaml_config(pool)
+            task.set_service(pool)
+
+        volume_mounts = config.pop('volume_mounts', None)
+        if volume_mounts is not None:
+            task.volume_mounts = []
+            for vol in volume_mounts:
+                common_utils.validate_schema(vol,
+                                             schemas.get_volume_mount_schema(),
+                                             'Invalid volume mount config: ')
+                volume_mount = volume_lib.VolumeMount.from_yaml_config(vol)
+                task.volume_mounts.append(volume_mount)
 
         assert not config, f'Invalid task args: {config.keys()}'
         return task
@@ -628,7 +796,8 @@ class Task:
             # TODO(zongheng): use
             #  https://github.com/yaml/pyyaml/issues/165#issuecomment-430074049
             # to raise errors on duplicate keys.
-            config = yaml.safe_load(f)
+            user_specified_yaml = f.read()
+            config = yaml_utils.safe_load(user_specified_yaml)
 
         if isinstance(config, str):
             with ux_utils.print_exception_no_traceback():
@@ -637,8 +806,101 @@ class Task:
 
         if config is None:
             config = {}
+        config['_user_specified_yaml'] = user_specified_yaml
         return Task.from_yaml_config(config)
 
+    def resolve_and_validate_volumes(self) -> None:
+        """Resolve volumes config to volume mounts and validate them.
+
+        Raises:
+            exceptions.VolumeNotFoundError: if any volume is not found.
+            exceptions.VolumeTopologyConflictError: if there is conflict in the
+              volumes and compute topology.
+        """
+        # Volumes has been resolved, a typical case is that the API server
+        # has resolved the volumes and the dag was then submitted to
+        # controllers.
+        if self.volume_mounts is not None:
+            return None
+        if not self._volumes:
+            return None
+        volume_mounts: List[volume_lib.VolumeMount] = []
+        for dst_path, vol in self._volumes.items():
+            self._validate_mount_path(dst_path, location='volumes')
+            # Shortcut for `dst_path: volume_name`
+            if isinstance(vol, str):
+                volume_mount = volume_lib.VolumeMount.resolve(dst_path, vol)
+            elif isinstance(vol, dict):
+                assert 'name' in vol, 'Volume name must be set.'
+                volume_mount = volume_lib.VolumeMount.resolve(
+                    dst_path, vol['name'])
+            else:
+                raise ValueError(f'Invalid volume config: {dst_path}: {vol}')
+            volume_mounts.append(volume_mount)
+        # Disable certain access modes
+        disabled_modes = {}
+        if self.num_nodes > 1:
+            disabled_modes[
+                volume_lib.VolumeAccessMode.READ_WRITE_ONCE.value] = (
+                    'access mode ReadWriteOnce is not supported for '
+                    'multi-node tasks.')
+            disabled_modes[
+                volume_lib.VolumeAccessMode.READ_WRITE_ONCE_POD.value] = (
+                    'access mode ReadWriteOncePod is not supported for '
+                    'multi-node tasks.')
+        # TODO(aylei): generalize access mode to all volume types
+        # Record the required topology and the volume that requires it, e.g.
+        # {'cloud': ('volume_name', 'aws')}
+        topology: Dict[str, Tuple[str, Optional[str]]] = {
+            'cloud': ('', None),
+            'region': ('', None),
+            'zone': ('', None),
+        }
+        for vol in volume_mounts:
+            # Check access mode
+            access_mode = vol.volume_config.config.get('access_mode', '')
+            if access_mode in disabled_modes:
+                raise ValueError(f'Volume {vol.volume_name} with '
+                                 f'{disabled_modes[access_mode]}')
+            # Check topology
+            for key, (vol_name, previous_req) in topology.items():
+                req = getattr(vol.volume_config, key)
+                if req is not None:
+                    if previous_req is not None and req != previous_req:
+                        raise exceptions.VolumeTopologyConflictError(
+                            f'Volume {vol.volume_name} can only be attached on '
+                            f'{key}:{req}, which conflicts with another volume '
+                            f'{vol_name} that requires {key}:{previous_req}.'
+                            f'Please use different volumes and retry.')
+                    topology[key] = (vol_name, req)
+        # Now we have the topology requirements from the intersection of all
+        # volumes. Check if there is topology conflict with the resources.
+        # Volume must have no conflict with ALL resources even if user
+        # specifies 'any_of' resources to ensure no resources will conflict
+        # with the volumes during failover.
+
+        for res in self.resources:
+            for key, (vol_name, vol_req) in topology.items():
+                req = getattr(res, key)
+                if (req is not None and vol_req is not None and
+                        str(req) != vol_req):
+                    raise exceptions.VolumeTopologyConflictError(
+                        f'The task requires {key}:{req}, which conflicts with '
+                        f'the volume constraint {key}:{vol_req}. Please '
+                        f'use different volumes and retry.')
+        # No topology conflict, we safely override the topology of resources to
+        # satisfy the volume constraints.
+        override_params = {}
+        for key, (vol_name, vol_req) in topology.items():
+            if vol_req is not None:
+                if key == 'cloud':
+                    override_params[key] = registry.CLOUD_REGISTRY.from_str(
+                        vol_req)
+                else:
+                    override_params[key] = vol_req
+        self.set_resources_override(override_params)
+        self.volume_mounts = volume_mounts
+
     @property
     def num_nodes(self) -> int:
         return self._num_nodes
@@ -653,10 +915,38 @@ class Task:
                     f'num_nodes should be a positive int. Got: {num_nodes}')
         self._num_nodes = num_nodes
 
+    @property
+    def metadata(self) -> Dict[str, Any]:
+        return self._metadata
+
+    @property
+    def metadata_json(self) -> str:
+        return json.dumps(self._metadata)
+
     @property
     def envs(self) -> Dict[str, str]:
         return self._envs
 
+    @property
+    def secrets(self) -> Dict[str, str]:
+        return self._secrets
+
+    @property
+    def volumes(self) -> Dict[str, str]:
+        return self._volumes
+
+    def set_volumes(self, volumes: Dict[str, str]) -> None:
+        """Sets the volumes for this task.
+
+        Args:
+          volumes: a dict of ``{mount_path: volume_name}``.
+        """
+        self._volumes = volumes
+
+    def update_volumes(self, volumes: Dict[str, str]) -> None:
+        """Updates the volumes for this task."""
+        self._volumes.update(volumes)
+
     def update_envs(
             self, envs: Union[None, List[Tuple[str, str]],
                               Dict[str, str]]) -> 'Task':
@@ -697,17 +987,70 @@ class Task:
         # If the update_envs() is called after set_resources(), we need to
         # manually update docker login config in task resources, in case the
         # docker login envs are newly added.
-        if _check_docker_login_config(self._envs):
+        if _check_docker_login_config(self._envs, self._secrets):
+            self.resources = _with_docker_login_config(self.resources,
+                                                       self._envs,
+                                                       self._secrets)
+        self.resources = _with_docker_username_for_runpod(
+            self.resources, self._envs, self._secrets)
+        return self
+
+    def update_secrets(
+            self, secrets: Union[None, List[Tuple[str, str]],
+                                 Dict[str, str]]) -> 'Task':
+        """Updates secret env vars for use inside the setup/run commands.
+
+        Args:
+          secrets: (optional) either a list of ``(secret_name, value)`` or a
+            dict ``{secret_name: value}``.
+
+        Returns:
+          self: The current task, with secrets updated.
+
+        Raises:
+          ValueError: if various invalid inputs errors are detected.
+        """
+        if secrets is None:
+            secrets = {}
+        if isinstance(secrets, (list, tuple)):
+            keys = set(secret[0] for secret in secrets)
+            if len(keys) != len(secrets):
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError('Duplicate secret keys provided.')
+            secrets = dict(secrets)
+        if isinstance(secrets, dict):
+            for key in secrets:
+                if not isinstance(key, str):
+                    with ux_utils.print_exception_no_traceback():
+                        raise ValueError('Secret keys must be strings.')
+                if not common_utils.is_valid_env_var(key):
+                    with ux_utils.print_exception_no_traceback():
+                        raise ValueError(f'Invalid secret key: {key}')
+        else:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    'secrets must be List[Tuple[str, str]] or Dict[str, str]: '
+                    f'{secrets}')
+        self._secrets.update(secrets)
+        # Validate Docker login configuration if needed
+        if _check_docker_login_config(self._envs, self._secrets):
             self.resources = _with_docker_login_config(self.resources,
-                                                       self._envs)
+                                                       self._envs,
+                                                       self._secrets)
         self.resources = _with_docker_username_for_runpod(
-            self.resources, self._envs)
+            self.resources, self._envs, self._secrets)
         return self
 
     @property
     def use_spot(self) -> bool:
         return any(r.use_spot for r in self.resources)
 
+    @property
+    def envs_and_secrets(self) -> Dict[str, str]:
+        envs = self.envs.copy()
+        envs.update(self.secrets)
+        return envs
+
     def set_inputs(self, inputs: str,
                    estimated_size_gigabytes: float) -> 'Task':
         # E.g., 's3://bucket', 'gs://bucket', or None.
@@ -749,7 +1092,7 @@ class Task:
     def set_resources(
         self, resources: Union['resources_lib.Resources',
                                List['resources_lib.Resources'],
-                               Set['resources_lib.Resources']]
+                               Set['resources_lib.Resources'], Dict[str, Any]]
     ) -> 'Task':
         """Sets the required resources to execute this task.
 
@@ -763,19 +1106,22 @@ class Task:
         Returns:
           self: The current task, with resources set.
         """
-        if isinstance(resources, sky.Resources):
+        if isinstance(resources, dict):
+            resources = resources_lib.Resources.from_yaml_config(resources)
+        elif isinstance(resources, resources_lib.Resources):
             resources = {resources}
         # TODO(woosuk): Check if the resources are None.
-        self.resources = _with_docker_login_config(resources, self.envs)
+        self.resources = _with_docker_login_config(resources, self.envs,
+                                                   self.secrets)
         # Only have effect on RunPod.
         self.resources = _with_docker_username_for_runpod(
-            self.resources, self.envs)
+            self.resources, self.envs, self.secrets)
 
         # Evaluate if the task requires FUSE and set the requires_fuse flag
         for _, storage_obj in self.storage_mounts.items():
             if storage_obj.mode in storage_lib.MOUNTABLE_STORAGE_MODES:
                 for r in self.resources:
-                    r.requires_fuse = True
+                    r.set_requires_fuse(True)
                 break
 
         return self
@@ -790,6 +1136,10 @@ class Task:
         self.set_resources(type(self.resources)(new_resources_list))
         return self
 
+    def get_resource_config(self) -> Dict[str, Any]:
+        return _resources_to_config(self.resources,
+                                    factor_out_common_fields=True)
+
     @property
     def service(self) -> Optional[service_spec.SkyServiceSpec]:
         return self._service
@@ -807,8 +1157,8 @@ class Task:
         self._service = service
         return self
 
-    def set_time_estimator(self, func: Callable[['sky.Resources'],
-                                                int]) -> 'Task':
+    def set_time_estimator(
+            self, func: Callable[['resources_lib.Resources'], int]) -> 'Task':
         """Sets a func mapping resources to estimated time (secs).
 
         This is EXPERIMENTAL.
@@ -864,7 +1214,7 @@ class Task:
 
         Different from set_file_mounts(), this function updates into the
         existing file_mounts (calls ``dict.update()``), rather than
-        overwritting it.
+        overwriting it.
 
         This should be called before provisioning in order to take effect.
 
@@ -931,7 +1281,7 @@ class Task:
             self.storage_mounts = {}
             # Clear the requires_fuse flag if no storage mounts are set.
             for r in self.resources:
-                r.requires_fuse = False
+                r.set_requires_fuse(False)
             return self
         for target, storage_obj in storage_mounts.items():
             # TODO(zhwu): /home/username/sky_workdir as the target path need
@@ -956,7 +1306,7 @@ class Task:
                 # If any storage is using MOUNT mode, we need to enable FUSE in
                 # the resources.
                 for r in self.resources:
-                    r.requires_fuse = True
+                    r.set_requires_fuse(True)
         # Storage source validation is done in Storage object
         self.storage_mounts = storage_mounts
         return self
@@ -1170,6 +1520,16 @@ class Task:
                     self.update_file_mounts({
                         mnt_path: blob_path,
                     })
+                elif store_type is storage_lib.StoreType.COREWEAVE:
+                    if storage.source is not None and not isinstance(
+                            storage.source,
+                            list) and storage.source.startswith('cw://'):
+                        blob_path = storage.source
+                    else:
+                        blob_path = 'cw://' + 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} '
@@ -1219,11 +1579,83 @@ class Task:
                 d[k] = v
         return d
 
-    def to_yaml_config(self) -> Dict[str, Any]:
+    def update_workdir(self, workdir: Optional[str], git_url: Optional[str],
+                       git_ref: Optional[str]) -> 'Task':
+        """Updates the task workdir.
+
+        Args:
+            workdir: The workdir to update.
+            git_url: The git url to update.
+            git_ref: The git ref to update.
+        """
+        if self.workdir is None or isinstance(self.workdir, str):
+            if workdir is not None:
+                self.workdir = workdir
+                return self
+            if git_url is not None:
+                self.workdir = {}
+                self.workdir['url'] = git_url
+                if git_ref is not None:
+                    self.workdir['ref'] = git_ref
+                return self
+            return self
+        if git_url is not None:
+            self.workdir['url'] = git_url
+        if git_ref is not None:
+            self.workdir['ref'] = git_ref
+        return self
+
+    def update_envs_and_secrets_from_workdir(self) -> 'Task':
+        """Updates the task envs and secrets from the workdir."""
+        if self.workdir is None:
+            return self
+        if not isinstance(self.workdir, dict):
+            return self
+        url = self.workdir['url']
+        ref = self.workdir.get('ref', '')
+        token = os.environ.get(git.GIT_TOKEN_ENV_VAR)
+        ssh_key_path = os.environ.get(git.GIT_SSH_KEY_PATH_ENV_VAR)
+        try:
+            git_repo = git.GitRepo(url, ref, token, ssh_key_path)
+            clone_info = git_repo.get_repo_clone_info()
+            if clone_info is None:
+                return self
+            self.envs[git.GIT_URL_ENV_VAR] = clone_info.url
+            if ref:
+                ref_type = git_repo.get_ref_type()
+                if ref_type == git.GitRefType.COMMIT:
+                    self.envs[git.GIT_COMMIT_HASH_ENV_VAR] = ref
+                elif ref_type == git.GitRefType.BRANCH:
+                    self.envs[git.GIT_BRANCH_ENV_VAR] = ref
+                elif ref_type == git.GitRefType.TAG:
+                    self.envs[git.GIT_TAG_ENV_VAR] = ref
+            if clone_info.token is None and clone_info.ssh_key is None:
+                return self
+            if clone_info.token is not None:
+                self.secrets[git.GIT_TOKEN_ENV_VAR] = clone_info.token
+            if clone_info.ssh_key is not None:
+                self.secrets[git.GIT_SSH_KEY_ENV_VAR] = clone_info.ssh_key
+        except exceptions.GitError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'{str(e)}') from None
+        return self
+
+    def to_yaml_config(self,
+                       use_user_specified_yaml: bool = False) -> Dict[str, Any]:
         """Returns a yaml-style dict representation of the task.
 
         INTERNAL: this method is internal-facing.
         """
+        if use_user_specified_yaml:
+            if self._user_specified_yaml is None:
+                return self._to_yaml_config(redact_secrets=True)
+            config = yaml_utils.safe_load(self._user_specified_yaml)
+            if config.get('secrets') is not None:
+                config['secrets'] = {k: '<redacted>' for k in config['secrets']}
+            return config
+        return self._to_yaml_config()
+
+    def _to_yaml_config(self, redact_secrets: bool = False) -> Dict[str, Any]:
         config = {}
 
         def add_if_not_none(key, value, no_empty: bool = False):
@@ -1234,15 +1666,7 @@ class Task:
 
         add_if_not_none('name', self.name)
 
-        tmp_resource_config = {}
-        if len(self.resources) > 1:
-            resource_list = []
-            for r in self.resources:
-                resource_list.append(r.to_yaml_config())
-            key = 'ordered' if isinstance(self.resources, list) else 'any_of'
-            tmp_resource_config[key] = resource_list
-        else:
-            tmp_resource_config = list(self.resources)[0].to_yaml_config()
+        tmp_resource_config = _resources_to_config(self.resources)
 
         add_if_not_none('resources', tmp_resource_config)
 
@@ -1263,8 +1687,15 @@ class Task:
         add_if_not_none('workdir', self.workdir)
         add_if_not_none('event_callback', self.event_callback)
         add_if_not_none('run', self.run)
+
+        # Add envs without redaction
         add_if_not_none('envs', self.envs, no_empty=True)
 
+        secrets = self.secrets
+        if secrets and redact_secrets:
+            secrets = {k: '<redacted>' for k in secrets}
+        add_if_not_none('secrets', secrets, no_empty=True)
+
         add_if_not_none('file_mounts', {})
 
         if self.file_mounts is not None:
@@ -1277,6 +1708,15 @@ class Task:
             })
 
         add_if_not_none('file_mounts_mapping', self.file_mounts_mapping)
+        add_if_not_none('volumes', self.volumes)
+        if self.volume_mounts is not None:
+            config['volume_mounts'] = [
+                volume_mount.to_yaml_config()
+                for volume_mount in self.volume_mounts
+            ]
+        # we manually check if its empty to not clog up the generated yaml
+        add_if_not_none('_metadata', self._metadata if self._metadata else None)
+        add_if_not_none('_user_specified_yaml', self._user_specified_yaml)
         return config
 
     def get_required_cloud_features(
@@ -1304,7 +1744,7 @@ class Task:
         return required_features
 
     def __rshift__(self, b):
-        sky.dag.get_current_dag().add_edge(self, b)
+        dag_lib.get_current_dag().add_edge(self, b)
 
     def __repr__(self):
         if isinstance(self.run, str):
@@ -1339,3 +1779,47 @@ class Task:
         else:
             s += '\n  resources: default instances'
         return s
+
+
+def _resources_to_config(
+        resources: Union[List['resources_lib.Resources'],
+                         Set['resources_lib.Resources']],
+        factor_out_common_fields: bool = False) -> Dict[str, Any]:
+    if len(resources) > 1:
+        resource_list: List[Dict[str, Union[str, int]]] = []
+        for r in resources:
+            resource_list.append(r.to_yaml_config())
+        group_key = 'ordered' if isinstance(resources, list) else 'any_of'
+        if factor_out_common_fields:
+            return _factor_out_common_resource_fields(resource_list, group_key)
+        return {group_key: resource_list}
+    else:
+        return list(resources)[0].to_yaml_config()
+
+
+def _factor_out_common_resource_fields(configs: List[Dict[str, Union[str,
+                                                                     int]]],
+                                       group_key: str) -> Dict[str, Any]:
+    """Factors out the fields that are common to all resources."""
+    return_config: Dict[str, Any] = configs[0].copy()
+    if len(configs) > 1:
+        for config in configs[1:]:
+            for key, value in config.items():
+                if key in return_config and return_config[key] != value:
+                    del return_config[key]
+    num_empty_configs = 0
+    for config in configs:
+        keys_to_delete = []
+        for key, value in config.items():
+            if key in return_config:
+                keys_to_delete.append(key)
+        for key in keys_to_delete:
+            del config[key]
+        if not config:
+            num_empty_configs += 1
+
+    if num_empty_configs == len(configs):
+        return return_config
+    if len(configs) > 0:
+        return_config[group_key] = configs
+    return return_config