From 83453f52702b0a667472af815c6044f1cd73f6f0 Mon Sep 17 00:00:00 2001 From: Rafal Skolasinski Date: Thu, 16 Jan 2025 11:57:09 +0000 Subject: [PATCH 1/3] feat: make get_or_create timeout configurable --- graphdatascience/session/aura_api.py | 10 ++++++- .../session/dedicated_sessions.py | 27 +++++++++++++++---- graphdatascience/session/gds_sessions.py | 13 ++++++--- 3 files changed, 41 insertions(+), 9 deletions(-) diff --git a/graphdatascience/session/aura_api.py b/graphdatascience/session/aura_api.py index c5e2f0561..8667e4d5b 100644 --- a/graphdatascience/session/aura_api.py +++ b/graphdatascience/session/aura_api.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import math import time import warnings from collections import defaultdict @@ -168,6 +169,8 @@ def wait_for_session_running( max_wait_time: float = 300, ) -> WaitResult: waited_time = 0.0 + if max_wait_time == -1: + max_wait_time = math.inf while waited_time < max_wait_time: session = self.get_session(session_id) if session is None: @@ -186,7 +189,12 @@ def wait_for_session_running( time.sleep(sleep_time) sleep_time = min(sleep_time * 2, max_sleep_time, max_wait_time - waited_time) - return WaitResult.from_error(f"Session `{session_id}` is not running after {waited_time} seconds") + return WaitResult.from_error( + f"Session `{session_id}` is not running after {waited_time} seconds.\n" + "\tThe session may become available at a later time.\n" + f'\tConsider running `sessions.delete(session_id="{session_id}")` ' + "to avoid resource leakage." + ) def delete_session(self, session_id: str) -> bool: response = self._request_session.delete( diff --git a/graphdatascience/session/dedicated_sessions.py b/graphdatascience/session/dedicated_sessions.py index 378435019..cef24c190 100644 --- a/graphdatascience/session/dedicated_sessions.py +++ b/graphdatascience/session/dedicated_sessions.py @@ -24,7 +24,10 @@ def __init__(self, aura_api: AuraApi) -> None: self._aura_api = aura_api def estimate( - self, node_count: int, relationship_count: int, algorithm_categories: Optional[list[AlgorithmCategory]] = None + self, + node_count: int, + relationship_count: int, + algorithm_categories: Optional[list[AlgorithmCategory]] = None, ) -> SessionMemory: if algorithm_categories is None: algorithm_categories = [] @@ -56,6 +59,7 @@ def get_or_create( db_connection: DbmsConnectionInfo, ttl: Optional[timedelta] = None, cloud_location: Optional[CloudLocation] = None, + timeout: Optional[int] = 300, ) -> AuraGraphDataScience: db_runner = Neo4jQueryRunner.create_for_db( endpoint=db_connection.uri, @@ -83,7 +87,9 @@ def get_or_create( connection_url = session_details.bolt_connection_url() if session_details.status != "Ready": - wait_result = self._aura_api.wait_for_session_running(session_id) + if timeout == 0: + raise RuntimeError(f"Failed to get or create session `{session_name}`.") + wait_result = self._aura_api.wait_for_session_running(session_id, max_wait_time=timeout) if err := wait_result.error: raise RuntimeError(f"Failed to get or create session `{session_name}`: {err}") @@ -93,7 +99,11 @@ def get_or_create( password=password, ) - return self._construct_client(session_id=session_id, session_connection=session_connection, db_runner=db_runner) + return self._construct_client( + session_id=session_id, + session_connection=session_connection, + db_runner=db_runner, + ) def delete(self, *, session_name: Optional[str] = None, session_id: Optional[str] = None) -> bool: if not session_name and not session_id: @@ -160,13 +170,20 @@ def _get_or_create_session( # If cloud location is provided we go for self managed DBs path if cloud_location: return self._aura_api.get_or_create_session( - name=session_name, pwd=pwd, memory=memory, ttl=ttl, cloud_location=cloud_location + name=session_name, + pwd=pwd, + memory=memory, + ttl=ttl, + cloud_location=cloud_location, ) else: return self._aura_api.get_or_create_session(name=session_name, dbid=dbid, pwd=pwd, memory=memory, ttl=ttl) def _construct_client( - self, session_id: str, session_connection: DbmsConnectionInfo, db_runner: Neo4jQueryRunner + self, + session_id: str, + session_connection: DbmsConnectionInfo, + db_runner: Neo4jQueryRunner, ) -> AuraGraphDataScience: return AuraGraphDataScience.create( gds_session_connection_info=session_connection, diff --git a/graphdatascience/session/gds_sessions.py b/graphdatascience/session/gds_sessions.py index 857d61d35..c200396c0 100644 --- a/graphdatascience/session/gds_sessions.py +++ b/graphdatascience/session/gds_sessions.py @@ -53,7 +53,10 @@ def __init__(self, api_credentials: AuraAPICredentials) -> None: self._impl: DedicatedSessions = DedicatedSessions(aura_api) def estimate( - self, node_count: int, relationship_count: int, algorithm_categories: Optional[list[AlgorithmCategory]] = None + self, + node_count: int, + relationship_count: int, + algorithm_categories: Optional[list[AlgorithmCategory]] = None, ) -> SessionMemory: """ Estimates the memory required for a session with the given node and relationship counts. @@ -86,6 +89,7 @@ def get_or_create( db_connection: DbmsConnectionInfo, ttl: Optional[timedelta] = None, cloud_location: Optional[CloudLocation] = None, + timeout: Optional[int] = 300, ) -> AuraGraphDataScience: """ Retrieves an existing session with the given session name and database connection, @@ -98,13 +102,16 @@ def get_or_create( session_name (str): The name of the session. memory (SessionMemory): The size of the session specified by memory. db_connection (DbmsConnectionInfo): The database connection information. - ttl: Optional[timedelta]: The sessions time to live after inactivity in seconds. + ttl: (Optional[timedelta]): The sessions time to live after inactivity in seconds. cloud_location (Optional[CloudLocation]): The cloud location. Required if the GDS session is for a self-managed database. + timeout (Optional[int]): The timeout when waiting for session to be ready. If session does not become ready an exception will be raised. It is user responsibility to ensure resource gets cleaned up in this situation. If value is set to `0` method will return immediately. If it is set to `-1` it will wait forever. Returns: AuraGraphDataScience: The session. """ - return self._impl.get_or_create(session_name, memory, db_connection, ttl=ttl, cloud_location=cloud_location) + return self._impl.get_or_create( + session_name, memory, db_connection, ttl=ttl, cloud_location=cloud_location, timeout=timeout + ) def delete(self, *, session_name: Optional[str] = None, session_id: Optional[str] = None) -> bool: """ From 9f6edcfb41df15a630cdb28d45fe7aac41cf5659 Mon Sep 17 00:00:00 2001 From: Rafal Skolasinski Date: Thu, 16 Jan 2025 12:45:48 +0000 Subject: [PATCH 2/3] code review --- graphdatascience/session/aura_api.py | 4 +--- graphdatascience/session/dedicated_sessions.py | 8 ++++---- graphdatascience/session/gds_sessions.py | 4 ++-- 3 files changed, 7 insertions(+), 9 deletions(-) diff --git a/graphdatascience/session/aura_api.py b/graphdatascience/session/aura_api.py index 8667e4d5b..fa28cf4b0 100644 --- a/graphdatascience/session/aura_api.py +++ b/graphdatascience/session/aura_api.py @@ -166,11 +166,9 @@ def wait_for_session_running( session_id: str, sleep_time: float = 0.2, max_sleep_time: float = 10, - max_wait_time: float = 300, + max_wait_time: float = math.inf, ) -> WaitResult: waited_time = 0.0 - if max_wait_time == -1: - max_wait_time = math.inf while waited_time < max_wait_time: session = self.get_session(session_id) if session is None: diff --git a/graphdatascience/session/dedicated_sessions.py b/graphdatascience/session/dedicated_sessions.py index cef24c190..3f28562b0 100644 --- a/graphdatascience/session/dedicated_sessions.py +++ b/graphdatascience/session/dedicated_sessions.py @@ -1,6 +1,7 @@ from __future__ import annotations import hashlib +import math import warnings from datetime import datetime, timedelta, timezone from typing import Optional @@ -59,7 +60,7 @@ def get_or_create( db_connection: DbmsConnectionInfo, ttl: Optional[timedelta] = None, cloud_location: Optional[CloudLocation] = None, - timeout: Optional[int] = 300, + timeout: Optional[int] = None, ) -> AuraGraphDataScience: db_runner = Neo4jQueryRunner.create_for_db( endpoint=db_connection.uri, @@ -87,9 +88,8 @@ def get_or_create( connection_url = session_details.bolt_connection_url() if session_details.status != "Ready": - if timeout == 0: - raise RuntimeError(f"Failed to get or create session `{session_name}`.") - wait_result = self._aura_api.wait_for_session_running(session_id, max_wait_time=timeout) + max_wait_time = float(timeout) if timeout is not None else math.inf + wait_result = self._aura_api.wait_for_session_running(session_id, max_wait_time=max_wait_time) if err := wait_result.error: raise RuntimeError(f"Failed to get or create session `{session_name}`: {err}") diff --git a/graphdatascience/session/gds_sessions.py b/graphdatascience/session/gds_sessions.py index c200396c0..275e933ee 100644 --- a/graphdatascience/session/gds_sessions.py +++ b/graphdatascience/session/gds_sessions.py @@ -89,7 +89,7 @@ def get_or_create( db_connection: DbmsConnectionInfo, ttl: Optional[timedelta] = None, cloud_location: Optional[CloudLocation] = None, - timeout: Optional[int] = 300, + timeout: Optional[int] = None, ) -> AuraGraphDataScience: """ Retrieves an existing session with the given session name and database connection, @@ -104,7 +104,7 @@ def get_or_create( db_connection (DbmsConnectionInfo): The database connection information. ttl: (Optional[timedelta]): The sessions time to live after inactivity in seconds. cloud_location (Optional[CloudLocation]): The cloud location. Required if the GDS session is for a self-managed database. - timeout (Optional[int]): The timeout when waiting for session to be ready. If session does not become ready an exception will be raised. It is user responsibility to ensure resource gets cleaned up in this situation. If value is set to `0` method will return immediately. If it is set to `-1` it will wait forever. + timeout (Optional[int]): Optional timeout when waiting for session to become ready. If unset the method will wait forever. If set and session does not become ready an exception will be raised. It is user responsibility to ensure resource gets cleaned up in this situation. Returns: AuraGraphDataScience: The session. From 408d30e2d8aa57b861837c09f878c4cdfeb88814 Mon Sep 17 00:00:00 2001 From: Rafal Skolasinski Date: Thu, 16 Jan 2025 14:15:37 +0000 Subject: [PATCH 3/3] Update graphdatascience/session/gds_sessions.py --- graphdatascience/session/gds_sessions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/graphdatascience/session/gds_sessions.py b/graphdatascience/session/gds_sessions.py index 275e933ee..47e08f5e6 100644 --- a/graphdatascience/session/gds_sessions.py +++ b/graphdatascience/session/gds_sessions.py @@ -104,7 +104,7 @@ def get_or_create( db_connection (DbmsConnectionInfo): The database connection information. ttl: (Optional[timedelta]): The sessions time to live after inactivity in seconds. cloud_location (Optional[CloudLocation]): The cloud location. Required if the GDS session is for a self-managed database. - timeout (Optional[int]): Optional timeout when waiting for session to become ready. If unset the method will wait forever. If set and session does not become ready an exception will be raised. It is user responsibility to ensure resource gets cleaned up in this situation. + timeout (Optional[int]): Optional timeout (in seconds) when waiting for session to become ready. If unset the method will wait forever. If set and session does not become ready an exception will be raised. It is user responsibility to ensure resource gets cleaned up in this situation. Returns: AuraGraphDataScience: The session.