Add missing dependency group (#1287)

mosaicml · Jun 17, 2024 · f8b2875 · f8b2875
1 parent ca528d5
commit f8b2875
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/scripts/data_prep/convert_delta_to_json.py b/scripts/data_prep/convert_delta_to_json.py
@@ -463,6 +463,7 @@ def validate_and_get_cluster_info(
         if res is None:
             raise ClusterDoesNotExistError(cluster_id)
 
+        assert res.spark_version is not None
         stripped_runtime = re.sub(
             r'[a-zA-Z]',
             '',

diff --git a/setup.py b/setup.py
@@ -54,7 +54,7 @@
 ]
 
 install_requires = [
-    'mosaicml[libcloud,wandb,oci,gcs]>=0.23.2,<0.24',
+    'mosaicml[libcloud,wandb,oci,gcs,mlflow]>=0.23.2,<0.24',
     'mlflow>=2.13.2,<2.14',
     'accelerate>=0.25,<0.26',  # for HF inference `device_map`
     'transformers>=4.40,<4.41',