adaptive-sdk 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

adaptive_sdk/graphql_client/fragments.py

@@ -240,6 +240,7 @@ class ModelData(BaseModel):
     key: str
     name: str
     online: ModelOnline
+    error: Optional[str]
     is_external: bool = Field(alias='isExternal')
     provider_name: ProviderName = Field(alias='providerName')
     is_adapter: bool = Field(alias='isAdapter')
@@ -340,6 +341,7 @@ class HarmonyGroupData(BaseModel):
     gpu_types: str = Field(alias='gpuTypes')
     created_at: int = Field(alias='createdAt')
     online_models: List['HarmonyGroupDataOnlineModels'] = Field(alias='onlineModels')
+    gpu_allocations: Optional[List['HarmonyGroupDataGpuAllocations']] = Field(alias='gpuAllocations')
 
 class HarmonyGroupDataComputePool(BaseModel):
     """@public"""
@@ -350,6 +352,15 @@ class HarmonyGroupDataOnlineModels(ModelData):
     """@public"""
     pass
 
+class HarmonyGroupDataGpuAllocations(BaseModel):
+    """@public"""
+    name: str
+    num_gpus: int = Field(alias='numGpus')
+    ranks: List[int]
+    created_at: int = Field(alias='createdAt')
+    user_name: Optional[str] = Field(alias='userName')
+    job_id: str = Field(alias='jobId')
+
 class JobData(BaseModel):
     """@public"""
     id: Any
@@ -566,6 +577,7 @@ class ModelDataAdmin(BaseModel):
     key: str
     name: str
     online: ModelOnline
+    error: Optional[str]
     use_cases: List['ModelDataAdminUseCases'] = Field(alias='useCases')
     is_external: bool = Field(alias='isExternal')
     provider_name: ProviderName = Field(alias='providerName')

adaptive_sdk/graphql_client/input_types.py

@@ -65,12 +65,6 @@ class AttachModel(BaseModel):
     wait: bool = False
     'Wait for the model to be deployed or not'
 
-class AzureProviderDataInput(BaseModel):
-    """@private"""
-    api_key: str = Field(alias='apiKey')
-    external_model_id: str = Field(alias='externalModelId')
-    endpoint: str
-
 class CancelAllocationInput(BaseModel):
     """@private"""
     harmony_group: str = Field(alias='harmonyGroup')
@@ -381,11 +375,9 @@ class ModelPlacementInput(BaseModel):
 
 class ModelProviderDataInput(BaseModel):
     """@private"""
-    azure: Optional['AzureProviderDataInput'] = None
     open_ai: Optional['OpenAIProviderDataInput'] = Field(alias='openAI', default=None)
     google: Optional['GoogleProviderDataInput'] = None
     anthropic: Optional['AnthropicProviderDataInput'] = None
-    nvidia: Optional['NvidiaProviderDataInput'] = None
 
 class ModelServiceDisconnect(BaseModel):
     """@private"""
@@ -397,11 +389,6 @@ class ModelServiceFilter(BaseModel):
     model: Optional[str] = None
     capabilities: Optional['CapabilityFilter'] = None
 
-class NvidiaProviderDataInput(BaseModel):
-    """@private"""
-    external_model_id: str = Field(alias='externalModelId')
-    endpoint: str
-
 class OpenAIProviderDataInput(BaseModel):
     """@private"""
     api_key: str = Field(alias='apiKey')
@@ -451,6 +438,10 @@ class SampleConfig(BaseModel):
     selection_type: SelectionTypeInput = Field(alias='selectionType')
     sample_size: Optional[int] = Field(alias='sampleSize', default=None)
 
+class SearchInput(BaseModel):
+    """@private"""
+    query: str
+
 class SystemPromptTemplateCreate(BaseModel):
     """@private"""
     name: str

adaptive_sdk/graphql_client/resize_inference_partition.py

@@ -0,0 +1,6 @@
+from pydantic import Field
+from .base_model import BaseModel
+
+class ResizeInferencePartition(BaseModel):
+    """@public"""
+    resize_inference_partition: str = Field(alias='resizeInferencePartition')

