datachain 0.30.5__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/query/session.py

@@ -1,21 +1,37 @@
 import atexit
-import gc
 import logging
+import os
 import re
 import sys
-from typing import TYPE_CHECKING, ClassVar, Optional
+import traceback
+from collections.abc import Callable
+from typing import TYPE_CHECKING, ClassVar
 from uuid import uuid4
+from weakref import WeakSet
 
 from datachain.catalog import get_catalog
-from datachain.error import TableMissingError
+from datachain.data_storage import JobQueryType, JobStatus
+from datachain.error import JobNotFoundError, TableMissingError
 
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
-    from datachain.dataset import DatasetRecord
+    from datachain.job import Job
 
 logger = logging.getLogger("datachain")
 
 
+def is_script_run() -> bool:
+    """
+    Returns True if this was ran as python script, e.g python my_script.py.
+    Otherwise (if interactive or module run) returns False.
+    """
+    try:
+        argv0 = sys.argv[0]
+    except (IndexError, AttributeError):
+        return False
+    return bool(argv0) and argv0 not in ("-c", "-m", "ipython")
+
+
 class Session:
     """
     Session is a context that keeps track of temporary DataChain datasets for a proper
@@ -39,10 +55,18 @@ class Session:
     catalog (Catalog): Catalog object.
     """
 
-    GLOBAL_SESSION_CTX: Optional["Session"] = None
+    GLOBAL_SESSION_CTX: "Session | None" = None
     SESSION_CONTEXTS: ClassVar[list["Session"]] = []
+    _ALL_SESSIONS: ClassVar[WeakSet["Session"]] = WeakSet()
     ORIGINAL_EXCEPT_HOOK = None
 
+    # Job management - class-level to ensure one job per process
+    _CURRENT_JOB: ClassVar["Job | None"] = None
+    _JOB_STATUS: ClassVar[JobStatus | None] = None
+    _OWNS_JOB: ClassVar[bool | None] = None
+    _JOB_HOOKS_REGISTERED: ClassVar[bool] = False
+    _JOB_FINALIZE_HOOK: ClassVar[Callable[[], None] | None] = None
+
     DATASET_PREFIX = "session_"
     GLOBAL_SESSION_NAME = "global"
     SESSION_UUID_LEN = 6
@@ -51,8 +75,8 @@ class Session:
     def __init__(
         self,
         name="",
-        catalog: Optional["Catalog"] = None,
-        client_config: Optional[dict] = None,
+        catalog: "Catalog | None" = None,
+        client_config: dict | None = None,
         in_memory: bool = False,
     ):
         if re.match(r"^[0-9a-zA-Z]*$", name) is None:
@@ -69,7 +93,7 @@ class Session:
         self.catalog = catalog or get_catalog(
             client_config=client_config, in_memory=in_memory
         )
-        self.dataset_versions: list[tuple[DatasetRecord, str, bool]] = []
+        Session._ALL_SESSIONS.add(self)
 
     def __enter__(self):
         # Push the current context onto the stack
@@ -78,9 +102,8 @@ class Session:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if exc_type:
-            self._cleanup_created_versions()
-
+        # Don't cleanup created versions on exception
+        # Datasets should persist even if the session fails
         self._cleanup_temp_datasets()
         if self.is_new_catalog:
             self.catalog.metastore.close_on_exit()
@@ -88,11 +111,116 @@ class Session:
 
         if Session.SESSION_CONTEXTS:
             Session.SESSION_CONTEXTS.pop()
+        Session._ALL_SESSIONS.discard(self)
 
