metaflow 2.12.33__py2.py3-none-any.whl → 2.12.35__py2.py3-none-any.whl

metaflow/plugins/events_decorator.py

@@ -1,9 +1,11 @@
 import re
+import json
 
 from metaflow import current
 from metaflow.decorators import FlowDecorator
 from metaflow.exception import MetaflowException
 from metaflow.util import is_stringish
+from metaflow.parameters import DeployTimeField, deploy_time_eval
 
 # TODO: Support dynamic parameter mapping through a context object that exposes
 #       flow name and user name similar to parameter context
@@ -68,6 +70,75 @@ class TriggerDecorator(FlowDecorator):
         "options": {},
     }
 
+    def process_event_name(self, event):
+        if is_stringish(event):
+            return {"name": str(event)}
+        elif isinstance(event, dict):
+            if "name" not in event:
+                raise MetaflowException(
+                    "The *event* attribute for *@trigger* is missing the *name* key."
+                )
+            if callable(event["name"]) and not isinstance(
+                event["name"], DeployTimeField
+            ):
+                event["name"] = DeployTimeField(
+                    "event_name", str, None, event["name"], False
+                )
+            event["parameters"] = self.process_parameters(event.get("parameters", {}))
+            return event
+        elif callable(event) and not isinstance(event, DeployTimeField):
+            return DeployTimeField("event", [str, dict], None, event, False)
+        else:
+            raise MetaflowException(
+                "Incorrect format for *event* attribute in *@trigger* decorator. "
+                "Supported formats are string and dictionary - \n"
+                "@trigger(event='foo') or @trigger(event={'name': 'foo', "
+                "'parameters': {'alpha': 'beta'}})"
+            )
+
+    def process_parameters(self, parameters):
+        new_param_values = {}
+        if isinstance(parameters, (list, tuple)):
+            for mapping in parameters:
+                if is_stringish(mapping):
+                    new_param_values[mapping] = mapping
+                elif callable(mapping) and not isinstance(mapping, DeployTimeField):
+                    mapping = DeployTimeField(
+                        "parameter_val", str, None, mapping, False
+                    )
+                    new_param_values[mapping] = mapping
+                elif isinstance(mapping, (list, tuple)) and len(mapping) == 2:
+                    if callable(mapping[0]) and not isinstance(
+                        mapping[0], DeployTimeField
+                    ):
+                        mapping[0] = DeployTimeField(
+                            "parameter_val", str, None, mapping[0], False
+                        )
+                    if callable(mapping[1]) and not isinstance(
+                        mapping[1], DeployTimeField
+                    ):
+                        mapping[1] = DeployTimeField(
+                            "parameter_val", str, None, mapping[1], False
+                        )
+                    new_param_values[mapping[0]] = mapping[1]
+                else:
+                    raise MetaflowException(
+                        "The *parameters* attribute for event is invalid. "
+                        "It should be a list/tuple of strings and lists/tuples of size 2"
+                    )
+        elif callable(parameters) and not isinstance(parameters, DeployTimeField):
+            return DeployTimeField(
+                "parameters", [list, dict, tuple], None, parameters, False
+            )
+        elif isinstance(parameters, dict):
+            for key, value in parameters.items():
+                if callable(key) and not isinstance(key, DeployTimeField):
+                    key = DeployTimeField("flow_parameter", str, None, key, False)
+                if callable(value) and not isinstance(value, DeployTimeField):
+                    value = DeployTimeField("signal_parameter", str, None, value, False)
+                new_param_values[key] = value
+        return new_param_values
+
     def flow_init(
         self,
         flow_name,
@@ -86,41 +157,9 @@ class TriggerDecorator(FlowDecorator):
                 "attributes in *@trigger* decorator."
             )
         elif self.attributes["event"]:
