datachain 0.27.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

datachain/cli/parser/job.py

@@ -109,6 +109,11 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
     studio_run_parser.add_argument(
         "--cron", action="store", help="Cron expression for the cron task."
     )
+    studio_run_parser.add_argument(
+        "--no-wait",
+        action="store_true",
+        help="Do not wait for the job to finish",
+    )
 
     studio_ls_help = "List jobs in Studio"
     studio_ls_description = "List jobs in Studio."

datachain/data_storage/job.py

@@ -12,10 +12,11 @@ class JobStatus(int, Enum):
     CANCELING = 7
     CANCELED = 8
     CANCELING_SCHEDULED = 9
+    TASK = 11
 
     @classmethod
     def finished(cls) -> tuple[int, ...]:
-        return cls.COMPLETE, cls.FAILED, cls.CANCELED
+        return cls.COMPLETE, cls.FAILED, cls.CANCELED, cls.TASK
 
 
 class JobQueryType(int, Enum):

datachain/studio.py

@@ -8,6 +8,7 @@ import dateparser
 import tabulate
 
 from datachain.config import Config, ConfigLevel
+from datachain.data_storage.job import JobStatus
 from datachain.dataset import QUERY_DATASET_PREFIX, parse_dataset_name
 from datachain.error import DataChainError
 from datachain.remote.studio import StudioClient
@@ -20,6 +21,8 @@ POST_LOGIN_MESSAGE = (
     "Once you've logged in, return here "
     "and you'll be ready to start using DataChain with Studio."
 )
+RETRY_MAX_TIMES = 10
+RETRY_SLEEP_SEC = 1
 
 
 def process_jobs_args(args: "Namespace"):
@@ -46,6 +49,7 @@ def process_jobs_args(args: "Namespace"):
             args.cluster,
             args.start_time,
             args.cron,
+            args.no_wait,
         )
 
     if args.cmd == "cancel":
@@ -287,17 +291,34 @@ def parse_start_time(start_time_str: Optional[str]) -> Optional[str]:
 def show_logs_from_client(client, job_id):
     # Sync usage
     async def _run():
+        retry_count = 0
         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:
-                latest_status = message["job"]["status"]
-                print(f"\n>>>> Job is now in {latest_status} status.")
+        processed_statuses = set()
+        while True:
+            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:
+                    latest_status = message["job"]["status"]
+                    if latest_status in processed_statuses:
+                        continue
+                    processed_statuses.add(latest_status)
+                    print(f"\n>>>> Job is now in {latest_status} status.")
+
+            try:
+                if retry_count > RETRY_MAX_TIMES or (
+                    latest_status and JobStatus[latest_status].finished()
+                ):
+                    break
+                await asyncio.sleep(RETRY_SLEEP_SEC)
+                retry_count += 1
+            except KeyError:
+                pass
+
         return latest_status
 
-    latest_status = asyncio.run(_run())
+    final_status = asyncio.run(_run())
 
     response = client.dataset_job_versions(job_id)
     if not response.ok:
@@ -314,9 +335,9 @@ def show_logs_from_client(client, job_id):
 
     exit_code_by_status = {
         "FAILED": 1,
-        "CANCELLED": 2,
+        "CANCELED": 2,
     }