-    def add_dataset_version(
-        self, dataset: "DatasetRecord", version: str, listing: bool = False
-    ) -> None:
-        self.dataset_versions.append((dataset, version, listing))
+    def get_or_create_job(self) -> "Job":
+        """
+        Get or create a Job for this process.
+
+        Returns:
+            Job: The active Job instance.
+
+        Behavior:
+            - If a job already exists, it is returned.
+            - If ``DATACHAIN_JOB_ID`` is set, the corresponding job is fetched.
+            - Otherwise, a new job is created:
+                * Name = absolute path to the Python script.
+                * Query = empty string.
+                * Parent = last job with the same name, if available.
+                * Status = "running".
+              Exit hooks are registered to finalize the job.
+
+        Note:
+            Job is shared across all Session instances to ensure one job per process.
+        """
+        if Session._CURRENT_JOB:
+            return Session._CURRENT_JOB
+
+        if env_job_id := os.getenv("DATACHAIN_JOB_ID"):
+            # SaaS run: just fetch existing job
+            Session._CURRENT_JOB = self.catalog.metastore.get_job(env_job_id)
+            if not Session._CURRENT_JOB:
+                raise JobNotFoundError(
+                    f"Job {env_job_id} from DATACHAIN_JOB_ID env not found"
+                )
+            Session._OWNS_JOB = False
+        else:
+            # Local run: create new job
+            if is_script_run():
+                script = os.path.abspath(sys.argv[0])
+            else:
+                # Interactive session or module run - use unique name to avoid
+                # linking unrelated sessions
+                script = str(uuid4())
+            python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+
+            # try to find the parent job
+            parent = self.catalog.metastore.get_last_job_by_name(script)
+
+            job_id = self.catalog.metastore.create_job(
+                name=script,
+                query="",
+                query_type=JobQueryType.PYTHON,
+                status=JobStatus.RUNNING,
+                python_version=python_version,
+                parent_job_id=parent.id if parent else None,
+            )
+            Session._CURRENT_JOB = self.catalog.metastore.get_job(job_id)
+            Session._OWNS_JOB = True
+            Session._JOB_STATUS = JobStatus.RUNNING
+
+            # register cleanup hooks only once
+            if not Session._JOB_HOOKS_REGISTERED:
+
+                def _finalize_success_hook() -> None:
+                    self._finalize_job_success()
+
+                Session._JOB_FINALIZE_HOOK = _finalize_success_hook
+                atexit.register(Session._JOB_FINALIZE_HOOK)
+                Session._JOB_HOOKS_REGISTERED = True
+
+        assert Session._CURRENT_JOB is not None
+        return Session._CURRENT_JOB
+
+    def _finalize_job_success(self):
+        """Mark the current job as completed."""
+        if (
+            Session._CURRENT_JOB
+            and Session._OWNS_JOB
+            and Session._JOB_STATUS == JobStatus.RUNNING
+        ):
+            self.catalog.metastore.set_job_status(
+                Session._CURRENT_JOB.id, JobStatus.COMPLETE
+            )
+            Session._JOB_STATUS = JobStatus.COMPLETE
+
+    def _finalize_job_as_canceled(self):
+        """Mark the current job as canceled."""
+        if (
+            Session._CURRENT_JOB
+            and Session._OWNS_JOB
+            and Session._JOB_STATUS == JobStatus.RUNNING
+        ):
+            self.catalog.metastore.set_job_status(
+                Session._CURRENT_JOB.id, JobStatus.CANCELED
+            )
+            Session._JOB_STATUS = JobStatus.CANCELED
+
+    def _finalize_job_as_failed(self, exc_type, exc_value, tb):
+        """Mark the current job as failed with error details."""
+        if (
+            Session._CURRENT_JOB
+            and Session._OWNS_JOB
+            and Session._JOB_STATUS == JobStatus.RUNNING
+        ):
+            error_stack = "".join(traceback.format_exception(exc_type, exc_value, tb))
+            self.catalog.metastore.set_job_status(
+                Session._CURRENT_JOB.id,
+                JobStatus.FAILED,
+                error_message=str(exc_value),
+                error_stack=error_stack,
+            )
+            Session._JOB_STATUS = JobStatus.FAILED
 
     def generate_temp_dataset_name(self) -> str:
         return self.get_temp_prefix() + uuid4().hex[: self.TEMP_TABLE_UUID_LEN]
