modin-project · data-makerman · Jul 19, 2024 · Jul 22, 2024 · Jul 29, 2024 · Jul 29, 2024
@@ -19,6 +19,25 @@
 from dask.distributed import wait
 from distributed import Future
 from distributed.client import default_client
+from distributed.worker import get_worker
+
+
+def get_dask_client():
+    """
+    Get the Dask client, reusing the worker's client if execution is on a Dask worker.
+
+    Returns
+    -------
+    distributed.Client
+        The Dask client.
+    """
+    try:
+        client = default_client()
+    except ValueError:
+        # We ought to be in a worker process
+        worker = get_worker()
+        client = worker.client
+    return client
 
 
 def _deploy_dask_func(func, *args, return_pandas_df=None, **kwargs):  # pragma: no cover
@@ -83,7 +102,7 @@
         list
             The result of ``func`` split into parts in accordance with ``num_returns``.
         """
-        client = default_client()
+        client = get_dask_client()
         args = [] if f_args is None else f_args
         kwargs = {} if f_kwargs is None else f_kwargs
         if callable(func):
@@ -137,7 +156,7 @@
         Any
             An object(s) from the distributed memory.
         """
-        client = default_client()
+        client = get_dask_client()
         return client.gather(future)
 
     @classmethod
@@ -164,7 +183,7 @@
             # {'sep': <Future: finished, type: pandas._libs.lib._NoDefault, key: sep>, \
             #  'delimiter': <Future: finished, type: NoneType, key: delimiter> ...
             data = UserDict(data)
-        client = default_client()
+        client = get_dask_client()
         return client.scatter(data, **kwargs)
 
     @classmethod

@@ -30,6 +30,17 @@
 def initialize_dask():
     """Initialize Dask environment."""
     from distributed.client import default_client
+    from distributed.worker import get_worker
+
+    try:
 class FactoryDispatcher(object): 
     """ 
     Class that routes IO-work to the factories. 
     This class is responsible for keeping selected factory up-to-date and dispatching 
     calls of IO-functions to its actual execution-specific implementations. 
     """ 
     __factory: factories.BaseFactory = None 
     @classmethod 
     def get_factory(cls) -> factories.BaseFactory: 
         """Get current factory.""" 
         if cls.__factory is None: 
             from modin.pandas import _update_engine 
             Engine.subscribe(_update_engine) 
             Engine.subscribe(cls._update_factory) 
             StorageFormat.subscribe(cls._update_factory) 
         return cls.__factory 
 if not ray.is_initialized() or override_is_cluster: 
 class FactoryDispatcher(object): 
     """ 
     Class that routes IO-work to the factories. 
  
     This class is responsible for keeping selected factory up-to-date and dispatching 
     calls of IO-functions to its actual execution-specific implementations. 
     """ 
  
     __factory: factories.BaseFactory = None 
  
     @classmethod 
     def get_factory(cls) -> factories.BaseFactory: 
         """Get current factory.""" 
         if cls.__factory is None: 
             from modin.pandas import _update_engine 
  
             Engine.subscribe(_update_engine) 
             Engine.subscribe(cls._update_factory) 
             StorageFormat.subscribe(cls._update_factory) 
         return cls.__factory 
 if not ray.is_initialized() or override_is_cluster: 
+        # Check if running within a Dask worker process
+        get_worker()
+        # If the above line does not raise an error, we are in a worker process
+        # and should not create a new client
+        return
+    except ValueError:
+        # Not in a Dask worker, proceed to check for or create a client
+        pass
 
     try:
         client = default_client()

@@ -2784,3 +2784,32 @@
         match="<function DataFrame.<property fget:_test_default_property>> is not currently supported",
     ):
         pd.DataFrame([[1]]).dataframe_test_default_property
+
+
+@pytest.mark.parametrize(
+    "modify_config",
+    [
+        {Engine: "Ray"},
+        {Engine: "Dask"},
+    ],
+    indirect=True,
+)
+def test_daemonic_worker_protection():
+    # Test for issue #7346, wherein some operations on Dask cause a second submission of a task to
+    # the Dask client from the worker scope, which should not cause a new client to be created
+
+    def submission_triggering_row_operation(row):
+        row_to_dict = row.to_dict()
+        dict_to_row = pd.Series(row_to_dict)
+        return dict_to_row
+
+    df = pd.DataFrame(
+        {
+            "A": ["a", "b", "c", "d"],
+            "B": [1, 2, 3, 4],
+            "C": [1, 2, 3, 4],
+            "D": [1, 2, 3, 4],
+        }
+    )
+
+    df = df.apply(submission_triggering_row_operation, axis=1)