Track total training time (#58)

* Update README

* Track total training time

* Fix test colocation

Author: krfricke

xgboost_ray/main.py

@@ -681,6 +681,8 @@ class _TrainingState:
     checkpoint: _Checkpoint
     additional_results: Dict
 
+    training_started_at: float = 0.
+
     placement_group: Optional[PlacementGroup] = None
 
     failed_actor_ranks: set = field(default_factory=set)
@@ -830,6 +832,8 @@ def handle_actor_failure(actor_id):
         _training_state.additional_results[
             "callback_returns"] = callback_returns
 
+    _training_state.training_started_at = time.time()
+
     # Trigger the train function
     training_futures = [
         actor.train.remote(rabit_args, params, dtrain, evals, *args, **kwargs)
@@ -919,9 +923,6 @@ def handle_actor_failure(actor_id):
 
     _training_state.additional_results["total_n"] = total_n
 
-    logger.info(f"[RayXGBoost] Finished XGBoost training on training data "
-                f"with total N={total_n:,}.")
-
     return bst, evals_result, _training_state.additional_results
 
 
@@ -1020,6 +1021,8 @@ def _wrapped(*args, **kwargs):
             additional_results.update(train_additional_results)
         return bst
 
+    start_time = time.time()
+
     ray_params = _validate_ray_params(ray_params)
 
     max_actor_restarts = ray_params.max_actor_restarts \
@@ -1104,13 +1107,15 @@ def _wrapped(*args, **kwargs):
 
     start_actor_ranks = set(range(ray_params.num_actors))  # Start these
 
+    total_training_time = 0.
     while tries <= max_actor_restarts:
         training_state = _TrainingState(
             actors=actors,
             queue=queue,
             stop_event=stop_event,
             checkpoint=checkpoint,
             additional_results=current_results,
+            training_started_at=0.,
             placement_group=pg,
             failed_actor_ranks=start_actor_ranks,
             pending_actors=pending_actors)
@@ -1126,8 +1131,14 @@ def _wrapped(*args, **kwargs):
                 gpus_per_actor=gpus_per_actor,
                 _training_state=training_state,
                 **kwargs)
+            if training_state.training_started_at > 0.:
+                total_training_time += time.time(
+                ) - training_state.training_started_at
             break
         except (RayActorError, RayTaskError) as exc:
+            if training_state.training_started_at > 0.:
+                total_training_time += time.time(
+                ) - training_state.training_started_at
             alive_actors = sum(1 for a in actors if a is not None)
             start_again = False
             if ray_params.elastic_training:
@@ -1186,6 +1197,16 @@ def _wrapped(*args, **kwargs):
                 ) from exc
             tries += 1
 
+    total_time = time.time() - start_time
+
+    train_additional_results["training_time_s"] = total_training_time
+    train_additional_results["total_time_s"] = total_time
+
+    logger.info("[RayXGBoost] Finished XGBoost training on training data "
+                "with total N={total_n:,} in {total_time_s:.2f} seconds "
+                "({training_time_s:.2f} pure XGBoost training time).".format(
+                    **train_additional_results))
+
     _shutdown(
         actors=actors,
         pending_actors=pending_actors,

xgboost_ray/tests/test_colocation.py

@@ -2,7 +2,7 @@
 import shutil
 import tempfile
 import unittest
-from unittest.mock import patch, DEFAULT
+from unittest.mock import patch
 import pytest
 
 import numpy as np
@@ -79,7 +79,7 @@ def _mock_train(*args, _training_state, **kwargs):
                 assert ray.get(
                     _training_state.stop_event.actor.get_node_id.remote()) == \
                     ray.state.current_node_id()
-                return DEFAULT, DEFAULT, DEFAULT
+                return _train(*args, _training_state=_training_state, **kwargs)
 
             with patch("xgboost_ray.main._train") as mocked:
                 mocked.side_effect = _mock_train