remove ray task post_execute because ray.shutdown will be automatically called

[yundai424]

Tracking issue

https://github.com/flyteorg/flyte/issues/

Why are the changes needed?

ray.shutdown() will always be automatically called (registered through atexit here). We should rather rely on it instead of explicitly calling it again especially we don't set _exiting_interpreter=True which may result in disconnecting too soon before worker logs are flushed to driver.

What changes were proposed in this pull request?

remove the post_execute of ray task

How was this patch tested?

test with a simple example:

import typing

from flytekit import ImageSpec, Resources, task, workflow

commit_hash = "244a85c876eab507742d72bd90916d0e8662ea5b"
ray_plugin = f"git+https://github.com/yundai424/flytekit.git@{commit_hash}#subdirectory=plugins/flytekit-ray"

custom_image = ImageSpec(
    name="ray-flyte-plugin",
    registry="localhost:30000",
    packages=[ray_plugin],
    apt_packages=['git', 'wget', 'vim'],
)

# if custom_image.is_container():
import ray
from flytekitplugins.ray import HeadNodeConfig, RayJobConfig, WorkerNodeConfig

@ray.remote
def f(x):
    print(f"ray task get {x}")
    return x * x

@task(
    task_config=RayJobConfig(
        head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True", "block": "true"}),
        worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=1)],
        runtime_env={"pip": []},
    ),
    requests=Resources(mem="4Gi", cpu="1"),
    container_image=custom_image,
)
def ray_task(n: int) -> typing.List[int]:
    futures = [f.remote(i) for i in range(n)]
    l = ray.get(futures)
    print(l, 2)
    return l

@workflow
def ray_workflow(n: int) -> typing.List[int]:
    return ray_task(n=n)

Setup process

Screenshots

Check all the applicable boxes

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

Related PRs

Docs link

[welcome]

Thank you for opening this pull request! 🙌

These tips will help get your PR across the finish line:

• Most of the repos have a PR template; if not, fill it out to the best of your knowledge.
• Sign off your commits (Reference: DCO Guide).