add support for default task resource overrides

Author: redartera

flytekit/clis/sdk_in_container/register.py

@@ -15,7 +15,8 @@
 from flytekit.configuration import ImageConfig
 from flytekit.configuration.default_images import DefaultImages
 from flytekit.constants import CopyFileDetection
-from flytekit.interaction.click_types import key_value_callback
+from flytekit.core.resources import ResourceSpec
+from flytekit.interaction.click_types import key_value_callback, resource_spec_callback
 from flytekit.loggers import logger
 from flytekit.tools import repo
 
@@ -134,6 +135,15 @@
     callback=key_value_callback,
     help="Environment variables to set in the container, of the format `ENV_NAME=ENV_VALUE`",
 )
+@click.option(
+    "--default-resources",
+    required=False,
+    type=str,
+    callback=resource_spec_callback,
+    help="Override default task resource requests and limits for tasks that have no statically defined resource request and limit. "
+    """Example usage: --default-resources 'cpu=1;mem=2Gi;gpu=1' for requests only or """
+    """--default-resources 'cpu=(0.5,1);mem=(2Gi,4Gi);gpu=1' to specify both requests and limits""",
+)
 @click.option(
     "--skip-errors",
     "--skip-error",
@@ -161,6 +171,7 @@ def register(
     dry_run: bool,
     activate_launchplans: bool,
     env: typing.Optional[typing.Dict[str, str]],
+    default_resources: typing.Optional[ResourceSpec],
     skip_errors: bool,
 ):
     """
@@ -225,6 +236,7 @@ def register(
         package_or_module=package_or_module,
         remote=remote,
         env=env,
+        default_resources=default_resources,
         dry_run=dry_run,
         activate_launchplans=activate_launchplans,
         skip_errors=skip_errors,

flytekit/clis/sdk_in_container/run.py

@@ -43,6 +43,7 @@
 from flytekit.core.artifact import ArtifactQuery
 from flytekit.core.base_task import PythonTask
 from flytekit.core.data_persistence import FileAccessProvider
+from flytekit.core.resources import ResourceSpec
 from flytekit.core.type_engine import TypeEngine
 from flytekit.core.workflow import PythonFunctionWorkflow, WorkflowBase
 from flytekit.exceptions.system import FlyteSystemException
@@ -51,6 +52,7 @@
     FlyteLiteralConverter,
     key_value_callback,
     labels_callback,
+    resource_spec_callback,
 )
 from flytekit.interaction.string_literals import literal_string_repr
 from flytekit.loggers import logger
@@ -197,6 +199,18 @@ class RunLevelParams(PyFlyteParams):
             help="Environment variables to set in the container, of the format `ENV_NAME=ENV_VALUE`",
         )
     )
+    default_resources: typing.Optional[ResourceSpec] = make_click_option_field(
+        click.Option(
+            param_decls=["--default-resources"],
+            required=False,
+            show_default=True,
+            type=str,
+            callback=resource_spec_callback,
+            help="During fast registration, will override default task resource requests and limits for tasks that have no statically defined resource request and limit. "
+            """Example usage: --default-resources 'cpu=1;mem=2Gi;gpu=1' for requests only or """
+            """--default-resources 'cpu=(0.5,1);mem=(2Gi,4Gi);gpu=1' to specify both requests and limits""",
+        )
+    )
     tags: typing.List[str] = make_click_option_field(
         click.Option(
             param_decls=["--tags", "--tag"],
@@ -745,6 +759,7 @@ def _run(*args, **kwargs):
                     source_path=run_level_params.computed_params.project_root,
                     module_name=run_level_params.computed_params.module,
                     fast_package_options=fast_package_options,
+                    default_resources=run_level_params.default_resources,
                 )
 
                 run_remote(

flytekit/configuration/__init__.py

@@ -148,6 +148,7 @@
 from flytekit.configuration import internal as _internal
 from flytekit.configuration.default_images import DefaultImages
 from flytekit.configuration.file import ConfigEntry, ConfigFile, get_config_file, read_file_if_exists, set_if_exists
+from flytekit.core.resources import ResourceSpec
 from flytekit.image_spec import ImageSpec
 from flytekit.image_spec.image_spec import ImageBuildEngine
 from flytekit.loggers import logger
@@ -824,6 +825,8 @@ class SerializationSettings(DataClassJsonMixin):
         version (str): The version (if any) with which to register entities under.
         image_config (ImageConfig): The image config used to define task container images.
         env (Optional[Dict[str, str]]): Environment variables injected into task container definitions.
+        default_resources (Optional[ResourceSpec]): The resources to request for the task - this is useful
+            if users need to override the default resource spec of an entity at registration time.
         flytekit_virtualenv_root (Optional[str]):  During out of container serialize the absolute path of the flytekit
             virtualenv at serialization time won't match the in-container value at execution time. This optional value
             is used to provide the in-container virtualenv path
@@ -842,6 +845,7 @@ class SerializationSettings(DataClassJsonMixin):
     domain: typing.Optional[str] = None
     version: typing.Optional[str] = None
     env: Optional[Dict[str, str]] = None
+    default_resources: Optional[ResourceSpec] = None
     git_repo: Optional[str] = None
     python_interpreter: str = DEFAULT_RUNTIME_PYTHON_INTERPRETER
     flytekit_virtualenv_root: Optional[str] = None
@@ -916,6 +920,7 @@ def new_builder(self) -> Builder:
             version=self.version,
             image_config=self.image_config,
             env=self.env.copy() if self.env else None,
+            default_resources=self.default_resources,
             git_repo=self.git_repo,
             flytekit_virtualenv_root=self.flytekit_virtualenv_root,
             python_interpreter=self.python_interpreter,
@@ -967,6 +972,7 @@ class Builder(object):
         version: str
         image_config: ImageConfig
         env: Optional[Dict[str, str]] = None
+        default_resources: Optional[ResourceSpec] = None
         git_repo: Optional[str] = None
         flytekit_virtualenv_root: Optional[str] = None
         python_interpreter: Optional[str] = None
@@ -984,6 +990,7 @@ def build(self) -> SerializationSettings:
                 version=self.version,
                 image_config=self.image_config,
                 env=self.env,
+                default_resources=self.default_resources,
                 git_repo=self.git_repo,
                 flytekit_virtualenv_root=self.flytekit_virtualenv_root,
                 python_interpreter=self.python_interpreter,

flytekit/core/python_auto_container.py

@@ -231,6 +231,17 @@ def _get_container(self, settings: SerializationSettings) -> _task_model.Contain
         for elem in (settings.env, self.environment):
             if elem:
                 env.update(elem)
+        # Override the task's resource spec if it was not set statically in the task definition
+
+        def _resources_unspecified(resources: ResourceSpec) -> bool:
+            return resources == ResourceSpec(
+                requests=Resources(),
+                limits=Resources(),
+            )
+
+        if isinstance(settings.default_resources, ResourceSpec) and _resources_unspecified(self.resources):
+            self._resources = settings.default_resources
+
         return _get_container_definition(
             image=self.get_image(settings),
             resource_spec=self.resources,

flytekit/interaction/click_types.py

@@ -20,6 +20,7 @@
 from flytekit import BlobType, FlyteContext, Literal, LiteralType, StructuredDataset
 from flytekit.core.artifact import ArtifactQuery
 from flytekit.core.data_persistence import FileAccessProvider
+from flytekit.core.resources import Resources, ResourceSpec
 from flytekit.core.type_engine import TypeEngine
 from flytekit.models.types import SimpleType
 from flytekit.remote.remote_fs import FlytePathResolver
@@ -81,6 +82,40 @@ def labels_callback(_: typing.Any, param: str, values: typing.List[str]) -> typi
     return result
 
 
+def resource_spec_callback(_: typing.Any, param: str, value: typing.Optional[str]) -> typing.Optional[ResourceSpec]:
+    """
+    Callback for click to parse a resource spec.
+    """
+    if not value:
+        return None
+
+    def _extract_pair(s: str) -> typing.Optional[typing.Tuple[str, str]]:
+        """Can extract the pair of values "0.5" and "1" from the string '(0.5,1)'"""
+        vals = s.strip("() ").split(",")
+        if len(vals) != 2:
+            return None
+        return vals[0].strip(), vals[1].strip()
+
+    items = value.split(";")
+    _allowed_keys = Resources.__annotations__.keys()
+    result = {}
+    for item in items:
+        kv_split = item.split("=")
+        if len(kv_split) != 2:
+            raise click.BadParameter(
+                f"Expected semicolon separated key-value pairs of the form 'key1=value1;key2=value2;...', got '{item}'"
+            )
+        k = kv_split[0].strip()
+        v = kv_split[1].strip()
+        if k not in _allowed_keys:
+            raise click.BadParameter(f"Expected key to be one of {list(_allowed_keys)}, got '{k}'")
+        if k in result:
+            raise click.BadParameter(f"Expected unique keys {list(_allowed_keys)}, got '{k}' multiple times")
+        result[k.strip()] = _extract_pair(v) or v
+
+    return ResourceSpec.from_multiple_resource(Resources(**result))
+
+
 class DirParamType(click.ParamType):
     name = "directory path"
 

flytekit/remote/remote.py

@@ -56,6 +56,7 @@
 )
 from flytekit.core.python_function_task import PythonFunctionTask
 from flytekit.core.reference_entity import ReferenceEntity, ReferenceSpec
+from flytekit.core.resources import ResourceSpec
 from flytekit.core.task import ReferenceTask
 from flytekit.core.tracker import extract_task_module
 from flytekit.core.type_engine import LiteralsResolver, TypeEngine, strict_type_hint_matching
@@ -1310,6 +1311,7 @@ def register_script(
         source_path: typing.Optional[str] = None,
         module_name: typing.Optional[str] = None,
         envs: typing.Optional[typing.Dict[str, str]] = None,
+        default_resources: typing.Optional[ResourceSpec] = None,
         fast_package_options: typing.Optional[FastPackageOptions] = None,
     ) -> typing.Union[FlyteWorkflow, FlyteTask, FlyteLaunchPlan, ReferenceEntity]:
         """
@@ -1326,6 +1328,7 @@ def register_script(
         :param source_path: The root of the project path
         :param module_name: the name of the module
         :param envs: Environment variables to be passed to the serialization
+        :param default_resources: Default resources to be passed to the serialization. These override the resource spec for any tasks that have no statically defined resource requests and limits.
         :param fast_package_options: Options to customize copy_all behavior, ignored when copy_all is False.
         :return:
         """
@@ -1364,6 +1367,7 @@ def register_script(
             image_config=image_config,
             git_repo=_get_git_repo_url(source_path),
             env=envs,
+            default_resources=default_resources,
             fast_serialization_settings=FastSerializationSettings(
                 enabled=True,
                 destination_dir=destination_dir,

flytekit/tools/repo.py

@@ -12,6 +12,7 @@
 from flytekit.configuration import FastSerializationSettings, ImageConfig, SerializationSettings
 from flytekit.constants import CopyFileDetection
 from flytekit.core.context_manager import FlyteContextManager
+from flytekit.core.resources import ResourceSpec
 from flytekit.loggers import logger
 from flytekit.models import launch_plan, task
 from flytekit.models.core.identifier import Identifier
@@ -251,6 +252,7 @@ def register(
     remote: FlyteRemote,
     copy_style: CopyFileDetection,
     env: typing.Optional[typing.Dict[str, str]],
+    default_resources: typing.Optional[ResourceSpec],
     dry_run: bool = False,
     activate_launchplans: bool = False,
     skip_errors: bool = False,
@@ -273,6 +275,7 @@ def register(
         image_config=image_config,
         fast_serialization_settings=None,  # should probably add incomplete fast settings
         env=env,
+        default_resources=default_resources,
     )
 
     if not version and copy_style == CopyFileDetection.NO_COPY:

tests/flytekit/integration/remote/test_remote.py

@@ -17,12 +17,15 @@
 import pytest
 from unittest import mock
 from dataclasses import dataclass
+import random
+import string
 
 from flytekit import LaunchPlan, kwtypes, WorkflowExecutionPhase, task, workflow
 from flytekit.configuration import Config, ImageConfig, SerializationSettings
 from flytekit.core.launch_plan import reference_launch_plan
 from flytekit.core.task import reference_task
 from flytekit.core.workflow import reference_workflow
+from flytekit.models import task as task_models
 from flytekit.exceptions.user import FlyteAssertion, FlyteEntityNotExistException
 from flytekit.extras.sqlite3.task import SQLite3Config, SQLite3Task
 from flytekit.remote.remote import FlyteRemote
@@ -1170,3 +1173,98 @@ def test_register_wf_twice(register):
         ]
     )
     assert out.returncode == 0
+
+
+def test_register_wf_with_default_resources_override(register):
+    # Save the version here to retrieve the created task later
+    version = str(uuid.uuid4())
+    # Register the workflow with overridden default resources
+    out = subprocess.run(
+        [
+            "pyflyte",
+            "--verbose",
+            "-c",
+            CONFIG,
+            "register",
+            "--default-resources",
+            "cpu=1300m;mem=1100Mi",
+            "--image",
+            IMAGE,
+            "--project",
+            PROJECT,
+            "--domain",
+            DOMAIN,
+            "--version",
+            version,
+            MODULE_PATH / "hello_world.py",
+        ]
+    )
+    assert out.returncode == 0
+
+    # Retrieve the created task
+    remote = FlyteRemote(Config.auto(config_file=CONFIG), PROJECT, DOMAIN)
+    task = remote.fetch_task(name="basic.hello_world.say_hello", version=version)
+    assert task.template.container is not None
+    assert task.template.container.resources == task_models.Resources(
+        requests=[
+            task_models.Resources.ResourceEntry(
+                name=task_models.Resources.ResourceName.CPU,
+                value="1300m",
+            ),
+            task_models.Resources.ResourceEntry(
+                name=task_models.Resources.ResourceName.MEMORY,
+                value="1100Mi",
+            ),
+        ],
+        limits=[],
+    )
+
+
+def test_run_wf_with_default_resources_override(register):
+    # Save the execution id here to retrieve the created execution later
+    prefix = random.choice(string.ascii_lowercase)
+    short_random_part = uuid.uuid4().hex[:8]
+    execution_id = f"{prefix}{short_random_part}"
+    # Register the workflow with overridden default resources
+    out = subprocess.run(
+        [
+            "pyflyte",
+            "--verbose",
+            "-c",
+            CONFIG,
+            "run",
+            "--remote",
+            "--default-resources",
+            "cpu=500m;mem=1Gi",
+            "--project",
+            PROJECT,
+            "--domain",
+            DOMAIN,
+            "--name",
+            execution_id,
+            MODULE_PATH / "hello_world.py",
+            "my_wf"
+        ]
+    )
+    assert out.returncode == 0
+
+    # Retrieve the created task
+    remote = FlyteRemote(Config.auto(config_file=CONFIG), PROJECT, DOMAIN)
+    execution = remote.fetch_execution(name=execution_id)
+    execution = remote.wait(execution=execution)
+    version = execution.spec.launch_plan.version
+    task = remote.fetch_task(name="basic.hello_world.say_hello", version=version)
+    assert task.template.container is not None
+    assert task.template.container.resources == task_models.Resources(
+        requests=[
+            task_models.Resources.ResourceEntry(
+                name=task_models.Resources.ResourceName.CPU,
+                value="500m",
+            ),
+            task_models.Resources.ResourceEntry(
+                name=task_models.Resources.ResourceName.MEMORY,
+                value="1Gi",
+            ),
+        ],
+        limits=[],
+    )

tests/flytekit/unit/core/test_context_manager.py

@@ -19,6 +19,7 @@
 )
 from flytekit.core import mock_stats, context_manager
 from flytekit.core.context_manager import ExecutionParameters, FlyteContext, FlyteContextManager, SecretsManager
+from flytekit.core.resources import ResourceSpec, Resources
 from flytekit.models.core import identifier as id_models
 
 
@@ -301,6 +302,7 @@ def test_serialization_settings_transport():
         domain="domain",
         version="version",
         env={"hello": "blah"},
+        default_resources=ResourceSpec(requests=Resources(cpu="1", mem="2Gi"), limits=Resources(cpu="1", mem="2Gi")),
         image_config=ImageConfig(
             default_image=default_img,
             images=[default_img],
@@ -322,7 +324,7 @@ def test_serialization_settings_transport():
     ss = SerializationSettings.from_transport(tp)
     assert ss is not None
     assert ss == serialization_settings
-    assert len(tp) == 408
+    assert len(tp) == 480
 
 
 def test_exec_params():