datachain 0.21.0__tar.gz → 0.21.1__tar.gz

{datachain-0.21.0 → datachain-0.21.1}/PKG-INFO

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

{datachain-0.21.0 → datachain-0.21.1}/src/datachain/cli/__init__.py

@@ -100,8 +100,10 @@ def handle_command(args, catalog, client_config) -> int:
 
     handler = command_handlers.get(args.command)
     if handler:
-        handler()
-        return 0
+        return_code = handler()
+        if return_code is None:
+            return 0
+        return return_code
     print(f"invalid command: {args.command}", file=sys.stderr)
     return 1
 

{datachain-0.21.0 → datachain-0.21.1}/src/datachain/studio.py

@@ -248,14 +248,17 @@ def save_config(hostname, token, level=ConfigLevel.GLOBAL):
 def show_logs_from_client(client, job_id):
     # Sync usage
     async def _run():
+        latest_status = None
         async for message in client.tail_job_logs(job_id):
             if "logs" in message:
                 for log in message["logs"]:
                     print(log["message"], end="")
             elif "job" in message:
-                print(f"\n>>>> Job is now in {message['job']['status']} status.")
+                latest_status = message["job"]["status"]
+                print(f"\n>>>> Job is now in {latest_status} status.")
+        return latest_status
 
-    asyncio.run(_run())
+    latest_status = asyncio.run(_run())
 
     response = client.dataset_job_versions(job_id)
     if not response.ok:
@@ -270,6 +273,12 @@ def show_logs_from_client(client, job_id):
     else:
         print("\n\nNo dataset versions created during the job.")
 
+    exit_code_by_status = {
+        "FAILED": 1,
+        "CANCELLED": 2,
+    }
+    return exit_code_by_status.get(latest_status.upper(), 0) if latest_status else 0
+
 
 def create_job(
     query_file: str,
@@ -326,7 +335,7 @@ def create_job(
     print("Open the job in Studio at", response.data.get("job", {}).get("url"))
     print("=" * 40)
 
-    show_logs_from_client(client, job_id)
+    return show_logs_from_client(client, job_id)
 
 
 def upload_files(client: StudioClient, files: list[str]) -> list[str]:
@@ -396,7 +405,7 @@ def show_job_logs(job_id: str, team_name: Optional[str]):
         )
 
     client = StudioClient(team=team_name)
-    show_logs_from_client(client, job_id)
+    return show_logs_from_client(client, job_id)
 
 
 def list_clusters(team_name: Optional[str]):

{datachain-0.21.0 → datachain-0.21.1}/src/datachain.egg-info/PKG-INFO

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

{datachain-0.21.0 → datachain-0.21.1}/tests/test_cli_studio.py

@@ -1,5 +1,6 @@
 from unittest.mock import MagicMock
 
+import pytest
 import requests_mock
 import websockets
 from dvc_studio_client.auth import AuthorizationExpiredError
@@ -429,3 +430,65 @@ def test_studio_run(capsys, mocker, tmp_dir):
         "priority": 5,
         "compute_cluster_name": "default",
     }
+
+
+@pytest.mark.parametrize(
+    "status,expected_exit_code", [("FAILED", 1), ("CANCELLED", 2), ("COMPLETED", 0)]
+)
+def test_studio_run_non_zero_exit_code(
+    capsys, mocker, tmp_dir, status, expected_exit_code
+):
+    # Mock tail_job_logs to return a status
+    async def mock_tail_job_logs(job_id):
+        # Simulate some log messages
+        yield {"logs": [{"message": "Starting job...\n"}]}
+        yield {"logs": [{"message": "Processing data...\n"}]}
+        # Return failed status
+        yield {"job": {"status": status}}
+
+    mocker.patch(
+        "datachain.remote.studio.StudioClient.tail_job_logs",
+        side_effect=mock_tail_job_logs,
+    )
+
+    with Config(ConfigLevel.GLOBAL).edit() as conf:
+        conf["studio"] = {"token": "isat_access_token", "team": "team_name"}
+
+    with requests_mock.mock() as m:
+        m.post(
+            f"{STUDIO_URL}/api/datachain/job",
+            json={"job": {"id": 1, "url": "https://example.com"}},
+        )
+        m.get(
+            f"{STUDIO_URL}/api/datachain/datasets/dataset_job_versions?job_id=1&team_name=team_name",
+            json={
+                "dataset_versions": [
+                    {"dataset_name": "dataset_name", "version": "1.0.0"}
+                ]
+            },
+        )
+
+        (tmp_dir / "example_query.py").write_text("print(1)")
+
+        assert (
+            main(
+                [
+                    "job",
+                    "run",
+                    "example_query.py",
+                ]
+            )
+            == expected_exit_code
+        )
+
+    out = capsys.readouterr().out
+    assert (
+        out.strip() == "Job 1 created\nOpen the job in Studio at https://example.com\n"
+        "========================================\n"
+        "Starting job...\n"
+        "Processing data...\n"
+        "\n"
+        f">>>> Job is now in {status} status.\n\n\n"
+        ">>>> Dataset versions created during the job:\n"
+        "    - dataset_name@v1.0.0"
+    )