extension to enable/disable remote signals (#2)

* extension to enable/disable remote signals

* code review

CHANGELOG.md

@@ -4,4 +4,4 @@
 - ...
 
 ### Changes
-- ...
+- Allow to disable handling of remote signals ([#2](https://github.com/neptune-ai/neptune-client-experimental/pull/2))

src/neptune_experimental/run.py

@@ -18,16 +18,34 @@
 from typing import Any
 
 import neptune
+from neptune.internal.backgroud_job_list import BackgroundJobList
+from neptune.internal.utils import verify_type
+from neptune.internal.websockets.websocket_signals_background_job import WebsocketSignalsBackgroundJob
 
 
-# That's just a boilerplate code to make sure that the extension is loaded
 class CustomRun(neptune.Run):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        print("That's custom class")
+        enable_remote_signals = kwargs.pop("enable_remote_signals", None)
 
-        kwargs["capture_hardware_metrics"] = False
-        kwargs["capture_stdout"] = False
-        kwargs["capture_stderr"] = False
-        kwargs["capture_traceback"] = False
+        if enable_remote_signals is None:
+            self._enable_remote_signals = True  # user did not pass this param in kwargs -> default value
+        else:
+
+            verify_type("enable_remote_signals", enable_remote_signals, bool)
+            self._enable_remote_signals = enable_remote_signals
 
         super().__init__(*args, **kwargs)
+
+    def _prepare_background_jobs(self) -> BackgroundJobList:
+        background_jobs = super()._prepare_background_jobs()
+
+        if not self._enable_remote_signals:
+            # filter-out websocket job
+            background_jobs._jobs = list(
+                filter(
+                    lambda x: not isinstance(x, WebsocketSignalsBackgroundJob),
+                    background_jobs._jobs,
+                )
+            )
+
+        return background_jobs

tests/unit/test_custom_run.py

@@ -13,11 +13,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+
 from neptune import Run
+from neptune.internal.websockets.websocket_signals_background_job import WebsocketSignalsBackgroundJob
 
 from neptune_experimental.run import CustomRun
 
 
 def test_custom_run():
     with Run(mode="debug") as run:
         assert isinstance(run, CustomRun)
+
+
+def test_disabled_remote_signals():
+    with Run(mode="debug", enable_remote_signals=False) as run:
+        assert run._enable_remote_signals is False
+        jobs = run._prepare_background_jobs()._jobs
+        assert not [job for job in jobs if isinstance(job, WebsocketSignalsBackgroundJob)]