Add NEPTUNE_SYNC_AFTER_STOP_TIMEOUT environment variable (#1260)

Co-authored-by: Rafal Jankowski <[email protected]>

CHANGELOG.md

@@ -3,6 +3,7 @@
 ### Features
 - Added ability to provide repository path with `GitRef` to `init_run` ([#1292](https://github.com/neptune-ai/neptune-client/pull/1292))
 - Added `SupportsNamespaces` interface in `neptune.typing` for proper type annotations of Handler and Neptune objects ([#1280](https://github.com/neptune-ai/neptune-client/pull/1280))
+- Added `NEPTUNE_SYNC_AFTER_STOP_TIMEOUT` environment variable ([#1260](https://github.com/neptune-ai/neptune-client/pull/1260))
 - `Run`, `Model`, `ModelVersion` and `Project` could be created with constructor in addition to `init_*` functions ([#1246](https://github.com/neptune-ai/neptune-client/pull/1246))
 
 ### Fixes

src/neptune/envs.py

@@ -26,6 +26,7 @@
     "NEPTUNE_SYNC_BATCH_TIMEOUT_ENV",
     "NEPTUNE_SUBPROCESS_KILL_TIMEOUT",
     "NEPTUNE_FETCH_TABLE_STEP_SIZE",
+    "NEPTUNE_SYNC_AFTER_STOP_TIMEOUT",
     "NEPTUNE_REQUEST_TIMEOUT",
 ]
 
@@ -54,6 +55,8 @@
 
 NEPTUNE_FETCH_TABLE_STEP_SIZE = "NEPTUNE_FETCH_TABLE_STEP_SIZE"
 
+NEPTUNE_SYNC_AFTER_STOP_TIMEOUT = "NEPTUNE_SYNC_AFTER_STOP_TIMEOUT"
+
 NEPTUNE_REQUEST_TIMEOUT = "NEPTUNE_REQUEST_TIMEOUT"
 
 S3_ENDPOINT_URL = "S3_ENDPOINT_URL"

src/neptune/internal/operation_processors/async_operation_processor.py

@@ -33,6 +33,7 @@
     ASYNC_DIRECTORY,
     NEPTUNE_DATA_DIRECTORY,
 )
+from neptune.envs import NEPTUNE_SYNC_AFTER_STOP_TIMEOUT
 from neptune.exceptions import NeptuneSynchronizationAlreadyStoppedException
 from neptune.internal.backends.neptune_backend import NeptuneBackend
 from neptune.internal.container_type import ContainerType
@@ -49,7 +50,7 @@
 
 class AsyncOperationProcessor(OperationProcessor):
     STOP_QUEUE_STATUS_UPDATE_FREQ_SECONDS = 30
-    STOP_QUEUE_MAX_TIME_NO_CONNECTION_SECONDS = 300
+    STOP_QUEUE_MAX_TIME_NO_CONNECTION_SECONDS = int(os.getenv(NEPTUNE_SYNC_AFTER_STOP_TIMEOUT, "300"))
 
     def __init__(
         self,
@@ -148,6 +149,9 @@ def _wait_for_queue_empty(self, initial_queue_size: int, seconds: Optional[float
 
         while True:
             if seconds is None:
+                if self._consumer.last_backoff_time == 0:
+                    # reset `waiting_start` on successful action
+                    waiting_start = monotonic()
                 wait_time = self.STOP_QUEUE_STATUS_UPDATE_FREQ_SECONDS
             else:
                 wait_time = max(