@@ -113,22 +241,12 @@ class Session:
         except TableMissingError:
             pass
 
-    def _cleanup_created_versions(self) -> None:
-        if not self.dataset_versions:
-            return
-
-        for dataset, version, listing in self.dataset_versions:
-            if not listing:
-                self.catalog.remove_dataset_version(dataset, version)
-
-        self.dataset_versions.clear()
-
     @classmethod
     def get(
         cls,
-        session: Optional["Session"] = None,
-        catalog: Optional["Catalog"] = None,
-        client_config: Optional[dict] = None,
+        session: "Session | None" = None,
+        catalog: "Catalog | None" = None,
+        client_config: dict | None = None,
         in_memory: bool = False,
     ) -> "Session":
         """Creates a Session() object from a catalog.
@@ -173,33 +291,72 @@ class Session:
 
     @staticmethod
     def except_hook(exc_type, exc_value, exc_traceback):
-        Session.GLOBAL_SESSION_CTX.__exit__(exc_type, exc_value, exc_traceback)
+        if Session.GLOBAL_SESSION_CTX:
+            # Handle KeyboardInterrupt specially - mark as canceled and exit with
+            # signal code
+            if exc_type is KeyboardInterrupt:
+                Session.GLOBAL_SESSION_CTX._finalize_job_as_canceled()
+            else:
+                Session.GLOBAL_SESSION_CTX._finalize_job_as_failed(
+                    exc_type, exc_value, exc_traceback
+                )
+            Session.GLOBAL_SESSION_CTX.__exit__(exc_type, exc_value, exc_traceback)
+
         Session._global_cleanup()
 
+        # Always delegate to original hook if it exists
         if Session.ORIGINAL_EXCEPT_HOOK:
             Session.ORIGINAL_EXCEPT_HOOK(exc_type, exc_value, exc_traceback)
 
+        if exc_type is KeyboardInterrupt:
+            # Exit with SIGINT signal code (128 + 2 = 130, or -2 in subprocess terms)
+            sys.exit(130)
+
     @classmethod
     def cleanup_for_tests(cls):
+        cls._close_all_contexts()
         if cls.GLOBAL_SESSION_CTX is not None:
             cls.GLOBAL_SESSION_CTX.__exit__(None, None, None)
             cls.GLOBAL_SESSION_CTX = None
             atexit.unregister(cls._global_cleanup)
 
+        # Reset job-related class variables
+        if cls._JOB_FINALIZE_HOOK:
+            try:
+                atexit.unregister(cls._JOB_FINALIZE_HOOK)
+            except ValueError:
+                pass  # Hook was not registered
+        cls._CURRENT_JOB = None
+        cls._JOB_STATUS = None
+        cls._OWNS_JOB = None
+        cls._JOB_HOOKS_REGISTERED = False
+        cls._JOB_FINALIZE_HOOK = None
+
         if cls.ORIGINAL_EXCEPT_HOOK:
             sys.excepthook = cls.ORIGINAL_EXCEPT_HOOK
 
     @staticmethod
     def _global_cleanup():
+        Session._close_all_contexts()
         if Session.GLOBAL_SESSION_CTX is not None:
             Session.GLOBAL_SESSION_CTX.__exit__(None, None, None)
 
-        for obj in gc.get_objects():  # Get all tracked objects
+        for session in list(Session._ALL_SESSIONS):
             try:
-                if isinstance(obj, Session):
-                    # Cleanup temp dataset for session variables.
-                    obj.__exit__(None, None, None)
+                session.__exit__(None, None, None)
             except ReferenceError:
                 continue  # Object has been finalized already
             except Exception as e:  # noqa: BLE001
                 logger.error(f"Exception while cleaning up session: {e}")  # noqa: G004
+
+    @classmethod
+    def _close_all_contexts(cls) -> None:
+        while cls.SESSION_CONTEXTS:
+            session = cls.SESSION_CONTEXTS.pop()
+            try:
+                session.__exit__(None, None, None)
+            except Exception as exc:  # noqa: BLE001
+                logger.error(
+                    "Exception while closing session context during cleanup: %s",
+                    exc,
+                )

datachain/query/udf.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Callable, Optional, TypedDict, Union
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, TypedDict
 
 if TYPE_CHECKING:
     from sqlalchemy import Select, Table
@@ -17,10 +18,11 @@ class UdfInfo(TypedDict):
     query: "Select"
     udf_fields: list[str]
     batching: "BatchingStrategy"
-    processes: Optional[int]
+    processes: int | None
     is_generator: bool
     cache: bool
     rows_total: int
+    batch_size: int
 
 
 class AbstractUDFDistributor(ABC):
@@ -32,13 +34,14 @@ class AbstractUDFDistributor(ABC):
         query: "Select",
         udf_data: bytes,
         batching: "BatchingStrategy",
-        workers: Union[bool, int],
-        processes: Union[bool, int],
+        workers: bool | int,
+        processes: bool | int,
         udf_fields: list[str],
         rows_total: int,
         use_cache: bool,
         is_generator: bool = False,
-        min_task_size: Optional[Union[str, int]] = None,
+        min_task_size: str | int | None = None,
+        batch_size: int | None = None,
     ) -> None: ...
 
     @abstractmethod
@@ -46,4 +49,4 @@ class AbstractUDFDistributor(ABC):
 
     @staticmethod
     @abstractmethod
-    def run_udf(fd: Optional[int] = None) -> int: ...
+    def run_udf() -> int: ...

datachain/remote/studio.py

@@ -1,16 +1,10 @@
-import base64
 import json
 import logging
 import os
 from collections.abc import AsyncIterator, Iterable, Iterator
 from datetime import datetime, timedelta, timezone
 from struct import unpack
-from typing import (
-    Any,
-    Generic,
-    Optional,
-    TypeVar,
-)
+from typing import Any, BinaryIO, Generic, TypeVar
 from urllib.parse import urlparse, urlunparse
 
 import websockets
@@ -22,16 +16,17 @@ from datachain.error import DataChainError
 from datachain.utils import STUDIO_URL, retry_with_backoff
 
 T = TypeVar("T")
-LsData = Optional[list[dict[str, Any]]]
-DatasetInfoData = Optional[dict[str, Any]]
-DatasetRowsData = Optional[Iterable[dict[str, Any]]]
-DatasetJobVersionsData = Optional[dict[str, Any]]
-DatasetExportStatus = Optional[dict[str, Any]]
-DatasetExportSignedUrls = Optional[list[str]]
-FileUploadData = Optional[dict[str, Any]]
-JobData = Optional[dict[str, Any]]
-JobListData = dict[str, Any]
-ClusterListData = dict[str, Any]
+LsData = list[dict[str, Any]] | None
+DatasetInfoData = dict[str, Any] | None
+DatasetRowsData = Iterable[dict[str, Any]] | None
+DatasetJobVersionsData = dict[str, Any] | None
+DatasetExportStatus = dict[str, Any] | None
+DatasetExportSignedUrls = list[str] | None
+FileUploadData = dict[str, Any] | None
+JobData = dict[str, Any] | None
+JobListData = list[dict[str, Any]]
+ClusterListData = list[dict[str, Any]]
+
 logger = logging.getLogger("datachain")
 
 DATASET_ROWS_CHUNK_SIZE = 8192
@@ -92,7 +87,7 @@ class Response(Generic[T]):
 
 
 class StudioClient:
-    def __init__(self, timeout: float = 3600.0, team: Optional[str] = None) -> None:
+    def __init__(self, timeout: float = 3600.0, team: str | None = None) -> None:
         self._check_dependencies()
         self.timeout = timeout
         self._config = None
@@ -153,7 +148,7 @@ class StudioClient:
             ) from None
 
     def _send_request_msgpack(
-        self, route: str, data: dict[str, Any], method: Optional[str] = "POST"
+        self, route: str, data: dict[str, Any], method: str | None = "POST"
     ) -> Response[Any]:
         import msgpack
         import requests
@@ -191,7 +186,7 @@ class StudioClient:
 
     @retry_with_backoff(retries=3, errors=(HTTPError, Timeout))
     def _send_request(
-        self, route: str, data: dict[str, Any], method: Optional[str] = "POST"
+        self, route: str, data: dict[str, Any], method: str | None = "POST"
     ) -> Response[Any]:
         """
         Function that communicate Studio API.
