Author: momento-github-actions-machine-user

src/momento/config/configurations.py

@@ -35,7 +35,7 @@ def v1() -> Configurations.Laptop:
             This configuration is guaranteed not to change in future releases of the Momento Python SDK.
             """
             return Configurations.Laptop(
-                StaticTransportStrategy(StaticGrpcConfiguration(timedelta(seconds=15))),
+                StaticTransportStrategy(StaticGrpcConfiguration(deadline=timedelta(seconds=15))),
                 FixedCountRetryStrategy(max_attempts=3),
             )
 
@@ -49,7 +49,14 @@ def latest() -> Configurations.Lambda:
             This configuration will be updated every time there is a new version of the Lambda configuration.
             """
             return Configurations.Lambda(
-                StaticTransportStrategy(StaticGrpcConfiguration(timedelta(milliseconds=1100))),
+                StaticTransportStrategy(
+                    StaticGrpcConfiguration(
+                        deadline=timedelta(milliseconds=1100),
+                        keepalive_permit_without_calls=False,
+                        keepalive_time=None,
+                        keepalive_timeout=None,
+                    )
+                ),
                 FixedCountRetryStrategy(max_attempts=3),
             )
 
@@ -81,7 +88,7 @@ def v1() -> Configurations.InRegion.Default:
                 This configuration is guaranteed not to change in future releases of the Momento Python SDK.
                 """
                 return Configurations.InRegion.Default(
-                    StaticTransportStrategy(StaticGrpcConfiguration(timedelta(milliseconds=1100))),
+                    StaticTransportStrategy(StaticGrpcConfiguration(deadline=timedelta(milliseconds=1100))),
                     FixedCountRetryStrategy(max_attempts=3),
                 )
 
@@ -108,6 +115,6 @@ def v1() -> Configurations.InRegion.LowLatency:
                 This configuration is guaranteed not to change in future releases of the Momento Python SDK.
                 """
                 return Configurations.InRegion.LowLatency(
-                    StaticTransportStrategy(StaticGrpcConfiguration(timedelta(milliseconds=500))),
+                    StaticTransportStrategy(StaticGrpcConfiguration(deadline=timedelta(milliseconds=500))),
                     FixedCountRetryStrategy(max_attempts=3),
                 )

src/momento/config/transport/grpc_configuration.py

@@ -22,3 +22,23 @@ def with_root_certificates_pem(self, root_certificates_pem_path: Path) -> GrpcCo
     @abstractmethod
     def get_root_certificates_pem(self) -> Optional[bytes]:
         pass
+
+    @abstractmethod
+    def get_max_send_message_length(self) -> Optional[int]:
+        pass
+
+    @abstractmethod
+    def get_max_receive_message_length(self) -> Optional[int]:
+        pass
+
+    @abstractmethod
+    def get_keepalive_permit_without_calls(self) -> Optional[int]:
+        pass
+
+    @abstractmethod
+    def get_keepalive_time(self) -> Optional[timedelta]:
+        pass
+
+    @abstractmethod
+    def get_keepalive_timeout(self) -> Optional[timedelta]:
+        pass

src/momento/config/transport/transport_strategy.py

@@ -48,9 +48,23 @@ def with_client_timeout(self, client_timeout: timedelta) -> TransportStrategy:
 
 
 class StaticGrpcConfiguration(GrpcConfiguration):
-    def __init__(self, deadline: timedelta, root_certificates_pem: Optional[bytes] = None):
+    def __init__(
+        self,
+        deadline: timedelta,
+        root_certificates_pem: Optional[bytes] = None,
+        max_send_message_length: Optional[int] = None,
+        max_receive_message_length: Optional[int] = None,
+        keepalive_permit_without_calls: Optional[bool] = True,
+        keepalive_time: Optional[timedelta] = timedelta(milliseconds=5000),
+        keepalive_timeout: Optional[timedelta] = timedelta(milliseconds=1000),
+    ):
         self._deadline = deadline
         self._root_certificates_pem = root_certificates_pem
