skypilot-nightly 1.0.0.dev20250624__py3-none-any.whl → 1.0.0.dev20250625__py3-none-any.whl

sky/server/stream_utils.py

@@ -155,9 +155,14 @@ async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
                 if request_task.status > requests_lib.RequestStatus.RUNNING:
                     if (request_task.status ==
                             requests_lib.RequestStatus.CANCELLED):
-                        buffer.append(
-                            f'{request_task.name!r} request {request_id}'
-                            ' cancelled\n')
+                        if request_task.should_retry:
+                            buffer.append(
+                                message_utils.encode_payload(
+                                    rich_utils.Control.RETRY.encode('')))
+                        else:
+                            buffer.append(
+                                f'{request_task.name!r} request {request_id}'
+                                ' cancelled\n')
                     break
             if not follow:
                 break

sky/server/uvicorn.py

@@ -3,17 +3,165 @@
 This module is a wrapper around uvicorn to customize the behavior of the
 server.
 """
-import functools
+import asyncio
 import os
+import signal
 import threading
-from typing import Optional
+import time
+from types import FrameType
+from typing import Optional, Union
 
+import filelock
 import uvicorn
 from uvicorn.supervisors import multiprocess
 
+from sky import sky_logging
+from sky.server import state
+from sky.server.requests import requests as requests_lib
+from sky.skylet import constants
 from sky.utils import context_utils
 from sky.utils import subprocess_utils
 
+logger = sky_logging.init_logger(__name__)
+
+# File lock path for coordinating graceful shutdown across processes
+_GRACEFUL_SHUTDOWN_LOCK_PATH = '/tmp/skypilot_graceful_shutdown.lock'
+
+# Interval to check for on-going requests.
+_WAIT_REQUESTS_INTERVAL_SECONDS = 5
+
+# Timeout for waiting for on-going requests to finish.
+try:
+    _WAIT_REQUESTS_TIMEOUT_SECONDS = int(
+        os.environ.get(constants.GRACE_PERIOD_SECONDS_ENV_VAR, '60'))
+except ValueError:
+    _WAIT_REQUESTS_TIMEOUT_SECONDS = 60
+
+# TODO(aylei): use decorator to register requests that need to be proactively
+# cancelled instead of hardcoding here.
+_RETRIABLE_REQUEST_NAMES = [
+    'sky.logs',
+    'sky.jobs.logs',
+    'sky.serve.logs',
+]
+
+
+class Server(uvicorn.Server):
+    """Server wrapper for uvicorn.
+
+    Extended functionalities:
+    - Handle exit signal and perform custom graceful shutdown.
+    - Run the server process with contextually aware.
+    """
+
+    def __init__(self, config: uvicorn.Config):
+        super().__init__(config=config)
+        self.exiting: bool = False
+
+    def handle_exit(self, sig: int, frame: Union[FrameType, None]) -> None:
+        """Handle exit signal.
+
+        When a server process receives a SIGTERM or SIGINT signal, a graceful
+        shutdown will be initiated. If a SIGINT signal is received again, the
+        server will be forcefully shutdown.
+        """
+        if self.exiting and sig == signal.SIGINT:
+            # The server has been siganled to exit and recieved a SIGINT again,
+            # do force shutdown.
+            logger.info('Force shutdown.')
+            self.should_exit = True
+            super().handle_exit(sig, frame)
+            return
+        if not self.exiting:
+            self.exiting = True
+            # Perform graceful shutdown in a separate thread to avoid blocking
+            # the main thread.
+            threading.Thread(target=self._graceful_shutdown,
+                             args=(sig, frame),
+                             daemon=True).start()
+
+    def _graceful_shutdown(self, sig: int, frame: Union[FrameType,
+                                                        None]) -> None:
+        """Perform graceful shutdown."""
+        # Block new requests so that we can wait until all on-going requests
+        # are finished. Note that /api/$verb operations are still allowed in
+        # this stage to ensure the client can still operate the on-going
+        # requests, e.g. /api/logs, /api/cancel, etc.
+        logger.info('Block new requests being submitted in worker '
+                    f'{os.getpid()}.')
+        state.set_block_requests(True)
+        # Ensure the shutting_down are set on all workers before next step.
+        # TODO(aylei): hacky, need a reliable solution.
+        time.sleep(1)
+
+        lock = filelock.FileLock(_GRACEFUL_SHUTDOWN_LOCK_PATH)
+        # Elect a coordinator process to handle on-going requests check
+        with lock.acquire():
+            logger.info(f'Worker {os.getpid()} elected as shutdown coordinator')
+            self._wait_requests()
+
+        logger.info('Shutting down server...')
+        self.should_exit = True
+        super().handle_exit(sig, frame)
+
+    def _wait_requests(self) -> None:
+        """Wait until all on-going requests are finished or cancelled."""
+        start_time = time.time()
+        while True:
+            statuses = [
+                requests_lib.RequestStatus.PENDING,
+                requests_lib.RequestStatus.RUNNING,
+            ]
+            reqs = requests_lib.get_request_tasks(status=statuses)
+            if not reqs:
+                break
+            logger.info(f'{len(reqs)} on-going requests '
+                        'found, waiting for them to finish...')
+            # Proactively cancel internal requests and logs requests since
+            # they can run for infinite time.
+            internal_request_ids = [
+                d.id for d in requests_lib.INTERNAL_REQUEST_DAEMONS
+            ]
+            if time.time() - start_time > _WAIT_REQUESTS_TIMEOUT_SECONDS:
+                logger.warning('Timeout waiting for on-going requests to '
+                               'finish, cancelling all on-going requests.')
+                for req in reqs:
+                    self.interrupt_request_for_retry(req.request_id)
+                break
+            interrupted = 0
+            for req in reqs:
+                if req.request_id in internal_request_ids:
+                    self.interrupt_request_for_retry(req.request_id)
+                    interrupted += 1
+                elif req.name in _RETRIABLE_REQUEST_NAMES:
+                    self.interrupt_request_for_retry(req.request_id)
+                    interrupted += 1
+                # TODO(aylei): interrupt pending requests to accelerate the
+                # shutdown.
+            # If some requests are not interrupted, wait for them to finish,
+            # otherwise we just check again immediately to accelerate the
+            # shutdown process.
+            if interrupted < len(reqs):
+                time.sleep(_WAIT_REQUESTS_INTERVAL_SECONDS)
+
+    def interrupt_request_for_retry(self, request_id: str) -> None:
+        """Interrupt a request for retry."""
+        with requests_lib.update_request(request_id) as req:
+            if req is None:
+                return
+            if req.pid is not None:
+                os.kill(req.pid, signal.SIGTERM)
+            req.status = requests_lib.RequestStatus.CANCELLED
+            req.should_retry = True
+        logger.info(
+            f'Request {request_id} interrupted and will be retried by client.')
+
+    def run(self, *args, **kwargs):
+        """Run the server process."""
+        context_utils.hijack_sys_attrs()
+        with self.capture_signals():
+            asyncio.run(self.serve(*args, **kwargs))
+
 
 def run(config: uvicorn.Config):
     """Run unvicorn server."""
@@ -22,28 +170,20 @@ def run(config: uvicorn.Config):
         # in uvicorn. Since we do not use reload now, simply
         # guard by an exception.
         raise ValueError('Reload is not supported yet.')
-    server = uvicorn.Server(config=config)
-    run_server_process = functools.partial(_run_server_process, server)
+    server = Server(config=config)
     try:
         if config.workers is not None and config.workers > 1:
             sock = config.bind_socket()
-            SlowStartMultiprocess(config,
-                                  target=run_server_process,
+            SlowStartMultiprocess(config, target=server.run,
                                   sockets=[sock]).run()
         else:
-            run_server_process()
+            server.run()
     finally:
         # Copied from unvicorn.run()
         if config.uds and os.path.exists(config.uds):
             os.remove(config.uds)
 
 
-def _run_server_process(server: uvicorn.Server, *args, **kwargs):
-    """Run the server process with contextually aware."""
-    context_utils.hijack_sys_attrs()
-    server.run(*args, **kwargs)
-
-
 class SlowStartMultiprocess(multiprocess.Multiprocess):
     """Uvicorn Multiprocess wrapper with slow start.
 

