Only require submodlib-py on linux systems

Subset selection only works on Linux and CUDA devices for now, so
don't add it as a requirement except on Linux machines.

Signed-off-by: Ben Browning <bbrownin@redhat.com>

Author: bbrowning

requirements.txt

@@ -14,7 +14,7 @@ openai>=1.13.3,<2.0.0
 numba
 sentencepiece>=0.2.0
 # Note: this dependency has to be built from source
-submodlib-py==0.0.1
+submodlib-py==0.0.1; sys_platform == 'linux'
 tabulate>=0.9.0
 
 # Note: this dependency goes along with langchain-text-splitters and may be

src/instructlab/sdg/subset_selection.py

@@ -14,7 +14,6 @@
 # Third Party
 from datasets import concatenate_datasets, load_dataset
 from jinja2 import BaseLoader, Environment
-from submodlib import FacilityLocationFunction
 from tqdm import tqdm
 import h5py
 import numpy as np
@@ -908,6 +907,10 @@ def process_folds_with_gpu(args):
         testing_mode,
     ) = args
 
+    # Third Party
+    # pylint: disable=import-error, import-outside-toplevel
+    from submodlib import FacilityLocationFunction
+
     try:
         if torch.cuda.is_available():
             torch.cuda.set_device(gpu_id)

tests/functional/test_subset_selection.py

@@ -14,9 +14,6 @@
 import pytest
 import torch
 
-# First Party
-from instructlab.sdg.subset_selection import subset_datasets
-
 
 def create_test_data(num_samples=50):
     """Create synthetic conversation data similar to the real dataset."""
@@ -85,8 +82,14 @@ def map(self, func, iterable):
         return [mock_process_folds_with_gpu(item) for item in iterable]
 
 
+@pytest.mark.gpu
 def test_subset_datasets_functional():
     """Functional test for subset_datasets."""
+
+    # Lazy import down here to not trigger the import except when we're running GPU tests
+    # First Party
+    from instructlab.sdg.subset_selection import subset_datasets
+
     logger = logging.getLogger(__name__)
 
     # Create a mock encoder class