Skip to content

Commit

Permalink
[core] Lazy import of pydantics (#41475)
Browse files Browse the repository at this point in the history
Closes #41338

The hypothesis is that we are importing some additional things when initializing the python workers, and changing the import to lazy seems to fix the issue.

See https://buildkite.com/ray-project/release/builds/2492#018c1923-c5d0-4e09-a67e-b45cf2c3b553

Master: tasks_per_seconds = 430
This PR: tasks_per_seconds = 530
  • Loading branch information
rickyyx authored Nov 30, 2023
1 parent 31260fb commit cef4c1f
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 2 deletions.
20 changes: 19 additions & 1 deletion python/ray/tests/test_basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ def f():
ray.get(bar.remote(*[f() for _ in range(200)]))


def test_default_worker_import_dependency():
def test_default_worker_import_dependency(shutdown_only):
"""
Test ray's python worker import doesn't import the not-allowed dependencies.
"""
Expand All @@ -129,6 +129,11 @@ def test_default_worker_import_dependency():
# See https://github.com/ray-project/ray/issues/33891
blocked_deps = ["numpy"]

# Ray should not be importing pydantic (used in serialization) eagerly.
# This introduces regression in worker start up time.
# https://github.com/ray-project/ray/issues/41338
blocked_deps += ["pydantic"]

# Remove the ray module and the blocked deps from sys.modules.
sys.modules.pop("ray", None)
assert "ray" not in sys.modules
Expand All @@ -146,6 +151,19 @@ def test_default_worker_import_dependency():
for dep in blocked_deps:
assert dep not in sys.modules

# Test starting a ray workers should not see unwanted deps loaded eagerly.
ray.init()

@ray.remote
def f():
import ray # noqa: F401

assert "ray" in sys.modules
for x in blocked_deps:
assert x not in sys.modules

ray.get(f.remote())


# https://github.com/ray-project/ray/issues/7287
def test_omp_threads_set(ray_start_cluster, monkeypatch):
Expand Down
3 changes: 2 additions & 1 deletion python/ray/util/serialization_addons.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
import sys

from ray.util.annotations import DeveloperAPI
from ray._private.pydantic_compat import register_pydantic_serializers


@DeveloperAPI
Expand All @@ -27,6 +26,8 @@ def register_starlette_serializer(serialization_context):

@DeveloperAPI
def apply(serialization_context):
from ray._private.pydantic_compat import register_pydantic_serializers

register_pydantic_serializers(serialization_context)
register_starlette_serializer(serialization_context)

Expand Down

0 comments on commit cef4c1f

Please sign in to comment.