@@ -239,6 +234,45 @@ class StudioClient:
 
         return Response(data, ok, message, response.status_code)
 
+    def _send_multipart_request(
+        self, route: str, files: dict[str, Any], params: dict[str, Any] | None = None
+    ) -> Response[Any]:
+        """
+        Function that communicates with Studio API using multipart/form-data.
+        It will raise an exception, and try to retry, if 5xx status code is
+        returned, or if Timeout exceptions is thrown from the requests lib
+        """
+        import requests
+
+        # Add team_name to params
+        request_params = {**(params or {}), "team_name": self.team}
+
+        response = requests.post(
+            url=f"{self.url}/{route}",
+            files=files,
+            params=request_params,
+            headers={
+                "Authorization": f"token {self.token}",
+            },
+            timeout=self.timeout,
+        )
+
+        ok = response.ok
+        try:
+            data = json.loads(response.content.decode("utf-8"))
+        except json.decoder.JSONDecodeError:
+            data = {}
+
+        if not ok:
+            if response.status_code == 403:
+                message = f"Not authorized for the team {self.team}"
+            else:
+                message = data.get("message", "")
+        else:
+            message = ""
+
+        return Response(data, ok, message, response.status_code)
+
     @staticmethod
     def _unpacker_hook(code, data):
         import msgpack