adaptive_sdk/resources/compute_pools.py

@@ -2,6 +2,9 @@
 
 from __future__ import annotations
 from typing import TYPE_CHECKING
+from pydantic import BaseModel
+
+from adaptive_sdk.graphql_client import ResizePartitionInput, HarmonyStatus
 
 from .base_resource import SyncAPIResource, AsyncAPIResource, UseCaseResource
 
@@ -9,6 +12,12 @@ if TYPE_CHECKING:
     from adaptive_sdk.client import Adaptive, AsyncAdaptive
 
 
+class ResizeResult(BaseModel):
+    harmony_group_key: str
+    success: bool
+    error: str | None = None
+
+
 class ComputePools(SyncAPIResource, UseCaseResource):
     """
     Resource to interact with compute pools.
@@ -19,7 +28,33 @@ class ComputePools(SyncAPIResource, UseCaseResource):
         UseCaseResource.__init__(self, client)
 
     def list(self):
-        return self._gql_client.list_compute_pools()
+        return self._gql_client.list_compute_pools().compute_pools
+
+    def resize_inference_partition(self, compute_pool_key: str, size: int) -> list[ResizeResult]:
+        """
+        Resize the inference partitions of all harmony groups in a compute pool.
+        """
+        cps = self.list()
+        found_cp = False
+        for cp in cps.compute_pools:
+            if cp.key == compute_pool_key:
+                selected_cp = cp
+                found_cp = True
+                break
+        if not found_cp:
+            raise ValueError(f"Compute pool with key {compute_pool_key} not found")
+
+        resize_results: list[ResizeResult] = []
+        for hg in selected_cp.harmony_groups:
+            if hg.status == HarmonyStatus.ONLINE:
+                input = ResizePartitionInput(harmonyGroup=hg.key, size=size)
+                try:
+                    _ = self._gql_client.resize_inference_partition(input)
+                    resize_results.append(ResizeResult(harmony_group_key=hg.key, success=True))
+                except Exception as e:
+                    resize_results.append(ResizeResult(harmony_group_key=hg.key, success=False, error_message=str(e)))
+
+        return resize_results
 
 
 class AsyncComputePools(AsyncAPIResource, UseCaseResource):
@@ -32,4 +67,30 @@ class AsyncComputePools(AsyncAPIResource, UseCaseResource):
         UseCaseResource.__init__(self, client)
 
     async def list(self):
-        return await self._gql_client.list_compute_pools()
+        return (await self._gql_client.list_compute_pools()).compute_pools
+
+    async def resize_inference_partition(self, compute_pool_key: str, size: int) -> list[ResizeResult]:
+        """
+        Resize the inference partitions of all harmony groups in a compute pool.
+        """
+        cps = await self.list()
+        found_cp = False
+        for cp in cps.compute_pools:
+            if cp.key == compute_pool_key:
+                selected_cp = cp
+                found_cp = True
+                break
+        if not found_cp:
+            raise ValueError(f"Compute pool with key {compute_pool_key} not found")
+
+        resize_results: list[ResizeResult] = []
+        for hg in selected_cp.harmony_groups:
+            if hg.status == HarmonyStatus.ONLINE:
+                input = ResizePartitionInput(harmonyGroup=hg.key, size=size)
+                try:
+                    _ = await self._gql_client.resize_inference_partition(input)
+                    resize_results.append(ResizeResult(harmony_group_key=hg.key, success=True))
+                except Exception as e:
+                    resize_results.append(ResizeResult(harmony_group_key=hg.key, success=False, error_message=str(e)))
+
+        return resize_results

adaptive_sdk/resources/datasets.py

@@ -2,22 +2,20 @@ from __future__ import annotations
 import json
 import math
 import os
+import time
 from pathlib import Path
-from typing import List, Literal, TYPE_CHECKING
+from typing import List, TYPE_CHECKING
 
 from adaptive_sdk.graphql_client import (
     DatasetCreate,
     Upload,
-    LoadDatasetCreateDataset,
     ListDatasetsDatasets,
     DatasetData,
     DatasetCreateFromMultipartUpload,
+    DatasetUploadProcessingStatusInput,
+    SessionStatus,
 )
-from adaptive_sdk.graphql_client.custom_mutations import Mutation
-from adaptive_sdk.graphql_client.custom_fields import (
-    DatasetFields,
-    DatasetUploadProcessingStatusFields,
-)
+
 from adaptive_sdk.rest import rest_types
 from adaptive_sdk.error_handling import rest_error_handler
 
@@ -29,23 +27,12 @@ if TYPE_CHECKING:
 MIN_CHUNK_SIZE_BYTES = 5 * 1024 * 1024  # 5MB
 MAX_CHUNK_SIZE_BYTES = 100 * 1024 * 1024  # 100MB
 MAX_PARTS_COUNT = 10000
+INIT_CHUNKED_UPLOAD_ROUTE = "/upload/init"
+UPLOAD_PART_ROUTE = "/upload/part"
+ABORT_CHUNKED_UPLOAD_ROUTE = "/upload/abort"
 
 
 def _calculate_upload_parts(file_size: int) -> tuple[int, int]:
-    """
-    Calculate optimal number of parts and chunk size for multipart upload.
-
-    Strategy: Scale chunk size based on file size for optimal performance
-
-    Args:
-        file_size: Size of the file in bytes
-
-    Returns:
-        Tuple of (total_parts, chunk_size_bytes)
-
-    Raises:
-        ValueError: If file is too large to upload with the given constraints
-    """
     if file_size < MIN_CHUNK_SIZE_BYTES:
         raise ValueError(f"File size ({file_size:,} bytes) is too small for chunked upload")
 
@@ -97,19 +84,15 @@ class Datasets(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
         dataset_key: str,
         name: str | None = None,
         use_case: str | None = None,
-    ) -> LoadDatasetCreateDataset:
+    ) -> DatasetData:
         """
