kleinkram 0.43.2.dev20250331124109__py3-none-any.whl → 0.58.0.dev20260110152317__py3-none-any.whl

kleinkram/api/file_transfer.py

@@ -8,6 +8,7 @@ from concurrent.futures import as_completed
 from enum import Enum
 from pathlib import Path
 from time import monotonic
+from time import sleep
 from typing import Dict
 from typing import NamedTuple
 from typing import Optional
@@ -17,6 +18,9 @@ from uuid import UUID
 import boto3.s3.transfer
 import botocore.config
 import httpx
+from rich.console import Console
+from tqdm import tqdm
+
 from kleinkram.api.client import AuthenticatedClient
 from kleinkram.config import get_config
 from kleinkram.errors import AccessDenied
@@ -27,21 +31,23 @@ from kleinkram.utils import format_bytes
 from kleinkram.utils import format_error
 from kleinkram.utils import format_traceback
 from kleinkram.utils import styled_string
-from rich.console import Console
-from tqdm import tqdm
 
 logger = logging.getLogger(__name__)
 
 UPLOAD_CREDS = "/files/temporaryAccess"
-UPLOAD_CONFIRM = "/queue/confirmUpload"
+UPLOAD_CONFIRM = "/files/upload/confirm"
 UPLOAD_CANCEL = "/files/cancelUpload"
 
 DOWNLOAD_CHUNK_SIZE = 1024 * 1024 * 16
 DOWNLOAD_URL = "/files/download"
 
+MAX_UPLOAD_RETRIES = 3
 S3_MAX_RETRIES = 60  # same as frontend
 S3_READ_TIMEOUT = 60 * 5  # 5 minutes
 
+RETRY_BACKOFF_BASE = 2  # exponential backoff base
+MAX_RETRIES = 5
+
 
 class UploadCredentials(NamedTuple):
     access_key: str
@@ -51,23 +57,20 @@ class UploadCredentials(NamedTuple):
     bucket: str
 
 
-def _confirm_file_upload(
-    client: AuthenticatedClient, file_id: UUID, file_hash: str
-) -> None:
+def _confirm_file_upload(client: AuthenticatedClient, file_id: UUID, file_hash: str) -> None:
     data = {
         "uuid": str(file_id),
         "md5": file_hash,
+        "source": "CLI",
     }
     resp = client.post(UPLOAD_CONFIRM, json=data)
     resp.raise_for_status()
 
 
-def _cancel_file_upload(
-    client: AuthenticatedClient, file_id: UUID, mission_id: UUID
-) -> None:
+def _cancel_file_upload(client: AuthenticatedClient, file_id: UUID, mission_id: UUID) -> None:
     data = {
-        "uuid": [str(file_id)],
-        "missionUUID": str(mission_id),
+        "uuids": [str(file_id)],
+        "missionUuid": str(mission_id),
     }
     resp = client.post(UPLOAD_CANCEL, json=data)
     resp.raise_for_status()
