Updating Session.close() (#3253)

* kill server on session close

* close desktop app

* cleanup

* lint

* lint

* is_open

* rm remote warning

fiftyone/core/session/client.py

@@ -9,7 +9,7 @@
 from dataclasses import asdict, dataclass
 import logging
 from retrying import retry
-from threading import Thread
+from threading import Thread, Event as ThreadEvent
 import time
 import typing as t
 
@@ -49,15 +49,27 @@ class Client:
     def __post_init__(self) -> None:
         self._subscription = str(uuid4())
         self._connected = True
+        self._closed = ThreadEvent()
+        self._closed.set()
         self._listeners: t.Dict[str, t.Set[t.Callable]] = defaultdict(set)
 
-    def run(self, state: fos.StateDescription) -> None:
-        """Runs the client subscription in a background thread
+    @property
+    def origin(self) -> str:
+        """The origin of the server"""
+        return f"http://{self.address}:{self.port}"
+
+    @property
+    def is_open(self) -> str:
+        """Whether the client is connected"""
+        return not self._closed.is_set()
+
+    def open(self, state: fos.StateDescription) -> None:
+        """Open the client connection
 
         Arg:
             state: the initial state description
         """
-        if hasattr(self, "_thread"):
+        if not self._closed.is_set():
             raise RuntimeError("Client is already running")
 
         def run_client() -> None:
@@ -97,32 +109,40 @@ def subscribe() -> None:
                     self._connected = True
                     subscribe()
                 except Exception as e:
-                    if foc.DEV_INSTALL:
+                    if logger.level == logging.DEBUG:
                         raise e
 
-                    self._connected = False
-                    print(
-                        "\r\nCould not connect session, trying again "
-                        "in 10 seconds\r\n"
-                    )
-                    time.sleep(10)
+                if self._closed.is_set():
+                    break
+
+                self._connected = False
+                print(
+                    "\r\nCould not connect session, trying again "
+                    "in 10 seconds\r\n"
+                )
+                time.sleep(10)
 
         self._thread = Thread(target=run_client, daemon=True)
         self._thread.start()
 
-    @property
-    def origin(self) -> str:
-        """The origin of the server"""
-        return f"http://{self.address}:{self.port}"
+    def close(self):
+        """Close the client connection"""
+        self._closed.set()
+        self._thread.join(timeout=0)
+        self._thread = None
 
     def send_event(self, event: EventType) -> None:
         """Sends an event to the server
 
         Args:
             event: the event
         """
+        if self._closed.is_set():
+            return
+
         if not self._connected:
             raise RuntimeError("Client is not connected")
+
         self._post_event(event)
         self._dispatch_event(event)
 
@@ -156,6 +176,9 @@ def _dispatch_event(self, event: EventType) -> None:
             listener(event)
 
     def _post_event(self, event: Event) -> None:
+        if self._closed.is_set():
+            return
+
         response = requests.post(
             f"{self.origin}/event",
             headers={"Content-type": "application/json"},

fiftyone/core/session/session.py

@@ -422,7 +422,7 @@ def __init__(
             remote=remote,
             start_time=self._get_time(),
         )
-        self._client.run(self._state)
+        self._client.open(self._state)
         _attach_listeners(self)
         _register_session(self)
 
@@ -515,6 +515,7 @@ def __del__(self) -> None:
                 # logger may already have been garbage-collected
                 print(_WAIT_INSTRUCTIONS)
 
+            self._client.close()
             _unregister_session(self)
         except:
             # e.g. globals were already garbage-collected
@@ -970,9 +971,10 @@ def open(self) -> None:
         -   Desktop: opens the desktop App, if necessary
         -   Other (non-remote): opens the App in a new browser tab
         """
-        if self.remote:
-            logger.warning("Remote sessions cannot open new App windows")
-            return
+        _register_session(self)
+
+        if not self._client.is_open:
+            self._client.open(self._state)
 
         if self.plots:
             self.plots.connect()
@@ -1102,13 +1104,14 @@ def wait(self, wait: float = 3) -> None:
 
     def close(self) -> None:
         """Closes the session and terminates the App, if necessary."""
-        if self.remote:
-            return
+        if self.desktop:
+            self._app_service.stop()
 
-        if self._client._connected and focx._get_context() == focx._NONE:
-            self._client.send_event(CloseSession())
+        if self._client.is_open and focx.is_notebook_context():
+            self.freeze()
 
         self.plots.disconnect()
+        self.__del__()
 
     def freeze(self) -> None:
         """Screenshots the active App cell, replacing it with a static image.