orca-sdk 0.0.96__py3-none-any.whl → 0.0.98__py3-none-any.whl

orca_sdk/conftest.py

@@ -8,6 +8,7 @@ from datasets import ClassLabel, Dataset, Features, Value
 
 from ._utils.auth import _create_api_key, _delete_org
 from .classification_model import ClassificationModel
+from .client import orca_api
 from .credentials import OrcaCredentials
 from .datasource import Datasource
 from .embedding_model import PretrainedEmbeddingModel
@@ -43,6 +44,13 @@ def _create_org_id():
     return "10e50000-0000-4000-a000-" + str(uuid4())[24:]
 
 
+@pytest.fixture()
+def base_url_reset():
+    original_base_url = orca_api.base_url
+    yield
+    orca_api.base_url = original_base_url
+
+
 @pytest.fixture(scope="session")
 def org_id():
     return _create_org_id()
@@ -192,6 +200,8 @@ def writable_memoryset(datasource: Datasource, api_key: str) -> Generator[Labele
         OrcaCredentials.set_api_key(api_key, check_validity=False)
 
         if LabeledMemoryset.exists("test_writable_memoryset"):
+            memoryset.refresh()
+
             memory_ids = [memoryset[i].memory_id for i in range(len(memoryset))]
 
             if memory_ids:

orca_sdk/credentials.py

@@ -1,17 +1,10 @@
 from datetime import datetime
 from typing import Literal, NamedTuple
 
-from ._generated_api_client.api import (
-    check_authentication,
-    create_api_key,
-    delete_api_key,
-    list_api_keys,
-)
-from ._generated_api_client.client import get_base_url, get_headers, set_headers
-from ._generated_api_client.models import (
-    CreateApiKeyRequest,
-    CreateApiKeyRequestScopeItem,
-)
+import httpx
+from httpx import ConnectError, Headers
+
+from .client import orca_api
 
 Scope = Literal["ADMINISTER", "PREDICT"]
 """
@@ -46,7 +39,7 @@ class OrcaCredentials:
         """
         Get the Orca API base URL that is currently being used
         """
-        return get_base_url()
+        return str(orca_api.base_url)
 
     @staticmethod
     def list_api_keys() -> list[ApiKeyInfo]:
@@ -57,8 +50,12 @@ class OrcaCredentials:
             A list of named tuples, with the name and creation date time of the API key
         """
         return [
-            ApiKeyInfo(name=api_key.name, created_at=api_key.created_at, scopes=set(s.value for s in api_key.scope))
-            for api_key in list_api_keys()
+            ApiKeyInfo(
+                name=api_key["name"],
+                created_at=datetime.fromisoformat(api_key["created_at"]),
+                scopes=set(api_key["scope"]),
+            )
+            for api_key in orca_api.GET("/auth/api_key")
         ]
 
     @staticmethod
@@ -70,7 +67,7 @@ class OrcaCredentials:
             True if you are authenticated, False otherwise
         """
         try:
-            return check_authentication()
+            return orca_api.GET("/auth")
         except ValueError as e:
             if "Invalid API key" in str(e):
                 return False
@@ -88,10 +85,11 @@ class OrcaCredentials:
         Returns:
             The secret value of the API key. Make sure to save this value as it will not be shown again.
         """
-        res = create_api_key(
-            body=CreateApiKeyRequest(name=name, scope=[CreateApiKeyRequestScopeItem(scope) for scope in scopes])
+        res = orca_api.POST(
+            "/auth/api_key",
+            json={"name": name, "scope": list(scopes)},
         )
-        return res.api_key
+        return res["api_key"]
 
     @staticmethod
     def revoke_api_key(name: str) -> None:
@@ -104,7 +102,21 @@ class OrcaCredentials:
         Raises:
             ValueError: if the API key is not found
         """
-        delete_api_key(name_or_id=name)
+        orca_api.DELETE("/auth/api_key/{name_or_id}", params={"name_or_id": name})
+
+    @staticmethod
+    def set_headers(headers: dict[str, str]):
+        """
+        Add or override default HTTP headers for all Orca API requests.
+
+        Args:
+            **kwargs: Header names with their string values
+
+        Notes:
+            New keys are merged into the existing headers, this will overwrite headers with the
+            same name, but leave other headers untouched.
+        """
+        orca_api.headers.update(Headers(headers))
 
     @staticmethod
     def set_api_key(api_key: str, check_validity: bool = True):
@@ -121,6 +133,46 @@ class OrcaCredentials:
         Raises:
             ValueError: if the API key is invalid and `check_validity` is True
         """
-        set_headers(get_headers() | {"Api-Key": api_key})
+        OrcaCredentials.set_headers({"Api-Key": api_key})
         if check_validity:
-            check_authentication()
+            orca_api.GET("/auth")
+
+    @staticmethod
+    def set_base_url(base_url: str, check_validity: bool = True):
+        """
+        Set the base URL for the Orca API
+
+        Args:
+            base_url: The base URL to set
+            check_validity: Whether to check if there is an API running at the given base URL
+
+        Raises:
+            ValueError: if there is no healthy API running at the given base URL and `check_validity` is True
+        """
+        # check if the base url is reachable before setting it
+        if check_validity:
+            try:
+                httpx.get(base_url, timeout=1)
+            except ConnectError as e:
+                raise ValueError(f"No API found at {base_url}") from e
+
+        orca_api.base_url = base_url
+
+        # check if the api passes the health check
+        if check_validity:
+            orca_api.GET("/")
+
+    @staticmethod
+    def is_healthy() -> bool:
+        """
+        Check whether the API is healthy
+
+        Returns:
+            True if the API is healthy, False otherwise
+        """
+        try:
+            orca_api.GET("/")
+            orca_api.GET("/gpu/")
+        except Exception:
+            return False
+        return True

orca_sdk/credentials_test.py

@@ -2,6 +2,7 @@ from uuid import uuid4
 
 import pytest
 
+from .client import orca_api
 from .credentials import OrcaCredentials
 
 
@@ -35,3 +36,22 @@ def test_set_invalid_api_key(api_key):
     with pytest.raises(ValueError, match="Invalid API key"):
         OrcaCredentials.set_api_key(str(uuid4()))
     assert not OrcaCredentials.is_authenticated()
+
+
+def test_set_base_url(base_url_reset):
+    OrcaCredentials.set_base_url("http://api.orcadb.ai")
+    assert str(orca_api.base_url) == "http://api.orcadb.ai"
+
+
+def test_set_invalid_base_url():
+    with pytest.raises(ValueError, match="No API found at http://localhost:1582"):
+        OrcaCredentials.set_base_url("http://localhost:1582")
+
+
+def test_is_healthy():
+    assert OrcaCredentials.is_healthy()
+
+
+def test_is_healthy_false(base_url_reset):
+    OrcaCredentials.set_base_url("http://localhost:1582", check_validity=False)
+    assert not OrcaCredentials.is_healthy()

orca_sdk/datasource.py

@@ -4,30 +4,86 @@ import logging
 import tempfile
 import zipfile
 from datetime import datetime
+from io import BytesIO
 from os import PathLike
 from pathlib import Path
-from typing import cast
+from typing import Literal, Union, cast
 
 import pandas as pd
 import pyarrow as pa
-from datasets import Dataset
+from datasets import Dataset, DatasetDict
+from httpx._types import FileTypes  # type: ignore
+from pyarrow import parquet
 from torch.utils.data import DataLoader as TorchDataLoader
 from torch.utils.data import Dataset as TorchDataset
 from tqdm.auto import tqdm
 
-from ._generated_api_client.api import (
-    delete_datasource,
-    get_datasource,
-    list_datasources,
-)
-from ._generated_api_client.api.datasource.create_datasource_datasource_post import (
-    _parse_response as parse_create_response,
-)
-from ._generated_api_client.client import get_client
-from ._generated_api_client.models import ColumnType, DatasourceMetadata
 from ._utils.common import CreateMode, DropMode
-from ._utils.data_parsing import hf_dataset_from_disk, hf_dataset_from_torch
+from ._utils.data_parsing import hf_dataset_from_torch
 from ._utils.tqdm_file_reader import TqdmFileReader
+from .client import DatasourceMetadata, orca_api
+
+
+def _upload_files_to_datasource(
+    name: str,
+    file_paths: list[Path],
+    description: str | None = None,
+) -> DatasourceMetadata:
+    """
+    Helper function to upload files to create a datasource using manual HTTP requests.
+
+    This bypasses the generated client because it doesn't handle file uploads properly.
+
+    Params:
+        name: Name for the datasource
+        file_paths: List of file paths to upload
+        description: Optional description for the datasource
+
+    Returns:
+        Metadata for the created datasource
+    """
+    files: list[tuple[Literal["files"], FileTypes]] = []
+
+    # Calculate total size for all files
+    total_size = sum(file_path.stat().st_size for file_path in file_paths)
+
+    with tqdm(total=total_size, unit="B", unit_scale=True, desc="Uploading") as pbar:
+        for file_path in file_paths:
+            buffered_reader = open(file_path, "rb")
+            tqdm_reader = TqdmFileReader(buffered_reader, pbar)
+            files.append(("files", (file_path.name, cast(bytes, tqdm_reader))))
+
+        # Use manual HTTP request for file uploads
+        metadata = orca_api.POST(
+            "/datasource/upload",
+            files=files,
+            data={"name": name, "description": description},
+        )
+
+    return metadata
+
+
+def _handle_existing_datasource(name: str, if_exists: CreateMode) -> Union["Datasource", None]:
+    """
+    Helper function to handle the common pattern of checking if a datasource exists
+    and taking action based on the if_exists parameter.
+
+    Params:
+        name: Name of the datasource to check
+        if_exists: What to do if a datasource with the same name already exists
+
+    Returns:
+        Datasource instance if opening existing, None if should proceed with creation
+
+    Raises:
+        ValueError: If the datasource already exists and if_exists is "error"
+    """
+    if Datasource.exists(name):
+        if if_exists == "error":
+            raise ValueError(f"Dataset with name {name} already exists")
+        elif if_exists == "open":
+            return Datasource.open(name)
+    return None
 
 
 class Datasource:
@@ -56,19 +112,19 @@ class Datasource:
 
     def __init__(self, metadata: DatasourceMetadata):
         # for internal use only, do not document
-        self.id = metadata.id
-        self.name = metadata.name
-        self.length = metadata.length
-        self.created_at = metadata.created_at
-        self.updated_at = metadata.updated_at
-        self.description = metadata.description
+        self.id = metadata["id"]
+        self.name = metadata["name"]
+        self.length = metadata["length"]
+        self.created_at = datetime.fromisoformat(metadata["created_at"])
+        self.updated_at = datetime.fromisoformat(metadata["updated_at"])
+        self.description = metadata["description"]
         self.columns = {
-            column.name: (
-                f"enum({', '.join(f'{option!r}' for option in column.enum_options) if column.enum_options else ''})"
-                if column.type == ColumnType.ENUM
-                else "str" if column.type == ColumnType.STRING else column.type.value.lower()
+            column["name"]: (
+                f"enum({', '.join(f'{option!r}' for option in column['enum_options'] or []) if 'enum_options' in column else ''})"
+                if column["type"] == "ENUM"
+                else "str" if column["type"] == "STRING" else column["type"].lower()
             )
-            for column in metadata.columns
+            for column in metadata["columns"]
         }
 
     def __eq__(self, other) -> bool:
@@ -94,27 +150,16 @@ class Datasource:
 
         Returns:
             None
-
-        Raises:
-            RuntimeError: If the download fails.
         """
-
-        output_path = Path(output_path)
-        client = get_client().get_httpx_client()
-        url = f"/datasource/{self.id}/download"
-        response = client.get(url)
-        if response.status_code == 404:
-            raise LookupError(f"Datasource {self.id} not found")
-        if response.status_code != 200:
-            raise RuntimeError(f"Failed to download datasource: {response.status_code} {response.text}")
-
+        # TODO: add progress bar to the download
+        response = orca_api.GET("/datasource/{name_or_id}/download", params={"name_or_id": self.id}, parse_as=None)
         with tempfile.NamedTemporaryFile(suffix=".zip") as tmp_zip:
-            tmp_zip.write(response.content)
+            tmp_zip.write(response)
             tmp_zip.flush()
             with zipfile.ZipFile(tmp_zip.name, "r") as zf:
-                output_path.mkdir(parents=True, exist_ok=True)
+                Path(output_path).mkdir(parents=True, exist_ok=True)
                 for file in zf.namelist():
-                    out_file = output_path / Path(file).name
+                    out_file = Path(output_path) / Path(file).name
                     with zf.open(file) as af:
                         out_file.write_bytes(af.read())
 
@@ -138,40 +183,54 @@ class Datasource:
         Raises:
             ValueError: If the datasource already exists and if_exists is `"error"`
         """
-        client = get_client()
-
-        if cls.exists(name):
-            if if_exists == "error":
-                raise ValueError(f"Dataset with name {name} already exists")
-            elif if_exists == "open":
-                return cls.open(name)
+        # Check if datasource already exists and handle accordingly
+        existing = _handle_existing_datasource(name, if_exists)
+        if existing is not None:
+            return existing
 
         with tempfile.TemporaryDirectory() as tmp_dir:
             dataset.save_to_disk(tmp_dir)
-            files = []
 
-            # Calculate total size for all files
+            # Get all file paths in the directory
             file_paths = list(Path(tmp_dir).iterdir())
-            total_size = sum(file_path.stat().st_size for file_path in file_paths)
-
-            with tqdm(total=total_size, unit="B", unit_scale=True, desc="Uploading") as pbar:
-                for file_path in file_paths:
-                    buffered_reader = open(file_path, "rb")
-                    tqdm_reader = TqdmFileReader(buffered_reader, pbar)
-                    files.append(("files", (file_path.name, tqdm_reader)))
-
-                # Do not use Generated client for this endpoint b/c it does not handle files properly
-                metadata = parse_create_response(
-                    response=client.get_httpx_client().request(
-                        method="post",
-                        url="/datasource/",
-                        files=files,
-                        data={"name": name, "description": description},
-                    )
-                )
 
+            # Use the helper function to upload files
+            metadata = _upload_files_to_datasource(name, file_paths, description)
             return cls(metadata=metadata)
 
+    @classmethod
+    def from_hf_dataset_dict(
+        cls,
+        name: str,
+        dataset_dict: DatasetDict,
+        if_exists: CreateMode = "error",
+        description: dict[str, str | None] | str | None = None,
+    ) -> dict[str, Datasource]:
+        """
+        Create datasources from a Hugging Face DatasetDict
+
+        Params:
+            name: Name prefix for the new datasources, will be suffixed with the dataset name
+            dataset_dict: The Hugging Face DatasetDict to create the datasources from
+            if_exists: What to do if a datasource with the same name already exists, defaults to
+                `"error"`. Other option is `"open"` to open the existing datasource.
+            description: Optional description for the datasources, can be a string or a dictionary of dataset names to descriptions
+
+        Returns:
+            A dictionary of datasource handles, keyed by the dataset name
+
+        Raises:
+            ValueError: If a datasource already exists and if_exists is `"error"`
+        """
+        if description is None or isinstance(description, str):
+            description = {dataset_name: description for dataset_name in dataset_dict.keys()}
+        return {
+            dataset_name: cls.from_hf_dataset(
+                f"{name}_{dataset_name}", dataset, if_exists=if_exists, description=description[dataset_name]
+            )
+            for dataset_name, dataset in dataset_dict.items()
+        }
+
     @classmethod
     def from_pytorch(
         cls,
@@ -225,8 +284,16 @@ class Datasource:
         Examples:
             >>> Datasource.from_list("my_datasource", [{"text": "Hello, world!", "label": 1}, {"text": "Goodbye", "label": 0}])
         """
-        hf_dataset = Dataset.from_list(data)
-        return cls.from_hf_dataset(name, hf_dataset, if_exists=if_exists, description=description)
+        # Check if datasource already exists and handle accordingly
+        existing = _handle_existing_datasource(name, if_exists)
+        if existing is not None:
+            return existing
+
+        metadata = orca_api.POST(
+            "/datasource",
+            json={"name": name, "description": description, "content": data},
+        )
+        return cls(metadata=metadata)
 
     @classmethod
     def from_dict(
@@ -251,8 +318,16 @@ class Datasource:
         Examples:
             >>> Datasource.from_dict("my_datasource", {"text": ["Hello, world!", "Goodbye"], "label": [1, 0]})
         """
-        hf_dataset = Dataset.from_dict(data)
-        return cls.from_hf_dataset(name, hf_dataset, if_exists=if_exists, description=description)
+        # Check if datasource already exists and handle accordingly
+        existing = _handle_existing_datasource(name, if_exists)
+        if existing is not None:
+            return existing
+
+        metadata = orca_api.POST(
+            "/datasource",
+            json={"name": name, "description": description, "content": data},
+        )
+        return cls(metadata=metadata)
 
     @classmethod
     def from_pandas(
@@ -274,8 +349,8 @@ class Datasource:
         Raises:
             ValueError: If the datasource already exists and if_exists is `"error"`
         """
-        hf_dataset = Dataset.from_pandas(dataframe)
-        return cls.from_hf_dataset(name, hf_dataset, if_exists=if_exists, description=description)
+        dataset = Dataset.from_pandas(dataframe)
+        return cls.from_hf_dataset(name, dataset, if_exists=if_exists, description=description)
 
     @classmethod
     def from_arrow(
@@ -297,8 +372,23 @@ class Datasource:
         Raises:
             ValueError: If the datasource already exists and if_exists is `"error"`
         """
-        hf_dataset = Dataset(pyarrow_table)
-        return cls.from_hf_dataset(name, hf_dataset, if_exists=if_exists, description=description)
+        # Check if datasource already exists and handle accordingly
+        existing = _handle_existing_datasource(name, if_exists)
+        if existing is not None:
+            return existing
+
+        # Write to bytes buffer
+        buffer = BytesIO()
+        parquet.write_table(pyarrow_table, buffer)
+        parquet_bytes = buffer.getvalue()
+
+        metadata = orca_api.POST(
+            "/datasource/upload",
+            files=[("files", ("data.parquet", parquet_bytes))],
+            data={"name": name, "description": description},
+        )
+
+        return cls(metadata=metadata)
 
     @classmethod
     def from_disk(
@@ -328,16 +418,31 @@ class Datasource:
         Raises:
             ValueError: If the datasource already exists and if_exists is `"error"`
         """
-        hf_dataset = hf_dataset_from_disk(file_path)
-        return cls.from_hf_dataset(name, cast(Dataset, hf_dataset), if_exists=if_exists, description=description)
+        # Check if datasource already exists and handle accordingly
+        existing = _handle_existing_datasource(name, if_exists)
+        if existing is not None:
+            return existing
+
+        file_path = Path(file_path)
+
+        # For dataset directories, use the upload endpoint with multiple files
+        if file_path.is_dir():
+            return cls.from_hf_dataset(
+                name, Dataset.load_from_disk(file_path), if_exists=if_exists, description=description
+            )
+
+        # For single files, use the helper function to upload files
+        metadata = _upload_files_to_datasource(name, [file_path], description)
+
+        return cls(metadata=metadata)
 
     @classmethod
-    def open(cls, name: str) -> Datasource:
+    def open(cls, name_or_id: str) -> Datasource:
         """
         Get a handle to a datasource by name or id in the OrcaCloud
 
         Params:
-            name: The name or unique identifier of the datasource to get
+            name_or_id: The name or unique identifier of the datasource to get
 
         Returns:
             A handle to the existing datasource in the OrcaCloud
@@ -345,7 +450,7 @@ class Datasource:
         Raises:
             LookupError: If the datasource does not exist
         """
-        return cls(get_datasource(name))
+        return cls(orca_api.GET("/datasource/{name_or_id}", params={"name_or_id": name_or_id}))
 
     @classmethod
     def exists(cls, name_or_id: str) -> bool:
@@ -372,7 +477,7 @@ class Datasource:
         Returns:
             A list of all datasource handles in the OrcaCloud
         """
-        return [cls(metadata) for metadata in list_datasources()]
+        return [cls(metadata) for metadata in orca_api.GET("/datasource")]
 
     @classmethod
     def drop(cls, name_or_id: str, if_not_exists: DropMode = "error") -> None:
@@ -388,7 +493,7 @@ class Datasource:
             LookupError: If the datasource does not exist and if_not_exists is `"error"`
         """
         try:
-            delete_datasource(name_or_id)
+            orca_api.DELETE("/datasource/{name_or_id}", params={"name_or_id": name_or_id})
             logging.info(f"Deleted datasource {name_or_id}")
         except LookupError:
             if if_not_exists == "error":