Dask Flow Integration (#374)

* Dask Flow Integration

* Remove explicit persist of data

* Updating error messages, some refactoring

---------

Co-authored-by: Cho <choeric@amazon.com>

Author: dewan-c

.gitignore

@@ -9,3 +9,4 @@ __pycache__
 .mypy_cache/
 .idea/
 .DS_Store
+test.parquet

src/sagemaker_xgboost_container/algorithm_mode/hyperparameter_validation.py

@@ -338,6 +338,8 @@ def interaction_constraints_validator(value, dependencies):
         hpv.CategoricalHyperparameter(name="deterministic_histogram", range=["true", "false"], required=False),
         hpv.CategoricalHyperparameter(name="sampling_method", range=["uniform", "gradient_based"], required=False),
         hpv.IntegerHyperparameter(name="prob_buffer_row", range=hpv.Interval(min_open=1.0), required=False),
+        # Not an XGB training HP, but is used to determine which distributed training framework to use by SM XGB.
+        hpv.CategoricalHyperparameter(name="use_dask_gpu_training", range=["true", "false"], required=False),
     )
 
     hyperparameters.declare_alias("eta", "learning_rate")

src/sagemaker_xgboost_container/algorithm_mode/train.py

@@ -12,25 +12,26 @@
 # language governing permissions and limitations under the License.
 import logging
 import os
-import signal
 
 import numpy as np
 import xgboost as xgb
 from sklearn.model_selection import RepeatedKFold, RepeatedStratifiedKFold
 
 from sagemaker_algorithm_toolkit import exceptions as exc
 from sagemaker_algorithm_toolkit.channel_validation import Channel
-from sagemaker_xgboost_container import checkpointing, distributed
+from sagemaker_xgboost_container import distributed
 from sagemaker_xgboost_container.algorithm_mode import channel_validation as cv
 from sagemaker_xgboost_container.algorithm_mode import hyperparameter_validation as hpv
 from sagemaker_xgboost_container.algorithm_mode import metrics as metrics_mod
 from sagemaker_xgboost_container.algorithm_mode import train_utils
-from sagemaker_xgboost_container.callback import add_debugging
-from sagemaker_xgboost_container.constants.sm_env_constants import SM_OUTPUT_DATA_DIR
+from sagemaker_xgboost_container.callback import add_debugging, get_callbacks
+from sagemaker_xgboost_container.constants.sm_env_constants import (
+    SM_NUM_GPUS,
+    SM_OUTPUT_DATA_DIR,
+)
 from sagemaker_xgboost_container.constants.xgb_constants import (
     CUSTOMER_ERRORS,
     MODEL_NAME,
-    XGB_MAXIMIZE_METRICS,
 )
 from sagemaker_xgboost_container.data_utils import (
     check_data_redundancy,
@@ -39,30 +40,12 @@
     get_size,
     validate_data_file_path,
 )
+from sagemaker_xgboost_container.distributed_gpu import distributed_gpu_training
 from sagemaker_xgboost_container.prediction_utils import ValidationPredictionRecorder
 
 logger = logging.getLogger(__name__)
 
