Merge branch 'main' into implement-search-for-PGVector

llama_stack/providers/registry/files.py

@@ -5,9 +5,11 @@
 # the root directory of this source tree.
 
 from llama_stack.providers.datatypes import (
+    AdapterSpec,
     Api,
     InlineProviderSpec,
     ProviderSpec,
+    remote_provider_spec,
 )
 from llama_stack.providers.utils.sqlstore.sqlstore import sql_store_pip_packages
 
@@ -23,4 +25,14 @@ def available_providers() -> list[ProviderSpec]:
             config_class="llama_stack.providers.inline.files.localfs.config.LocalfsFilesImplConfig",
             description="Local filesystem-based file storage provider for managing files and documents locally.",
         ),
+        remote_provider_spec(
+            api=Api.files,
+            adapter=AdapterSpec(
+                adapter_type="s3",
+                pip_packages=["boto3"] + sql_store_pip_packages,
+                module="llama_stack.providers.remote.files.s3",
+                config_class="llama_stack.providers.remote.files.s3.config.S3FilesImplConfig",
+                description="AWS S3-based file storage provider for scalable cloud file management with metadata persistence.",
+            ),
+        ),
     ]

llama_stack/providers/registry/inference.py

@@ -40,8 +40,9 @@ def available_providers() -> list[ProviderSpec]:
         InlineProviderSpec(
             api=Api.inference,
             provider_type="inline::sentence-transformers",
+            # CrossEncoder depends on torchao.quantization
             pip_packages=[
-                "torch torchvision --index-url https://download.pytorch.org/whl/cpu",
+                "torch torchvision torchao>=0.12.0 --extra-index-url https://download.pytorch.org/whl/cpu",
                 "sentence-transformers --no-deps",
             ],
             module="llama_stack.providers.inline.inference.sentence_transformers",

llama_stack/providers/registry/post_training.py

@@ -5,27 +5,50 @@
 # the root directory of this source tree.
 
 
+from typing import cast
+
 from llama_stack.providers.datatypes import AdapterSpec, Api, InlineProviderSpec, ProviderSpec, remote_provider_spec
 
+# We provide two versions of these providers so that distributions can package the appropriate version of torch.
+# The CPU version is used for distributions that don't have GPU support -- they result in smaller container images.
+torchtune_def = dict(
+    api=Api.post_training,
+    pip_packages=["numpy"],
+    module="llama_stack.providers.inline.post_training.torchtune",
+    config_class="llama_stack.providers.inline.post_training.torchtune.TorchtunePostTrainingConfig",
+    api_dependencies=[
+        Api.datasetio,
+        Api.datasets,
+    ],
+    description="TorchTune-based post-training provider for fine-tuning and optimizing models using Meta's TorchTune framework.",
+)
+
 
 def available_providers() -> list[ProviderSpec]:
     return [
         InlineProviderSpec(
-            api=Api.post_training,
-            provider_type="inline::torchtune",
-            pip_packages=["torch", "torchtune==0.5.0", "torchao==0.8.0", "numpy"],
-            module="llama_stack.providers.inline.post_training.torchtune",
-            config_class="llama_stack.providers.inline.post_training.torchtune.TorchtunePostTrainingConfig",
-            api_dependencies=[
-                Api.datasetio,
-                Api.datasets,
-            ],
-            description="TorchTune-based post-training provider for fine-tuning and optimizing models using Meta's TorchTune framework.",
+            **{  # type: ignore
+                **torchtune_def,
+                "provider_type": "inline::torchtune-cpu",
+                "pip_packages": (
+                    cast(list[str], torchtune_def["pip_packages"])
+                    + ["torch torchtune>=0.5.0 torchao>=0.12.0 --extra-index-url https://download.pytorch.org/whl/cpu"]
+                ),
+            },
+        ),
+        InlineProviderSpec(
+            **{  # type: ignore
+                **torchtune_def,
+                "provider_type": "inline::torchtune-gpu",
+                "pip_packages": (
+                    cast(list[str], torchtune_def["pip_packages"]) + ["torch torchtune>=0.5.0 torchao>=0.12.0"]
+                ),
+            },
         ),
         InlineProviderSpec(
             api=Api.post_training,
-            provider_type="inline::huggingface",
-            pip_packages=["torch", "trl", "transformers", "peft", "datasets"],
+            provider_type="inline::huggingface-gpu",
+            pip_packages=["trl", "transformers", "peft", "datasets", "torch"],
             module="llama_stack.providers.inline.post_training.huggingface",
             config_class="llama_stack.providers.inline.post_training.huggingface.HuggingFacePostTrainingConfig",
             api_dependencies=[