Upgrade ray tune component (#270)

* add more optimizer/distribution for ray tune component and add ray tune distributed component.
* revise lightgbm_training to use ray tune distributed component.
* use distributed component to set up ray cluster with more than 1 node.
* update requirements for ray tune components.
* provide ray tune example yaml.
* provide ray tune distributed example yaml file.

Author: ruizhuanguw

.github/workflows/docs.yml

@@ -33,8 +33,9 @@ jobs:
     - name: pip install
       run: |
         python -m pip install --upgrade pip==21.3.1
-        python -m pip install markdown-include==0.6.0 mkdocstrings==0.15.0 mkdocs-material==7.1.3 livereload==2.6.3
-        
+
+        python -m pip install markdown-include==0.7.0 mkdocstrings==0.19.0 mkdocstrings-python==0.7.1 mkdocs-material==8.4.2 livereload==2.6.3
+
         # NOTE: we need requirements to be able to parse reference docs scripts
         sudo apt-get install libopenmpi-dev
         python -m pip install -r ./requirements.txt

conf/experiments/lightgbm_training/raytune.yaml

@@ -0,0 +1,81 @@
+# to execute:
+# > python src/pipelines/azureml/lightgbm_training.py --exp-config conf/experiments/lightgbm_training/raytune.yaml
+
+defaults:
+  - aml: lightgbm-benchmark-eus2
+  - compute: lightgbm-benchmark-eus2
+
+### CUSTOM PARAMETERS ###
+
+experiment:
+  name: "dev_lightgbm_ray_tune"
+  description: "something interesting to say about this"
+
+lightgbm_training_config:
+  # name of your particular benchmark
+  benchmark_name: "lightgbm-ray-tune" # override this with a unique name
+
+  # list all the train/test pairs to train on
+  tasks:
+    - train:
+        name: "data-synthetic-headercsv-regression-10cols-100000samples-train"
+      test:
+        name: "data-synthetic-headercsv-regression-10cols-10000samples-test"
+      task_key: "dev_ray" # optional, user to register outputs
+
+  # NOTE: this example uses only 1 training (reference)
+  # see other config files for creating training variants
+  reference:
+    framework: lightgbm_ray_tune
+
+    # input parameters
+    data:
+      auto_partitioning: True # inserts partitioning to match expected number of partitions (if nodes*processes > 1)
+      pre_convert_to_binary: False # inserts convertion of train/test data into binary to speed up training (not compatible with auto_partitioning yet)
+      header: true # IMPORTANT
+      label_column: "0"
+      group_column: null
+
+    # lightgbm training parameters
+    training:
+      objective: "regression"
+      metric: "rmse"
+      boosting: "gbdt"
+      tree_learner: "data"
+      num_iterations: "choice([30,40,50,60])"
+      num_leaves: "31"
+      min_data_in_leaf: "20"
+      learning_rate: "0.1"
+      max_bin: "255"
+      feature_fraction: "1.0"
+
+      # compute parameters
+      device_type: "cpu"
+
+      # you can add anything under custom_params, it will be sent as a dictionary
+      # to the lightgbm training module to override its parameters (see lightgbm docs for list)
+      custom_params:
+          deterministic: True
+          use_two_round_loading: True
+
+    # compute parameters
+    runtime:
+      target: "linux-cpu-ds14v2" # optional: force target for this training job
+      nodes: 1
+      processes: 1
+
+    # model registration
+    # naming convention: "{register_model_prefix}-{task_key}-{num_iterations}trees-{num_leaves}leaves-{register_model_suffix}"
+    output:
+      register_model: False
+      #register_model_prefix: "model"
+      #register_model_suffix: null
+
+    raytune:
+      mode: "min"
+      search_alg: "BasicVariantGenerator"
+      scheduler: "FIFOScheduler"
+      num_samples: 5
+      time_budget: 1800
+      concurrent_trials: 0
+      cpus_per_trial: 16

conf/experiments/lightgbm_training/raytune_distributed.yaml

@@ -0,0 +1,84 @@
+# to execute:
+# > python src/pipelines/azureml/lightgbm_training.py --exp-config conf/experiments/lightgbm_training/raytune_distributed.yaml
+
+defaults:
+  - aml: lightgbm-benchmark-eus2
+  - compute: lightgbm-benchmark-eus2
+
+### CUSTOM PARAMETERS ###
+
+experiment:
+  name: "dev_lightgbm_ray_tune"
+  description: "something interesting to say about this"
+
+lightgbm_training_config:
+  # name of your particular benchmark
+  benchmark_name: "lightgbm-ray-tune" # override this with a unique name
+
+  # list all the train/test pairs to train on
+  tasks:
+    - train:
+        name: "data-synthetic-headercsv-regression-10cols-100000samples-train"
+      test:
+        name: "data-synthetic-headercsv-regression-10cols-10000samples-test"
+      task_key: "dev_ray" # optional, user to register outputs
+
+  # NOTE: this example uses only 1 training (reference)
+  # see other config files for creating training variants
+  reference:
+    framework: lightgbm_ray_tune_distributed
+
+    # input parameters
+    data:
+      auto_partitioning: True # inserts partitioning to match expected number of partitions (if nodes*processes > 1)
+      pre_convert_to_binary: False # inserts convertion of train/test data into binary to speed up training (not compatible with auto_partitioning yet)
+      header: true # IMPORTANT
+      label_column: "0"
+      group_column: null
+      train_data_format: 'CSV'
+      test_data_format: 'CSV'
+
+    # lightgbm training parameters
+    training:
+      objective: "regression"
+      metric: "rmse"
+      boosting: "gbdt"
+      tree_learner: "data"
+      num_iterations: "choice([30,40,50,60])"
+      num_leaves: "31"
+      min_data_in_leaf: "20"
+      learning_rate: "0.1"
+      max_bin: "255"
+      feature_fraction: "1.0"
+
+      # compute parameters
+      device_type: "cpu"
+
+      # you can add anything under custom_params, it will be sent as a dictionary
+      # to the lightgbm training module to override its parameters (see lightgbm docs for list)
+      custom_params:
+          deterministic: True
+          use_two_round_loading: True
+
+    # compute parameters
+    runtime:
+      target: "linux-cpu-ds14v2" # optional: force target for this training job
+      nodes: 4
+      processes: 1
+
+    # model registration
+    # naming convention: "{register_model_prefix}-{task_key}-{num_iterations}trees-{num_leaves}leaves-{register_model_suffix}"
+    output:
+      register_model: False
+      #register_model_prefix: "model"
+      #register_model_suffix: null
+
+    raytune:
+      mode: "min"
+      search_alg: "BasicVariantGenerator"
+      scheduler: "FIFOScheduler"
+      num_samples: 5
+      time_budget: 1800
+      concurrent_trials: 2
+      lightgbm_ray_actors: 2
+      cpus_per_actor: 16

requirements.txt

@@ -12,6 +12,9 @@ lightgbm==3.3.1
 treelite==2.1.0
 treelite_runtime==2.1.0
 flaml==0.9.6
+hpbandster==0.7.4
+ConfigSpace==0.5.0
+optuna==2.8.0
 
 # pipelines
 shrike[pipeline]==1.14.7

src/__init__.py

@@ -111,7 +111,7 @@ def apply_sweep_settings(step, sweep_settings_config):
             )
         elif sweep_settings_config.early_termination.policy_type == "truncation_selection":
             step.runsettings.sweep.early_termination.configure(
-                policy_type="bandit",
+                policy_type="truncation_selection",
                 truncation_percentage=sweep_settings_config.early_termination.truncation_percentage,
                 evaluation_interval=sweep_settings_config.early_termination.evaluation_interval,
                 delay_evaluation=sweep_settings_config.early_termination.delay_evaluation

src/common/aml.py

@@ -4,38 +4,25 @@
 """
 Parses Ray Tune parameters from text arguments (cli or yaml)
 """
-import re
 import argparse
 import logging
-from azureml.core import Workspace, Datastore, Dataset
-import ray
-from ray import tune
 from ray.tune import (
     uniform,
     quniform,
-    loguniform,
-    qloguniform,
-    randn,
-    qrandn,
     randint,
     qrandint,
-    lograndint,
-    qlograndint,
     choice,
 )
 
 
 class RayTuneParameterParser():
 
+    # TODO: allow more distributions
     DISTRIBUTIONS_MAP = {"choice": choice,
                          "uniform": uniform,
-                         # "loguniform": loguniform,
-                         # "normal": normal,
-                         # "lognormal": lognormal,
-                         # "quniform": quniform,
-                         # "qloguniform": qloguniform,
-                         # "qnormal": qnormal,
-                         # "qlognormal": qlognormal,
+                         "quniform": quniform,
+                         "randint": randint,
+                         "qrandint": qrandint,
                          }
 
     def __init__(self, tunable_parameters):

src/common/raytune_param.py

@@ -117,6 +117,10 @@ class lightgbm_training_data_variant_parameters:
     group_column: Optional[str] = None
     construct: bool = True
 
+    # data formats
+    train_data_format: Optional[str] = None
+    test_data_format: Optional[str] = None
+
 @dataclass
 class lightgbm_training_environment_variant_parameters:
     # COMPUTE

src/common/tasks.py

@@ -86,6 +86,9 @@ class lightgbm_training_config: # pragma: no cover
 # load ray tune module.
 lightgbm_ray_tune_module = Component.from_yaml(yaml_file=os.path.join(COMPONENTS_ROOT, "training", "ray_tune", "spec.yaml"))
 
+# load ray tune distributed module.
+lightgbm_ray_tune_distributed_module = Component.from_yaml(yaml_file=os.path.join(
+    COMPONENTS_ROOT, "training", "ray_tune_distributed", "spec.yaml"))
 ### PIPELINE SPECIFIC CODE ###
 
 def process_sweep_parameters(params_dict, sweep_algorithm):
@@ -182,8 +185,20 @@ def lightgbm_training_pipeline_function(config,
 
         # if we're using multinode, add partitioning
         if variant_params.data.auto_partitioning and (variant_params.training.tree_learner == "data" or variant_params.training.tree_learner == "voting"):
+            # if training is distributed to multiple nodes using ray:
+            if variant_params.framework == 'lightgbm_ray_tune_distributed' and variant_params.raytune.lightgbm_ray_actors > 1:
+                partition_data_step = partition_data_module(
+                    input_data=train_dataset,
+                    mode="roundrobin",
+                    number=variant_params.raytune.lightgbm_ray_actors,
+                    header=variant_params.data.header,
+                    verbose=variant_params.training.verbose
+                )
+                partition_data_step.runsettings.configure(
+                    target=config.compute.linux_cpu)
+                partitioned_train_data = partition_data_step.outputs.output_data
             # if using data parallel, train data has to be partitioned first
-            if (variant_params.runtime.nodes * variant_params.runtime.processes) > 1:
+            elif (variant_params.runtime.nodes * variant_params.runtime.processes) > 1:
                 partition_data_step = partition_data_module(
                     input_data=train_dataset,
                     mode="roundrobin",
@@ -309,20 +324,50 @@ def lightgbm_training_pipeline_function(config,
             training_params['num_samples'] = variant_params.raytune.num_samples
             training_params['time_budget'] = variant_params.raytune.time_budget
             training_params['concurrent_trials'] = variant_params.raytune.concurrent_trials
+            training_params['cpus_per_trial'] = variant_params.raytune.cpus_per_trial
+            if 'low_num_iterations' in variant_params.raytune:
+                training_params['low_num_iterations'] = variant_params.raytune.low_num_iterations
+            if 'low_num_leaves' in variant_params.raytune:
+                training_params['low_num_leaves'] = variant_params.raytune.low_num_leaves
+
+            # remove arguments that are not in lightgbm_ray_tune component
+            if 'multinode_driver' in training_params:
+                del training_params['multinode_driver']
+            if 'custom_properties' in training_params:
+                del training_params['custom_properties']
+            if 'verbose' in training_params:
+                del training_params['verbose']
+
+        elif variant_params.framework == 'lightgbm_ray_tune_distributed':
+            lightgbm_train_module = lightgbm_ray_tune_distributed_module
+            use_sweep = False
+
+            # manually add parameters for lightgbm_ray
+            training_params['train_data_format'] = variant_params.data.train_data_format
+            training_params['test_data_format'] = variant_params.data.test_data_format
+
+            # manually add ray tune parameters.
+            training_params['mode'] = variant_params.raytune.mode
+            training_params['search_alg'] = variant_params.raytune.search_alg
+            training_params['scheduler'] = variant_params.raytune.scheduler
+            training_params['num_samples'] = variant_params.raytune.num_samples
+            training_params['time_budget'] = variant_params.raytune.time_budget
+            training_params['concurrent_trials'] = variant_params.raytune.concurrent_trials
+            training_params['lightgbm_ray_actors'] = variant_params.raytune.lightgbm_ray_actors
+            training_params['cpus_per_actor'] = variant_params.raytune.cpus_per_actor
 
             # remove arguments that are not in lightgbm_ray_tune component
             if 'multinode_driver' in training_params:
                 del training_params['multinode_driver']
-            if 'header' in training_params:
-                del training_params['header']
             if 'construct' in training_params:
                 del training_params['construct']
             if 'custom_properties' in training_params:
                 del training_params['custom_properties']
             if 'verbose' in training_params:
                 del training_params['verbose']
         else:
-            raise NotImplementedError(f"training framework {variant_params.framework} hasn't been implemented yet.")
+            raise NotImplementedError(
+                f"training framework {variant_params.framework} hasn't been implemented yet.")
 
         # configure the training module
         lightgbm_train_step = lightgbm_train_module(
@@ -452,13 +497,38 @@ def main():
         "```"
     ])
 
+    # add pipeline tags
+    autotags = {}
+    reference_configs = config.lightgbm_training_config.reference
+    print(f"tags: {autotags}")
+    # add the information of the reference
+    if reference_configs.raytune:
+        autotags.update({
+            'search_algo': reference_configs.raytune.search_alg,
+            'scheduler': reference_configs.raytune.scheduler,
+            'concurrent_trials': str(reference_configs.raytune.concurrent_trials),
+            'time_minutes': str(reference_configs.raytune.time_budget/60),
+            'cluster_nodes': str(reference_configs.runtime.nodes),
+        })
+    if reference_configs.sweep:
+        autotags.update({
+            "search_algo": reference_configs.sweep.algorithm,
+            "truncate_percentage": str(reference_configs.sweep.early_termination.truncation_percentage),
+            "concurrent_trials": str(reference_configs.sweep.limits.max_concurrent_trials),
+            "time_minutes": str(reference_configs.sweep.limits.timeout_minutes),
+            'cluster_nodes': str(reference_configs.runtime.nodes * reference_configs.sweep.limits.max_concurrent_trials),
+        })
+
+    print(f"tags: {autotags}")
     # validate/submit the pipeline (if run.submit=True)
     pipeline_submit(
         workspace,
         config,
         pipeline_instance,
-        experiment_description=experiment_description
+        experiment_description=experiment_description,
+        tags=autotags
     )
 
+
 if __name__ == "__main__":
     main()