@@ -305,7 +339,7 @@ class StudioClient:
             response = self._send_request_msgpack("datachain/ls", {"source": path})
             yield path, response
 
-    def ls_datasets(self, prefix: Optional[str] = None) -> Response[LsData]:
+    def ls_datasets(self, prefix: str | None = None) -> Response[LsData]:
         return self._send_request(
             "datachain/datasets", {"prefix": prefix}, method="GET"
         )
@@ -315,9 +349,9 @@ class StudioClient:
         name: str,
         namespace: str,
         project: str,
-        new_name: Optional[str] = None,
-        description: Optional[str] = None,
-        attrs: Optional[list[str]] = None,
+        new_name: str | None = None,
+        description: str | None = None,
+        attrs: list[str] | None = None,
     ) -> Response[DatasetInfoData]:
         body = {
             "new_name": new_name,
@@ -338,8 +372,8 @@ class StudioClient:
         name: str,
         namespace: str,
         project: str,
-        version: Optional[str] = None,
-        force: Optional[bool] = False,
+        version: str | None = None,
+        force: bool | None = False,
     ) -> Response[DatasetInfoData]:
         return self._send_request(
             "datachain/datasets",
@@ -409,29 +443,30 @@ class StudioClient:
             method="GET",
         )
 
-    def upload_file(self, content: bytes, file_name: str) -> Response[FileUploadData]:
-        data = {
-            "file_content": base64.b64encode(content).decode("utf-8"),
-            "file_name": file_name,
-        }
-        return self._send_request("datachain/upload-file", data)
+    def upload_file(
+        self, file_obj: BinaryIO, file_name: str
+    ) -> Response[FileUploadData]:
+        # Prepare multipart form data
+        files = {"file": (file_name, file_obj, "application/octet-stream")}
+
+        return self._send_multipart_request("datachain/jobs/files", files)
 
     def create_job(
         self,
         query: str,
         query_type: str,
-        environment: Optional[str] = None,
-        workers: Optional[int] = None,
-        query_name: Optional[str] = None,
-        files: Optional[list[str]] = None,
-        python_version: Optional[str] = None,
-        requirements: Optional[str] = None,
-        repository: Optional[str] = None,
-        priority: Optional[int] = None,
-        cluster: Optional[str] = None,
-        start_time: Optional[str] = None,
-        cron: Optional[str] = None,
-        credentials_name: Optional[str] = None,
+        environment: str | None = None,
+        workers: int | None = None,
+        query_name: str | None = None,
+        files: list[str] | None = None,
+        python_version: str | None = None,
+        requirements: str | None = None,
+        repository: str | None = None,
+        priority: int | None = None,
+        cluster: str | None = None,
+        start_time: str | None = None,
+        cron: str | None = None,
+        credentials_name: str | None = None,
     ) -> Response[JobData]:
         data = {
             "query": query,
@@ -449,25 +484,27 @@ class StudioClient:
             "cron_expression": cron,
             "credentials_name": credentials_name,
         }
