datachain 0.30.7__py3-none-any.whl → 0.31.0__py3-none-any.whl

datachain/remote/studio.py

@@ -1,4 +1,3 @@
-import base64
 import json
 import logging
 import os
@@ -7,6 +6,7 @@ from datetime import datetime, timedelta, timezone
 from struct import unpack
 from typing import (
     Any,
+    BinaryIO,
     Generic,
     Optional,
     TypeVar,
@@ -30,8 +30,9 @@ DatasetExportStatus = Optional[dict[str, Any]]
 DatasetExportSignedUrls = Optional[list[str]]
 FileUploadData = Optional[dict[str, Any]]
 JobData = Optional[dict[str, Any]]
-JobListData = dict[str, Any]
-ClusterListData = dict[str, Any]
+JobListData = list[dict[str, Any]]
+ClusterListData = list[dict[str, Any]]
+
 logger = logging.getLogger("datachain")
 
 DATASET_ROWS_CHUNK_SIZE = 8192
@@ -239,6 +240,45 @@ class StudioClient:
 
         return Response(data, ok, message, response.status_code)
 
+    def _send_multipart_request(
+        self, route: str, files: dict[str, Any], params: Optional[dict[str, Any]] = None
+    ) -> Response[Any]:
+        """
+        Function that communicates with Studio API using multipart/form-data.
+        It will raise an exception, and try to retry, if 5xx status code is
+        returned, or if Timeout exceptions is thrown from the requests lib
+        """
+        import requests
+
+        # Add team_name to params
+        request_params = {**(params or {}), "team_name": self.team}
+
+        response = requests.post(
+            url=f"{self.url}/{route}",
+            files=files,
+            params=request_params,
+            headers={
+                "Authorization": f"token {self.token}",
+            },
+            timeout=self.timeout,
+        )
+
+        ok = response.ok
+        try:
+            data = json.loads(response.content.decode("utf-8"))
+        except json.decoder.JSONDecodeError:
+            data = {}
+
+        if not ok:
+            if response.status_code == 403:
+                message = f"Not authorized for the team {self.team}"
+            else:
+                message = data.get("message", "")
+        else:
+            message = ""
+
+        return Response(data, ok, message, response.status_code)
+
     @staticmethod
     def _unpacker_hook(code, data):
         import msgpack
@@ -409,12 +449,13 @@ class StudioClient:
             method="GET",
         )
 
-    def upload_file(self, content: bytes, file_name: str) -> Response[FileUploadData]:
-        data = {
-            "file_content": base64.b64encode(content).decode("utf-8"),
-            "file_name": file_name,
-        }
-        return self._send_request("datachain/upload-file", data)
+    def upload_file(
+        self, file_obj: BinaryIO, file_name: str
+    ) -> Response[FileUploadData]:
+        # Prepare multipart form data
+        files = {"file": (file_name, file_obj, "application/octet-stream")}
+
+        return self._send_multipart_request("datachain/jobs/files", files)
 
     def create_job(
         self,
@@ -449,25 +490,27 @@ class StudioClient:
             "cron_expression": cron,
             "credentials_name": credentials_name,
         }
-        return self._send_request("datachain/job", data)
+        return self._send_request("datachain/jobs/", data)
 
     def get_jobs(
         self,
         status: Optional[str] = None,
         limit: int = 20,
+        job_id: Optional[str] = None,
     ) -> Response[JobListData]:
-        return self._send_request(
-            "datachain/jobs",
-            {"status": status, "limit": limit} if status else {"limit": limit},
-            method="GET",
-        )
+        params: dict[str, Any] = {"limit": limit}
+        if status is not None:
+            params["status"] = status
+        if job_id is not None:
+            params["job_id"] = job_id
+        return self._send_request("datachain/jobs/", params, method="GET")
 
     def cancel_job(
         self,
         job_id: str,
     ) -> Response[JobData]:
-        url = f"datachain/job/{job_id}/cancel"
+        url = f"datachain/jobs/{job_id}/cancel"
         return self._send_request(url, data={}, method="POST")
 
     def get_clusters(self) -> Response[ClusterListData]:
-        return self._send_request("datachain/clusters", {}, method="GET")
+        return self._send_request("datachain/clusters/", {}, method="GET")

datachain/studio.py