-
-def add_sigterm_handler(model_dir, is_master):
-    """Stop training and cleanup model directory when SIGTERM is received.
-
-    Model directory is only cleaned if is_master is True. Otherwise program terminates.
-
-    :param model_dir: Directory where model is saved
-    :param is_master: True if single node training, or the current node is the master node in distributed training
-    """
-
-    def _terminate():
-        os._exit(0)
-
-    def _cleanup_files(signo, frame):
-        if is_master:
-            train_utils.cleanup_dir(model_dir, MODEL_NAME)
-
-        _terminate()
-
-    signal.signal(signal.SIGTERM, _cleanup_files)
+DOCUMENTATION_LINK = "https://docs.aws.amazon.com/sagemaker/latest/dg/xgboost.html"
 
 
 def get_validated_dmatrices(
@@ -169,50 +152,86 @@ def sagemaker_train(
     # Obtain information about training resources to determine which distributed setup to use, if needed.
     num_hosts = len(sm_hosts)
 
-    train_dmatrix, val_dmatrix, train_val_dmatrix = get_validated_dmatrices(
-        train_path, val_path, file_type, csv_weights, is_pipe, combine_train_val
-    )
     checkpoint_dir = checkpoint_config.get("LocalPath", None)
 
-    train_args = dict(
-        train_cfg=validated_train_config,
-        train_dmatrix=train_dmatrix,
-        val_dmatrix=val_dmatrix,
-        train_val_dmatrix=train_val_dmatrix,
-        model_dir=model_dir,
-        checkpoint_dir=checkpoint_dir,
-    )
+    num_gpus = int(os.getenv(SM_NUM_GPUS, 0))
+    logging.info(f"Determined {num_gpus} GPU(s) available on the instance.")
+    tree_method_hp = validated_train_config.get("tree_method")
+
+    is_dask_job = validated_train_config.pop("use_dask_gpu_training", "false")
 
-    if num_hosts > 1:
-        # Wait for hosts to find each other
-        logging.info("Distributed node training with {} hosts: {}".format(num_hosts, sm_hosts))
-        distributed.wait_hostname_resolution(sm_hosts)
+    if is_dask_job == "true":
+        gpu_train_validation_errors = distributed_gpu_training.validate_gpu_train_configuration(
+            tree_method_hp=tree_method_hp,
+            num_hosts=num_hosts,
+            num_gpus=num_gpus,
+            input_mode=input_mode,
+            input_format=file_type,
+            data_config=validated_data_config,
+        )
 
-        if not train_dmatrix:
+        if gpu_train_validation_errors:
+            raise exc.UserError(f"Some configurations unsuitable for Dask GPU training were found: "
+                                f"{'. '.join(gpu_train_validation_errors)}")
+
+        logging.info("Going to run distributed GPU training through Dask.")
+        distributed_gpu_training.run_training_with_dask(
+            hyperparameters=validated_train_config,
+            train_path=train_path,
+            validation_path=val_path,
+            model_dir=model_dir,
+            content_type=file_type,
+            sm_hosts=sm_hosts,
+            current_host=sm_current_host,
+            checkpoint_dir=checkpoint_dir,
+            num_gpus=num_gpus,
+        )
+    else:
+        if num_gpus > 1:
             logging.warning(
-                "Host {} does not have data. Will broadcast to cluster and will not be used in distributed"
-                " training.".format(sm_current_host)
+                f"If you're using GPU training, not all GPUs on the instance will be used. "
+                f"See how to use all GPUs at {DOCUMENTATION_LINK}"
             )
-        distributed.rabit_run(
-            exec_fun=train_job,
-            args=train_args,
-            include_in_training=(train_dmatrix is not None),
-            hosts=sm_hosts,
-            current_host=sm_current_host,
-            update_rabit_args=True,
+
+        train_dmatrix, val_dmatrix, train_val_dmatrix = get_validated_dmatrices(
+            train_path, val_path, file_type, csv_weights, is_pipe, combine_train_val
         )
-    elif num_hosts == 1:
-        if train_dmatrix:
-            if validation_channel:
-                if not val_dmatrix:
+        train_args = dict(
+            train_cfg=validated_train_config,
+            train_dmatrix=train_dmatrix,
+            val_dmatrix=val_dmatrix,
+            train_val_dmatrix=train_val_dmatrix,
+            model_dir=model_dir,
+            checkpoint_dir=checkpoint_dir,
+        )
+        if num_hosts > 1:
+            # Wait for hosts to find each other
+            logging.info("Distributed node training with {} hosts: {}".format(num_hosts, sm_hosts))
+            distributed.wait_hostname_resolution(sm_hosts)
+            if not train_dmatrix:
+                logging.warning(
+                    "Host {} does not have data. Will broadcast to cluster and will not be used in distributed"
+                    " training.".format(sm_current_host)
+                )
+            distributed.rabit_run(
+                exec_fun=train_job,
+                args=train_args,
+                include_in_training=(train_dmatrix is not None),
+                hosts=sm_hosts,
+                current_host=sm_current_host,
+                update_rabit_args=True,
+            )
+        elif num_hosts == 1:
+            if train_dmatrix:
+                if validation_channel and not val_dmatrix:
                     raise exc.UserError("No data in validation channel path {}".format(val_path))
-            logging.info("Single node training.")
-            train_args.update({"is_master": True})
-            train_job(**train_args)
+                logging.info("Single node training.")
+                train_args.update({"is_master": True})
+                train_job(**train_args)
+            else:
+                raise exc.UserError("No data in training channel path {}".format(train_path))
         else:
-            raise exc.UserError("No data in training channel path {}".format(train_path))
-    else:
-        raise exc.PlatformError("Number of hosts should be an int greater than or equal to 1")
+            raise exc.PlatformError("Number of hosts should be an int greater than or equal to 1")
 
 
 def train_job(train_cfg, train_dmatrix, val_dmatrix, train_val_dmatrix, model_dir, checkpoint_dir, is_master):
@@ -259,11 +278,12 @@ def train_job(train_cfg, train_dmatrix, val_dmatrix, train_val_dmatrix, model_di
 
     try:
         kfold = train_cfg.pop("_kfold", None)
+        watchlist = [(train_dmatrix, "train")]
+        if val_dmatrix is not None:
+            watchlist.append((val_dmatrix, "validation"))
 
         if kfold is None:
-            xgb_model, iteration, callbacks, watchlist = get_callbacks_watchlist(
-                train_dmatrix=train_dmatrix,
-                val_dmatrix=val_dmatrix,
+            xgb_model, iteration, callbacks = get_callbacks(
                 model_dir=model_dir,
                 checkpoint_dir=checkpoint_dir,
                 early_stopping_data_name=early_stopping_data_name,
@@ -322,9 +342,7 @@ def train_job(train_cfg, train_dmatrix, val_dmatrix, train_val_dmatrix, model_di
                 cv_train_dmatrix = train_val_dmatrix.slice(train_idx)
                 cv_val_dmatrix = train_val_dmatrix.slice(val_idx)
 
-                xgb_model, iteration, callbacks, watchlist = get_callbacks_watchlist(
-                    train_dmatrix=cv_train_dmatrix,
-                    val_dmatrix=cv_val_dmatrix,
+                xgb_model, iteration, callbacks = get_callbacks(
                     model_dir=model_dir,
                     checkpoint_dir=checkpoint_dir,
                     early_stopping_data_name=early_stopping_data_name,
@@ -391,61 +409,6 @@ def train_job(train_cfg, train_dmatrix, val_dmatrix, train_val_dmatrix, model_di
                 logging.debug("Stored trained model {} at {}".format(fold, model_location))
 
 
-def get_callbacks_watchlist(
-    train_dmatrix,
-    val_dmatrix,
-    model_dir,
-    checkpoint_dir,
-    early_stopping_data_name,
-    early_stopping_metric,
-    early_stopping_rounds,
-    save_model_on_termination,
-    is_master,
-    fold=None,
-):
-    if checkpoint_dir and fold is not None:
-        checkpoint_dir = os.path.join(checkpoint_dir, f"model-{fold}")
-
-    # Set callbacks
-    xgb_model, iteration = checkpointing.load_checkpoint(checkpoint_dir)
-    if xgb_model is not None:
-        if fold is not None:
-            xgb_model = f"{xgb_model}-{fold}"
-        logging.info("Checkpoint loaded from %s", xgb_model)
-        logging.info("Resuming from iteration %s", iteration)
-
-    callbacks = []
-    callbacks.append(xgb.callback.EvaluationMonitor())
-    if checkpoint_dir:
-        save_checkpoint = xgb.callback.TrainingCheckPoint(
-            directory=checkpoint_dir, iterations=iteration, name=checkpointing.CHECKPOINT_FILENAME
-        )
-        callbacks.append(save_checkpoint)
-
-    if save_model_on_termination == "true":
-        model_name = f"{MODEL_NAME}-{fold}" if fold is not None else MODEL_NAME
-        save_intermediate_model = checkpointing.SaveIntermediateModelCallBack(model_dir, model_name, is_master)
-        callbacks.append(save_intermediate_model)
-        add_sigterm_handler(model_dir, is_master)
-
-    if early_stopping_data_name and early_stopping_metric and early_stopping_rounds:
-        maximize = early_stopping_metric in XGB_MAXIMIZE_METRICS
-        early_stop = xgb.callback.EarlyStopping(
-            rounds=early_stopping_rounds,
-            data_name=early_stopping_data_name,
-            metric_name=early_stopping_metric,
-            maximize=maximize,
-            save_best=True,
-        )
-        callbacks.append(early_stop)
-
-    watchlist = [(train_dmatrix, "train")]
-    if val_dmatrix is not None:
-        watchlist.append((val_dmatrix, "validation"))
-
-    return xgb_model, iteration, callbacks, watchlist
-
-
 def print_cv_metric(num_round, evals_results):
     cv_eval_report = f"[{num_round}]"
     for metric_name in evals_results[0]["train"]:

src/sagemaker_xgboost_container/callback.py

@@ -1,6 +1,11 @@
 import logging
-
+import os
+import signal
 import xgboost as xgb
+
+from sagemaker_xgboost_container import checkpointing
+from sagemaker_xgboost_container.algorithm_mode import train_utils
+from sagemaker_xgboost_container.constants.xgb_constants import MODEL_NAME, XGB_MAXIMIZE_METRICS
 from smdebug.xgboost import Hook
 
 logger = logging.getLogger(__name__)
@@ -45,3 +50,73 @@ def add_debugging(callbacks, hyperparameters, train_dmatrix, val_dmatrix=None, j
         logging.debug("Failed to create debug hook", e)
     else:
         callbacks.append(hook)
+
+
+def add_sigterm_handler(model_dir, is_master):
+    """Stop training and cleanup model directory when SIGTERM is received.
+
+    Model directory is only cleaned if is_master is True. Otherwise program terminates.
+
+    :param model_dir: Directory where model is saved
+    :param is_master: True if single node training, or the current node is the master node in distributed training
+    """
+
+    def _terminate():
+        os._exit(0)
+
+    def _cleanup_files(signo, frame):
+        if is_master:
+            train_utils.cleanup_dir(model_dir, MODEL_NAME)
+
+        _terminate()
+
+    signal.signal(signal.SIGTERM, _cleanup_files)
+
+
+def get_callbacks(
+    model_dir,
+    checkpoint_dir,
+    early_stopping_data_name,
+    early_stopping_metric,
+    early_stopping_rounds,
+    save_model_on_termination,
+    is_master,
+    fold=None,
+):
+    if checkpoint_dir and fold is not None:
+        checkpoint_dir = os.path.join(checkpoint_dir, f"model-{fold}")
+
+    # Set callbacks
+    xgb_model, iteration = checkpointing.load_checkpoint(checkpoint_dir)
+    if xgb_model is not None:
+        if fold is not None:
+            xgb_model = f"{xgb_model}-{fold}"
+        logging.info("Checkpoint loaded from %s", xgb_model)
+        logging.info("Resuming from iteration %s", iteration)
+
+    callbacks = []
+    callbacks.append(xgb.callback.EvaluationMonitor())
+    if checkpoint_dir:
+        save_checkpoint = xgb.callback.TrainingCheckPoint(
+            directory=checkpoint_dir, iterations=iteration, name=checkpointing.CHECKPOINT_FILENAME
+         )
+        callbacks.append(save_checkpoint)
+
+    if save_model_on_termination == "true":
+        model_name = f"{MODEL_NAME}-{fold}" if fold is not None else MODEL_NAME
+        save_intermediate_model = checkpointing.SaveIntermediateModelCallBack(model_dir, model_name, is_master)
+        callbacks.append(save_intermediate_model)
+        add_sigterm_handler(model_dir, is_master)
+
+    if early_stopping_data_name and early_stopping_metric and early_stopping_rounds:
+        maximize = early_stopping_metric in XGB_MAXIMIZE_METRICS
+        early_stop = xgb.callback.EarlyStopping(
+            rounds=early_stopping_rounds,
+            data_name=early_stopping_data_name,
+            metric_name=early_stopping_metric,
+            maximize=maximize,
+            save_best=True,
+        )
+        callbacks.append(early_stop)
+
+    return xgb_model, iteration, callbacks

src/sagemaker_xgboost_container/constants/xgb_constants.py

@@ -93,3 +93,7 @@
 MULTI_SOFTPROB = "multi:softprob"
 
 MODEL_NAME = "xgboost-model"
+GPU_TREE_METHOD = "gpu_hist"
+
+FULLY_REPLICATED = "FullyReplicated"
+PIPE_MODE = "Pipe"