-         Upload a dataset from a file. File must be jsonl, where each line should match structure in example below.
+        Upload a dataset from a file. File must be jsonl, where each line should match supported structure.
 
         Args:
             file_path: Path to jsonl file.
             dataset_key: New dataset key.
             name: Optional name to render in UI; if `None`, defaults to same as `dataset_key`.
 
-        Example:
-        ```
-        {"messages": [{"role": "system", "content": "<optional system prompt>"}, {"role": "user", "content": "<user content>"}, {"role": "assistant", "content": "<assistant answer>"}], "completion": "hey"}
-        ```
         """
         file_size = os.path.getsize(file_path)
 
@@ -124,12 +107,8 @@ class Datasets(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
             )
             filename = Path(file_path).stem
             with open(file_path, "rb") as f:
-                file_upload = Upload(
-                    filename=filename, content=f, content_type="application/jsonl"
-                )
-                return self._gql_client.load_dataset(
-                    input=input, file=file_upload
-                ).create_dataset
+                file_upload = Upload(filename=filename, content=f, content_type="application/jsonl")
+                return self._gql_client.load_dataset(input=input, file=file_upload).create_dataset
 
     def _chunked_upload(
         self,
@@ -137,78 +116,73 @@ class Datasets(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
         dataset_key: str,
         name: str | None = None,
         use_case: str | None = None,
-    ) -> LoadDatasetCreateDataset:
+    ) -> DatasetData:
         """Upload large files using chunked upload via REST API."""
         file_size = os.path.getsize(file_path)
         total_parts, chunk_size = _calculate_upload_parts(file_size)
 
-        # Step 1: Initialize chunked upload session
         init_request = rest_types.InitChunkedUploadRequest(
             content_type="application/jsonl",
             metadata=None,
             total_parts_count=total_parts,
         )
-        response = self._rest_client.post(
-            "/upload/init", json=init_request.model_dump()
-        )
+        response = self._rest_client.post(INIT_CHUNKED_UPLOAD_ROUTE, json=init_request.model_dump())
         rest_error_handler(response)
-        init_response = rest_types.InitChunkedUploadResponse.model_validate(
-            response.json()
-        )
+        init_response = rest_types.InitChunkedUploadResponse.model_validate(response.json())
         session_id = init_response.session_id
 
         try:
-            # Step 2: Upload each part
             with open(file_path, "rb") as f:
                 for part_number in range(1, total_parts + 1):
                     chunk_data = f.read(chunk_size)
 
                     response = self._rest_client.post(
-                        "/upload/part",
+                        UPLOAD_PART_ROUTE,
                         params={"session_id": session_id, "part_number": part_number},
                         content=chunk_data,
                         headers={"Content-Type": "application/octet-stream"},
                     )
                     rest_error_handler(response)
 
-            # Step 3: Finalize upload by creating dataset from multipart upload
             input = DatasetCreateFromMultipartUpload(
                 useCase=self.use_case_key(use_case),
                 name=name if name else dataset_key,
                 key=dataset_key,
                 uploadSessionId=session_id,
             )
+            create_dataset_result = self._gql_client.create_dataset_from_multipart_upload(
+                input=input
+            ).create_dataset_from_multipart_upload
+
+            upload_done = False
+            while not upload_done:
+                check_progress_result = self._gql_client.dataset_upload_processing_status(
+                    input=DatasetUploadProcessingStatusInput(
+                        useCase=self.use_case_key(use_case), datasetId=create_dataset_result.dataset_id
+                    )
+                ).dataset_upload_processing_status
+                if check_progress_result.status == SessionStatus.DONE:
+                    upload_done = True
+                elif check_progress_result.status == SessionStatus.ERROR:
+                    raise Exception(f"Upload failed: {check_progress_result.error}")
+                else:
+                    time.sleep(2)
 
-            mutation_field = Mutation.create_dataset_from_multipart_upload(input=input)
-            mutation_field.fields(
-                DatasetUploadProcessingStatusFields.dataset_id,
-                DatasetUploadProcessingStatusFields.total_parts,
-                DatasetUploadProcessingStatusFields.processed_parts,
-                DatasetUploadProcessingStatusFields.progress,
-                DatasetUploadProcessingStatusFields.error,
-                DatasetUploadProcessingStatusFields.status,
-            )
-            result = self._gql_client.mutation(
-                mutation_field, operation_name="CreateDatasetFromMultipartUpload"
-            )
-            return LoadDatasetCreateDataset.model_validate(
-                result["createDatasetFromMultipartUpload"]
-            )
+            dataset_data = self.get(create_dataset_result.dataset_id, use_case=self.use_case_key(use_case))
+            assert dataset_data is not None
+
+            return dataset_data
 
-        except Exception as e:
-            # Abort the upload session on error
+        except Exception:
             try:
-                abort_request = rest_types.AbortChunkedUploadRequest(
-                    session_id=session_id
-                )
-                # delete with body works in runtime but not in type checking
-                self._rest_client.delete(  # type: ignore[call-arg]
-                    "/upload/abort",
-                    content=json.dumps(abort_request.model_dump()),
+                abort_request = rest_types.AbortChunkedUploadRequest(session_id=session_id)
+                self._rest_client.delete(
+                    ABORT_CHUNKED_UPLOAD_ROUTE,
+                    content=json.dumps(abort_request.model_dump()),  # type: ignore[call-arg]
                     headers={"Content-Type": "application/json"},
                 )
             except Exception:
-                pass  # Best effort cleanup
+                pass
             raise
 
     def list(
@@ -227,15 +201,11 @@ class Datasets(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
         Args:
             key: Dataset key.
         """