sky/setup_files/dependencies.py

@@ -62,6 +62,8 @@ install_requires = [
     # the client-side actually not importing them.
     'casbin',
     'sqlalchemy_adapter',
+    # Required for API server metrics
+    'prometheus_client>=0.8.0',
     'passlib',
 ]
 

sky/skylet/constants.py

@@ -413,6 +413,8 @@ SKY_USER_FILE_PATH = '~/.sky/generated'
 # Environment variable that is set to 'true' if this is a skypilot server.
 ENV_VAR_IS_SKYPILOT_SERVER = 'IS_SKYPILOT_SERVER'
 
+# Environment variable that is set to 'true' if metrics are enabled.
+ENV_VAR_SERVER_METRICS_ENABLED = 'SKY_API_SERVER_METRICS_ENABLED'
 # Environment variable that is set to 'true' if basic
 # authentication is enabled in the API server.
 ENV_VAR_ENABLE_BASIC_AUTH = 'ENABLE_BASIC_AUTH'
@@ -449,20 +451,29 @@ TIME_PATTERN: str = (
 
 MEMORY_SIZE_UNITS = {
     'kb': 2**10,
+    'ki': 2**10,
     'mb': 2**20,
+    'mi': 2**20,
     'gb': 2**30,
+    'gi': 2**30,
     'tb': 2**40,
+    'ti': 2**40,
     'pb': 2**50,
+    'pi': 2**50,
 }
 
 MEMORY_SIZE_PATTERN = (
     '^[0-9]+('
-    f'{"|".join([unit.lower() for unit in MEMORY_SIZE_UNITS])}'
-    ')?$/i')
-MEMORY_SIZE_PLUS_PATTERN = f'{MEMORY_SIZE_PATTERN[:-3]}+?$/i'
+    f'{"|".join([unit.lower() for unit in MEMORY_SIZE_UNITS])}|'
+    f'{"|".join([unit.upper() for unit in MEMORY_SIZE_UNITS])}|'
+    f'{"|".join([unit[0].upper() + unit[1:] for unit in MEMORY_SIZE_UNITS if len(unit) > 1])}'  # pylint: disable=line-too-long
+    ')?$')
+
+LAST_USE_TRUNC_LENGTH = 25
 
 MIN_PRIORITY = -1000
 MAX_PRIORITY = 1000
 DEFAULT_PRIORITY = 0
 
