Add request and limit to ray config #3087

pingsutw · 2025-01-24T23:49:35Z

Tracking issue

NA

Why are the changes needed?

Improve the devx: people can easily override the request and limit for the head/worker node

What changes were proposed in this pull request?

Replace K8sPod with PodTemplate to make it consistent with the parameter in @task
Add resource and limit to HeadNodeConfig and WorkerNodeConfig

How was this patch tested?

Setup process

import typing

from flytekit import ImageSpec, Resources, task, workflow
import ray
from flytekitplugins.ray import HeadNodeConfig, RayJobConfig, WorkerNodeConfig

flytekit_hash = "234603bde982c4ece390dd77c33390c6d8db93cf"
new_flytekitplugins_ray = f"git+https://github.com/flyteorg/flytekit.git@{flytekit_hash}#subdirectory=plugins/flytekit-ray"
new_flytekit = f"git+https://github.com/flyteorg/flytekit.git@{flytekit_hash}"


custom_image = ImageSpec(
    registry="pingsutw",
    packages=[new_flytekitplugins_ray, "kubernetes", new_flytekit],
    apt_packages=["wget", "git"],
)


@ray.remote
def f(x):
    return x * x


ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}, requests=Resources(cpu="1.7")),
    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2, requests=Resources(cpu="1.8"))],
    runtime_env={"pip": ["numpy", "pandas"]},
    enable_autoscaling=False,
    shutdown_after_job_finishes=True,
    ttl_seconds_after_finished=3,
)


@task(
    task_config=ray_config,
    requests=Resources(mem="2Gi", cpu="2"),
    container_image=custom_image,
)
def ray_task(n: int) -> typing.List[int]:
    futures = [f.remote(i) for i in range(n)]
    return ray.get(futures)


@workflow
def ray_workflow(n: int = 3):
    for i in range(1):
        ray_task(n=n)

Screenshots

Check all the applicable boxes

I updated the documentation accordingly.
All new and existing tests passed.
All commits are signed-off.

Related PRs

#2943

Docs link

NA

Summary by Bito

Enhanced Ray configuration system with resource management capabilities for head and worker nodes, replacing K8sPod with PodTemplate. Added support for CPU/memory specifications in HeadNodeConfig and WorkerNodeConfig, improved documentation quality, and introduced baseline documentation error tracking. The changes include support for pip extra arguments and improved URI modifications in dataclasses. Updated resource configuration implementation and flytekit version requirement with improved naming conventions.

Unit tests added: True

Estimated effort to review (1-5, lower is better): 4

Signed-off-by: Kevin Su <[email protected]>

flyte-bot · 2025-01-24T23:49:48Z

Code Review Agent Run #5148c9

Actionable Suggestions - 6

plugins/flytekit-ray/flytekitplugins/ray/task.py - 3
- Incomplete validation of resource specifications · Line 32-39
- Consider enhancing resource validation checks · Line 53-56
- Consider constant for worker container name · Line 123-123
plugins/flytekit-ray/tests/test_ray.py - 1
- Consider extracting pod template creation logic · Line 67-73
flytekit/core/resources.py - 2
- Consider making container name required · Line 105-106
- Consider visibility of pod spec function · Line 105-105

Review Details

Files reviewed - 4 · Commit Range: 234603b..234603b
- flytekit/core/resources.py
- flytekit/models/task.py
- plugins/flytekit-ray/flytekitplugins/ray/task.py
- plugins/flytekit-ray/tests/test_ray.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

Signed-off-by: Kevin Su <[email protected]>

flyte-bot · 2025-01-24T23:55:45Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
Feature Improvement - Enhanced Ray Resource Configuration	- `resources.py` - Refactored pod spec generation with improved type hints and naming - `task.py` - Added PodTemplate conversion methods and type checking - `task.py` - Added resource requests/limits support to Ray head and worker node configs - `setup.py` - Updated flytekit version requirement - `test_ray.py` - Added tests for new resource configuration options - `test_resources.py` - Updated resource tests with new naming conventions

plugins/flytekit-ray/flytekitplugins/ray/task.py

plugins/flytekit-ray/tests/test_ray.py

flytekit/core/resources.py

Signed-off-by: Kevin Su <[email protected]>

Sovietaced

I agree that the ux around the pod template is not great. Folks have to know the primary container name for the head/worker which is easy to miss and is the primary source of errors. I just think that folks will want these types of helpers for any future things they will need to configure. So in the future this code may need to be refactored to manipulate a pod template across several helpers.

Given that this doesn't change the backend I dont have too strong of opinions so LGTM.

flyte-bot · 2025-01-25T00:57:23Z

Code Review Agent Run #c9c4ac

Actionable Suggestions - 0

Review Details