-        return self._gql_client.describe_dataset(
-            key, self.use_case_key(use_case)
-        ).dataset
+        return self._gql_client.describe_dataset(key, self.use_case_key(use_case)).dataset
 
     def delete(self, key: str, use_case: str | None = None) -> bool:
         """Delete dataset."""
-        return self._gql_client.delete_dataset(
-            id_or_key=key, use_case=self.use_case_key(use_case)
-        ).delete_dataset
+        return self._gql_client.delete_dataset(id_or_key=key, use_case=self.use_case_key(use_case)).delete_dataset
 
 
 class AsyncDatasets(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
@@ -249,7 +219,7 @@ class AsyncDatasets(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
         dataset_key: str,
         name: str | None = None,
         use_case: str | None = None,
-    ) -> LoadDatasetCreateDataset:
+    ) -> DatasetData:
         """
         Upload a dataset from a file. File must be jsonl, where each line should match structure in example below.
 
@@ -277,12 +247,8 @@ class AsyncDatasets(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
             )
             filename = Path(file_path).stem
             with open(file_path, "rb") as f:
-                file_upload = Upload(
-                    filename=filename, content=f, content_type="application/jsonl"
-                )
-                upload_result = await self._gql_client.load_dataset(
-                    input=input, file=file_upload
-                )
+                file_upload = Upload(filename=filename, content=f, content_type="application/jsonl")
+                upload_result = await self._gql_client.load_dataset(input=input, file=file_upload)
                 return upload_result.create_dataset
 
     async def _chunked_upload(
@@ -291,34 +257,28 @@ class AsyncDatasets(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
         dataset_key: str,
         name: str | None = None,
         use_case: str | None = None,
-    ) -> LoadDatasetCreateDataset:
+    ) -> DatasetData:
         """Upload large files using chunked upload via REST API."""
         file_size = os.path.getsize(file_path)
         total_parts, chunk_size = _calculate_upload_parts(file_size)
 
-        # Step 1: Initialize chunked upload session
         init_request = rest_types.InitChunkedUploadRequest(
             content_type="application/jsonl",
             metadata=None,
             total_parts_count=total_parts,
         )
-        response = await self._rest_client.post(
-            "/upload/init", json=init_request.model_dump()
-        )
+        response = await self._rest_client.post(INIT_CHUNKED_UPLOAD_ROUTE, json=init_request.model_dump())
         rest_error_handler(response)
-        init_response = rest_types.InitChunkedUploadResponse.model_validate(
-            response.json()
-        )
+        init_response = rest_types.InitChunkedUploadResponse.model_validate(response.json())
         session_id = init_response.session_id
 
         try:
-            # Step 2: Upload each part
             with open(file_path, "rb") as f:
                 for part_number in range(1, total_parts + 1):
                     chunk_data = f.read(chunk_size)
 
                     response = await self._rest_client.post(
-                        "/upload/part",
+                        UPLOAD_PART_ROUTE,
                         params={"session_id": session_id, "part_number": part_number},
                         content=chunk_data,
                         headers={"Content-Type": "application/octet-stream"},
@@ -331,33 +291,37 @@ class AsyncDatasets(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
                 key=dataset_key,
                 uploadSessionId=session_id,
             )
+            create_dataset_result = (
+                await self._gql_client.create_dataset_from_multipart_upload(input=input)
+            ).create_dataset_from_multipart_upload
+
+            upload_done = False
+            while not upload_done:
+                check_progress_result = (
+                    await self._gql_client.dataset_upload_processing_status(
+                        input=DatasetUploadProcessingStatusInput(
+                            useCase=self.use_case_key(use_case), datasetId=create_dataset_result.dataset_id
+                        )
+                    )
+                ).dataset_upload_processing_status
+                if check_progress_result.status == SessionStatus.DONE:
+                    upload_done = True
+                elif check_progress_result.status == SessionStatus.ERROR:
+                    raise Exception(f"Upload failed: {check_progress_result.error}")
+                else:
+                    time.sleep(2)
 
-            mutation_field = Mutation.create_dataset_from_multipart_upload(input=input)
-            mutation_field.fields(
-                DatasetUploadProcessingStatusFields.dataset_id,
-                DatasetUploadProcessingStatusFields.total_parts,
-                DatasetUploadProcessingStatusFields.processed_parts,
-                DatasetUploadProcessingStatusFields.progress,
-                DatasetUploadProcessingStatusFields.error,
-                DatasetUploadProcessingStatusFields.status,
-            )
-            result = await self._gql_client.mutation(
-                mutation_field, operation_name="CreateDatasetFromMultipartUpload"
-            )
-            return LoadDatasetCreateDataset.model_validate(
-                result["createDatasetFromMultipartUpload"]
-            )
+            dataset_data = await self.get(create_dataset_result.dataset_id, use_case=self.use_case_key(use_case))
+            assert dataset_data is not None
 
-        except Exception as e:
-            # Abort the upload session on error
+            return dataset_data
+
+        except Exception:
             try:
-                abort_request = rest_types.AbortChunkedUploadRequest(
-                    session_id=session_id
-                )
-                # delete with body works in runtime but not in type checking
-                await self._rest_client.delete(  # type: ignore[call-arg]
-                    "/upload/abort",
-                    content=json.dumps(abort_request.model_dump()),
+                abort_request = rest_types.AbortChunkedUploadRequest(session_id=session_id)
+                _ = await self._rest_client.delete(
+                    ABORT_CHUNKED_UPLOAD_ROUTE,
+                    content=json.dumps(abort_request.model_dump()),  # type: ignore[call-arg]
                     headers={"Content-Type": "application/json"},
                 )
             except Exception:
@@ -381,15 +345,11 @@ class AsyncDatasets(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
         Args:
             key: Dataset key.
         """
