datachain 0.32.2__tar.gz → 0.33.0__tar.gz

{datachain-0.32.2 → datachain-0.33.0}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.13.0'
+    rev: 'v0.13.1'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.32.2 → datachain-0.33.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.32.2
+Version: 0.33.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -86,6 +86,7 @@ Requires-Dist: psycopg2-binary>=2.9.0; extra == "postgres"
 Provides-Extra: tests
 Requires-Dist: datachain[audio,hf,postgres,remote,torch,vector,video]; extra == "tests"
 Requires-Dist: pytest<9,>=8; extra == "tests"
+Requires-Dist: pytest-asyncio; extra == "tests"
 Requires-Dist: pytest-sugar>=0.9.6; extra == "tests"
 Requires-Dist: pytest-cov>=4.1.0; extra == "tests"
 Requires-Dist: pytest-mock>=3.12.0; extra == "tests"

{datachain-0.32.2 → datachain-0.33.0}/pyproject.toml

@@ -102,6 +102,7 @@ postgres = [
 tests = [
   "datachain[torch,audio,remote,vector,hf,video,postgres]",
   "pytest>=8,<9",
+  "pytest-asyncio",
   "pytest-sugar>=0.9.6",
   "pytest-cov>=4.1.0",
   "pytest-mock>=3.12.0",

datachain-0.33.0/src/datachain/checkpoint.py

@@ -0,0 +1,44 @@
+import uuid
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Union
+
+
+@dataclass
+class Checkpoint:
+    """
+    Represents a checkpoint within a job run.
+
+    A checkpoint marks a successfully completed stage of execution. In the event
+    of a failure, the job can resume from the most recent checkpoint rather than
+    starting over from the beginning.
+
+    Checkpoints can also be created in a "partial" mode, which indicates that the
+    work at this stage was only partially completed. For example, if a failure
+    occurs halfway through running a UDF, already computed results can still be
+    saved, allowing the job to resume from that partially completed state on
+    restart.
+    """
+
+    id: str
+    job_id: str
+    hash: str
+    partial: bool
+    created_at: datetime
+
+    @classmethod
+    def parse(
+        cls,
+        id: Union[str, uuid.UUID],
+        job_id: str,
+        _hash: str,
+        partial: bool,
+        created_at: datetime,
+    ) -> "Checkpoint":
+        return cls(
+            str(id),
+            job_id,
+            _hash,
+            bool(partial),
+            created_at,
+        )

{datachain-0.32.2 → datachain-0.33.0}/src/datachain/client/fsspec.py

@@ -93,10 +93,11 @@ class Client(ABC):
         self.uri = self.get_uri(self.name)
 
     @staticmethod
-    def get_implementation(url: Union[str, os.PathLike[str]]) -> type["Client"]:
+    def get_implementation(url: Union[str, os.PathLike[str]]) -> type["Client"]:  # noqa: PLR0911
         from .azure import AzureClient
         from .gcs import GCSClient
         from .hf import HfClient
+        from .http import HTTPClient, HTTPSClient
         from .local import FileClient
         from .s3 import ClientS3
 
@@ -114,6 +115,10 @@ class Client(ABC):
             return FileClient
         if protocol == HfClient.protocol:
             return HfClient
+        if protocol == HTTPClient.protocol:
+            return HTTPClient
+        if protocol == HTTPSClient.protocol:
+            return HTTPSClient
 
         raise NotImplementedError(f"Unsupported protocol: {protocol}")
 

datachain-0.33.0/src/datachain/client/http.py

@@ -0,0 +1,157 @@
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, cast
+from urllib.parse import urlparse
+
+from fsspec.implementations.http import HTTPFileSystem
+
+from datachain.dataset import StorageURI
+from datachain.lib.file import File
+
+from .fsspec import Client
+
+if TYPE_CHECKING:
+    from datachain.cache import Cache
+
+
+class HTTPClient(Client):
+    FS_CLASS = HTTPFileSystem
+    PREFIX: ClassVar[str] = "http://"
+    protocol: ClassVar[str] = "http"
+
+    @classmethod
+    def create_fs(cls, **kwargs) -> HTTPFileSystem:
+        # Configure HTTPFileSystem options
+        kwargs.setdefault("simple_links", True)
+        kwargs.setdefault("same_scheme", True)
+        kwargs.setdefault("cache_type", "bytes")
+
+        kwargs.pop("version_aware", None)
+
+        fs = cls.FS_CLASS(**kwargs)
+        fs.invalidate_cache()
+        return cast("HTTPFileSystem", fs)
+
+    @classmethod
+    def from_name(
+        cls,
+        name: str,
+        cache: "Cache",
+        kwargs: dict[str, Any],
+    ) -> "HTTPClient":
+        parsed = urlparse(name)
+
+        if parsed.scheme:
+            name = parsed.netloc + parsed.path
+
+        return cls(name, kwargs, cache)
+
+    @classmethod
+    def split_url(cls, url: str) -> tuple[str, str]:
+        """Split HTTP/HTTPS URL into domain (bucket equivalent) and path."""
+        parsed = urlparse(url)
+        domain = parsed.netloc
+        path = parsed.path.lstrip("/")
+
+        if parsed.query:
+            path += f"?{parsed.query}"
+        if parsed.fragment:
+            path += f"#{parsed.fragment}"
+
+        return domain, path
+
+    @classmethod
+    def get_uri(cls, name: str) -> "StorageURI":
+        if not name.startswith(("http://", "https://")):
+            return StorageURI(f"{cls.PREFIX}{name}")
+        return StorageURI(name)
+
+    @classmethod
+    def is_root_url(cls, url: str) -> bool:
+        parsed = urlparse(url)
+        return parsed.path in ("", "/") and not parsed.query and not parsed.fragment
+
+    def get_full_path(self, rel_path: str, version_id: Optional[str] = None) -> str:
+        if self.name.startswith(("http://", "https://")):
+            base_url = self.name
+        else:
+            if rel_path and "/" in rel_path:
+                first_part = rel_path.split("/")[0]
+                if "." in first_part and not first_part.startswith("."):
+                    return f"{self.protocol}://{rel_path}"
+
+            base_url = f"{self.protocol}://{self.name}"
+
+        if rel_path:
+            if not base_url.endswith("/") and not rel_path.startswith("/"):
+                base_url += "/"
+            full_url = base_url + rel_path
+        else:
+            full_url = base_url
+
+        return full_url
+
+    def url(self, path: str, expires: int = 3600, **kwargs) -> str:
+        """
+        Generate URL for the given path.
+        Note: HTTP URLs don't support signed/expiring URLs.
+        """
+        return self.get_full_path(path, kwargs.pop("version_id", None))
+
+    def info_to_file(self, v: dict[str, Any], path: str) -> File:
+        etag = v.get("ETag", "").strip('"')
+        last_modified = v.get("last_modified")
+        if last_modified:
+            if isinstance(last_modified, str):
+                try:
+                    from email.utils import parsedate_to_datetime
+
+                    last_modified = parsedate_to_datetime(last_modified)
+                except (ValueError, TypeError):
+                    last_modified = datetime.now(timezone.utc)
+            elif isinstance(last_modified, (int, float)):
+                last_modified = datetime.fromtimestamp(last_modified, timezone.utc)
+        else:
+            last_modified = datetime.now(timezone.utc)
+
+        return File(
+            source=self.uri,
+            path=path,
+            size=v.get("size", 0),
+            etag=etag,
+            version="",
+            is_latest=True,
+            last_modified=last_modified,
+        )
+
+    def upload(self, data: bytes, path: str) -> "File":
+        raise NotImplementedError(
+            "HTTP/HTTPS client is read-only. Upload operations are not supported."
+        )
+
+    def get_file_info(self, path: str, version_id: Optional[str] = None) -> "File":
+        info = self.fs.info(self.get_full_path(path))
+        return self.info_to_file(info, path)
+
+    def open_object(self, file: "File", use_cache: bool = True, cb=None):
+        from datachain.client.fileslice import FileWrapper
+
+        if use_cache and (cache_path := self.cache.get_path(file)):
+            return open(cache_path, mode="rb")
+
+        assert not file.location
+        return FileWrapper(
+            self.fs.open(self.get_full_path(file.get_path_normalized())),
+            cb or (lambda x: None),
+        )
+
+    async def get_file(self, lpath, rpath, callback, version_id: Optional[str] = None):
+        return await self.fs._get_file(lpath, rpath, callback=callback)
+
+    async def _fetch_dir(self, prefix: str, pbar, result_queue) -> set[str]:
+        full_url = self.get_full_path(prefix)
+        raise NotImplementedError(f"Cannot download file from {full_url}")
+
+
+class HTTPSClient(HTTPClient):
+    protocol = "https"
+    PREFIX = "https://"

{datachain-0.32.2 → datachain-0.33.0}/src/datachain/data_storage/metastore.py

@@ -13,6 +13,7 @@ from uuid import uuid4
 from sqlalchemy import (
     JSON,
     BigInteger,
+    Boolean,
     Column,
     DateTime,
     ForeignKey,
@@ -24,6 +25,7 @@ from sqlalchemy import (
 )
 from sqlalchemy.sql import func as f
 
+from datachain.checkpoint import Checkpoint
 from datachain.data_storage import JobQueryType, JobStatus
 from datachain.data_storage.serializer import Serializable
 from datachain.dataset import (
@@ -36,6 +38,7 @@ from datachain.dataset import (
     StorageURI,
 )
 from datachain.error import (
+    CheckpointNotFoundError,
     DatasetNotFoundError,
     DatasetVersionNotFoundError,
     NamespaceDeleteNotAllowedError,
@@ -75,6 +78,7 @@ class AbstractMetastore(ABC, Serializable):
     dataset_list_version_class: type[DatasetListVersion] = DatasetListVersion
     dependency_class: type[DatasetDependency] = DatasetDependency
     job_class: type[Job] = Job
+    checkpoint_class: type[Checkpoint] = Checkpoint
 
     def __init__(
         self,
@@ -431,6 +435,35 @@ class AbstractMetastore(ABC, Serializable):
     def get_job_status(self, job_id: str) -> Optional[JobStatus]:
         """Returns the status of the given job."""
 
+    #
+    # Checkpoints
+    #
+
+    @abstractmethod
+    def list_checkpoints(self, job_id: str, conn=None) -> Iterator["Checkpoint"]:
+        """Returns all checkpoints related to some job"""
+
+    @abstractmethod
+    def get_checkpoint_by_id(self, checkpoint_id: str, conn=None) -> Checkpoint:
+        """Gets single checkpoint by id"""
+
+    def find_checkpoint(
+        self, job_id: str, _hash: str, partial: bool = False, conn=None
+    ) -> Optional[Checkpoint]:
+        """
+        Tries to find checkpoint for a job with specific hash and optionally partial
+        """
+
+    @abstractmethod
+    def create_checkpoint(
+        self,
+        job_id: str,
+        _hash: str,
+        partial: bool = False,
+        conn: Optional[Any] = None,
+    ) -> Checkpoint:
+        """Creates new checkpoint"""
+
 
 class AbstractDBMetastore(AbstractMetastore):
     """
@@ -446,6 +479,7 @@ class AbstractDBMetastore(AbstractMetastore):
     DATASET_VERSION_TABLE = "datasets_versions"
     DATASET_DEPENDENCY_TABLE = "datasets_dependencies"
     JOBS_TABLE = "jobs"
+    CHECKPOINTS_TABLE = "checkpoints"
 
     db: "DatabaseEngine"
 
@@ -689,9 +723,6 @@ class AbstractDBMetastore(AbstractMetastore):
             return self._projects.select()
         return select(*columns)
 
-    def _projects_update(self) -> "Update":
-        return self._projects.update()
-
     def _projects_delete(self) -> "Delete":
         return self._projects.delete()
 
@@ -839,6 +870,16 @@ class AbstractDBMetastore(AbstractMetastore):
 
         return self.get_project(name, namespace.name)
 
+    def _projects_base_query(self) -> "Select":
+        n = self._namespaces
+        p = self._projects
+
+        query = self._projects_select(
+            *(getattr(n.c, f) for f in self._namespaces_fields),
+            *(getattr(p.c, f) for f in self._projects_fields),
+        )
+        return query.select_from(n.join(p, n.c.id == p.c.namespace_id))
+
     def get_project(
         self, name: str, namespace_name: str, create: bool = False, conn=None
     ) -> Project:
@@ -854,11 +895,7 @@ class AbstractDBMetastore(AbstractMetastore):
             create = True
             validate = False
 
-        query = self._projects_select(
-            *(getattr(n.c, f) for f in self._namespaces_fields),
-            *(getattr(p.c, f) for f in self._projects_fields),
-        )
-        query = query.select_from(n.join(p, n.c.id == p.c.namespace_id)).where(
+        query = self._projects_base_query().where(
             p.c.name == name, n.c.name == namespace_name
         )
 
@@ -873,16 +910,9 @@ class AbstractDBMetastore(AbstractMetastore):
 
     def get_project_by_id(self, project_id: int, conn=None) -> Project:
         """Gets a single project by id"""
-        n = self._namespaces
         p = self._projects
 
-        query = self._projects_select(
-            *(getattr(n.c, f) for f in self._namespaces_fields),
-            *(getattr(p.c, f) for f in self._projects_fields),
-        )
-        query = query.select_from(n.join(p, n.c.id == p.c.namespace_id)).where(
-            p.c.id == project_id
-        )
+        query = self._projects_base_query().where(p.c.id == project_id)
 
         rows = list(self.db.execute(query, conn=conn))
         if not rows:
@@ -891,7 +921,8 @@ class AbstractDBMetastore(AbstractMetastore):
 
     def count_projects(self, namespace_id: Optional[int] = None) -> int:
         p = self._projects
-        query = self._projects_select()
+
+        query = self._projects_base_query()
         if namespace_id:
             query = query.where(p.c.namespace_id == namespace_id)
 
@@ -917,17 +948,12 @@ class AbstractDBMetastore(AbstractMetastore):
         """
         Gets a list of projects inside some namespace, or in all namespaces
         """
-        n = self._namespaces
         p = self._projects
 
-        query = self._projects_select(
-            *(getattr(n.c, f) for f in self._namespaces_fields),
-            *(getattr(p.c, f) for f in self._projects_fields),
-        )
-        query = query.select_from(n.join(p, n.c.id == p.c.namespace_id))
+        query = self._projects_base_query()
 
         if namespace_id:
-            query = query.where(n.c.id == namespace_id)
+            query = query.where(p.c.namespace_id == namespace_id)
 
         rows = list(self.db.execute(query, conn=conn))
 
@@ -1671,3 +1697,106 @@ class AbstractDBMetastore(AbstractMetastore):
         if not results:
             return None
         return results[0][0]
+
+    #
+    # Checkpoints
+    #
+
+    @staticmethod
+    def _checkpoints_columns() -> "list[SchemaItem]":
+        return [
+            Column(
+                "id",
+                Text,
+                default=uuid4,
+                primary_key=True,
+                nullable=False,
+            ),
+            Column("job_id", Text, nullable=True),
+            Column("hash", Text, nullable=False),
+            Column("partial", Boolean, default=False),
+            Column("created_at", DateTime(timezone=True), nullable=False),
+            UniqueConstraint("job_id", "hash"),
+        ]
+
+    @cached_property
+    def _checkpoints_fields(self) -> list[str]:
+        return [c.name for c in self._checkpoints_columns() if c.name]  # type: ignore[attr-defined]
+
+    @cached_property
+    def _checkpoints(self) -> "Table":
+        return Table(
+            self.CHECKPOINTS_TABLE,
+            self.db.metadata,
+            *self._checkpoints_columns(),
+        )
+
+    @abstractmethod
+    def _checkpoints_insert(self) -> "Insert": ...
+
+    def _checkpoints_select(self, *columns) -> "Select":
+        if not columns:
+            return self._checkpoints.select()
+        return select(*columns)
+
+    def _checkpoints_delete(self) -> "Delete":
+        return self._checkpoints.delete()
+
+    def _checkpoints_query(self):
+        return self._checkpoints_select(
+            *[getattr(self._checkpoints.c, f) for f in self._checkpoints_fields]
+        )
+
+    def create_checkpoint(
+        self,
+        job_id: str,
+        _hash: str,
+        partial: bool = False,
+        conn: Optional[Any] = None,
+    ) -> Checkpoint:
+        """
+        Creates a new job query step.
+        """
+        checkpoint_id = str(uuid4())
+        self.db.execute(
+            self._checkpoints_insert().values(
+                id=checkpoint_id,
+                job_id=job_id,
+                hash=_hash,
+                partial=partial,
+                created_at=datetime.now(timezone.utc),
+            ),
+            conn=conn,
+        )
+        return self.get_checkpoint_by_id(checkpoint_id)
+
+    def list_checkpoints(self, job_id: str, conn=None) -> Iterator["Checkpoint"]:
+        """List checkpoints by job id."""
+        query = self._checkpoints_query().where(self._checkpoints.c.job_id == job_id)
+        rows = list(self.db.execute(query, conn=conn))
+
+        yield from [self.checkpoint_class.parse(*r) for r in rows]
+
+    def get_checkpoint_by_id(self, checkpoint_id: str, conn=None) -> Checkpoint:
+        """Returns the checkpoint with the given ID."""
+        ch = self._checkpoints
+        query = self._checkpoints_select(ch).where(ch.c.id == checkpoint_id)
+        rows = list(self.db.execute(query, conn=conn))
+        if not rows:
+            raise CheckpointNotFoundError(f"Checkpoint {checkpoint_id} not found")
+        return self.checkpoint_class.parse(*rows[0])
+
+    def find_checkpoint(
+        self, job_id: str, _hash: str, partial: bool = False, conn=None
+    ) -> Optional[Checkpoint]:
+        """
+        Tries to find checkpoint for a job with specific hash and optionally partial
+        """
+        ch = self._checkpoints
+        query = self._checkpoints_select(ch).where(
+            ch.c.job_id == job_id, ch.c.hash == _hash, ch.c.partial == partial
+        )
+        rows = list(self.db.execute(query, conn=conn))
+        if not rows:
+            return None
+        return self.checkpoint_class.parse(*rows[0])

{datachain-0.32.2 → datachain-0.33.0}/src/datachain/data_storage/schema.py

@@ -51,7 +51,7 @@ def dedup_columns(columns: Iterable[sa.Column]) -> list[sa.Column]:
     """
     c_set: dict[str, sa.Column] = {}
     for c in columns:
-        if (ec := c_set.get(c.name, None)) is not None:
+        if (ec := c_set.get(c.name)) is not None:
             if str(ec.type) != str(c.type):
                 raise ValueError(
                     f"conflicting types for column {c.name}:{c.type!s} and {ec.type!s}"

{datachain-0.32.2 → datachain-0.33.0}/src/datachain/data_storage/sqlite.py

@@ -459,6 +459,8 @@ class SQLiteMetastore(AbstractDBMetastore):
         self.default_table_names.append(self._datasets_dependencies.name)
         self.db.create_table(self._jobs, if_not_exists=True)
         self.default_table_names.append(self._jobs.name)
+        self.db.create_table(self._checkpoints, if_not_exists=True)
+        self.default_table_names.append(self._checkpoints.name)
 
     def _init_namespaces_projects(self) -> None:
         """
@@ -543,6 +545,12 @@ class SQLiteMetastore(AbstractDBMetastore):
     def _jobs_insert(self) -> "Insert":
         return sqlite.insert(self._jobs)
 
+    #
+    # Checkpoints
+    #
+    def _checkpoints_insert(self) -> "Insert":
+        return sqlite.insert(self._checkpoints)
+
     #
     # Namespaces
     #

{datachain-0.32.2 → datachain-0.33.0}/src/datachain/error.py

@@ -97,3 +97,7 @@ class TableMissingError(DataChainError):
 
 class OutdatedDatabaseSchemaError(DataChainError):
     pass
+
+
+class CheckpointNotFoundError(NotFoundError):
+    pass

{datachain-0.32.2 → datachain-0.33.0}/src/datachain/lib/dc/datachain.py

@@ -1,4 +1,5 @@
 import copy
+import hashlib
 import os
 import os.path
 import sys
@@ -18,6 +19,7 @@ from typing import (
     cast,
     overload,
 )
+from uuid import uuid4
 
 import sqlalchemy
 import ujson as json
@@ -665,7 +667,7 @@ class DataChain:
                     name, namespace=namespace_name, project=project_name, **kwargs
                 )
 
-        return self._evolve(
+        result = self._evolve(
             query=self._query.save(
                 name=name,
                 version=version,
@@ -678,6 +680,16 @@ class DataChain:
             )
         )
 
+        if job_id := os.getenv("DATACHAIN_JOB_ID"):
+            catalog.metastore.create_checkpoint(
+                job_id,  # type: ignore[arg-type]
+                _hash=hashlib.sha256(  # TODO this will be replaced with self.hash()
+                    str(uuid4()).encode()
+                ).hexdigest(),
+            )
+
+        return result
+
     def apply(self, func, *args, **kwargs):
         """Apply any function to the chain.
 

{datachain-0.32.2 → datachain-0.33.0}/src/datachain/lib/file.py

@@ -332,7 +332,10 @@ class File(DataModel):
 
     @classmethod
     def upload(
-        cls, data: bytes, path: str, catalog: Optional["Catalog"] = None
+        cls,
+        data: bytes,
+        path: Union[str, os.PathLike[str]],
+        catalog: Optional["Catalog"] = None,
     ) -> "Self":
         if catalog is None:
             from datachain.catalog.loader import get_catalog
@@ -340,8 +343,10 @@ class File(DataModel):
             catalog = get_catalog()
         from datachain.client.fsspec import Client
 
-        client_cls = Client.get_implementation(path)
-        source, rel_path = client_cls.split_url(path)
+        path_str = stringify_path(path)
+
+        client_cls = Client.get_implementation(path_str)
+        source, rel_path = client_cls.split_url(path_str)
 
         client = catalog.get_client(client_cls.get_uri(source))
         file = client.upload(data, rel_path)
@@ -351,7 +356,9 @@ class File(DataModel):
         return file
 
     @classmethod
-    def at(cls, uri: str, session: Optional["Session"] = None) -> "Self":
+    def at(
+        cls, uri: Union[str, os.PathLike[str]], session: Optional["Session"] = None
+    ) -> "Self":
         """Construct a File from a full URI in one call.
 
         Example:
@@ -364,9 +371,10 @@ class File(DataModel):
         if session is None:
             session = Session.get()
         catalog = session.catalog
+        uri_str = stringify_path(uri)
 
-        client_cls = Client.get_implementation(uri)
-        source, rel_path = client_cls.split_url(uri)
+        client_cls = Client.get_implementation(uri_str)
+        source, rel_path = client_cls.split_url(uri_str)
         file = cls(source=client_cls.get_uri(source), path=rel_path)
         file._set_stream(catalog)
         return file

{datachain-0.32.2 → datachain-0.33.0}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.32.2
+Version: 0.33.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -86,6 +86,7 @@ Requires-Dist: psycopg2-binary>=2.9.0; extra == "postgres"
 Provides-Extra: tests
 Requires-Dist: datachain[audio,hf,postgres,remote,torch,vector,video]; extra == "tests"
 Requires-Dist: pytest<9,>=8; extra == "tests"
+Requires-Dist: pytest-asyncio; extra == "tests"
 Requires-Dist: pytest-sugar>=0.9.6; extra == "tests"
 Requires-Dist: pytest-cov>=4.1.0; extra == "tests"
 Requires-Dist: pytest-mock>=3.12.0; extra == "tests"