datachain 0.8.3__py3-none-any.whl → 0.8.5__py3-none-any.whl

datachain/asyn.py

@@ -8,12 +8,14 @@ from collections.abc import (
     Iterable,
     Iterator,
 )
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, wait
 from heapq import heappop, heappush
 from typing import Any, Callable, Generic, Optional, TypeVar
 
 from fsspec.asyn import get_loop
 
+from datachain.utils import safe_closing
+
 ASYNC_WORKERS = 20
 
 InputT = TypeVar("InputT", contravariant=True)  # noqa: PLC0105
@@ -56,6 +58,7 @@ class AsyncMapper(Generic[InputT, ResultT]):
         self.pool = ThreadPoolExecutor(workers)
         self._tasks: set[asyncio.Task] = set()
         self._shutdown_producer = threading.Event()
+        self._producer_is_shutdown = threading.Event()
 
     def start_task(self, coro: Coroutine) -> asyncio.Task:
         task = self.loop.create_task(coro)
@@ -64,11 +67,16 @@ class AsyncMapper(Generic[InputT, ResultT]):
         return task
 
     def _produce(self) -> None:
-        for item in self.iterable:
-            if self._shutdown_producer.is_set():
-                return
-            fut = asyncio.run_coroutine_threadsafe(self.work_queue.put(item), self.loop)
-            fut.result()  # wait until the item is in the queue
+        try:
+            with safe_closing(self.iterable):
+                for item in self.iterable:
+                    if self._shutdown_producer.is_set():
+                        return
+                    coro = self.work_queue.put(item)
+                    fut = asyncio.run_coroutine_threadsafe(coro, self.loop)
+                    fut.result()  # wait until the item is in the queue
+        finally:
+            self._producer_is_shutdown.set()
 
     async def produce(self) -> None:
         await self.to_thread(self._produce)
@@ -179,6 +187,8 @@ class AsyncMapper(Generic[InputT, ResultT]):
             self.shutdown_producer()
             if not async_run.done():
                 async_run.cancel()
+                wait([async_run])
+            self._producer_is_shutdown.wait()
 
     def __iter__(self):
         return self.iterate()

datachain/cache.py

@@ -1,8 +1,12 @@
 import os
+from collections.abc import Iterator
+from contextlib import contextmanager
+from tempfile import mkdtemp
 from typing import TYPE_CHECKING, Optional
 
 from dvc_data.hashfile.db.local import LocalHashFileDB
 from dvc_objects.fs.local import LocalFileSystem
+from dvc_objects.fs.utils import remove
 from fsspec.callbacks import Callback, TqdmCallback
 
 from .progress import Tqdm
@@ -20,6 +24,23 @@ def try_scandir(path):
         pass
 
 
+def get_temp_cache(tmp_dir: str, prefix: Optional[str] = None) -> "DataChainCache":
+    cache_dir = mkdtemp(prefix=prefix, dir=tmp_dir)
+    return DataChainCache(cache_dir, tmp_dir=tmp_dir)
+
+
+@contextmanager
+def temporary_cache(
+    tmp_dir: str, prefix: Optional[str] = None, delete: bool = True
+) -> Iterator["DataChainCache"]:
+    cache = get_temp_cache(tmp_dir, prefix=prefix)
+    try:
+        yield cache
+    finally:
+        if delete:
+            cache.destroy()
+
+
 class DataChainCache:
     def __init__(self, cache_dir: str, tmp_dir: str):
         self.odb = LocalHashFileDB(
@@ -28,6 +49,9 @@ class DataChainCache:
             tmp_dir=tmp_dir,
         )
 
+    def __eq__(self, other) -> bool:
+        return self.odb == other.odb
+
     @property
     def cache_dir(self):
         return self.odb.path
@@ -63,7 +87,7 @@ class DataChainCache:
         if size < 0:
             size = await client.get_size(from_path, version_id=file.version)
         cb = callback or TqdmCallback(
-            tqdm_kwargs={"desc": odb_fs.name(from_path), "bytes": True},
+            tqdm_kwargs={"desc": odb_fs.name(from_path), "bytes": True, "leave": False},
             tqdm_cls=Tqdm,
             size=size,
         )
