chore: import upstream snapshot with attribution

nemo/utils/loggers/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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
+#
+#     http://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.
+
+from nemo.utils.loggers.clearml_logger import ClearMLLogger, ClearMLParams
+from nemo.utils.loggers.dllogger import DLLogger, DLLoggerParams
+from nemo.utils.loggers.mlflow_logger import MLFlowParams

nemo/utils/loggers/clearml_logger.py

@@ -0,0 +1,191 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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
+#
+#     http://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.
+
+import os
+from argparse import Namespace
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, List, Literal, Mapping, Optional, Union
+
+import pandas as pd
+from lightning.pytorch.callbacks import Checkpoint
+from lightning.pytorch.loggers import Logger
+from lightning.pytorch.utilities.parsing import AttributeDict
+from lightning_utilities.core.apply_func import apply_to_collection
+from omegaconf import DictConfig, ListConfig, OmegaConf
+from torch import Tensor
+
+from nemo.utils import logging
+
+try:
+    from clearml import OutputModel, Task
+
+    HAVE_CLEARML_LOGGER = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_CLEARML_LOGGER = False
+
+
+@dataclass
+class ClearMLParams:  # pylint: disable=C0115
+    project: Optional[str] = None
+    task: Optional[str] = None
+    connect_pytorch: Optional[bool] = False
+    model_name: Optional[str] = None
+    tags: Optional[List[str]] = None
+    log_model: Optional[bool] = False
+    log_cfg: Optional[bool] = False
+    log_metrics: Optional[bool] = False
+
+
+class ClearMLLogger(Logger):  # pylint: disable=C0115
+    @property
+    def name(self) -> str:  # pylint: disable=C0116
+        return self.clearml_task.name
+
+    @property
+    def version(self) -> str:  # pylint: disable=C0116
+        return self.clearml_task.id
+
+    def __init__(
+        self, clearml_cfg: DictConfig, log_dir: str, prefix: str, save_best_model: bool, postfix: str = ".nemo"
+    ) -> None:  # pylint: disable=C0116
+        if not HAVE_CLEARML_LOGGER:
+            raise ImportError(
+                "Found create_clearml_logger is True."
+                "But ClearML not found. Please see the README for installation instructions:"
+                "https://github.com/clearml/clearml"
+            )
+
+        self.clearml_task = None
+        self.clearml_model = None
+        self.clearml_cfg = clearml_cfg
+        self.path_nemo_model = os.path.abspath(
+            os.path.expanduser(os.path.join(log_dir, "checkpoints", prefix + postfix))
+        )
+        self.save_best_model = save_best_model
+        self.prefix = prefix
+        self.previos_best_model_path = None
+        self.last_metrics = None
+        self.save_blocked = True
+
+        self.project_name = os.getenv("CLEARML_PROJECT", clearml_cfg.project if clearml_cfg.project else "NeMo")
+        self.task_name = os.getenv("CLEARML_TASK", clearml_cfg.task if clearml_cfg.task else f"Trainer {self.prefix}")
+
+        tags = ["NeMo"]
+        if clearml_cfg.tags:
+            tags.extend(clearml_cfg.tags)
+
+        self.clearml_task: Task = Task.init(
+            project_name=self.project_name,
+            task_name=self.task_name,
+            auto_connect_frameworks={"pytorch": clearml_cfg.connect_pytorch},
+            output_uri=True,
+            tags=tags,
+        )
+
+        if clearml_cfg.model_name:
+            model_name = clearml_cfg.model_name
+        elif self.prefix:
+            model_name = self.prefix
+        else:
+            model_name = self.task_name
+
+        if clearml_cfg.log_model:
+            self.clearml_model: OutputModel = OutputModel(
+                name=model_name, task=self.clearml_task, tags=tags, framework="NeMo"
+            )
+
+    def log_hyperparams(self, params, *args, **kwargs) -> None:  # pylint: disable=C0116
+        if self.clearml_model and self.clearml_cfg.log_cfg:
+            if isinstance(params, Namespace):
+                params = vars(params)
+            elif isinstance(params, AttributeDict):
+                params = dict(params)
+            params = apply_to_collection(params, (DictConfig, ListConfig), OmegaConf.to_container, resolve=True)
+            params = apply_to_collection(params, Path, str)
+            params = OmegaConf.to_yaml(params)
+            self.clearml_model.update_design(config_text=params)
+
+    def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None) -> None:  # pylint: disable=C0116
+        if self.clearml_model and self.clearml_cfg.log_metrics:
+            metrics = {
+                k: {
+                    "value": str(v.item() if type(v) == Tensor else v),
+                    "type": str(type(v.item() if type(v) == Tensor else v)),
+                }
+                for k, v in metrics.items()
+            }
+            self.last_metrics = metrics
+
+    # pylint: disable=C0116
+    def log_table(
+        self,
+        key: str,
+        columns: List[str] = None,
+        data: List[List[Any]] = None,
+        dataframe: Any = None,
+        step: Optional[int] = None,
+    ) -> None:
+        table: Optional[Union[pd.DataFrame, List[List[Any]]]] = None
+
+        if dataframe is not None:
+            table = dataframe
+            if columns is not None:
+                table.columns = columns
+
+        if data is not None:
+            table = data
+            assert len(columns) == len(table[0]), "number of column names should match the total number of columns"
+            table.insert(0, columns)
+
+        if table is not None:
+            self.clearml_task.logger.report_table(title=key, series=key, iteration=step, table_plot=table)
+
+    def after_save_checkpoint(self, checkpoint_callback: Checkpoint) -> None:  # pylint: disable=C0116
+        if self.clearml_model:
+            if self.save_best_model:
+                if self.save_blocked:
+                    self.save_blocked = False
+                    return None
+                if not os.path.exists(checkpoint_callback.best_model_path):
+                    return None
+                if self.previos_best_model_path == checkpoint_callback.best_model_path:
+                    return None
+                self.previos_best_model_path = checkpoint_callback.best_model_path
+            self._log_model(self.path_nemo_model)
+
+    def finalize(self, status: Literal["success", "failed", "aborted"] = "success") -> None:  # pylint: disable=C0116
+        if status == "success":
+            self.clearml_task.mark_completed()
+        elif status == "failed":
+            self.clearml_task.mark_failed()
+        elif status == "aborted":
+            self.clearml_task.mark_stopped()
+
+    def _log_model(self, save_path: str) -> None:  # pylint: disable=C0116
+        if self.clearml_model:
+            if os.path.exists(save_path):
+                self.clearml_model.update_weights(
+                    weights_filename=save_path,
+                    upload_uri=self.clearml_task.storage_uri or self.clearml_task._get_default_report_storage_uri(),
+                    auto_delete_file=False,
+                    is_package=True,
+                )
+
+                if self.clearml_cfg.log_metrics and self.last_metrics:
+                    self.clearml_model.set_all_metadata(self.last_metrics)
+
+                self.save_blocked = True
+            else:
+                logging.warning((f"Logging model enabled, but cant find .nemo file!" f" Path: {save_path}"))

