datacrunch 1.14.0__py3-none-any.whl → 1.16.0__py3-none-any.whl

datacrunch/InferenceClient/inference_client.py

@@ -9,14 +9,15 @@ from enum import Enum
 
 class InferenceClientError(Exception):
     """Base exception for InferenceClient errors."""
+
     pass
 
 
 class AsyncStatus(str, Enum):
-    Initialized = "Initialized"
-    Queue = "Queue"
-    Inference = "Inference"
-    Completed = "Completed"
+    Initialized = 'Initialized'
+    Queue = 'Queue'
+    Inference = 'Inference'
+    Completed = 'Completed'
 
 
 @dataclass_json(undefined=Undefined.EXCLUDE)
@@ -38,28 +39,29 @@ class InferenceResponse:
             bool: True if the response is likely a stream, False otherwise
         """
         # Standard chunked transfer encoding
-        is_chunked_transfer = headers.get(
-            'Transfer-Encoding', '').lower() == 'chunked'
+        is_chunked_transfer = headers.get('Transfer-Encoding', '').lower() == 'chunked'
         # Server-Sent Events content type
-        is_event_stream = headers.get(
-            'Content-Type', '').lower() == 'text/event-stream'
+        is_event_stream = headers.get('Content-Type', '').lower() == 'text/event-stream'
         # NDJSON
-        is_ndjson = headers.get(
-            'Content-Type', '').lower() == 'application/x-ndjson'
+        is_ndjson = headers.get('Content-Type', '').lower() == 'application/x-ndjson'
         # Stream JSON
-        is_stream_json = headers.get(
-            'Content-Type', '').lower() == 'application/stream+json'
+        is_stream_json = headers.get('Content-Type', '').lower() == 'application/stream+json'
         # Keep-alive
-        is_keep_alive = headers.get(
-            'Connection', '').lower() == 'keep-alive'
+        is_keep_alive = headers.get('Connection', '').lower() == 'keep-alive'
         # No content length
         has_no_content_length = 'Content-Length' not in headers
 
         # No Content-Length with keep-alive often suggests streaming (though not definitive)
         is_keep_alive_and_no_content_length = is_keep_alive and has_no_content_length
 
-        return (self._stream or is_chunked_transfer or is_event_stream or is_ndjson or
-                is_stream_json or is_keep_alive_and_no_content_length)
+        return (
+            self._stream
+            or is_chunked_transfer
+            or is_event_stream
+            or is_ndjson
+            or is_stream_json
+            or is_keep_alive_and_no_content_length
+        )
 
     def output(self, is_text: bool = False) -> Any:
         try:
@@ -70,9 +72,9 @@ class InferenceResponse:
             # if the response is a stream (check headers), raise relevant error
             if self._is_stream_response(self._original_response.headers):
                 raise InferenceClientError(
-                    f"Response might be a stream, use the stream method instead")
-            raise InferenceClientError(
-                f"Failed to parse response as JSON: {str(e)}")
+                    'Response might be a stream, use the stream method instead'
+                )
+            raise InferenceClientError(f'Failed to parse response as JSON: {str(e)}')
 
     def stream(self, chunk_size: int = 512, as_text: bool = True) -> Generator[Any, None, None]:
         """Stream the response content.
@@ -95,7 +97,9 @@ class InferenceResponse:
 
 
 class InferenceClient:
-    def __init__(self, inference_key: str, endpoint_base_url: str, timeout_seconds: int = 60 * 5) -> None:
+    def __init__(
+        self, inference_key: str, endpoint_base_url: str, timeout_seconds: int = 60 * 5
+    ) -> None:
         """
         Initialize the InferenceClient.
 
@@ -108,23 +112,21 @@ class InferenceClient:
             InferenceClientError: If the parameters are invalid
         """
         if not inference_key:
-            raise InferenceClientError("inference_key cannot be empty")
+            raise InferenceClientError('inference_key cannot be empty')
 
         parsed_url = urlparse(endpoint_base_url)
         if not parsed_url.scheme or not parsed_url.netloc:
-            raise InferenceClientError("endpoint_base_url must be a valid URL")
+            raise InferenceClientError('endpoint_base_url must be a valid URL')
 
         self.inference_key = inference_key
         self.endpoint_base_url = endpoint_base_url.rstrip('/')