-        return self._send_request("datachain/job", data)
+        return self._send_request("datachain/jobs/", data)
 
     def get_jobs(
         self,
-        status: Optional[str] = None,
+        status: str | None = None,
         limit: int = 20,
+        job_id: str | None = None,
     ) -> Response[JobListData]:
-        return self._send_request(
-            "datachain/jobs",
-            {"status": status, "limit": limit} if status else {"limit": limit},
-            method="GET",
-        )
+        params: dict[str, Any] = {"limit": limit}
+        if status is not None:
+            params["status"] = status
+        if job_id is not None:
+            params["job_id"] = job_id
+        return self._send_request("datachain/jobs/", params, method="GET")
 
     def cancel_job(
         self,
         job_id: str,
     ) -> Response[JobData]:
-        url = f"datachain/job/{job_id}/cancel"
+        url = f"datachain/jobs/{job_id}/cancel"
         return self._send_request(url, data={}, method="POST")
 
     def get_clusters(self) -> Response[ClusterListData]:
-        return self._send_request("datachain/clusters", {}, method="GET")
+        return self._send_request("datachain/clusters/", {}, method="GET")

datachain/script_meta.py

@@ -1,6 +1,6 @@
 import re
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 
 try:
     import tomllib
@@ -59,23 +59,23 @@ class ScriptConfig:
 
     """
 
-    python_version: Optional[str]
+    python_version: str | None
     dependencies: list[str]
     attachments: dict[str, str]
     params: dict[str, Any]
     inputs: dict[str, Any]
     outputs: dict[str, Any]
-    num_workers: Optional[int] = None
+    num_workers: int | None = None
 
     def __init__(
         self,
-        python_version: Optional[str] = None,
-        dependencies: Optional[list[str]] = None,
-        attachments: Optional[dict[str, str]] = None,
-        params: Optional[dict[str, Any]] = None,
-        inputs: Optional[dict[str, Any]] = None,
-        outputs: Optional[dict[str, Any]] = None,
-        num_workers: Optional[int] = None,
+        python_version: str | None = None,
+        dependencies: list[str] | None = None,
+        attachments: dict[str, str] | None = None,
+        params: dict[str, Any] | None = None,
+        inputs: dict[str, Any] | None = None,
+        outputs: dict[str, Any] | None = None,
+        num_workers: int | None = None,
     ):
         self.python_version = python_version
         self.dependencies = dependencies or []
@@ -98,7 +98,7 @@ class ScriptConfig:
         return self.attachments.get(name, default)
 
     @staticmethod
-    def read(script: str) -> Optional[dict]:
+    def read(script: str) -> dict | None:
         """Converts inline script metadata to dict with all found data"""
         regex = (
             r"(?m)^# \/\/\/ (?P<type>[a-zA-Z0-9-]+)[ \t]*$[\r\n|\r|\n]"
@@ -119,7 +119,7 @@ class ScriptConfig:
         return None
 
     @staticmethod
-    def parse(script: str) -> Optional["ScriptConfig"]:
+    def parse(script: str) -> "ScriptConfig | None":
         """
         Method that is parsing inline script metadata from datachain script and
         instantiating ScriptConfig class with found data. If no inline metadata is