From 40d4048f0fde686a07154c14fa2e531c62e7d55c Mon Sep 17 00:00:00 2001 From: Aleksey Veresov Date: Mon, 21 Oct 2024 19:22:11 +0200 Subject: [PATCH 1/2] Add finalizer to Connection --- python/hopsworks/connection.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/python/hopsworks/connection.py b/python/hopsworks/connection.py index 1fe984030..8d21a7f47 100644 --- a/python/hopsworks/connection.py +++ b/python/hopsworks/connection.py @@ -16,14 +16,15 @@ import os import re -import warnings import sys +import warnings +import weakref -from requests.exceptions import ConnectionError - -from hopsworks.decorators import connected, not_connected from hopsworks import client, version from hopsworks.core import project_api, secret_api, variable_api +from hopsworks.decorators import connected, not_connected +from requests.exceptions import ConnectionError + HOPSWORKS_PORT_DEFAULT = 443 HOSTNAME_VERIFICATION_DEFAULT = True @@ -241,6 +242,7 @@ def connect(self): """ client.stop() self._connected = True + finalizer = weakref.finalize(self, self.close) try: # init client if client.base.Client.REST_ENDPOINT not in os.environ: @@ -263,6 +265,7 @@ def connect(self): self._variable_api = variable_api.VariableApi() except (TypeError, ConnectionError): self._connected = False + finalizer.detach() raise print( "Connected. Call `.close()` to terminate connection gracefully.", @@ -278,7 +281,7 @@ def close(self): This will clean up any materialized certificates on the local file system of external environments such as AWS SageMaker. - Usage is recommended but optional. + Usage is optional. """ from hsfs import client as hsfs_client from hsfs import engine as hsfs_engine @@ -299,6 +302,9 @@ def close(self): except: # noqa: E722 pass + if not self._connected: + return # the connection is already closed + client.stop() self._connected = False From 1a254761d4aac7055ff02bd4cd9b07afbd7aeb58 Mon Sep 17 00:00:00 2001 From: Aleksey Veresov Date: Tue, 22 Oct 2024 09:55:57 +0200 Subject: [PATCH 2/2] Ruff fix --- python/hopsworks/connection.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/hopsworks/connection.py b/python/hopsworks/connection.py index 8d21a7f47..c16566232 100644 --- a/python/hopsworks/connection.py +++ b/python/hopsworks/connection.py @@ -211,7 +211,8 @@ def _check_compatibility(self): warnings.warn( "The installed hopsworks client version {0} may not be compatible with the connected Hopsworks backend version {1}. \nTo ensure compatibility please install the latest bug fix release matching the minor version of your backend ({2}) by running 'pip install hopsworks=={2}.*'".format( client_version, backend_version, major_minor_backend - ) + ), + stacklevel=1, ) sys.stderr.flush()