Repository: modin-project/modin

Issue:
BUG: Benchmark Mode is evaluated only once when each function is defined
### Modin version checks

- [X] I have checked that this issue has not already been reported.

- [X] I have confirmed this bug exists on the latest released version of Modin.

- [X] I have confirmed this bug exists on the main branch of Modin. (In order to do this you can follow [this guide](https://modin.readthedocs.io/en/stable/getting_started/installation.html#installing-from-the-github-master-branch).)


### Reproducible Example

```python
import time
import modin.pandas as pd
from modin.config import BenchmarkMode

df = pd.DataFrame(list(range(32)))
BenchmarkMode.put(True)
%time result = df.applymap(lambda col: time.sleep(0.3))
%time print(result)
```


### Issue Description

wait_computations_if_benchmark_mode decides [when each function is defined](https://github.com/modin-project/modin/blob/2660aa33b79c79af828bce3b8f5c6573bc4d7c56/modin/core/dataframe/pandas/partitioning/partition_manager.py#L54) whether benchmark mode is on. Consequently, it's not possible to set benchmark mode after the methods have been defined. In the example above, turning benchmark mode on has no effect, even though the user turns benchmark mode on before the `applymap`. The `applymap` finishes in 2.39 ms because it's still async, and then `print` finishes in 13 seconds because it blocks on the `applymap`.

### Expected Behavior

Users should be able to set benchmark mode on or off at any point in the main process.

### Error Logs

N/A

### Installed Versions

<details>

INSTALLED VERSIONS
------------------
commit           : f727c045729b969de0e4a876e14209cde77166b9
python           : 3.10.4.final.0
python-bits      : 64
OS               : Darwin
OS-release       : 21.5.0
Version          : Darwin Kernel Version 21.5.0: Tue Apr 26 21:08:22 PDT 2022; root:xnu-8020.121.3~4/RELEASE_X86_64
machine          : x86_64
processor        : i386
byteorder        : little
LC_ALL           : None
LANG             : en_US.UTF-8
LOCALE           : en_US.UTF-8

Modin dependencies
------------------
modin            : 0.15.0+134.gf727c0457
ray              : 2.0.0
dask             : 2022.7.1
distributed      : 2022.7.1
omniscidbe       : None

pandas dependencies
-------------------
pandas           : 1.4.4
numpy            : 1.23.2
pytz             : 2022.2.1
dateutil         : 2.8.2
setuptools       : 61.2.0
pip              : 22.2.2
Cython           : None
pytest           : 7.1.2
hypothesis       : None
sphinx           : 5.1.1
blosc            : None
feather          : 0.4.1
xlsxwriter       : None
lxml.etree       : 4.9.1
html5lib         : None
pymysql          : None
psycopg2         : 2.9.3
jinja2           : 3.1.2
IPython          : 8.4.0
pandas_datareader: None
bs4              : 4.11.1
bottleneck       : None
brotli           : 1.0.9
fastparquet      : 0.8.1
fsspec           : 2022.7.1
gcsfs            : None
markupsafe       : 2.1.1
matplotlib       : 3.5.2
numba            : None
numexpr          : 2.8.3
odfpy            : None
openpyxl         : 3.0.10
pandas_gbq       : 0.17.7
pyarrow          : 8.0.0
pyreadstat       : None
pyxlsb           : None
s3fs             : 2022.7.1
scipy            : 1.9.0
snappy           : None
sqlalchemy       : 1.4.39
tables           : 3.7.0
tabulate         : None
xarray           : 2022.6.0
xlrd             : 2.0.1
xlwt             : None
zstandard        : None

</details>



Discussion:

Resolving patch:
diff --git a/modin/core/dataframe/pandas/partitioning/partition_manager.py b/modin/core/dataframe/pandas/partitioning/partition_manager.py
--- a/modin/core/dataframe/pandas/partitioning/partition_manager.py
+++ b/modin/core/dataframe/pandas/partitioning/partition_manager.py
@@ -51,12 +51,12 @@ def wait_computations_if_benchmark_mode(func):
     -----
     `func` should return NumPy array with partitions.
     """
-    if BenchmarkMode.get():
 
-        @wraps(func)
-        def wait(cls, *args, **kwargs):
-            """Wait for computation results."""
-            result = func(cls, *args, **kwargs)
+    @wraps(func)
+    def wait(cls, *args, **kwargs):
+        """Wait for computation results."""
+        result = func(cls, *args, **kwargs)
+        if BenchmarkMode.get():
             if isinstance(result, tuple):
                 partitions = result[0]
             else:
@@ -70,13 +70,10 @@ def wait(cls, *args, **kwargs):
             [part.drain_call_queue() for part in partitions.flatten()]
             # The partition manager invokes the relevant .wait() method under
             # the hood, which should wait in parallel for all computations to finish
-            # (We can't just add a `cls` argument to this `wait` function, since doing so
-            # seems to be incompatible with the way function decorators work)
             cls.wait_partitions(partitions.flatten())
-            return result
+        return result
 
-        return wait
-    return func
+    return wait
 
 
 class PandasDataframePartitionManager(ClassLogger, ABC):


Test patch:
diff --git a/modin/pandas/test/internals/test_benchmark_mode.py b/modin/pandas/test/internals/test_benchmark_mode.py
--- a/modin/pandas/test/internals/test_benchmark_mode.py
+++ b/modin/pandas/test/internals/test_benchmark_mode.py
@@ -11,21 +11,55 @@
 # ANY KIND, either express or implied. See the License for the specific language
 # governing permissions and limitations under the License.
 
-# test BenchmarkMode == True
+import unittest.mock as mock
 
-from contextlib import nullcontext
 import modin.pandas as pd
 from modin.pandas.test.utils import test_data_values
-from modin.config import BenchmarkMode, StorageFormat
-from modin.test.test_utils import warns_that_defaulting_to_pandas
+from modin.config import BenchmarkMode, Engine
 
 
-def test_syncronous_mode():
+engine = Engine.get()
+
+# We have to explicitly mock subclass implementations of wait_partitions.
+if engine == "Ray":
+    wait_method = (
+        "modin.core.execution.ray.implementations.pandas_on_ray."
+        + "partitioning.partition_manager."
+        + "PandasOnRayDataframePartitionManager.wait_partitions"
+    )
+elif engine == "Dask":
+    wait_method = (
+        "modin.core.execution.dask.implementations."
+        + "pandas_on_dask.partitioning.partition_manager."
+        + "PandasOnDaskDataframePartitionManager.wait_partitions"
+    )
+else:
+    wait_method = (
+        "modin.core.dataframe.pandas.partitioning."
+        + "partition_manager.PandasDataframePartitionManager.wait_partitions"
+    )
+
+
+def test_from_environment_variable():
     assert BenchmarkMode.get()
-    # On Omnisci storage, transpose() defaults to Pandas.
-    with (
-        warns_that_defaulting_to_pandas()
-        if StorageFormat.get() == "Omnisci"
-        else nullcontext()
-    ):
+    with mock.patch(wait_method) as wait:
         pd.DataFrame(test_data_values[0]).mean()
+
+    wait.assert_called()
+
+
+def test_turn_off():
+    df = pd.DataFrame([0])
+    BenchmarkMode.put(False)
+    with mock.patch(wait_method) as wait:
+        df.dropna()
+    wait.assert_not_called()
+
+
+def test_turn_on():
+    BenchmarkMode.put(False)
+    df = pd.DataFrame([0])
+    BenchmarkMode.put(True)
+    with mock.patch(wait_method) as wait:
+        df.dropna()
+    wait.assert_called()