Initial implementation of the end-to-end autotrain module

ludwig/api.py

@@ -36,7 +36,7 @@
 
 from ludwig.backend import Backend, initialize_backend
 from ludwig.callbacks import Callback
-from ludwig.constants import FULL, PREPROCESSING, TEST, TRAINING, VALIDATION
+from ludwig.constants import FULL, PREPROCESSING, TEST, TRAINING, VALIDATION, LEARNING_RATE, BATCH_SIZE, AUTO
 from ludwig.data.dataset.base import Dataset
 from ludwig.data.postprocessing import convert_predictions, postprocess
 from ludwig.data.preprocessing import (load_metadata,
@@ -481,24 +481,24 @@ def train(
                     )
 
                 # auto tune batch size
-                if self.config[TRAINING]["batch_size"] == "auto":
+                if self.config[TRAINING][BATCH_SIZE] == AUTO:
                     # TODO (ASN): add support for substitute_with_max parameter
                     tuned_batch_size = trainer.tune_batch_size(
                         self.config,
                         training_set,
                         random_seed=random_seed
                     )
-                    self.config[TRAINING]['batch_size'] = tuned_batch_size
+                    self.config[TRAINING][BATCH_SIZE] = tuned_batch_size
 
                 # auto tune learning rate
-                if self.config[TRAINING]["learning_rate"] == "auto":
+                if self.config[TRAINING][LEARNING_RATE] == AUTO:
                     new_learning_rate = trainer.tune_learning_rate(
                         self.config,
                         LudwigModel.create_model(self.config, random_seed),
                         training_set,
                         random_seed=random_seed
                     )
-                    self.config[TRAINING]['learning_rate'] = new_learning_rate
+                    self.config[TRAINING][LEARNING_RATE] = new_learning_rate
 
                 # train model
                 if self.backend.is_coordinator():

ludwig/automl/automl.py

@@ -8,15 +8,29 @@
 (2) Tunes config based on resource constraints
 (3) Runs hyperparameter optimization experiment
 """
-from logging import raiseExceptions
+import logging
+import sys
 from typing import Dict, Union
 
-import dask.dataframe as dd
 import numpy as np
 import pandas as pd
 from ludwig.automl.base_config import create_default_config
 from ludwig.hyperopt.run import hyperopt
 
+logger = logging.getLogger(__name__)
+
+
+try:
+    import dask.dataframe as dd
+    import ray
+except ImportError:
+    logger.error(
+        ' ray is not installed. '
+        'In order to use auto_train please run '
+        'pip install ludwig[ray]'
+    )
+    sys.exit(-1)
+
 OUTPUT_DIR = "."
 
 
@@ -26,14 +40,15 @@ def model_select(default_configs):
     Note: Current implementation returns tabnet by default. This will be
         improved in subsequent iterations
     """
-    return default_configs['tabnet'], 'tabnet'
+    return default_configs['tabnet']
 
 
 def auto_train(
     dataset: Union[str, pd.DataFrame, dd.core.DataFrame],
     target: str,
     time_limit_s: Union[int, float],
-    output_dir: str = OUTPUT_DIR
+    output_dir: str = OUTPUT_DIR,
+    config=None,
 ):
     """
     Main auto train API that first builds configs for each model type
@@ -51,13 +66,12 @@ def auto_train(
     # Returns
     :return: (str) path to best trained model
     """
-
-    default_configs = create_default_config(dataset, target, time_limit_s)
-    model_config, model_name = model_select(default_configs)
-    hyperopt_results = _train(model_config, dataset,
+    if config is None:
+        config = _create_auto_config(dataset, target, time_limit_s)
+    model_name = config['combiner']['type']
+    hyperopt_results = _train(config, dataset,
                               output_dir, model_name=model_name)
     experiment_analysis = hyperopt_results.experiment_analysis
-
     # catch edge case where metric_score is nan
     # TODO (ASN): Decide how we want to proceed if at least one trial has
     # completed
@@ -76,6 +90,12 @@ def auto_train(
     return autotrain_results
 
 
+def _create_auto_config(dataset, target, time_limit_s) -> dict:
+    default_configs = create_default_config(dataset, target, time_limit_s)
+    model_config = model_select(default_configs)
+    return model_config
+
+
 def _train(
     config: Dict,
     dataset: Union[str, pd.DataFrame, dd.core.DataFrame],

ludwig/automl/base_config.py

@@ -14,15 +14,29 @@
     (base implementation -- # CPU, # GPU)
 """
 
+import logging
 import os
-from typing import Dict, List, Union
+import sys
+from typing import List, Union
 
 import pandas as pd
-import dask.dataframe as dd
-from ludwig.automl.utils import (FieldInfo, get_available_resources,
-                                 avg_num_tokens)
+from ludwig.automl.utils import (FieldInfo, avg_num_tokens,
+                                 get_available_resources)
+from ludwig.constants import BINARY, CATEGORY, CONFIG, NUMERICAL, TEXT, TYPE
 from ludwig.utils.data_utils import load_yaml
 
+logger = logging.getLogger(__name__)
+try:
+    import dask.dataframe as dd
+    import ray
+except ImportError:
+    logger.error(
+        ' ray is not installed. '
+        'In order to use auto_train please run '
+        'pip install ludwig[ray]'
+    )
+    sys.exit(-1)
+
 PATH_HERE = os.path.abspath(os.path.dirname(__file__))
 CONFIG_DIR = os.path.join(PATH_HERE, 'defaults')
 
@@ -33,16 +47,16 @@
 }
 
 
-def allocate_experiment_resources(resources: Dict) -> Dict:
+def allocate_experiment_resources(resources: dict) -> dict:
     """
     Allocates ray trial resources based on available resources
 
     # Inputs
-    :param resources (Dict) specifies all available GPUs, CPUs and associated
+    :param resources (dict) specifies all available GPUs, CPUs and associated
         metadata of the machines (i.e. memory)
 
     # Return
-    :return: (Dict) gpu and cpu resources per trial
+    :return: (dict) gpu and cpu resources per trial
     """
     # TODO (ASN):
     # (1) expand logic to support multiple GPUs per trial (multi-gpu training)
@@ -51,18 +65,23 @@ def allocate_experiment_resources(resources: Dict) -> Dict:
     experiment_resources = {
         'cpu_resources_per_trial': 1
     }
-    if resources['gpu'] > 0:
+    gpu_count, cpu_count = resources['gpu'], resources['cpu']
+    if gpu_count > 0:
         experiment_resources.update({
             'gpu_resources_per_trial': 1
         })
+        if cpu_count > 1:
+            cpus_per_trial = int(cpu_count/gpu_count)
+            experiment_resources['cpu_resources_per_trial'] = cpus_per_trial
+
     return experiment_resources
 
 
 def create_default_config(
     dataset: Union[str, dd.core.DataFrame, pd.DataFrame],
     target_name: str = None,
     time_limit_s: Union[int, float] = None
-) -> Dict:
+) -> dict:
     """
     Returns auto_train configs for three available combiner models. 
     Coordinates the following tasks:
@@ -81,7 +100,7 @@ def create_default_config(
                                     as the stopping parameter
 
     # Return
-    :return: (Dict) dictionaries contain auto train config files for all available
+    :return: (dict) dictionaries contain auto train config files for all available
     combiner types
 
     """
@@ -105,7 +124,7 @@ def create_default_config(
 
 def get_field_info(dataset: str):
     """
-    Constructs FeildInfo objects for each feature in dataset. These objects
+    Constructs FieldInfo objects for each feature in dataset. These objects
     are used for downstream type inference
 
     # Inputs
@@ -138,9 +157,9 @@ def get_field_info(dataset: str):
 def get_features_config(
     fields: List[FieldInfo],
     row_count: int,
-    resources: Dict,
+    resources: dict,
     target_name: str = None,
-) -> Dict:
+) -> dict:
     """
     Constructs FeildInfo objects for each feature in dataset. These objects
     are used for downstream type inference
@@ -151,7 +170,7 @@ def get_features_config(
     :param target_name (str) name of target feature
 
     # Return
-    :return: (Dict) section of auto_train config for input_features and output_features 
+    :return: (dict) section of auto_train config for input_features and output_features 
     """
     metadata = get_field_metadata(fields, row_count, resources, target_name)
     return get_config_from_metadata(metadata, target_name)
@@ -163,11 +182,11 @@ def get_config_from_metadata(metadata: list, target_name: str = None) -> dict:
     metadata
 
     # Inputs
-    :param metadata: (List[Dict]) field descriptions
+    :param metadata: (List[dict]) field descriptions
     :param target_name (str) name of target feature
 
     # Return
-    :return: (Dict) section of auto_train config for input_features and output_features
+    :return: (dict) section of auto_train config for input_features and output_features
     """
     config = {
         "input_features": [],
@@ -176,15 +195,15 @@ def get_config_from_metadata(metadata: list, target_name: str = None) -> dict:
 
     for field_meta in metadata:
         if field_meta["name"] == target_name:
-            config["output_features"].append(field_meta["config"])
+            config["output_features"].append(field_meta[CONFIG])
         elif not field_meta["excluded"] and field_meta["mode"] == "input":
-            config["input_features"].append(field_meta["config"])
+            config["input_features"].append(field_meta[CONFIG])
 
     return config
 
 
 def get_field_metadata(
-    fields: List[FieldInfo], row_count: int, resources: Dict, target_name: str = None
+    fields: List[FieldInfo], row_count: int, resources: dict, target_name: str = None
 ) -> list:
     """
     Computes metadata for each field in dataset
@@ -221,7 +240,7 @@ def get_field_metadata(
         sum(
             not meta["excluded"]
             and meta["mode"] == "input"
-            and meta["config"]["type"] != "text"
+            and meta[CONFIG][TYPE] != TEXT
             for meta in metadata
         )
         - 1
@@ -230,7 +249,7 @@ def get_field_metadata(
     # Exclude text fields if no GPUs are available
     if resources['gpu'] == 0:
         for meta in metadata:
-            if input_count > 2 and meta["config"]["type"] == "text":
+            if input_count > 2 and meta[CONFIG][TYPE] == TEXT:
                 # By default, exclude text inputs when there are other candidate inputs
                 meta["excluded"] = True
 
@@ -255,20 +274,20 @@ def infer_type(
     if distinct_values == 2 and (
         missing_value_percent == 0 or field.name == target_name
     ):
-        return "binary"
+        return BINARY
 
     if distinct_values < 20:
         # TODO (tgaddair): come up with something better than this, maybe attempt to fit to Gaussian
         # NOTE (ASN): edge case -- there are less than 20 samples in dataset
-        return "category"
+        return CATEGORY
 
     # add criteria for number of spaces
     if field.avg_words and field.avg_words > 2:
-        return "text"
+        return TEXT
 
     # TODO (ASN): add other modalities (image, etc. )
 
-    return "numerical"
+    return NUMERICAL
 
 
 def should_exclude(field: FieldInfo, row_count: int, target_name: str) -> bool:

ludwig/automl/utils.py

@@ -1,9 +1,23 @@
+import logging
+import sys
 from dataclasses import dataclass
 
-import ray
 from dataclasses_json import LetterCase, dataclass_json
 from pandas import Series
 
+logger = logging.getLogger(__name__)
+
+
+try:
+    import ray
+except ImportError:
+    logger.error(
+        ' ray is not installed. '
+        'In order to use auto_train please run '
+        'pip install ludwig[ray]'
+    )
+    sys.exit(-1)
+
 
 @dataclass_json(letter_case=LetterCase.CAMEL)
 @dataclass

ludwig/constants.py

@@ -124,3 +124,8 @@
 TFRECORD = "tfrecord"
 
 SRC = 'dataset_src'
+
+BATCH_SIZE = 'batch_size'
+LEARNING_RATE = 'learning_rate'
+AUTO = 'auto'
+CONFIG = 'config'