@@ -403,14 +403,14 @@ def create_job(
     if not response.data:
         raise DataChainError("Failed to create job")
 
-    job_id = response.data.get("job", {}).get("id")
+    job_id = response.data.get("id")
 
     if parsed_start_time or cron:
         print(f"Job {job_id} is scheduled as a task in Studio.")
         return 0
 
     print(f"Job {job_id} created")
-    print("Open the job in Studio at", response.data.get("job", {}).get("url"))
+    print("Open the job in Studio at", response.data.get("url"))
     print("=" * 40)
 
     return 0 if no_wait else show_logs_from_client(client, job_id)
@@ -421,16 +421,14 @@ def upload_files(client: StudioClient, files: list[str]) -> list[str]:
     for file in files:
         file_name = os.path.basename(file)
         with open(file, "rb") as f:
-            file_content = f.read()
-        response = client.upload_file(file_content, file_name)
+            response = client.upload_file(f, file_name)
         if not response.ok:
             raise DataChainError(response.message)
 
         if not response.data:
             raise DataChainError(f"Failed to upload file {file_name}")
 
-        file_id = response.data.get("blob", {}).get("id")
-        if file_id:
+        if file_id := response.data.get("id"):
             file_ids.append(str(file_id))
     return file_ids
 
@@ -456,7 +454,7 @@ def list_jobs(status: Optional[str], team_name: Optional[str], limit: int):
     if not response.ok:
         raise DataChainError(response.message)
 
-    jobs = response.data.get("jobs", [])
+    jobs = response.data or []
     if not jobs:
         print("No jobs found")
         return
@@ -492,7 +490,7 @@ def list_clusters(team_name: Optional[str]):
     if not response.ok:
         raise DataChainError(response.message)
 
-    clusters = response.data.get("clusters", [])
+    clusters = response.data or []
     if not clusters:
         print("No clusters found")
         return
@@ -505,6 +503,7 @@ def list_clusters(team_name: Optional[str]):
             "Cloud Provider": cluster.get("cloud_provider"),
             "Cloud Credentials": cluster.get("cloud_credentials"),
             "Is Active": cluster.get("is_active"),
+            "Is Default": cluster.get("default"),
             "Max Workers": cluster.get("max_workers"),
         }
         for cluster in clusters

{datachain-0.30.7.dist-info → datachain-0.31.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.30.7
+Version: 0.31.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.30.7.dist-info → datachain-0.31.0.dist-info}/RECORD

@@ -17,7 +17,7 @@ datachain/project.py,sha256=90D4GpJSA3t0fayYZbzrL3sk4U7EJhQo8psnWvdI7_o,2280
 datachain/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/script_meta.py,sha256=V-LaFOZG84pD0Zc0NvejYdzwDgzITv6yHvAHggDCnuY,4978
 datachain/semver.py,sha256=UB8GHPBtAP3UJGeiuJoInD7SK-DnB93_Xd1qy_CQ9cU,2074
-datachain/studio.py,sha256=27750qCSNxIChEzhV02damIFreLMfr7UdiWqMFyk8AA,15361
+datachain/studio.py,sha256=IS8o4BZnhUo73Bd8m4CJxFc5utdmh2miIs25WswkFBA,15283
 datachain/telemetry.py,sha256=0A4IOPPp9VlP5pyW9eBfaTK3YhHGzHl7dQudQjUAx9A,994
 datachain/utils.py,sha256=5ehFeqXau7MFmGUQRsjRyPfDMPoOF1ojpfVciYUo5fE,15659
 datachain/catalog/__init__.py,sha256=9NBaywvAOaXdkyqiHjbBEiXs7JImR1OJsY9r8D5Q16g,403
@@ -136,7 +136,7 @@ datachain/query/session.py,sha256=gKblltJAVQAVSTswAgWGDgGbpmFlFzFVkIQojDCjgXM,68
 datachain/query/udf.py,sha256=jqutTpvkT6eHl96ZEgYiiTMAhI7vmTQA6JH9y4WCibI,1405
 datachain/query/utils.py,sha256=a2PTBZ3qsG6XlUcp9XsoGiQfKkca4Q3m-VzFgiGQPAc,1230
 datachain/remote/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datachain/remote/studio.py,sha256=pDThxvEEpIKVGfa9rmtz_zeqHwrgzh0Lv-Pd4wzDx5k,15448
+datachain/remote/studio.py,sha256=amjcV0B8qumsVBnxPQnt8oSrnfMK2vAdOurVMA9L_zA,16868
 datachain/sql/__init__.py,sha256=8D2omsBiATt8bjLjGo6jBEtaKEkOlnlNFWhVryHMDv0,388
 datachain/sql/postgresql_dialect.py,sha256=pDTfH8xaXz5xZsq8O1aQUvWLRIv_ogYeAqtmKlPp3Rw,280
 datachain/sql/postgresql_types.py,sha256=ryb_0lzuA9UOJ_B6nW9Yb8nJjzeSmEItAL_Ceue65lc,627
@@ -160,9 +160,9 @@ datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR
 datachain/toolkit/__init__.py,sha256=eQ58Q5Yf_Fgv1ZG0IO5dpB4jmP90rk8YxUWmPc1M2Bo,68
 datachain/toolkit/split.py,sha256=ktGWzY4kyzjWyR86dhvzw-Zhl0lVk_LOX3NciTac6qo,2914
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.30.7.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.30.7.dist-info/METADATA,sha256=d6ClkSVhY7AFkjh7jgUFEwHpTa7LhpJU75_M8ufegcI,13898
-datachain-0.30.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-datachain-0.30.7.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.30.7.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.30.7.dist-info/RECORD,,
+datachain-0.31.0.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.31.0.dist-info/METADATA,sha256=hY_KVFdUHZmZcxRiy5e-GY6CXI-sY0oKAtrvNakApdY,13898
+datachain-0.31.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+datachain-0.31.0.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.31.0.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.31.0.dist-info/RECORD,,