Merge branch 'main' into enable-vector-stores-files-api-tests

tests/integration/agents/test_openai_responses.py

@@ -6,7 +6,7 @@
 import pytest
 from openai import BadRequestError, OpenAI
 
-from llama_stack.distribution.library_client import LlamaStackAsLibraryClient
+from llama_stack.core.library_client import LlamaStackAsLibraryClient
 
 
 @pytest.fixture

tests/integration/files/test_files.py

@@ -10,8 +10,8 @@ from unittest.mock import patch
 import pytest
 from openai import OpenAI
 
-from llama_stack.distribution.datatypes import User
-from llama_stack.distribution.library_client import LlamaStackAsLibraryClient
+from llama_stack.core.datatypes import User
+from llama_stack.core.library_client import LlamaStackAsLibraryClient
 
 
 def test_openai_client_basic_operations(compat_client, client_with_models):

tests/integration/fixtures/common.py

@@ -20,7 +20,7 @@ from llama_stack_client import LlamaStackClient
 from openai import OpenAI
 
 from llama_stack import LlamaStackAsLibraryClient
-from llama_stack.distribution.stack import run_config_from_adhoc_config_spec
+from llama_stack.core.stack import run_config_from_adhoc_config_spec
 from llama_stack.env import get_env_or_fail
 
 DEFAULT_PORT = 8321

tests/integration/inference/test_openai_completion.py

@@ -14,7 +14,7 @@ from openai import OpenAI
 from reportlab.lib.pagesizes import letter
 from reportlab.pdfgen import canvas
 
-from llama_stack.distribution.library_client import LlamaStackAsLibraryClient
+from llama_stack.core.library_client import LlamaStackAsLibraryClient
 
 from ..test_cases.test_case import TestCase
 

tests/integration/inference/test_openai_embeddings.py

@@ -10,7 +10,7 @@ import struct
 import pytest
 from openai import OpenAI
 
-from llama_stack.distribution.library_client import LlamaStackAsLibraryClient
+from llama_stack.core.library_client import LlamaStackAsLibraryClient
 
 
 def decode_base64_to_floats(base64_string: str) -> list[float]:

tests/integration/post_training/test_post_training.py

@@ -13,6 +13,9 @@ import pytest
 
 from llama_stack.apis.post_training import (
     DataConfig,
+    DatasetFormat,
+    DPOAlignmentConfig,
+    DPOLossType,
     LoraFinetuningConfig,
     TrainingConfig,
 )
@@ -51,6 +54,7 @@ sys.stdout.reconfigure(line_buffering=True)
 #
 
 
+# SFT test
 class TestPostTraining:
     @pytest.mark.integration
     @pytest.mark.parametrize(
@@ -90,7 +94,7 @@ class TestPostTraining:
             dataset_id=dataset.identifier,
             batch_size=1,
             shuffle=False,
-            data_format="instruct",
+            data_format=DatasetFormat.instruct,
         )
 
         # setup training config with minimal settings
@@ -132,6 +136,8 @@ class TestPostTraining:
         artifacts = llama_stack_client.post_training.job.artifacts(job_uuid=job_uuid)
         logger.info(f"Job artifacts: {artifacts}")
 
+        logger.info(f"Registered dataset with ID: {dataset.identifier}")
+
     # TODO: Fix these tests to properly represent the Jobs API in training
     #
     # async def test_get_training_jobs(self, post_training_stack):
@@ -159,3 +165,78 @@ class TestPostTraining:
     #     assert job_artifacts.checkpoints[0].identifier == "instructlab/granite-7b-lab"
     #     assert job_artifacts.checkpoints[0].epoch == 0
     # assert "/.llama/checkpoints/Llama3.2-3B-Instruct-sft-0" in job_artifacts.checkpoints[0].path
