From 6250a44ab02e762d2de5c25083313cbe8fed4c88 Mon Sep 17 00:00:00 2001 From: Vincent Chen Date: Thu, 18 Apr 2024 22:25:29 -0700 Subject: [PATCH] optional --- llmfoundry/utils/config_utils.py | 2 +- tests/utils/test_mlflow_logging.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py index 0ae1986398..64134fc205 100644 --- a/llmfoundry/utils/config_utils.py +++ b/llmfoundry/utils/config_utils.py @@ -236,7 +236,7 @@ def parse_source_dataset(cfg: DictConfig): return data_paths -def log_dataset_uri(cfg: DictConfig) -> mlflow.data.meta_dataset.MetaDataset: +def log_dataset_uri(cfg: DictConfig) -> Optional[mlflow.data.meta_dataset.MetaDataset]: """Logs dataset tracking information to MLflow.""" if mlflow is None: log.warning('MLflow is not installed. Skipping dataset logging.') diff --git a/tests/utils/test_mlflow_logging.py b/tests/utils/test_mlflow_logging.py index 356064f9e7..6b0abf83df 100644 --- a/tests/utils/test_mlflow_logging.py +++ b/tests/utils/test_mlflow_logging.py @@ -1,10 +1,10 @@ # Copyright 2024 MosaicML LLM Foundry authors # SPDX-License-Identifier: Apache-2.0 +from typing import Any from unittest.mock import patch import pytest -from typing import Any from omegaconf import OmegaConf from llmfoundry.utils.config_utils import log_dataset_uri, parse_source_dataset