nemo/utils/loggers/dllogger.py

@@ -0,0 +1,104 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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
+#
+#     http://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.
+
+from argparse import Namespace
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+from lightning.pytorch.loggers import Logger
+from lightning.pytorch.utilities import rank_zero_only
+from lightning.pytorch.utilities.parsing import AttributeDict
+from lightning_utilities.core.apply_func import apply_to_collection
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+from nemo.utils import logging
+
+try:
+    import dllogger
+    from dllogger import Verbosity
+
+    HAVE_DLLOGGER = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_DLLOGGER = False
+
+try:
+    from lightning.fabric.utilities.logger import _convert_params, _flatten_dict, _sanitize_callable_params
+
+    PL_LOGGER_UTILITIES = True
+except (ImportError, ModuleNotFoundError):
+    PL_LOGGER_UTILITIES = False
+
+
+@dataclass
+class DLLoggerParams:
+    verbose: Optional[bool] = False
+    stdout: Optional[bool] = False
+    json_file: Optional[str] = "./dllogger.json"
+
+
+class DLLogger(Logger):
+    @property
+    def name(self):
+        return self.__class__.__name__
+
+    @property
+    def version(self):
+        return None
+
+    def __init__(self, stdout: bool, verbose: bool, json_file: str):
+        if not HAVE_DLLOGGER:
+            raise ImportError(
+                "DLLogger was not found. Please see the README for installation instructions: "
+                "https://github.com/NVIDIA/dllogger"
+            )
+        if not PL_LOGGER_UTILITIES:
+            raise ImportError(
+                "DLLogger utilities were not found. You probably need to update PyTorch Lightning>=1.9.0. "
+                "pip install pytorch-lightning -U"
+            )
+        verbosity = Verbosity.VERBOSE if verbose else Verbosity.DEFAULT
+        backends = []
+        if json_file:
+            Path(json_file).parent.mkdir(parents=True, exist_ok=True)
+            backends.append(dllogger.JSONStreamBackend(verbosity, json_file))
+        if stdout:
+            backends.append(dllogger.StdOutBackend(verbosity))
+
+        if not backends:
+            logging.warning(
+                "Neither stdout nor json_file DLLogger parameters were specified." "DLLogger will not log anything."
+            )
+        dllogger.init(backends=backends)
+
+    @rank_zero_only
+    def log_hyperparams(self, params, *args, **kwargs):
+        if isinstance(params, Namespace):
+            params = vars(params)
+        elif isinstance(params, AttributeDict):
+            params = dict(params)
+        params = apply_to_collection(params, (DictConfig, ListConfig), OmegaConf.to_container, resolve=True)
+        params = apply_to_collection(params, Path, str)
+        params = _sanitize_callable_params(_flatten_dict(_convert_params(params)))
+        dllogger.log(step="PARAMETER", data=params)
+
+    @rank_zero_only
+    def log_metrics(self, metrics, step=None):
+        if step is None:
+            step = tuple()
+
+        dllogger.log(step=step, data=metrics)
+
+    def save(self):
+        dllogger.flush()

nemo/utils/loggers/mlflow_logger.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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
+#
+#     http://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.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class MLFlowParams:
+    """ML Flow Configuration Dataclass."""
+
+    # name of experiment, if none, defaults to the globally set experiment name
+    experiment_name: Optional[str] = None
+    run_name: Optional[str] = None
+    # if no run_name is set, it's set by version
+    # local or remote tracking seerver. If tracking_uri is not set, it defaults to save_dir
+    tracking_uri: Optional[str] = None
+    tags: Optional[Dict[str, Any]] = None
+    save_dir: Optional[str] = "./mlruns"
+    prefix: str = ""
+    artifact_location: Optional[str] = None
+    # provide run_id if resuming a previously started run
+    run_id: Optional[str] = None
+    # Log checkpoints created by ModelCheckpoint as MLFlow artifacts.
+    log_model: bool = False