zipher-sdk 0.1.2__tar.gz → 0.1.4__tar.gz

{zipher-sdk-0.1.2 → zipher-sdk-0.1.4}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.1
 Name: zipher-sdk
-Version: 0.1.2
-Summary: Python SDK for accessing the Zipher API
+Version: 0.1.4
+Summary: Python SDK for programmatic access to the Zipher API
 Home-page: UNKNOWN
 Author: Zipher Inc.
 Author-email: yoav@zipher.cloud
@@ -12,5 +12,5 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 License-File: LICENSE
 
-UNKNOWN
+Python SDK for programmatic access to the Zipher API
 

zipher-sdk-0.1.4/README.md

@@ -0,0 +1,58 @@
+
+# Zipher SDK
+
+The Zipher SDK is a Python library for interacting with Zipher's APIs.
+
+## Installation
+
+You can install the Zipher SDK using pip:
+
+```bash
+pip install zipher-sdk
+```
+
+## Usage
+
+Here are some basic examples of how you can use the Zipher SDK to optimize your databricks clusters using Zipher's ML-powered optimization engine:
+
+### Update Existing Configuration
+
+You can update an existing configuration by initializing a zipher Client and sending a JSON payload to the `update_existing_conf` function. Here's how you can do it:
+
+```python
+from zipher import Client
+
+client = Client(customer_id="my_customer_id")  # assuming the zipher API key is stored in ZIPHER_API_KEY environment variable
+
+# Your existing cluster config:
+config_payload = {
+    "new_cluster": {
+        "autoscale": {
+            "min_workers": 1,
+            "max_workers": 30
+        },
+        "cluster_name": "my-cluster",
+        "spark_version": "10.4.x-scala2.12",
+        "spark_conf": {
+            "spark.driver.maxResultSize": "4g"
+        },
+        "aws_attributes": {
+            "first_on_demand": 0,
+            "availability": "SPOT",
+            "zone_id": "auto",
+            "spot_bid_price_percent": 100,
+            "ebs_volume_count": 0
+        },
+        "node_type_id": "rd-fleet.2xlarge",
+        "driver_node_type_id": "rd-fleet.xlarge",
+        "spark_env_vars": {},
+        "enable_elastic_disk": "false"
+    }
+}
+
+# Update configuration
+optimized_cluster = client.update_existing_conf(job_id="my-job-id", existing_conf=config_payload)
+
+# Continue with sending the optimized configuration to Databricks via the Databricks python SDK, Airflow operator, etc.
+
+```

{zipher-sdk-0.1.2 → zipher-sdk-0.1.4}/setup.py

