Merge branch 'main' into nvidia-e2e-notebook

tests/integration/providers/nvidia/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

tests/integration/providers/nvidia/conftest.py

@@ -0,0 +1,14 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+
+import pytest
+
+# Skip all tests in this directory when running in GitHub Actions
+in_github_actions = os.environ.get("GITHUB_ACTIONS") == "true"
+if in_github_actions:
+    pytest.skip("Skipping NVIDIA tests in GitHub Actions environment", allow_module_level=True)

tests/integration/providers/nvidia/test_datastore.py

@@ -0,0 +1,47 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+
+import pytest
+
+# How to run this test:
+#
+# LLAMA_STACK_CONFIG="nvidia" pytest -v tests/integration/providers/nvidia/test_datastore.py
+
+
+# nvidia provider only
+@pytest.mark.parametrize(
+    "provider_id",
+    [
+        "nvidia",
+    ],
+)
+def test_register_and_unregister(llama_stack_client, provider_id):
+    purpose = "eval/messages-answer"
+    source = {
+        "type": "uri",
+        "uri": "hf://datasets/llamastack/simpleqa?split=train",
+    }
+    dataset_id = f"test-dataset-{provider_id}"
+    dataset = llama_stack_client.datasets.register(
+        dataset_id=dataset_id,
+        purpose=purpose,
+        source=source,
+        metadata={"provider_id": provider_id, "format": "json", "description": "Test dataset description"},
+    )
+    assert dataset.identifier is not None
+    assert dataset.provider_id == provider_id
+    assert dataset.identifier == dataset_id
+
+    dataset_list = llama_stack_client.datasets.list()
+    provider_datasets = [d for d in dataset_list if d.provider_id == provider_id]
+    assert any(provider_datasets)
+    assert any(d.identifier == dataset_id for d in provider_datasets)
+
+    llama_stack_client.datasets.unregister(dataset.identifier)
+    dataset_list = llama_stack_client.datasets.list()
+    provider_datasets = [d for d in dataset_list if d.identifier == dataset.identifier]
+    assert not any(provider_datasets)

tests/unit/providers/nvidia/test_datastore.py

@@ -0,0 +1,138 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+import unittest
+from unittest.mock import patch
+
+import pytest
+
+from llama_stack.apis.datasets import Dataset, DatasetPurpose, URIDataSource
+from llama_stack.providers.remote.datasetio.nvidia.config import NvidiaDatasetIOConfig
+from llama_stack.providers.remote.datasetio.nvidia.datasetio import NvidiaDatasetIOAdapter
+
+
+class TestNvidiaDatastore(unittest.TestCase):
+    def setUp(self):
+        os.environ["NVIDIA_DATASETS_URL"] = "http://nemo.test/datasets"
+
+        config = NvidiaDatasetIOConfig(
+            datasets_url=os.environ["NVIDIA_DATASETS_URL"], dataset_namespace="default", project_id="default"
+        )
+        self.adapter = NvidiaDatasetIOAdapter(config)
+        self.make_request_patcher = patch(
+            "llama_stack.providers.remote.datasetio.nvidia.datasetio.NvidiaDatasetIOAdapter._make_request"
+        )
+        self.mock_make_request = self.make_request_patcher.start()
+
+    def tearDown(self):
+        self.make_request_patcher.stop()
+
+    @pytest.fixture(autouse=True)
+    def inject_fixtures(self, run_async):
+        self.run_async = run_async
+
+    def _assert_request(self, mock_call, expected_method, expected_path, expected_json=None):
+        """Helper method to verify request details in mock calls."""
+        call_args = mock_call.call_args
+
+        assert call_args[0][0] == expected_method
+        assert call_args[0][1] == expected_path
+
+        if expected_json:
+            for key, value in expected_json.items():
+                assert call_args[1]["json"][key] == value
+
+    def test_register_dataset(self):
+        self.mock_make_request.return_value = {
+            "id": "dataset-123456",
+            "name": "test-dataset",
+            "namespace": "default",
+        }
+
+        dataset_def = Dataset(
+            identifier="test-dataset",
+            type="dataset",
+            provider_resource_id="",
+            provider_id="",
+            purpose=DatasetPurpose.post_training_messages,
+            source=URIDataSource(uri="https://example.com/data.jsonl"),
+            metadata={"provider_id": "nvidia", "format": "jsonl", "description": "Test dataset description"},
+        )
+
+        self.run_async(self.adapter.register_dataset(dataset_def))
+
+        self.mock_make_request.assert_called_once()
+        self._assert_request(
+            self.mock_make_request,
+            "POST",
+            "/v1/datasets",
+            expected_json={
+                "name": "test-dataset",
+                "namespace": "default",
+                "files_url": "https://example.com/data.jsonl",
+                "project": "default",
+                "format": "jsonl",
+                "description": "Test dataset description",
+            },
+        )
+
+    def test_unregister_dataset(self):
+        self.mock_make_request.return_value = {
+            "message": "Resource deleted successfully.",
+            "id": "dataset-81RSQp7FKX3rdBtKvF9Skn",
+            "deleted_at": None,
+        }
+        dataset_id = "test-dataset"
+
+        self.run_async(self.adapter.unregister_dataset(dataset_id))
+
+        self.mock_make_request.assert_called_once()
+        self._assert_request(self.mock_make_request, "DELETE", "/v1/datasets/default/test-dataset")
+
+    def test_register_dataset_with_custom_namespace_project(self):
+        custom_config = NvidiaDatasetIOConfig(
+            datasets_url=os.environ["NVIDIA_DATASETS_URL"],
+            dataset_namespace="custom-namespace",
+            project_id="custom-project",
+        )
+        custom_adapter = NvidiaDatasetIOAdapter(custom_config)
+
+        self.mock_make_request.return_value = {
+            "id": "dataset-123456",
+            "name": "test-dataset",
+            "namespace": "custom-namespace",
+        }
+
+        dataset_def = Dataset(
+            identifier="test-dataset",
+            type="dataset",
+            provider_resource_id="",
+            provider_id="",
+            purpose=DatasetPurpose.post_training_messages,
+            source=URIDataSource(uri="https://example.com/data.jsonl"),
+            metadata={"format": "jsonl"},
+        )
+
+        self.run_async(custom_adapter.register_dataset(dataset_def))
+
+        self.mock_make_request.assert_called_once()
+        self._assert_request(
+            self.mock_make_request,
+            "POST",
+            "/v1/datasets",
+            expected_json={
+                "name": "test-dataset",
+                "namespace": "custom-namespace",
+                "files_url": "https://example.com/data.jsonl",
+                "project": "custom-project",
+                "format": "jsonl",
+            },
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()