datachain 0.30.7__tar.gz → 0.31.0__tar.gz

{datachain-0.30.7 → datachain-0.31.0}/PKG-INFO

@@ -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 → datachain-0.31.0}/src/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-0.30.7 → datachain-0.31.0}/src/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 → datachain-0.31.0}/src/datachain.egg-info/PKG-INFO

@@ -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 → datachain-0.31.0}/tests/test_cli_studio.py

@@ -336,7 +336,7 @@ def test_studio_rm_dataset(capsys, mocker):
 def test_studio_cancel_job(capsys, mocker):
     job_id = "8bddde6c-c3ca-41b0-9d87-ee945bfdce70"
     with requests_mock.mock() as m:
-        m.post(f"{STUDIO_URL}/api/datachain/job/{job_id}/cancel", json={})
+        m.post(f"{STUDIO_URL}/api/datachain/jobs/{job_id}/cancel", json={})
 
         # Studio token is required
         assert main(["job", "cancel", job_id]) == 1
@@ -359,10 +359,12 @@ def test_studio_run(capsys, mocker, tmp_dir):
         conf["studio"] = {"token": "isat_access_token", "team": "team_name"}
 
     with requests_mock.mock() as m:
-        m.post(f"{STUDIO_URL}/api/datachain/upload-file", json={"blob": {"id": 1}})
         m.post(
-            f"{STUDIO_URL}/api/datachain/job",
-            json={"job": {"id": 1, "url": "https://example.com"}},
+            f"{STUDIO_URL}/api/datachain/jobs/files?team_name=team_name", json={"id": 1}
+        )
+        m.post(
+            f"{STUDIO_URL}/api/datachain/jobs/",
+            json={"id": 1, "url": "https://example.com"},
         )
         m.get(
             f"{STUDIO_URL}/api/datachain/datasets/dataset_job_versions?job_id=1&team_name=team_name",
@@ -421,15 +423,17 @@ def test_studio_run(capsys, mocker, tmp_dir):
     second_request = m.request_history[1]
 
     assert first_request.method == "POST"
-    assert first_request.url == f"{STUDIO_URL}/api/datachain/upload-file"
-    assert first_request.json() == {
-        "file_content": "ZmlsZSBjb250ZW50",
-        "file_name": "file.txt",
-        "team_name": "team_name",
-    }
+    assert (
+        first_request.url
+        == f"{STUDIO_URL}/api/datachain/jobs/files?team_name=team_name"
+    )
+    # Check that it's multipart/form-data request
+    assert "multipart/form-data" in first_request.headers.get("Content-Type", "")
+    # Check query parameters
+    assert first_request.qs["team_name"] == ["team_name"]
 
     assert second_request.method == "POST"
-    assert second_request.url == f"{STUDIO_URL}/api/datachain/job"
+    assert second_request.url == f"{STUDIO_URL}/api/datachain/jobs/"
     assert second_request.json() == {
         "query": "print(1)",
         "query_type": "PYTHON",
@@ -456,8 +460,8 @@ def test_studio_run_task(capsys, mocker, tmp_dir, studio_token):
 
     with requests_mock.mock() as m:
         m.post(
-            f"{STUDIO_URL}/api/datachain/job",
-            json={"job": {"id": 1, "url": "https://example.com"}},
+            f"{STUDIO_URL}/api/datachain/jobs/",
+            json={"id": 1, "url": "https://example.com"},
         )
         m.get(
             f"{STUDIO_URL}/api/datachain/datasets/dataset_job_versions?job_id=1&team_name=team_name",
@@ -485,7 +489,7 @@ def test_studio_run_task(capsys, mocker, tmp_dir, studio_token):
         )
     first_request = m.request_history[0]
     assert first_request.method == "POST"
-    assert first_request.url == f"{STUDIO_URL}/api/datachain/job"
+    assert first_request.url == f"{STUDIO_URL}/api/datachain/jobs/"
     request_json = first_request.json()
     assert request_json["start_after"] is not None
     assert request_json["cron_expression"] is not None
@@ -518,8 +522,8 @@ def test_studio_run_non_zero_exit_code(
 
     with requests_mock.mock() as m:
         m.post(
-            f"{STUDIO_URL}/api/datachain/job",
-            json={"job": {"id": 1, "url": "https://example.com"}},
+            f"{STUDIO_URL}/api/datachain/jobs/",
+            json={"id": 1, "url": "https://example.com"},
         )
         m.get(
             f"{STUDIO_URL}/api/datachain/datasets/dataset_job_versions?job_id=1&team_name=team_name",