@@ -4,8 +4,9 @@ from setuptools import setup, find_packages
 
 setup(
     name='zipher-sdk',
-    version='0.1.2',
-    description='Python SDK for accessing the Zipher API',
+    version='0.1.4',
+    description='Python SDK for programmatic access to the Zipher API',
+    long_description='Python SDK for programmatic access to the Zipher API',
     author='Zipher Inc.',
     author_email='yoav@zipher.cloud',
     packages=find_packages(),

zipher-sdk-0.1.4/tests/test_client.py

@@ -0,0 +1,39 @@
+import unittest
+from unittest.mock import patch, MagicMock
+from zipher.client import Client
+from zipher.models import ConfFetcherRequest
+
+
+class TestClient(unittest.TestCase):
+
+    def setUp(self):
+        self.customer_id = "test_id"
+        self.api_key = "test_api_key"
+        self.client = Client(customer_id=self.customer_id, zipher_api_key=self.api_key)
+
+    def test_initialization(self):
+        self.assertEqual(self.client.customer_id, self.customer_id)
+        self.assertEqual(self.client.zipher_api_key, self.api_key)
+        # Test initialization with API key as env var
+        with patch('os.getenv', return_value='env_api_key'):
+            client = Client(customer_id=self.customer_id)
+            self.assertEqual(client.zipher_api_key, 'env_api_key')
+
+    @patch('requests.get')
+    def test_get_optimized_config_success(self, mock_get):
+        mock_get.return_value.ok = True
+        mock_get.return_value.json = MagicMock(return_value={'body': '{"key": "value"}'})
+
+        response = self.client.get_optimized_config(job_id="1234")
+        self.assertEqual(response, {"key": "value"})
+
+    @patch('requests.get')
+    def test_get_optimized_config_failure(self, mock_get):
+        mock_get.return_value.ok = False
+        mock_get.return_value.raise_for_status = MagicMock(side_effect=ConnectionError("API Failure"))
+        with self.assertRaises(ConnectionError):
+            self.client.get_optimized_config(job_id="1234")
+
+
+if __name__ == '__main__':
+    unittest.main()

zipher-sdk-0.1.4/zipher/client.py

@@ -0,0 +1,48 @@
+import os
+import json
+import requests
+from zipher.conf import Conf
+from typing import Dict
+from databricks.sdk.service.compute import ClusterSpec
+from zipher.models import ConfFetcherRequest
+
+
+class Client:
+    def __init__(self, customer_id: str, zipher_api_key: str = ''):
+        self.zipher_api_key = zipher_api_key or os.getenv(Conf.zipher_api_key_env_var)
+        self.customer_id = customer_id
+
+    def _call_conf_fetcher_api(self, params: ConfFetcherRequest):
+        headers = {
+            'x-api-key': self.zipher_api_key
+        }
+        response = requests.get(
+            url=Conf.zipher_config_fetcher_api_endpoint,
+            headers=headers,
+            params=params.model_dump()
+        )
+        if response.ok:
+            return json.loads(response.json()['body'])
+        else:
+            response.raise_for_status()
+
+    def get_optimized_config(self, job_id: str) -> Dict:
+        params = ConfFetcherRequest(customer_id=self.customer_id, job_id=job_id)
+        return self._call_conf_fetcher_api(params)
+
+    def get_optimized_config_as_dbx_cluster_spec(self, job_id: str) -> ClusterSpec:
+        return ClusterSpec.from_dict(self.get_optimized_config(job_id))
+
+    def update_existing_conf(self, job_id: str, existing_conf: Dict) -> Dict:
+        if 'new_cluster' in existing_conf:
+            existing_conf = existing_conf['new_cluster']
+        params = ConfFetcherRequest(customer_id=self.customer_id, job_id=job_id, merge_with=json.dumps(existing_conf))
+        return self._call_conf_fetcher_api(params)
+
+        # optimized_spec = self.get_optimized_config_as_dbx_cluster_spec(job_id=job_id)
+        # if 'new_cluster' in existing_conf:
+        #     existing_conf = existing_conf['new_cluster']
+        # existing_spec = ClusterSpec.from_dict(existing_conf)
+        # merged_spec = self.cluster_spec_merger.merge(orig_spec=existing_spec, optimized_spec=optimized_spec)
+        # merged_dict = merged_spec.as_dict()
+

zipher-sdk-0.1.4/zipher/models.py

@@ -0,0 +1,8 @@
+from typing import Optional
+from pydantic import BaseModel
+
+
+class ConfFetcherRequest(BaseModel):
+    job_id: str
+    customer_id: str
+    merge_with: Optional[str] = None

{zipher-sdk-0.1.2 → zipher-sdk-0.1.4}/zipher_sdk.egg-info/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.1
 Name: zipher-sdk
-Version: 0.1.2
-Summary: Python SDK for accessing the Zipher API
+Version: 0.1.4
+Summary: Python SDK for programmatic access to the Zipher API
 Home-page: UNKNOWN
 Author: Zipher Inc.
 Author-email: yoav@zipher.cloud
@@ -12,5 +12,5 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 License-File: LICENSE
 
-UNKNOWN
+Python SDK for programmatic access to the Zipher API
 

{zipher-sdk-0.1.2 → zipher-sdk-0.1.4}/zipher_sdk.egg-info/SOURCES.txt

@@ -1,10 +1,12 @@
 LICENSE
 README.md
 setup.py
+tests/__init__.py
+tests/test_client.py
 zipher/__init__.py
 zipher/client.py
-zipher/cluster_spec_merger.py
 zipher/conf.py
+zipher/models.py
 zipher_sdk.egg-info/PKG-INFO
 zipher_sdk.egg-info/SOURCES.txt
 zipher_sdk.egg-info/dependency_links.txt

{zipher-sdk-0.1.2 → zipher-sdk-0.1.4}/zipher_sdk.egg-info/top_level.txt

@@ -1 +1,2 @@
+tests
 zipher

zipher-sdk-0.1.2/README.md

@@ -1,2 +0,0 @@
-# zipher-sdk
-Zipher's python sdk provides programmatic access to Zipher's api

zipher-sdk-0.1.2/zipher/client.py

@@ -1,40 +0,0 @@
-import os
-import json
-import requests
-from zipher.conf import Conf
-from zipher.cluster_spec_merger import ClusterSpecMerger
-from typing import Dict
-from databricks.sdk.service.compute import ClusterSpec
-
-
-class Client:
-    def __init__(self, customer_id: str, zipher_api_key: str = ''):
-        self.zipher_api_key = zipher_api_key or os.getenv(Conf.zipher_api_key_env_var)
-        self.customer_id = customer_id
-        self.cluster_spec_merger = ClusterSpecMerger()
-
-    def get_optimized_config(self, job_id: str) -> Dict:
-        headers = {
-            'x-api-key': self.zipher_api_key
-        }
-        params = {
-            'customer_id': self.customer_id,
-            'job_id': job_id
-        }
-        response = requests.get(url=Conf.zipher_config_fetcher_api_endpoint, headers=headers, params=params)
-        if response.ok:
-            return json.loads(response.json()['body'])
-        else:
-            response.raise_for_status()
-
-    def get_optimized_config_as_dbx_cluster_spec(self, job_id: str) -> ClusterSpec:
-        return ClusterSpec.from_dict(self.get_optimized_config(job_id))
-
-    def update_existing_conf(self, job_id: str, existing_conf: Dict) -> Dict:
-        optimized_spec = self.get_optimized_config_as_dbx_cluster_spec(job_id=job_id)
-        if 'new_cluster' in existing_conf:
-            existing_conf = existing_conf['new_cluster']
-        existing_spec = ClusterSpec.from_dict(existing_conf)
-        merged_spec = self.cluster_spec_merger.merge(orig_spec=existing_spec, optimized_spec=optimized_spec)
-        return merged_spec.as_dict()
-

zipher-sdk-0.1.2/zipher/cluster_spec_merger.py

@@ -1,64 +0,0 @@
-from databricks.sdk.service.compute import ClusterSpec, AwsAttributes
-
-
-class ClusterSpecMerger:
-    def __init__(self):
-        self.orig_spec: ClusterSpec = None
-        self.optimized_spec: ClusterSpec = None
-
-    def merge(self, orig_spec: ClusterSpec, optimized_spec: ClusterSpec) -> ClusterSpec:
-        self.orig_spec = orig_spec
-        self.optimized_spec = optimized_spec
-
-        self._merge_autoscaling_config()
-        self._merge_spark_cong_config()
-        self._merge_aws_attrs()
-        self._merge_node_types()
-        self._merge_spark_env_vars()
-
-        return self.orig_spec
-
-    def _merge_autoscaling_config(self):
-        if self.orig_spec.autoscale is None:
-            self.orig_spec.autoscale = self.optimized_spec.autoscale
-        else:
-            if self.optimized_spec.autoscale is not None:
-                self.orig_spec.autoscale.min_workers = self.optimized_spec.autoscale.min_workers
-                self.orig_spec.autoscale.max_workers = self.optimized_spec.autoscale.max_workers
-
-    def _merge_spark_cong_config(self):
-        self.orig_spec.spark_conf = self.orig_spec.spark_conf or {}
-        if self.optimized_spec.spark_conf:
-            self.orig_spec.spark_conf.update(self.optimized_spec.spark_conf)
-
-    def _merge_aws_attrs(self):
-        self.orig_spec.aws_attributes = self.orig_spec.aws_attributes or AwsAttributes()
-        if self.optimized_spec.aws_attributes:
-            self.orig_spec.aws_attributes.ebs_volume_count = (self.optimized_spec.aws_attributes.ebs_volume_count or
-                                                              self.orig_spec.aws_attributes.ebs_volume_count)
-            self.orig_spec.aws_attributes.ebs_volume_size = (self.optimized_spec.aws_attributes.ebs_volume_size or
-                                                              self.orig_spec.aws_attributes.ebs_volume_size)
-            self.orig_spec.aws_attributes.ebs_volume_throughput = (
-                    self.optimized_spec.aws_attributes.ebs_volume_throughput or
-                    self.orig_spec.aws_attributes.ebs_volume_throughput)
-            self.orig_spec.aws_attributes.ebs_volume_iops = (
-                    self.optimized_spec.aws_attributes.ebs_volume_iops or
-                    self.orig_spec.aws_attributes.ebs_volume_iops)
-            self.orig_spec.aws_attributes.ebs_volume_type = (
-                    self.optimized_spec.aws_attributes.ebs_volume_type or
-                    self.orig_spec.aws_attributes.ebs_volume_type)
-            self.orig_spec.aws_attributes.zone_id = (
-                    self.optimized_spec.aws_attributes.zone_id or
-                    self.orig_spec.aws_attributes.zone_id)
-            self.orig_spec.aws_attributes.first_on_demand = (
-                    self.optimized_spec.aws_attributes.first_on_demand or
-                    self.orig_spec.aws_attributes.first_on_demand)
-
-    def _merge_node_types(self):
-        self.orig_spec.node_type_id = self.optimized_spec.node_type_id or self.orig_spec.node_type_id
-        self.orig_spec.driver_node_type_id = self.optimized_spec.driver_node_type_id or self.orig_spec.driver_node_type_id
-
-    def _merge_spark_env_vars(self):
-        self.orig_spec.spark_env_vars = self.orig_spec.spark_env_vars or {}
-        if self.optimized_spec.spark_env_vars:
-            self.orig_spec.spark_env_vars.update(self.optimized_spec.spark_env_vars)