-        self.base_domain = self.endpoint_base_url[:self.endpoint_base_url.rindex(
-            '/')]
-        self.deployment_name = self.endpoint_base_url[self.endpoint_base_url.rindex(
-            '/')+1:]
+        self.base_domain = self.endpoint_base_url[: self.endpoint_base_url.rindex('/')]
+        self.deployment_name = self.endpoint_base_url[self.endpoint_base_url.rindex('/') + 1 :]
         self.timeout_seconds = timeout_seconds
         self._session = requests.Session()
         self._global_headers = {
             'Authorization': f'Bearer {inference_key}',
-            'Content-Type': 'application/json'
+            'Content-Type': 'application/json',
         }
 
     def __enter__(self):
@@ -174,9 +176,11 @@ class InferenceClient:
 
     def _build_url(self, path: str) -> str:
         """Construct the full URL by joining the base URL with the path."""
-        return f"{self.endpoint_base_url}/{path.lstrip('/')}"
+        return f'{self.endpoint_base_url}/{path.lstrip("/")}'
 
-    def _build_request_headers(self, request_headers: Optional[Dict[str, str]] = None) -> Dict[str, str]:
+    def _build_request_headers(
+        self, request_headers: Optional[Dict[str, str]] = None
+    ) -> Dict[str, str]:
         """
         Build the final headers by merging global headers with request-specific headers.
 
@@ -211,20 +215,26 @@ class InferenceClient:
             response = self._session.request(
                 method=method,
                 url=self._build_url(path),
-                headers=self._build_request_headers(
-                    kwargs.pop('headers', None)),
+                headers=self._build_request_headers(kwargs.pop('headers', None)),
                 timeout=timeout,
-                **kwargs
+                **kwargs,
             )
             response.raise_for_status()
             return response
         except requests.exceptions.Timeout:
-            raise InferenceClientError(
-                f"Request to {path} timed out after {timeout} seconds")
+            raise InferenceClientError(f'Request to {path} timed out after {timeout} seconds')
         except requests.exceptions.RequestException as e:
-            raise InferenceClientError(f"Request to {path} failed: {str(e)}")
-
-    def run_sync(self, data: Dict[str, Any], path: str = "", timeout_seconds: int = 60 * 5, headers: Optional[Dict[str, str]] = None, http_method: str = "POST", stream: bool = False):
+            raise InferenceClientError(f'Request to {path} failed: {str(e)}')
+
+    def run_sync(
+        self,
+        data: Dict[str, Any],
+        path: str = '',
+        timeout_seconds: int = 60 * 5,
+        headers: Optional[Dict[str, str]] = None,
+        http_method: str = 'POST',
+        stream: bool = False,
+    ):
         """Make a synchronous request to the inference endpoint.
 
         Args:
@@ -242,16 +252,30 @@ class InferenceClient:
             InferenceClientError: If the request fails
         """
         response = self._make_request(
-            http_method, path, json=data, timeout_seconds=timeout_seconds, headers=headers, stream=stream)
+            http_method,
+            path,
+            json=data,
+            timeout_seconds=timeout_seconds,
+            headers=headers,
+            stream=stream,
+        )
 
         return InferenceResponse(
             headers=response.headers,
             status_code=response.status_code,
             status_text=response.reason,
-            _original_response=response
+            _original_response=response,
         )
 
-    def run(self, data: Dict[str, Any], path: str = "", timeout_seconds: int = 60 * 5, headers: Optional[Dict[str, str]] = None, http_method: str = "POST", no_response: bool = False):
+    def run(
+        self,
+        data: Dict[str, Any],
+        path: str = '',
+        timeout_seconds: int = 60 * 5,
+        headers: Optional[Dict[str, str]] = None,
+        http_method: str = 'POST',
+        no_response: bool = False,
+    ):
         """Make an asynchronous request to the inference endpoint.
 
         Args:
@@ -275,44 +299,143 @@ class InferenceClient:
             # If no_response is True, use the "Prefer: respond-async-proxy" header to run async and don't wait for the response
             headers['Prefer'] = 'respond-async-proxy'
             self._make_request(
-                http_method, path, json=data, timeout_seconds=timeout_seconds, headers=headers)
+                http_method,
+                path,
+                json=data,
+                timeout_seconds=timeout_seconds,
+                headers=headers,
+            )
             return
         # Add the "Prefer: respond-async" header to the request, to run async and wait for the response
         headers['Prefer'] = 'respond-async'
 
         response = self._make_request(
-            http_method, path, json=data, timeout_seconds=timeout_seconds, headers=headers)
+            http_method,
+            path,
+            json=data,
+            timeout_seconds=timeout_seconds,
+            headers=headers,
+        )
 
         result = response.json()
         execution_id = result['Id']
 
         return AsyncInferenceExecution(self, execution_id, AsyncStatus.Initialized)
 
-    def get(self, path: str, params: Optional[Dict[str, Any]] = None, headers: Optional[Dict[str, str]] = None, timeout_seconds: Optional[int] = None) -> requests.Response:
-        return self._make_request('GET', path, params=params, headers=headers, timeout_seconds=timeout_seconds)
+    def get(
+        self,
+        path: str,
+        params: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
+        timeout_seconds: Optional[int] = None,
+    ) -> requests.Response:
+        return self._make_request(
+            'GET', path, params=params, headers=headers, timeout_seconds=timeout_seconds
+        )
 
-    def post(self, path: str, json: Optional[Dict[str, Any]] = None, data: Optional[Union[str, Dict[str, Any]]] = None,
-             params: Optional[Dict[str, Any]] = None, headers: Optional[Dict[str, str]] = None, timeout_seconds: Optional[int] = None) -> requests.Response:
-        return self._make_request('POST', path, json=json, data=data, params=params, headers=headers, timeout_seconds=timeout_seconds)
+    def post(
+        self,
+        path: str,
+        json: Optional[Dict[str, Any]] = None,
+        data: Optional[Union[str, Dict[str, Any]]] = None,
+        params: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
+        timeout_seconds: Optional[int] = None,
+    ) -> requests.Response:
+        return self._make_request(
+            'POST',
+            path,
+            json=json,
+            data=data,
+            params=params,
+            headers=headers,
+            timeout_seconds=timeout_seconds,
+        )
 
-    def put(self, path: str, json: Optional[Dict[str, Any]] = None, data: Optional[Union[str, Dict[str, Any]]] = None,
-            params: Optional[Dict[str, Any]] = None, headers: Optional[Dict[str, str]] = None, timeout_seconds: Optional[int] = None) -> requests.Response:
-        return self._make_request('PUT', path, json=json, data=data, params=params, headers=headers, timeout_seconds=timeout_seconds)
+    def put(
+        self,
+        path: str,
+        json: Optional[Dict[str, Any]] = None,
+        data: Optional[Union[str, Dict[str, Any]]] = None,
+        params: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
+        timeout_seconds: Optional[int] = None,
+    ) -> requests.Response:
+        return self._make_request(
+            'PUT',
+            path,
+            json=json,
+            data=data,
+            params=params,
+            headers=headers,
+            timeout_seconds=timeout_seconds,
+        )
 
-    def delete(self, path: str, params: Optional[Dict[str, Any]] = None, headers: Optional[Dict[str, str]] = None, timeout_seconds: Optional[int] = None) -> requests.Response:
-        return self._make_request('DELETE', path, params=params, headers=headers, timeout_seconds=timeout_seconds)
+    def delete(
+        self,
+        path: str,
+        params: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
+        timeout_seconds: Optional[int] = None,
+    ) -> requests.Response:
+        return self._make_request(
+            'DELETE',
+            path,
+            params=params,
+            headers=headers,
+            timeout_seconds=timeout_seconds,
+        )
 
-    def patch(self, path: str, json: Optional[Dict[str, Any]] = None, data: Optional[Union[str, Dict[str, Any]]] = None,
-              params: Optional[Dict[str, Any]] = None, headers: Optional[Dict[str, str]] = None, timeout_seconds: Optional[int] = None) -> requests.Response:
-        return self._make_request('PATCH', path, json=json, data=data, params=params, headers=headers, timeout_seconds=timeout_seconds)
+    def patch(
+        self,
+        path: str,
+        json: Optional[Dict[str, Any]] = None,
+        data: Optional[Union[str, Dict[str, Any]]] = None,
+        params: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
+        timeout_seconds: Optional[int] = None,
+    ) -> requests.Response:
+        return self._make_request(
+            'PATCH',
+            path,
+            json=json,
+            data=data,
+            params=params,
+            headers=headers,
+            timeout_seconds=timeout_seconds,
+        )
 
