deltacat 1.1.5__tar.gz → 1.1.6__tar.gz

{deltacat-1.1.5 → deltacat-1.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: deltacat
-Version: 1.1.5
+Version: 1.1.6
 Summary: A scalable, fast, ACID-compliant Data Catalog powered by Ray.
 Home-page: https://github.com/ray-project/deltacat
 Author: Ray Team

{deltacat-1.1.5 → deltacat-1.1.6}/deltacat/__init__.py

@@ -44,7 +44,7 @@ from deltacat.types.tables import TableWriteMode
 
 deltacat.logs.configure_deltacat_logger(logging.getLogger(__name__))
 
-__version__ = "1.1.5"
+__version__ = "1.1.6"
 
 
 __all__ = [

{deltacat-1.1.5 → deltacat-1.1.6}/deltacat/aws/constants.py

@@ -1,8 +1,10 @@
-from typing import List
+from typing import List, Set
 
 from deltacat.utils.common import env_integer, env_string
 
 DAFT_MAX_S3_CONNECTIONS_PER_FILE = env_integer("DAFT_MAX_S3_CONNECTIONS_PER_FILE", 8)
 BOTO_MAX_RETRIES = env_integer("BOTO_MAX_RETRIES", 5)
 TIMEOUT_ERROR_CODES: List[str] = ["ReadTimeoutError", "ConnectTimeoutError"]
+RETRYABLE_PUT_OBJECT_ERROR_CODES: Set[str] = {"Throttling", "SlowDown"}
 AWS_REGION = env_string("AWS_REGION", "us-east-1")
+RETRY_STOP_AFTER_DELAY = env_integer("RETRY_STOP_AFTER_DELAY", 10 * 60)

{deltacat-1.1.5 → deltacat-1.1.6}/deltacat/aws/s3u.py

@@ -4,14 +4,18 @@ from functools import partial
 from typing import Any, Callable, Dict, Generator, List, Optional, Union
 from uuid import uuid4
 from botocore.config import Config
-from deltacat.aws.constants import BOTO_MAX_RETRIES
+from deltacat.aws.constants import (
+    BOTO_MAX_RETRIES,
+    RETRY_STOP_AFTER_DELAY,
+    RETRYABLE_PUT_OBJECT_ERROR_CODES,
+)
 
 import pyarrow as pa
 import ray
 import s3fs
 from boto3.resources.base import ServiceResource
 from botocore.client import BaseClient
-from botocore.exceptions import ClientError
+from botocore.exceptions import ClientError, NoCredentialsError
 from ray.data.block import Block, BlockAccessor, BlockMetadata
 from ray.data.datasource import BlockWritePathProvider
 from ray.types import ObjectRef
@@ -315,7 +319,6 @@ def upload_sliced_table(
                 **s3_client_kwargs,
             )
             manifest_entries.extend(slice_entries)
-
     return manifest_entries
 
 
@@ -504,41 +507,85 @@ def download_manifest_entries_distributed(
 
 def upload(s3_url: str, body, **s3_client_kwargs) -> Dict[str, Any]:
 
-    # TODO (pdames): add tenacity retrying
     parsed_s3_url = parse_s3_url(s3_url)
     s3 = s3_client_cache(None, **s3_client_kwargs)
-    return s3.put_object(
-        Body=body,
-        Bucket=parsed_s3_url.bucket,
-        Key=parsed_s3_url.key,
+    retrying = Retrying(
+        wait=wait_random_exponential(multiplier=1, max=15),
+        stop=stop_after_delay(RETRY_STOP_AFTER_DELAY),
+        retry=retry_if_exception_type(RetryableError),
+    )
+    return retrying(
+        _put_object,
+        s3,
+        body,
+        parsed_s3_url.bucket,
+        parsed_s3_url.key,
     )
 
 
+def _put_object(
+    s3_client, body: Any, bucket: str, key: str, **s3_put_object_kwargs
+) -> Dict[str, Any]:
+    try:
+        return s3_client.put_object(
+            Body=body, Bucket=bucket, Key=key, **s3_put_object_kwargs
+        )
+    except ClientError as e:
+        if e.response["Error"]["Code"] in RETRYABLE_PUT_OBJECT_ERROR_CODES:
+            raise RetryableError(
+                f"Retry upload for: {bucket}/{key} after receiving {e.response['Error']['Code']}"
+            ) from e
+        raise NonRetryableError(f"Failed table upload to: {bucket}/{key}") from e
+    except NoCredentialsError as e:
+        raise RetryableError(
+            f"Failed to fetch credentials when putting object into: {bucket}/{key}"
+        ) from e
+    except BaseException as e:
+        logger.error(
+            f"Upload has failed for {bucket}/{key}. Error: {e}",
+            exc_info=True,
+        )
+        raise NonRetryableError(f"Failed table upload to: {bucket}/{key}") from e
+
+
 def download(
     s3_url: str, fail_if_not_found: bool = True, **s3_client_kwargs
 ) -> Optional[Dict[str, Any]]:
 
-    # TODO (pdames): add tenacity retrying
     parsed_s3_url = parse_s3_url(s3_url)
     s3 = s3_client_cache(None, **s3_client_kwargs)
+    retrying = Retrying(
+        wait=wait_random_exponential(multiplier=1, max=15),
+        stop=stop_after_delay(RETRY_STOP_AFTER_DELAY),
+        retry=retry_if_exception_type(RetryableError),
+    )
+    return retrying(
+        _get_object,
+        s3,
+        parsed_s3_url.bucket,
+        parsed_s3_url.key,
+        fail_if_not_found=fail_if_not_found,
+    )
+
+
+def _get_object(s3_client, bucket: str, key: str, fail_if_not_found: bool = True):
     try:
-        return s3.get_object(
-            Bucket=parsed_s3_url.bucket,
-            Key=parsed_s3_url.key,
+        return s3_client.get_object(
+            Bucket=bucket,
+            Key=key,
         )
     except ClientError as e:
-        if fail_if_not_found:
-            raise
-        else:
-            if e.response["Error"]["Code"] != "404":
-                if e.response["Error"]["Code"] != "NoSuchKey":
-                    raise
-            logger.info(f"file not found: {s3_url}")
-    except s3.exceptions.NoSuchKey:
-        if fail_if_not_found:
-            raise
-        else:
-            logger.info(f"file not found: {s3_url}")
+        if e.response["Error"]["Code"] == "NoSuchKey":
+            if fail_if_not_found:
+                raise NonRetryableError(
+                    f"Failed get object from: {bucket}/{key}"
+                ) from e
+            logger.info(f"file not found: {bucket}/{key}")
+    except NoCredentialsError as e:
+        raise RetryableError(
+            f"Failed to fetch credentials when getting object from: {bucket}/{key}"
+        ) from e
+
     return None
 
 

{deltacat-1.1.5 → deltacat-1.1.6}/deltacat/compute/compactor_v2/utils/task_options.py

@@ -20,7 +20,6 @@ from deltacat.compute.compactor_v2.utils.primary_key_index import (
 from deltacat.compute.compactor_v2.constants import (
     PARQUET_TO_PYARROW_INFLATION,
 )
-
 from daft.exceptions import DaftTransientError
 
 
@@ -65,7 +64,11 @@ def get_task_options(
     cpu: float, memory: float, ray_custom_resources: Optional[Dict] = None
 ) -> Dict:
 
-    task_opts = {"num_cpus": cpu, "memory": memory}
+    # NOTE: With DEFAULT scheduling strategy in Ray 2.20.0, autoscaler does
+    # not spin up enough nodes fast and hence we see only approximately
+    # 20 tasks get scheduled out of 100 tasks in queue.
+    # https://docs.ray.io/en/latest/ray-core/scheduling/index.html
+    task_opts = {"num_cpus": cpu, "memory": memory, "scheduling_strategy": "SPREAD"}
 
     if ray_custom_resources:
         task_opts["resources"] = ray_custom_resources

deltacat-1.1.6/deltacat/tests/aws/test_s3u.py

@@ -0,0 +1,124 @@
+import unittest
+from deltacat.aws.s3u import UuidBlockWritePathProvider, CapturedBlockWritePaths
+
+
+import os
+from unittest import mock
+from unittest.mock import patch
+
+import boto3
+import pytest
+from boto3.resources.base import ServiceResource
+from botocore.exceptions import ClientError, NoCredentialsError
+from deltacat.exceptions import NonRetryableError
+from moto import mock_s3
+from tenacity import RetryError
+
+from deltacat.aws import s3u
+
+
+class TestUuidBlockWritePathProvider(unittest.TestCase):
+    def test_uuid_block_write_provider_sanity(self):
+        capture_object = CapturedBlockWritePaths()
+        provider = UuidBlockWritePathProvider(capture_object=capture_object)
+
+        result = provider("base_path")
+
+        self.assertRegex(result, r"^base_path/[\w-]{36}$")
+
+
+class TestDownloadUpload(unittest.TestCase):
+    TEST_S3_BUCKET_NAME = "TEST_S3_BUCKET"
+    TEST_S3_KEY = "TEST_S3_KEY"
+
+    @pytest.fixture(autouse=True)
+    def mock_aws_credential(self):
+        os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+        os.environ["AWS_SECRET_ACCESS_ID"] = "testing"
+        os.environ["AWS_SECURITY_TOKEN"] = "testing"
+        os.environ["AWS_SESSION_TOKEN"] = "testing"
+        os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
+        yield
+
+    @pytest.fixture(autouse=True)
+    def setup_s3_resource(self):
+        with mock_s3():
+            yield boto3.resource("s3")
+
+    @pytest.fixture(autouse=True)
+    def setup_test_s3_bucket(self, setup_s3_resource: ServiceResource):
+        setup_s3_resource.create_bucket(
+            ACL="authenticated-read",
+            Bucket=self.TEST_S3_BUCKET_NAME,
+        )
+        yield
+
+    def test_sanity(self):
+        uri = f"s3://{self.TEST_S3_BUCKET_NAME}/{self.TEST_S3_KEY}"
+        body = "test-body"
+        uploaded_file = s3u.upload(uri, body)
+        assert uploaded_file is not None
+        assert uploaded_file["ResponseMetadata"]["HTTPStatusCode"] == 200
+        downloaded_file = s3u.download(uri)
+        downloaded_body = downloaded_file["Body"].read().decode("utf-8")
+        assert downloaded_file["ResponseMetadata"]["HTTPStatusCode"] == 200
+        assert downloaded_body == body
+
+    @patch("deltacat.aws.s3u.RETRY_STOP_AFTER_DELAY", 1)
+    @patch("deltacat.aws.s3u.s3_client_cache")
+    def test_upload_throttled(self, mock_s3_client_cache):
+        uri = f"s3://{self.TEST_S3_BUCKET_NAME}/{self.TEST_S3_KEY}"
+        body = "test-body"
+        throttling_err = ClientError({"Error": {"Code": "Throttling"}}, "put_object")
+        mock_s3_client_cache.return_value = mock_s3 = mock.MagicMock()
+        mock_s3.put_object.side_effect = throttling_err
+        with pytest.raises(RetryError):
+            s3u.upload(uri, body)
+
+        slowdown_err = ClientError({"Error": {"Code": "SlowDown"}}, "put_object")
+        mock_s3.put_object.side_effect = slowdown_err
+        with pytest.raises(RetryError):
+            s3u.upload(uri, body)
+
+        no_credentials_err = NoCredentialsError()
+        mock_s3.put_object.side_effect = no_credentials_err
+        with pytest.raises(RetryError):
+            s3u.upload(uri, body)
+
+        assert mock_s3.put_object.call_count > 3
+
+    @patch("deltacat.aws.s3u.s3_client_cache")
+    def test_upload_unexpected_error_code(self, mock_s3_client_cache):
+        uri = f"s3://{self.TEST_S3_BUCKET_NAME}/{self.TEST_S3_KEY}"
+        body = "test-body"
+        err = ClientError({"Error": {"Code": "UnexpectedError"}}, "put_object")
+        mock_s3_client_cache.return_value = mock_s3 = mock.MagicMock()
+        mock_s3.put_object.side_effect = err
+        file = None
+        with pytest.raises(NonRetryableError):
+            s3u.upload(uri, body)
+        assert file is None
+        assert mock_s3.put_object.call_count == 1
+
+    @patch("deltacat.aws.s3u.RETRY_STOP_AFTER_DELAY", 1)
+    @patch("deltacat.aws.s3u.s3_client_cache")
+    def test_download_throttled(self, mock_s3_client_cache):
+        uri = f"s3://{self.TEST_S3_BUCKET_NAME}/{self.TEST_S3_KEY}"
+        no_credentials_err = NoCredentialsError()
+        mock_s3_client_cache.return_value = mock_s3 = mock.MagicMock()
+        mock_s3.get_object.side_effect = no_credentials_err
+        file = None
+        with pytest.raises(RetryError):
+            file = s3u.download(uri)
+        assert file is None
+        assert mock_s3.get_object.call_count > 1
+
+    def test_download_not_exists(self):
+        uri = f"s3://{self.TEST_S3_BUCKET_NAME}/key-not-exists"
+        file = None
+        with pytest.raises(NonRetryableError):
+            file = s3u.download(uri)
+        assert file is None
+
+        file = s3u.download(uri, fail_if_not_found=False)
+        assert file is None

{deltacat-1.1.5 → deltacat-1.1.6}/deltacat.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: deltacat
-Version: 1.1.5
+Version: 1.1.6
 Summary: A scalable, fast, ACID-compliant Data Catalog powered by Ray.
 Home-page: https://github.com/ray-project/deltacat
 Author: Ray Team

deltacat-1.1.5/deltacat/tests/aws/test_s3u.py

@@ -1,12 +0,0 @@
-import unittest
-from deltacat.aws.s3u import UuidBlockWritePathProvider, CapturedBlockWritePaths
-
-
-class TestUuidBlockWritePathProvider(unittest.TestCase):
-    def test_uuid_block_write_provider_sanity(self):
-        capture_object = CapturedBlockWritePaths()
-        provider = UuidBlockWritePathProvider(capture_object=capture_object)
-
-        result = provider("base_path")
-
-        self.assertRegex(result, r"^base_path/[\w-]{36}$")