datachain 0.14.3__py3-none-any.whl → 0.14.4__py3-none-any.whl

datachain/lib/dataset_info.py

@@ -12,6 +12,7 @@ from datachain.dataset import (
 )
 from datachain.job import Job
 from datachain.lib.data_model import DataModel
+from datachain.query.session import Session
 from datachain.utils import TIME_ZERO
 
 if TYPE_CHECKING:
@@ -32,6 +33,10 @@ class DatasetInfo(DataModel):
     error_message: str = Field(default="")
     error_stack: str = Field(default="")
 
+    @property
+    def is_temp(self) -> bool:
+        return Session.is_temp_dataset(self.name)
+
     @staticmethod
     def _validate_dict(
         v: Optional[Union[str, dict]],

datachain/lib/dc/datasets.py

@@ -140,6 +140,8 @@ def datasets(
         )
     ]
 
+    datasets_values = [d for d in datasets_values if not d.is_temp]
+
     return read_values(
         session=session,
         settings=settings,

datachain/lib/udf.py

@@ -16,7 +16,6 @@ from datachain.lib.convert.flatten import flatten
 from datachain.lib.data_model import DataValue
 from datachain.lib.file import File
 from datachain.lib.utils import AbstractUDF, DataChainError, DataChainParamsError
-from datachain.progress import CombinedDownloadCallback
 from datachain.query.batch import (
     Batch,
     BatchingStrategy,
@@ -327,8 +326,9 @@ def _prefetch_inputs(
 
     if after_prefetch is None:
         after_prefetch = noop
-        if isinstance(download_cb, CombinedDownloadCallback):
-            after_prefetch = download_cb.increment_file_count
+        if download_cb and hasattr(download_cb, "increment_file_count"):
+            increment_file_count: Callable[[], None] = download_cb.increment_file_count
+            after_prefetch = increment_file_count
 
     f = partial(_prefetch_input, download_cb=download_cb, after_prefetch=after_prefetch)
     mapper = AsyncMapper(f, prepared_inputs, workers=prefetch)

datachain/query/dispatch.py

@@ -4,9 +4,8 @@ from itertools import chain
 from multiprocessing import cpu_count
 from sys import stdin
 from threading import Timer
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Literal, Optional
 
-import attrs
 import multiprocess
 from cloudpickle import load, loads
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
@@ -60,6 +59,7 @@ def udf_entrypoint() -> int:
     query = udf_info["query"]
     rows_total = udf_info["rows_total"]
     batching = udf_info["batching"]
+    is_generator = udf_info["is_generator"]
     n_workers = udf_info["processes"]
     if n_workers is True:
         n_workers = None  # Use default number of CPUs (cores)
@@ -72,17 +72,20 @@ def udf_entrypoint() -> int:
     ) as udf_inputs:
         download_cb = get_download_callback()
         processed_cb = get_processed_callback()
+        generated_cb = get_generated_callback(is_generator)
         try:
             dispatch.run_udf_parallel(
                 udf_inputs,
                 rows_total=rows_total,
                 n_workers=n_workers,
-                processed_cb=processed_cb,
                 download_cb=download_cb,
+                processed_cb=processed_cb,
+                generated_cb=generated_cb,
             )
         finally:
             download_cb.close()
             processed_cb.close()
+            generated_cb.close()
 
     return 0
 
@@ -128,7 +131,6 @@ class UDFDispatcher:
             self.done_queue,
             self.query,
             self.table,
-            self.is_generator,
             self.is_batching,
             self.cache,
             self.udf_fields,
@@ -152,16 +154,14 @@ class UDFDispatcher:
         for _ in range(n_workers):
             put_into_queue(task_queue, STOP_SIGNAL)
 
-    def create_input_queue(self):
-        return self.ctx.Queue()
-
     def run_udf_parallel(  # noqa: C901, PLR0912
         self,
         input_rows: Iterable[RowsOutput],
         rows_total: int,
         n_workers: Optional[int] = None,
-        processed_cb: Callback = DEFAULT_CALLBACK,
         download_cb: Callback = DEFAULT_CALLBACK,
+        processed_cb: Callback = DEFAULT_CALLBACK,
+        generated_cb: Callback = DEFAULT_CALLBACK,
     ) -> None:
         n_workers = get_n_workers_from_arg(n_workers)
 
@@ -214,6 +214,8 @@ class UDFDispatcher:
                     download_cb.relative_update(downloaded)
                 if processed := result.get("processed"):
                     processed_cb.relative_update(processed)
+                if generated := result.get("generated"):
+                    generated_cb.relative_update(generated)
 
                 status = result["status"]
                 if status in (OK_STATUS, NOTIFY_STATUS):
@@ -260,46 +262,61 @@ class UDFDispatcher:
                 p.join()
 
 
-class WorkerCallback(Callback):
-    def __init__(self, queue: "multiprocess.Queue"):
+class DownloadCallback(Callback):
+    def __init__(self, queue: "multiprocess.Queue") -> None:
         self.queue = queue
         super().__init__()
 
     def relative_update(self, inc: int = 1) -> None:
+        # This callback is used to notify the size of the downloaded files
+        pass
+
+    def increment_file_count(self, inc: int = 1) -> None:
         put_into_queue(self.queue, {"status": NOTIFY_STATUS, "downloaded": inc})
 
 
 class ProcessedCallback(Callback):
-    def __init__(self):
-        self.processed_rows: Optional[int] = None
+    def __init__(
+        self,
+        name: Literal["processed", "generated"],
+        queue: "multiprocess.Queue",
+    ) -> None:
+        self.name = name
+        self.queue = queue
         super().__init__()
 
     def relative_update(self, inc: int = 1) -> None:
-        self.processed_rows = inc
+        put_into_queue(self.queue, {"status": NOTIFY_STATUS, self.name: inc})
 
 
-@attrs.define
 class UDFWorker:
-    catalog: "Catalog"
-    udf: "UDFAdapter"
-    task_queue: "multiprocess.Queue"
-    done_queue: "multiprocess.Queue"
-    query: "Select"
-    table: "Table"
-    is_generator: bool
-    is_batching: bool
-    cache: bool
-    udf_fields: Sequence[str]
-    cb: Callback = attrs.field()
-
-    @cb.default
-    def _default_callback(self) -> WorkerCallback:
-        return WorkerCallback(self.done_queue)
+    def __init__(
+        self,
+        catalog: "Catalog",
+        udf: "UDFAdapter",
+        task_queue: "multiprocess.Queue",
+        done_queue: "multiprocess.Queue",
+        query: "Select",
+        table: "Table",
+        is_batching: bool,
+        cache: bool,
+        udf_fields: Sequence[str],
+    ) -> None:
+        self.catalog = catalog
+        self.udf = udf
+        self.task_queue = task_queue
+        self.done_queue = done_queue
+        self.query = query
+        self.table = table
+        self.is_batching = is_batching
+        self.cache = cache
+        self.udf_fields = udf_fields
+
+        self.download_cb = DownloadCallback(self.done_queue)
+        self.processed_cb = ProcessedCallback("processed", self.done_queue)
+        self.generated_cb = ProcessedCallback("generated", self.done_queue)
 
     def run(self) -> None:
-        processed_cb = ProcessedCallback()
-        generated_cb = get_generated_callback(self.is_generator)
-
         prefetch = self.udf.prefetch
         with _get_cache(self.catalog.cache, prefetch, use_cache=self.cache) as _cache:
             catalog = clone_catalog_with_cache(self.catalog, _cache)
@@ -308,29 +325,22 @@ class UDFWorker:
                 self.get_inputs(),
                 catalog,
                 self.cache,
-                download_cb=self.cb,
-                processed_cb=processed_cb,
+                download_cb=self.download_cb,
+                processed_cb=self.processed_cb,
             )
             with safe_closing(udf_results):
                 process_udf_outputs(
                     catalog.warehouse,
                     self.table,
-                    self.notify_and_process(udf_results, processed_cb),
+                    self.notify_and_process(udf_results),
                     self.udf,
-                    cb=generated_cb,
+                    cb=self.generated_cb,
                 )
+        put_into_queue(self.done_queue, {"status": FINISHED_STATUS})
 
-        put_into_queue(
-            self.done_queue,
-            {"status": FINISHED_STATUS, "processed": processed_cb.processed_rows},
-        )
-
-    def notify_and_process(self, udf_results, processed_cb):
+    def notify_and_process(self, udf_results):
         for row in udf_results:
-            put_into_queue(
-                self.done_queue,
-                {"status": OK_STATUS, "processed": processed_cb.processed_rows},
-            )
+            put_into_queue(self.done_queue, {"status": OK_STATUS})
             yield row
 
     def get_inputs(self):

datachain/query/session.py

@@ -100,6 +100,10 @@ class Session:
     def get_temp_prefix(self) -> str:
         return f"{self.DATASET_PREFIX}{self.name}_"
 
+    @classmethod
+    def is_temp_dataset(cls, name) -> bool:
+        return name.startswith(cls.DATASET_PREFIX)
+
     def _cleanup_temp_datasets(self) -> None:
         prefix = self.get_temp_prefix()
         try:

{datachain-0.14.3.dist-info → datachain-0.14.4.dist-info}/METADATA

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

{datachain-0.14.3.dist-info → datachain-0.14.4.dist-info}/RECORD

@@ -69,7 +69,7 @@ datachain/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/lib/arrow.py,sha256=9UBCF-lftQaz0yxdsjbLKbyzVSmrF_QSWdhp2oBDPqs,9486
 datachain/lib/clip.py,sha256=lm5CzVi4Cj1jVLEKvERKArb-egb9j1Ls-fwTItT6vlI,6150
 datachain/lib/data_model.py,sha256=ZwBXELtqROEdLL4DmxTipnwUZmhQvMz_UVDzyf7nQ9Y,2899
-datachain/lib/dataset_info.py,sha256=IjdF1E0TQNOq9YyynfWiCFTeZpbyGfyJvxgJY4YN810,2493
+datachain/lib/dataset_info.py,sha256=Jnjy7vq4iNVkq1e-SYjqxdojlxIDXvZ352NCLLZg59k,2633
 datachain/lib/file.py,sha256=HLQXS_WULm7Y-fkHMy0WpibVAcrkLPRS6CrZy6rwFe0,30450
 datachain/lib/hf.py,sha256=gjxuStZBlKtNk3-4yYSlWZDv9zBGblOdvEy_Lwap5hA,5882
 datachain/lib/image.py,sha256=butvUY_33PVEYPKX2nVCPeJjJVcBaptZwsE9REQsTS8,3247
@@ -82,7 +82,7 @@ datachain/lib/settings.py,sha256=ZELRCTLbi5vzRPiDX6cQ9LLg9TefJ_A05gIGni0lll8,253
 datachain/lib/signal_schema.py,sha256=DRatqSG7OVtCUCWyZvMXe4m7r7XFO6NCfzsJRDErMtg,35185
 datachain/lib/tar.py,sha256=3WIzao6yD5fbLqXLTt9GhPGNonbFIs_fDRu-9vgLgsA,1038
 datachain/lib/text.py,sha256=UNHm8fhidk7wdrWqacEWaA6I9ykfYqarQ2URby7jc7M,1261
-datachain/lib/udf.py,sha256=h38a457xg-4wO2XcxPs4pzDq8JxTmYm4N84iAf0HRzY,16168
+datachain/lib/udf.py,sha256=JJwjvy41N65PtWGUAq7TYnhdOOR6RiMDUJEKl5xtwLc,16199
 datachain/lib/udf_signature.py,sha256=2EtsOPDNSPqcOlYwqbCdy6RF5MldI-7smii8aLy8p7Y,7543
 datachain/lib/utils.py,sha256=QrjVs_oLRXEotOPUYurBJypBFi_ReTJmxcnJeH4j2Uk,1596
 datachain/lib/video.py,sha256=suH_8Mi8VYk4-IVb1vjSduF_njs64ji1WGKHxDLnGYw,6629
@@ -97,7 +97,7 @@ datachain/lib/convert/values_to_tuples.py,sha256=EFfIGBiVVltJQG8blzsQ1dGXneh4D3w
 datachain/lib/dc/__init__.py,sha256=6rKKHS6MA3mS6UJXiysrv4TURs4R_UWAQK2tJ2t1QMs,743
 datachain/lib/dc/csv.py,sha256=d0ULzpsTTeqp_eM-2jVHb1kYHQN2lJFf4O6LWd5tOJw,4401
 datachain/lib/dc/datachain.py,sha256=hwuAElfEhRLyh-Uvuc7YIpFx6nsI_B90xwnMqgkkgrI,76390
-datachain/lib/dc/datasets.py,sha256=hTzq18Ij9kpOAJOU-VN4-VyThTTxLSWLfVIk3bgzAPs,4329
+datachain/lib/dc/datasets.py,sha256=IYa8fixaqYasl0iE_cEaMDYfVCM4Dn4eOimz73YJTBY,4398
 datachain/lib/dc/hf.py,sha256=I1vFNOa1C87lBuBj5FHENLY2jTaQ8erngiX0cyBmOp4,2170
 datachain/lib/dc/json.py,sha256=9ei9ZNzWVXZWD4HNGTfBhcoLPnXBBDywKV-3Wi1mT28,2725
 datachain/lib/dc/listings.py,sha256=qPy1DTvYkbNICT1ujo8LwezzMEW8E3dln1knw7Jwl0I,1044
@@ -119,12 +119,12 @@ datachain/model/ultralytics/segment.py,sha256=koq1HASo29isf0in6oSlzmU4IzsmOXe87F
 datachain/query/__init__.py,sha256=7DhEIjAA8uZJfejruAVMZVcGFmvUpffuZJwgRqNwe-c,263
 datachain/query/batch.py,sha256=6w8gzLTmLeylststu-gT5jIqEfi4-djS7_yTYyeo-fw,4190
 datachain/query/dataset.py,sha256=caUsFzaVZXOz8NmeTMeOdyRQLQP8KCnxYMxF-pG4yFQ,58712
-datachain/query/dispatch.py,sha256=T4vdJE0k3Ff1osaQzYTC_2gOOkT0mXwKzNy-5aZcrTE,12300
+datachain/query/dispatch.py,sha256=ErdK-biHYhRLDsm7on6vAHSjX-hAHgEHsBRHmuMS_4E,12979
 datachain/query/metrics.py,sha256=r5b0ygYhokbXp8Mg3kCH8iFSRw0jxzyeBe-C-J_bKFc,938
 datachain/query/params.py,sha256=O_j89mjYRLOwWNhYZl-z7mi-rkdP7WyFmaDufsdTryE,863
 datachain/query/queue.py,sha256=waqM_KzavU8C-G95-4211Nd4GXna_u2747Chgwtgz2w,3839
 datachain/query/schema.py,sha256=b_KnVy6B26Ol4nYG0LqNNpeQ1QYPk95YRGUjXfdaQWs,6606
-datachain/query/session.py,sha256=I1KG8jDIaxGAfRfDRucMx8DqsANf_VYWtwtXjeD19lI,6399
+datachain/query/session.py,sha256=wNdOHAi4HrsEihfzdcTlfB5i1xyj0dw6rlUz84StOoU,6512
 datachain/query/udf.py,sha256=ljAYaF-J77t7iS4zc1-g1ssYd4c6Q-ccKGEc3VQQmeM,1322
 datachain/query/utils.py,sha256=u0A_BwG9PNs0DxoDcvSWgWLpj3ByTUv8CqH13CIuGag,1293
 datachain/remote/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -150,9 +150,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.14.3.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.14.3.dist-info/METADATA,sha256=hTVICGrF_sALHSr7uuAipFr5HVrgydiP2JcDlZ-Q_w0,11338
-datachain-0.14.3.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-datachain-0.14.3.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.14.3.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.14.3.dist-info/RECORD,,
+datachain-0.14.4.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.14.4.dist-info/METADATA,sha256=CmIX9sam07khHa1Wrw5eMcSUo06rQs10kF0zF9flz28,11338
+datachain-0.14.4.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+datachain-0.14.4.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.14.4.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.14.4.dist-info/RECORD,,