-        result = await self._gql_client.describe_dataset(
-            key, self.use_case_key(use_case)
-        )
+        result = await self._gql_client.describe_dataset(key, self.use_case_key(use_case))
         return result.dataset
 
     async def delete(self, key: str, use_case: str | None = None) -> bool:
         """Delete dataset."""
         return (
-            await self._gql_client.delete_dataset(
-                id_or_key=key, use_case=self.use_case_key(use_case)
-            )
+            await self._gql_client.delete_dataset(id_or_key=key, use_case=self.use_case_key(use_case))
         ).delete_dataset

adaptive_sdk/resources/models.py

@@ -5,7 +5,6 @@ from adaptive_sdk.graphql_client import (
     OpenAIModel,
     OpenAIProviderDataInput,
     GoogleProviderDataInput,
-    AzureProviderDataInput,
     ModelProviderDataInput,
     AddExternalModelInput,
     ExternalModelProviderName,
@@ -40,13 +39,6 @@ provider_config = {
             google=GoogleProviderDataInput(apiKey=api_key, externalModelId=model_id)
         ),
     },
-    "azure": {
-        "provider_data": lambda api_key, model_id, endpoint: ModelProviderDataInput(
-            azure=AzureProviderDataInput(
-                apiKey=api_key, externalModelId=model_id, endpoint=endpoint
-            )
-        )
-    },
 }
 
 SupportedHFModels = Literal[

{adaptive_sdk-0.1.8.dist-info → adaptive_sdk-0.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: adaptive-sdk
-Version: 0.1.8
+Version: 0.1.10
 Summary: Python SDK for Adaptive Engine
 Author-email: Vincent Debergue <vincent@adaptive-ml.com>, Joao Moura <joao@adaptive-ml.com>, Yacine Bouraoui <yacine@adaptive-ml.com>
 Requires-Python: >=3.10
@@ -10,7 +10,7 @@ Requires-Dist: httpx-sse >= 0.4.0
 Requires-Dist: gql >= 3.5.0
 Requires-Dist: pydantic[email] >= 2.9.0
 Requires-Dist: pyhumps >= 3.8.0
-Requires-Dist: fastapi >= 0.115.8
+Requires-Dist: fastapi >= 0.121.1
 Requires-Dist: uvicorn >= 0.34.0
 Requires-Dist: jsonschema==4.24.0
 Requires-Dist: websockets==15.0.1