-    return exit_code_by_status.get(latest_status.upper(), 0) if latest_status else 0
+    return exit_code_by_status.get(final_status.upper(), 0) if final_status else 0
 
 
 def create_job(
@@ -334,6 +355,7 @@ def create_job(
     cluster: Optional[str] = None,
     start_time: Optional[str] = None,
     cron: Optional[str] = None,
+    no_wait: Optional[bool] = False,
 ):
     query_type = "PYTHON" if query_file.endswith(".py") else "SHELL"
     with open(query_file) as f:
@@ -388,7 +410,7 @@ def create_job(
     print("Open the job in Studio at", response.data.get("job", {}).get("url"))
     print("=" * 40)
 
-    return show_logs_from_client(client, job_id)
+    return 0 if no_wait else show_logs_from_client(client, job_id)
 
 
 def upload_files(client: StudioClient, files: list[str]) -> list[str]:

{datachain-0.27.0.dist-info → datachain-0.28.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.27.0
+Version: 0.28.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -45,7 +45,7 @@ Requires-Dist: datamodel-code-generator>=0.25
 Requires-Dist: Pillow<12,>=10.0.0
 Requires-Dist: msgpack<2,>=1.0.4
 Requires-Dist: psutil
-Requires-Dist: huggingface_hub
+Requires-Dist: huggingface_hub<0.34.0
 Requires-Dist: iterative-telemetry>=0.0.10
 Requires-Dist: platformdirs
 Requires-Dist: dvc-studio-client<1,>=0.21

{datachain-0.27.0.dist-info → datachain-0.28.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=RCpVZdHRX-ClEddXaAsZDGFy5o-SOqVCa5NhLj8337s,14486
+datachain/studio.py,sha256=-BmKLVNBLPFveUgVVE2So3aaiGndO2jK2qbHZ0zBDd8,15239
 datachain/telemetry.py,sha256=0A4IOPPp9VlP5pyW9eBfaTK3YhHGzHl7dQudQjUAx9A,994
 datachain/utils.py,sha256=DNqOi-Ydb7InyWvD9m7_yailxz6-YGpZzh00biQaHNo,15305
 datachain/catalog/__init__.py,sha256=cMZzSz3VoUi-6qXSVaHYN-agxQuAcz2XSqnEPZ55crE,353
@@ -35,7 +35,7 @@ datachain/cli/commands/misc.py,sha256=c0DmkOLwcDI2YhA8ArOuLJk6aGzSMZCiKL_E2JGibV
 datachain/cli/commands/query.py,sha256=Xzfgh14nPVH-sclqX1tpZqgfdTugw5s_44v0D33z6FA,1505
 datachain/cli/commands/show.py,sha256=Cf8wBs12h-xtdOzjU5GTDy2C8rF5HJSF0hDJYER1zH8,1606
 datachain/cli/parser/__init__.py,sha256=NPB6ssP4CCt7G1SWZ_8oNQEH2C1lktWgkyHYXDQJZNc,15073
-datachain/cli/parser/job.py,sha256=iytBZaCcQUhaOcRlYZFeAJsscN2T2XcEY7MibTeuZhg,5786
+datachain/cli/parser/job.py,sha256=g6ozI3pnV0ly79L7M9mikCeYTPgKlG5gR0D144R82tk,5928
 datachain/cli/parser/studio.py,sha256=Bo__LKM7qhJGgkyX8M_bCvgZ2Gvqq6r_X4t1NdtaBIY,3881
 datachain/cli/parser/utils.py,sha256=rETdD-9Hq9A4OolgfT7jQw4aoawtbfmkdtH6E7nkhpI,2888
 datachain/client/__init__.py,sha256=1kDpCPoibMXi1gExR4lTLc5pi-k6M5TANiwtXkPoLhU,49
@@ -48,7 +48,7 @@ datachain/client/local.py,sha256=0J52Wzvw25hSucVlzBvLuMRAZwrAHZAYDvD1mNBqf4c,460
 datachain/client/s3.py,sha256=6DNVGLg-woPS1DVlYVX2rIlunNblsuxyOnI1rSzhW3k,7515
 datachain/data_storage/__init__.py,sha256=9Wit-oe5P46V7CJQTD0BJ5MhOa2Y9h3ddJ4VWTe-Lec,273
 datachain/data_storage/db_engine.py,sha256=n8ojCbvVMPY2e3SG8fUaaD0b9GkVfpl_Naa_6EiHfWg,3788
-datachain/data_storage/job.py,sha256=9r0OGwh22bHNIvLHqg8_-eJSP1YYB-BN5HOla5TdCxw,402
+datachain/data_storage/job.py,sha256=ZkeXCNUj_VCkoKYx29hqB4AcfVUielnRjY-GYUcUxt4,426
 datachain/data_storage/metastore.py,sha256=Qw332arvhgXB4UY0yX-Hu8Vgl3smU12l6bvxrL9Q-vo,53810
 datachain/data_storage/schema.py,sha256=o3JbURKXRg3IJyIVA4QjHHkn6byRuz7avbydU2FlvNY,9897
 datachain/data_storage/serializer.py,sha256=6G2YtOFqqDzJf1KbvZraKGXl2XHZyVml2krunWUum5o,927
@@ -158,9 +158,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.27.0.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.27.0.dist-info/METADATA,sha256=PWZ_EWTpk1OvWlQZe__5SCjFem6BD1AtYmTxJ5wV3iY,13759
-datachain-0.27.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-datachain-0.27.0.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.27.0.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.27.0.dist-info/RECORD,,
+datachain-0.28.0.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.28.0.dist-info/METADATA,sha256=lA3lv9RX2NeQPobrEjoEbAwg5K3zmnAnbDJ_hjR8KLw,13766
+datachain-0.28.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+datachain-0.28.0.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.28.0.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.28.0.dist-info/RECORD,,