[Learn] Fix lightgbm machines (#3351)

* add worker ip meta for chunk meta

* fix expect worker

* fix CollectPorts

* fix test_collect_ports

* Optimize lightgbm train

* Fix meta setting

* Add get_node_ip_address

* Fix test_local_ranker

* Ignore ip when no lightgbm

* Lint

* add lightgbm tests running on ray

---------

Co-authored-by: 不涸 <[email protected]>

Author: zhongchun

.github/workflows/platform-ci.yml

@@ -108,6 +108,7 @@ jobs:
               pip install "xgboost_ray<0.1.14" "protobuf<4"
               # Ray Datasets need pyarrow>=6.0.1
               pip install "pyarrow>=6.0.1"
+              pip install lightgbm
             fi
             if [ -n "$RUN_DASK" ]; then
               pip install "dask[complete]" "mimesis<9.0.0" scikit-learn

mars/deploy/oscar/tests/test_ray.py

@@ -23,6 +23,7 @@
 
 from .... import tensor as mt
 from .... import dataframe as md
+from ....learn.contrib.lightgbm.tests import test_classifier
 from ....oscar.errors import ReconstructWorkerError
 from ....session import get_default_session, new_session
 from ....tests.core import require_ray, mock, DICT_NOT_EMPTY
@@ -335,3 +336,8 @@ def test_init_metrics_on_ray(ray_start_regular_shared, create_cluster):
     assert api._metric_backend == "ray"
 
     client.session.stop_server()
+
+
+@require_ray
+def test_lightgbm_classifier_on_ray(ray_start_regular_shared, create_cluster):
+    test_classifier.test_local_classifier(create_cluster)

mars/learn/contrib/lightgbm/_train.py

@@ -211,8 +211,19 @@ def __call__(self):
     @staticmethod
     def _get_data_chunks_workers(ctx, data):
         # data_chunk.inputs is concat, and concat's input is the co-allocated chunks
-        metas = ctx.get_chunks_meta([c.key for c in data.chunks], fields=["bands"])
-        return [m["bands"][0][0] for m in metas]
+        metas = ctx.get_chunks_meta(
+            [c.key for c in data.chunks], fields=["ip", "bands"]
+        )
+
+        ips = []
+        ip_to_worker = {}
+        for m in metas:
+            ip = m["ip"]
+            assert ip, "There is meta {meta} who doesn't contain ip."
+            ips.append(ip)
+            bands = m["bands"]
+            ip_to_worker[ip] = bands[0][0] if bands else None
+        return ips, ip_to_worker
 
     @staticmethod
     def _concat_chunks_by_worker(chunks, chunk_workers):
@@ -230,23 +241,24 @@ def tile(cls, op: "LGBMTrain"):
         data = op.data
         worker_to_args = defaultdict(dict)
 
-        workers = cls._get_data_chunks_workers(ctx, data)
+        # Note: Mars worker is band address, and LGBMTrain worker is machine ip.
+        ips, ip_to_worker = cls._get_data_chunks_workers(ctx, data)
 
         for arg in ["_data", "_label", "_sample_weight", "_init_score"]:
             if getattr(op, arg) is not None:
                 for worker, chunk in cls._concat_chunks_by_worker(
-                    getattr(op, arg).chunks, workers
+                    getattr(op, arg).chunks, ips
                 ).items():
                     worker_to_args[worker][arg] = chunk
 
         if op.eval_datas:
             eval_workers_list = [
-                cls._get_data_chunks_workers(ctx, d) for d in op.eval_datas
+                cls._get_data_chunks_workers(ctx, d)[0] for d in op.eval_datas
             ]
             extra_workers = reduce(
                 operator.or_, (set(w) for w in eval_workers_list)
-            ) - set(workers)
-            worker_remap = dict(zip(extra_workers, itertools.cycle(workers)))
+            ) - set(ips)
+            worker_remap = dict(zip(extra_workers, itertools.cycle(ips)))
             if worker_remap:
                 eval_workers_list = [
                     [worker_remap.get(w, w) for w in wl] for wl in eval_workers_list
@@ -270,10 +282,11 @@ def tile(cls, op: "LGBMTrain"):
                             worker_to_args[worker][arg].append(chunk)
 
         out_chunks = []
-        workers = list(set(workers))
-        for worker_id, worker in enumerate(workers):
+        ips = list(set(ips))
+        workers = list(ip_to_worker.values())
+        for worker_id, worker in enumerate(ips):
             chunk_op = op.copy().reset_key()
-            chunk_op.expect_worker = worker
+            chunk_op.expect_worker = ip_to_worker[worker]
 
             input_chunks = []
             concat_args = worker_to_args.get(worker, {})
@@ -301,7 +314,7 @@ def tile(cls, op: "LGBMTrain"):
             ).chunks[0]
             input_chunks.append(worker_ports_chunk)
 