+GRACE_PERIOD_SECONDS_ENV_VAR = SKYPILOT_ENV_VAR_PREFIX + 'GRACE_PERIOD_SECONDS'
 COST_REPORT_DEFAULT_DAYS = 30

sky/task.py

@@ -24,6 +24,7 @@ from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import schemas
 from sky.utils import ux_utils
+from sky.volumes import volume as volume_lib
 
 if typing.TYPE_CHECKING:
     import yaml
@@ -246,12 +247,14 @@ class Task:
         secrets: Optional[Dict[str, str]] = None,
         workdir: Optional[str] = None,
         num_nodes: Optional[int] = None,
+        volumes: Optional[Dict[str, str]] = 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,
+        volume_mounts: Optional[List[volume_lib.VolumeMount]] = None,
     ):
         """Initializes a Task.
 
@@ -319,6 +322,7 @@ class Task:
         self.setup = setup
         self._envs = envs or {}
         self._secrets = secrets or {}
+        self._volumes = volumes or {}
 
         # Validate Docker login configuration early if both envs and secrets
         # contain Docker variables
@@ -361,7 +365,9 @@ class Task:
         self.best_resources: Optional[sky.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)
 
         dag = sky.dag.get_current_dag()
         if dag is not None:
@@ -442,12 +448,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(
@@ -462,17 +465,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.
@@ -587,6 +598,7 @@ class Task:
             secrets=config.pop('secrets', None),
             event_callback=config.pop('event_callback', None),
             file_mounts_mapping=config.pop('file_mounts_mapping', None),
+            volumes=config.pop('volumes', None),
         )
 
         # Create lists to store storage objects inlined in file_mounts.
@@ -711,6 +723,16 @@ class Task:
             service = service_spec.SkyServiceSpec.from_yaml_config(service)
         task.set_service(service)
 
+        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
 
@@ -745,6 +767,97 @@ class Task:
             config = {}
         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] = sky.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
@@ -767,6 +880,10 @@ class Task:
     def secrets(self) -> Dict[str, str]:
         return self._secrets
 
+    @property
+    def volumes(self) -> Dict[str, str]:
+        return self._volumes
+
     def update_envs(
             self, envs: Union[None, List[Tuple[str, str]],
                               Dict[str, str]]) -> 'Task':
@@ -1453,6 +1570,12 @@ 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
+            ]
         return config
 
     def get_required_cloud_features(

sky/templates/kubernetes-ray.yml.j2

@@ -243,6 +243,22 @@ provider:
             # This selector must match the head node pod's selector below.
             selector:
                 component: {{cluster_name_on_cloud}}-head
+      # Headless service mapping hostnames to rest of the worker nodes
+      {% for worker_id in range(1, num_nodes) %}
+      - apiVersion: v1
+        kind: Service
+        metadata:
+          labels:
+            parent: skypilot
+            skypilot-cluster: {{cluster_name_on_cloud}}
+            skypilot-user: {{ user }}
+          name: {{cluster_name_on_cloud}}-worker{{ worker_id }}
+        spec:
+          selector:
+              component: {{cluster_name_on_cloud}}-worker{{ worker_id }}
+          clusterIP: None
+      {% endfor %}
+
 
 # Specify the pod type for the ray head node (as configured below).
 head_node_type: ray_head_default
@@ -255,7 +271,7 @@ available_node_types:
       metadata:
         # name will be filled in the provisioner
         # head node name will be {{cluster_name_on_cloud}}-head, which will match the head node service selector above if a head node
-        # service is required.
+        # service is required. Worker nodes are named {{cluster_name_on_cloud}}-worker{{ node_id }}
         labels:
             parent: skypilot
             # component will be set for the head node pod to be the same as the head node service selector above if a
@@ -287,6 +303,10 @@ available_node_types:
         serviceAccountName: {{k8s_service_account_name}}
         automountServiceAccountToken: {{k8s_automount_sa_token}}
         restartPolicy: {{ "Always" if high_availability else "Never" }}
+        {% if volume_mounts %}
+        securityContext:
+          fsGroup: 1000
+        {% endif %}
 
         # Add node selector if GPU/TPUs are requested:
         {% if (k8s_topology_label_key is not none and k8s_topology_label_value is not none) or (k8s_spot_label_key is not none) %}
@@ -365,6 +385,11 @@ available_node_types:
           persistentVolumeClaim:
             claimName: {{cluster_name_on_cloud}}-{{k8s_high_availability_deployment_volume_mount_name}}
         {% endif %}
+        {% for volume_mount in volume_mounts %}
+        - name: {{volume_mount.name}}
+          persistentVolumeClaim:
+            claimName: {{volume_mount.volume_name_on_cloud}}
+        {% endfor %}
         containers:
         - name: ray-node
           imagePullPolicy: IfNotPresent
@@ -734,6 +759,10 @@ available_node_types:
           - name: fusermount-shared-dir
             mountPath: {{k8s_fusermount_shared_dir}}
           {% endif %}
+          {% for volume_mount in volume_mounts %}
+          - name: {{volume_mount.name}}
+            mountPath: {{volume_mount.path}}
+          {% endfor %}
           resources:
             requests:
               cpu: {{cpus}}

sky/users/permission.py

@@ -18,6 +18,8 @@ from sky.utils import common_utils
 
 logging.getLogger('casbin.policy').setLevel(sky_logging.ERROR)
 logging.getLogger('casbin.role').setLevel(sky_logging.ERROR)
+logging.getLogger('casbin.model').setLevel(sky_logging.ERROR)
+logging.getLogger('casbin.rbac').setLevel(sky_logging.ERROR)
 logger = sky_logging.init_logger(__name__)
 
 # Filelocks for the policy update.

sky/utils/context.py

@@ -254,7 +254,9 @@ class Popen(subprocess.Popen):
     def __init__(self, *args, **kwargs):
         env = kwargs.pop('env', None)
         if env is None:
-            env = os.environ
+            # Pass a copy of current context.environ to avoid race condition
+            # when the context is updated after the Popen is created.
+            env = os.environ.copy()
         super().__init__(*args, env=env, **kwargs)
 
 

sky/utils/resources_utils.py

@@ -8,6 +8,7 @@ import typing
 from typing import Dict, List, Optional, Set, Union
 
 from sky import skypilot_config
+from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import registry
 from sky.utils import ux_utils
@@ -331,3 +332,68 @@ def make_launchables_for_valid_region_zones(
             # Batch the requests at the granularity of a single region.
             launchables.append(launchable_resources.copy(region=region.name))
     return launchables
+
+
+def parse_memory_resource(resource_qty_str: Union[str, int, float],
+                          field_name: str,
+                          ret_type: type = int,
+                          unit: str = 'gb',
+                          allow_plus: bool = False,
+                          allow_x: bool = False,
+                          allow_rounding: bool = False) -> str:
+    """Returns memory size in chosen units given a resource quantity string.
+
+    Args:
+        resource_qty_str: Resource quantity string
+        unit: Unit to convert to
+        allow_plus: Whether to allow '+' prefix
+        allow_x: Whether to allow 'x' suffix
+    """
+    assert unit in constants.MEMORY_SIZE_UNITS, f'Invalid unit: {unit}'
+
+    error_msg = (f'"{field_name}" field should be a '
+                 f'{constants.MEMORY_SIZE_PATTERN}+?,'
+                 f' got {resource_qty_str}')
+
+    resource_str = str(resource_qty_str)
+
+    # Handle plus and x suffixes, x is only used internally for jobs controller
+    plus = ''
+    if resource_str.endswith('+'):
+        if allow_plus:
+            resource_str = resource_str[:-1]
+            plus = '+'
+        else:
+            raise ValueError(error_msg)
+
+    x = ''
+    if resource_str.endswith('x'):
+        if allow_x:
+            resource_str = resource_str[:-1]
+            x = 'x'
+        else:
+            raise ValueError(error_msg)
+
+    try:
+        # We assume it is already in the wanted units to maintain backwards
+        # compatibility
+        ret_type(resource_str)
+        return f'{resource_str}{plus}{x}'
+    except ValueError:
+        pass
+
+    resource_str = resource_str.lower()
+    for mem_unit, multiplier in constants.MEMORY_SIZE_UNITS.items():
+        if resource_str.endswith(mem_unit):
+            try:
+                value = ret_type(resource_str[:-len(mem_unit)])
+                converted = (value * multiplier /
+                             constants.MEMORY_SIZE_UNITS[unit])
+                if not allow_rounding and ret_type(converted) != converted:
+                    raise ValueError(error_msg)
+                converted = ret_type(converted)
+                return f'{converted}{plus}{x}'
+            except ValueError:
+                continue
+
+    raise ValueError(error_msg)