+        self._max_send_message_length = max_send_message_length
+        self._max_receive_message_length = max_receive_message_length
+        self._keepalive_permit_without_calls = keepalive_permit_without_calls
+        self._keepalive_time = keepalive_time
+        self._keepalive_timeout = keepalive_timeout
 
     def get_deadline(self) -> timedelta:
         return self._deadline
@@ -71,6 +85,21 @@ def with_root_certificates_pem(self, root_certificates_pem_path: Path) -> GrpcCo
     def get_root_certificates_pem(self) -> Optional[bytes]:
         return self._root_certificates_pem
 
+    def get_max_send_message_length(self) -> Optional[int]:
+        return self._max_send_message_length
+
+    def get_max_receive_message_length(self) -> Optional[int]:
+        return self._max_receive_message_length
+
+    def get_keepalive_permit_without_calls(self) -> Optional[int]:
+        return int(self._keepalive_permit_without_calls) if self._keepalive_permit_without_calls is not None else None
+
+    def get_keepalive_time(self) -> Optional[timedelta]:
+        return self._keepalive_time
+
+    def get_keepalive_timeout(self) -> Optional[timedelta]:
+        return self._keepalive_timeout
+
 
 class StaticTransportStrategy(TransportStrategy):
     def __init__(self, grpc_configuration: GrpcConfiguration):

src/momento/config/vector_index_configurations.py