-            chunk_op._workers = workers
+            chunk_op._workers = ips
             chunk_op._worker_ports = worker_ports_chunk
             chunk_op._worker_id = worker_id
 
@@ -357,9 +370,8 @@ def execute(cls, ctx, op: "LGBMTrain"):
         # if model is trained, remove unsupported parameters
         params.pop("out_dtype_", None)
         worker_ports = ctx[op.worker_ports.key]
-        worker_ips = [worker.split(":", 1)[0] for worker in op.workers]
         worker_endpoints = [
-            f"{worker}:{port}" for worker, port in zip(worker_ips, worker_ports)
+            f"{worker}:{port}" for worker, port in zip(op.workers, worker_ports)
         ]
 
         params["machines"] = ",".join(worker_endpoints)

mars/learn/contrib/lightgbm/tests/test_classifier.py

@@ -21,6 +21,7 @@
 
 from ..... import tensor as mt
 from ..... import dataframe as md
+from .....deploy.oscar.local import new_cluster
 
 try:
     import lightgbm
@@ -45,8 +46,22 @@
 X_sparse = mt.tensor(x_sparse, chunk_size=chunk_size).tosparse(missing=np.nan)[filter]
 
 
+@pytest.mark.parametrize(indirect=True)
+@pytest.fixture
+async def create_cluster():
+    start_method = os.environ.get("POOL_START_METHOD", None)
+    client = await new_cluster(
+        subprocess_start_method=start_method,
+        n_worker=2,
+        n_cpu=4,
+        use_uvloop=False,
+    )
+    async with client:
+        yield client
+
+
 @pytest.mark.skipif(lightgbm is None, reason="LightGBM not installed")
-def test_local_classifier(setup):
+def test_local_classifier(create_cluster):
     y_data = (y * 10).astype(mt.int32)
     classifier = LGBMClassifier(n_estimators=2)
     classifier.fit(X, y_data, eval_set=[(X, y_data)], verbose=True)

mars/services/meta/core.py

@@ -68,6 +68,7 @@ class _ChunkMeta(_CommonMeta):
     bands: List[BandType] = None
     # needed by ray ownership to keep object alive when worker died.
     object_refs: List[Any] = None
+    ip: str = None
 
     def merge_from(self, value: "_ChunkMeta"):
         if value.bands:

mars/services/subtask/worker/processor.py

@@ -32,7 +32,12 @@
 from ....optimization.physical import optimize
 from ....serialization import AioSerializer
 from ....typing import BandType, ChunkType
-from ....utils import get_chunk_key_to_data_keys, calc_data_size
+from ....utils import (
+    calc_data_size,
+    get_chunk_key_to_data_keys,
+    get_node_ip_address,
+    lazy_import,
+)
 from ...context import ThreadedServiceContext
 from ...meta.api import MetaAPI, WorkerMetaAPI
 from ...session import SessionAPI
@@ -41,6 +46,9 @@
 from ..core import Subtask, SubtaskStatus, SubtaskResult
 from ..utils import iter_input_data_keys, iter_output_data, get_mapper_data_keys
 
+
+lightgbm = lazy_import("lightgbm")
+
 logger = logging.getLogger(__name__)
 
 
@@ -464,6 +472,10 @@ async def _store_meta(
                         bands=[self._band],
                         chunk_key=chunk_key,
                         exclude_fields=["object_ref"],
+                        # Note: Why add the `ip` field?
+                        # lightgbm needs the machine addresses where the data are
+                        # to implement distributed learning.
+                        ip=get_node_ip_address() if lightgbm else None,
                     )
                 )
             # for supervisor, only save basic meta that is small like memory_size etc

mars/utils.py

@@ -1877,3 +1877,38 @@ def retry_call(*args, **kwargs):
             raise ex  # pylint: disable-msg=E0702
 
     return retry_call
+
+
+# `get_node_ip_address` is taken from Ray.
+# https://github.com/ray-project/ray/blob/master/python/ray/_private/services.py#L617
+def get_node_ip_address(address="8.8.8.8:53"):
+    """Determine the IP address of the local node.
+
+    Args:
+        address (str): The IP address and port of any known live service on the
+            network you care about.
+
+    Returns:
+        The IP address of the current node.
+    """
+    ip_address, port = address.split(":")
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        # This command will raise an exception if there is no internet
+        # connection.
+        s.connect((ip_address, int(port)))
+        node_ip_address = s.getsockname()[0]
+    except Exception as e:  # pragma: no cover
+        node_ip_address = "127.0.0.1"
+        # [Errno 101] Network is unreachable
+        if e.errno == 101:
+            try:
+                # try get node ip address from host name
+                host_name = socket.getfqdn(socket.gethostname())
+                node_ip_address = socket.gethostbyname(host_name)
+            except Exception:  # noqa: E722  # nosec  # pylint: disable=bare-except
+                pass
+    finally:
+        s.close()
+
+    return node_ip_address