From 6cb00208f7ecc99dd6c38b79093bb65a901aca52 Mon Sep 17 00:00:00 2001 From: SangBin Cho Date: Thu, 8 Oct 2020 16:08:28 -0700 Subject: [PATCH] [Placement Group] Export bundle reservation check method only once. (#11153) * Export bundle reservation check method only once. * Addressed code review. --- python/ray/util/placement_group.py | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/python/ray/util/placement_group.py b/python/ray/util/placement_group.py index b0c346407..6e3d49189 100644 --- a/python/ray/util/placement_group.py +++ b/python/ray/util/placement_group.py @@ -5,6 +5,25 @@ from typing import (List, Dict, Optional) import ray from ray._raylet import PlacementGroupID, ObjectRef +bundle_reservation_check = None + + +# We need to import this method to use for ready API. +# But ray.remote is only available in runtime, and +# if we define this method inside ready method, this function is +# exported whenever ready is called, which can impact performance, +# https://github.com/ray-project/ray/issues/6240. +def _export_bundle_reservation_check_method_if_needed(): + global bundle_reservation_check + if bundle_reservation_check: + return + + @ray.remote(num_cpus=0, max_calls=0) + def bundle_reservation_check_func(placement_group): + return placement_group + + bundle_reservation_check = bundle_reservation_check_func + class PlacementGroup: """A handle to a placement group.""" @@ -33,9 +52,7 @@ class PlacementGroup: """ self._fill_bundle_cache_if_needed() - @ray.remote(num_cpus=0, max_calls=0) - def bundle_reservation_check(placement_group): - return placement_group + _export_bundle_reservation_check_method_if_needed() assert len(self.bundle_cache) != 0, ( "ready() cannot be called on placement group object with a "