+
+    # DPO test
+    @pytest.mark.integration
+    @pytest.mark.parametrize(
+        "purpose, source",
+        [
+            (
+                "post-training/messages",
+                {
+                    "type": "uri",
+                    "uri": "huggingface://datasets/trl-internal-testing/hh-rlhf-helpful-base-trl-style?split=train[:20]",
+                },
+            ),
+        ],
+    )
+    @pytest.mark.timeout(360)
+    def test_preference_optimize(self, llama_stack_client, purpose, source):
+        logger.info("Starting DPO preference optimization test")
+
+        # register preference dataset to train
+        dataset = llama_stack_client.datasets.register(
+            purpose=purpose,
+            source=source,
+        )
+        logger.info(f"Registered preference dataset with ID: {dataset.identifier}")
+
+        # DPO algorithm configuration
+        algorithm_config = DPOAlignmentConfig(
+            beta=0.1,
+            loss_type=DPOLossType.sigmoid,
+        )
+
+        data_config = DataConfig(
+            dataset_id=dataset.identifier,
+            batch_size=1,
+            shuffle=False,
+            data_format=DatasetFormat.dialog,  # DPO datasets often use dialog format
+        )
+
+        # setup training config with minimal settings for DPO
+        training_config = TrainingConfig(
+            n_epochs=1,
+            data_config=data_config,
+            max_steps_per_epoch=1,  # Just 2 steps for quick testing
+            gradient_accumulation_steps=1,
+        )
+
+        job_uuid = f"test-dpo-job-{uuid.uuid4()}"
+        logger.info(f"Starting DPO training job with UUID: {job_uuid}")
+
+        # train with HuggingFace DPO implementation
+        _ = llama_stack_client.post_training.preference_optimize(
+            job_uuid=job_uuid,
+            finetuned_model="distilgpt2",  # Much smaller model for faster CI testing
+            algorithm_config=algorithm_config,
+            training_config=training_config,
+            hyperparam_search_config={},
+            logger_config={},
+        )
+
+        while True:
+            status = llama_stack_client.post_training.job.status(job_uuid=job_uuid)
+            if not status:
+                logger.error("DPO job not found")
+                break
+
+            logger.info(f"Current DPO status: {status}")
+            if status.status == "completed":
+                break
+
+            logger.info("Waiting for DPO job to complete...")
+            time.sleep(10)  # Increased sleep time to reduce polling frequency
+
+        artifacts = llama_stack_client.post_training.job.artifacts(job_uuid=job_uuid)
+        logger.info(f"DPO job artifacts: {artifacts}")

tests/integration/providers/utils/sqlstore/test_authorized_sqlstore.py

@@ -10,8 +10,8 @@ from unittest.mock import patch
 
 import pytest
 
-from llama_stack.distribution.access_control.access_control import default_policy
-from llama_stack.distribution.datatypes import User
+from llama_stack.core.access_control.access_control import default_policy
+from llama_stack.core.datatypes import User
 from llama_stack.providers.utils.sqlstore.api import ColumnType
 from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
 from llama_stack.providers.utils.sqlstore.sqlstore import PostgresSqlStoreConfig, SqliteSqlStoreConfig, sqlstore_impl
@@ -186,7 +186,7 @@ async def test_authorized_store_attributes(mock_get_authenticated_user, authoriz
 @patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
 async def test_user_ownership_policy(mock_get_authenticated_user, authorized_store, request):
     """Test that 'user is owner' policies work correctly with record ownership"""
-    from llama_stack.distribution.access_control.datatypes import AccessRule, Action, Scope
+    from llama_stack.core.access_control.datatypes import AccessRule, Action, Scope
 
     backend_name = request.node.callspec.id
 

tests/integration/tool_runtime/test_mcp.py

@@ -10,7 +10,7 @@ import pytest
 from llama_stack_client import Agent
 
 from llama_stack import LlamaStackAsLibraryClient
-from llama_stack.distribution.datatypes import AuthenticationRequiredError
+from llama_stack.core.datatypes import AuthenticationRequiredError
 
 AUTH_TOKEN = "test-token"
 

tests/integration/vector_io/test_openai_vector_stores.py

@@ -14,7 +14,7 @@ from openai import BadRequestError as OpenAIBadRequestError
 from openai import OpenAI
 
 from llama_stack.apis.vector_io import Chunk
-from llama_stack.distribution.library_client import LlamaStackAsLibraryClient
+from llama_stack.core.library_client import LlamaStackAsLibraryClient
 
 logger = logging.getLogger(__name__)