feat: add sync cli

.pre-commit-config.yaml

@@ -33,5 +33,6 @@ repos:
           - neptune-api==0.11.0
           - more-itertools
           - backoff
+          - types-click
 default_language_version:
   python: python3

pyproject.toml

@@ -15,6 +15,8 @@ neptune-api = "^0.11.0"
 more-itertools = "^10.0.0"
 psutil = "^5.0.0"
 backoff = "^2.0.0"
+click = ">=7.0"
+tqdm = "^4.21.0"
 
 [tool.poetry]
 name = "neptune-scale"
@@ -95,3 +97,6 @@ show_error_codes = "True"
 
 [tool.pytest.ini_options]
 addopts = "--doctest-modules -n auto"
+
+[tool.poetry.scripts]
+neptune = "neptune_scale.cli.commands:main"

src/neptune_scale/cli/commands.py

@@ -0,0 +1,56 @@
+#
+# Copyright (c) 2022, Neptune Labs Sp. z o.o.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+__all__ = ["sync"]
+
+import os
+from pathlib import Path
+from typing import Optional
+
+import click
+
+from neptune_scale.cli.sync import sync_all
+from neptune_scale.exceptions import NeptuneApiTokenNotProvided
+from neptune_scale.util.envs import API_TOKEN_ENV_NAME
+
+
+@click.group()
+def main() -> None:
+    pass
+
+
+@main.command()
+@click.argument(
+    "run_log_file",
+    type=click.Path(exists=True, dir_okay=False, resolve_path=True),
+    metavar="<run-log-file>",
+)
+@click.option(
+    "--api-token",
+    "api_token",
+    multiple=False,
+    default=os.environ.get(API_TOKEN_ENV_NAME),
+    metavar="<api-token>",
+    help="API token for authentication. Overrides NEPTUNE_API_TOKEN environment variable",
+)
+def sync(
+    run_log_file: Path,
+    api_token: Optional[str],
+) -> None:
+    if api_token is None:
+        raise NeptuneApiTokenNotProvided()
+
+    sync_all(run_log_file, api_token)

src/neptune_scale/cli/sync.py

@@ -0,0 +1,141 @@
+#
+# Copyright (c) 2022, Neptune Labs Sp. z o.o.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+__all__ = ["sync_all"]
+
+from pathlib import Path
+from typing import Optional
+
+from tqdm import tqdm
+
+from neptune_scale.sync import sync_process
+from neptune_scale.sync.errors_tracking import (
+    ErrorsMonitor,
+    ErrorsQueue,
+)
+from neptune_scale.sync.operations_repository import (
+    OperationsRepository,
+    SequenceId,
+)
+from neptune_scale.sync.sync_process import SyncProcess
+from neptune_scale.util import (
+    SharedFloat,
+    SharedInt,
+    get_logger,
+)
+
+logger = get_logger()
+
+
+def sync_all(
+    run_log_file: Path,
+    api_token: str,
+) -> None:
+    runner = SyncRunner(api_token=api_token, run_log_file=run_log_file)
+
+    try:
+        runner.start()
+        runner.wait()
+    finally:
+        runner.stop()
+
+
+class SyncRunner:
+    def __init__(
+        self,
+        api_token: str,
+        run_log_file: Path,
+    ) -> None:
+        self._api_token: str = api_token
+        self._run_log_file: Path = run_log_file
+        self._operations_repository: OperationsRepository = OperationsRepository(db_path=run_log_file)
+        self._process_link = sync_process.ProcessLink()
+        self._errors_queue: ErrorsQueue = ErrorsQueue()
+        self._last_queued_seq = SharedInt(-1)
+        self._last_ack_seq = SharedInt(-1)
+        self._last_ack_timestamp = SharedFloat(-1)
+        self._log_seq_id_range: Optional[tuple[SequenceId, SequenceId]] = None
+        self._sync_process: Optional[SyncProcess] = None
+        self._errors_monitor: Optional[ErrorsMonitor] = None
+
+    def start(
+        self,
+    ) -> None:
+        self._log_seq_id_range = self._operations_repository.get_sequence_id_range()
+        if self._log_seq_id_range is None:
+            logger.info("No operations to process")
+            return
+
+        metadata = self._operations_repository.get_metadata()
+        if metadata is None:
+            logger.error("No run metadata found in log")
+            return
+
+        self._sync_process = sync_process.SyncProcess(
+            operations_repository_path=self._run_log_file,
+            errors_queue=self._errors_queue,
+            process_link=self._process_link,
+            api_token=self._api_token,
+            project=metadata.project,
+            family=metadata.run_id,
+            mode="async",
+            last_queued_seq=self._last_queued_seq,
+            last_ack_seq=self._last_ack_seq,
+            last_ack_timestamp=self._last_ack_timestamp,
+        )
+        self._errors_monitor = ErrorsMonitor(errors_queue=self._errors_queue)
+
+        self._sync_process.start()
+        self._process_link.start()
+
+        self._errors_monitor.start()
+
+    def wait(self, progress_bar_enabled: bool = True, wait_time: float = 0.1) -> None:
+        if self._log_seq_id_range is None:
+            return
+
+        total_count = self._log_seq_id_range[1] - self._log_seq_id_range[0] + 1
+        with tqdm(
+            desc="Syncing operations", total=total_count, unit="op", disable=not progress_bar_enabled
+        ) as progress_bar:
+            while True:
+                try:
+                    with self._last_ack_seq:
+                        self._last_ack_seq.wait(timeout=wait_time)
+                        last_ack_seq_id = self._last_ack_seq.value
+
+                    if last_ack_seq_id != -1:
+                        acked_count = last_ack_seq_id - self._log_seq_id_range[0] + 1
+                        progress_bar.update(acked_count - progress_bar.n)
+
+                    if last_ack_seq_id >= self._log_seq_id_range[1]:
+                        break
+                except KeyboardInterrupt:
+                    logger.warning("Waiting interrupted by user")
+                    return
+
+    def stop(self) -> None:
+        if self._errors_monitor is not None:
+            self._errors_monitor.interrupt()
+            self._errors_monitor.join()
+
+        if self._sync_process is not None:
+            self._sync_process.terminate()
+            self._sync_process.join()
+            self._process_link.stop()
+
+        self._operations_repository.close()
+        self._errors_queue.close()