-    def head(self, path: str, params: Optional[Dict[str, Any]] = None, headers: Optional[Dict[str, str]] = None, timeout_seconds: Optional[int] = None) -> requests.Response:
-        return self._make_request('HEAD', path, params=params, headers=headers, timeout_seconds=timeout_seconds)
+    def head(
+        self,
+        path: str,
+        params: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
+        timeout_seconds: Optional[int] = None,
+    ) -> requests.Response:
+        return self._make_request(
+            'HEAD',
+            path,
+            params=params,
+            headers=headers,
+            timeout_seconds=timeout_seconds,
+        )
 
-    def options(self, path: str, params: Optional[Dict[str, Any]] = None, headers: Optional[Dict[str, str]] = None, timeout_seconds: Optional[int] = None) -> requests.Response:
-        return self._make_request('OPTIONS', path, params=params, headers=headers, timeout_seconds=timeout_seconds)
+    def options(
+        self,
+        path: str,
+        params: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
+        timeout_seconds: Optional[int] = None,
+    ) -> requests.Response:
+        return self._make_request(
+            'OPTIONS',
+            path,
+            params=params,
+            headers=headers,
+            timeout_seconds=timeout_seconds,
+        )
 
-    def health(self, healthcheck_path: str = "/health") -> requests.Response:
+    def health(self, healthcheck_path: str = '/health') -> requests.Response:
         """
         Check the health status of the API.
 
@@ -325,7 +448,7 @@ class InferenceClient:
         try:
             return self.get(healthcheck_path)
         except InferenceClientError as e:
-            raise InferenceClientError(f"Health check failed: {str(e)}")
+            raise InferenceClientError(f'Health check failed: {str(e)}')
 
 
 @dataclass_json(undefined=Undefined.EXCLUDE)
@@ -351,9 +474,15 @@ class AsyncInferenceExecution:
         Returns:
             Dict[str, Any]: The status response containing the execution status and other metadata
         """
-        url = f'{self._inference_client.base_domain}/status/{self._inference_client.deployment_name}'
+        url = (
+            f'{self._inference_client.base_domain}/status/{self._inference_client.deployment_name}'
+        )
         response = self._inference_client._session.get(
-            url, headers=self._inference_client._build_request_headers({self.INFERENCE_ID_HEADER: self.id}))
+            url,
+            headers=self._inference_client._build_request_headers(
+                {self.INFERENCE_ID_HEADER: self.id}
+            ),
+        )
 
         response_json = response.json()
         self._status = AsyncStatus(response_json['Status'])
@@ -366,9 +495,15 @@ class AsyncInferenceExecution:
         Returns:
             Dict[str, Any]: The results of the inference execution
         """
-        url = f'{self._inference_client.base_domain}/result/{self._inference_client.deployment_name}'
+        url = (
+            f'{self._inference_client.base_domain}/result/{self._inference_client.deployment_name}'
+        )
         response = self._inference_client._session.get(
-            url, headers=self._inference_client._build_request_headers({self.INFERENCE_ID_HEADER: self.id}))
+            url,
+            headers=self._inference_client._build_request_headers(
+                {self.INFERENCE_ID_HEADER: self.id}
+            ),
+        )
 
         if response.headers['Content-Type'] == 'application/json':
             return response.json()

datacrunch/__init__.py

@@ -1 +1,3 @@
 from datacrunch.datacrunch import DataCrunchClient
+
+from datacrunch._version import __version__

datacrunch/_version.py

@@ -0,0 +1,6 @@
+try:
+    from importlib.metadata import version
+
+    __version__ = version('datacrunch')
+except Exception:
+    __version__ = '0.0.0+dev'  # fallback for development

datacrunch/authentication/authentication.py

@@ -34,13 +34,12 @@ class AuthenticationService:
         """
         url = self._base_url + TOKEN_ENDPOINT
         payload = {
-            "grant_type": CLIENT_CREDENTIALS,
-            "client_id": self._client_id,
-            "client_secret": self._client_secret
+            'grant_type': CLIENT_CREDENTIALS,
+            'client_id': self._client_id,
+            'client_secret': self._client_secret,
         }
 
