Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

don't override timeout on with_overrides if not specified #3097

Open
wants to merge 5 commits into
base: master
Choose a base branch
from
Open

don't override timeout on with_overrides if not specified #3097

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1fae040
don't override timeout on with_overrides if not specified
pvditt Jan 27, 2025
f205c41
lint
pvditt Jan 27, 2025
4276012
clean up
pvditt Feb 6, 2025
1e6bb36
clean up
pvditt Feb 6, 2025
f43f286
clean up
pvditt Feb 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 14 additions & 10 deletions flytekit/core/node.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@
ID, which from the registration step
"""

TIMEOUT_OVERRIDE_SENTINEL = object()

def __init__(
self,
id: str,
Expand Down Expand Up @@ -127,7 +129,7 @@
def _override_node_metadata(
self,
name,
timeout: Optional[Union[int, datetime.timedelta]] = None,
timeout: Optional[Union[int, datetime.timedelta, object]] = TIMEOUT_OVERRIDE_SENTINEL,
retries: Optional[int] = None,
interruptible: typing.Optional[bool] = None,
cache: typing.Optional[bool] = None,
Expand All @@ -142,14 +144,16 @@
else:
node_metadata = self._metadata

if timeout is None:
node_metadata._timeout = datetime.timedelta()
elif isinstance(timeout, int):
node_metadata._timeout = datetime.timedelta(seconds=timeout)
elif isinstance(timeout, datetime.timedelta):
node_metadata._timeout = timeout
else:
raise ValueError("timeout should be duration represented as either a datetime.timedelta or int seconds")
if timeout is not Node.TIMEOUT_OVERRIDE_SENTINEL:

Check warning on line 147 in flytekit/core/node.py

View check run for this annotation

Codecov / codecov/patch

flytekit/core/node.py#L147 

Added line #L147 was not covered by tests
if timeout is None:
node_metadata._timeout = 0
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider using timedelta for timeout value

Consider using datetime.timedelta() instead of 0 for timeout value to maintain consistency with the type system and other code paths.

Code suggestion
Check the AI-generated fix before applying
Suggested change
node_metadata._timeout = 0
node_metadata._timeout = datetime.timedelta()

Code Review Run #01a35d

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

  • it was incorrectly flagged

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@pvditt For typing, I think this needs to be a timedelta object.

elif isinstance(timeout, int):
node_metadata._timeout = datetime.timedelta(seconds=timeout)

Check warning on line 151 in flytekit/core/node.py

View check run for this annotation

Codecov / codecov/patch

flytekit/core/node.py#L151 

Added line #L151 was not covered by tests
elif isinstance(timeout, datetime.timedelta):
node_metadata._timeout = timeout

Check warning on line 153 in flytekit/core/node.py

View check run for this annotation

Codecov / codecov/patch

flytekit/core/node.py#L153 

Added line #L153 was not covered by tests
else:
raise ValueError("timeout should be duration represented as either a datetime.timedelta or int seconds")

Check warning on line 155 in flytekit/core/node.py

View check run for this annotation

Codecov / codecov/patch

flytekit/core/node.py#L155 

Added line #L155 was not covered by tests

if retries is not None:
assert_not_promise(retries, "retries")
node_metadata._retries = (
Expand Down Expand Up @@ -181,7 +185,7 @@
aliases: Optional[Dict[str, str]] = None,
requests: Optional[Resources] = None,
limits: Optional[Resources] = None,
timeout: Optional[Union[int, datetime.timedelta]] = None,
timeout: Optional[Union[int, datetime.timedelta, object]] = TIMEOUT_OVERRIDE_SENTINEL,
retries: Optional[int] = None,
interruptible: Optional[bool] = None,
name: Optional[str] = None,
Expand Down
2 changes: 1 addition & 1 deletion flytekit/core/promise.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -586,7 +586,7 @@ def with_overrides(
aliases: Optional[Dict[str, str]] = None,
requests: Optional[Resources] = None,
limits: Optional[Resources] = None,
timeout: Optional[Union[int, datetime.timedelta]] = None,
timeout: Optional[Union[int, datetime.timedelta, object]] = Node.TIMEOUT_OVERRIDE_SENTINEL,
retries: Optional[int] = None,
interruptible: Optional[bool] = None,
name: Optional[str] = None,
Expand Down
39 changes: 33 additions & 6 deletions tests/flytekit/unit/core/test_node_creation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -302,18 +302,42 @@ def my_wf(a: typing.List[str]) -> typing.List[str]:
]


preset_timeout = datetime.timedelta(seconds=100)


@pytest.mark.parametrize(
"timeout,expected",
[(None, datetime.timedelta()), (10, datetime.timedelta(seconds=10))],
"timeout,t1_expected_timeout_overridden, t1_expected_timeout_unset, t2_expected_timeout_overridden, "
"t2_expected_timeout_unset",
[
(None, 0, 0, 0, preset_timeout),
(10, datetime.timedelta(seconds=10), 0,
datetime.timedelta(seconds=10), preset_timeout)
],
)
def test_timeout_override(timeout, expected):
def test_timeout_override(
timeout,
t1_expected_timeout_overridden,
t1_expected_timeout_unset,
t2_expected_timeout_overridden,
t2_expected_timeout_unset,
):
Comment on lines +319 to +323
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider more descriptive parameter names

Consider using more descriptive parameter names. The current names t1_expected_timeout_overridden, t1_expected_timeout_unset, etc. could be renamed to better indicate their purpose, such as task1_timeout_with_override and task1_timeout_without_override.

Code suggestion
Check the AI-generated fix before applying
Suggested change
t1_expected_timeout_overridden,
t1_expected_timeout_unset,
t2_expected_timeout_overridden,
t2_expected_timeout_unset,
):
task1_timeout_with_override,
task1_timeout_without_override,
task2_timeout_with_override,
task2_timeout_without_override,
):

Code Review Run #01a35d

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

  • it was incorrectly flagged

@task
def t1(a: str) -> str:
return f"*~*~*~{a}*~*~*~"

@task(
timeout=preset_timeout
)
def t2(a: str) -> str:
return f"*~*~*~{a}*~*~*~"

@workflow
def my_wf(a: str) -> str:
return t1(a=a).with_overrides(timeout=timeout)
s = t1(a=a).with_overrides(timeout=timeout)
s1 = t1(a=s).with_overrides()
s2 = t2(a=s1).with_overrides(timeout=timeout)
s3 = t2(a=s2).with_overrides()
return s3

serialization_settings = flytekit.configuration.SerializationSettings(
project="test_proj",
Expand All @@ -323,8 +347,11 @@ def my_wf(a: str) -> str:
env={},
)
wf_spec = get_serializable(OrderedDict(), serialization_settings, my_wf)
assert len(wf_spec.template.nodes) == 1
assert wf_spec.template.nodes[0].metadata.timeout == expected
assert len(wf_spec.template.nodes) == 4
assert wf_spec.template.nodes[0].metadata.timeout == t1_expected_timeout_overridden
assert wf_spec.template.nodes[1].metadata.timeout == t1_expected_timeout_unset
assert wf_spec.template.nodes[2].metadata.timeout == t2_expected_timeout_overridden
assert wf_spec.template.nodes[3].metadata.timeout == t2_expected_timeout_unset


def test_timeout_override_invalid_value():
Expand Down
Loading