[BUG] Task errors are not directly surfaced in unit tests in 1.12.0

[charliemoriarty]

Describe the bug

Since upgrading flyte from 1.11.0 -> 1.12.0, unit tests to verify that tasks raise expected errors are no longer passing. Instead of surfacing the error, the tests fail with a SystemExit.

from flytekit import task

@task
def f(x: int) -> int:
    if x > 0:
        return x + 2
    else:
        raise ValueError("x must be positive")
    

# Passes in 1.11.0, fails in 1.12.0
def test_f_raises():
    with pytest.raises(ValueError):
        f(-1)

We've been able to bypass this issue by accessing the task_function value directly in the tests, but this is non-intuitive and required a lot of manual updates to our test suite:

# Passes in 1.12.0
def test_f_raises():
    with pytest.raises(ValueError):
        _f = f.task_function
        _f(-1)

Slack thread: https://flyte-org.slack.com/archives/CP2HDHKE1/p1716391591793619

Expected behavior

I would expect to be able to test that tasks raise errors directly, without needing to access the internal task function:

# Should pass
def test_f_raises():
    with pytest.raises(ValueError):
        f(-1)

Additional context to reproduce

No response

Screenshots

No response

Are you sure this issue hasn't been raised already?

• Yes

Have you read the Code of Conduct?

• Yes

[runllm]

The issue you're experiencing seems to be related to changes in how Flyte handles task execution and error propagation between versions 1.11.0 and 1.12.0. This behavior might be due to internal changes in Flytekit's exception handling mechanism.

To better understand and address this issue, could you please provide some additional information?

1. Could you share the full stack trace of the SystemExit error you're encountering?
2. Are there any other relevant changes in your environment or dependencies that coincided with the Flyte upgrade?
3. Have you tried downgrading to Flytekit 1.11.0 to confirm that the issue is indeed related to the version change?

This information will help us understand the problem more clearly and find a solution.

Sources:

• Flytekit cant raise exception for Pydantic v2 ValidationError
• ShellTask error is recognize as system failure

---

This is an AI-generated response and your feedback is appreciated! Please leave a 👍 if this is helpful and 👎 if it is not.

Mention @RunLLM to ask follow-up questions. See runllm.com for more!

[wild-endeavor]

I'm not able to repro actually, this test is passing for me.

@task
def f(x: int) -> int:
    if x > 0:
        return x + 2
    else:
        raise ValueError("x must be positive")

def test_xyz():
    with pytest.raises(ValueError):
        f(x=-1)

I did change the function call to add x= since flytekit still doesn't support positional args (but that's a separate issue we're hoping to address soon).

───────────────────── Traceback (most recent call last) ──────────────────────╮
│ /Users/ytong/go/src/github.com/flyteorg/flytekit/tests/flytekit/unit/core/te │
│ st_array_node_map_task.py:355 in f                                           │
│                                                                              │
│   352 │   if x > 0:                                                          │
│   353 │   │   return x + 2                                                   │
│   354 │   else:                                                              │
│ ❱ 355 │   │   raise ValueError("x must be positive")                         │
│   356                                                                        │
│   357 def test_xyz():                                                        │
│   358                                                                        │
╰──────────────────────────────────────────────────────────────────────────────╯
ValueError: x must be positive
============================== 1 passed in 0.15s ===============================

[eapolinario]

flyteorg/flytekit#2358 is the culprit. If you want to unblock immediately, in your tests, set the environment variable FLYTE_EXIT_ON_USER_EXCEPTION to 0 prior to running the tests (similar to how we do in our unit tests).

That said, we plan to revisit that change, since expecting users to set that specific env var to get that behavior is not reasonable.