-        response = requests.post(
-            url, json=payload, headers=self._generate_headers())
+        response = requests.post(url, json=payload, headers=self._generate_headers())
         handle_error(response)
 
         auth_data = response.json()
@@ -71,13 +70,9 @@ class AuthenticationService:
         """
         url = self._base_url + TOKEN_ENDPOINT
 
-        payload = {
-            "grant_type": REFRESH_TOKEN,
-            "refresh_token": self._refresh_token
-        }
+        payload = {'grant_type': REFRESH_TOKEN, 'refresh_token': self._refresh_token}
 
-        response = requests.post(
-            url, json=payload, headers=self._generate_headers())
+        response = requests.post(url, json=payload, headers=self._generate_headers())
 
         # if refresh token is also expired, authenticate again:
         if response.status_code == 401 or response.status_code == 400:
@@ -98,9 +93,7 @@ class AuthenticationService:
     def _generate_headers(self):
         # get the first 10 chars of the client id
         client_id_truncated = self._client_id[:10]
-        headers = {
-            'User-Agent': 'datacrunch-python-' + client_id_truncated
-        }
+        headers = {'User-Agent': 'datacrunch-python-' + client_id_truncated}
         return headers
 
     def is_expired(self) -> bool:

datacrunch/balance/balance.py

@@ -1,5 +1,3 @@
-from typing import Dict
-
 BALANCE_ENDPOINT = '/balance'
 
 
@@ -49,4 +47,4 @@ class BalanceService:
         :rtype: Balance
         """
         balance = self._http_client.get(BALANCE_ENDPOINT).json()
-        return Balance(balance["amount"], balance["currency"])
+        return Balance(balance['amount'], balance['currency'])

datacrunch/constants.py

@@ -36,12 +36,12 @@ class InstanceStatus:
 
 
 class VolumeStatus:
-    ORDERED = "ordered"
-    CREATING = "creating"
-    ATTACHED = "attached"
-    DETACHED = "detached"
-    DELETING = "deleting"
-    DELETED = "deleted"
+    ORDERED = 'ordered'
+    CREATING = 'creating'
+    ATTACHED = 'attached'
+    DETACHED = 'detached'
+    DELETING = 'deleting'
+    DELETED = 'deleted'
     CLONING = 'cloning'
 
     def __init__(self):
@@ -49,29 +49,31 @@ class VolumeStatus:
 
 
 class VolumeTypes:
-    NVMe = "NVMe"
-    HDD = "HDD"
+    NVMe = 'NVMe'
+    HDD = 'HDD'
 
     def __init__(self):
         return
 
 
 class Locations:
-    FIN_01: str = "FIN-01"
-    ICE_01: str = "ICE-01"
+    FIN_01: str = 'FIN-01'
+    FIN_02: str = 'FIN-02'
+    FIN_03: str = 'FIN-03'
+    ICE_01: str = 'ICE-01'
 
     def __init__(self):
         return
 
 
 class ErrorCodes:
-    INVALID_REQUEST = "invalid_request"
-    UNAUTHORIZED_REQUEST = "unauthorized_request"
-    INSUFFICIENT_FUNDS = "insufficient_funds"
-    FORBIDDEN_ACTION = "forbidden_action"
-    NOT_FOUND = "not_found"
-    SERVER_ERROR = "server_error"
-    SERVICE_UNAVAILABLE = "service_unavailable"
+    INVALID_REQUEST = 'invalid_request'
+    UNAUTHORIZED_REQUEST = 'unauthorized_request'
+    INSUFFICIENT_FUNDS = 'insufficient_funds'
+    FORBIDDEN_ACTION = 'forbidden_action'
+    NOT_FOUND = 'not_found'
+    SERVER_ERROR = 'server_error'
+    SERVICE_UNAVAILABLE = 'service_unavailable'
 
     def __init__(self):
         return