Function cached_remote_fn

python/ray/data/_internal/remote_fn.py:8–44

Lazily defines a ray.remote function. This is used in Datasets to avoid circular import issues with ray.remote. (ray imports ray.data in order to allow ``ray.data.read_foo()`` to work, which means ray.remote cannot be used top-level in ray.data). NOTE: Dynamic arguments should not

(fn: Any, **ray_remote_args)

def cached_remote_fn(fn: Any, **ray_remote_args) -> Any:
    """Lazily defines a ray.remote function.

    This is used in Datasets to avoid circular import issues with ray.remote.
    (ray imports ray.data in order to allow ``ray.data.read_foo()`` to work,
    which means ray.remote cannot be used top-level in ray.data).

    NOTE: Dynamic arguments should not be passed in directly,
    and should be set with ``options`` instead:
    ``cached_remote_fn(fn, **static_args).options(**dynamic_args)``.
    """

    # NOTE: Hash of the passed in arguments guarantees that we're caching
    #       complete instantiation of the Ray's remote method
    #
    # To compute the hash of passed in arguments and make sure it's deterministic
    #   - Sort all KV-pairs by the keys
    #   - Convert sorted list into tuple
    #   - Compute hash of the resulting tuple
    hashable_args = _make_hashable(ray_remote_args)
    args_hash = hash(hashable_args)

    if (fn, args_hash) not in CACHED_FUNCTIONS:
        default_ray_remote_args = {
            # Use the default scheduling strategy for all tasks so that we will
            # not inherit a placement group from the caller, if there is one.
            # The caller of this function may override the scheduling strategy
            # as needed.
            "scheduling_strategy": "DEFAULT",
            "max_retries": -1,
        }
        ray_remote_args = {**default_ray_remote_args, **ray_remote_args}
        _add_system_error_to_retry_exceptions(ray_remote_args)

        CACHED_FUNCTIONS[(fn, args_hash)] = ray.remote(**ray_remote_args)(fn)

    return CACHED_FUNCTIONS[(fn, args_hash)]


def _make_hashable(obj):