deltacat 1.1.5__py3-none-any.whl → 1.1.6__py3-none-any.whl

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/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/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/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/tests/aws/test_s3u.py

@@ -2,6 +2,21 @@ 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()
@@ -10,3 +25,100 @@ class TestUuidBlockWritePathProvider(unittest.TestCase):
         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.dist-info → deltacat-1.1.6.dist-info}/METADATA

@@ -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.dist-info → deltacat-1.1.6.dist-info}/RECORD

@@ -1,11 +1,11 @@
-deltacat/__init__.py,sha256=i3bpmX1O0ZJD9XFiuqF0RHURR1lLKvdTBzVOP90fKQA,1777
+deltacat/__init__.py,sha256=01om7qgj3agAF60Q1qwZXAzsUtP7cabwc_1RXqRr0vw,1777
 deltacat/constants.py,sha256=_6oRI-3yp5c8J1qKGQZrt89I9-ttT_gSSvVsJ0h8Duc,1939
 deltacat/exceptions.py,sha256=xqZf8CwysNYP2d39pf27OnXGStPREgBgIM-e2Tts-TI,199
 deltacat/logs.py,sha256=6g16VkEFidbaMjgenAjggE1r2l664drMVhreRs8B1IQ,8438
 deltacat/aws/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 deltacat/aws/clients.py,sha256=VgddlV3AEjlBGIFmhhHxokYzwJ-lXnmHAeprVyADduI,6948
-deltacat/aws/constants.py,sha256=aAhOKeLVgtpekA3h9otHUrHqY2bLDWs2QlL7GrdI63g,352
-deltacat/aws/s3u.py,sha256=lgoE6es6N4xfzwyydxmVspROP1hrNfanB6JqjyBoRb4,24859
+deltacat/aws/constants.py,sha256=OnRbtfFdu4buJEsl39Kg5cH-7A-dEL_ESeBSAlR_1Cs,501
+deltacat/aws/s3u.py,sha256=qZL5Omz1onW79vB_KrPHQ2Mox4sNPrLXkNxpFl9HFHM,26525
 deltacat/aws/redshift/__init__.py,sha256=7SvjG-dqox8zZUhFicTsUvpG5vXYDl_QQ3ohlHOgTKc,342
 deltacat/aws/redshift/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 deltacat/aws/redshift/model/manifest.py,sha256=ThgpdwzaWz493Zz9e8HSWwuxEheA1nDuypM3pe4vozk,12987
@@ -74,7 +74,7 @@ deltacat/compute/compactor_v2/utils/delta.py,sha256=8hjkDeIIkSX-gAQ2utQSp2sZcO2t
 deltacat/compute/compactor_v2/utils/io.py,sha256=autXlE3uHICdCCuJoS7mfdeJbRRiz2_xlz-3izlccB4,5264
 deltacat/compute/compactor_v2/utils/merge.py,sha256=hK4Y7acrtgfvWWTz-fAGznEg6qn6dBYu8blQUQVHhs0,5244
 deltacat/compute/compactor_v2/utils/primary_key_index.py,sha256=MAscmL35WfwN7Is72aFlD_cGhxtZgjRwwR5kS9Yn2uU,11393
-deltacat/compute/compactor_v2/utils/task_options.py,sha256=WK-1-1gncUSZI9yxBEG1AokDZKTBBhisYcYGijlitHY,13857
+deltacat/compute/compactor_v2/utils/task_options.py,sha256=n1zKOFmAg2cL7CDpT9y9h-J0aYzTMtOdUjkDm1svo9k,14160
 deltacat/compute/merge_on_read/__init__.py,sha256=ckbgngmqPjYBYz_NySsR1vNTOb_hNpeL1sYkZKvBI9M,214
 deltacat/compute/merge_on_read/daft.py,sha256=1oC38u5ig_aTrq7EzyWBo8Ui54rb6yERYMk-vEFbpxM,1400
 deltacat/compute/merge_on_read/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -136,7 +136,7 @@ deltacat/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 deltacat/tests/test_logs.py,sha256=6BEMw8VApFg2msFwCAVosz8NWJYATtX5furHyz8UluM,3828
 deltacat/tests/aws/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 deltacat/tests/aws/test_clients.py,sha256=23GMWfz27WWBDXSqphG9mfputsyS7j3I5P_HRk4YoKE,3790
-deltacat/tests/aws/test_s3u.py,sha256=QflXbR94o7WobGBm6jhQDK5lJJD2Pd9z2uvi4J7WEJg,437
+deltacat/tests/aws/test_s3u.py,sha256=aFvUa9f63hFU8T4r_cuKYxcFg6jVUoJWygiPwDUd09s,4654
 deltacat/tests/catalog/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 deltacat/tests/catalog/test_default_catalog_impl.py,sha256=9srCU5yQ159oZ9_PoJ_mWMzVUW5bKV0mnmPJc5zKCQQ,3125
 deltacat/tests/compute/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -206,8 +206,8 @@ deltacat/utils/ray_utils/concurrency.py,sha256=JDVwMiQWrmuSlyCWAoiq9ctoJ0XADEfDD
 deltacat/utils/ray_utils/dataset.py,sha256=SIljK3UkSqQ6Ntit_iSiYt9yYjN_gGrCTX6_72XdQ3w,3244
 deltacat/utils/ray_utils/performance.py,sha256=d7JFM7vTXHzkGx9qNQcZzUWajnqINvYRwaM088_FpsE,464
 deltacat/utils/ray_utils/runtime.py,sha256=5eaBWTDm0IXVoc5Y6aacoVB-f0Mnv-K2ewyTSjHKHwM,5009
-deltacat-1.1.5.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-deltacat-1.1.5.dist-info/METADATA,sha256=ZXriHPvPca0lz3LrEfynDq-u4dyxhGvhKaM3m9veiT0,1780
-deltacat-1.1.5.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-deltacat-1.1.5.dist-info/top_level.txt,sha256=RWdIcid4Bv2i2ozLVh-70kJpyB61xEKXod9XXGpiono,9
-deltacat-1.1.5.dist-info/RECORD,,
+deltacat-1.1.6.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+deltacat-1.1.6.dist-info/METADATA,sha256=C5eD7a_S7Zxm5W6A5dBUGPKKBnwttmcu2qHELs6YImw,1780
+deltacat-1.1.6.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+deltacat-1.1.6.dist-info/top_level.txt,sha256=RWdIcid4Bv2i2ozLVh-70kJpyB61xEKXod9XXGpiono,9
+deltacat-1.1.6.dist-info/RECORD,,