@@ -82,20 +106,18 @@ class DataChainCache:
             os.unlink(tmp_info)
 
     def store_data(self, file: "File", contents: bytes) -> None:
-        checksum = file.get_hash()
-        dst = self.path_from_checksum(checksum)
-        if not os.path.exists(dst):
-            # Create the file only if it's not already in cache
-            os.makedirs(os.path.dirname(dst), exist_ok=True)
-            with open(dst, mode="wb") as f:
-                f.write(contents)
-
-    def clear(self):
+        self.odb.add_bytes(file.get_hash(), contents)
+
+    def clear(self) -> None:
         """
         Completely clear the cache.
         """
         self.odb.clear()
 
+    def destroy(self) -> None:
+        # `clear` leaves the prefix directory structure intact.
+        remove(self.cache_dir)
+
     def get_total_size(self) -> int:
         total = 0
         for subdir in try_scandir(self.odb.path):

datachain/catalog/catalog.py

@@ -405,6 +405,7 @@ def get_download_bar(bar_format: str, total_size: int):
         unit_scale=True,
         unit_divisor=1000,
         total=total_size,
+        leave=False,
     )
 
 
@@ -429,6 +430,7 @@ def instantiate_node_groups(
             unit_scale=True,
             unit_divisor=1000,
             total=total_files,
+            leave=False,
         )
     )
 