-            # event attribute supports the following formats -
-            #     1. event='table.prod_db.members'
-            #     2. event={'name': 'table.prod_db.members',
-            #               'parameters': {'alpha': 'member_weight'}}
-            if is_stringish(self.attributes["event"]):
-                self.triggers.append({"name": str(self.attributes["event"])})
-            elif isinstance(self.attributes["event"], dict):
-                if "name" not in self.attributes["event"]:
-                    raise MetaflowException(
-                        "The *event* attribute for *@trigger* is missing the "
-                        "*name* key."
-                    )
-                param_value = self.attributes["event"].get("parameters", {})
-                if isinstance(param_value, (list, tuple)):
-                    new_param_value = {}
-                    for mapping in param_value:
-                        if is_stringish(mapping):
-                            new_param_value[mapping] = mapping
-                        elif isinstance(mapping, (list, tuple)) and len(mapping) == 2:
-                            new_param_value[mapping[0]] = mapping[1]
-                        else:
-                            raise MetaflowException(
-                                "The *parameters* attribute for event '%s' is invalid. "
-                                "It should be a list/tuple of strings and lists/tuples "
-                                "of size 2" % self.attributes["event"]["name"]
-                            )
-                    self.attributes["event"]["parameters"] = new_param_value
-                self.triggers.append(self.attributes["event"])
-            else:
-                raise MetaflowException(
-                    "Incorrect format for *event* attribute in *@trigger* decorator. "
-                    "Supported formats are string and dictionary - \n"
-                    "@trigger(event='foo') or @trigger(event={'name': 'foo', "
-                    "'parameters': {'alpha': 'beta'}})"
-                )
+            event = self.attributes["event"]
+            processed_event = self.process_event_name(event)
+            self.triggers.append(processed_event)
         elif self.attributes["events"]:
             # events attribute supports the following formats -
             #     1. events=[{'name': 'table.prod_db.members',
@@ -128,43 +167,17 @@ class TriggerDecorator(FlowDecorator):
             #                {'name': 'table.prod_db.metadata',
             #               'parameters': {'beta': 'grade'}}]
             if isinstance(self.attributes["events"], list):
+                # process every event in events
                 for event in self.attributes["events"]:
-                    if is_stringish(event):
-                        self.triggers.append({"name": str(event)})
-                    elif isinstance(event, dict):
-                        if "name" not in event:
-                            raise MetaflowException(
-                                "One or more events in *events* attribute for "
-                                "*@trigger* are missing the *name* key."
-                            )
-                        param_value = event.get("parameters", {})
-                        if isinstance(param_value, (list, tuple)):
-                            new_param_value = {}
-                            for mapping in param_value:
-                                if is_stringish(mapping):
-                                    new_param_value[mapping] = mapping
-                                elif (
-                                    isinstance(mapping, (list, tuple))
-                                    and len(mapping) == 2
-                                ):
-                                    new_param_value[mapping[0]] = mapping[1]
-                                else:
-                                    raise MetaflowException(
-                                        "The *parameters* attribute for event '%s' is "
-                                        "invalid. It should be a list/tuple of strings "
-                                        "and lists/tuples of size 2" % event["name"]
-                                    )
-                            event["parameters"] = new_param_value
-                        self.triggers.append(event)
-                    else:
-                        raise MetaflowException(
-                            "One or more events in *events* attribute in *@trigger* "
-                            "decorator have an incorrect format. Supported format "
-                            "is dictionary - \n"
-                            "@trigger(events=[{'name': 'foo', 'parameters': {'alpha': "
-                            "'beta'}}, {'name': 'bar', 'parameters': "
-                            "{'gamma': 'kappa'}}])"
-                        )
+                    processed_event = self.process_event_name(event)
+                    self.triggers.append(processed_event)
+            elif callable(self.attributes["events"]) and not isinstance(
+                self.attributes["events"], DeployTimeField
+            ):
+                trig = DeployTimeField(
+                    "events", list, None, self.attributes["events"], False
+                )
+                self.triggers.append(trig)
             else:
                 raise MetaflowException(
                     "Incorrect format for *events* attribute in *@trigger* decorator. "
@@ -178,7 +191,12 @@ class TriggerDecorator(FlowDecorator):
             raise MetaflowException("No event(s) specified in *@trigger* decorator.")
 
         # same event shouldn't occur more than once
-        names = [x["name"] for x in self.triggers]
+        names = [
+            x["name"]
+            for x in self.triggers
+            if not isinstance(x, DeployTimeField)
+            and not isinstance(x["name"], DeployTimeField)
+        ]
         if len(names) != len(set(names)):
             raise MetaflowException(
                 "Duplicate event names defined in *@trigger* decorator."
@@ -188,6 +206,104 @@ class TriggerDecorator(FlowDecorator):
 
         # TODO: Handle scenario for local testing using --trigger.
 
+    def format_deploytime_value(self):
+        new_triggers = []
+        for trigger in self.triggers:
+            # Case where trigger is a function that returns a list of events
+            # Need to do this bc we need to iterate over list later
+            if isinstance(trigger, DeployTimeField):
+                evaluated_trigger = deploy_time_eval(trigger)
+                if isinstance(evaluated_trigger, dict):
+                    trigger = evaluated_trigger
+                elif isinstance(evaluated_trigger, str):
+                    trigger = {"name": evaluated_trigger}
+                if isinstance(evaluated_trigger, list):
+                    for trig in evaluated_trigger:
+                        if is_stringish(trig):
+                            new_triggers.append({"name": trig})
+                        else:  # dict or another deploytimefield
+                            new_triggers.append(trig)
+                else:
+                    new_triggers.append(trigger)
+            else:
+                new_triggers.append(trigger)
+
+        self.triggers = new_triggers
+        for trigger in self.triggers:
+            old_trigger = trigger
+            trigger_params = trigger.get("parameters", {})
+            # Case where param is a function (can return list or dict)
+            if isinstance(trigger_params, DeployTimeField):
+                trigger_params = deploy_time_eval(trigger_params)
+            # If params is a list of strings, convert to dict with same key and value
+            if isinstance(trigger_params, (list, tuple)):
+                new_trigger_params = {}
+                for mapping in trigger_params:
+                    if is_stringish(mapping) or callable(mapping):
+                        new_trigger_params[mapping] = mapping
+                    elif callable(mapping) and not isinstance(mapping, DeployTimeField):
+                        mapping = DeployTimeField(
+                            "parameter_val", str, None, mapping, False
+                        )
+                        new_trigger_params[mapping] = mapping
+                    elif isinstance(mapping, (list, tuple)) and len(mapping) == 2:
+                        if callable(mapping[0]) and not isinstance(
+                            mapping[0], DeployTimeField
+                        ):
+                            mapping[0] = DeployTimeField(
+                                "parameter_val",
+                                str,
+                                None,
+                                mapping[1],
+                                False,
+                            )
+                        if callable(mapping[1]) and not isinstance(
+                            mapping[1], DeployTimeField
+                        ):
+                            mapping[1] = DeployTimeField(
+                                "parameter_val",
+                                str,
+                                None,
+                                mapping[1],
+                                False,
+                            )
+
+                        new_trigger_params[mapping[0]] = mapping[1]
+                    else:
+                        raise MetaflowException(
+                            "The *parameters* attribute for event '%s' is invalid. "
+                            "It should be a list/tuple of strings and lists/tuples "
+                            "of size 2" % self.attributes["event"]["name"]
+                        )
+                trigger_params = new_trigger_params
+            trigger["parameters"] = trigger_params
+
+            trigger_name = trigger.get("name")
+            # Case where just the name is a function (always a str)
+            if isinstance(trigger_name, DeployTimeField):
+                trigger_name = deploy_time_eval(trigger_name)
+                trigger["name"] = trigger_name
+
+            # Third layer
+            # {name:, parameters:[func, ..., ...]}
+            # {name:, parameters:{func : func2}}
+            for trigger in self.triggers:
+                old_trigger = trigger
+                trigger_params = trigger.get("parameters", {})
+                new_trigger_params = {}
+                for key, value in trigger_params.items():
+                    if isinstance(value, DeployTimeField) and key is value:
+                        evaluated_param = deploy_time_eval(value)
+                        new_trigger_params[evaluated_param] = evaluated_param
+                    elif isinstance(value, DeployTimeField):
+                        new_trigger_params[key] = deploy_time_eval(value)
+                    elif isinstance(key, DeployTimeField):
+                        new_trigger_params[deploy_time_eval(key)] = value
+                    else:
+                        new_trigger_params[key] = value
+                trigger["parameters"] = new_trigger_params
+            self.triggers[self.triggers.index(old_trigger)] = trigger
+
 
 class TriggerOnFinishDecorator(FlowDecorator):
     """
@@ -312,6 +428,13 @@ class TriggerOnFinishDecorator(FlowDecorator):
                             "The *project_branch* attribute of the *flow* is not a string"
                         )
                 self.triggers.append(result)
+            elif callable(self.attributes["flow"]) and not isinstance(
+                self.attributes["flow"], DeployTimeField
+            ):
+                trig = DeployTimeField(
+                    "fq_name", [str, dict], None, self.attributes["flow"], False
+                )
+                self.triggers.append(trig)
             else:
                 raise MetaflowException(
                     "Incorrect type for *flow* attribute in *@trigger_on_finish* "
@@ -369,6 +492,13 @@ class TriggerOnFinishDecorator(FlowDecorator):
                             "Supported type is string or Dict[str, str]- \n"
                             "@trigger_on_finish(flows=['FooFlow', 'BarFlow']"
                         )
+            elif callable(self.attributes["flows"]) and not isinstance(
+                self.attributes["flows"], DeployTimeField
+            ):
+                trig = DeployTimeField(
+                    "flows", list, None, self.attributes["flows"], False
+                )
+                self.triggers.append(trig)
             else:
                 raise MetaflowException(
                     "Incorrect type for *flows* attribute in *@trigger_on_finish* "
@@ -383,6 +513,8 @@ class TriggerOnFinishDecorator(FlowDecorator):
 
         # Make triggers @project aware
         for trigger in self.triggers:
+            if isinstance(trigger, DeployTimeField):
+                continue
             if trigger["fq_name"].count(".") == 0:
                 # fully qualified name is just the flow name
                 trigger["flow"] = trigger["fq_name"]
@@ -427,5 +559,54 @@ class TriggerOnFinishDecorator(FlowDecorator):
                 run_objs.append(run_obj)
             current._update_env({"trigger": Trigger.from_runs(run_objs)})
 
+    def _parse_fq_name(self, trigger):
+        if isinstance(trigger, DeployTimeField):
+            trigger["fq_name"] = deploy_time_eval(trigger["fq_name"])
+        if trigger["fq_name"].count(".") == 0:
+            # fully qualified name is just the flow name
+            trigger["flow"] = trigger["fq_name"]
+        elif trigger["fq_name"].count(".") >= 2:
+            # fully qualified name is of the format - project.branch.flow_name
+            trigger["project"], tail = trigger["fq_name"].split(".", maxsplit=1)
+            trigger["branch"], trigger["flow"] = tail.rsplit(".", maxsplit=1)
+        else:
+            raise MetaflowException(
+                "Incorrect format for *flow* in *@trigger_on_finish* "
+                "decorator. Specify either just the *flow_name* or a fully "
+                "qualified name like *project_name.branch_name.flow_name*."
+            )
+        if not re.match(r"^[A-Za-z0-9_]+$", trigger["flow"]):
+            raise MetaflowException(
+                "Invalid flow name *%s* in *@trigger_on_finish* "
+                "decorator. Only alphanumeric characters and "
+                "underscores(_) are allowed." % trigger["flow"]
+            )
+        return trigger
+
+    def format_deploytime_value(self):
+        for trigger in self.triggers:
+            # Case were trigger is a function that returns a list
+            # Need to do this bc we need to iterate over list and process
+            if isinstance(trigger, DeployTimeField):
+                deploy_value = deploy_time_eval(trigger)
+                if isinstance(deploy_value, list):
+                    self.triggers = deploy_value
+            else:
+                break
+        for trigger in self.triggers:
+            # Entire trigger is a function (returns either string or dict)
+            old_trig = trigger
+            if isinstance(trigger, DeployTimeField):
+                trigger = deploy_time_eval(trigger)
+            if isinstance(trigger, dict):
+                trigger["fq_name"] = trigger.get("name")
+                trigger["project"] = trigger.get("project")
+                trigger["branch"] = trigger.get("project_branch")
+            # We also added this bc it won't be formatted yet
+            if isinstance(trigger, str):
+                trigger = {"fq_name": trigger}
+                trigger = self._parse_fq_name(trigger)
+            self.triggers[self.triggers.index(old_trig)] = trigger
+
     def get_top_level_options(self):
         return list(self._option_values.items())

metaflow/plugins/kubernetes/kube_utils.py

@@ -23,3 +23,32 @@ def parse_cli_options(flow_name, run_id, user, my_runs, echo):
             raise CommandException("A previous run id was not found. Specify --run-id.")
 
     return flow_name, run_id, user
+
+
+def qos_requests_and_limits(qos: str, cpu: int, memory: int, storage: int):
+    "return resource requests and limits for the kubernetes pod based on the given QoS Class"
+    # case insensitive matching for QoS class
+    qos = qos.lower()
+    # Determine the requests and limits to define chosen QoS class
+    qos_limits = {}
+    qos_requests = {}
+    if qos == "guaranteed":
+        # Guaranteed - has both cpu/memory limits. requests not required, as these will be inferred.
+        qos_limits = {
+            "cpu": str(cpu),
+            "memory": "%sM" % str(memory),
+            "ephemeral-storage": "%sM" % str(storage),
+        }
+        # NOTE: Even though Kubernetes will produce matching requests for the specified limits, this happens late in the lifecycle.
+        # We specify them explicitly here to make some K8S tooling happy, in case they rely on .resources.requests being present at time of submitting the job.
+        qos_requests = qos_limits
+    else:
+        # Burstable - not Guaranteed, and has a memory/cpu limit or request
+        qos_requests = {
+            "cpu": str(cpu),
+            "memory": "%sM" % str(memory),
+            "ephemeral-storage": "%sM" % str(storage),
+        }
+    # TODO: Add support for BestEffort once there is a use case for it.
+    # BestEffort - no limit or requests for cpu/memory
+    return qos_requests, qos_limits

metaflow/plugins/kubernetes/kubernetes.py

@@ -196,6 +196,7 @@ class Kubernetes(object):
         shared_memory=None,
         port=None,
         num_parallel=None,
+        qos=None,
     ):
         name = "js-%s" % str(uuid4())[:6]
         jobset = (
@@ -228,6 +229,7 @@ class Kubernetes(object):
                 shared_memory=shared_memory,
                 port=port,
                 num_parallel=num_parallel,
+                qos=qos,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
@@ -488,6 +490,7 @@ class Kubernetes(object):
         shared_memory=None,
         port=None,
         name_pattern=None,
+        qos=None,
     ):
         if env is None:
             env = {}
@@ -528,6 +531,7 @@ class Kubernetes(object):
                 persistent_volume_claims=persistent_volume_claims,
                 shared_memory=shared_memory,
                 port=port,
+                qos=qos,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -126,6 +126,12 @@ def kubernetes():
     type=int,
     help="Number of parallel nodes to run as a multi-node job.",
 )
+@click.option(
+    "--qos",
+    default=None,
+    type=str,
+    help="Quality of Service class for the Kubernetes pod",
+)
 @click.pass_context
 def step(
     ctx,
@@ -154,6 +160,7 @@ def step(
     shared_memory=None,
     port=None,
     num_parallel=None,
+    qos=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -294,6 +301,7 @@ def step(
                 shared_memory=shared_memory,
                 port=port,
                 num_parallel=num_parallel,
+                qos=qos,
             )
     except Exception as e:
         traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -26,6 +26,7 @@ from metaflow.metaflow_config import (
     KUBERNETES_SERVICE_ACCOUNT,
     KUBERNETES_SHARED_MEMORY,
     KUBERNETES_TOLERATIONS,
+    KUBERNETES_QOS,
 )
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
@@ -41,6 +42,8 @@ except NameError:
     unicode = str
     basestring = str
 
+SUPPORTED_KUBERNETES_QOS_CLASSES = ["Guaranteed", "Burstable"]
+
 
 class KubernetesDecorator(StepDecorator):
     """
@@ -109,6 +112,8 @@ class KubernetesDecorator(StepDecorator):
     hostname_resolution_timeout: int, default 10 * 60
         Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
         Only applicable when @parallel is used.
+    qos: str, default: Burstable
+        Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
     """
 
     name = "kubernetes"
@@ -136,6 +141,7 @@ class KubernetesDecorator(StepDecorator):
         "compute_pool": None,
         "executable": None,
         "hostname_resolution_timeout": 10 * 60,
+        "qos": KUBERNETES_QOS,
     }
     package_url = None
     package_sha = None
@@ -259,6 +265,17 @@ class KubernetesDecorator(StepDecorator):
         self.step = step
         self.flow_datastore = flow_datastore
 
+        if (
+            self.attributes["qos"] is not None
+            # case insensitive matching.
+            and self.attributes["qos"].lower()
+            not in [c.lower() for c in SUPPORTED_KUBERNETES_QOS_CLASSES]
+        ):
+            raise MetaflowException(
+                "*%s* is not a valid Kubernetes QoS class. Choose one of the following: %s"
+                % (self.attributes["qos"], ", ".join(SUPPORTED_KUBERNETES_QOS_CLASSES))
+            )
+
         if any([deco.name == "batch" for deco in decos]):
             raise MetaflowException(
                 "Step *{step}* is marked for execution both on AWS Batch and "

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -15,6 +15,8 @@ from .kubernetes_jobsets import (
     KubernetesJobSet,
 )  # We need this import for Kubernetes Client.
 
+from .kube_utils import qos_requests_and_limits
+
 
 class KubernetesJobException(MetaflowException):
     headline = "Kubernetes job error"
@@ -74,6 +76,13 @@ class KubernetesJob(object):
             if self._kwargs["shared_memory"]
             else None
         )
+        qos_requests, qos_limits = qos_requests_and_limits(
+            self._kwargs["qos"],
+            self._kwargs["cpu"],
+            self._kwargs["memory"],
+            self._kwargs["disk"],
+        )
+
         return client.V1JobSpec(
             # Retries are handled by Metaflow when it is responsible for
             # executing the flow. The responsibility is moved to Kubernetes
@@ -154,20 +163,18 @@ class KubernetesJob(object):
                             image_pull_policy=self._kwargs["image_pull_policy"],
                             name=self._kwargs["step_name"].replace("_", "-"),
                             resources=client.V1ResourceRequirements(
-                                requests={
-                                    "cpu": str(self._kwargs["cpu"]),
-                                    "memory": "%sM" % str(self._kwargs["memory"]),
-                                    "ephemeral-storage": "%sM"
-                                    % str(self._kwargs["disk"]),
-                                },
+                                requests=qos_requests,
                                 limits={
-                                    "%s.com/gpu".lower()
-                                    % self._kwargs["gpu_vendor"]: str(
-                                        self._kwargs["gpu"]
-                                    )
-                                    for k in [0]
-                                    # Don't set GPU limits if gpu isn't specified.
-                                    if self._kwargs["gpu"] is not None
+                                    **qos_limits,
+                                    **{
+                                        "%s.com/gpu".lower()
+                                        % self._kwargs["gpu_vendor"]: str(
+                                            self._kwargs["gpu"]
+                                        )
+                                        for k in [0]
+                                        # Don't set GPU limits if gpu isn't specified.
+                                        if self._kwargs["gpu"] is not None
+                                    },
                                 },
                             ),
                             volume_mounts=(

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -9,6 +9,8 @@ from metaflow.metaflow_config import KUBERNETES_JOBSET_GROUP, KUBERNETES_JOBSET_
 from metaflow.tracing import inject_tracing_vars
 from metaflow.metaflow_config import KUBERNETES_SECRETS
 
+from .kube_utils import qos_requests_and_limits
+
 
 class KubernetesJobsetException(MetaflowException):
     headline = "Kubernetes jobset error"
@@ -554,7 +556,12 @@ class JobSetSpec(object):
             if self._kwargs["shared_memory"]
             else None
         )
-
+        qos_requests, qos_limits = qos_requests_and_limits(
+            self._kwargs["qos"],
+            self._kwargs["cpu"],
+            self._kwargs["memory"],
+            self._kwargs["disk"],
+        )
         return dict(
             name=self.name,
             template=client.api_client.ApiClient().sanitize_for_serialization(
@@ -653,21 +660,18 @@ class JobSetSpec(object):
                                             "_", "-"
                                         ),
                                         resources=client.V1ResourceRequirements(
-                                            requests={
-                                                "cpu": str(self._kwargs["cpu"]),
-                                                "memory": "%sM"
-                                                % str(self._kwargs["memory"]),
-                                                "ephemeral-storage": "%sM"
-                                                % str(self._kwargs["disk"]),
-                                            },
+                                            requests=qos_requests,
                                             limits={
-                                                "%s.com/gpu".lower()
-                                                % self._kwargs["gpu_vendor"]: str(
-                                                    self._kwargs["gpu"]
-                                                )
-                                                for k in [0]
-                                                # Don't set GPU limits if gpu isn't specified.
-                                                if self._kwargs["gpu"] is not None
+                                                **qos_limits,
+                                                **{
+                                                    "%s.com/gpu".lower()
+                                                    % self._kwargs["gpu_vendor"]: str(
+                                                        self._kwargs["gpu"]
+                                                    )
+                                                    for k in [0]
+                                                    # Don't set GPU limits if gpu isn't specified.
+                                                    if self._kwargs["gpu"] is not None
+                                                },
                                             },
                                         ),
                                         volume_mounts=(