Merge branch 'master' of github.com:ludwig-ai/ludwig into automl_whylogs

.pre-commit-config.yaml

@@ -17,7 +17,6 @@ repos:
       - id: check-ast
       - id: check-byte-order-marker
       - id: check-case-conflict
-      - id: check-docstring-first
       - id: check-executables-have-shebangs
       - id: check-json
       - id: check-toml

LICENSE

@@ -216,3 +216,37 @@ distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
+
+--------------------------------------------------------------------------
+
+Code in ludwig/utils/structural_warnings.py adapted from
+https://github.com/ray-project/ray (Apache-2.0 License)
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+--------------------------------------------------------------------------
+
+Code in ludwig/utils/logging_utils.py adapted from
+https://github.com/ray-project/ray (Apache-2.0 License)
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

ludwig/api.py

@@ -73,6 +73,7 @@
 )
 from ludwig.models.registry import model_type_registry
 from ludwig.schema.model_config import ModelConfig
+from ludwig.types import ModelConfigDict, TrainingSetMetadataDict
 from ludwig.utils import metric_utils
 from ludwig.utils.config_utils import get_preprocessing_params
 from ludwig.utils.data_utils import (
@@ -155,7 +156,7 @@ class PreprocessedDataset:
     training_set: Dataset
     validation_set: Dataset
     test_set: Dataset
-    training_set_metadata: Dict[str, Any]
+    training_set_metadata: TrainingSetMetadataDict
 
     # TODO(daniel): deprecate multiple return value unpacking and indexed access
     def __iter__(self):
@@ -444,7 +445,9 @@ def train(
                 output_directory = model_resume_path
             else:
                 if self.backend.is_coordinator():
-                    logger.info("Model resume path does not exists, " "starting training from scratch")
+                    logger.info(
+                        f"Model resume path '{model_resume_path}' does not exist, starting training from scratch"
+                    )
                 model_resume_path = None
 
         if model_resume_path is None:
@@ -1656,12 +1659,12 @@ def set_logging_level(logging_level: int) -> None:
             set_disable_progressbar(False)
 
     @property
-    def config(self) -> Dict[str, Any]:
+    def config(self) -> ModelConfigDict:
         """Returns the fully-rendered config of this model including default values."""
         return self.config_obj.to_dict()
 
     @config.setter
-    def config(self, user_config: Dict[str, Any]):
+    def config(self, user_config: ModelConfigDict):
         """Updates the config of this model.
 
         WARNING: this can have unexpected results on an already trained model.

ludwig/backend/base.py

@@ -32,6 +32,7 @@
 from ludwig.data.dataset.pandas import PandasDatasetManager
 from ludwig.models.base import BaseModel
 from ludwig.schema.trainer import ECDTrainerConfig, GBMTrainerConfig
+from ludwig.types import HyperoptConfigDict
 from ludwig.utils.fs_utils import get_bytes_obj_from_path
 from ludwig.utils.misc_utils import get_from_registry
 from ludwig.utils.system_utils import Resources
@@ -117,7 +118,7 @@ def get_available_resources(self) -> Resources:
         raise NotImplementedError()
 
     @abstractmethod
-    def max_concurrent_trials(self, hyperopt_config: Dict[str, Any]) -> Union[int, None]:
+    def max_concurrent_trials(self, hyperopt_config: HyperoptConfigDict) -> Union[int, None]:
         raise NotImplementedError()
 
 
@@ -217,7 +218,7 @@ def num_nodes(self) -> int:
     def get_available_resources(self) -> Resources:
         return Resources(cpus=psutil.cpu_count(), gpus=torch.cuda.device_count())
 
-    def max_concurrent_trials(self, hyperopt_config: Dict[str, Any]) -> Union[int, None]:
+    def max_concurrent_trials(self, hyperopt_config: HyperoptConfigDict) -> Union[int, None]:
         # Every trial will be run with Pandas and NO Ray Datasets. Allow Ray Tune to use all the
         # trial resources it wants, because there is no Ray Datasets process to compete with it for CPUs.
         return None

ludwig/backend/horovod.py

@@ -15,7 +15,7 @@
 # ==============================================================================
 
 import time
-from typing import Any, Dict, Union
+from typing import Union
 
 import psutil
 import torch
@@ -26,6 +26,7 @@
 from ludwig.models.base import BaseModel
 from ludwig.models.predictor import Predictor
 from ludwig.trainers.trainer import Trainer
+from ludwig.types import HyperoptConfigDict
 from ludwig.utils.horovod_utils import initialize_horovod
 from ludwig.utils.system_utils import Resources
 from ludwig.utils.torch_utils import initialize_pytorch
@@ -85,6 +86,6 @@ def get_available_resources(self) -> Resources:
 
         return Resources(cpus=cpus, gpus=gpus)
 
-    def max_concurrent_trials(self, hyperopt_config: Dict[str, Any]) -> Union[int, None]:
+    def max_concurrent_trials(self, hyperopt_config: HyperoptConfigDict) -> Union[int, None]:
         # Return None since there is no Ray component
         return None

ludwig/backend/ray.py

@@ -61,6 +61,13 @@
 from ludwig.schema.trainer import ECDTrainerConfig
 from ludwig.trainers.registry import ray_trainers_registry, register_ray_trainer
 from ludwig.trainers.trainer import BaseTrainer, RemoteTrainer
+from ludwig.types import (
+    FeatureConfigDict,
+    HyperoptConfigDict,
+    ModelConfigDict,
+    TrainerConfigDict,
+    TrainingSetMetadataDict,
+)
 from ludwig.utils.data_utils import use_credentials
 from ludwig.utils.dataframe_utils import set_index_name
 from ludwig.utils.fs_utils import get_fs_and_path
@@ -111,7 +118,7 @@ def initialize_horovod():
     return _initialize_horovod()
 
 
-def get_trainer_kwargs(**kwargs) -> Dict[str, Any]:
+def get_trainer_kwargs(**kwargs) -> TrainerConfigDict:
     # Horovod an optional import, so avoid importing at the top.
     from ray.train.horovod import HorovodConfig
 
@@ -186,7 +193,7 @@ def _get_df_engine(processor):
 def train_fn(
     executable_kwargs: Dict[str, Any] = None,
     model_ref: ObjectRef = None,  # noqa: F821
-    training_set_metadata: Dict[str, Any] = None,
+    training_set_metadata: TrainingSetMetadataDict = None,
     features: Dict[str, Dict] = None,
     **kwargs,
 ):
@@ -252,8 +259,8 @@ def tune_batch_size_fn(
     data_loader_kwargs: Dict[str, Any] = None,
     executable_kwargs: Dict[str, Any] = None,
     model: ECD = None,  # noqa: F821
-    ludwig_config: Dict[str, Any] = None,
-    training_set_metadata: Dict[str, Any] = None,
+    ludwig_config: ModelConfigDict = None,
+    training_set_metadata: TrainingSetMetadataDict = None,
     features: Dict[str, Dict] = None,
     **kwargs,
 ) -> int:
@@ -282,11 +289,11 @@ def tune_batch_size_fn(
 @ray.remote(max_calls=1)
 def tune_learning_rate_fn(
     dataset: RayDataset,
-    config: Dict[str, Any],
+    config: ModelConfigDict,
     data_loader_kwargs: Dict[str, Any] = None,
     executable_kwargs: Dict[str, Any] = None,
     model: ECD = None,  # noqa: F821
-    training_set_metadata: Dict[str, Any] = None,
+    training_set_metadata: TrainingSetMetadataDict = None,
     features: Dict[str, Dict] = None,
     **kwargs,
 ) -> float:
@@ -448,7 +455,7 @@ def train_online(self, *args, **kwargs):
 
     def tune_batch_size(
         self,
-        config: Dict[str, Any],
+        config: ModelConfigDict,
         training_set: RayDataset,
         **kwargs,
     ) -> int:
@@ -530,8 +537,8 @@ def shutdown(self):
 def legacy_train_fn(
     trainer: RemoteTrainer = None,
     remote_model: "LudwigModel" = None,  # noqa: F821
-    training_set_metadata: Dict[str, Any] = None,
-    features: Dict[str, Dict] = None,
+    training_set_metadata: TrainingSetMetadataDict = None,
+    features: Dict[str, FeatureConfigDict] = None,
     train_shards: List[DatasetPipeline] = None,
     val_shards: List[DatasetPipeline] = None,
     test_shards: List[DatasetPipeline] = None,
@@ -638,7 +645,7 @@ def shutdown(self):
 def eval_fn(
     predictor_kwargs: Dict[str, Any] = None,
     model_ref: ObjectRef = None,  # noqa: F821
-    training_set_metadata: Dict[str, Any] = None,
+    training_set_metadata: TrainingSetMetadataDict = None,
     features: Dict[str, Dict] = None,
     **kwargs,
 ):
@@ -785,7 +792,7 @@ def get_batch_infer_model(
         predictor_kwargs: Dict[str, Any],
         output_columns: List[str],
         features: Dict[str, Dict],
-        training_set_metadata: Dict[str, Any],
+        training_set_metadata: TrainingSetMetadataDict,
         *args,
         **kwargs,
     ):
@@ -1025,7 +1032,7 @@ def get_available_resources(self) -> Resources:
         resources = ray.cluster_resources()
         return Resources(cpus=resources.get("CPU", 0), gpus=resources.get("GPU", 0))
 
-    def max_concurrent_trials(self, hyperopt_config: Dict[str, Any]) -> Union[int, None]:
+    def max_concurrent_trials(self, hyperopt_config: HyperoptConfigDict) -> Union[int, None]:
         cpus_per_trial = hyperopt_config[EXECUTOR].get(CPU_RESOURCES_PER_TRIAL, 1)
         num_cpus_available = self.get_available_resources().cpus
 

ludwig/benchmarking/artifacts.py

@@ -2,6 +2,7 @@
 from dataclasses import dataclass
 from typing import Any, Dict
 
+from ludwig.types import ModelConfigDict, TrainingSetMetadataDict
 from ludwig.utils.data_utils import load_json, load_yaml
 
 
@@ -14,7 +15,7 @@ class BenchmarkingResult:
     experiment_config: Dict[str, Any]
 
     # The Ludwig config used to run the experiment.
-    ludwig_config: Dict[str, Any]
+    ludwig_config: ModelConfigDict
 
     # The python script that is used to process the config before being used.
     process_config_file: str
@@ -35,7 +36,7 @@ class BenchmarkingResult:
     training_progress: Dict[str, Any]
 
     # Loaded `training_set_metadata.json` file.
-    training_set_metadata: Dict[str, Any]
+    training_set_metadata: TrainingSetMetadataDict
 
 
 def build_benchmarking_result(benchmarking_config: dict, experiment_idx: int):

ludwig/benchmarking/summary_dataclasses.py

@@ -3,12 +3,13 @@
 import os
 from dataclasses import dataclass
 from statistics import mean
-from typing import Any, Dict, List, Optional, Set, Union
+from typing import Dict, List, Optional, Set, Union
 
 import ludwig.modules.metric_modules  # noqa: F401
 from ludwig.benchmarking.utils import format_memory, format_time
 from ludwig.globals import MODEL_HYPERPARAMETERS_FILE_NAME
 from ludwig.modules.metric_registry import get_metric_classes, metric_feature_registry  # noqa: F401
+from ludwig.types import ModelConfigDict
 from ludwig.utils.data_utils import load_json
 
 logger = logging.getLogger()
@@ -82,7 +83,7 @@ class MetricsSummary:
     experiment_local_directory: str
 
     # Full Ludwig config.
-    config: Dict[str, Any]
+    config: ModelConfigDict
 
     # LudwigModel output feature type.
     output_feature_type: str

ludwig/callbacks.py

@@ -18,6 +18,7 @@
 from typing import Any, Callable, Dict, List, Union
 
 from ludwig.api_annotations import PublicAPI
+from ludwig.types import HyperoptConfigDict, ModelConfigDict, TrainingSetMetadataDict
 
 
 @PublicAPI
@@ -30,14 +31,14 @@ def on_cmdline(self, cmd: str, *args: List[str]):
         """
         pass
 
-    def on_preprocess_start(self, config: Dict[str, Any]):
+    def on_preprocess_start(self, config: ModelConfigDict):
         """Called before preprocessing starts.
 
         :param config: The config dictionary.
         """
         pass
 
-    def on_preprocess_end(self, training_set, validation_set, test_set, training_set_metadata: Dict[str, Any]):
+    def on_preprocess_end(self, training_set, validation_set, test_set, training_set_metadata: TrainingSetMetadataDict):
         """Called after preprocessing ends.
 
         :param training_set: The training set.
@@ -95,14 +96,14 @@ def on_hyperopt_finish(self, experiment_name: str):
         # TODO(travis): remove in favor of on_hyperopt_end for naming consistency
         pass
 
-    def on_hyperopt_trial_start(self, parameters: Dict[str, Any]):
+    def on_hyperopt_trial_start(self, parameters: HyperoptConfigDict):
         """Called before the start of each hyperparameter optimization trial.
 
         :param parameters: The complete dictionary of parameters for this hyperparameter optimization experiment.
         """
         pass
 
-    def on_hyperopt_trial_end(self, parameters: Dict[str, Any]):
+    def on_hyperopt_trial_end(self, parameters: HyperoptConfigDict):
         """Called after the end of each hyperparameter optimization trial.
 
         :param parameters: The complete dictionary of parameters for this hyperparameter optimization experiment.
@@ -118,7 +119,7 @@ def should_stop_hyperopt(self):
 
     def on_train_init(
         self,
-        base_config: Dict[str, Any],
+        base_config: ModelConfigDict,
         experiment_directory: str,
         experiment_name: str,
         model_name: str,
@@ -139,7 +140,7 @@ def on_train_init(
     def on_train_start(
         self,
         model,
-        config: Dict[str, Any],
+        config: ModelConfigDict,
         config_fp: Union[str, None],
     ):
         """Called after creation of trainer, before the start of training.

ludwig/collect.py

@@ -28,7 +28,7 @@
 from ludwig.constants import FULL, TEST, TRAINING, VALIDATION
 from ludwig.contrib import add_contrib_callback_args
 from ludwig.globals import LUDWIG_VERSION
-from ludwig.utils.print_utils import logging_level_registry, print_boxed, print_ludwig
+from ludwig.utils.print_utils import get_logging_level_registry, print_boxed, print_ludwig
 from ludwig.utils.strings_utils import make_safe_filename
 
 logger = logging.getLogger(__name__)
@@ -302,7 +302,7 @@ def cli_collect_activations(sys_argv):
     for callback in args.callbacks:
         callback.on_cmdline("collect_activations", *sys_argv)
 
-    args.logging_level = logging_level_registry[args.logging_level]
+    args.logging_level = get_logging_level_registry()[args.logging_level]
     logging.getLogger("ludwig").setLevel(args.logging_level)
     global logger
     logger = logging.getLogger("ludwig.collect")
@@ -358,7 +358,7 @@ def cli_collect_weights(sys_argv):
     for callback in args.callbacks:
         callback.on_cmdline("collect_weights", *sys_argv)
 
-    args.logging_level = logging_level_registry[args.logging_level]
+    args.logging_level = get_logging_level_registry()[args.logging_level]
     logging.getLogger("ludwig").setLevel(args.logging_level)
     global logger
     logger = logging.getLogger("ludwig.collect")
@@ -406,7 +406,7 @@ def cli_collect_summary(sys_argv):
     for callback in args.callbacks:
         callback.on_cmdline("collect_summary", *sys_argv)
 
-    args.logging_level = logging_level_registry[args.logging_level]
+    args.logging_level = get_logging_level_registry()[args.logging_level]
     logging.getLogger("ludwig").setLevel(args.logging_level)
     global logger
     logger = logging.getLogger("ludwig.collect")