@@ -534,6 +536,12 @@ def find_column_to_str(  # noqa: PLR0911
     return ""
 
 
+def clone_catalog_with_cache(catalog: "Catalog", cache: "DataChainCache") -> "Catalog":
+    clone = catalog.copy()
+    clone.cache = cache
+    return clone
+
+
 class Catalog:
     def __init__(
         self,
@@ -1242,10 +1250,17 @@ class Catalog:
         path: str,
         version_id: Optional[str] = None,
         client_config=None,
+        content_disposition: Optional[str] = None,
+        **kwargs,
     ) -> str:
         client_config = client_config or self.client_config
         client = Client.get_client(source, self.cache, **client_config)
-        return client.url(path, version_id=version_id)
+        return client.url(
+            path,
+            version_id=version_id,
+            content_disposition=content_disposition,
+            **kwargs,
+        )
 
     def export_dataset_table(
         self,
@@ -1437,6 +1452,7 @@ class Catalog:
             unit_scale=True,
             unit_divisor=1000,
             total=ds_stats.num_objects,  # type: ignore [union-attr]
+            leave=False,
         )
 
         schema = DatasetRecord.parse_schema(remote_ds_version.schema)

datachain/cli/__init__.py

@@ -0,0 +1,311 @@
+import logging
+import os
+import sys
+import traceback
+from multiprocessing import freeze_support
+from typing import Optional
+
+from datachain.cli.utils import get_logging_level
+from datachain.telemetry import telemetry
+
+from .commands import (
+    clear_cache,
+    completion,
+    dataset_stats,
+    du,
+    edit_dataset,
+    garbage_collect,
+    index,
+    list_datasets,
+    ls,
+    query,
+    rm_dataset,
+    show,
+)
+from .parser import get_parser
+
+logger = logging.getLogger("datachain")
+
+
+def main(argv: Optional[list[str]] = None) -> int:
+    from datachain.catalog import get_catalog
+
+    # Required for Windows multiprocessing support
+    freeze_support()
+
+    datachain_parser = get_parser()
+    args = datachain_parser.parse_args(argv)
+
+    if args.command in ("internal-run-udf", "internal-run-udf-worker"):
+        return handle_udf(args.command)
+
+    logger.addHandler(logging.StreamHandler())
+    logging_level = get_logging_level(args)
+    logger.setLevel(logging_level)
+
+    client_config = {
+        "aws_endpoint_url": args.aws_endpoint_url,
+        "anon": args.anon,
+    }
+
+    if args.debug_sql:
+        # This also sets this environment variable for any subprocesses
+        os.environ["DEBUG_SHOW_SQL_QUERIES"] = "True"
+
+    error = None
+
+    try:
+        catalog = get_catalog(client_config=client_config)
+        return handle_command(args, catalog, client_config)
+    except BrokenPipeError as exc:
+        error, return_code = handle_broken_pipe_error(exc)
+        return return_code
+    except (KeyboardInterrupt, Exception) as exc:
+        error, return_code = handle_general_exception(exc, args, logging_level)
+        return return_code
+    finally:
+        telemetry.send_cli_call(args.command, error=error)
+
+
+def handle_command(args, catalog, client_config) -> int:
+    """Handle the different CLI commands."""
+    from datachain.studio import process_jobs_args, process_studio_cli_args
+
+    command_handlers = {
+        "cp": lambda: handle_cp_command(args, catalog),
+        "clone": lambda: handle_clone_command(args, catalog),
+        "dataset": lambda: handle_dataset_command(args, catalog),
+        "ds": lambda: handle_dataset_command(args, catalog),
+        "ls": lambda: handle_ls_command(args, client_config),
+        "show": lambda: handle_show_command(args, catalog),
+        "du": lambda: handle_du_command(args, catalog, client_config),
+        "find": lambda: handle_find_command(args, catalog),
+        "index": lambda: handle_index_command(args, catalog),
+        "completion": lambda: handle_completion_command(args),
+        "query": lambda: handle_query_command(args, catalog),
+        "clear-cache": lambda: clear_cache(catalog),
+        "gc": lambda: garbage_collect(catalog),
+        "studio": lambda: process_studio_cli_args(args),
+        "job": lambda: process_jobs_args(args),
+    }
+
+    handler = command_handlers.get(args.command)
+    if handler:
+        handler()
+        return 0
+    print(f"invalid command: {args.command}", file=sys.stderr)
+    return 1
+
+
+def handle_cp_command(args, catalog):
+    catalog.cp(
+        args.sources,
+        args.output,
+        force=bool(args.force),
+        update=bool(args.update),
+        recursive=bool(args.recursive),
+        edatachain_file=None,
+        edatachain_only=False,
+        no_edatachain_file=True,
+        no_glob=args.no_glob,
+    )
+
+
+def handle_clone_command(args, catalog):
+    catalog.clone(
+        args.sources,
+        args.output,
+        force=bool(args.force),
+        update=bool(args.update),
+        recursive=bool(args.recursive),
+        no_glob=args.no_glob,
+        no_cp=args.no_cp,
+        edatachain=args.edatachain,
+        edatachain_file=args.edatachain_file,
+    )
+
+
+def handle_dataset_command(args, catalog):
+    dataset_commands = {
+        "pull": lambda: catalog.pull_dataset(
+            args.dataset,
+            args.output,
+            local_ds_name=args.local_name,
+            local_ds_version=args.local_version,
+            cp=args.cp,
+            force=bool(args.force),
+            edatachain=args.edatachain,
+            edatachain_file=args.edatachain_file,
+        ),
+        "edit": lambda: edit_dataset(
+            catalog,
+            args.name,
+            new_name=args.new_name,
+            description=args.description,
+            labels=args.labels,
+            studio=args.studio,
+            local=args.local,
+            all=args.all,
+            team=args.team,
+        ),
+        "ls": lambda: list_datasets(
+            catalog=catalog,
+            studio=args.studio,
+            local=args.local,
+            all=args.all,
+            team=args.team,
+        ),
+        "rm": lambda: rm_dataset(
+            catalog,
+            args.name,
+            version=args.version,
+            force=args.force,
+            studio=args.studio,
+            local=args.local,
+            all=args.all,
+            team=args.team,
+        ),
+        "remove": lambda: rm_dataset(
+            catalog,
+            args.name,
+            version=args.version,
+            force=args.force,
+            studio=args.studio,
+            local=args.local,
+            all=args.all,
+            team=args.team,
+        ),
+        "stats": lambda: dataset_stats(
+            catalog,
+            args.name,
+            args.version,
+            show_bytes=args.bytes,
+            si=args.si,
+        ),
+    }
+
+    handler = dataset_commands.get(args.datasets_cmd)
+    if handler:
+        return handler()
+    raise Exception(f"Unexpected command {args.datasets_cmd}")
+
+
+def handle_ls_command(args, client_config):
+    ls(
+        args.sources,
+        long=bool(args.long),
+        studio=args.studio,
+        local=args.local,
+        all=args.all,
+        team=args.team,
+        update=bool(args.update),
+        client_config=client_config,
+    )
+
+
+def handle_show_command(args, catalog):
+    show(
+        catalog,
+        args.name,
+        args.version,
+        limit=args.limit,
+        offset=args.offset,
+        columns=args.columns,
+        no_collapse=args.no_collapse,
+        schema=args.schema,
+    )
+
+
+def handle_du_command(args, catalog, client_config):
+    du(
+        catalog,
+        args.sources,
+        show_bytes=args.bytes,
+        depth=args.depth,
+        si=args.si,
+        update=bool(args.update),
+        client_config=client_config,
+    )
+
+
+def handle_find_command(args, catalog):
+    results_found = False
+    for result in catalog.find(
+        args.sources,
+        update=bool(args.update),
+        names=args.name,
+        inames=args.iname,
+        paths=args.path,
+        ipaths=args.ipath,
+        size=args.size,
+        typ=args.type,
+        columns=args.columns,
+    ):
+        print(result)
+        results_found = True
+    if not results_found:
+        print("No results")
+
+
+def handle_index_command(args, catalog):
+    index(
+        catalog,
+        args.sources,
+        update=bool(args.update),
+    )
+
+
+def handle_completion_command(args):
+    print(completion(args.shell))
+
+
+def handle_query_command(args, catalog):
+    query(
+        catalog,
+        args.script,
+        parallel=args.parallel,
+        params=args.param,
+    )
+
+
+def handle_broken_pipe_error(exc):
+    # Python flushes standard streams on exit; redirect remaining output
+    # to devnull to avoid another BrokenPipeError at shutdown
+    # See: https://docs.python.org/3/library/signal.html#note-on-sigpipe
+    error = str(exc)
+    devnull = os.open(os.devnull, os.O_WRONLY)
+    os.dup2(devnull, sys.stdout.fileno())
+    return error, 141  # 128 + 13 (SIGPIPE)
+
+
+def handle_general_exception(exc, args, logging_level):
+    error = str(exc)
+    if isinstance(exc, KeyboardInterrupt):
+        msg = "Operation cancelled by the user"
+    else:
+        msg = str(exc)
+    print("Error:", msg, file=sys.stderr)
+    if logging_level <= logging.DEBUG:
+        traceback.print_exception(
+            type(exc),
+            exc,
+            exc.__traceback__,
+            file=sys.stderr,
+        )
+    if args.pdb:
+        import pdb  # noqa: T100
+
+        pdb.post_mortem()
+    return error, 1
+
+
+def handle_udf(command):
+    if command == "internal-run-udf":
+        from datachain.query.dispatch import udf_entrypoint
+
+        return udf_entrypoint()
+
+    if command == "internal-run-udf-worker":
+        from datachain.query.dispatch import udf_worker_entrypoint
+
+        return udf_worker_entrypoint()

datachain/cli/commands/__init__.py

@@ -0,0 +1,29 @@
+from .datasets import (
+    dataset_stats,
+    edit_dataset,
+    list_datasets,
+    list_datasets_local,
+    rm_dataset,
+)
+from .du import du
+from .index import index
+from .ls import ls
+from .misc import clear_cache, completion, garbage_collect
+from .query import query
+from .show import show
+
+__all__ = [
+    "clear_cache",
+    "completion",
+    "dataset_stats",
+    "du",
+    "edit_dataset",
+    "garbage_collect",
+    "index",
+    "list_datasets",
+    "list_datasets_local",
+    "ls",
+    "query",
+    "rm_dataset",
+    "show",
+]

datachain/cli/commands/datasets.py

@@ -0,0 +1,129 @@
+import sys
+from typing import TYPE_CHECKING, Optional
+
+from tabulate import tabulate
+
+from datachain import utils
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+from datachain.cli.utils import determine_flavors
+from datachain.config import Config
+from datachain.error import DatasetNotFoundError
+
+
+def list_datasets(
+    catalog: "Catalog",
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
+):
+    from datachain.studio import list_datasets
+
+    token = Config().read().get("studio", {}).get("token")
+    all, local, studio = determine_flavors(studio, local, all, token)
+
+    local_datasets = set(list_datasets_local(catalog)) if all or local else set()
+    studio_datasets = (
+        set(list_datasets(team=team)) if (all or studio) and token else set()
+    )
+
+    rows = [
+        _datasets_tabulate_row(
+            name=name,
+            version=version,
+            both=(all or (local and studio)) and token,
+            local=(name, version) in local_datasets,
+            studio=(name, version) in studio_datasets,
+        )
+        for name, version in local_datasets.union(studio_datasets)
+    ]
+
+    print(tabulate(rows, headers="keys"))
+
+
+def list_datasets_local(catalog: "Catalog"):
+    for d in catalog.ls_datasets():
+        for v in d.versions:
+            yield (d.name, v.version)
+
+
+def _datasets_tabulate_row(name, version, both, local, studio):
+    row = {
+        "Name": name,
+        "Version": version,
+    }
+    if both:
+        row["Studio"] = "\u2714" if studio else "\u2716"
+        row["Local"] = "\u2714" if local else "\u2716"
+    return row
+
+
+def rm_dataset(
+    catalog: "Catalog",
+    name: str,
+    version: Optional[int] = None,
+    force: Optional[bool] = False,
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
+):
+    from datachain.studio import remove_studio_dataset
+
+    token = Config().read().get("studio", {}).get("token")
+    all, local, studio = determine_flavors(studio, local, all, token)
+
+    if all or local:
+        try:
+            catalog.remove_dataset(name, version=version, force=force)
+        except DatasetNotFoundError:
+            print("Dataset not found in local", file=sys.stderr)
+
+    if (all or studio) and token:
+        remove_studio_dataset(team, name, version, force)
+
+
+def edit_dataset(
+    catalog: "Catalog",
+    name: str,
+    new_name: Optional[str] = None,
+    description: Optional[str] = None,
+    labels: Optional[list[str]] = None,
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
+):
+    from datachain.studio import edit_studio_dataset
+
+    token = Config().read().get("studio", {}).get("token")
+    all, local, studio = determine_flavors(studio, local, all, token)
+
+    if all or local:
+        try:
+            catalog.edit_dataset(name, new_name, description, labels)
+        except DatasetNotFoundError:
+            print("Dataset not found in local", file=sys.stderr)
+
+    if (all or studio) and token:
+        edit_studio_dataset(team, name, new_name, description, labels)
+
+
+def dataset_stats(
+    catalog: "Catalog",
+    name: str,
+    version: int,
+    show_bytes=False,
+    si=False,
+):
+    stats = catalog.dataset_stats(name, version)
+
+    if stats:
+        print(f"Number of objects: {stats.num_objects}")
+        if show_bytes:
+            print(f"Total objects size: {stats.size}")
+        else:
+            print(f"Total objects size: {utils.sizeof_fmt(stats.size, si=si): >7}")

datachain/cli/commands/du.py

@@ -0,0 +1,14 @@
+from typing import TYPE_CHECKING
+
+from datachain import utils
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+
+def du(catalog: "Catalog", sources, show_bytes=False, si=False, **kwargs):
+    for path, size in catalog.du(sources, **kwargs):
+        if show_bytes:
+            print(f"{size} {path}")
+        else:
+            print(f"{utils.sizeof_fmt(size, si=si): >7} {path}")

datachain/cli/commands/index.py

@@ -0,0 +1,12 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+
+def index(
+    catalog: "Catalog",
+    sources,
+    **kwargs,
+):
+    catalog.index(sources, **kwargs)