Mlflow datasets

llmfoundry/utils/config_utils.py

@@ -7,7 +7,8 @@
 import warnings
 from typing import Any, Dict, Literal, Mapping, Optional, Tuple, Union
 
-from composer.utils import dist
+import mlflow
+from composer.utils import dist, parse_uri
 from omegaconf import DictConfig, ListConfig
 from omegaconf import OmegaConf as om
 
@@ -184,3 +185,85 @@ def log_config(cfg: DictConfig) -> None:
             raise e
         if mlflow.active_run():
             mlflow.log_params(params=om.to_container(cfg, resolve=True))
+            log_dataset_uri(cfg)
+
+
+def parse_source_dataset(cfg: DictConfig):
+    """Parse a run config for dataset information."""
+    data_paths = set()
+
+    for data_split in ['train', 'eval']:
+        split = cfg.get(f'{data_split}_loader', {}).get('dataset',
+                                                        {}).get('split', None)
+        source_dataset_path = cfg.get(f'source_dataset_{data_split}', {})
+
+        # check for Delta table
+        if source_dataset_path and len(source_dataset_path.split('.')) >= 3:
+            data_paths.add(('delta_table', source_dataset_path, data_split))
+        # check for UC volume
+        elif source_dataset_path and source_dataset_path.startswith('/Volumes'):
+            data_paths.add(('uc_volume', source_dataset_path, data_split))
+        # check for HF path
+        elif cfg.get(f'{data_split}_loader', {}).get('dataset',
+                                                     {}).get('hf_name'):
+            hf_path = cfg.get(f'{data_split}_loader', {}).get('dataset',
+                                                              {}).get('hf_name')
+            backend, _, _ = parse_uri(hf_path)
+            if backend:
+                hf_path = f'{hf_path.rstrip("/")}/{split}/' if split else hf_path
+                data_paths.add((backend, hf_path, data_split))
+            else:
+                data_paths.add(('hf', hf_path, data_split))
+        # check for remote path
+        elif cfg.get(f'{data_split}_loader', {}).get('dataset',
+                                                     {}).get('remote', None):
+            remote_path = cfg.get(f'{data_split}_loader',
+                                  {}).get('dataset', {}).get('remote', None)
+            backend, _, _ = parse_uri(remote_path)
+            remote_path = f'{remote_path.rstrip("/")}/{split}/' if split else remote_path
+            data_paths.add((backend, remote_path, data_split))
+        # check for local path
+        elif cfg.get(f'{data_split}_loader', {}).get('dataset',
+                                                     {}).get('local', None):
+            local_path = cfg.get(f'{data_split}_loader',
+                                 {}).get('dataset', {}).get('local', None)
+            split = cfg.get(f'{data_split}_loader',
+                            {}).get('dataset', {}).get('split', None)
+            data_paths.add(('local', local_path, data_split))
+
+    return data_paths
+
+
+def log_dataset_uri(cfg: DictConfig) -> mlflow.data.meta_dataset.MetaDataset:
+    """Extracts dataset information from the provided configuration."""
+    # Figure out which data source to use
+    data_paths = parse_source_dataset(cfg)
+
+    dataset_source_mapping = {
+        's3': mlflow.data.http_dataset_source.HTTPDatasetSource,
+        'oci': mlflow.data.http_dataset_source.HTTPDatasetSource,
+        'https': mlflow.data.http_dataset_source.HTTPDatasetSource,
+        'hf': mlflow.data.huggingface_dataset_source.HuggingFaceDatasetSource,
+        'delta_table': mlflow.data.delta_dataset_source.DeltaDatasetSource,
+        'uc_volume': mlflow.data.uc_volume_dataset_source.UCVolumeDatasetSource,
+        'local': mlflow.data.http_dataset_source.HTTPDatasetSource,
+    }
+
+    for dataset_type, path, split in data_paths:
+        source_class = dataset_source_mapping.get(dataset_type)
+
+        if source_class:
+            if dataset_type == 'delta_table':
+                source = source_class(delta_table_name=path)
+            elif dataset_type == 'hf' or dataset_type == 'uc_volume':
+                source = source_class(path=path)
+            else:
+                source = source_class(url=path)
+        else:
+            log.info(
+                f'{dataset_type} unknown, defaulting to http dataset source')
+            source = mlflow.data.http_dataset_source.HTTPDatasetSource(uri=path)
+
+        mlflow.log_input(
+            mlflow.data.meta_dataset.MetaDataset(
+                source, name=f'{dataset_type} | {split}'))

setup.py

@@ -52,7 +52,7 @@
 
 install_requires = [
     'mosaicml[libcloud,wandb,oci,gcs]>=0.21.1,<0.22',
-    'mlflow>=2.10,<3',
+    'mlflow>=2.12.1,<3',
     'accelerate>=0.25,<0.26',  # for HF inference `device_map`
     'transformers>=4.40,<4.41',
     'mosaicml-streaming>=0.7.5,<0.8',