@@ -25,6 +25,6 @@ def latest() -> VectorIndexConfigurations.Default:
             This configuration will be updated every time there is a new version of the laptop configuration.
             """
             return VectorIndexConfigurations.Default(
-                # This is high to account for time-intensive adds.
-                StaticTransportStrategy(StaticGrpcConfiguration(timedelta(seconds=120)))
+                # The deadline is high to account for time-intensive adds.
+                StaticTransportStrategy(StaticGrpcConfiguration(deadline=timedelta(seconds=120)))
             )

src/momento/internal/_utilities/__init__.py

@@ -14,6 +14,7 @@
     _validate_ttl,
 )
 from ._momento_version import momento_version
+from ._time import _timedelta_to_ms
 from ._vector_index_validation import (
     _validate_index_name,
     _validate_num_dimensions,

src/momento/internal/_utilities/_grpc_channel_options.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from typing import Sequence, Tuple, Union
+
+from momento.config.transport.grpc_configuration import GrpcConfiguration
+from momento.config.transport.transport_strategy import StaticGrpcConfiguration
+from momento.internal._utilities import _timedelta_to_ms
+
+DEFAULT_MAX_MESSAGE_SIZE = 5_243_000  # bytes
+
+ChannelArguments = Sequence[Tuple[str, Union[int, None]]]
+
+
+def grpc_data_channel_options_from_grpc_config(grpc_config: GrpcConfiguration) -> ChannelArguments:
+    """Create gRPC channel options from a GrpcConfiguration.
+
+    Args:
+        grpc_config (GrpcConfiguration): the gRPC configuration.
+
+    Returns:
+        grpc.aio.ChannelArgumentType: a list of gRPC channel options as key-value tuples.
+    """
+    channel_options = []
+
+    max_send_length = grpc_config.get_max_send_message_length()
+    channel_options.append(
+        ("grpc.max_send_message_length", max_send_length if max_send_length is not None else DEFAULT_MAX_MESSAGE_SIZE)
+    )
+
+    max_receive_length = grpc_config.get_max_receive_message_length()
+    channel_options.append(
+        (
+            "grpc.max_receive_message_length",
+            max_receive_length if max_receive_length is not None else DEFAULT_MAX_MESSAGE_SIZE,
+        )
+    )
+
+    keepalive_permit = grpc_config.get_keepalive_permit_without_calls()
+    if keepalive_permit is not None:
+        channel_options.append(("grpc.keepalive_permit_without_calls", keepalive_permit))
+
+    keepalive_time = grpc_config.get_keepalive_time()
+    if keepalive_time is not None:
+        channel_options.append(("grpc.keepalive_time_ms", _timedelta_to_ms(keepalive_time)))
+
+    keepalive_timeout = grpc_config.get_keepalive_timeout()
+    if keepalive_timeout is not None:
+        channel_options.append(("grpc.keepalive_timeout_ms", _timedelta_to_ms(keepalive_timeout)))
+
+    return channel_options
+
+
+def grpc_control_channel_options_from_grpc_config(grpc_config: GrpcConfiguration) -> ChannelArguments:
+    """Create gRPC channel options from a GrpcConfiguration, but disable keepalives.
+
+    Args:
+        grpc_config (GrpcConfiguration): the gRPC configuration.
+
+    Returns:
+        grpc.aio.ChannelArgumentType: a list of gRPC channel options as key-value tuples.
+    """
+    # Override the keepalive options to disable keepalives
+    control_grpc_config = StaticGrpcConfiguration(
+        deadline=grpc_config.get_deadline(),
+        root_certificates_pem=grpc_config.get_root_certificates_pem(),
+        max_send_message_length=grpc_config.get_max_send_message_length(),
+        max_receive_message_length=grpc_config.get_max_receive_message_length(),
+        keepalive_permit_without_calls=None,
+        keepalive_time=None,
+        keepalive_timeout=None,
+    )
+    return grpc_data_channel_options_from_grpc_config(control_grpc_config)

src/momento/internal/_utilities/_time.py

@@ -0,0 +1,15 @@
+from datetime import timedelta
+
+
+def _timedelta_to_ms(delta: timedelta) -> int:
+    """Expresses a timedelta as milliseconds.
+
+    Note: truncates the microseconds.
+
+    Args:
+        delta (timedelta): The timedelta to convert.
+
+    Returns:
+        int: The total duration of the timedelta in milliseconds.
+    """
+    return int(delta.total_seconds() * 1000)

src/momento/internal/aio/_scs_data_client.py

@@ -16,6 +16,7 @@
     _gen_dictionary_items_as_bytes,
     _gen_list_as_bytes,
     _gen_set_input_as_bytes,
+    _timedelta_to_ms,
     _validate_cache_name,
     _validate_dictionary_name,
     _validate_list_name,
@@ -154,7 +155,7 @@ def __init__(self, configuration: Configuration, credential_provider: Credential
         self._endpoint = endpoint
 
         default_deadline: timedelta = configuration.get_transport_strategy().get_grpc_configuration().get_deadline()
-        self._default_deadline_seconds = int(default_deadline.total_seconds())
+        self._default_deadline_seconds = default_deadline.total_seconds()
 
         self._grpc_manager = _DataGrpcManager(configuration, credential_provider)
         _validate_ttl(default_ttl)
@@ -1107,7 +1108,7 @@ def _ttl_or_default_milliseconds(self, ttl: Optional[timedelta]) -> int:
         if ttl is not None:
             which_ttl = ttl
 
-        return int(which_ttl.total_seconds() * 1000)
+        return _timedelta_to_ms(which_ttl)
 
     def _build_stub(self) -> cache_grpc.ScsStub:
         return self._grpc_manager.async_stub()

src/momento/internal/aio/_scs_grpc_manager.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+from datetime import timedelta
 from typing import Optional
 
 import grpc
@@ -10,10 +11,15 @@
 
 from momento.auth import CredentialProvider
 from momento.config import Configuration, TopicConfiguration
+from momento.config.transport.transport_strategy import StaticGrpcConfiguration
 from momento.internal._utilities import momento_version
 from momento.internal._utilities._channel_credentials import (
     channel_credentials_from_root_certs_or_default,
 )
+from momento.internal._utilities._grpc_channel_options import (
+    grpc_control_channel_options_from_grpc_config,
+    grpc_data_channel_options_from_grpc_config,
+)
 from momento.retry import RetryStrategy
 
 from ... import logs
@@ -35,6 +41,9 @@ def __init__(self, configuration: Configuration, credential_provider: Credential
             target=credential_provider.control_endpoint,
             credentials=channel_credentials_from_root_certs_or_default(configuration),
             interceptors=_interceptors(credential_provider.auth_token, configuration.get_retry_strategy()),
+            options=grpc_control_channel_options_from_grpc_config(
+                grpc_config=configuration.get_transport_strategy().get_grpc_configuration(),
+            ),
         )
 
     async def close(self) -> None:
@@ -64,11 +73,14 @@ def __init__(self, configuration: Configuration, credential_provider: Credential
             # https://grpc.github.io/grpc/python/grpc.html
             # https://grpc.github.io/grpc/python/glossary.html#term-channel_arguments
             # https://github.com/grpc/grpc/blob/v1.46.x/include/grpc/impl/codegen/grpc_types.h#L140
-            options=[
-                # ('grpc.max_concurrent_streams', 1000),
-                # ('grpc.use_local_subchannel_pool', 1),
-                # (experimental.ChannelOptions.SingleThreadedUnaryStream, 1)
-            ],
+            # options=[
+            #     ('grpc.max_concurrent_streams', 1000),
+            #     ('grpc.use_local_subchannel_pool', 1),
+            #     (experimental.ChannelOptions.SingleThreadedUnaryStream, 1)
+            # ],
+            options=grpc_data_channel_options_from_grpc_config(
+                configuration.get_transport_strategy().get_grpc_configuration()
+            ),
         )
 
     async def eagerly_connect(self, timeout_seconds: float) -> None:
@@ -79,6 +91,9 @@ async def eagerly_connect(self, timeout_seconds: float) -> None:
             await asyncio.wait_for(self.wait_for_ready(), timeout_seconds)
         except Exception as error:
             self._logger.debug(f"Failed to connect to the server within the given timeout. {error}")
+            raise RuntimeError(
+                f"Failed to connect to Momento's server within given eager connection timeout {error}"
+            ) from error
 
     async def wait_for_ready(self) -> None:
         latest_state = self._secure_channel.get_state(True)  # try_to_connect
@@ -117,10 +132,14 @@ class _PubsubGrpcManager:
     version = momento_version
 
     def __init__(self, configuration: TopicConfiguration, credential_provider: CredentialProvider):
+        # NOTE: This is hard-coded for now but we may want to expose it via TopicConfiguration in the future, as we do with some of the other clients.
+        grpc_config = StaticGrpcConfiguration(deadline=timedelta(milliseconds=1100))
+
         self._secure_channel = grpc.aio.secure_channel(
             target=credential_provider.cache_endpoint,
             credentials=grpc.ssl_channel_credentials(),
             interceptors=_interceptors(credential_provider.auth_token, None),
+            options=grpc_data_channel_options_from_grpc_config(grpc_config),
         )
 
     async def close(self) -> None:
@@ -136,10 +155,14 @@ class _PubsubGrpcStreamManager:
     version = momento_version
 
     def __init__(self, configuration: TopicConfiguration, credential_provider: CredentialProvider):
+        # NOTE: This is hard-coded for now but we may want to expose it via TopicConfiguration in the future, as we do with some of the other clients.
+        grpc_config = StaticGrpcConfiguration(deadline=timedelta(milliseconds=1100))
+
         self._secure_channel = grpc.aio.secure_channel(
             target=credential_provider.cache_endpoint,
             credentials=grpc.ssl_channel_credentials(),
             interceptors=_stream_interceptors(credential_provider.auth_token),
+            options=grpc_data_channel_options_from_grpc_config(grpc_config),
         )
 
     async def close(self) -> None:

src/momento/internal/aio/_vector_index_data_client.py

@@ -45,7 +45,7 @@ def __init__(self, configuration: VectorIndexConfiguration, credential_provider:
         self._endpoint = endpoint
 
         default_deadline: timedelta = configuration.get_transport_strategy().get_grpc_configuration().get_deadline()
-        self._default_deadline_seconds = int(default_deadline.total_seconds())
+        self._default_deadline_seconds = default_deadline.total_seconds()
 
         self._grpc_manager = _VectorIndexDataGrpcManager(configuration, credential_provider)