Adjustments after no synchronization callbacks

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ### Changes
 - Added support for download progress update handling in fetching API ([#25](https://github.com/neptune-ai/neptune-client-experimental/pull/25))
+- Adjustments after no synchronization callbacks fix ([#20](https://github.com/neptune-ai/neptune-client-experimental/pull/20))
 
 
 ## neptune-experimental 0.2.0

pyproject.toml

@@ -12,7 +12,7 @@ pattern = "default-unprefixed"
 python = "^3.7"
 
 # Base neptune package
-neptune = "^1.8.4"
+neptune = "^1.8.5"
 
 # For 3.7 compatibility
 typing_extensions = "^4.0.0"

src/neptune_experimental/partitioned_operation_processor.py

@@ -18,9 +18,10 @@
 import shutil
 import threading
 from datetime import datetime
+from queue import Queue
 from typing import (
+    TYPE_CHECKING,
     Any,
-    Callable,
     List,
     Optional,
 )
@@ -29,15 +30,15 @@
 from neptune.internal.backends.neptune_backend import NeptuneBackend
 from neptune.internal.container_type import ContainerType
 from neptune.internal.id_formats import UniqueId
-from neptune.internal.init.parameters import (
-    ASYNC_LAG_THRESHOLD,
-    ASYNC_NO_PROGRESS_THRESHOLD,
-)
 from neptune.internal.operation import Operation
 from neptune.internal.operation_processors.async_operation_processor import AsyncOperationProcessor
 from neptune.internal.operation_processors.operation_processor import OperationProcessor
 from neptune.internal.operation_processors.operation_storage import get_container_dir
 
+if TYPE_CHECKING:
+    from neptune.internal.signals_processing.signals import Signal
+
+
 _logger = logging.getLogger(__name__)
 
 
@@ -48,12 +49,9 @@ def __init__(
         container_type: ContainerType,
         backend: NeptuneBackend,
         lock: threading.RLock,
+        queue: "Queue[Signal]",
         batch_size: int,
         sleep_time: float = 5,
-        async_lag_callback: Optional[Callable[[], None]] = None,
-        async_lag_threshold: float = ASYNC_LAG_THRESHOLD,
-        async_no_progress_callback: Optional[Callable[[], None]] = None,
-        async_no_progress_threshold: float = ASYNC_NO_PROGRESS_THRESHOLD,
         partitions: int = 5,
     ):
         self._data_path = self._init_data_path(container_id, container_type)
@@ -64,12 +62,9 @@ def __init__(
                 container_type=container_type,
                 backend=backend,
                 lock=lock,
+                queue=queue,
                 sleep_time=sleep_time,
                 batch_size=batch_size,
-                async_lag_callback=async_lag_callback,
-                async_lag_threshold=async_lag_threshold,
-                async_no_progress_callback=async_no_progress_callback,
-                async_no_progress_threshold=async_no_progress_threshold,
                 data_path=self._data_path / f"partition-{partition_id}",
                 should_print_logs=False,
             )

src/neptune_experimental/remote_signals.py

@@ -19,21 +19,22 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    List,
 )
 
 from neptune_experimental.utils import wrap_method
 
 if TYPE_CHECKING:
     from neptune import Run
-    from neptune.internal.backgroud_job_list import BackgroundJobList
+    from neptune.internal.background_job import BackgroundJob
 
 
 def initialize() -> None:
     from neptune import Run
 
     # Monkey patching
     wrap_method(obj=Run, method="__init__", wrapper=init_with_enable_remote_signals)
-    wrap_method(obj=Run, method="_prepare_background_jobs", wrapper=prepare_background_jobs)
+    wrap_method(obj=Run, method="_get_background_jobs", wrapper=get_background_jobs)
 
 
 def init_with_enable_remote_signals(self: "Run", *args: Any, original: Callable[..., Any], **kwargs: Any) -> None:
@@ -51,17 +52,17 @@ def init_with_enable_remote_signals(self: "Run", *args: Any, original: Callable[
     original(self, *args, **kwargs)
 
 
-def prepare_background_jobs(self: "Run", original: Callable[..., Any]) -> "BackgroundJobList":
+def get_background_jobs(self: "Run", original: Callable[..., Any]) -> List["BackgroundJob"]:
     from neptune.internal.websockets.websocket_signals_background_job import WebsocketSignalsBackgroundJob
 
-    background_jobs = original(self)
+    background_jobs: List["BackgroundJob"] = original(self)
 
     if not self._enable_remote_signals:
         # filter-out websocket job
-        background_jobs._jobs = list(
+        background_jobs = list(
             filter(
                 lambda x: not isinstance(x, WebsocketSignalsBackgroundJob),
-                background_jobs._jobs,
+                background_jobs,
             )
         )
 

tests/unit/test_remote_signals.py

@@ -21,5 +21,5 @@
 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
+        jobs = run._get_background_jobs()
         assert not [job for job in jobs if isinstance(job, WebsocketSignalsBackgroundJob)]