[data] ray_tqdm does not work with numba

[rynewang]

What happened + What you expected to happen

numba is a JIT lib to compile python code to accelerate. It tries to type all functions and if it can't it complains (and fails). One dynamic function we have is print; numba somehow can handle it. However in ray.experimental.tqdm_ray we switched it to a modified print; and it breaks numba functions.

The tqdm is used by Ray Data, so if someone used numba funcs in a ray data api, it fails.

Workaround: set RAY_TQDM_PATCH_PRINT=0. But this corrupts progress bars. Maybe we can know how numba treats print and do the same thing to our print.

Versions / Dependencies

master

Reproduction script

import numba
import numpy as np

@numba.njit
def do_sum(A):
    acc = 0.
    for x in A:
        acc += np.sqrt(x)
        print(x)
    return acc

# works
print(do_sum(list(range(100))))

import ray
ray.init()
@ray.remote
def f():
    @numba.njit
    def do_sum(A):
        acc = 0.
        for x in A:
            acc += np.sqrt(x)
            print(x)
        return acc

    return do_sum(list(range(100)))

# works
print(ray.get(f.remote()))

@ray.remote
def f_tqdm():
    # used by ray data to modify builtin.print
    from ray.experimental import tqdm_ray
    tqdm_ray.instance()

    @numba.njit
    def do_sum(A):
        acc = 0.
        for x in A:
            acc += np.sqrt(x)
            print(x)
        return acc

    return do_sum(list(range(100)))

# error out
# print(ray.get(f_tqdm.remote()))

# works
print(ray.get(f_tqdm.options(runtime_env={"env_vars":{"RAY_TQDM_PATCH_PRINT":"0"}}).remote())) 

Issue Severity

Medium: It is a significant difficulty but I can work around it.

[rynewang]

@c21 you are the last one touched that file. Would you mind take a look? Thanks.