Files reviewed - 4 · Commit Range: 234603b..18a5fef
- flytekit/core/resources.py
- plugins/flytekit-ray/flytekitplugins/ray/task.py
- plugins/flytekit-ray/tests/test_ray.py
- tests/flytekit/unit/core/test_resources.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

fiedlerNr9 · 2025-01-25T01:04:50Z

Love it! LGTM

Signed-off-by: Kevin Su <[email protected]>

eapolinario · 2025-01-25T01:23:10Z

#3090 is going to fix the lint CI check.

Change looks good!

flyte-bot · 2025-01-25T02:17:22Z

Code Review Agent Run #89224b

Actionable Suggestions - 1

plugins/flytekit-ray/setup.py - 1
- Consider adding upper bound to flytekit · Line 7-7

Review Details

Files reviewed - 2 · Commit Range: 18a5fef..49de42e
- plugins/flytekit-ray/setup.py
- tests/flytekit/unit/core/test_resources.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2025-01-25T02:20:18Z

plugins/flytekit-ray/setup.py

@@ -4,7 +4,7 @@

 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"

-plugin_requires = ["ray[default]", "flytekit>=1.3.0b2,<2.0.0", "flyteidl>=1.13.6"]
+plugin_requires = ["ray[default]", "flytekit>1.14.5", "flyteidl>=1.13.6"]


Consider adding upper bound to flytekit

Consider using a more specific version constraint for flytekit. The current constraint >1.14.5 has no upper bound which could lead to compatibility issues with future major versions. Consider using something like >1.14.5,<2.0.0 to ensure compatibility.

Code suggestion

Check the AI-generated fix before applying

Suggested change

plugin_requires = ["ray[default]", "flytekit>1.14.5", "flyteidl>=1.13.6"]

plugin_requires = ["ray[default]", "flytekit>1.14.5,<2.0.0", "flyteidl>=1.13.6"]

Code Review Run #89224b

Is this a valid issue, or was it incorrectly flagged by the Agent?

it was incorrectly flagged

Signed-off-by: Eduardo Apolinario <[email protected]>

codecov · 2025-01-27T17:26:51Z

Codecov Report

Attention: Patch coverage is 50.00000% with 5 lines in your changes missing coverage. Please review.

Project coverage is 51.69%. Comparing base (f0ba47f) to head (88d3c86).
Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
flytekit/models/task.py	57.14%	3 Missing ⚠️
flytekit/core/resources.py	33.33%	2 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (f0ba47f) and HEAD (88d3c86). Click for more details.

HEAD has 8 uploads less than BASE

Flag BASE (f0ba47f) HEAD (88d3c86)

9 1

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #3087       +/-   ##
===========================================
- Coverage   76.46%   51.69%   -24.78%     
===========================================
  Files         218      202       -16     
  Lines       22505    21398     -1107     
  Branches     2766     2765        -1     
===========================================
- Hits        17209    11062     -6147     
- Misses       4483     9731     +5248     
+ Partials      813      605      -208

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

flyte-bot · 2025-01-27T18:04:25Z

Code Review Agent Run #7ec23d

Actionable Suggestions - 0

Additional Suggestions - 2

flytekit/image_spec/image_spec.py - 1
- Consider adding pip args validation · Line 86-86
tests/flytekit/unit/extras/pydantic_transformer/test_pydantic_basemodel_transformer.py - 1
- Consider simplifying mock resolver function · Line 1011-1015

Review Details

Files reviewed - 12 · Commit Range: 49de42e..88d3c86
- .pre-commit-config.yaml
- flytekit/core/type_engine.py
- flytekit/image_spec/default_builder.py
- flytekit/image_spec/image_spec.py
- flytekit/remote/remote_fs.py
- flytekit/types/structured/structured_dataset.py
- pydoclint-errors-baseline.txt
- pyproject.toml
- tests/flytekit/unit/core/image_spec/test_default_builder.py
- tests/flytekit/unit/core/test_dataclass.py
- tests/flytekit/unit/extras/pydantic_transformer/test_pydantic_basemodel_transformer.py
- tests/flytekit/unit/types/structured_dataset/test_structured_dataset.py
Files skipped - 1
- .github/workflows/pythonbuild.yml - Reason: Filter setting
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

* Add request and limit to ray config Signed-off-by: Kevin Su <[email protected]> * lint Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * update flytekit version Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> --------- Signed-off-by: Kevin Su <[email protected]> Signed-off-by: Eduardo Apolinario <[email protected]> Co-authored-by: Eduardo Apolinario <[email protected]> Signed-off-by: Umer Ahmad <[email protected]>

Add request and limit to ray config

234603b

Signed-off-by: Kevin Su <[email protected]>

pingsutw requested review from wild-endeavor, kumare3, eapolinario, cosmicBboy, samhita-alla, thomasjpfan and Future-Outlier as code owners January 24, 2025 23:49

lint

6808e75

Signed-off-by: Kevin Su <[email protected]>