deltacat 1.1.38__py3-none-any.whl → 2.0.0__py3-none-any.whl

deltacat/compute/jobs/client.py

@@ -0,0 +1,417 @@
+# from deltacat.compute import index
+import subprocess
+import socket
+import os
+import time
+import re
+
+import deltacat as dc
+
+from dataclasses import dataclass
+
+from typing import Set, Optional, Dict, Any, Union
+
+from ray.job_submission import JobSubmissionClient, JobStatus
+
+from deltacat.utils.performance import timed_invocation
+
+
+def _run_cmd(cmd: str) -> None:
+    exit_code = int(os.system(cmd))
+    assert exit_code == 0, f"`{cmd}` failed. Exit code: {exit_code}"
+
+
+def _ray_up(
+    cluster_cfg: str, cluster_name_override: str = None, restart_only: bool = False
+) -> None:
+    restart_flag = "--no-restart" if not restart_only else "--restart-only"
+    cluster_name_option = (
+        f"-n '{cluster_name_override}'" if cluster_name_override else ""
+    )
+    print(f"Starting Ray cluster from '{cluster_cfg}'")
+    _run_cmd(
+        f"ray up '{cluster_cfg}' -y --no-config-cache {restart_flag} {cluster_name_option} --disable-usage-stats"
+    )
+    print(f"Started Ray cluster from '{cluster_cfg}'")
+
+
+def _is_port_in_use(port: Union[int, str]) -> bool:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        return s.connect_ex(("localhost", int(port))) == 0
+
+
+def _is_dashboard_running(port: Union[int, str]) -> bool:
+    return _is_port_in_use(port)
+
+
+def _ray_dashboard_up(
+    cluster_cfg: str, port: Union[str, int], timeout_seconds=15
+) -> None:
+    print(f"Starting Ray Dashboard for Ray cluster '{cluster_cfg}'")
+    _run_cmd(f"ray dashboard '{cluster_cfg}' --port {port} &")
+    start = time.monotonic()
+    dashboard_is_up = False
+    while time.monotonic() - start <= timeout_seconds:
+        if _is_dashboard_running(port):
+            dashboard_is_up = True
+            break
+        time.sleep(0.1)
+    if not dashboard_is_up:
+        raise TimeoutError(
+            f"Timed out after waiting {timeout_seconds} seconds for dashboard "
+            f"to establish connection on port {port}."
+        )
+    print(f"Started Ray Dashboard for Ray cluster '{cluster_cfg}'")
+
+
+def _get_head_node_ip(cluster_cfg: str) -> str:
+    print(f"Getting Ray cluster head node IP for '{cluster_cfg}'")
+    cmd = f"ray get-head-ip '{cluster_cfg}'"
+    proc = subprocess.run(
+        cmd,
+        shell=True,
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    # the head node IP should be the last line printed to stdout
+    # TODO(pdames): add IPv6 support
+    head_node_ip = proc.stdout.splitlines()[-1]
+    if not re.match(
+        r"^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$",
+        head_node_ip,
+    ):
+        print(
+            f"Failed to find Ray Head Node IP Address in `{cmd}` "
+            f"output: {proc.stdout}"
+        )
+        raise RuntimeError("No Ray Head Node IP Address Found")
+    print(f"Ray cluster head node IP for '{cluster_cfg}': {head_node_ip}")
+    return head_node_ip
+
+
+def _ray_down_cmd(cluster_cfg: str) -> str:
+    return f"ray down '{cluster_cfg}' -y"
+
+
+def _ray_down(cluster_cfg: str) -> None:
+    print(f"Destroying Ray cluster for '{cluster_cfg}'")
+    _run_cmd(_ray_down_cmd(cluster_cfg))
+    print(f"Destroyed Ray cluster for '{cluster_cfg}'")
+
+
+def _ray_cluster_running(cluster_cfg: str) -> bool:
+    try:
+        _get_head_node_ip(cluster_cfg)
+    except Exception as e:
+        print(f"Get Head Node IP Failed with Exception: {e}")
+        print(f"Assuming Ray Cluster is Not Running")
+        return False
+    return True
+
+
+@dataclass(frozen=True)
+class DeltaCatJobRunResult:
+    job_id: str
+    job_status: JobStatus
+    job_logs: Any
+
+
+class DeltaCatJobClient(JobSubmissionClient):
+    @staticmethod
+    def of(
+        cluster_cfg_file_path: str = "./deltacat.yaml",
+        *,
+        launch_cluster: bool = True,
+        start_dashboard: bool = True,
+        restart_ray: bool = False,
+        head_node_ip: str = None,
+        dashboard_wait_time_seconds: int = 30,
+        port: Union[int, str] = "8265",
+        cluster_name_override: str = None,
+    ):
+        job_submission_client_url = None
+        try:
+            # launch Ray cluster if necessary
+            if cluster_cfg_file_path:
+                if launch_cluster:
+                    if not _ray_cluster_running(cluster_cfg_file_path) or restart_ray:
+                        _ray_up(cluster_cfg_file_path, cluster_name_override)
+                elif restart_ray:
+                    if _ray_cluster_running(cluster_cfg_file_path):
+                        _ray_up(
+                            cluster_cfg_file_path, restart_ray, cluster_name_override
+                        )
+                    else:
+                        raise RuntimeError(
+                            f"Cannot Restart Ray: Ray Cluster for "
+                            f"`{cluster_cfg_file_path}` not found."
+                        )
+                dashboard_running = _is_dashboard_running(port)
+                if not dashboard_running and start_dashboard:
+                    _ray_dashboard_up(
+                        cluster_cfg=cluster_cfg_file_path,
+                        port=port,
+                        timeout_seconds=dashboard_wait_time_seconds,
+                    )
+                    dashboard_running = True
+                if not head_node_ip:
+                    head_node_ip = (
+                        "127.0.0.1"
+                        # use dashboard port forwarding on localhost
+                        if dashboard_running
+                        # fetch the remote head node IP
+                        else _get_head_node_ip(cluster_cfg_file_path)
+                    )
+            else:
+                head_node_ip = "127.0.0.1"
+            job_submission_client_url = f"http://{head_node_ip}:{port}"
+            print(
+                f"Initializing Ray Job Submission Client with URL: "
+                f"{job_submission_client_url}"
+            )
+            client = JobSubmissionClient(f"http://{head_node_ip}:{port}")
+            # the below class change is safe as long as we only add new methods
+            # to the wrapped JobSubmissionClient that don't alter its internal
+            # state
+            client.__class__ = DeltaCatJobClient
+            return client
+        except Exception as e:
+            print(f"Unexpected error while initializing Ray Job Client: {e}")
+            if job_submission_client_url:
+                print(
+                    f"Please ensure that Ray was installed with a job server "
+                    f'enabled via `pip install -U "ray[default]"` and '
+                    f"that http://{head_node_ip}:{port} is accessible. You "
+                    f"can optionally run `ray dashboard` to forward the "
+                    f"remote Ray head node port to a local port (default 8265) "
+                    f'then run `ray_job_client("127.0.0.1", 8265)` '
+                    f"to connect via localhost."
+                )
+            if cluster_cfg_file_path:
+                print(
+                    f"If you're done submitting jobs, ensure that the remote "
+                    f"Ray Cluster is shut down by running: "
+                    f"{_ray_down_cmd(cluster_cfg_file_path)}"
+                )
+            raise e
+
+    def run_job(
+        self,
+        *,
+        entrypoint: str,
+        runtime_env: Optional[Dict[str, Any]] = None,
+        timeout_seconds: int = 600,
+        **kwargs,
+    ) -> DeltaCatJobRunResult:
+        """
+        Synchronously submit and run a Ray job. This method combines Ray job submission and monitoring by submitting
+        the job to the Ray Job Server, waiting for the job to complete,
+        validating the job's terminal status, retrieving and returning job run
+        result information if successful.
+
+        Args:
+            entrypoint: The entry point for the job to be executed (module
+                or script to run)
+            runtime_env: Runtime environment configuration for the job.
+                Some commonly used keys include `working_dir` (directory
+                containing the job code), `pip` (list of pip packages to
+                install), and `env_vars` (environment variables for the job).
+            timeout_seconds: Maximum time in seconds to wait for job completion.
+                Default to 600 seconds (10 minutes).
+            kwargs: Additional keyword arguments to pass to the job submission.
+
+        Returns:
+            Final results from the successful job run execution.
+
+        Raises:
+            RuntimeError: If the job fails or terminates with status other
+                than SUCCEEDED.
+            TimeoutError: If the job doesn't complete within the specified
+                timeout period
+
+        Example:
+            >>> client = job_client()
+            >>> logs = client.run_job(
+            ...     # Shell command to run job
+            ...     entrypoint="my_script.py",
+            ...     runtime_env={
+            ...         # Path to the local directory containing my_script.py
+            ...         "working_dir": "./",
+            ...         # Pip dependencies to install
+            ...         "pip": ["pandas", "numpy"],
+            ...         # System environment variables to set
+            ...         "env_vars": {"DATA_PATH": "/path/to/data"},
+            ...     },
+            ...     timeout_seconds=1200
+            ... )
+        """
+
+        job_id = self.submit_job(
+            entrypoint=entrypoint,
+            runtime_env=runtime_env,
+            **kwargs,
+        )
+        job_status, latency = timed_invocation(
+            self.await_job,
+            job_id,
+            timeout_seconds=timeout_seconds,
+        )
+        job_logs = self.get_job_logs(job_id)
+        if job_status != JobStatus.SUCCEEDED:
+            print(f"Job `{job_id}` logs: ")
+            print(job_logs)
+            raise RuntimeError(f"Job `{job_id}` terminated with status: {job_status}")
+        return DeltaCatJobRunResult(
+            job_id=job_id,
+            job_status=job_status,
+            job_logs=job_logs,
+        )
+
+    def await_job(
+        self,
+        job_id: str,
+        await_status: Set[JobStatus] = {
+            JobStatus.SUCCEEDED,
+            JobStatus.STOPPED,
+            JobStatus.FAILED,
+        },
+        *,
+        timeout_seconds: int = 600,
+    ) -> JobStatus:
+        """
+        Polls a job's status until it matches the desired status or times out.
+
+        This function continuously checks the status of a specified job using the
+        provided client. It will keep polling until either the desired status is
+        reached or the timeout period expires.
+
+        Args:
+            job_id: The unique identifier of the job to monitor.
+            await_status: Set of :class:`ray.job_submission.JobStatus` to wait for.
+                The function will return when the job reaches any of these states.
+            timeout_seconds: Maximum time to wait in seconds.
+                Defaults to 600 seconds (10 minutes).
+
+        Returns:
+            The final status of the job.
+
+        Raises:
+            TimeoutError: If the desired status is not reached within the
+            specified timeout period.
+
+        Example:
+            >>>
+            >>> client = job_client()
+            >>> job_id = client.submit_job(
+            >>>     # Shell command to run job
+            >>>     entrypoint=f"python copy.py --source '{source}' --dest '{dest}'",
+            >>>     # Path to the local directory containing copy.py
+            >>>     runtime_env={"working_dir": "./"},
+            >>> )
+            >>> # wait for the job to reach a terminal state
+            >>> client.await_job(job_id)
+        """
+        start = time.monotonic()
+        terminal_status = None
+        while time.monotonic() - start <= timeout_seconds:
+            status = self.get_job_status(job_id)
+            if status in await_status:
+                terminal_status = status
+                break
+            time.sleep(0.1)
+        if not terminal_status:
+            self.stop_job(job_id)
+            raise TimeoutError(
+                f"Timed out after waiting {timeout_seconds} seconds for job "
+                f"`{job_id}` status: {status}"
+            )
+        return terminal_status
+
+
+def local_job_client(*args, **kwargs) -> DeltaCatJobClient:
+    """
+    Create a DeltaCAT Job Client that can be used to submit jobs to a local Ray
+    cluster. Initializes Ray if it's not already running.
+
+    Args:
+        *args: Positional arguments to pass to `deltacat.init()`.
+        **kwargs: Keyword arguments to pass to `deltacat.init()`.
+    Returns:
+        DeltaCatJobClient: A client instance that can be used to submit and
+            manage local Ray jobs.
+
+    Raises:
+        RuntimeError: If a local Ray Job Server cannot be found.
+    """
+    # force reinitialization to ensure that we can get the Ray context
+    kwargs["force"] = True
+    context = dc.init(*args, **kwargs)
+    if context is None:
+        raise RuntimeError("Failed to retrieve Ray context.")
+    if context.dashboard_url:
+        head_node_ip, port = context.dashboard_url.split(":")
+    else:
+        # the Ray Dashboard URL is also the Ray Job Server URL
+        raise RuntimeError(
+            "Ray Job Server not found! Please reinstall Ray using "
+            "`pip install -U `ray[default]`"
+        )
+    return DeltaCatJobClient.of(
+        None,
+        launch_cluster=False,
+        start_dashboard=False,
+        head_node_ip=head_node_ip,
+        port=port,
+    )
+
+
+def job_client(
+    cluster_cfg_file_path: str = "./deltacat.yaml",
+    *,
+    launch_cluster: bool = True,
+    start_dashboard: bool = True,
+    restart_ray: bool = False,
+    head_node_ip: str = None,
+    dashboard_wait_time_seconds: int = 15,
+    port: Union[str, int] = "8265",
+    cluster_name_override: str = None,
+) -> DeltaCatJobClient:
+    """
+    Create a DeltaCAT Job Client that can be used to submit jobs to a remote
+    Ray cluster.
+
+    Args:
+        cluster_cfg_file_path: Path to the Ray Cluster Launcher
+            Config file. Defaults to "./deltacat.yaml".
+        launch_cluster : Whether to launch a new Ray cluster.
+            Defaults to True.
+        start_dashboard: Whether to start the Ray dashboard.
+            Defaults to True.
+        restart_ray: Whether to restart Ray if it's already
+            running. Defaults to False.
+        head_node_ip: IP address of the Ray cluster head node.
+            If None, will use the configuration from the cluster config file.
+            Defaults to None.
+        dashboard_wait_time_seconds: Time in seconds to wait for the Ray
+            dashboard to start if `start_dashboard` is True.
+        port: Port number for the Ray
+            dashboard/job server. Defaults to "8265".
+
+    Returns:
+        DeltaCatJobClient: A client instance that can be used to submit and
+            manage jobs on the Ray cluster.
+
+    Raises:
+        RuntimeError: If the Ray Job Server is not found.
+    """
+    return DeltaCatJobClient.of(
+        cluster_cfg_file_path,
+        launch_cluster=launch_cluster,
+        start_dashboard=start_dashboard,
+        restart_ray=restart_ray,
+        head_node_ip=head_node_ip,
+        dashboard_wait_time_seconds=dashboard_wait_time_seconds,
+        port=port,
+        cluster_name_override=cluster_name_override,
+    )

deltacat/compute/resource_estimation/delta.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, List
 from deltacat import logs
 from deltacat.storage import (
     Delta,
@@ -61,6 +61,7 @@ def _estimate_resources_required_to_process_delta_using_previous_inflation(
 def _estimate_resources_required_to_process_delta_using_type_params(
     delta: Delta,
     operation_type: OperationType,
+    all_column_names: List[str],
     estimate_resources_params: EstimateResourcesParams,
     deltacat_storage: unimplemented_deltacat_storage,
     deltacat_storage_kwargs: Dict[str, Any],
@@ -113,6 +114,7 @@ def _estimate_resources_required_to_process_delta_using_type_params(
     """
     appended = append_content_type_params(
         delta=delta,
+        all_column_names=all_column_names,
         deltacat_storage=deltacat_storage,
         deltacat_storage_kwargs=deltacat_storage_kwargs,
         file_reader_kwargs_provider=file_reader_kwargs_provider,
@@ -289,10 +291,18 @@ def estimate_resources_required_to_process_delta(
         estimate_resources_params.resource_estimation_method
     )
 
+    all_column_names = deltacat_storage.get_table_version_column_names(
+        delta.locator.namespace,
+        delta.locator.table_name,
+        delta.locator.table_version,
+        **deltacat_storage_kwargs,
+    )
+
     for func in functions:
         resources = func(
             delta=delta,
             operation_type=operation_type,
+            all_column_names=all_column_names,
             estimate_resources_params=estimate_resources_params,
             deltacat_storage=deltacat_storage,
             deltacat_storage_kwargs=deltacat_storage_kwargs,

deltacat/constants.py

@@ -1,4 +1,10 @@
-from deltacat.utils.common import env_string
+from __future__ import annotations
+
+import botocore.exceptions
+
+from daft.exceptions import DaftTransientError
+from deltacat.utils.common import env_string, env_bool
+from deltacat.utils.common import env_integer
 
 # Environment variables
 DELTACAT_SYS_LOG_LEVEL = env_string("DELTACAT_SYS_LOG_LEVEL", "DEBUG")
@@ -30,6 +36,25 @@ DELTACAT_APP_DEBUG_LOG_BASE_FILE_NAME = env_string(
 )
 # A json context which will be logged along with other context args.
 DELTACAT_LOGGER_CONTEXT = env_string("DELTACAT_LOGGER_CONTEXT", None)
+DELTACAT_LOGGER_USE_SINGLE_HANDLER = env_bool(
+    "DELTACAT_LOGGER_USE_SINGLE_HANDLER",
+    False,
+)
+DELTACAT_ROOT = env_string(
+    "DELTACAT_ROOT",
+    "",
+)
+
+# CLI Args
+METAFILE_FORMAT_JSON = "json"
+METAFILE_FORMAT_MSGPACK = "msgpack"
+METAFILE_FORMAT = env_string("METAFILE_FORMAT", METAFILE_FORMAT_MSGPACK)
+SUPPORTED_METAFILE_FORMATS = [METAFILE_FORMAT_JSON, METAFILE_FORMAT_MSGPACK]
+METAFILE_EXT = {
+    "json": ".json",
+    "msgpack": ".mpk",
+}[METAFILE_FORMAT]
+
 
 # Byte Units
 BYTES_PER_KIBIBYTE = 2**10
@@ -41,6 +66,11 @@ BYTES_PER_PEBIBYTE = 2**50
 SIGNED_INT64_MIN_VALUE = -(2**63)
 SIGNED_INT64_MAX_VALUE = 2**63 - 1
 
+# Time Units
+NANOS_PER_SEC = 1_000_000_000
+MICROS_PER_SEC = 1_000_000
+MILLIS_PER_SEC = 1000
+
 # Inflation multiplier from snappy-compressed parquet to pyarrow.
 # This should be kept larger than actual average inflation multipliers.
 # Note that this is a very rough guess since actual observed pyarrow
@@ -58,3 +88,62 @@ MEMORY_TO_HASH_BUCKET_COUNT_RATIO = 0.0512 * BYTES_PER_TEBIBYTE
 
 # The number of bytes allocated to null values in string physical type in parquet
 NULL_SIZE_BYTES = 4
+
+# Metastore Constants
+REVISION_DIR_NAME: str = "rev"
+TXN_DIR_NAME: str = "txn"
+RUNNING_TXN_DIR_NAME: str = "running"
+FAILED_TXN_DIR_NAME: str = "failed"
+PAUSED_TXN_DIR_NAME: str = "paused"
+SUCCESS_TXN_DIR_NAME: str = "success"
+DATA_FILE_DIR_NAME: str = "data"
+REV_DIR_NAME: str = "rev"
+TXN_PART_SEPARATOR = "_"
+
+# Storage interface defaults
+# These defaults should be applied in catalog interface implementations
+# Storage interface implementations should be agnostic to defaults and require full information
+DEFAULT_CATALOG = "default"
+DEFAULT_NAMESPACE = "default"
+DEFAULT_TABLE_VERSION = "1"
+DEFAULT_STREAM_ID = "stream"
+DEFAULT_PARTITION_ID = "partition"
+DEFAULT_PARTITION_VALUES = ["default"]
+
+# Transaction Status constants
+SUCCESSFULLY_CLEANED = "cleaned"
+CURRENTLY_CLEANING = "cleaning"
+TIMEOUT_TXN = "timedout"
+
+# operation timeout constants
+OPERATION_TIMEOUTS = {
+    "create": 5,
+    "update": 3,
+    "delete": 4,
+    "read_siblings": 2,
+    "read_children": 2,
+    "read_latest": 3,
+    "read_exists": 1,
+}
+# Upload/Download Retry Defaults
+UPLOAD_DOWNLOAD_RETRY_STOP_AFTER_DELAY = env_integer(
+    "UPLOAD_DOWNLOAD_RETRY_STOP_AFTER_DELAY", 10 * 60
+)
+UPLOAD_SLICED_TABLE_RETRY_STOP_AFTER_DELAY = env_integer(
+    "UPLOAD_SLICED_TABLE_RETRY_STOP_AFTER_DELAY", 30 * 60
+)
+DOWNLOAD_MANIFEST_ENTRY_RETRY_STOP_AFTER_DELAY = env_integer(
+    "DOWNLOAD_MANIFEST_ENTRY_RETRY_STOP_AFTER_DELAY", 30 * 60
+)
+DEFAULT_FILE_READ_TIMEOUT_MS = env_integer(
+    "DEFAULT_FILE_READ_TIMEOUT_MS", 300_000
+)  # 5 mins
+RETRYABLE_TRANSIENT_ERRORS = (
+    OSError,
+    botocore.exceptions.ConnectionError,
+    botocore.exceptions.HTTPClientError,
+    botocore.exceptions.NoCredentialsError,
+    botocore.exceptions.ConnectTimeoutError,
+    botocore.exceptions.ReadTimeoutError,
+    DaftTransientError,
+)