datachain 0.8.10__py3-none-any.whl → 0.8.11__py3-none-any.whl

datachain/cache.py

@@ -22,15 +22,15 @@ def try_scandir(path):
         pass
 
 
-def get_temp_cache(tmp_dir: str, prefix: Optional[str] = None) -> "DataChainCache":
+def get_temp_cache(tmp_dir: str, prefix: Optional[str] = None) -> "Cache":
     cache_dir = mkdtemp(prefix=prefix, dir=tmp_dir)
-    return DataChainCache(cache_dir, tmp_dir=tmp_dir)
+    return Cache(cache_dir, tmp_dir=tmp_dir)
 
 
 @contextmanager
 def temporary_cache(
     tmp_dir: str, prefix: Optional[str] = None, delete: bool = True
-) -> Iterator["DataChainCache"]:
+) -> Iterator["Cache"]:
     cache = get_temp_cache(tmp_dir, prefix=prefix)
     try:
         yield cache
@@ -39,7 +39,7 @@ def temporary_cache(
             cache.destroy()
 
 
-class DataChainCache:
+class Cache:
     def __init__(self, cache_dir: str, tmp_dir: str):
         self.odb = LocalHashFileDB(
             LocalFileSystem(),

datachain/catalog/__init__.py

@@ -3,7 +3,6 @@ from .catalog import (
     QUERY_SCRIPT_CANCELED_EXIT_CODE,
     QUERY_SCRIPT_INVALID_LAST_STATEMENT_EXIT_CODE,
     Catalog,
-    parse_edatachain_file,
 )
 from .loader import get_catalog
 
@@ -13,5 +12,4 @@ __all__ = [
     "QUERY_SCRIPT_INVALID_LAST_STATEMENT_EXIT_CODE",
     "Catalog",
     "get_catalog",
-    "parse_edatachain_file",
 ]

datachain/catalog/catalog.py

@@ -4,6 +4,7 @@ import logging
 import os
 import os.path
 import posixpath
+import signal
 import subprocess
 import sys
 import time
@@ -26,11 +27,10 @@ from uuid import uuid4
 
 import requests
 import sqlalchemy as sa
-import yaml
 from sqlalchemy import Column
 from tqdm.auto import tqdm
 
-from datachain.cache import DataChainCache
+from datachain.cache import Cache
 from datachain.client import Client
 from datachain.dataset import (
     DATASET_PREFIX,
@@ -57,7 +57,7 @@ from datachain.node import DirType, Node, NodeWithPath
 from datachain.nodes_thread_pool import NodesThreadPool
 from datachain.remote.studio import StudioClient
 from datachain.sql.types import DateTime, SQLType
-from datachain.utils import DataChainDir, datachain_paths_join
+from datachain.utils import DataChainDir
 
 from .datasource import DataSource
 
@@ -73,7 +73,6 @@ if TYPE_CHECKING:
 logger = logging.getLogger("datachain")
 
 DEFAULT_DATASET_DIR = "dataset"
-DATASET_FILE_SUFFIX = ".edatachain"
 
 TTL_INT = 4 * 60 * 60
 
@@ -99,6 +98,47 @@ def noop(_: str):
     pass
 
 
+class TerminationSignal(RuntimeError):  # noqa: N818
+    def __init__(self, signal):
+        self.signal = signal
+        super().__init__("Received termination signal", signal)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.signal})"
+
+
+if sys.platform == "win32":
+    SIGINT = signal.CTRL_C_EVENT
+else:
+    SIGINT = signal.SIGINT
+
+
+def shutdown_process(
+    proc: subprocess.Popen,
+    interrupt_timeout: Optional[int] = None,
+    terminate_timeout: Optional[int] = None,
+) -> int:
+    """Shut down the process gracefully with SIGINT -> SIGTERM -> SIGKILL."""
+
+    logger.info("sending interrupt signal to the process %s", proc.pid)
+    proc.send_signal(SIGINT)
+
+    logger.info("waiting for the process %s to finish", proc.pid)
+    try:
+        return proc.wait(interrupt_timeout)
+    except subprocess.TimeoutExpired:
+        logger.info(
+            "timed out waiting, sending terminate signal to the process %s", proc.pid
+        )
+        proc.terminate()
+        try:
+            return proc.wait(terminate_timeout)
+        except subprocess.TimeoutExpired:
+            logger.info("timed out waiting, killing the process %s", proc.pid)
+            proc.kill()
+            return proc.wait()
+
+
 def _process_stream(stream: "IO[bytes]", callback: Callable[[str], None]) -> None:
     buffer = b""
     while byt := stream.read(1):  # Read one byte at a time
@@ -247,7 +287,6 @@ class NodeGroup:
     # The source path within the bucket
     # (not including the bucket name or s3:// prefix)
     source_path: str = ""
-    is_edatachain: bool = False
     dataset_name: Optional[str] = None
     dataset_version: Optional[int] = None
     instantiated_nodes: Optional[list[NodeWithPath]] = None
@@ -272,55 +311,11 @@ class NodeGroup:
             self.client.fetch_nodes(self.iternodes(recursive), shared_progress_bar=pbar)
 
 
-def check_output_dataset_file(
-    output: str,
-    force: bool = False,
-    dataset_filename: Optional[str] = None,
-    skip_check_edatachain: bool = False,
-) -> str:
-    """
-    Checks the dataset filename for existence or if it should be force-overwritten.
-    """
-    dataset_file = (
-        dataset_filename if dataset_filename else output + DATASET_FILE_SUFFIX
-    )
-    if not skip_check_edatachain and os.path.exists(dataset_file):
-        if force:
-            os.remove(dataset_file)
-        else:
-            raise RuntimeError(f"Output dataset file already exists: {dataset_file}")
-    return dataset_file
-
-
-def parse_edatachain_file(filename: str) -> list[dict[str, Any]]:
-    with open(filename, encoding="utf-8") as f:
-        contents = yaml.safe_load(f)
-
-    if not isinstance(contents, list):
-        contents = [contents]
-
-    for entry in contents:
-        if not isinstance(entry, dict):
-            raise TypeError(
-                "Failed parsing EDataChain file, "
-                "each data source entry must be a dictionary"
-            )
-        if "data-source" not in entry or "files" not in entry:
-            raise ValueError(
-                "Failed parsing EDataChain file, "
-                "each data source entry must contain the "
-                '"data-source" and "files" keys'
-            )
-
-    return contents
-
-
 def prepare_output_for_cp(
     node_groups: list[NodeGroup],
     output: str,
     force: bool = False,
-    edatachain_only: bool = False,
-    no_edatachain_file: bool = False,
+    no_cp: bool = False,
 ) -> tuple[bool, Optional[str]]:
     total_node_count = 0
     for node_group in node_groups:
@@ -333,7 +328,7 @@ def prepare_output_for_cp(
     always_copy_dir_contents = False
     copy_to_filename = None
 
-    if edatachain_only:
+    if no_cp:
         return always_copy_dir_contents, copy_to_filename
 
     if not os.path.isdir(output):
@@ -358,10 +353,6 @@ def prepare_output_for_cp(
                 copy_to_filename = output
         else:
             raise FileNotFoundError(f"Is not a directory: {output}")
-
-    if copy_to_filename and not no_edatachain_file:
-        raise RuntimeError("File to file cp not supported with .edatachain files!")
-
     return always_copy_dir_contents, copy_to_filename
 
 
@@ -465,8 +456,6 @@ def instantiate_node_groups(
                 copy_to_filename,
                 recursive,
                 copy_dir_contents,
-                source_path,
-                node_group.is_edatachain,
                 node_group.is_dataset,
             )
             if not virtual_only:
@@ -484,24 +473,6 @@ def instantiate_node_groups(
         instantiate_progress_bar.close()
 
 
-def compute_metafile_data(node_groups) -> list[dict[str, Any]]:
-    metafile_data = []
-    for node_group in node_groups:
-        if not node_group.sources:
-            continue
-        listing: Listing = node_group.listing
-        metafile_group = {"data-source": {"uri": listing.uri}, "files": []}
-        for node in node_group.instantiated_nodes:
-            if not node.n.is_dir:
-                metafile_group["files"].append(  # type: ignore [attr-defined]
-                    node.get_metafile_data()
-                )
-        if metafile_group["files"]:
-            metafile_data.append(metafile_group)
-
-    return metafile_data
-
-
 def find_column_to_str(  # noqa: PLR0911
     row: tuple[Any, ...], field_lookup: dict[str, int], src: DataSource, column: str
 ) -> str:
@@ -536,7 +507,7 @@ def find_column_to_str(  # noqa: PLR0911
     return ""
 
 
-def clone_catalog_with_cache(catalog: "Catalog", cache: "DataChainCache") -> "Catalog":
+def clone_catalog_with_cache(catalog: "Catalog", cache: "Cache") -> "Catalog":
     clone = catalog.copy()
     clone.cache = cache
     return clone
@@ -559,7 +530,7 @@ class Catalog:
         datachain_dir.init()
         self.metastore = metastore
         self._warehouse = warehouse
-        self.cache = DataChainCache(datachain_dir.cache, datachain_dir.tmp)
+        self.cache = Cache(datachain_dir.cache, datachain_dir.tmp)
         self.client_config = client_config if client_config is not None else {}
         self._init_params = {
             "cache_dir": cache_dir,
@@ -703,22 +674,8 @@ class Catalog:
         enlisted_sources: list[tuple[bool, bool, Any]] = []
         client_config = client_config or self.client_config
         for src in sources:  # Opt: parallel
-            if src.endswith(DATASET_FILE_SUFFIX) and os.path.isfile(src):
-                # TODO: Also allow using EDataChain files from cloud locations?
-                edatachain_data = parse_edatachain_file(src)
-                indexed_sources = []
-                for ds in edatachain_data:
-                    listing, _, source_path = self.enlist_source(
-                        ds["data-source"]["uri"],
-                        update,
-                        client_config=client_config,
-                    )
-                    paths = datachain_paths_join(
-                        source_path, (f["name"] for f in ds["files"])
-                    )
-                    indexed_sources.append((listing, source_path, paths))
-                enlisted_sources.append((True, False, indexed_sources))
-            elif src.startswith("ds://"):
+            listing: Optional[Listing]
+            if src.startswith("ds://"):
                 ds_name, ds_version = parse_dataset_uri(src)
                 dataset = self.get_dataset(ds_name)
                 if not ds_version:
@@ -796,7 +753,6 @@ class Catalog:
                             listing.client,
                             dsrc,
                             source_path,
-                            is_edatachain=True,
                         )
                     )
             else:
@@ -1360,8 +1316,6 @@ class Catalog:
         local_ds_version: Optional[int] = None,
         cp: bool = False,
         force: bool = False,
-        edatachain: bool = False,
-        edatachain_file: Optional[str] = None,
         *,
         client_config=None,
     ) -> None:
@@ -1373,8 +1327,6 @@ class Catalog:
                 [ds_uri],
                 output,
                 force=force,
-                no_edatachain_file=not edatachain,
-                edatachain_file=edatachain_file,
                 client_config=client_config,
             )
             print(f"Dataset {ds_uri} instantiated locally to {output}")
@@ -1541,8 +1493,6 @@ class Catalog:
         recursive: bool = False,
         no_glob: bool = False,
         no_cp: bool = False,
-        edatachain: bool = False,
-        edatachain_file: Optional[str] = None,
         *,
         client_config=None,
     ) -> None:
@@ -1551,9 +1501,8 @@ class Catalog:
         them into the dataset folder.
         It also adds those files to a dataset in database, which is
         created if doesn't exist yet
-        Optionally, it creates a .edatachain file
         """
-        if not no_cp or edatachain:
+        if not no_cp:
             self.cp(
                 sources,
                 output,
@@ -1561,9 +1510,7 @@ class Catalog:
                 update=update,
                 recursive=recursive,
                 no_glob=no_glob,
-                edatachain_only=no_cp,
-                no_edatachain_file=not edatachain,
-                edatachain_file=edatachain_file,
+                no_cp=no_cp,
                 client_config=client_config,
             )
         else:
@@ -1588,6 +1535,8 @@ class Catalog:
         output_hook: Callable[[str], None] = noop,
         params: Optional[dict[str, str]] = None,
         job_id: Optional[str] = None,
+        interrupt_timeout: Optional[int] = None,
+        terminate_timeout: Optional[int] = None,
     ) -> None:
         cmd = [python_executable, "-c", query_script]
         env = dict(env or os.environ)
@@ -1601,13 +1550,48 @@ class Catalog:
         if capture_output:
             popen_kwargs = {"stdout": subprocess.PIPE, "stderr": subprocess.STDOUT}
 
+        def raise_termination_signal(sig: int, _: Any) -> NoReturn:
+            raise TerminationSignal(sig)
+
+        thread: Optional[Thread] = None
         with subprocess.Popen(cmd, env=env, **popen_kwargs) as proc:  # noqa: S603
-            if capture_output:
-                args = (proc.stdout, output_hook)
-                thread = Thread(target=_process_stream, args=args, daemon=True)
-                thread.start()
-                thread.join()  # wait for the reader thread
+            logger.info("Starting process %s", proc.pid)
+
+            orig_sigint_handler = signal.getsignal(signal.SIGINT)
+            # ignore SIGINT in the main process.
+            # In the terminal, SIGINTs are received by all the processes in
+            # the foreground process group, so the script will receive the signal too.
+            # (If we forward the signal to the child, it will receive it twice.)
+            signal.signal(signal.SIGINT, signal.SIG_IGN)
 
+            orig_sigterm_handler = signal.getsignal(signal.SIGTERM)
+            signal.signal(signal.SIGTERM, raise_termination_signal)
+            try:
+                if capture_output:
+                    args = (proc.stdout, output_hook)
+                    thread = Thread(target=_process_stream, args=args, daemon=True)
+                    thread.start()
+
+                proc.wait()
+            except TerminationSignal as exc:
+                signal.signal(signal.SIGTERM, orig_sigterm_handler)
+                signal.signal(signal.SIGINT, orig_sigint_handler)
+                logging.info("Shutting down process %s, received %r", proc.pid, exc)
+                # Rather than forwarding the signal to the child, we try to shut it down
+                # gracefully. This is because we consider the script to be interactive
+                # and special, so we give it time to cleanup before exiting.
+                shutdown_process(proc, interrupt_timeout, terminate_timeout)
+                if proc.returncode:
+                    raise QueryScriptCancelError(
+                        "Query script was canceled by user", return_code=proc.returncode
+                    ) from exc
+            finally:
+                signal.signal(signal.SIGTERM, orig_sigterm_handler)
+                signal.signal(signal.SIGINT, orig_sigint_handler)
+                if thread:
+                    thread.join()  # wait for the reader thread
+
+        logging.info("Process %s exited with return code %s", proc.pid, proc.returncode)
         if proc.returncode == QUERY_SCRIPT_CANCELED_EXIT_CODE:
             raise QueryScriptCancelError(
                 "Query script was canceled by user",
@@ -1626,17 +1610,14 @@ class Catalog:
         force: bool = False,
         update: bool = False,
         recursive: bool = False,
-        edatachain_file: Optional[str] = None,
-        edatachain_only: bool = False,
-        no_edatachain_file: bool = False,
+        no_cp: bool = False,
         no_glob: bool = False,
         *,
-        client_config=None,
-    ) -> list[dict[str, Any]]:
+        client_config: Optional["dict"] = None,
+    ) -> None:
         """
         This function copies files from cloud sources to local destination directory
         If cloud source is not indexed, or has expired index, it runs indexing
-        It also creates .edatachain file by default, if not specified differently
         """
         client_config = client_config or self.client_config
         node_groups = self.enlist_sources_grouped(
@@ -1647,17 +1628,11 @@ class Catalog:
         )
 
         always_copy_dir_contents, copy_to_filename = prepare_output_for_cp(
-            node_groups, output, force, edatachain_only, no_edatachain_file
+            node_groups, output, force, no_cp
         )
-        dataset_file = check_output_dataset_file(
-            output, force, edatachain_file, no_edatachain_file
-        )
-
         total_size, total_files = collect_nodes_for_cp(node_groups, recursive)
-
-        if total_files == 0:
-            # Nothing selected to cp
-            return []
+        if not total_files:
+            return
 
         desc_max_len = max(len(output) + 16, 19)
         bar_format = (
@@ -1667,7 +1642,7 @@ class Catalog:
             "[{elapsed}<{remaining}, {rate_fmt:>8}]"
         )
 
-        if not edatachain_only:
+        if not no_cp:
             with get_download_bar(bar_format, total_size) as pbar:
                 for node_group in node_groups:
                     node_group.download(recursive=recursive, pbar=pbar)
@@ -1679,21 +1654,10 @@ class Catalog:
             total_files,
             force,
             recursive,
-            edatachain_only,
+            no_cp,
             always_copy_dir_contents,
             copy_to_filename,
         )
-        if no_edatachain_file:
-            return []
-
-        metafile_data = compute_metafile_data(node_groups)
-        if metafile_data:
-            # Don't write the metafile if nothing was copied
-            print(f"Creating '{dataset_file}'")
-            with open(dataset_file, "w", encoding="utf-8") as fd:
-                yaml.dump(metafile_data, fd, sort_keys=False)
-
-        return metafile_data
 
     def du(
         self,

datachain/cli/__init__.py

@@ -47,9 +47,13 @@ def main(argv: Optional[list[str]] = None) -> int:
     logging_level = get_logging_level(args)
     logger.setLevel(logging_level)
 
-    client_config = {
-        "anon": args.anon,
-    }
+    client_config = (
+        {
+            "anon": args.anon,
+        }
+        if getattr(args, "anon", False)
+        else {}
+    )
 
     if args.debug_sql:
         # This also sets this environment variable for any subprocesses
@@ -107,9 +111,6 @@ def handle_cp_command(args, catalog):
         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,
     )
 

datachain/cli/parser/__init__.py

@@ -8,7 +8,14 @@ from datachain.cli.utils import BooleanOptionalAction, KeyValueArgs
 
 from .job import add_jobs_parser
 from .studio import add_auth_parser
-from .utils import FIND_COLUMNS, add_show_args, add_sources_arg, find_columns_type
+from .utils import (
+    FIND_COLUMNS,
+    add_anon_arg,
+    add_show_args,
+    add_sources_arg,
+    add_update_arg,
+    find_columns_type,
+)
 
 
 def get_parser() -> ArgumentParser:  # noqa: PLR0915
@@ -32,19 +39,6 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         "-q", "--quiet", action="count", default=0, help="Be quiet"
     )
 
-    parent_parser.add_argument(
-        "--anon",
-        action="store_true",
-        help="Use anonymous access to storage",
-    )
-    parent_parser.add_argument(
-        "-u",
-        "--update",
-        action="count",
-        default=0,
-        help="Update cached list of files for the sources",
-    )
-
     parent_parser.add_argument(
         "--debug-sql",
         action="store_true",
@@ -92,6 +86,8 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         action="store_true",
         help="Do not expand globs (such as * or ?)",
     )
+    add_anon_arg(parse_cp)
+    add_update_arg(parse_cp)
 
     parse_clone = subp.add_parser(
         "clone", parents=[parent_parser], description="Copy data files from the cloud."
@@ -127,6 +123,8 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         action="store_true",
         help="Do not copy files, just create a dataset",
     )
+    add_anon_arg(parse_clone)
+    add_update_arg(parse_clone)
 
     add_auth_parser(subp, parent_parser)
     add_jobs_parser(subp, parent_parser)
@@ -137,6 +135,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         parents=[parent_parser],
         description="Commands for managing datasets.",
     )
+    add_anon_arg(datasets_parser)
     datasets_subparser = datasets_parser.add_subparsers(
         dest="datasets_cmd",
         help="Use `datachain dataset CMD --help` to display command-specific help",
@@ -336,6 +335,8 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     parse_ls = subp.add_parser(
         "ls", parents=[parent_parser], description="List storage contents."
     )
+    add_anon_arg(parse_ls)
+    add_update_arg(parse_ls)
     add_sources_arg(parse_ls, nargs="*")
     parse_ls.add_argument(
         "-l",
@@ -375,6 +376,8 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         "du", parents=[parent_parser], description="Display space usage."
     )
     add_sources_arg(parse_du)
+    add_anon_arg(parse_du)
+    add_update_arg(parse_du)
     parse_du.add_argument(
         "-b",
         "--bytes",
@@ -404,6 +407,8 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     parse_find = subp.add_parser(
         "find", parents=[parent_parser], description="Search in a directory hierarchy."
     )
+    add_anon_arg(parse_find)
+    add_update_arg(parse_find)
     add_sources_arg(parse_find)
     parse_find.add_argument(
         "--name",
@@ -457,6 +462,8 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     parse_index = subp.add_parser(
         "index", parents=[parent_parser], description="Index storage location."
     )
+    add_anon_arg(parse_index)
+    add_update_arg(parse_index)
     add_sources_arg(parse_index)
 
     show_parser = subp.add_parser(
@@ -480,6 +487,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         parents=[parent_parser],
         description="Create a new dataset with a query script.",
     )
+    add_anon_arg(query_parser)
     query_parser.add_argument(
         "script", metavar="<script.py>", type=str, help="Filepath for script"
     )
@@ -504,14 +512,17 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="Query parameters",
     )
 
-    subp.add_parser(
+    parse_clear_cache = subp.add_parser(
         "clear-cache",
         parents=[parent_parser],
         description="Clear the local file cache.",
     )
-    subp.add_parser(
+    add_anon_arg(parse_clear_cache)
+
+    parse_gc = subp.add_parser(
         "gc", parents=[parent_parser], description="Garbage collect temporary tables."
     )
+    add_anon_arg(parse_gc)
 
     subp.add_parser("internal-run-udf", parents=[parent_parser])
     subp.add_parser("internal-run-udf-worker", parents=[parent_parser])

datachain/cli/parser/studio.py

@@ -1,9 +1,8 @@
 def add_auth_parser(subparsers, parent_parser) -> None:
+    from dvc_studio_client.auth import AVAILABLE_SCOPES
+
     auth_help = "Manage Studio authentication"
-    auth_description = (
-        "Manage authentication and settings for Studio. "
-        "Configure tokens for sharing datasets and using Studio features."
-    )
+    auth_description = "Manage authentication and settings for Studio. "
 
     auth_parser = subparsers.add_parser(
         "auth",
@@ -19,8 +18,10 @@ def add_auth_parser(subparsers, parent_parser) -> None:
     auth_login_help = "Authenticate with Studio"
     auth_login_description = (
         "Authenticate with Studio using default scopes. "
-        "A random name will be assigned as the token name if not specified."
+        "A random name will be assigned if the token name is not specified."
     )
+
+    allowed_scopes = ", ".join(AVAILABLE_SCOPES)
     login_parser = auth_subparser.add_parser(
         "login",
         parents=[parent_parser],
@@ -40,7 +41,7 @@ def add_auth_parser(subparsers, parent_parser) -> None:
         "--scopes",
         action="store",
         default=None,
-        help="Authentication token scopes",
+        help=f"Authentication token scopes. Allowed scopes: {allowed_scopes}",
     )
 
     login_parser.add_argument(

datachain/cli/parser/utils.py

@@ -34,6 +34,24 @@ def add_sources_arg(parser: ArgumentParser, nargs: Union[str, int] = "+") -> Act
     )
 
 
+def add_anon_arg(parser: ArgumentParser) -> None:
+    parser.add_argument(
+        "--anon",
+        action="store_true",
+        help="Use anonymous access to storage",
+    )
+
+
+def add_update_arg(parser: ArgumentParser) -> None:
+    parser.add_argument(
+        "-u",
+        "--update",
+        action="count",
+        default=0,
+        help="Update cached list of files for the sources",
+    )
+
+
 def add_show_args(parser: ArgumentParser) -> None:
     parser.add_argument(
         "--limit",