src/neptune_scale/sync/operations_repository.py

@@ -262,6 +262,26 @@ def get_metadata(self) -> Optional[Metadata]:
                 version=version, project=project, run_id=run_id, parent_run_id=parent_run_id, fork_step=fork_step
             )
 
+    def get_sequence_id_range(self) -> Optional[tuple[SequenceId, SequenceId]]:
+        with self._get_connection() as conn:  # type: ignore
+            cursor = conn.cursor()
+
+            cursor.execute(
+                """
+                SELECT MIN(sequence_id), MAX(sequence_id)
+                FROM run_operations
+                """
+            )
+
+            row = cursor.fetchone()
+            if not row:
+                return None
+
+            min_seq_id, max_seq_id = row
+            if min_seq_id is None or max_seq_id is None:
+                return None
+            return SequenceId(min_seq_id), SequenceId(max_seq_id)
+
     def close(self) -> None:
         with self._lock:
             if self._connection is not None:

tests/unit/sync/test_operations_repository.py

@@ -160,6 +160,37 @@ def test_get_operations_empty_db(operations_repo):
     assert len(operations) == 0
 
 
+def test_get_sequence_id_range_single(operations_repo):
+    # Given
+    snapshots = [UpdateRunSnapshot(assign={"key": Value(string="a")})]
+    operations_repo.save_update_run_snapshots(snapshots)
+
+    # When
+    start_end = operations_repo.get_sequence_id_range()
+
+    # Then
+    assert start_end == (1, 1)
+
+
+def test_get_sequence_id_range_multiple(operations_repo):
+    # Given
+    for i in range(5):
+        snapshots = [UpdateRunSnapshot(assign={f"key_{i}": Value(string=f"value_{i}")})]
+        operations_repo.save_update_run_snapshots(snapshots)
+
+    # When
+    start_end = operations_repo.get_sequence_id_range()
+
+    # Then
+    assert start_end == (1, 5)
+
+
+def test_get_sequence_id_range_empty_db(operations_repo):
+    # Given
+    start_end = operations_repo.get_sequence_id_range()
+    assert start_end is None
+
+
 def test_delete_operations(operations_repo, temp_db_path):
     # Given
     snapshots = []