Skip to content

Commit fa207a6

Browse files
authored
[Learn] Fix lightgbm train TypeError (#3349)
* fix lightgbm classifier when multiple machines * change log level for metrics record * add a test
1 parent adf9fc8 commit fa207a6

File tree

3 files changed

+68
-2
lines changed

3 files changed

+68
-2
lines changed

mars/learn/contrib/lightgbm/_train.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -436,7 +436,8 @@ def train(params, train_set, eval_sets=None, **kwargs):
436436
sample_weights.append(train_kw["sample_weight"])
437437
init_scores.append(train_kw["init_score"])
438438

439-
op = LGBMTrain(
439+
train_cls = kwargs.pop("train_cls", LGBMTrain)
440+
op = train_cls(
440441
params=params,
441442
data=datas[0],
442443
label=labels[0],
@@ -452,6 +453,10 @@ def train(params, train_set, eval_sets=None, **kwargs):
452453
kwds=kwargs,
453454
)
454455
ret = op().execute(session=session, **run_kwargs).fetch(session=session)
456+
# Note: There may be multiple results if running on clusters, so we should
457+
# combine the results if so.
458+
if isinstance(ret, list):
459+
ret = b"".join(ret)
455460

456461
bst = pickle.loads(ret)
457462
evals_result.update(bst.evals_result_ or {})

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

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,3 +170,64 @@ def test_local_classifier_from_to_parquet(setup):
170170
expected = classifier.predict(X)
171171
expected = np.stack([1 - expected, expected]).argmax(axis=0)
172172
np.testing.assert_array_equal(ret, expected)
173+
174+
175+
@pytest.mark.skipif(lightgbm is None, reason="LightGBM not installed")
176+
def test_classifier_on_multiple_machines(setup):
177+
from .._train import LGBMTrain
178+
179+
class MockLGMBTrain(LGBMTrain):
180+
@classmethod
181+
def execute(cls, ctx, op: "LGBMTrain"):
182+
super().execute(ctx, op)
183+
# Note: There may be a list result when running on multiple
184+
# machines, here just make an array of length 1 to simulate
185+
# this scenario.
186+
ctx[op.outputs[0].key] = [ctx[op.outputs[0].key]]
187+
188+
from ..core import LGBMModelType
189+
from .._train import train
190+
from ....utils import check_consistent_length
191+
192+
class MockLGBMClassifier(LGBMClassifier, lightgbm.LGBMClassifier):
193+
def fit(
194+
self,
195+
X,
196+
y,
197+
sample_weight=None,
198+
init_score=None,
199+
eval_set=None,
200+
eval_sample_weight=None,
201+
eval_init_score=None,
202+
session=None,
203+
run_kwargs=None,
204+
**kwargs,
205+
):
206+
check_consistent_length(X, y, session=session, run_kwargs=run_kwargs)
207+
params = self.get_params(True)
208+
model = train(
209+
params,
210+
self._wrap_train_tuple(X, y, sample_weight, init_score),
211+
eval_sets=self._wrap_eval_tuples(
212+
eval_set, eval_sample_weight, eval_init_score
213+
),
214+
model_type=LGBMModelType.CLASSIFIER,
215+
session=session,
216+
run_kwargs=run_kwargs,
217+
train_cls=MockLGMBTrain,
218+
**kwargs,
219+
)
220+
221+
self.set_params(**model.get_params())
222+
self._copy_extra_params(model, self)
223+
return self
224+
225+
y_data = (y * 10).astype(mt.int32)
226+
classifier = MockLGBMClassifier(n_estimators=2)
227+
classifier.fit(X, y_data, eval_set=[(X, y_data)], verbose=True)
228+
prediction = classifier.predict(X)
229+
230+
assert prediction.ndim == 1
231+
assert prediction.shape[0] == len(X)
232+
233+
assert isinstance(prediction, mt.Tensor)

mars/metrics/api.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ def record(self, value=1, tags: Optional[Dict[str, str]] = None):
119119
self._metric.record(value, tags)
120120
elif not self._log_not_init_error:
121121
self._log_not_init_error = True
122-
logger.warning(
122+
logger.info(
123123
"Metric is not initialized, please call `init_metrics()` before using metrics."
124124
)
125125

0 commit comments

Comments
 (0)