@@ -91,9 +94,16 @@ def _get_upload_creditials(
     dct = {
         "filenames": [internal_filename],
         "missionUUID": str(mission_id),
+        "source": "CLI",
     }
-    resp = client.post(UPLOAD_CREDS, json=dct)
-    resp.raise_for_status()
+    try:
+        resp = client.post(UPLOAD_CREDS, json=dct)
+        resp.raise_for_status()
+    except httpx.HTTPStatusError as e:
+        # 409 Conflict means file already exists
+        if e.response.status_code == 409:
+            return None
+        raise
 
     data = resp.json()["data"][0]
 
@@ -151,6 +161,33 @@ class UploadState(Enum):
     CANCELED = 3
 
 
+def _get_upload_credentials_with_retry(client, pbar, filename, mission_id, max_attempts=5):
+    """
+    Retrieves upload credentials with retry logic.
+
+    Args:
+        client: The client object used for retrieving credentials.
+        filename: The internal filename.
+        mission_id: The mission ID.
+        max_attempts: Maximum number of retry attempts.
+
+    Returns:
+        The upload credentials or None if retrieval fails after all attempts.
+    """
+    attempt = 0
+    while attempt < max_attempts:
+        creds = _get_upload_creditials(client, internal_filename=filename, mission_id=mission_id)
+        if creds is not None:
+            return creds
+
+        attempt += 1
+        if attempt < max_attempts:
+            delay = 2**attempt  # Exponential backoff (2, 4, 8, 16...)
+            sleep(delay)
+
+    return None
+
+
 # TODO: i dont want to handle errors at this level
 def upload_file(
     client: AuthenticatedClient,
@@ -161,79 +198,119 @@ def upload_file(
     verbose: bool = False,
     s3_endpoint: Optional[str] = None,
 ) -> Tuple[UploadState, int]:
-    """\
+    """
     returns UploadState and bytes uploaded (0 if not uploaded)
+    Retries up to 3 times on failure.
     """
     if s3_endpoint is None:
         s3_endpoint = get_config().endpoint.s3
 
     total_size = path.stat().st_size
-    with tqdm(
-        total=total_size,
-        unit="B",
-        unit_scale=True,
-        desc=f"uploading {path}...",
-        leave=False,
-        disable=not verbose,
-    ) as pbar:
-        # get per file upload credentials
-        creds = _get_upload_creditials(
-            client, internal_filename=filename, mission_id=mission_id
-        )
-        if creds is None:
-            return UploadState.EXISTS, 0
+    for attempt in range(MAX_UPLOAD_RETRIES):
+        with tqdm(
+            total=total_size,
+            unit="B",
+            unit_scale=True,
+            desc=f"uploading {path}...",
+            leave=False,
+            disable=not verbose,
+        ) as pbar:
+
+            # get per file upload credentials
+            creds = _get_upload_credentials_with_retry(
+                client, pbar, filename, mission_id, max_attempts=5 if attempt > 0 else 1
+            )
+
+            if creds is None:
+                return UploadState.EXISTS, 0
 
-        try:
-            _s3_upload(path, endpoint=s3_endpoint, credentials=creds, pbar=pbar)
-        except Exception as e:
-            logger.error(format_traceback(e))
             try:
-                _cancel_file_upload(client, creds.file_id, mission_id)
-            except Exception as cancel_e:
-                logger.error(f"Failed to cancel upload for {creds.file_id}: {cancel_e}")
-            raise e from e
-        else:
-            _confirm_file_upload(client, creds.file_id, b64_md5(path))
-            return UploadState.UPLOADED, total_size
+                _s3_upload(path, endpoint=s3_endpoint, credentials=creds, pbar=pbar)
+            except Exception as e:
+                logger.error(format_traceback(e))
+                try:
+                    _cancel_file_upload(client, creds.file_id, mission_id)
+                except Exception as cancel_e:
+                    logger.error(f"Failed to cancel upload for {creds.file_id}: {cancel_e}")
+
+                if attempt < 2:  # Retry if not the last attempt
+                    pbar.update(0)
+                    logger.error(f"Retrying upload for {attempt + 1}")
+                    continue
+                else:
+                    logger.error(f"Cancelling upload for {attempt}")
+                    raise e from e
+
+            else:
+                _confirm_file_upload(client, creds.file_id, b64_md5(path))
+                return UploadState.UPLOADED, total_size
 
 
 def _get_file_download(client: AuthenticatedClient, id: UUID) -> str:
     """\
     get the download url for a file by file id
     """
-    resp = client.get(DOWNLOAD_URL, params={"uuid": str(id), "expires": True})
+    resp = client.get(DOWNLOAD_URL, params={"uuid": str(id), "expires": True, "preview_only": False})
 
     if 400 <= resp.status_code < 500:
         raise AccessDenied(
-            f"Failed to download file: {resp.json()['message']}"
-            f" Status Code: {resp.status_code}",
+            f"Failed to download file: {resp.json()['message']}" f" Status Code: {resp.status_code}",
         )
 
     resp.raise_for_status()
 
-    return resp.text
+    return resp.json()["url"]
 
 
-def _url_download(
-    url: str, *, path: Path, size: int, overwrite: bool = False, verbose: bool = False
-) -> None:
-    if path.exists() and not overwrite:
-        raise FileExistsError(f"file already exists: {path}")
-
-    with httpx.stream("GET", url) as response:
-        response.raise_for_status()
-        with open(path, "wb") as f:
-            with tqdm(
-                total=size,
-                desc=f"downloading {path.name}",
-                unit="B",
-                unit_scale=True,
-                leave=False,
-                disable=not verbose,
-            ) as pbar:
-                for chunk in response.iter_bytes(chunk_size=DOWNLOAD_CHUNK_SIZE):
-                    f.write(chunk)
-                    pbar.update(len(chunk))
+def _url_download(url: str, *, path: Path, size: int, overwrite: bool = False, verbose: bool = False) -> None:
+    if path.exists():
+        if overwrite:
+            path.unlink()
+            downloaded = 0
+        else:
+            downloaded = path.stat().st_size
+            if downloaded >= size:
+                raise FileExistsError(f"file already exists and is complete: {path}")
+    else:
+        downloaded = 0
+
+    attempt = 0
+    while downloaded < size:
+        try:
+            headers = {"Range": f"bytes={downloaded}-"}
+            with httpx.stream("GET", url, headers=headers, timeout=S3_READ_TIMEOUT) as response:
+                # Accept both 206 Partial Content and 200 OK if starting from 0
+                if not (response.status_code == 206 or (downloaded == 0 and response.status_code == 200)):
+                    response.raise_for_status()
+                    raise RuntimeError(f"Expected 206 Partial Content, got {response.status_code}")
+
+                mode = "ab" if downloaded > 0 else "wb"
+                with open(path, mode) as f:
+                    with tqdm(
+                        total=size,
+                        initial=downloaded,
+                        desc=f"downloading {path.name}",
+                        unit="B",
+                        unit_scale=True,
+                        leave=False,
+                        disable=not verbose,
+                    ) as pbar:
+                        for chunk in response.iter_bytes(chunk_size=DOWNLOAD_CHUNK_SIZE):
+                            attempt = 0  # reset attempt counter on successful download of non-empty chunk
+                            if not chunk:
+                                break
+                            f.write(chunk)
+                            downloaded += len(chunk)
+                            pbar.update(len(chunk))
+            break  # download complete
+        except Exception as e:
+            logger.info(f"Error: {e}, retrying...")
+            attempt += 1
+            if attempt > MAX_RETRIES:
+                raise RuntimeError(f"Download failed after {MAX_RETRIES} retries due to {e}") from e
+            if verbose:
+                print(f"{e} on attempt {attempt}/{MAX_RETRIES}, retrying after backoff...")
+            sleep(RETRY_BACKOFF_BASE**attempt)
 
 
 class DownloadState(Enum):
@@ -242,6 +319,7 @@ class DownloadState(Enum):
     DOWNLOADED_INVALID_HASH = 3
     SKIPPED_INVALID_HASH = 4
     SKIPPED_INVALID_REMOTE_STATE = 5
+    SKIPPED_FILE_SIZE_MISMATCH = 6
 
 
 def download_file(
@@ -260,18 +338,24 @@ def download_file(
         return DownloadState.SKIPPED_INVALID_REMOTE_STATE, 0
 
     if path.exists():
-        local_hash = b64_md5(path)
-        if local_hash != file.hash and not overwrite and file.hash is not None:
-            return DownloadState.SKIPPED_INVALID_HASH
 
-        elif local_hash == file.hash:
-            return DownloadState.SKIPPED_OK
+        # compare file size
+        if file.size == path.stat().st_size:
+            local_hash = b64_md5(path)
+            if local_hash != file.hash and not overwrite and file.hash is not None:
+                return DownloadState.SKIPPED_INVALID_HASH, 0
 
-        # this has to be here
-        if verbose:
-            tqdm.write(
-                styled_string(f"overwriting {path}, hash missmatch", style="yellow")
-            )
+            elif local_hash == file.hash:
+                return DownloadState.SKIPPED_OK, 0
+
+            elif verbose:
+                tqdm.write(styled_string(f"overwriting {path}, hash mismatch", style="yellow"))
+
+        elif not overwrite and file.size is not None:
+            return DownloadState.SKIPPED_FILE_SIZE_MISMATCH, 0
+
+        elif verbose:
+            tqdm.write(styled_string(f"overwriting {path}, file size mismatch", style="yellow"))
 
     # request a download url
     download_url = _get_file_download(client, file.id)
@@ -301,6 +385,10 @@ def download_file(
 
     observed_hash = b64_md5(path)
     if file.hash is not None and observed_hash != file.hash:
+        print(
+            f"HASH MISMATCH: {path} expected={file.hash} observed={observed_hash}",
+            file=sys.stderr,
+        )
         # Download completed but hash failed
         return (
             DownloadState.DOWNLOADED_INVALID_HASH,
@@ -317,9 +405,7 @@ UPLOAD_STATE_COLOR = {
 }
 
 
-def _upload_handler(
-    future: Future[Tuple[UploadState, int]], path: Path, *, verbose: bool = False
-) -> int:
+def _upload_handler(future: Future[Tuple[UploadState, int]], path: Path, *, verbose: bool = False) -> int:
     """Returns bytes uploaded successfully."""
     state = UploadState.CANCELED  # Default to canceled if exception occurs
     size_bytes = 0
@@ -328,7 +414,7 @@ def _upload_handler(
     except Exception as e:
         logger.error(format_traceback(e))
         if verbose:
-            tqdm.write(format_error(f"error uploading", e, verbose=verbose))
+            tqdm.write(format_error("error uploading", e, verbose=verbose))
         else:
             print(f"ERROR: {path.absolute()}: {e}", file=sys.stderr)
         return 0  # Return 0 bytes on error
@@ -353,6 +439,7 @@ DOWNLOAD_STATE_COLOR = {
     DownloadState.SKIPPED_OK: "green",
     DownloadState.DOWNLOADED_INVALID_HASH: "red",
     DownloadState.SKIPPED_INVALID_HASH: "yellow",
+    DownloadState.SKIPPED_FILE_SIZE_MISMATCH: "yellow",
     DownloadState.SKIPPED_INVALID_REMOTE_STATE: "purple",
 }
 
@@ -385,6 +472,8 @@ def _download_handler(
         msg = f"skipped {path} already downloaded (hash ok)"
     elif state == DownloadState.SKIPPED_INVALID_HASH:
         msg = f"skipped {path}, exists with hash mismatch (use --overwrite?)"
+    elif state == DownloadState.SKIPPED_FILE_SIZE_MISMATCH:
+        msg = f"skipped {path}, exists with file size mismatch (use --overwrite?)"
     elif state == DownloadState.SKIPPED_INVALID_REMOTE_STATE:
         msg = f"skipped {path}, remote file has invalid state ({file.state.value})"
     else:
@@ -395,11 +484,7 @@ def _download_handler(
     elif state not in (DownloadState.DOWNLOADED_OK, DownloadState.SKIPPED_OK):
         print(f"SKIP/FAIL: {path.absolute()} ({state.name})", file=sys.stderr)
 
-    return (
-        size_bytes
-        if state in (DownloadState.DOWNLOADED_OK, DownloadState.SKIPPED_OK)
-        else 0
-    )
+    return size_bytes if state in (DownloadState.DOWNLOADED_OK, DownloadState.SKIPPED_OK) else 0
 
 
 def upload_files(
@@ -416,16 +501,17 @@ def upload_files(
         unit="files",
         desc="Uploading files",
         disable=not verbose,
-        leave=False,
+        leave=True,
     ) as pbar:
         start = monotonic()
         futures: Dict[Future[Tuple[UploadState, int]], Path] = {}
+
+        skipped_files = 0
+        failed_files = 0
         with ThreadPoolExecutor(max_workers=n_workers) as executor:
             for name, path in files.items():
                 if not path.is_file():
-                    console.print(
-                        f"[yellow]Skipping non-existent file: {path}[/yellow]"
-                    )
+                    console.print(f"[yellow]Skipping non-existent file: {path}[/yellow]")
                     pbar.update()
                     continue
 
@@ -441,20 +527,38 @@ def upload_files(
 
             total_uploaded_bytes = 0
             for future in as_completed(futures):
+
+                if future.exception():
+                    failed_files += 1
+
+                if future.exception() is None and future.result()[0] == UploadState.EXISTS:
+                    skipped_files += 1
+
                 path = futures[future]
                 uploaded_bytes = _upload_handler(future, path, verbose=verbose)
                 total_uploaded_bytes += uploaded_bytes
                 pbar.update()
 
-        end = monotonic()
-        elapsed_time = end - start
+    end = monotonic()
+    elapsed_time = end - start
 
-        avg_speed_bps = total_uploaded_bytes / elapsed_time if elapsed_time > 0 else 0
+    avg_speed_bps = total_uploaded_bytes / elapsed_time if elapsed_time > 0 else 0
 
+    if verbose:
+        console.print()
         console.print(f"Upload took {elapsed_time:.2f} seconds")
         console.print(f"Total uploaded: {format_bytes(total_uploaded_bytes)}")
         console.print(f"Average speed: {format_bytes(avg_speed_bps, speed=True)}")
 
+        if failed_files > 0:
+            console.print(
+                f"\nUploaded {len(files) - failed_files - skipped_files} files, "
+                f"{skipped_files} skipped, {failed_files} uploads failed",
+                style="red",
+            )
+        else:
+            console.print(f"\nUploaded {len(files) - skipped_files} files, {skipped_files} skipped")
+
 
 def download_files(
     client: AuthenticatedClient,
@@ -470,7 +574,7 @@ def download_files(
         unit="files",
         desc="Downloading files",
         disable=not verbose,
-        leave=False,
+        leave=True,
     ) as pbar:
 
         start = monotonic()
@@ -490,18 +594,15 @@ def download_files(
             total_downloaded_bytes = 0
             for future in as_completed(futures):
                 file, path = futures[future]
-                downloaded_bytes = _download_handler(
-                    future, file, path, verbose=verbose
-                )
+                downloaded_bytes = _download_handler(future, file, path, verbose=verbose)
                 total_downloaded_bytes += downloaded_bytes
                 pbar.update()
 
-        end = monotonic()
-        elapsed_time = end - start
-        avg_speed_bps = total_downloaded_bytes / elapsed_time if elapsed_time > 0 else 0
+    end = monotonic()
+    elapsed_time = end - start
+    avg_speed_bps = total_downloaded_bytes / elapsed_time if elapsed_time > 0 else 0
 
-        console.print(f"Download took {elapsed_time:.2f} seconds")
-        console.print(
-            f"Total downloaded/verified: {format_bytes(total_downloaded_bytes)}"
-        )
-        console.print(f"Average speed: {format_bytes(avg_speed_bps, speed=True)}")
+    console.print()
+    console.print(f"Download took {elapsed_time:.2f} seconds")
+    console.print(f"Total downloaded/verified: {format_bytes(total_downloaded_bytes)}")
+    console.print(f"Average speed: {format_bytes(avg_speed_bps, speed=True)}")

kleinkram/api/pagination.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from enum import Enum
 from typing import Any
 from typing import Dict
 from typing import Generator
@@ -17,6 +16,7 @@ DataPage = Dict[str, Any]
 PAGE_SIZE = 128
 SKIP = "skip"
 TAKE = "take"
+EXACT_MATCH = "exactMatch"
 
 
 def paginated_request(
@@ -25,6 +25,7 @@ def paginated_request(
     params: Optional[Mapping[str, Any]] = None,
     max_entries: Optional[int] = None,
     page_size: int = PAGE_SIZE,
+    exact_match: bool = False,
 ) -> Generator[DataPage, None, None]:
     total_entries_count = 0
 
@@ -32,10 +33,18 @@ def paginated_request(
 
     params[TAKE] = page_size
     params[SKIP] = 0
+    if exact_match:
+        params[EXACT_MATCH] = str(exact_match).lower()  # pass string rather than bool
 
     while True:
         resp = client.get(endpoint, params=params)
-        resp.raise_for_status()  # TODO: this is fine for now
+
+        # explicitly handle 404 if json contains message
+        if resp.status_code == 404 and "message" in resp.json():
+            raise ValueError(resp.json()["message"])
+
+        # raise for other errors
+        resp.raise_for_status()
 
         paged_data = resp.json()
         data_page = cast(List[DataPage], paged_data["data"])

kleinkram/api/query.py

@@ -41,6 +41,14 @@ class FileQuery:
     mission_query: MissionQuery = field(default_factory=MissionQuery)
 
 
+@dataclass
+class RunQuery:
+    mission_ids: List[UUID] = field(default_factory=list)
+    mission_patterns: List[str] = field(default_factory=list)
+    project_ids: List[UUID] = field(default_factory=list)
+    project_patterns: List[str] = field(default_factory=list)
+
+
 def check_mission_query_is_creatable(query: MissionQuery) -> str:
     """\
     check if a query is unique and can be used to create a mission
@@ -87,11 +95,7 @@ def mission_query_is_unique(query: MissionQuery) -> bool:
         return True
 
     # a single mission name a unique project spec are specified
-    if (
-        project_query_is_unique(query.project_query)
-        and len(query.patterns) == 1
-        and _pattern_is_unique(query.patterns[0])
-    ):
+    if project_query_is_unique(query.project_query) and len(query.patterns) == 1 and _pattern_is_unique(query.patterns[0]):
         return True
     return False
 
@@ -102,10 +106,6 @@ def file_query_is_unique(query: FileQuery) -> bool:
         return True
 
     # a single file name a unique mission spec are specified
-    if (
-        mission_query_is_unique(query.mission_query)
-        and len(query.patterns) == 1
-        and _pattern_is_unique(query.patterns[0])
-    ):
+    if mission_query_is_unique(query.mission_query) and len(query.patterns) == 1 and _pattern_is_unique(query.patterns[0]):
         return True
     return False