From c7fdf1d690617697462b0ea13e9bf30fd7292cc1 Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Mon, 30 Sep 2024 13:08:03 -0700
Subject: [PATCH] Remove debug prints

---
 python/distributed-ucxx/distributed_ucxx/ucxx.py | 14 ++------------
 1 file changed, 2 insertions(+), 12 deletions(-)

diff --git a/python/distributed-ucxx/distributed_ucxx/ucxx.py b/python/distributed-ucxx/distributed_ucxx/ucxx.py
index 7234825a..605df8a3 100644
--- a/python/distributed-ucxx/distributed_ucxx/ucxx.py
+++ b/python/distributed-ucxx/distributed_ucxx/ucxx.py
@@ -98,7 +98,6 @@ def _stop_notifier_thread_and_progress_tasks():
     """
     ctx = ucxx.core._get_ctx()
     if len(ctx._dask_resources) == 0:
-        print(f"[{os.getpid()}] _stop_notifier_thread_and_progress_tasks", flush=True)
         ucxx.stop_notifier_thread()
         ucxx.core._get_ctx().progress_tasks.clear()
 
@@ -117,7 +116,6 @@ def _register_dask_resource(resource, name=None):
 
     with ctx._dask_resources_lock:
         ctx._dask_resources.add(resource)
-    print(f"[{os.getpid()}] {name} registered: {ctx._dask_resources=}", flush=True)
 
 
 def _deregister_dask_resource(resource, name=None):
@@ -132,6 +130,8 @@ def _deregister_dask_resource(resource, name=None):
 
     ctx = ucxx.core._get_ctx()
 
+    # Check if the attribute exists first, in tests the UCXX context may have
+    # been reset before some resources are deregistered.
     if hasattr(ctx, "_dask_resources_lock"):
         with ctx._dask_resources_lock:
             try:
@@ -139,12 +139,6 @@ def _deregister_dask_resource(resource, name=None):
             except KeyError:
                 pass
             _stop_notifier_thread_and_progress_tasks()
-            print(f"[{os.getpid()}] {name} deregistered: {ctx._dask_resources=}")
-    else:
-        print(
-            f"[{os.getpid()}] {name}: {ctx=} has no _dask_resources_lock, "
-            "the context was probably already reset."
-        )
 
 
 def _allocate_dask_resources_tracker() -> None:
@@ -161,10 +155,6 @@ def _allocate_dask_resources_tracker() -> None:
 
         ctx._dask_resources = set()
         ctx._dask_resources_lock = Lock()
-        print(
-            f"[{os.getpid()}] init_once: {ctx=}, {ctx._dask_resources}, "
-            f"{ctx.progress_tasks=}, {ctx.notifier_thread=}"
-        )
 
 
 def init_once():