diff --git a/.gitignore b/.gitignore index 81f08936..25aa99bd 100644 --- a/.gitignore +++ b/.gitignore @@ -173,4 +173,5 @@ data/ bioimageio_test_reports/ # MacOS -.DS_STORE \ No newline at end of file +.DS_STORE +demo-dataset/ \ No newline at end of file diff --git a/bioengine_apps/bia_resolve_url_proxy/README.md b/bioengine_apps/bia_resolve_url_proxy/README.md new file mode 100644 index 00000000..54e25671 --- /dev/null +++ b/bioengine_apps/bia_resolve_url_proxy/README.md @@ -0,0 +1,5 @@ +# BIA Resolve URL Proxy + +Minimal ray-serve app that exposes a single `resolve_url` RPC method. + +It is used by `cellpose_finetuning` to fetch BioImage Archive assets through Hypha when direct fetches are blocked. diff --git a/bioengine_apps/bia_resolve_url_proxy/main.py b/bioengine_apps/bia_resolve_url_proxy/main.py new file mode 100644 index 00000000..b1537666 --- /dev/null +++ b/bioengine_apps/bia_resolve_url_proxy/main.py @@ -0,0 +1,136 @@ +import base64 +import os +from typing import Any +from urllib.parse import urlparse + +import httpx +from hypha_rpc.utils.schema import schema_method +from ray import serve + +_DEFAULT_ALLOWED_HOSTS = "beta.bioimagearchive.org,www.ebi.ac.uk,ftp.ebi.ac.uk" + + +def _allowed_hosts() -> set[str]: + raw_value = os.environ.get("RESOLVE_URL_ALLOWED_HOSTS", _DEFAULT_ALLOWED_HOSTS) + return {part.strip().lower() for part in raw_value.split(",") if part.strip()} + + +def _normalize_headers(headers: dict[str, Any] | None) -> dict[str, str]: + normalized: dict[str, str] = {} + if not isinstance(headers, dict): + return normalized + for key, value in headers.items(): + if not isinstance(key, str): + continue + if isinstance(value, str): + normalized[key] = value + elif value is not None: + normalized[key] = str(value) + return normalized + + +@serve.deployment( + ray_actor_options={ + "num_cpus": 0.5, + "num_gpus": 0, + "runtime_env": { + "pip": [ + "httpx", + ], + }, + } +) +class ResolveUrlProxy: + @schema_method + async def resolve_url( + self, + url: str, + method: str = "GET", + headers: dict[str, Any] | None = None, + timeout: float = 60.0, + body: str | dict[str, Any] | list[Any] | None = None, + context: dict[str, Any] | None = None, + ) -> dict[str, Any]: + try: + parsed = urlparse(str(url)) + except Exception as exp: + return { + "ok": False, + "status_code": 400, + "url": str(url), + "error": f"Invalid URL: {exp}", + } + + if parsed.scheme.lower() != "https": + return { + "ok": False, + "status_code": 400, + "url": str(url), + "error": "Only https URLs are allowed", + } + + host = (parsed.hostname or "").lower() + if host not in _allowed_hosts(): + return { + "ok": False, + "status_code": 403, + "url": str(url), + "error": f"Host '{host}' is not allowed", + } + + method_value = str(method or "GET").upper() + if method_value not in { + "GET", + "POST", + "PUT", + "PATCH", + "DELETE", + "HEAD", + "OPTIONS", + }: + return { + "ok": False, + "status_code": 400, + "url": str(url), + "error": f"Unsupported method: {method_value}", + } + + request_kwargs: dict[str, Any] = { + "method": method_value, + "url": str(url), + "headers": _normalize_headers(headers), + } + request_kwargs["headers"].setdefault( + "User-Agent", "bioengine-bia-resolve-url-proxy/1.0" + ) + + if body is not None: + if isinstance(body, (dict, list)): + request_kwargs["json"] = body + else: + request_kwargs["content"] = str(body) + + try: + async with httpx.AsyncClient( + timeout=max(1.0, float(timeout)), + follow_redirects=True, + ) as client: + response = await client.request(**request_kwargs) + except Exception as exp: + return { + "ok": False, + "status_code": 502, + "url": str(url), + "error": str(exp), + } + + status_code = int(response.status_code) + ok = 200 <= status_code < 300 + return { + "ok": ok, + "status_code": status_code, + "url": str(response.url), + "headers": dict(response.headers), + "content_base64": base64.b64encode(response.content).decode("ascii"), + "error": "" if ok else f"Upstream returned HTTP {status_code}", + } diff --git a/bioengine_apps/bia_resolve_url_proxy/manifest.yaml b/bioengine_apps/bia_resolve_url_proxy/manifest.yaml new file mode 100644 index 00000000..9dc5269e --- /dev/null +++ b/bioengine_apps/bia_resolve_url_proxy/manifest.yaml @@ -0,0 +1,16 @@ +name: BIA Resolve URL Proxy +id: bia-resolve-url-proxy +id_emoji: '🌐' +covers: [] +description: Resolve URL proxy for BioImage Archive downloads (resolve_url only). +documentation: README.md +format_version: 0.5.0 +git_repo: https://github.com/aicell-lab/bioengine-worker +links: [] +tags: [bioimage-archive, proxy, resolve-url] +type: ray-serve +version: 0.0.1 +deployments: + - main:ResolveUrlProxy +authorized_users: + - "*" diff --git a/bioengine_apps/cellpose_finetuning/README.md b/bioengine_apps/cellpose_finetuning/README.md index 1bc52da7..05b36f12 100644 --- a/bioengine_apps/cellpose_finetuning/README.md +++ b/bioengine_apps/cellpose_finetuning/README.md @@ -191,9 +191,11 @@ Start asynchronous model fine-tuning. - `metadata_dir` (str, optional): Folder containing metadata JSON files with image/annotation paths (e.g. `image_path` + `mask_path`). If provided, explicit `train_annotations` can be omitted. - `test_images` (str, optional): Optional test images path (same format as train_images). Providing test data enables per-epoch pixel-level validation metrics and end-of-training instance segmentation metrics (AP@0.5/0.75/0.9). - `test_annotations` (str, optional): Optional test annotations (label masks). Required together with `test_images`. +- `split_mode` (str): Dataset split mode (`"manual"` or `"auto"`, default: `"manual"`). In `"auto"`, the service creates train/validation split from the selected sample pool. +- `train_split_ratio` (float): Train ratio used when `split_mode="auto"` (default: `0.8`, i.e. 80% train / 20% validation). - `model` (str): Pretrained model to start from (default: "cpsam") - `n_epochs` (int): Number of training epochs (default: 10) -- `n_samples` (int, optional): Limit number of samples to use +- `n_samples` (int | float, optional): Limit sample usage. If `n_samples > 1`, treated as absolute count. If `0 < n_samples <= 1`, treated as decimal fraction of available samples (e.g. `0.05` = 5%). - `learning_rate` (float): Learning rate (default: 1e-6) - `weight_decay` (float): Weight decay (default: 0.0001) - `min_train_masks` (int): Minimum number of masks per training batch (default: 5). Lower values speed up training. @@ -376,6 +378,31 @@ Notes: - Sessions with stale `waiting`/`preparing`/`running` states after a redeploy are normalized to `stopped` (interrupted). - Running-session indicators in the UI only track sessions that are currently active. +### `delete_training_session()` + +Delete a non-running training session and its local artifacts. + +**Parameters:** +- `session_id` (str): Session ID to delete +- `force_stop_if_blocked` (bool, optional): If true, stop a blocked `running`/`preparing` session first, then delete + +**Returns:** Dict with `deleted` (bool) and `session_id` (str) + +### `preflight_training_dataset()` + +Run a lightweight backend preflight over training path configuration without starting training. + +**Parameters:** +- `artifact` (str): Artifact ID or BioImage Archive URL +- `train_images`, `train_annotations` (str, optional): Training path specs +- `metadata_dir` (str, optional): Metadata mode directory +- `test_images`, `test_annotations` (str, optional): Optional test path specs +- `split_mode` (str): `manual` or `auto` +- `n_samples` (int | float, optional): Optional sample limit/count +- `max_candidates` (int): Cap for path matching candidates during preflight + +**Returns:** Dict including `ok`, pair counts (`train_pair_count`, `test_pair_count`), and an explanatory `message`. + Validation metrics are only produced when both `test_images` and `test_annotations` are provided. ## Inference Parameters Guide @@ -430,7 +457,7 @@ The recommended way to deploy (whether for the first time or to update) is to us source .env # Deploy to production -python tests/cellpose_legacy_scripts/redeploy_cellpose.py \ +python tests/cellpose/redeploy_cellpose.py \ --artifact-id ri-scale/cellpose-finetuning \ --application-id cellpose-finetuning ``` diff --git a/bioengine_apps/cellpose_finetuning/index.html b/bioengine_apps/cellpose_finetuning/index.html index a1e84b85..97622be9 100644 --- a/bioengine_apps/cellpose_finetuning/index.html +++ b/bioengine_apps/cellpose_finetuning/index.html @@ -53,7 +53,7 @@

Cellpose Fine-Tuning<
Connected Disconnected - • {{ workspace }} + • [[ workspace ]]
@@ -86,7 +86,7 @@

Cellpose Fine-Tuning< @@ -128,7 +128,7 @@

Cellpose Fine-Tuning< @click="currentTab = tab.id" :class="currentTab === tab.id ? 'bg-blue-600 text-white shadow-lg shadow-blue-900/50' : 'hover:bg-slate-800 hover:text-white'" class="w-full text-left px-4 py-3 rounded-lg flex items-center gap-3 transition-all duration-200 font-medium"> - {{ tab.name }} + [[ tab.name ]] @@ -141,7 +141,7 @@

Cellpose Fine-Tuning<

- {{ monitoredSession.status_type === 'running' ? 'Active Training' : 'Training Status' }} + [[ monitoredSession.status_type === 'running' ? 'Active Training' : 'Training Status' ]]

- {{ monitoringSessionId.slice(0, 8) }}... + [[ monitoringSessionId.slice(0, 8) ]]... 'bg-red-900 text-red-300': monitoredSession.status_type === 'failed', 'bg-gray-700 text-gray-300': monitoredSession.status_type === 'stopped' || monitoredSession.status_type === 'waiting' }"> - {{ monitoredSession.status_type || 'Run' }} + [[ monitoredSession.status_type || 'Run' ]]
@@ -171,12 +171,12 @@

:style="{ width: ((monitoredSession.current_epoch || 0) / (monitoredSession.total_epochs || 1) * 100) + '%' }">

- Epoch {{ monitoredSession.current_epoch || 0 }} - {{ monitoredSession.total_epochs || '?' }} Total + Epoch [[ monitoredSession.current_epoch || 0 ]] + [[ monitoredSession.total_epochs || '?' ]] Total
- {{ monitoredSession.message }} + [[ monitoredSession.message ]]
@@ -195,12 +195,12 @@

Welcome to Cellpose Fine-Tunin

Select a Cellpose Fine-Tuning Service to continue.

- No services found in workspace "{{ workspace }}". + No services found in workspace "[[ workspace ]]".
Make sure the Worker is running and you are logged into the correct workspace.
- {{ autoConnectError }} + [[ autoConnectError ]]
+ -
+ @@ -235,9 +255,20 @@

Dashboard

+ -
+ + Session ID Status Progress
+ +
- {{ id.slice(0, 8) }} + [[ id.slice(0, 8) ]]
@@ -248,12 +279,12 @@

Dashboard

'bg-red-100 text-red-800': session.status_type === 'failed', 'bg-gray-100 text-gray-800': session.status_type === 'waiting' || session.status_type === 'stopped' }"> - {{ session.status_type }} + [[ session.status_type ]]
-
{{ session.current_epoch || 0 }}{{ session.total_epochs }}
+
[[ session.current_epoch || 0 ]][[ session.total_epochs ]]
@@ -261,18 +292,42 @@

Dashboard

-
- {{ getLatestLoss(session) }} + [[ getLatestLoss(session) ]] - {{ formatDate(session.start_time) }} + [[ formatDate(session.start_time) ]] - - + + +
+
📭

No training sessions found.

@@ -295,62 +350,258 @@

New Training Session

Dataset & Model

-
-
+
+ +
+ + + +
+
+ +
-
- - + + +
+
+ [[ suggestion ]] + Tab +
+
+
+

[[ artifactInputError ]]

+

Checking artifact…

+

Artifact found

+
+
+ +
+ +
+
+ + [[ trainingParams.uploaded_artifact ]] +
+
+ No dataset uploaded yet +
+ +
+

Use Browse Files to pick image/annotation folders, same as Artifact source.

+
+ +
+ +
+ + + Search EBI + +
+ +
-
- -
- cpsam - cyto3 -
-
+
-
+
- +
Examples: +
- + +
+
+
+ + Auto-detect paths by scanning this artifact +
+ +
+
+ [[ pathAutoError ]] +
+
+ + [[ pathValidationMessage ]]
- +

Use this to load image/annotation pairs from metadata JSON files; then explicit annotation masks are optional.

+
+ + Path settings changed. Click Validate paths to unlock split and sample controls. +
+
+
+ BioImage Archive mode discovers image/mask pairs from the archive URL automatically. +
+ +
- - + +
+ +
+
+ +
+ + % +
+
+
+ + +
+
+
+

+ Selected samples: [[ selectedTotalSampleCount ?? '-' ]] + / [[ availableTotalSampleCount ?? '-' ]] total +

+
- - + +
+ + + BioImage Archive uses auto split +
+ +
+
+ +
+
+ +
+ + % +
+
+
+ + +
+
+
+

Estimated split from selected samples: train [[ selectedTrainSampleCount ?? '-' ]], validation [[ selectedValidationSampleCount ?? '-' ]].

+
+ +
+
+ + +
+
+ + +
+

Estimated split from selected samples: train [[ selectedTrainSampleCount ?? '-' ]], validation [[ selectedValidationSampleCount ?? '-' ]].

+
+
@@ -373,15 +624,19 @@

- +

- {{ trainingError }} + [[ trainingError ]]
@@ -396,7 +651,7 @@

Session Analysis - {{ selectedSessionId }} + [[ selectedSessionId ]]

@@ -408,44 +663,51 @@

Stop +

Status
-
{{ currentSession?.status_type }}
+
[[ currentSession?.status_type ]]
Progress
-
{{ currentSession?.current_epoch || 0 }} / {{ currentSession?.total_epochs }} Epochs
+
[[ currentSession?.current_epoch || 0 ]] / [[ currentSession?.total_epochs ]] Epochs
Samples
-
{{ currentSession?.n_train || '-' }} Train
+
[[ currentSession?.n_train || '-' ]] Train
Duration
-
{{ formatDuration(currentSession?.elapsed_seconds) }}
+
[[ formatDuration(currentSession?.elapsed_seconds) ]]

Training Hyperparameters

-
Model
{{ currentSession?.model || '-' }}
-
Epochs
{{ currentSession?.n_epochs ?? currentSession?.total_epochs ?? '-' }}
-
Samples
{{ currentSession?.n_samples ?? 'all' }}
-
Learning Rate
{{ currentSession?.learning_rate ?? '-' }}
-
Weight Decay
{{ currentSession?.weight_decay ?? '-' }}
-
Min Train Masks
{{ currentSession?.min_train_masks ?? '-' }}
-
Validation Interval
{{ currentSession?.validation_interval ?? 'default (10)' }}
+
Model
[[ currentSession?.model || '-' ]]
+
Epochs
[[ currentSession?.n_epochs ?? currentSession?.total_epochs ?? '-' ]]
+
Samples
[[ currentSession?.n_samples ?? 'all' ]]
+
Learning Rate
[[ currentSession?.learning_rate ?? '-' ]]
+
Weight Decay
[[ currentSession?.weight_decay ?? '-' ]]
+
Min Train Masks
[[ currentSession?.min_train_masks ?? '-' ]]
+
Validation Interval
[[ currentSession?.validation_interval ?? 'default (10)' ]]
@@ -470,11 +732,11 @@

Validation Metrics

SYSTEM LOG - {{ currentSession?.status_type === 'running' ? 'LIVE' : 'OFFLINE' }} + [[ currentSession?.status_type === 'running' ? 'LIVE' : 'OFFLINE' ]]
-

Session ID: {{ selectedSessionId }}

-

> {{ currentSession?.message }}

+

Session ID: [[ selectedSessionId ]]

+

> [[ currentSession?.message ]]

Error: Process terminated unexpectedly.

@@ -490,23 +752,23 @@

Ready for Export

Export Success
- {{ exportResult.model_name }} + [[ exportResult.model_name ]]
- -
+ +

Live Inference

- -
+ +
@@ -518,11 +780,11 @@

Configuration

@@ -538,7 +800,7 @@

Configuration

- {{ runningTrainingSessionsCount }} training session(s) are running. Inference may fail with GPU memory errors while training is active. + [[ runningTrainingSessionsCount ]] training session(s) are running. Inference may fail with GPU memory errors while training is active.
@@ -550,7 +812,7 @@

Configuration

Click or Drag Image Here

- {{ inferenceFile.name }} + [[ inferenceFile.name ]]
@@ -559,7 +821,7 @@

Configuration

@@ -579,11 +841,24 @@

Configuration

- Found {{ inferenceObjectCount }} objects + Found [[ inferenceObjectCount ]] objects
+
+
+ +
+

Initializing inference service

+

+ The worker connection is active, but the inference backend is not ready yet. +

+ +
@@ -596,7 +871,7 @@

Configuration

Artifact Explorer

-

{{ trainingParams.artifact || 'No artifact selected' }}

+

[[ getSelectedArtifactId() || 'No artifact selected' ]]

@@ -604,7 +879,7 @@

Artifact Explorer

/ - {{ browserPath }} + [[ browserPath ]]
@@ -612,7 +887,7 @@

Artifact Explorer


Loading files...
-
{{ browserError }} +
[[ browserError ]]
  • @@ -624,7 +899,7 @@

    Artifact Explorer

    - {{ item.name }} + [[ item.name ]]
    @@ -640,10 +915,10 @@

    Artifact Explorer

    - Current: /{{ browserPath }} + Current: /[[ browserPath ]] - {{ browserFeedback }} + [[ browserFeedback ]]
    @@ -668,19 +943,30 @@

    Upload Dataset

    - + + +

    + + + [[ selectedFolderName ]] ([[ selectedFileCount ]] files) + + No folder selected +

    Folder should contain images and mask files.

    Uploading... - {{ uploadProgress }}% + [[ uploadProgress ]]%
    -
    {{ uploadStatus }}
    +
    [[ uploadStatus ]]
    @@ -694,9 +980,9 @@

    Upload Dataset

    \ No newline at end of file diff --git a/bioengine_apps/cellpose_finetuning/main.py b/bioengine_apps/cellpose_finetuning/main.py index ba6a4442..b1832b8f 100644 --- a/bioengine_apps/cellpose_finetuning/main.py +++ b/bioengine_apps/cellpose_finetuning/main.py @@ -8,7 +8,9 @@ import asyncio import base64 +import csv import fnmatch +import hashlib import io import json import logging @@ -47,6 +49,13 @@ TRAINING_PARAMS_FILENAME = "training_params.json" STOP_REQUESTED_FILENAME = "stop.requested" STATUS_STALE_SECONDS = 300 +BIA_FTS_ENDPOINT = "https://beta.bioimagearchive.org/search/search/fts" +BIA_IMAGE_ENDPOINT = "https://beta.bioimagearchive.org/search/search/fts/image" +MIN_FREE_GPU_MEMORY_TO_START_BYTES = 2 * GB +BIA_RESOLVE_URL_SERVICE_ID = os.environ.get("BIA_RESOLVE_URL_SERVICE_ID") +BIA_RESOLVE_URL_APPLICATION_ID = os.environ.get( + "BIA_RESOLVE_URL_APPLICATION_ID", "bia-resolve-url-proxy" +) # Model template for BioImage.io export MODEL_TEMPLATE_PY = '''"""BioImage.io Model Wrapper for Cellpose 4.0.7 (Cellpose-SAM). @@ -256,12 +265,14 @@ class TrainingParams(TypedDict): metadata_dir: str | None test_images: str | None test_annotations: str | None + split_mode: str + train_split_ratio: float model: str | Path n_epochs: int learning_rate: float weight_decay: float server_url: str - n_samples: int | None + n_samples: int | float | None session_id: str min_train_masks: int validation_interval: int | None @@ -542,7 +553,7 @@ class SessionStatus(TypedDict, total=False): exported_artifact_id: str # Artifact ID if model has been exported model_modified: bool # Flag indicating if model was modified since last export model: str - n_samples: int | None + n_samples: int | float | None n_epochs: int learning_rate: float weight_decay: float @@ -576,7 +587,7 @@ class SessionStatusWithId(TypedDict, total=False): exported_artifact_id: str model_modified: bool model: str - n_samples: int | None + n_samples: int | float | None n_epochs: int learning_rate: float weight_decay: float @@ -634,6 +645,9 @@ def get_url_and_artifact_id(artifact_id: str | Any) -> tuple[str, str]: parsed = urlparse(artifact_id) if parsed.scheme in ("http", "https"): + if _is_bioimage_archive_url(artifact_id): + return DEFAULT_SERVER_URL, artifact_id + path_parts = parsed.path.lstrip("/").split("/") if path_parts[1] != "artifacts": msg = ( @@ -760,7 +774,7 @@ def update_status( current_batch: int | None = None, total_batches: int | None = None, model: str | None = None, - n_samples: int | None = None, + n_samples: int | float | None = None, n_epochs: int | None = None, learning_rate: float | None = None, weight_decay: float | None = None, @@ -1103,6 +1117,12 @@ def train_seg_with_callbacks( nimg_test_val if nimg_test_per_epoch is None else nimg_test_per_epoch ) + if nimg_per_epoch_val <= 0: + raise ValueError( + "No training samples available after dataset filtering. " + "Try increasing n_samples or lowering min_train_masks." + ) + # learning rate schedule learning_rate_schedule = np.linspace(0, learning_rate, 10) learning_rate_schedule = np.append( @@ -1399,8 +1419,87 @@ def run_blocking_task( """Run the blocking training task.""" import time + from cellpose import io as cp_io # type: ignore + session_id = training_params["session_id"] + def _is_readable_pair( + image_path: Path, label_path: Path + ) -> tuple[bool, str | None]: + try: + image_array = cp_io.imread(str(image_path)) + label_array = cp_io.imread(str(label_path)) + except Exception as exc: + return False, str(exc) + + if image_array is None: + return False, "image could not be decoded" + if label_array is None: + return False, "annotation could not be decoded" + + if not hasattr(image_array, "ndim") or int(image_array.ndim) < 2: + return False, "image has invalid ndim" + if not hasattr(label_array, "ndim") or int(label_array.ndim) < 2: + return False, "annotation has invalid ndim" + + return True, None + + train_files = list(dataset_split["train_files"]) + train_labels_files = list(dataset_split["train_labels_files"]) + test_files = list(dataset_split["test_files"] or []) + test_labels_files = list(dataset_split["test_labels_files"] or []) + + filtered_train_files: list[Path] = [] + filtered_train_labels: list[Path] = [] + dropped_train: list[str] = [] + for image_path, label_path in zip(train_files, train_labels_files): + ok, reason = _is_readable_pair(image_path, label_path) + if ok: + filtered_train_files.append(image_path) + filtered_train_labels.append(label_path) + else: + dropped_train.append(f"{image_path} <-> {label_path} ({reason})") + + filtered_test_files: list[Path] = [] + filtered_test_labels: list[Path] = [] + dropped_test: list[str] = [] + for image_path, label_path in zip(test_files, test_labels_files): + ok, reason = _is_readable_pair(image_path, label_path) + if ok: + filtered_test_files.append(image_path) + filtered_test_labels.append(label_path) + else: + dropped_test.append(f"{image_path} <-> {label_path} ({reason})") + + if dropped_train or dropped_test: + append_info( + session_id, + ( + f"Skipped unreadable pairs: train={len(dropped_train)}, " + f"test={len(dropped_test)}" + ), + with_time=True, + ) + logger.warning( + "Session %s: skipped unreadable pairs (train=%d, test=%d)", + session_id, + len(dropped_train), + len(dropped_test), + ) + + if len(filtered_train_files) == 0: + raise ValueError( + "No readable training pairs remain after file validation. " + "Please verify that selected image/annotation files are valid TIFF/PNG/JPEG masks and images." + ) + + dataset_split = DatasetSplit( + train_files=filtered_train_files, + train_labels_files=filtered_train_labels, + test_files=filtered_test_files if filtered_test_files else None, + test_labels_files=filtered_test_labels if filtered_test_labels else None, + ) + # Calculate dataset sizes n_train = len(dataset_split["train_files"]) n_test = ( @@ -1851,7 +1950,7 @@ def load_model( If `identifier` points to an existing file, it is treated as a path to a finetuned model. Otherwise, it is treated as a builtin model name - (e.g., "cyto3"). + (e.g., "cpsam"). """ from cellpose import core, models # type: ignore @@ -1987,6 +2086,26 @@ def _contains_glob(path_pattern: str) -> bool: return "*" in path_pattern +SUPPORTED_DATASET_IMAGE_SUFFIXES = ( + ".ome.tiff", + ".ome.tif", + ".tiff", + ".tif", + ".png", + ".jpg", + ".jpeg", +) + + +def _is_supported_dataset_image_path(path_value: str) -> bool: + normalized = _normalize_artifact_relpath(path_value).lower().rstrip("/") + return normalized.endswith(SUPPORTED_DATASET_IMAGE_SUFFIXES) + + +def _filter_supported_dataset_image_paths(paths: list[str]) -> list[str]: + return [path for path in paths if _is_supported_dataset_image_path(path)] + + def _glob_base_folder(path_pattern: str) -> str: normalized = _normalize_artifact_relpath(path_pattern) wildcard_index = normalized.find("*") @@ -2003,18 +2122,537 @@ def _glob_base_folder(path_pattern: str) -> str: return normalized[: slash_before + 1] +def _is_bioimage_archive_url(candidate: str) -> bool: + parsed = urlparse(str(candidate or "")) + host = parsed.netloc.lower() + path = parsed.path.lower() + return parsed.scheme in {"http", "https"} and ( + "bioimagearchive.org" in host + or ("ebi.ac.uk" in host and "/biostudies/bioimages" in path) + ) + + +async def _get_bia_resolver_service() -> tuple[Any, Any]: + """Connect to Hypha and resolve the BIA URL resolver service.""" + from hypha_rpc import connect_to_server + + server_url = os.environ.get("HYPHA_SERVER_URL", DEFAULT_SERVER_URL) + workspace = os.environ.get("BIA_RESOLVE_URL_WORKSPACE", "ri-scale") + token = ( + os.environ.get("RI_SCALE_TOKEN") + or os.environ.get("HYPHA_TOKEN") + or os.environ.get("BIOENGINE_HYPHA_TOKEN") + ) + if not token: + raise RuntimeError( + "Missing token for resolver service connection. " + "Set RI_SCALE_TOKEN or HYPHA_TOKEN in environment." + ) + + server: Any = await connect_to_server( + { + "server_url": server_url, + "token": token, + "workspace": workspace, + } + ) + + resolved_service_id = BIA_RESOLVE_URL_SERVICE_ID + if not resolved_service_id: + worker_id = os.environ.get("HYPHA_WORKER_SERVICE_ID", "bioimage-io/bioengine-worker") + worker: Any = await server.get_service(worker_id) + app_status = await worker.get_application_status( + application_ids=[BIA_RESOLVE_URL_APPLICATION_ID] + ) + app = (app_status or {}).get(BIA_RESOLVE_URL_APPLICATION_ID, {}) + service_ids = app.get("service_ids") or [] + for item in service_ids: + if isinstance(item, dict) and item.get("websocket_service_id"): + resolved_service_id = item.get("websocket_service_id") + break + if isinstance(item, str) and item: + resolved_service_id = item + break + + if not isinstance(resolved_service_id, str) or not resolved_service_id: + await server.disconnect() + raise RuntimeError( + "Could not resolve BIA resolver service ID. " + "Set BIA_RESOLVE_URL_SERVICE_ID or ensure application " + f"'{BIA_RESOLVE_URL_APPLICATION_ID}' is running on the worker." + ) + + service: Any = await server.get_service(resolved_service_id) + return server, service + + +async def _download_remote_bytes_via_resolver( + resolver_service: Any, + url: str, + *, + timeout_seconds: float = 120.0, +) -> bytes: + """Download binary content through resolver service.""" + payload = await resolver_service.resolve_url( + url=url, + method="GET", + timeout=timeout_seconds, + ) + if not isinstance(payload, dict): + raise RuntimeError(f"Resolver returned invalid payload type for {url}") + + if not payload.get("ok"): + error = payload.get("error") or f"HTTP {payload.get('status_code', 'unknown')}" + raise RuntimeError(f"Resolver request failed for {url}: {error}") + + content_base64 = payload.get("content_base64") + if isinstance(content_base64, str) and content_base64: + return base64.b64decode(content_base64) + + text_value = payload.get("text") + if isinstance(text_value, str): + return text_value.encode("utf-8") + + raise RuntimeError(f"Resolver returned no content for {url}") + + +def _extract_bia_accession(url: str) -> str | None: + match = re.search(r"(S-BIAD\d+)", str(url or ""), flags=re.IGNORECASE) + return match.group(1).upper() if match else None + + +def _bia_ftp_base_url(accession: str) -> str: + accession_number = int(accession.split("S-BIAD", 1)[1]) + suffix = accession_number % 1000 + return f"https://ftp.ebi.ac.uk/biostudies/fire/S-BIAD/{suffix}/{accession}" + + +def _extract_tsv_pairs(tsv_content: str) -> list[tuple[str, str]]: + pairs: list[tuple[str, str]] = [] + reader = csv.DictReader(io.StringIO(tsv_content), delimiter="\t") + for row in reader: + label_path = str(row.get("Files") or "").strip() + source_image = str(row.get("Source image") or "").strip() + if not label_path or not source_image: + continue + pairs.append((source_image, label_path)) + return pairs + + +def _bia_pair_cap_from_n_samples(n_samples: Any) -> int | None: + if isinstance(n_samples, bool): + return None + if isinstance(n_samples, int) and n_samples > 0: + return max(20, n_samples * 8) + if isinstance(n_samples, float) and float(n_samples).is_integer() and n_samples > 0: + as_int = int(n_samples) + return max(20, as_int * 8) + return None + + +async def _fetch_bia_tsv_pairs( + client: Any, + accession: str, + max_pairs: int | None = None, +) -> list[tuple[str, str]]: + base_url = _bia_ftp_base_url(accession) + candidate_paths = [ + f"{base_url}/Files/ps_ovule_labels.tsv", + f"{base_url}/Files/labels.tsv", + f"{base_url}/Files/{accession}_labels.tsv", + ] + + for tsv_url in candidate_paths: + try: + pairs: list[tuple[str, str]] = [] + async with client.stream("GET", tsv_url) as response: + if response.status_code // 100 != 2: + continue + + header: list[str] | None = None + label_idx: int | None = None + source_idx: int | None = None + + async for line in response.aiter_lines(): + if not line: + continue + + row = next(csv.reader([line], delimiter="\t"), []) + if not row: + continue + + if header is None: + header = [str(col).strip().lower() for col in row] + try: + label_idx = header.index("files") + source_idx = header.index("source image") + except ValueError: + label_idx = None + source_idx = None + break + continue + + if label_idx is None or source_idx is None: + continue + if len(row) <= max(label_idx, source_idx): + continue + + label_path = str(row[label_idx] or "").strip() + source_image = str(row[source_idx] or "").strip() + if not label_path or not source_image: + continue + + pairs.append((source_image, label_path)) + if max_pairs is not None and len(pairs) >= max_pairs: + break + + if not pairs: + continue + + return [ + ( + f"{base_url}/Files/{image_rel.lstrip('/')}", + f"{base_url}/Files/{label_rel.lstrip('/')}", + ) + for image_rel, label_rel in pairs + ] + except Exception: + continue + + return [] + + +def _iter_string_values(payload: Any) -> list[str]: + out: list[str] = [] + stack = [payload] + seen: set[int] = set() + + while stack: + current = stack.pop() + current_id = id(current) + if current_id in seen: + continue + seen.add(current_id) + + if isinstance(current, str): + out.append(current) + continue + if isinstance(current, dict): + stack.extend(current.values()) + continue + if isinstance(current, list): + stack.extend(current) + continue + + return out + + +def _looks_like_image_url(url: str) -> bool: + path = urlparse(url).path.lower() + return path.endswith( + (".tif", ".tiff", ".ome.tif", ".ome.tiff", ".png", ".jpg", ".jpeg") + ) + + +def _looks_like_mask_url(url: str) -> bool: + path = urlparse(url).path.lower() + return any( + token in path + for token in ( + "_mask", + "-mask", + "_label", + "-label", + "annotation", + "segmentation", + ) + ) + + +def _pair_key_from_url(url: str, *, is_mask: bool) -> str: + path = urlparse(url).path.lower().strip("/") + for suffix in ( + ".ome.tiff", + ".ome.tif", + ".tiff", + ".tif", + ".png", + ".jpg", + ".jpeg", + ): + if path.endswith(suffix): + path = path[: -len(suffix)] + break + + if is_mask: + for marker in ( + "_mask", + "-mask", + "_label", + "-label", + "_annotation", + "-annotation", + ): + if path.endswith(marker): + path = path[: -len(marker)] + break + + for marker in ("/images/", "/annotations/", "/masks/", "/labels/"): + if marker in path: + path = path.split(marker, 1)[1] + break + + return path + + +async def _fetch_bia_payload(client: Any, endpoint: str, accession: str) -> Any: + attempts: list[tuple[str, dict[str, Any]]] = [ + ("get", {"params": {"q": accession, "size": 1000}}), + ("get", {"params": {"query": accession, "size": 1000}}), + ("post", {"json": {"query": accession, "size": 1000}}), + ("post", {"json": {"q": accession, "size": 1000}}), + ] + + for method, kwargs in attempts: + try: + if method == "get": + response = await client.get(endpoint, **kwargs) + else: + response = await client.post(endpoint, **kwargs) + if response.status_code // 100 == 2: + return response.json() + except Exception: + continue + + return None + + +def _local_path_for_remote_url(root: Path, category: str, remote_url: str) -> Path: + parsed = urlparse(remote_url) + basename = Path(parsed.path).name or "asset.bin" + stem = Path(basename).stem + suffix = "".join(Path(basename).suffixes) or ".bin" + digest = hashlib.md5(remote_url.encode("utf-8")).hexdigest()[:10] + filename = f"{stem}_{digest}{suffix}" + return root / category / filename + + +def _is_test_url(url: str) -> bool: + path = urlparse(url).path.lower() + return "/test/" in path or "_test" in path or "-test" in path + + +async def make_training_pairs_from_bioimage_archive_url( + config: TrainingParams, + save_path: Path, +) -> tuple[list[TrainingPair], list[TrainingPair]]: + import httpx + + archive_url = config["artifact_id"] + accession = _extract_bia_accession(archive_url) + if not accession: + raise ValueError( + "Could not parse BioImage Archive accession (e.g. S-BIAD1234) from URL." + ) + + async with httpx.AsyncClient(timeout=120) as client: + pair_cap = _bia_pair_cap_from_n_samples(config.get("n_samples")) + paired_urls = await _fetch_bia_tsv_pairs(client, accession, max_pairs=pair_cap) + + if not paired_urls: + payloads = [ + await _fetch_bia_payload(client, BIA_FTS_ENDPOINT, accession), + await _fetch_bia_payload(client, BIA_IMAGE_ENDPOINT, accession), + ] + + candidates: set[str] = set() + for payload in payloads: + if payload is None: + continue + for value in _iter_string_values(payload): + if isinstance(value, str) and value.startswith( + ("http://", "https://") + ): + if _looks_like_image_url(value) and accession in value.upper(): + candidates.add(value) + + if not candidates: + raise ValueError( + "No downloadable image assets were found from BioImage Archive for this accession." + ) + + image_urls = sorted( + [url for url in candidates if not _looks_like_mask_url(url)] + ) + mask_urls = sorted([url for url in candidates if _looks_like_mask_url(url)]) + + if not image_urls or not mask_urls: + raise ValueError( + "Found BioImage Archive assets, but could not identify both image and mask files." + ) + + mask_map: dict[str, str] = {} + for mask_url in mask_urls: + key = _pair_key_from_url(mask_url, is_mask=True) + if key and key not in mask_map: + mask_map[key] = mask_url + + paired_urls = [] + for image_url in image_urls: + key = _pair_key_from_url(image_url, is_mask=False) + mask_url = mask_map.get(key) + if mask_url: + paired_urls.append((image_url, mask_url)) + + if not paired_urls: + raise ValueError( + "No image/mask pairs could be inferred from BioImage Archive assets." + ) + + train_url_pairs: list[tuple[str, str]] = [] + test_url_pairs: list[tuple[str, str]] = [] + for image_url, mask_url in paired_urls: + if _is_test_url(image_url) or _is_test_url(mask_url): + test_url_pairs.append((image_url, mask_url)) + else: + train_url_pairs.append((image_url, mask_url)) + + if not train_url_pairs and test_url_pairs: + train_url_pairs = test_url_pairs + test_url_pairs = [] + + if not train_url_pairs: + raise ValueError( + "No training pairs found after downloading BioImage Archive assets" + ) + + split_mode = str(config.get("split_mode", "manual") or "manual").lower() + train_split_ratio = float(config.get("train_split_ratio", 0.8)) + requested_total = resolve_requested_sample_count( + config.get("n_samples"), + len(train_url_pairs) + len(test_url_pairs), + ) + + selected_train_url_pairs = list(train_url_pairs) + selected_test_url_pairs = list(test_url_pairs) + + if requested_total is not None: + if split_mode == "auto": + combined_pairs = [*selected_train_url_pairs, *selected_test_url_pairs] + if requested_total < len(combined_pairs): + indices = np.random.default_rng().permutation(len(combined_pairs))[ # type: ignore[index] + :requested_total + ] + combined_pairs = [combined_pairs[i] for i in indices] + + if len(combined_pairs) <= 1: + selected_train_url_pairs = combined_pairs + selected_test_url_pairs = [] + else: + ratio = max(0.05, min(0.95, float(train_split_ratio))) + n_train = int(round(len(combined_pairs) * ratio)) + n_train = max(1, min(len(combined_pairs) - 1, n_train)) + shuffled = np.random.default_rng().permutation(len(combined_pairs)) + train_idx = set(shuffled[:n_train]) + selected_train_url_pairs = [ + pair + for idx, pair in enumerate(combined_pairs) + if idx in train_idx + ] + selected_test_url_pairs = [ + pair + for idx, pair in enumerate(combined_pairs) + if idx not in train_idx + ] + else: + train_target, test_target = proportional_manual_sample_counts( + len(selected_train_url_pairs), + len(selected_test_url_pairs), + requested_total, + ) + if train_target < len(selected_train_url_pairs): + train_idx = np.random.default_rng().permutation( + len(selected_train_url_pairs) + )[:train_target] + selected_train_url_pairs = [ + selected_train_url_pairs[i] for i in train_idx + ] + if test_target < len(selected_test_url_pairs): + test_idx = np.random.default_rng().permutation( + len(selected_test_url_pairs) + )[:test_target] + selected_test_url_pairs = [ + selected_test_url_pairs[i] for i in test_idx + ] + + selected_train_lookup = set(selected_train_url_pairs) + + bia_cache_root = save_path / "bia_download" + train_pairs: list[TrainingPair] = [] + test_pairs: list[TrainingPair] = [] + + resolver_server, resolver_service = await _get_bia_resolver_service() + try: + for image_url, mask_url in [ + *selected_train_url_pairs, + *selected_test_url_pairs, + ]: + local_image = _local_path_for_remote_url( + bia_cache_root, "images", image_url + ) + local_mask = _local_path_for_remote_url( + bia_cache_root, "annotations", mask_url + ) + local_image.parent.mkdir(parents=True, exist_ok=True) + local_mask.parent.mkdir(parents=True, exist_ok=True) + + if not local_image.exists() or local_image.stat().st_size <= 0: + image_bytes = await _download_remote_bytes_via_resolver( + resolver_service, + image_url, + ) + local_image.write_bytes(image_bytes) + + if not local_mask.exists() or local_mask.stat().st_size <= 0: + mask_bytes = await _download_remote_bytes_via_resolver( + resolver_service, + mask_url, + ) + local_mask.write_bytes(mask_bytes) + + pair = TrainingPair(image=local_image, annotation=local_mask) + if (image_url, mask_url) in selected_train_lookup: + train_pairs.append(pair) + else: + test_pairs.append(pair) + finally: + await resolver_server.disconnect() + + if not train_pairs: + raise ValueError( + "No training pairs found after downloading BioImage Archive assets" + ) + + return train_pairs, test_pairs + + async def list_artifact_files_recursive( artifact: AsyncHyphaArtifact, folder_path: str, + max_results: int | None = None, ) -> list[str]: """Recursively list artifact file paths (relative paths, not folders).""" + limit = None + if isinstance(max_results, int) and max_results > 0: + limit = max_results + root = _normalize_artifact_relpath(folder_path) if root and not root.endswith("/"): root += "/" queue = [root] visited: set[str] = set() - collected: list[str] = [] + collected: set[str] = set() while queue: current = queue.pop(0) @@ -2042,6 +2680,19 @@ async def list_artifact_files_recursive( "directory", "dir", } + if not is_dir and not entry_type and normalized: + # Some artifact backends return directory entries without trailing '/' + # and without explicit type metadata. Probe such paths as directories. + probe_dir = normalized if normalized.endswith("/") else normalized + "/" + if probe_dir not in visited: + try: + probe_entries = await artifact.ls(probe_dir) + if isinstance(probe_entries, list): + queue.append(probe_dir) + continue + except Exception: + pass + if is_dir: next_dir = normalized if normalized.endswith("/") else normalized + "/" if next_dir not in visited: @@ -2049,22 +2700,153 @@ async def list_artifact_files_recursive( continue if normalized: - collected.append(normalized) + collected.add(normalized) + if limit is not None and len(collected) >= limit: + return sorted(collected) - return sorted(set(collected)) + return sorted(collected) async def list_matching_artifact_paths( artifact: AsyncHyphaArtifact, path_pattern: str, + max_results: int | None = None, ) -> list[str]: """List artifact files that match a folder path or glob path pattern.""" normalized_pattern = _normalize_artifact_relpath(path_pattern) + limit = max_results if isinstance(max_results, int) and max_results > 0 else None + + def _cap(values: list[str]) -> list[str]: + if limit is None: + return values + return values[:limit] + + def _path_depth(path: str) -> int: + stripped = path.strip("/") + if not stripped: + return 0 + return len([segment for segment in stripped.split("/") if segment]) + + async def _list_matching_artifact_dirs(dir_pattern: str) -> list[str]: + base_folder = _glob_base_folder(dir_pattern) + target_depth = _path_depth(dir_pattern) + queue = [base_folder] + visited: set[str] = set() + matched: set[str] = set() + + while queue: + current_dir = queue.pop(0) + if current_dir in visited: + continue + visited.add(current_dir) + + try: + entries = await artifact.ls(current_dir) + except Exception: + continue + + for entry in entries: + entry_type = "" + if isinstance(entry, dict): + entry_type = str(entry.get("type") or "").lower() + raw_path = str(entry.get("path") or entry.get("name") or "") + else: + raw_path = str(entry) + + if raw_path and "/" not in raw_path and current_dir: + raw_path = f"{current_dir.rstrip('/')}/{raw_path}" + + normalized = _normalize_artifact_relpath(raw_path) + if not normalized: + continue + + candidate_dir = ( + normalized if normalized.endswith("/") else normalized + "/" + ) + is_dir = entry_type in {"directory", "dir"} or normalized.endswith("/") + if not is_dir: + try: + probe_entries = await artifact.ls(candidate_dir) + is_dir = ( + isinstance(probe_entries, list) and len(probe_entries) > 0 + ) + except Exception: + is_dir = False + + if not is_dir: + continue + + if fnmatch.fnmatch(candidate_dir, dir_pattern): + matched.add(candidate_dir) + + if ( + _path_depth(candidate_dir) < target_depth + and candidate_dir not in visited + ): + queue.append(candidate_dir) + + return sorted(matched) if normalized_pattern.endswith("/"): - folder = normalized_pattern - names = await list_artifact_files(artifact, folder) - return [_normalize_artifact_relpath(f"{folder}{name}") for name in names] + # Folder-like patterns should match files recursively to stay consistent + # with UI validation/path suggestions. + if not _contains_glob(normalized_pattern): + scan_limit = limit * 5 if limit is not None else None + candidates = await list_artifact_files_recursive( + artifact, + normalized_pattern, + max_results=scan_limit, + ) + return _cap(_filter_supported_dataset_image_paths(candidates)) + + matched_dirs = await _list_matching_artifact_dirs(normalized_pattern) + if matched_dirs: + matched_files: list[str] = [] + seen: set[str] = set() + for folder in matched_dirs: + try: + names = await list_artifact_files(artifact, folder) + except Exception: + continue + + folder_prefix = folder if folder.endswith("/") else f"{folder}/" + for name in names: + candidate = _normalize_artifact_relpath(f"{folder_prefix}{name}") + if not candidate or not _is_supported_dataset_image_path(candidate): + continue + if candidate in seen: + continue + seen.add(candidate) + matched_files.append(candidate) + if limit is not None and len(matched_files) >= limit: + return _cap(sorted(matched_files)) + + if matched_files: + return _cap(sorted(matched_files)) + + return _cap(matched_dirs) + + base_folder = _glob_base_folder(normalized_pattern) + scan_limit = limit * 5 if limit is not None else None + candidates = await list_artifact_files_recursive( + artifact, + base_folder, + max_results=scan_limit, + ) + matched = [] + for candidate in candidates: + parent = Path(candidate).parent.as_posix() + parent_pattern_value = parent + "/" if parent and parent != "." else "" + if fnmatch.fnmatch(parent_pattern_value, normalized_pattern): + matched.append(candidate) + if limit is not None and len(matched) >= limit: + break + + filtered_files = sorted(set(_filter_supported_dataset_image_paths(matched))) + if filtered_files: + return _cap(filtered_files) + + return [] if not _contains_glob(normalized_pattern): # If caller passed a directory without trailing slash (common in UI/manual input), @@ -2078,21 +2860,29 @@ async def list_matching_artifact_paths( if normalized_pattern.endswith("/") else normalized_pattern + "/" ) - return [ + candidates = [ _normalize_artifact_relpath(f"{folder}{name}") for name in names ] + return _cap(_filter_supported_dataset_image_paths(candidates)) except Exception: pass - return [normalized_pattern] + if _is_supported_dataset_image_path(normalized_pattern): + return [normalized_pattern] + return [] base_folder = _glob_base_folder(normalized_pattern) - candidates = await list_artifact_files_recursive(artifact, base_folder) + scan_limit = limit * 5 if limit is not None else None + candidates = await list_artifact_files_recursive( + artifact, + base_folder, + max_results=scan_limit, + ) matched = [ candidate for candidate in candidates if fnmatch.fnmatch(candidate, normalized_pattern) ] - return sorted(set(matched)) + return _cap(sorted(set(_filter_supported_dataset_image_paths(matched)))) async def list_artifact_files( @@ -2209,12 +2999,79 @@ def _image_key(path: str) -> str: def _annotation_key(path: str) -> str: basename = Path(path).name key = _strip_known_suffixes(basename).lower() - for marker in ("_mask", "-mask", "_label", "-label", "_annotation", "-annotation"): + for marker in ( + "_mask", + "-mask", + "_label", + "-label", + "_annotation", + "-annotation", + ): if key.endswith(marker): key = key[: -len(marker)] break return key + image_root = _normalize_artifact_relpath( + _glob_base_folder(normalized_image_pattern) + ) + annotation_root = _normalize_artifact_relpath( + _glob_base_folder(normalized_annotation_pattern) + ) + + def _relative_key(path: str, root: str, remove_label_suffix: bool) -> str: + normalized = _normalize_artifact_relpath(path).lower() + relative = normalized + if root and normalized.startswith(root): + relative = normalized[len(root) :] + + for suffix in ( + ".ome.tiff", + ".ome.tif", + ".tiff", + ".tif", + ".png", + ".jpg", + ".jpeg", + ): + if relative.endswith(suffix): + relative = relative[: -len(suffix)] + break + + if remove_label_suffix: + for marker in ( + "_mask", + "-mask", + "_label", + "-label", + "_annotation", + "-annotation", + ): + if relative.endswith(marker): + relative = relative[: -len(marker)] + break + + return relative.strip("/") + + relative_ann_map: dict[str, str] = {} + for annot_file in annotation_files: + key = _relative_key(annot_file, annotation_root, True) + if key and key not in relative_ann_map: + relative_ann_map[key] = annot_file + + for image_file in image_files: + key = _relative_key(image_file, image_root, False) + annot_file = relative_ann_map.get(key) + if annot_file: + pairs.append((image_file, annot_file)) + + if pairs: + logger.info( + f"Matched {len(pairs)} pairs from {len(image_files)} images " + f"and {len(annotation_files)} annotations" + ) + return pairs + # Fallback matcher for mixed conventions like image '*.tif' vs annotation '*_mask.ome.tif'. fallback_ann_map: dict[str, str] = {} for annot_file in annotation_files: @@ -2398,7 +3255,7 @@ async def make_training_pairs_from_metadata( artifact: AsyncHyphaArtifact, metadata_dir: str, save_path: Path, - n_samples: int | None, + n_samples: int | float | None, ) -> tuple[list[TrainingPair], list[TrainingPair]]: metadata_root = _normalize_artifact_relpath(metadata_dir) if metadata_root and not metadata_root.endswith("/"): @@ -2445,11 +3302,17 @@ async def make_training_pairs_from_metadata( "Expected keys like image_path/mask_path (also supports camelCase variants)." ) - if n_samples is not None and n_samples < len(train_pairs_raw): - subset_idx = np.random.default_rng().permutation(len(train_pairs_raw))[ - :n_samples - ] - train_pairs_raw = [train_pairs_raw[i] for i in subset_idx] + requested_total = resolve_requested_sample_count( + n_samples, + len(train_pairs_raw) + len(test_pairs_raw), + ) + train_pairs_raw, test_pairs_raw = sample_pair_lists( + train_pairs_raw, + test_pairs_raw, + requested_total, + split_mode="manual", + train_split_ratio=0.8, + ) train_pairs = await download_pairs_from_artifact( artifact, @@ -2489,10 +3352,31 @@ async def download_pairs_from_artifact( logger.info("Downloading %d files from artifact", len(missing_rpaths)) try: # Add a timeout to prevent indefinite hanging (10 minutes) - await asyncio.wait_for( - artifact.get(missing_rpaths, missing_lpaths, on_error="ignore"), - timeout=600, - ) + try: + await asyncio.wait_for( + artifact.get(missing_rpaths, missing_lpaths, on_error="ignore"), + timeout=600, + ) + except Exception as first_error: + err_msg = str(first_error).lower() + needs_recursive = ( + "path is a directory" in err_msg or "recursive" in err_msg + ) + if needs_recursive: + logger.info( + "Retrying artifact download with recursive=True for directory paths" + ) + await asyncio.wait_for( + artifact.get( + missing_rpaths, + missing_lpaths, + on_error="ignore", + recursive=True, + ), + timeout=600, + ) + else: + raise logger.info("Download completed successfully") except asyncio.TimeoutError: logger.error("Download timed out after 600 seconds") @@ -2577,6 +3461,149 @@ def get_training_subset( return image_paths, annotation_paths +def resolve_requested_sample_count( + n_samples: int | float | None, + total_count: int, +) -> int | None: + """Resolve requested sample usage into an absolute count. + + If ``0 < n_samples <= 1``, it is treated as a decimal fraction of + ``total_count``. Values greater than 1 are treated as absolute counts. + """ + if n_samples is None: + return None + if total_count <= 0: + return 0 + + requested = float(n_samples) + if requested <= 0: + raise ValueError("n_samples must be > 0") + + if requested <= 1.0: + resolved = int(round(total_count * requested)) + else: + resolved = int(round(requested)) + + return max(1, min(total_count, resolved)) + + +def random_subset_pairs(pairs: list[TrainingPair], count: int) -> list[TrainingPair]: + """Return a random subset of ``pairs`` of size ``count``.""" + if count >= len(pairs): + return pairs + idx = np.random.default_rng().permutation(len(pairs))[:count] + return [pairs[i] for i in idx] + + +def proportional_manual_sample_counts( + train_available: int, + test_available: int, + requested_total: int, +) -> tuple[int, int]: + """Allocate requested manual-split samples across train/test pools.""" + total_available = train_available + test_available + if total_available <= 0: + return 0, 0 + + requested_total = max(0, min(requested_total, total_available)) + if requested_total == 0: + return 0, 0 + + if test_available <= 0: + return min(train_available, requested_total), 0 + if train_available <= 0: + return 0, min(test_available, requested_total) + + train_target = int(round(requested_total * (train_available / total_available))) + train_target = max(0, min(train_available, train_target)) + test_target = requested_total - train_target + + if test_target > test_available: + overflow = test_target - test_available + test_target = test_available + train_target = min(train_available, train_target + overflow) + if train_target > train_available: + overflow = train_target - train_available + train_target = train_available + test_target = min(test_available, test_target + overflow) + + assigned = train_target + test_target + if assigned < requested_total: + remaining = requested_total - assigned + train_room = train_available - train_target + add_train = min(train_room, remaining) + train_target += add_train + remaining -= add_train + if remaining > 0: + test_room = test_available - test_target + add_test = min(test_room, remaining) + test_target += add_test + + return train_target, test_target + + +def sample_pair_lists( + train_pairs: list[TrainingPair], + test_pairs: list[TrainingPair], + requested_total: int | None, + *, + split_mode: str, + train_split_ratio: float, +) -> tuple[list[TrainingPair], list[TrainingPair]]: + """Apply random sample limits to pair lists.""" + if requested_total is None: + return train_pairs, test_pairs + + total_pairs = len(train_pairs) + len(test_pairs) + if requested_total >= total_pairs: + return train_pairs, test_pairs + + if split_mode == "auto": + combined_pairs = [*train_pairs, *test_pairs] + combined_pairs = random_subset_pairs(combined_pairs, requested_total) + return split_training_pairs(combined_pairs, train_split_ratio) + + train_target, test_target = proportional_manual_sample_counts( + len(train_pairs), + len(test_pairs), + requested_total, + ) + return ( + random_subset_pairs(train_pairs, train_target), + random_subset_pairs(test_pairs, test_target), + ) + + +def split_training_pairs( + train_pairs: list[TrainingPair], + train_ratio: float, +) -> tuple[list[TrainingPair], list[TrainingPair]]: + """Split training pairs into train/test subsets using the provided ratio. + + Keeps at least one sample in train; for datasets with >=2 samples, keeps at + least one sample in test. + """ + if not train_pairs: + return [], [] + + n_total = len(train_pairs) + if n_total == 1: + return train_pairs, [] + + ratio = float(train_ratio) + ratio = max(0.05, min(0.95, ratio)) + + n_train = int(round(n_total * ratio)) + n_train = max(1, min(n_total - 1, n_train)) + + indices = np.random.default_rng().permutation(n_total) + train_idx = set(indices[:n_train]) + + split_train = [pair for i, pair in enumerate(train_pairs) if i in train_idx] + split_test = [pair for i, pair in enumerate(train_pairs) if i not in train_idx] + return split_train, split_test + + async def make_training_pairs( config: TrainingParams, save_path: Path, @@ -2586,8 +3613,13 @@ async def make_training_pairs( Returns: Tuple of (train_pairs, test_pairs). test_pairs is empty if test folders not specified. """ + artifact_id = config["artifact_id"] + + if _is_bioimage_archive_url(artifact_id): + return await make_training_pairs_from_bioimage_archive_url(config, save_path) + artifact = await make_artifact_client( - config["artifact_id"], + artifact_id, config["server_url"], ) @@ -2618,12 +3650,30 @@ async def make_training_pairs( "train_annotations must be provided." ) - train_image_files = await list_matching_artifact_paths(artifact, train_images) + requested_by_user = config.get("n_samples") + scan_limit: int | None = None + if isinstance(requested_by_user, (int, float)): + requested_float = float(requested_by_user) + if requested_float >= 1.0: + requested_count = max(1, int(round(requested_float))) + # For low-sample interactive runs, avoid scanning entire artifacts. + # Keep a generous buffer so wildcard matching still finds enough pairs. + scan_limit = max(100, requested_count * 40) + + train_image_files = await list_matching_artifact_paths( + artifact, + train_images, + max_results=scan_limit, + ) train_annotation_files = await list_matching_artifact_paths( artifact, train_annotations, + max_results=scan_limit, ) + split_mode = str(config.get("split_mode", "manual") or "manual").lower() + train_split_ratio = float(config.get("train_split_ratio", 0.8)) + # Match training pairs train_matched = match_image_annotation_pairs( train_image_files, @@ -2631,37 +3681,23 @@ async def make_training_pairs( train_images, train_annotations, ) - - # Build full paths for training files - train_image_paths = [Path(img) for img, _ in train_matched] - train_annotation_paths = [Path(ann) for _, ann in train_matched] - - # Apply n_samples if specified - if config["n_samples"] is not None and config["n_samples"] < len(train_image_paths): - train_image_paths, train_annotation_paths = get_training_subset( - train_image_paths, - train_annotation_paths, - config["n_samples"], - ) - - # Download training pairs - train_pairs = await download_pairs_from_artifact( - artifact, - save_path, - train_image_paths, - train_annotation_paths, - ) + train_pairs_raw = [ + TrainingPair(image=Path(img), annotation=Path(ann)) + for img, ann in train_matched + ] # Handle test pairs if specified - test_pairs: list[TrainingPair] = [] + test_pairs_raw: list[TrainingPair] = [] if config["test_images"] and config["test_annotations"]: test_image_files = await list_matching_artifact_paths( artifact, config["test_images"], + max_results=scan_limit, ) test_annotation_files = await list_matching_artifact_paths( artifact, config["test_annotations"], + max_results=scan_limit, ) # Match test pairs @@ -2672,18 +3708,44 @@ async def make_training_pairs( config["test_annotations"], ) - # Build full paths for test files - test_image_paths = [Path(img) for img, _ in test_matched] - test_annotation_paths = [Path(ann) for _, ann in test_matched] + test_pairs_raw = [ + TrainingPair(image=Path(img), annotation=Path(ann)) + for img, ann in test_matched + ] + + requested_total = resolve_requested_sample_count( + config.get("n_samples"), + len(train_pairs_raw) + len(test_pairs_raw), + ) + train_pairs_raw, test_pairs_raw = sample_pair_lists( + train_pairs_raw, + test_pairs_raw, + requested_total, + split_mode=split_mode, + train_split_ratio=train_split_ratio, + ) + + # Download selected training pairs only + train_pairs = await download_pairs_from_artifact( + artifact, + save_path, + [pair["image"] for pair in train_pairs_raw], + [pair["annotation"] for pair in train_pairs_raw], + ) - # Download test pairs + # Download selected test pairs only + test_pairs: list[TrainingPair] = [] + if test_pairs_raw: test_pairs = await download_pairs_from_artifact( artifact, save_path, - test_image_paths, - test_annotation_paths, + [pair["image"] for pair in test_pairs_raw], + [pair["annotation"] for pair in test_pairs_raw], ) + if split_mode == "auto" and not test_pairs: + train_pairs, test_pairs = split_training_pairs(train_pairs, train_split_ratio) + return train_pairs, test_pairs @@ -3163,6 +4225,51 @@ def __init__(self) -> None: self.tasks = {} self._session_lock = asyncio.Lock() + def _has_active_training_task(self) -> bool: + """Return True when any known training task is still running.""" + for task in self.tasks.values(): + if task is not None and not task.done(): + return True + return False + + def _get_free_gpu_memory_bytes(self) -> int | None: + """Return free GPU memory in bytes, or None if unavailable.""" + try: + import torch + except Exception: + return None + + try: + if not torch.cuda.is_available(): + return None + free_bytes, _total_bytes = torch.cuda.mem_get_info() + return int(free_bytes) + except Exception: + return None + + def _admission_denial_message(self) -> str | None: + """Return a human-readable reason why a new training run should not start.""" + if self._has_active_training_task(): + return ( + "Deferred start to avoid GPU contention: another training session is already active. " + "Use restart_training(session_id=...) when resources are available." + ) + + free_gpu_bytes = self._get_free_gpu_memory_bytes() + if free_gpu_bytes is None: + return None + + if free_gpu_bytes < MIN_FREE_GPU_MEMORY_TO_START_BYTES: + free_gb = free_gpu_bytes / GB + required_gb = MIN_FREE_GPU_MEMORY_TO_START_BYTES / GB + return ( + "Deferred start due to low free GPU memory " + f"({free_gb:.2f} GB available, need at least {required_gb:.2f} GB). " + "Use restart_training(session_id=...) after active workloads finish." + ) + + return None + def get_model_id(self, model: str) -> str | Path: """Return a model identifier suitable for loading by Cellpose.""" # print all contents of get_session_path(model): @@ -3176,6 +4283,277 @@ def get_model_id(self, model: str) -> str | Path: ) raise ValueError(msg) + @schema_method(arbitrary_types_allowed=True) # type: ignore + async def preflight_training_dataset( + self, + artifact: str = Field( + description="Artifact identifier or BioImage Archive URL", + examples=["ri-scale/zarr-demo"], + ), + train_images: str | None = Field( + None, + description="Training images path or glob pattern", + ), + train_annotations: str | None = Field( + None, + description="Training annotations path or glob pattern", + ), + metadata_dir: str | None = Field( + None, + description="Optional metadata directory path", + ), + test_images: str | None = Field( + None, + description="Optional test images path or glob pattern", + ), + test_annotations: str | None = Field( + None, + description="Optional test annotations path or glob pattern", + ), + split_mode: str = Field( + "manual", + description="Dataset split mode: manual or auto", + ), + n_samples: int | float | None = Field( + None, + description="Optional sample count or fraction", + ), + max_candidates: int = Field( + 500, + description=( + "Maximum matching candidates to inspect per path spec. " + "Keeps preflight fast for large artifacts." + ), + ), + ) -> dict[str, Any]: + """Preflight-check dataset specs with lightweight backend matching. + + Returns counts and readiness information without starting training. + """ + if isinstance(artifact, dict): + wrapped = artifact + artifact = wrapped.get( + "artifact", wrapped.get("artifact_id", wrapped.get("id", artifact)) + ) + train_images = wrapped.get("train_images", train_images) + train_annotations = wrapped.get("train_annotations", train_annotations) + metadata_dir = wrapped.get("metadata_dir", metadata_dir) + test_images = wrapped.get("test_images", test_images) + test_annotations = wrapped.get("test_annotations", test_annotations) + split_mode = wrapped.get("split_mode", split_mode) + n_samples = wrapped.get("n_samples", n_samples) + max_candidates = wrapped.get("max_candidates", max_candidates) + + artifact = normalize_optional_param(artifact) + train_images = normalize_optional_param(train_images) + train_annotations = normalize_optional_param(train_annotations) + metadata_dir = normalize_optional_param(metadata_dir) + test_images = normalize_optional_param(test_images) + test_annotations = normalize_optional_param(test_annotations) + split_mode = normalize_optional_param(split_mode) + n_samples = normalize_optional_param(n_samples) + max_candidates = int(max_candidates) if max_candidates is not None else 500 + max_candidates = max(10, min(max_candidates, 5000)) + + split_mode_value = str(split_mode or "manual").lower() + if split_mode_value not in {"manual", "auto"}: + split_mode_value = "manual" + + result: dict[str, Any] = { + "ok": False, + "artifact_id": None, + "mode": "artifact", + "split_mode": split_mode_value, + "train_image_count": 0, + "train_annotation_count": 0, + "train_pair_count": 0, + "test_image_count": 0, + "test_annotation_count": 0, + "test_pair_count": 0, + "sampled_total_count": None, + "requested_total_count": None, + "message": "", + } + + if not isinstance(artifact, str) or not artifact: + result["message"] = "artifact must be a non-empty string" + return result + + is_bia = _is_bioimage_archive_url(artifact) + if is_bia: + result["ok"] = True + result["mode"] = "bioimage-archive" + result["artifact_id"] = artifact + result["split_mode"] = "auto" + result["message"] = ( + "BioImage Archive mode detected. Pair matching happens during preparation." + ) + return result + + server_url, artifact_id = get_url_and_artifact_id(artifact) + result["artifact_id"] = artifact_id + + try: + artifact_client = await make_artifact_client(artifact_id, server_url) + except Exception as e: + result["message"] = f"Failed to access artifact '{artifact_id}': {e}" + return result + + if metadata_dir: + metadata_dir_value = str(metadata_dir) + if not metadata_dir_value.endswith("/"): + metadata_dir_value += "/" + try: + entries = await artifact_client.ls(metadata_dir_value) + except Exception as e: + result["message"] = ( + f"Metadata directory '{metadata_dir_value}' is not readable: {e}" + ) + return result + + metadata_file_count = 0 + for entry in entries or []: + raw_path = "" + if isinstance(entry, dict): + raw_path = str(entry.get("path") or entry.get("name") or "") + else: + raw_path = str(entry) + if raw_path.lower().endswith(".json"): + metadata_file_count += 1 + + if metadata_file_count <= 0: + result["message"] = ( + f"No metadata JSON files found under '{metadata_dir_value}'." + ) + return result + + result["ok"] = True + result["mode"] = "metadata" + result["metadata_file_count"] = metadata_file_count + result["message"] = ( + f"Metadata mode ready. Found {metadata_file_count} metadata JSON files." + ) + return result + + if not isinstance(train_images, str) or not train_images: + result["message"] = ( + "train_images must be provided when metadata_dir is not used" + ) + return result + if not isinstance(train_annotations, str) or not train_annotations: + result["message"] = ( + "train_annotations must be provided when metadata_dir is not used" + ) + return result + + try: + train_image_paths = await list_matching_artifact_paths( + artifact_client, + train_images, + max_results=max_candidates, + ) + train_annotation_paths = await list_matching_artifact_paths( + artifact_client, + train_annotations, + max_results=max_candidates, + ) + except Exception as e: + result["message"] = f"Failed to list training paths: {e}" + return result + + train_image_count = len( + [ + path + for path in train_image_paths + if _is_supported_dataset_image_path(path) + ] + ) + train_annotation_count = len( + [ + path + for path in train_annotation_paths + if _is_supported_dataset_image_path(path) + ] + ) + train_pair_count = min(train_image_count, train_annotation_count) + + result["train_image_count"] = train_image_count + result["train_annotation_count"] = train_annotation_count + result["train_pair_count"] = train_pair_count + + if train_pair_count <= 0: + result["message"] = ( + "No training pairs found from current training path specs. " + f"images={train_image_count}, annotations={train_annotation_count}." + ) + return result + + test_pair_count = 0 + if split_mode_value == "manual" and test_images and test_annotations: + try: + test_image_paths = await list_matching_artifact_paths( + artifact_client, + test_images, + max_results=max_candidates, + ) + test_annotation_paths = await list_matching_artifact_paths( + artifact_client, + test_annotations, + max_results=max_candidates, + ) + except Exception as e: + result["message"] = f"Failed to list test paths: {e}" + return result + + test_image_count = len( + [ + path + for path in test_image_paths + if _is_supported_dataset_image_path(path) + ] + ) + test_annotation_count = len( + [ + path + for path in test_annotation_paths + if _is_supported_dataset_image_path(path) + ] + ) + test_pair_count = min(test_image_count, test_annotation_count) + + result["test_image_count"] = test_image_count + result["test_annotation_count"] = test_annotation_count + result["test_pair_count"] = test_pair_count + + if test_pair_count <= 0: + result["message"] = ( + "Manual split selected but no test pairs were found from test path specs. " + f"images={test_image_count}, annotations={test_annotation_count}." + ) + return result + + total_pairs = train_pair_count + test_pair_count + requested_total = None + if n_samples is not None: + try: + requested_total = resolve_requested_sample_count( + float(n_samples), total_pairs + ) + except Exception as e: + result["message"] = f"Invalid n_samples value: {e}" + return result + + result["requested_total_count"] = requested_total + result["sampled_total_count"] = ( + requested_total if requested_total is not None else total_pairs + ) + result["ok"] = True + result["message"] = ( + "Dataset preflight passed. " + f"train_pairs={train_pair_count}, test_pairs={test_pair_count}, total={total_pairs}." + ) + return result + @schema_method(arbitrary_types_allowed=True) # type: ignore async def start_training( self, @@ -3239,6 +4617,21 @@ async def start_training( ), examples=["annotations/test/", "annotations/test/*_mask.ome.tif"], ), + split_mode: str = Field( + "manual", + description=( + "Dataset split mode. 'manual' uses test_images/test_annotations; " + "'auto' creates test split from training pairs using train_split_ratio." + ), + examples=["manual", "auto"], + ), + train_split_ratio: float = Field( + 0.8, + description=( + "Train ratio used only when split_mode='auto'. " + "Example: 0.8 means 80% train / 20% test." + ), + ), model: str = Field( PretrainedModel.CPSAM.value, description=( @@ -3251,11 +4644,13 @@ async def start_training( *PretrainedModel.values(), ], ), - n_samples: int | None = Field( + n_samples: int | float | None = Field( None, description=( - "Optional number of samples to use from the dataset. If None, " - "all available samples are used." + "Optional number of samples to use from the dataset. " + "If 0 < value <= 1, treated as decimal fraction of total samples " + "(e.g., 0.6 means 60%). If value > 1, treated as absolute count. " + "If None, all available samples are used." ), ), n_epochs: int = Field(10, description="Number of training epochs"), @@ -3297,6 +4692,8 @@ async def start_training( metadata_dir = wrapped.get("metadata_dir", metadata_dir) test_images = wrapped.get("test_images", test_images) test_annotations = wrapped.get("test_annotations", test_annotations) + split_mode = wrapped.get("split_mode", split_mode) + train_split_ratio = wrapped.get("train_split_ratio", train_split_ratio) model = wrapped.get("model", model) n_samples = wrapped.get("n_samples", n_samples) n_epochs = wrapped.get("n_epochs", n_epochs) @@ -3313,6 +4710,8 @@ async def start_training( metadata_dir = normalize_optional_param(metadata_dir) test_images = normalize_optional_param(test_images) test_annotations = normalize_optional_param(test_annotations) + split_mode = normalize_optional_param(split_mode) + train_split_ratio = normalize_optional_param(train_split_ratio) model = normalize_optional_param(model) n_samples = normalize_optional_param(n_samples) n_epochs = normalize_optional_param(n_epochs) @@ -3322,22 +4721,42 @@ async def start_training( validation_interval = normalize_optional_param(validation_interval) if metadata_dir is None: - if not isinstance(train_images, str) or not train_images: - raise ValueError( - "train_images must be a non-empty string when metadata_dir is not provided" - ) - if not isinstance(train_annotations, str) or not train_annotations: - raise ValueError( - "train_annotations must be a non-empty string when metadata_dir is not provided" - ) + is_bia = isinstance(artifact, str) and _is_bioimage_archive_url(artifact) + if not is_bia: + if not isinstance(train_images, str) or not train_images: + raise ValueError( + "train_images must be a non-empty string when metadata_dir is not provided" + ) + if not isinstance(train_annotations, str) or not train_annotations: + raise ValueError( + "train_annotations must be a non-empty string when metadata_dir is not provided" + ) else: if not isinstance(metadata_dir, str) or not metadata_dir: raise ValueError("metadata_dir must be a non-empty string") - if (test_images is None) ^ (test_annotations is None): - raise ValueError( - "test_images and test_annotations must be provided together" - ) + is_bia = isinstance(artifact, str) and _is_bioimage_archive_url(artifact) + + if not isinstance(split_mode, str) or split_mode not in {"manual", "auto"}: + split_mode = "manual" + if is_bia: + split_mode = "auto" + + if train_split_ratio is None: + train_split_ratio = 0.8 + train_split_ratio = float(train_split_ratio) + if not (0.0 < train_split_ratio < 1.0): + raise ValueError("train_split_ratio must be between 0 and 1") + + if split_mode == "manual": + if (test_images is None) ^ (test_annotations is None): + raise ValueError( + "test_images and test_annotations must be provided together in manual split mode" + ) + else: + # In auto split mode, ignore explicitly provided test inputs. + test_images = None + test_annotations = None if not isinstance(model, str) or not model: model = PretrainedModel.CPSAM.value @@ -3350,7 +4769,9 @@ async def start_training( if min_train_masks is None: min_train_masks = 5 if n_samples is not None: - n_samples = int(n_samples) + n_samples = float(n_samples) + if n_samples <= 0: + raise ValueError("n_samples must be > 0") if validation_interval is not None: validation_interval = int(validation_interval) @@ -3385,6 +4806,8 @@ async def start_training( metadata_dir=metadata_dir, test_images=test_images, test_annotations=test_annotations, + split_mode=split_mode, + train_split_ratio=train_split_ratio, model=model_id, n_epochs=n_epochs, learning_rate=learning_rate, @@ -3403,6 +4826,29 @@ async def start_training( training_params_path.write_text(json.dumps(params_dict, indent=2)) async with self._session_lock: + denial_message = self._admission_denial_message() + if denial_message is not None: + update_status( + session_id=session_id, + status_type=StatusType.STOPPED, + message=denial_message, + dataset_artifact_id=artifact_id, + model=model_id, + n_samples=n_samples, + n_epochs=n_epochs, + learning_rate=learning_rate, + weight_decay=weight_decay, + min_train_masks=min_train_masks, + validation_interval=validation_interval, + ) + append_info( + session_id, + denial_message, + with_time=True, + ) + status = get_status(session_id) + return SessionStatusWithId(**status, session_id=session_id) + executor = ThreadPoolExecutor(max_workers=1) self.executors[session_id] = executor task = asyncio.create_task(launch_training_task(training_params, executor)) @@ -3618,6 +5064,8 @@ async def restart_training( metadata_dir=params.get("metadata_dir"), test_images=params.get("test_images"), test_annotations=params.get("test_annotations"), + split_mode=params.get("split_mode", "manual"), + train_split_ratio=params.get("train_split_ratio", 0.8), model=restart_model, n_samples=params.get("n_samples"), n_epochs=restart_epochs, @@ -3629,6 +5077,72 @@ async def restart_training( restarted["restarted_from"] = session_id return restarted + @schema_method(arbitrary_types_allowed=True) # type: ignore + async def delete_training_session( + self, + session_id: str = Field(description="Session ID to delete"), + force_stop_if_blocked: bool = Field( + False, + description="If true, stop blocked running/preparing session first, then delete.", + ), + ) -> dict[str, Any]: + """Delete a non-running training session and its local artifacts.""" + import shutil + + if isinstance(session_id, dict): + wrapped = session_id + session_id = wrapped.get("session_id", wrapped.get("id", session_id)) + + session_id = str(session_id).strip().replace("\\", "/") + if session_id.endswith("/status.json"): + session_id = session_id[: -len("/status.json")] + session_id = Path(session_id).name + + status_path = get_status_path(session_id) + if not status_path.exists(): + raise ValueError(f"Session {session_id} does not exist") + + try: + status_mtime = status_path.stat().st_mtime + except OSError: + status_mtime = time.time() + + status = self._normalize_session_status( + session_id=session_id, + session_data=dict(get_status(session_id)), + status_mtime=status_mtime, + ) + status_type = str(status.get("status_type") or "").lower() + blocked = {StatusType.RUNNING.value, StatusType.PREPARING.value} + if status_type in blocked: + if force_stop_if_blocked: + await self.stop_training(session_id=session_id) + else: + raise ValueError( + f"Session {session_id} has status '{status_type}'. " + "Running or preparing sessions cannot be deleted. Stop the session first." + ) + + task = self.tasks.get(session_id) + if task and not task.done(): + task.cancel() + if session_id in self.tasks: + del self.tasks[session_id] + + executor = self.executors.get(session_id) + if executor is not None: + executor.shutdown(wait=False, cancel_futures=True) + del self.executors[session_id] + + session_path = get_session_path(session_id) + if session_path.exists(): + shutil.rmtree(session_path, ignore_errors=False) + + return { + "deleted": True, + "session_id": session_id, + } + def _list_sessions_sync( self, status_types: list[str] | None, @@ -4459,7 +5973,7 @@ async def test_cellpose_finetune() -> None: # Reverting to original call style but casting: cellpose_tuner = CellposeFinetune.func_or_class() - await infer(cellpose_tuner, model="cyto3") + await infer(cellpose_tuner, model="cpsam") session_status = await cellpose_tuner.start_training( artifact="ri-scale/zarr-demo", diff --git a/bioengine_apps/cellpose_finetuning/tests/test_metadata_and_glob.py b/bioengine_apps/cellpose_finetuning/tests/test_metadata_and_glob.py deleted file mode 100644 index 8e479f60..00000000 --- a/bioengine_apps/cellpose_finetuning/tests/test_metadata_and_glob.py +++ /dev/null @@ -1,184 +0,0 @@ -from __future__ import annotations - -import asyncio -from pathlib import Path - -from bioengine_apps.cellpose_finetuning.main import ( - list_matching_artifact_paths, - make_training_pairs_from_metadata, - match_image_annotation_pairs, -) - - -def test_match_image_annotation_pairs_with_nested_globs() -> None: - image_files = [ - "images/a/t0001.tif", - "images/b/t0002.tif", - "images/c/extra.tif", - ] - annotation_files = [ - "annotations/a/t0001_mask.ome.tif", - "annotations/b/t0002_mask.ome.tif", - "annotations/x/other_mask.ome.tif", - ] - - pairs = match_image_annotation_pairs( - image_files, - annotation_files, - "images/*/*.tif", - "annotations/*/*_mask.ome.tif", - ) - - assert pairs == [ - ("images/a/t0001.tif", "annotations/a/t0001_mask.ome.tif"), - ("images/b/t0002.tif", "annotations/b/t0002_mask.ome.tif"), - ] - - -def test_match_image_annotation_pairs_with_mixed_ome_suffix_convention() -> None: - image_files = [ - "images/108bb69d-2e52-4382-8100-e96173db24ee/t0000.ome.tif", - "images/108bb69d-2e52-4382-8100-e96173db24ee/t0001.ome.tif", - ] - annotation_files = [ - "annotations/108bb69d-2e52-4382-8100-e96173db24ee/t0000_mask.ome.tif", - "annotations/108bb69d-2e52-4382-8100-e96173db24ee/t0001_mask.ome.tif", - ] - - pairs = match_image_annotation_pairs( - image_files, - annotation_files, - "images/108bb69d-2e52-4382-8100-e96173db24ee/*.tif", - "annotations/108bb69d-2e52-4382-8100-e96173db24ee/*_mask.ome.tif", - ) - - assert pairs == [ - ( - "images/108bb69d-2e52-4382-8100-e96173db24ee/t0000.ome.tif", - "annotations/108bb69d-2e52-4382-8100-e96173db24ee/t0000_mask.ome.tif", - ), - ( - "images/108bb69d-2e52-4382-8100-e96173db24ee/t0001.ome.tif", - "annotations/108bb69d-2e52-4382-8100-e96173db24ee/t0001_mask.ome.tif", - ), - ] - - -class _FakeArtifact: - async def ls(self, folder_path: str): - await asyncio.sleep(0) - folder = folder_path.rstrip("/") + "/" - if folder == "metadata/": - return [{"path": "metadata/sample.json", "type": "file"}] - if folder == "images/train/": - return [ - {"path": "images/train/t0001.tif", "type": "file"}, - {"path": "images/train/t0002.tif", "type": "file"}, - ] - return [] - - async def get(self, remote_paths, local_paths, on_error="ignore"): - await asyncio.sleep(0) - for remote, local in zip(remote_paths, local_paths): - local_path = Path(local) - local_path.parent.mkdir(parents=True, exist_ok=True) - - if str(remote).endswith("metadata/sample.json"): - local_path.write_text( - """ -[ - {"image_path": "images/train/t0001.tif", "mask_path": "annotations/train/t0001_mask.ome.tif", "split": "train"}, - {"image_path": "images/test/t0002.tif", "mask_path": "annotations/test/t0002_mask.ome.tif", "split": "test"} -] - """.strip(), - encoding="utf-8", - ) - else: - local_path.write_bytes(b"dummy") - - -class _FakeArtifactNestedMetadata: - async def ls(self, folder_path: str): - await asyncio.sleep(0) - folder = folder_path.rstrip("/") + "/" - if folder == "metadata/": - return [{"path": "metadata/records.json", "type": "file"}] - return [] - - async def get(self, remote_paths, local_paths, on_error="ignore"): - await asyncio.sleep(0) - for remote, local in zip(remote_paths, local_paths): - local_path = Path(local) - local_path.parent.mkdir(parents=True, exist_ok=True) - - if str(remote).endswith("metadata/records.json"): - local_path.write_text( - """ -{ - "payload": { - "items": [ - { - "imagePath": "images/train/t0010.tif", - "maskPath": "annotations/train/t0010_mask.ome.tif", - "split": "train" - }, - { - "source": {"path": "images/test/t0011.tif"}, - "target": {"path": "annotations/test/t0011_mask.ome.tif"}, - "subset": "validation" - } - ] - } -} - """.strip(), - encoding="utf-8", - ) - else: - local_path.write_bytes(b"dummy") - - -def test_make_training_pairs_from_metadata(tmp_path: Path) -> None: - artifact = _FakeArtifact() - train_pairs, test_pairs = asyncio.run( - make_training_pairs_from_metadata( - artifact=artifact, - metadata_dir="metadata/", - save_path=tmp_path, - n_samples=None, - ) - ) - - assert len(train_pairs) == 1 - assert len(test_pairs) == 1 - assert train_pairs[0]["image"].exists() - assert train_pairs[0]["annotation"].exists() - assert test_pairs[0]["image"].exists() - assert test_pairs[0]["annotation"].exists() - - -def test_list_matching_artifact_paths_directory_without_trailing_slash() -> None: - artifact = _FakeArtifact() - matches = asyncio.run(list_matching_artifact_paths(artifact, "images/train")) - assert matches == [ - "images/train/t0001.tif", - "images/train/t0002.tif", - ] - - -def test_make_training_pairs_from_nested_metadata_formats(tmp_path: Path) -> None: - artifact = _FakeArtifactNestedMetadata() - train_pairs, test_pairs = asyncio.run( - make_training_pairs_from_metadata( - artifact=artifact, - metadata_dir="metadata/", - save_path=tmp_path, - n_samples=None, - ) - ) - - assert len(train_pairs) == 1 - assert len(test_pairs) == 1 - assert train_pairs[0]["image"].as_posix().endswith("images/train/t0010.tif") - assert train_pairs[0]["annotation"].as_posix().endswith("annotations/train/t0010_mask.ome.tif") - assert test_pairs[0]["image"].as_posix().endswith("images/test/t0011.tif") - assert test_pairs[0]["annotation"].as_posix().endswith("annotations/test/t0011_mask.ome.tif") diff --git a/scripts/tmp_probe_bia_fetch_only.py b/scripts/tmp_probe_bia_fetch_only.py new file mode 100644 index 00000000..ee1ac181 --- /dev/null +++ b/scripts/tmp_probe_bia_fetch_only.py @@ -0,0 +1,140 @@ +import argparse +import asyncio +import csv +import io +import time +from dataclasses import dataclass + +import httpx + + +@dataclass +class FetchResult: + label: str + url: str + ok: bool + status_code: int | None + elapsed_s: float + size_bytes: int | None = None + error: str | None = None + + +def _bia_ftp_base_url(accession: str) -> str: + accession_number = int(accession.split("S-BIAD", 1)[1]) + suffix = accession_number % 1000 + return f"https://ftp.ebi.ac.uk/biostudies/fire/S-BIAD/{suffix}/{accession}" + + +def _extract_tsv_pairs(tsv_content: str) -> list[tuple[str, str]]: + pairs: list[tuple[str, str]] = [] + reader = csv.DictReader(io.StringIO(tsv_content), delimiter="\t") + for row in reader: + label_path = str(row.get("Files") or "").strip() + source_image = str(row.get("Source image") or "").strip() + if not label_path or not source_image: + continue + pairs.append((source_image, label_path)) + return pairs + + +async def _timed_get(client: httpx.AsyncClient, label: str, url: str) -> FetchResult: + started = time.perf_counter() + try: + resp = await client.get(url) + elapsed = time.perf_counter() - started + return FetchResult( + label=label, + url=url, + ok=200 <= resp.status_code < 300, + status_code=resp.status_code, + elapsed_s=elapsed, + size_bytes=len(resp.content), + error=None if 200 <= resp.status_code < 300 else f"HTTP {resp.status_code}", + ) + except Exception as exp: + elapsed = time.perf_counter() - started + return FetchResult( + label=label, + url=url, + ok=False, + status_code=None, + elapsed_s=elapsed, + size_bytes=None, + error=str(exp), + ) + + +def _print_result(result: FetchResult) -> None: + status = result.status_code if result.status_code is not None else "ERR" + size = result.size_bytes if result.size_bytes is not None else "-" + msg = f"[{result.label}] ok={result.ok} status={status} time={result.elapsed_s:.2f}s size={size}" + if result.error: + msg += f" error={result.error}" + print(msg, flush=True) + + +async def main() -> None: + parser = argparse.ArgumentParser(description="BIA-only fetch probe (no training calls).") + parser.add_argument("--accession", default="S-BIAD1392") + parser.add_argument("--sample-pairs", type=int, default=3) + args = parser.parse_args() + + accession = str(args.accession).upper() + base_url = _bia_ftp_base_url(accession) + + study_url = f"https://www.ebi.ac.uk/biostudies/bioimages/studies/{accession}" + tsv_candidates = [ + f"{base_url}/Files/ps_ovule_labels.tsv", + f"{base_url}/Files/labels.tsv", + f"{base_url}/Files/{accession}_labels.tsv", + ] + + timeout = httpx.Timeout(connect=20.0, read=120.0, write=20.0, pool=20.0) + async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client: + print(f"Probing accession: {accession}") + + study_result = await _timed_get(client, "study_page", study_url) + _print_result(study_result) + + tsv_result: FetchResult | None = None + tsv_text: str | None = None + for tsv_url in tsv_candidates: + result = await _timed_get(client, "tsv", tsv_url) + _print_result(result) + if result.ok: + resp = await client.get(tsv_url) + tsv_text = resp.text + tsv_result = result + break + + if not tsv_result or not tsv_text: + print("No TSV source succeeded. BIA pair listing is blocked/slow at source.") + return + + pairs = _extract_tsv_pairs(tsv_text) + print(f"Parsed TSV pairs: {len(pairs)}") + if not pairs: + print("TSV loaded but contained no usable image/mask pairs.") + return + + sample_count = max(1, min(int(args.sample_pairs), len(pairs))) + sample_pairs = pairs[:sample_count] + + urls: list[tuple[str, str]] = [] + for idx, (image_rel, mask_rel) in enumerate(sample_pairs, start=1): + image_url = f"{base_url}/Files/{image_rel.lstrip('/')}" + mask_url = f"{base_url}/Files/{mask_rel.lstrip('/')}" + urls.append((f"pair{idx}_image", image_url)) + urls.append((f"pair{idx}_mask", mask_url)) + + print(f"Fetching {len(urls)} sample files from FTP...") + results = await asyncio.gather(*[_timed_get(client, label, url) for label, url in urls]) + for result in results: + _print_result(result) + + ok_count = sum(1 for item in results if item.ok) + print(f"Sample fetch success: {ok_count}/{len(results)}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/cellpose/cellpose_training_agent_startup.pyodide.py b/tests/cellpose/cellpose_training_agent_startup.pyodide.py new file mode 100644 index 00000000..77eeb88f --- /dev/null +++ b/tests/cellpose/cellpose_training_agent_startup.pyodide.py @@ -0,0 +1,391 @@ +import micropip + +await micropip.install(["hypha-rpc", "httpx"]) + +import json +import os +from typing import Any + +from hypha_rpc import connect_to_server, login + +SERVER_URL = os.environ.get("HYPHA_SERVER_URL", "https://hypha.aicell.io") +DEFAULT_SERVICE = "bioimage-io/cellpose-finetuning" + +_state: dict[str, Any] = { + "server": None, + "workspace": None, + "service": None, + "service_id": None, + "schemas": {}, +} + + +async def connect_cellpose( + server_url: str | None = None, + workspace: str | None = None, + token: str | None = None, +): + target_server = server_url or SERVER_URL + target_token = token or os.environ.get("HYPHA_TOKEN") + if not target_token: + target_token = await login({"server_url": target_server}) + + config: dict[str, Any] = {"server_url": target_server, "token": target_token} + if workspace: + config["workspace"] = workspace + + server = await connect_to_server(config) + _state["server"] = server + _state["workspace"] = server.config.workspace + return { + "server_url": target_server, + "workspace": server.config.workspace, + "user_id": server.config.user.get("id"), + } + + +async def connect_cellpose_service(): + return await select_cellpose_service(DEFAULT_SERVICE) + + +async def list_cellpose_services(): + if not _state["server"]: + await connect_cellpose() + + server = _state["server"] + workspace = _state.get("workspace") + candidates: list[str] = [] + + env_service = os.environ.get("HYPHA_TEST_SERVICE_ID") + if env_service: + candidates.append(env_service) + + if workspace: + candidates.append(f"{workspace}/cellpose-finetuning") + + candidates.extend( + [ + DEFAULT_SERVICE, + "bioimage-io/cellpose-finetuning", + ] + ) + + # Worker-derived service ids (best source after redeploys) + worker_candidates = ["bioimage-io/bioengine-worker"] + if workspace: + worker_candidates.append(f"{workspace}/bioengine-worker") + + for worker_id in worker_candidates: + try: + worker = await server.get_service(worker_id) + app_status = await worker.get_application_status( + application_ids=["cellpose-finetuning"] + ) + app_data = ( + app_status.get("cellpose-finetuning", {}) + if isinstance(app_status, dict) + else {} + ) + for entry in app_data.get("service_ids", []): + if isinstance(entry, dict): + service_id = entry.get("websocket_service_id") + else: + service_id = entry + if isinstance(service_id, str) and service_id: + candidates.append(service_id) + except Exception: + continue + + deduped: list[str] = [] + seen: set[str] = set() + for service_id in candidates: + if service_id and service_id not in seen: + seen.add(service_id) + deduped.append(service_id) + + services: list[dict[str, Any]] = [] + for service_id in deduped: + available = False + error = None + try: + await server.get_service(service_id) + available = True + except Exception as e: + error = str(e) + services.append( + {"service_id": service_id, "available": available, "error": error} + ) + + return services + + +async def select_cellpose_service(service_id: str | None = None): + if not _state["server"]: + await connect_cellpose() + + server = _state["server"] + target = service_id + if not target: + services = await list_cellpose_services() + for service_info in services: + if service_info.get("available"): + target = service_info.get("service_id") + break + + if not target: + raise RuntimeError( + "Could not resolve an available Cellpose service. " + "Call list_cellpose_services() and pass an explicit service_id." + ) + + service = await server.get_service(target) + _state["service"] = service + _state["service_id"] = target + await _refresh_service_schemas() + return {"service_id": target, "workspace": _state.get("workspace")} + + +async def _refresh_service_schemas(): + if not _state["service"]: + await connect_cellpose_service() + + service = _state["service"] + method_names = [ + "start_training", + "get_training_status", + "list_training_sessions", + "stop_training", + "restart_training", + "infer", + "export_model", + ] + + schemas: dict[str, Any] = {} + for name in method_names: + method = getattr(service, name, None) + if method is None: + continue + schema = getattr(method, "__schema__", None) + if schema is not None: + schemas[name] = schema + + _state["schemas"] = schemas + return schemas + + +async def start_cellpose_training(**kwargs): + if not _state["service"]: + await connect_cellpose_service() + return await _state["service"].start_training(**kwargs) + + +async def get_training_status(session_id: str): + if not _state["service"]: + await connect_cellpose_service() + return await _state["service"].get_training_status(session_id=session_id) + + +async def list_training_sessions(status_types: list[str] | None = None): + if not _state["service"]: + await connect_cellpose_service() + if status_types: + return await _state["service"].list_training_sessions( + status_types=status_types, + ) + return await _state["service"].list_training_sessions() + + +async def stop_training(session_id: str): + if not _state["service"]: + await connect_cellpose_service() + return await _state["service"].stop_training(session_id=session_id) + + +async def restart_training(session_id: str, n_epochs: int | None = None): + if not _state["service"]: + await connect_cellpose_service() + payload: dict[str, Any] = {"session_id": session_id} + if n_epochs is not None: + payload["n_epochs"] = int(n_epochs) + return await _state["service"].restart_training(**payload) + + +async def infer_cellpose(**kwargs): + if not _state["service"]: + await connect_cellpose_service() + return await _state["service"].infer(**kwargs) + + +async def export_model( + session_id: str, + model_name: str | None = None, + collection: str = "bioimage-io/colab-annotations", +): + if not _state["service"]: + await connect_cellpose_service() + payload: dict[str, Any] = { + "session_id": session_id, + "collection": collection, + } + if model_name: + payload["model_name"] = model_name + return await _state["service"].export_model(**payload) + + +async def get_cellpose_schemas(refresh: bool = False): + if refresh or not _state.get("schemas"): + await _refresh_service_schemas() + return _state.get("schemas", {}) + + +async def print_cellpose_runtime_docs(include_schemas: bool = True): + docs = RUNTIME_GUIDE + if include_schemas: + schemas = await get_cellpose_schemas(refresh=True) + if schemas: + docs += "\n\nLive service schemas (from service..__schema__):\n" + docs += json.dumps(schemas, indent=2, sort_keys=True) + else: + docs += ( + "\n\nLive service schemas are unavailable. " + "Connect/select a service first, then call print_cellpose_runtime_docs()." + ) + print(docs) + + +RUNTIME_GUIDE = """ +You are a Cellpose training operations assistant for BioEngine. + +Primary goal: +- Help users run reliable Cellpose fine-tuning workflows end-to-end: connect, select service, start training, monitor, troubleshoot, stop/restart, infer, export. + +Recommended execution order: +1) Connect: + - await connect_cellpose(server_url=None, workspace=None, token=None) +2) Discover and select service (important after redeploys): + - await list_cellpose_services() + - await select_cellpose_service(service_id=None) +3) Start training: + - await start_cellpose_training(...) +4) Monitor: + - await get_training_status(session_id) + - optional: await list_training_sessions(status_types=[...]) +5) Control: + - await stop_training(session_id) + - await restart_training(session_id, n_epochs=None) +6) Post-training: + - await infer_cellpose(...) + - await export_model(session_id, model_name=None, collection='bioimage-io/colab-annotations') + +Function reference (wrapper-level): +- connect_cellpose(server_url=None, workspace=None, token=None) + - Connects to Hypha server and caches server/workspace in _state. + - If token missing, invokes interactive login(). + - Returns {server_url, workspace, user_id}. + +- list_cellpose_services() + - Lists candidate Cellpose service IDs from defaults + worker app status. + - Returns [{service_id, available, error}]. + +- select_cellpose_service(service_id=None) + - Selects explicit service_id or first available from list_cellpose_services(). + - Caches service + service_id in _state. + - Refreshes live __schema__ docs. + - Returns {service_id, workspace}. + +- start_cellpose_training(**kwargs) + - Forwards kwargs to service.start_training. + - Required: artifact and either + a) train_images + train_annotations, or + b) metadata_dir, or + c) BioImage Archive URL in artifact. + +- get_training_status(session_id) + - Returns latest status including status_type/message/progress/losses/metrics. + +- list_training_sessions(status_types=None) + - Lists known sessions, optionally filtered by status. + +- stop_training(session_id) + - Requests graceful stop for running session. + +- restart_training(session_id, n_epochs=None) + - Starts a new run from previous session checkpoint/params. + +- infer_cellpose(**kwargs) + - Forwards infer parameters to service.infer. + +- export_model(session_id, model_name=None, collection='bioimage-io/colab-annotations') + - Exports completed training to BioImage.IO artifact. + +- get_cellpose_schemas(refresh=False) + - Returns cached or live-refreshed method schemas from service..__schema__. + +- print_cellpose_runtime_docs(include_schemas=True) + - Prints this guide plus live schemas. + +Training parameter guidance: +- split_mode: use 'auto' for BioImage Archive URL artifacts. +- train_split_ratio: used only in auto split (e.g. 0.8 for 80/20). +- n_samples: + - if user gives percentage p%, pass n_samples=p/100. + - valid decimal range: 0 < n_samples <= 1. + - examples: 0.005 = 0.5%, 0.02 = 2%, 1.0 = 100%. +- validation metrics require test data (manual split) or auto split that creates test data. + +Monitoring loop checklist: +- Poll every 2-5 seconds. +- Report: status_type, message, current_epoch/total_epochs, elapsed_seconds. +- Include latest available train/test losses and validation metrics when present. +- Stop polling when status_type is completed/failed/stopped. + +Error triage checklist: +- Service resolution failure: + - call list_cellpose_services(); pick available websocket service ID. +- Artifact/path mismatch: + - verify artifact workspace and glob patterns. +- Token/workspace mismatch: + - reconnect with correct token/workspace. +- Resource pressure (VRAM/OOM): + - stop competing sessions, retry with smaller n_samples and/or lower load. + +Examples: +1) Start from artifact paths: + await start_cellpose_training( + artifact='workspace/dataset', + train_images='images/*/*.tif', + train_annotations='annotations/*/*_mask.ome.tif', + split_mode='auto', + train_split_ratio=0.8, + n_samples=0.005, + model='cpsam', + n_epochs=1, + learning_rate=1e-5, + min_train_masks=1, + validation_interval=1, + ) + +2) Start from metadata: + await start_cellpose_training( + artifact='workspace/dataset', + metadata_dir='metadata/', + model='cpsam', + n_epochs=1, + ) + +3) BioImage Archive URL: + await start_cellpose_training( + artifact='https://www.ebi.ac.uk/biostudies/bioimages/studies/S-BIAD1392', + split_mode='auto', + train_split_ratio=0.8, + n_samples=0.005, + n_epochs=1, + model='cpsam', + ) + +Schema usage note: +- To inspect exact callable schemas at runtime, run: + await print_cellpose_runtime_docs(include_schemas=True) +""" + +print(RUNTIME_GUIDE) diff --git a/tests/cellpose/live_test_utils.py b/tests/cellpose/live_test_utils.py new file mode 100644 index 00000000..e8daca80 --- /dev/null +++ b/tests/cellpose/live_test_utils.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import asyncio +import os +from typing import Any + +import pytest +from hypha_rpc import connect_to_server + +SERVER_URL = os.environ.get("HYPHA_SERVER_URL", "https://hypha.aicell.io") + + +def is_transient_connection_error(err: Exception) -> bool: + text = str(err).lower() + return ( + "websocket" in text + or "http 404" in text + or "http 502" in text + or "server rejected websocket connection" in text + or "method call timed out" in text + ) + + +def is_unavailable_workspace_error(err: Exception) -> bool: + text = str(err).lower() + return ( + "workspace" in text + and ("does not exist" in text or "not accessible" in text) + ) + + +async def with_live_server( + token: str, + workspace: str, + callback: Any, + *, + retries: int = 6, + method_timeout: int = 240000, +) -> Any: + last_error: Exception | None = None + for attempt in range(retries): + try: + async with connect_to_server( + { + "server_url": SERVER_URL, + "token": token, + "workspace": workspace, + "method_timeout": method_timeout, + } + ) as server: + return await callback(server) + except Exception as e: + last_error = e + if attempt < retries - 1 and is_transient_connection_error(e): + await asyncio.sleep(10) + continue + if is_transient_connection_error(e): + pytest.skip( + "Live Hypha websocket endpoint is temporarily unavailable " + "(HTTP 404/502/timeouts during connection)." + ) + if is_unavailable_workspace_error(e): + pytest.skip( + "Live workspace is currently not accessible for this token/session." + ) + raise + + if last_error is not None: + raise last_error + return None + + +def _supports_split_mode(service: Any) -> bool: + try: + schema = getattr(service.start_training, "__schema__", None) + if isinstance(schema, dict): + params = schema.get("parameters") + if isinstance(params, dict): + return "split_mode" in params.get("properties", {}) + if isinstance(params, list): + return any( + isinstance(p, dict) and p.get("name") == "split_mode" + for p in params + ) + except Exception: + return False + return False + + +async def resolve_cellpose_service( + server: Any, + workspace: str, + *, + application_id: str = "cellpose-finetuning", + requested_service_id: str | None = None, +) -> Any: + if requested_service_id: + service = await server.get_service(requested_service_id) + if _supports_split_mode(service): + return service + + worker_candidates = [ + "bioimage-io/bioengine-worker", + f"{workspace}/bioengine-worker", + ] + for worker_id in worker_candidates: + try: + worker = await server.get_service(worker_id) + status = await worker.get_application_status( + application_ids=[application_id] + ) + except Exception: + continue + + app_data = status.get(application_id, {}) if isinstance(status, dict) else {} + service_entries = app_data.get("service_ids", []) + for entry in reversed(service_entries): + if isinstance(entry, dict): + websocket_id = ( + entry.get("websocket_service_id") + or entry.get("service_id") + or entry.get("id") + ) + else: + websocket_id = entry + + if isinstance(websocket_id, str) and websocket_id: + try: + service = await server.get_service(websocket_id) + if _supports_split_mode(service): + return service + except Exception: + continue + + candidate_ids = [ + os.environ.get("HYPHA_TEST_SERVICE_ID"), + f"{workspace}/{application_id}", + "bioimage-io/cellpose-finetuning", + ] + for candidate in candidate_ids: + if not candidate: + continue + try: + service = await server.get_service(candidate) + if _supports_split_mode(service): + return service + except Exception: + continue + + raise RuntimeError( + "Could not resolve cellpose service id from workspace or worker status. " + "Verify app deployment and websocket service ids in worker status." + ) diff --git a/tests/cellpose/test_metadata_and_glob.py b/tests/cellpose/test_metadata_and_glob.py new file mode 100644 index 00000000..835c5083 --- /dev/null +++ b/tests/cellpose/test_metadata_and_glob.py @@ -0,0 +1,475 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path + +from bioengine_apps.cellpose_finetuning.main import ( + CellposeFinetune, + _extract_bia_accession, + _pair_key_from_url, + list_matching_artifact_paths, + make_training_pairs_from_metadata, + match_image_annotation_pairs, + proportional_manual_sample_counts, + resolve_requested_sample_count, + sample_pair_lists, +) + + +def test_match_image_annotation_pairs_with_nested_globs() -> None: + image_files = [ + "images/a/t0001.tif", + "images/b/t0002.tif", + "images/c/extra.tif", + ] + annotation_files = [ + "annotations/a/t0001_mask.ome.tif", + "annotations/b/t0002_mask.ome.tif", + "annotations/x/other_mask.ome.tif", + ] + + pairs = match_image_annotation_pairs( + image_files, + annotation_files, + "images/*/*.tif", + "annotations/*/*_mask.ome.tif", + ) + + assert pairs == [ + ("images/a/t0001.tif", "annotations/a/t0001_mask.ome.tif"), + ("images/b/t0002.tif", "annotations/b/t0002_mask.ome.tif"), + ] + + +def test_match_image_annotation_pairs_with_mixed_ome_suffix_convention() -> None: + image_files = [ + "images/108bb69d-2e52-4382-8100-e96173db24ee/t0000.ome.tif", + "images/108bb69d-2e52-4382-8100-e96173db24ee/t0001.ome.tif", + ] + annotation_files = [ + "annotations/108bb69d-2e52-4382-8100-e96173db24ee/t0000_mask.ome.tif", + "annotations/108bb69d-2e52-4382-8100-e96173db24ee/t0001_mask.ome.tif", + ] + + pairs = match_image_annotation_pairs( + image_files, + annotation_files, + "images/108bb69d-2e52-4382-8100-e96173db24ee/*.tif", + "annotations/108bb69d-2e52-4382-8100-e96173db24ee/*_mask.ome.tif", + ) + + assert pairs == [ + ( + "images/108bb69d-2e52-4382-8100-e96173db24ee/t0000.ome.tif", + "annotations/108bb69d-2e52-4382-8100-e96173db24ee/t0000_mask.ome.tif", + ), + ( + "images/108bb69d-2e52-4382-8100-e96173db24ee/t0001.ome.tif", + "annotations/108bb69d-2e52-4382-8100-e96173db24ee/t0001_mask.ome.tif", + ), + ] + + +def test_match_image_annotation_pairs_with_many_nested_folders() -> None: + image_files = [ + "images/folder1/sub1/img001.tif", + "images/folder2/sub2/img002.tif", + ] + annotation_files = [ + "annotations/folder1/sub1/img001_mask.ome.tif", + "annotations/folder2/sub2/img002_mask.ome.tif", + ] + + pairs = match_image_annotation_pairs( + image_files, + annotation_files, + "images/*/*.tif", + "annotations/*/*_mask.ome.tif", + ) + + assert pairs == [ + ( + "images/folder1/sub1/img001.tif", + "annotations/folder1/sub1/img001_mask.ome.tif", + ), + ( + "images/folder2/sub2/img002.tif", + "annotations/folder2/sub2/img002_mask.ome.tif", + ), + ] + + +class _FakeArtifact: + async def ls(self, folder_path: str): + await asyncio.sleep(0) + folder = folder_path.rstrip("/") + "/" + if folder == "metadata/": + return [{"path": "metadata/sample.json", "type": "file"}] + if folder == "images/train/": + return [ + {"path": "images/train/t0001.tif", "type": "file"}, + {"path": "images/train/t0002.tif", "type": "file"}, + ] + return [] + + async def get(self, remote_paths, local_paths, on_error="ignore"): + await asyncio.sleep(0) + for remote, local in zip(remote_paths, local_paths): + local_path = Path(local) + local_path.parent.mkdir(parents=True, exist_ok=True) + + if str(remote).endswith("metadata/sample.json"): + local_path.write_text( + """ +[ + {"image_path": "images/train/t0001.tif", "mask_path": "annotations/train/t0001_mask.ome.tif", "split": "train"}, + {"image_path": "images/test/t0002.tif", "mask_path": "annotations/test/t0002_mask.ome.tif", "split": "test"} +] + """.strip(), + encoding="utf-8", + ) + else: + local_path.write_bytes(b"dummy") + + +class _FakeArtifactNestedMetadata: + async def ls(self, folder_path: str): + await asyncio.sleep(0) + folder = folder_path.rstrip("/") + "/" + if folder == "metadata/": + return [{"path": "metadata/records.json", "type": "file"}] + return [] + + async def get(self, remote_paths, local_paths, on_error="ignore"): + await asyncio.sleep(0) + for remote, local in zip(remote_paths, local_paths): + local_path = Path(local) + local_path.parent.mkdir(parents=True, exist_ok=True) + + if str(remote).endswith("metadata/records.json"): + local_path.write_text( + """ +{ + "payload": { + "items": [ + { + "imagePath": "images/train/t0010.tif", + "maskPath": "annotations/train/t0010_mask.ome.tif", + "split": "train" + }, + { + "source": {"path": "images/test/t0011.tif"}, + "target": {"path": "annotations/test/t0011_mask.ome.tif"}, + "subset": "validation" + } + ] + } +} + """.strip(), + encoding="utf-8", + ) + else: + local_path.write_bytes(b"dummy") + + +class _FakeArtifactAmbiguousDirs: + async def ls(self, folder_path: str): + await asyncio.sleep(0) + folder = folder_path.rstrip("/") + "/" + if folder == "images/": + return [{"path": "images/108bb69d-2e52-4382-8100-e96173db24ee"}] + if folder == "images/108bb69d-2e52-4382-8100-e96173db24ee/": + return [ + { + "path": "images/108bb69d-2e52-4382-8100-e96173db24ee/t0000.ome.tif", + "type": "file", + }, + { + "path": "images/108bb69d-2e52-4382-8100-e96173db24ee/t0001.ome.tif", + "type": "file", + }, + ] + return [] + + async def get(self, remote_paths, local_paths, on_error="ignore"): + await asyncio.sleep(0) + for local in local_paths: + local_path = Path(local) + local_path.parent.mkdir(parents=True, exist_ok=True) + local_path.write_bytes(b"dummy") + + +class _FakeArtifactFolderGlobDirs: + async def ls(self, folder_path: str): + await asyncio.sleep(0) + folder = folder_path.rstrip("/") + "/" + if folder == "images/": + return [ + {"path": "images/sample-a"}, + {"path": "images/sample-b"}, + {"path": "images/readme.txt", "type": "file"}, + ] + if folder in {"images/sample-a/", "images/sample-b/"}: + return [ + {"path": f"{folder}0", "type": "file"}, + {"path": f"{folder}1", "type": "file"}, + ] + return [] + + async def get(self, remote_paths, local_paths, on_error="ignore"): + await asyncio.sleep(0) + for local in local_paths: + local_path = Path(local) + local_path.parent.mkdir(parents=True, exist_ok=True) + local_path.write_bytes(b"dummy") + + +def test_make_training_pairs_from_metadata(tmp_path: Path) -> None: + artifact = _FakeArtifact() + train_pairs, test_pairs = asyncio.run( + make_training_pairs_from_metadata( + artifact=artifact, + metadata_dir="metadata/", + save_path=tmp_path, + n_samples=None, + ) + ) + + assert len(train_pairs) == 1 + assert len(test_pairs) == 1 + assert train_pairs[0]["image"].exists() + assert train_pairs[0]["annotation"].exists() + assert test_pairs[0]["image"].exists() + assert test_pairs[0]["annotation"].exists() + + +def test_list_matching_artifact_paths_directory_without_trailing_slash() -> None: + artifact = _FakeArtifact() + matches = asyncio.run(list_matching_artifact_paths(artifact, "images/train")) + assert matches == [ + "images/train/t0001.tif", + "images/train/t0002.tif", + ] + + +def test_list_matching_artifact_paths_with_ambiguous_directory_entries() -> None: + artifact = _FakeArtifactAmbiguousDirs() + matches = asyncio.run(list_matching_artifact_paths(artifact, "images/*/*.tif")) + assert matches == [ + "images/108bb69d-2e52-4382-8100-e96173db24ee/t0000.ome.tif", + "images/108bb69d-2e52-4382-8100-e96173db24ee/t0001.ome.tif", + ] + + +def test_list_matching_artifact_paths_folder_glob_returns_directories() -> None: + artifact = _FakeArtifactFolderGlobDirs() + matches = asyncio.run(list_matching_artifact_paths(artifact, "images/*/")) + assert matches == [ + "images/sample-a/", + "images/sample-b/", + ] + + +def test_make_training_pairs_from_nested_metadata_formats(tmp_path: Path) -> None: + artifact = _FakeArtifactNestedMetadata() + train_pairs, test_pairs = asyncio.run( + make_training_pairs_from_metadata( + artifact=artifact, + metadata_dir="metadata/", + save_path=tmp_path, + n_samples=None, + ) + ) + + assert len(train_pairs) == 1 + assert len(test_pairs) == 1 + assert train_pairs[0]["image"].as_posix().endswith("images/train/t0010.tif") + assert ( + train_pairs[0]["annotation"] + .as_posix() + .endswith("annotations/train/t0010_mask.ome.tif") + ) + assert test_pairs[0]["image"].as_posix().endswith("images/test/t0011.tif") + assert ( + test_pairs[0]["annotation"] + .as_posix() + .endswith("annotations/test/t0011_mask.ome.tif") + ) + + +def test_extract_bia_accession_from_gallery_url() -> None: + url = "https://beta.bioimagearchive.org/bioimage-archive/galleries/ai/ai-ready-study/S-BIAD1392" + assert _extract_bia_accession(url) == "S-BIAD1392" + + +def test_pair_key_from_url_handles_mask_suffix() -> None: + img = "https://example.org/data/images/folder1/t0001.ome.tif" + ann = "https://example.org/data/annotations/folder1/t0001_mask.ome.tif" + assert _pair_key_from_url(img, is_mask=False) == _pair_key_from_url( + ann, is_mask=True + ) + + +def test_resolve_requested_sample_count_supports_decimal_fraction() -> None: + assert resolve_requested_sample_count(0.5, 10) == 5 + assert resolve_requested_sample_count(0.01, 10) == 1 + assert resolve_requested_sample_count(1.0, 10) == 10 + + +def test_resolve_requested_sample_count_supports_absolute_count() -> None: + assert resolve_requested_sample_count(4, 10) == 4 + assert resolve_requested_sample_count(20, 10) == 10 + + +def test_resolve_requested_sample_count_rejects_non_positive_values() -> None: + try: + resolve_requested_sample_count(0, 10) + assert False, "Expected ValueError for n_samples=0" + except ValueError as e: + assert "n_samples must be > 0" in str(e) + + try: + resolve_requested_sample_count(-0.2, 10) + assert False, "Expected ValueError for n_samples<0" + except ValueError as e: + assert "n_samples must be > 0" in str(e) + + +def test_proportional_manual_sample_counts_basic_allocation() -> None: + train_count, test_count = proportional_manual_sample_counts( + train_available=8, + test_available=2, + requested_total=5, + ) + assert train_count + test_count == 5 + assert train_count == 4 + assert test_count == 1 + + +def test_proportional_manual_sample_counts_respects_capacity_limits() -> None: + train_count, test_count = proportional_manual_sample_counts( + train_available=1, + test_available=9, + requested_total=6, + ) + assert train_count + test_count == 6 + assert train_count == 1 + assert test_count == 5 + + +def test_sample_pair_lists_auto_mode_subsets_and_splits() -> None: + train_pairs = [ + {"image": Path(f"images/t{i}.tif"), "annotation": Path(f"annotations/t{i}.tif")} + for i in range(10) + ] + test_pairs = [] + + sampled_train, sampled_test = sample_pair_lists( + train_pairs, + test_pairs, + requested_total=4, + split_mode="auto", + train_split_ratio=0.75, + ) + + assert len(sampled_train) + len(sampled_test) == 4 + assert len(sampled_train) == 3 + assert len(sampled_test) == 1 + + +def test_sample_pair_lists_manual_mode_proportional() -> None: + train_pairs = [ + {"image": Path(f"images/t{i}.tif"), "annotation": Path(f"annotations/t{i}.tif")} + for i in range(8) + ] + test_pairs = [ + {"image": Path(f"images/v{i}.tif"), "annotation": Path(f"annotations/v{i}.tif")} + for i in range(2) + ] + + sampled_train, sampled_test = sample_pair_lists( + train_pairs, + test_pairs, + requested_total=5, + split_mode="manual", + train_split_ratio=0.8, + ) + + assert len(sampled_train) + len(sampled_test) == 5 + assert len(sampled_train) == 4 + assert len(sampled_test) == 1 + + +def test_preflight_training_dataset_reports_pair_counts(monkeypatch) -> None: + class _FakeArtifactForPreflight: + async def ls(self, _folder_path: str): + await asyncio.sleep(0) + return [] + + async def _fake_make_artifact_client(_artifact_id: str, _server_url: str): + return _FakeArtifactForPreflight() + + async def _fake_list_matching_paths(_artifact, path_pattern: str, max_results=None): + if "images" in path_pattern: + return [f"images/t{i}.tif" for i in range(6)] + if "annotations" in path_pattern: + return [f"annotations/t{i}_mask.tif" for i in range(5)] + return [] + + monkeypatch.setattr( + "bioengine_apps.cellpose_finetuning.main.make_artifact_client", + _fake_make_artifact_client, + ) + monkeypatch.setattr( + "bioengine_apps.cellpose_finetuning.main.list_matching_artifact_paths", + _fake_list_matching_paths, + ) + + service = CellposeFinetune.func_or_class() + result = asyncio.run( + service.preflight_training_dataset( + artifact="ri-scale/zarr-demo", + train_images="images/*/", + train_annotations="annotations/*/", + split_mode="auto", + n_samples=5, + ) + ) + + assert result["ok"] is True + assert result["train_image_count"] == 6 + assert result["train_annotation_count"] == 5 + assert result["train_pair_count"] == 5 + assert result["sampled_total_count"] == 5 + + +def test_preflight_training_dataset_metadata_mode(monkeypatch) -> None: + class _FakeArtifactWithMetadata: + async def ls(self, folder_path: str): + await asyncio.sleep(0) + if folder_path == "metadata/": + return [ + {"path": "metadata/records.json", "type": "file"}, + {"path": "metadata/readme.txt", "type": "file"}, + ] + return [] + + async def _fake_make_artifact_client(_artifact_id: str, _server_url: str): + return _FakeArtifactWithMetadata() + + monkeypatch.setattr( + "bioengine_apps.cellpose_finetuning.main.make_artifact_client", + _fake_make_artifact_client, + ) + + service = CellposeFinetune.func_or_class() + result = asyncio.run( + service.preflight_training_dataset( + artifact="ri-scale/zarr-demo", + metadata_dir="metadata/", + ) + ) + + assert result["ok"] is True + assert result["mode"] == "metadata" + assert result["metadata_file_count"] == 1 diff --git a/tests/cellpose/test_service.py b/tests/cellpose/test_service.py index 8df44ced..a0dc84e7 100644 --- a/tests/cellpose/test_service.py +++ b/tests/cellpose/test_service.py @@ -39,12 +39,17 @@ import asyncio import logging import os +import sys from datetime import datetime, timezone +from pathlib import Path from typing import TYPE_CHECKING # from dotenv import load_dotenv from hypha_rpc import connect_to_server, login +sys.path.append(str(Path(__file__).resolve().parent)) +from live_test_utils import resolve_cellpose_service + if TYPE_CHECKING: from hypha_rpc.rpc import RemoteService @@ -67,9 +72,22 @@ async def get_cellpose_service( resolved_service_id = service_id or f"{workspace}/{application_id}" logger.info("Using cellpose service id: %s", resolved_service_id) - cellpose_service = await server.get_service(resolved_service_id) - logger.info("Obtained Cellpose Fine-Tuning service") - return cellpose_service + try: + cellpose_service = await resolve_cellpose_service( + server, + workspace, + application_id=application_id, + requested_service_id=resolved_service_id, + ) + logger.info("Obtained Cellpose Fine-Tuning service") + return cellpose_service + except Exception as original_error: + msg = ( + f"Could not resolve Cellpose service '{resolved_service_id}'. " + "Verify deployment and websocket service ids in " + "worker.get_application_status(...)[application_id]['service_ids']." + ) + raise RuntimeError(msg) from original_error async def infer( @@ -318,7 +336,9 @@ async def start_training( test_images: str | None, test_annotations: str | None, n_epochs: int, - n_samples: int | None, + n_samples: float | None, + split_mode: str, + train_split_ratio: float, learning_rate: float, min_train_masks: int, validation_interval: int | None, @@ -333,6 +353,8 @@ async def start_training( "learning_rate": learning_rate, "n_samples": n_samples, "min_train_masks": min_train_masks, + "split_mode": split_mode, + "train_split_ratio": train_split_ratio, } if validation_interval is not None: kwargs["validation_interval"] = validation_interval @@ -379,7 +401,19 @@ def _parse_args() -> argparse.Namespace: p.add_argument("--test-annotations", default=None) p.add_argument("--n-epochs", type=int, default=10) - p.add_argument("--n-samples", type=int, default=None) + p.add_argument("--n-samples", type=float, default=None) + p.add_argument( + "--split-mode", + default="manual", + choices=["manual", "auto"], + help="Dataset split mode for training.", + ) + p.add_argument( + "--train-split-ratio", + type=float, + default=0.8, + help="Train split ratio when --split-mode auto.", + ) p.add_argument("--learning-rate", type=float, default=1e-5) p.add_argument("--min-train-masks", type=int, default=5) p.add_argument( @@ -471,6 +505,8 @@ async def main() -> None: test_annotations=args.test_annotations, n_epochs=args.n_epochs, n_samples=args.n_samples, + split_mode=args.split_mode, + train_split_ratio=args.train_split_ratio, learning_rate=args.learning_rate, min_train_masks=args.min_train_masks, validation_interval=args.validation_interval, diff --git a/tests/cellpose/test_training_e2e_small_percentage.py b/tests/cellpose/test_training_e2e_small_percentage.py new file mode 100644 index 00000000..fddf9654 --- /dev/null +++ b/tests/cellpose/test_training_e2e_small_percentage.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import asyncio +import os +import sys +import time +from pathlib import Path + +import pytest +from hypha_rpc import connect_to_server + +sys.path.append(str(Path(__file__).resolve().parent)) +from live_test_utils import resolve_cellpose_service + +SERVER_URL = os.environ.get("HYPHA_SERVER_URL", "https://hypha.aicell.io") +RUN_LIVE_REGRESSION_TESTS = os.environ.get("RUN_LIVE_REGRESSION_TESTS") == "1" + + +def _require_live_env() -> tuple[str, str, str]: + if not RUN_LIVE_REGRESSION_TESTS: + pytest.skip("Set RUN_LIVE_REGRESSION_TESTS=1 to run live e2e tests") + + token = os.environ.get("HYPHA_TOKEN") + if not token: + pytest.skip("HYPHA_TOKEN is required for live e2e tests") + + workspace = os.environ.get("HYPHA_TEST_WORKSPACE", "ri-scale") + service_id = os.environ.get("HYPHA_TEST_SERVICE_ID", "bioimage-io/cellpose-finetuning") + return token, workspace, service_id + + +@pytest.mark.asyncio +async def test_small_percentage_auto_split_progresses() -> None: + token, workspace, service_id = _require_live_env() + + async with connect_to_server( + {"server_url": SERVER_URL, "token": token, "workspace": workspace} + ) as server: + service = await resolve_cellpose_service( + server, + workspace, + requested_service_id=service_id, + ) + + started = await service.start_training( + artifact="ri-scale/zarr-demo", + train_images="images/*/", + train_annotations="annotations/*/", + split_mode="auto", + train_split_ratio=0.8, + n_samples=0.02, + n_epochs=2, + min_train_masks=1, + ) + + session_id = started["session_id"] + assert session_id + + reached_non_preparing = False + deadline = time.time() + 240 + + try: + while time.time() < deadline: + status = await service.get_training_status(session_id) + status_type = str(status.get("status_type", "")).lower() + + if status_type != "preparing": + reached_non_preparing = True + break + + await asyncio.sleep(3) + finally: + latest = await service.get_training_status(session_id) + latest_type = str(latest.get("status_type", "")).lower() + if latest_type in {"preparing", "running", "waiting"}: + await service.stop_training(session_id=session_id) + + assert reached_non_preparing, "Training remained stuck in 'preparing' state" + + +@pytest.mark.asyncio +async def test_infer_pretrained_model_e2e() -> None: + token, workspace, service_id = _require_live_env() + + async with connect_to_server( + {"server_url": SERVER_URL, "token": token, "workspace": workspace} + ) as server: + service = await resolve_cellpose_service( + server, + workspace, + requested_service_id=service_id, + ) + + result = await service.infer( + artifact="ri-scale/zarr-demo", + image_paths=["images/108bb69d-2e52-4382-8100-e96173db24ee/t0000.ome.tif"], + model="cpsam", + diameter=40, + json_safe=True, + ) + + assert isinstance(result, list) + assert len(result) == 1 + output = result[0].get("output") + assert output is not None + assert output.get("encoding") in {"ndarray_base64", "mask_png_base64"} diff --git a/tests/cellpose/test_training_e2e_three_sources.py b/tests/cellpose/test_training_e2e_three_sources.py new file mode 100644 index 00000000..4b55fff5 --- /dev/null +++ b/tests/cellpose/test_training_e2e_three_sources.py @@ -0,0 +1,409 @@ +from __future__ import annotations + +import asyncio +import os +import sys +import uuid +from pathlib import Path +from typing import Any + +import httpx +import pytest + +sys.path.append(str(Path(__file__).resolve().parent)) +from live_test_utils import resolve_cellpose_service, with_live_server + +RUN_LIVE_REGRESSION_TESTS = os.environ.get("RUN_LIVE_REGRESSION_TESTS") == "1" +ROOT_DIR = Path(__file__).resolve().parents[2] +DEMO_DATASET_DIR = ROOT_DIR / "demo-dataset" + + +def _require_live_env() -> tuple[str, str]: + if not RUN_LIVE_REGRESSION_TESTS: + pytest.skip("Set RUN_LIVE_REGRESSION_TESTS=1 to run live e2e tests") + + token = os.environ.get("HYPHA_TOKEN") + if not token: + pytest.skip("HYPHA_TOKEN is required for live e2e tests") + + workspace = os.environ.get("HYPHA_TEST_WORKSPACE", "ri-scale") + return token, workspace + + +async def _upload_demo_dataset(server: Any, workspace: str, local_dir: Path) -> str: + artifact_manager = await server.get_service("public/artifact-manager") + alias = f"demo-dataset-e2e-{uuid.uuid4().hex[:8]}" + + artifact = await artifact_manager.create( + type="dataset", + alias=alias, + manifest={ + "name": alias, + "description": "Temporary demo dataset for cellpose e2e tests", + "type": "dataset", + }, + config={"permissions": {"@": "*"}}, + stage=True, + ) + + all_files = [p for p in local_dir.rglob("*") if p.is_file()] + timeout = httpx.Timeout(120.0) + async with httpx.AsyncClient(timeout=timeout) as client: + for file_path in all_files: + rel_path = file_path.relative_to(local_dir).as_posix() + put_url = await artifact_manager.put_file(artifact.id, file_path=rel_path) + with open(file_path, "rb") as f: + response = await client.put(put_url, content=f.read()) + response.raise_for_status() + + await artifact_manager.commit(artifact.id) + return artifact.id + + +async def _run_training_to_finish( + service: Any, + *, + artifact: str, + train_images: str | None, + train_annotations: str | None, + n_samples: float = 5, + retry_sample_fractions: tuple[float, ...] = (8, 10), + train_split_ratio: float = 0.8, + n_epochs: int = 2, + timeout_seconds: int = 360, +) -> dict[str, Any]: + print( + f"[TRAIN] artifact={artifact} n_samples={n_samples} n_epochs={n_epochs}", + flush=True, + ) + + async def _get_status_with_retries( + session_id: str, + attempts: int = 10, + delay_seconds: float = 5.0, + ) -> dict[str, Any]: + last_error: Exception | None = None + for _ in range(attempts): + try: + return await service.get_training_status(session_id=session_id) + except Exception as e: + last_error = e + text = str(e).lower() + is_timeout = "timed out" in text or isinstance(e, TimeoutError) + if not is_timeout: + raise + await asyncio.sleep(delay_seconds) + if last_error is not None: + return { + "status_type": "preparing", + "message": f"Transient status polling timeout: {last_error}", + } + return {} + + async def _stop_active_sessions() -> None: + try: + sessions = await service.list_training_sessions( + status_types=["waiting", "preparing", "running"] + ) + except Exception: + return + + if not isinstance(sessions, dict): + return + + for existing_session_id, session_status in sessions.items(): + status_type = str((session_status or {}).get("status_type", "")).lower() + if status_type in {"waiting", "preparing", "running"}: + try: + await service.stop_training(session_id=existing_session_id) + except Exception: + continue + + await asyncio.sleep(2) + + candidate_fractions = (n_samples, *retry_sample_fractions) + last_status: dict[str, Any] = {} + + for sample_fraction in candidate_fractions: + for _network_retry in range(2): + print( + f"[TRAIN] attempt sample_fraction={sample_fraction} network_retry={_network_retry}", + flush=True, + ) + await _stop_active_sessions() + + kwargs: dict[str, Any] = { + "artifact": artifact, + "split_mode": "auto", + "train_split_ratio": train_split_ratio, + "n_samples": sample_fraction, + "n_epochs": n_epochs, + "validation_interval": 1, + "min_train_masks": 0, + } + if train_images is not None: + kwargs["train_images"] = train_images + if train_annotations is not None: + kwargs["train_annotations"] = train_annotations + + try: + started = await service.start_training(**kwargs) + session_id = started["session_id"] + print( + f"[TRAIN] started session_id={session_id} status={started.get('status_type')} msg={started.get('message')}", + flush=True, + ) + + started_type = str(started.get("status_type", "")).lower() + started_message = str(started.get("message", "")).lower() + deferred_start = ( + started_type == "stopped" + and "deferred start" in started_message + and "gpu contention" in started_message + ) + if deferred_start: + restarted = await service.restart_training( + session_id=session_id, n_epochs=n_epochs + ) + session_id = restarted["session_id"] + print( + f"[TRAIN] restarted deferred session -> {session_id}", + flush=True, + ) + + assert session_id + + deadline = asyncio.get_running_loop().time() + timeout_seconds + latest_status: dict[str, Any] = {} + + try: + while asyncio.get_running_loop().time() < deadline: + latest_status = await _get_status_with_retries( + session_id=session_id + ) + status_type = str(latest_status.get("status_type", "")).lower() + print( + f"[TRAIN] session_id={session_id} status={status_type} msg={latest_status.get('message')}", + flush=True, + ) + if status_type in {"completed", "failed", "stopped"}: + break + await asyncio.sleep(3) + finally: + status_type = str(latest_status.get("status_type", "")).lower() + if status_type in {"waiting", "preparing", "running"}: + try: + await service.stop_training(session_id=session_id) + except Exception: + pass + + latest_status = await _get_status_with_retries(session_id=session_id) + last_status = latest_status + final_type = str(latest_status.get("status_type", "")).lower() + print( + f"[TRAIN] final session_id={session_id} status={final_type} msg={latest_status.get('message')}", + flush=True, + ) + if final_type == "completed": + latest_status["n_samples_used"] = sample_fraction + latest_status["session_id"] = session_id + return latest_status + + message = str(latest_status.get("message") or "").lower() + retriable = ( + "no training samples available" in message + or "no training pairs found" in message + or "float division by zero" in message + ) + if not retriable: + break + except Exception as e: + err = str(e).lower() + transient = ( + "websocket" in err + or "http 404" in err + or "http 502" in err + or "method call timed out" in err + ) + if transient and _network_retry == 0: + await asyncio.sleep(5) + continue + raise + + final_type = str(last_status.get("status_type", "")).lower() + assert final_type == "completed", ( + f"Training did not complete for artifact={artifact}. " + f"final status={final_type}, message={last_status.get('message')}" + ) + return last_status + + +async def _infer_with_session_model( + service: Any, + *, + session_id: str, + artifact: str, + image_path: str, +) -> dict[str, Any]: + print( + f"[INFER] model_session={session_id} artifact={artifact} image={image_path}", + flush=True, + ) + result = await service.infer( + model=session_id, + artifact=artifact, + image_paths=[image_path], + diameter=40, + json_safe=True, + ) + assert isinstance(result, list) + assert len(result) == 1 + output = result[0].get("output") + assert output is not None + assert output.get("encoding") in {"ndarray_base64", "mask_png_base64"} + print( + f"[INFER] ok model_session={session_id} encoding={output.get('encoding')}", + flush=True, + ) + return result[0] + + +@pytest.mark.asyncio +async def test_training_bia_url_small_percentage_auto_split() -> None: + token, workspace = _require_live_env() + + async def _run(server: Any) -> None: + service = await resolve_cellpose_service(server, workspace) + await _run_training_to_finish( + service, + artifact="https://www.ebi.ac.uk/biostudies/bioimages/studies/S-BIAD1392", + train_images="images/*/", + train_annotations="annotations/*/", + n_samples=0.005, + train_split_ratio=0.8, + timeout_seconds=420, + ) + + await with_live_server(token, workspace, _run) + + +@pytest.mark.asyncio +async def test_training_ri_scale_zarr_demo_small_percentage_auto_split() -> None: + token, workspace = _require_live_env() + + async def _run(server: Any) -> None: + service = await resolve_cellpose_service(server, workspace) + await _run_training_to_finish( + service, + artifact="ri-scale/zarr-demo", + train_images="images/108bb69d-2e52-4382-8100-e96173db24ee/*.ome.tif", + train_annotations="annotations/108bb69d-2e52-4382-8100-e96173db24ee/*_mask.ome.tif", + n_samples=0.005, + train_split_ratio=0.8, + timeout_seconds=1200, + ) + + await with_live_server(token, workspace, _run) + + +@pytest.mark.asyncio +async def test_training_local_demo_dataset_small_percentage_auto_split() -> None: + token, workspace = _require_live_env() + + if not DEMO_DATASET_DIR.exists(): + pytest.skip(f"demo dataset not found at {DEMO_DATASET_DIR}") + + async def _run(server: Any) -> None: + service = await resolve_cellpose_service(server, workspace) + artifact_id = await _upload_demo_dataset(server, workspace, DEMO_DATASET_DIR) + + await _run_training_to_finish( + service, + artifact=artifact_id, + train_images="images/*/*.ome.tif", + train_annotations="annotations/*/*_mask.ome.tif", + n_samples=0.005, + train_split_ratio=0.8, + timeout_seconds=1200, + ) + + await with_live_server(token, workspace, _run) + + +@pytest.mark.asyncio +async def test_three_sources_sequential_with_inference() -> None: + token, workspace = _require_live_env() + + async def _run(server: Any) -> None: + service = await resolve_cellpose_service(server, workspace) + + if not DEMO_DATASET_DIR.exists(): + pytest.skip(f"demo dataset not found at {DEMO_DATASET_DIR}") + + shared_infer_artifact = "ri-scale/zarr-demo" + shared_infer_image = "images/108bb69d-2e52-4382-8100-e96173db24ee/t0000.ome.tif" + + # 1) Artifact source (pre-existing artifact id) using uploaded demo dataset + print("[E2E] Source 1/3: artifact-id source", flush=True) + artifact_source_id = await _upload_demo_dataset(server, workspace, DEMO_DATASET_DIR) + artifact_status = await _run_training_to_finish( + service, + artifact=artifact_source_id, + train_images="images/*/*.ome.tif", + train_annotations="annotations/*/*_mask.ome.tif", + n_samples=5, + n_epochs=2, + timeout_seconds=1200, + ) + artifact_session_id = str(artifact_status.get("session_id") or "") + assert artifact_session_id + await _infer_with_session_model( + service, + session_id=artifact_session_id, + artifact=shared_infer_artifact, + image_path=shared_infer_image, + ) + + # 2) BioImage Archive URL source + print("[E2E] Source 2/3: BioImage Archive URL source", flush=True) + bia_status = await _run_training_to_finish( + service, + artifact="https://www.ebi.ac.uk/biostudies/bioimages/studies/S-BIAD1392", + train_images="images/*/", + train_annotations="annotations/*/", + n_samples=5, + n_epochs=2, + timeout_seconds=1200, + ) + bia_session_id = str(bia_status.get("session_id") or "") + assert bia_session_id + await _infer_with_session_model( + service, + session_id=bia_session_id, + artifact=shared_infer_artifact, + image_path=shared_infer_image, + ) + + # 3) Uploaded local demo dataset source + print("[E2E] Source 3/3: uploaded local dataset source", flush=True) + uploaded_artifact_id = await _upload_demo_dataset(server, workspace, DEMO_DATASET_DIR) + uploaded_status = await _run_training_to_finish( + service, + artifact=uploaded_artifact_id, + train_images="images/*/*.ome.tif", + train_annotations="annotations/*/*_mask.ome.tif", + n_samples=5, + n_epochs=2, + timeout_seconds=1200, + ) + uploaded_session_id = str(uploaded_status.get("session_id") or "") + assert uploaded_session_id + await _infer_with_session_model( + service, + session_id=uploaded_session_id, + artifact=shared_infer_artifact, + image_path=shared_infer_image, + ) + print("[E2E] Completed all three sources with inference", flush=True) + + await with_live_server(token, workspace, _run) diff --git a/tests/cellpose/test_ui_e2e.py b/tests/cellpose/test_ui_e2e.py index ef1a576a..52bdcc86 100644 --- a/tests/cellpose/test_ui_e2e.py +++ b/tests/cellpose/test_ui_e2e.py @@ -8,18 +8,22 @@ SERVER_URL = "https://hypha.aicell.io" APP_WORKSPACE = os.environ.get("HYPHA_TEST_WORKSPACE", "ri-scale") -APP_URL = f"{SERVER_URL}/{APP_WORKSPACE}/view/cellpose-finetuning" +APP_ID = os.environ.get("HYPHA_TEST_APP_ID", "cellpose-finetuning") +APP_URL = f"{SERVER_URL}/{APP_WORKSPACE}/view/{APP_ID}" +APP_FALLBACK_URL = f"{SERVER_URL}/{APP_WORKSPACE}/view/cellpose-finetuning-test" RUN_LIVE_REGRESSION_TESTS = os.environ.get("RUN_LIVE_REGRESSION_TESTS") == "1" def open_and_connect(page: Page) -> None: token = os.environ.get("HYPHA_TOKEN") if token: - page.add_init_script( - f"window.localStorage.setItem('hypha_token', {token!r});" - ) + page.add_init_script(f"window.localStorage.setItem('hypha_token', {token!r});") response = page.goto(APP_URL, timeout=60000) - assert response and response.ok, "Failed to load app" + if not (response and response.ok) and APP_URL != APP_FALLBACK_URL: + response = page.goto(APP_FALLBACK_URL, timeout=60000) + assert ( + response and response.ok + ), f"Failed to load app at {APP_URL} (fallback attempted: {APP_FALLBACK_URL})" expect(page).to_have_title(re.compile(r"Cellpose Fine-Tuning"), timeout=15000) expect(page.locator("text=Connected").first).to_be_visible(timeout=30000) @@ -34,57 +38,82 @@ def _handle(dialog): page.on("dialog", _handle) return captured + +def capture_console_logs(page: Page, marker: str): + captured = [] + + def _on_console(msg): + try: + text = msg.text + except Exception: + text = str(msg) + if marker in text: + captured.append(text) + + page.on("console", _on_console) + return captured + + def test_app_loads(page: Page): """App loads and connects to Hypha.""" open_and_connect(page) + def test_navigation(page: Page): """Basic tab navigation works.""" open_and_connect(page) - + # Default is Dashboard expect(page.locator("h2:has-text('Dashboard')")).to_be_visible() - + # Click New Training page.click("text=New Training") expect(page.locator("h2:has-text('New Training Session')")).to_be_visible() - + # Click Inference page.click("text=Inference") expect(page.locator("h2:has-text('Live Inference')")).to_be_visible() + def test_file_browser(page: Page): """File browser opens and lists artifact content.""" open_and_connect(page) - + # Go to Training page.click("text=New Training") - + artifact_id = "ri-scale/cellpose-finetuning" page.fill("input[placeholder='workspace/dataset-alias']", artifact_id) - + # Open Browser page.click("button[title='Browse Files']") - + # Wait for modal expect(page.locator("text=Artifact Explorer")).to_be_visible() - + # Wait for loading to finish expect(page.locator("text=Loading files...")).not_to_be_visible() - + # Should see "index.html" expect(page.locator("text=index.html")).to_be_visible() - + page.locator("div:has(h3:has-text('Artifact Explorer')) button").last.click() - + + def test_start_training_button_and_result(page: Page): """Clicks Start Training and validates the returned UI result.""" open_and_connect(page) page.click("text=New Training") page.fill("input[placeholder='workspace/dataset-alias']", "ri-scale/zarr-demo") - page.fill("input[placeholder='e.g. images/*/*.tif']", "images/108bb69d-2e52-4382-8100-e96173db24ee/*.ome.tif") - page.fill("input[placeholder='e.g. annotations/*/*_mask.ome.tif']", "annotations/108bb69d-2e52-4382-8100-e96173db24ee/*_mask.ome.tif") + page.fill( + "input[placeholder='e.g. images/*/*.tif']", + "images/108bb69d-2e52-4382-8100-e96173db24ee/*.ome.tif", + ) + page.fill( + "input[placeholder='e.g. annotations/*/*_mask.ome.tif']", + "annotations/108bb69d-2e52-4382-8100-e96173db24ee/*_mask.ome.tif", + ) page.fill("input[placeholder='Optional']", "") page.locator("input[placeholder='Optional']").nth(1).fill("") page.fill("input[placeholder='Pretrained ID or Session ID']", "cpsam") @@ -107,13 +136,17 @@ def test_infer_button_and_result(page: Page): open_and_connect(page) page.click("text=Inference") - image_path = Path("/Users/hugokallander/github-repos/bioengine-worker/tests/cellpose_legacy_scripts/t0000.ome.tif") + image_path = Path( + "/Users/hugokallander/github-repos/bioengine-worker/tests/cellpose_legacy_scripts/t0000.ome.tif" + ) assert image_path.exists(), f"Inference image not found: {image_path}" infer_input = page.locator("input[type='file']").first infer_input.set_input_files(str(image_path)) expect(page.locator("text=t0000.ome.tif")).to_be_visible(timeout=10000) - expect(page.locator("img[src^='blob:'], img[src^='data:image/png;base64']").first).to_be_visible(timeout=10000) + expect( + page.locator("img[src^='blob:'], img[src^='data:image/png;base64']").first + ).to_be_visible(timeout=10000) dialogs = capture_dialog(page) page.click("button:has-text('Run Segmentation')") @@ -129,7 +162,9 @@ def test_infer_button_and_result(page: Page): def test_no_raw_template_markers(page: Page): """The rendered UI should not leak template raw markers.""" if not RUN_LIVE_REGRESSION_TESTS: - pytest.skip("Set RUN_LIVE_REGRESSION_TESTS=1 to run live template-marker regression") + pytest.skip( + "Set RUN_LIVE_REGRESSION_TESTS=1 to run live template-marker regression" + ) open_and_connect(page) html = page.content() @@ -140,14 +175,22 @@ def test_no_raw_template_markers(page: Page): def test_stop_training_persists_after_reload(page: Page): """Start then stop a session, and ensure status does not revert to running after refresh/reload.""" if not RUN_LIVE_REGRESSION_TESTS: - pytest.skip("Set RUN_LIVE_REGRESSION_TESTS=1 to run live stop-persistence regression") + pytest.skip( + "Set RUN_LIVE_REGRESSION_TESTS=1 to run live stop-persistence regression" + ) open_and_connect(page) page.click("text=New Training") page.fill("input[placeholder='workspace/dataset-alias']", "ri-scale/zarr-demo") - page.fill("input[placeholder='e.g. images/*/*.tif']", "images/108bb69d-2e52-4382-8100-e96173db24ee/*.ome.tif") - page.fill("input[placeholder='e.g. annotations/*/*_mask.ome.tif']", "annotations/108bb69d-2e52-4382-8100-e96173db24ee/*_mask.ome.tif") + page.fill( + "input[placeholder='e.g. images/*/*.tif']", + "images/108bb69d-2e52-4382-8100-e96173db24ee/*.ome.tif", + ) + page.fill( + "input[placeholder='e.g. annotations/*/*_mask.ome.tif']", + "annotations/108bb69d-2e52-4382-8100-e96173db24ee/*_mask.ome.tif", + ) page.fill("input[placeholder='Optional']", "") page.locator("input[placeholder='Optional']").nth(1).fill("") page.fill("input[placeholder='Pretrained ID or Session ID']", "cpsam") @@ -161,7 +204,9 @@ def test_stop_training_persists_after_reload(page: Page): assert start_dialogs["message"] is not None, "No Start Training dialog received" match = re.search(r"Session ID:\s*([a-zA-Z0-9-]+)", start_dialogs["message"] or "") - assert match is not None, f"Could not parse session id from: {start_dialogs['message']}" + assert ( + match is not None + ), f"Could not parse session id from: {start_dialogs['message']}" session_id = match.group(1) page.click("text=Dashboard") @@ -190,6 +235,162 @@ def test_stop_training_persists_after_reload(page: Page): page.click("text=Dashboard") row_after_reload = page.locator("tr", has=page.locator(f"text={session_id[:8]}")) expect(row_after_reload.first).to_be_visible(timeout=20000) - status_after_reload = row_after_reload.first.locator("td").nth(1).inner_text().strip().lower() - assert "running" not in status_after_reload, f"Status reverted to running: {status_after_reload}" + status_after_reload = ( + row_after_reload.first.locator("td").nth(1).inner_text().strip().lower() + ) + assert ( + "running" not in status_after_reload + ), f"Status reverted to running: {status_after_reload}" + + +def test_dataset_source_options(page: Page): + """Test switching between Artifact, Upload, and BioImage Archive options.""" + open_and_connect(page) + page.click("text=New Training") + + # Locate the dataset source dropdown + # Try finding by label text and getting the select in the parent container + # Or simply finding the select that contains 'BioImage Archive' option + dropdown = page.locator("select:has(option:has-text('BioImage Archive'))") + expect(dropdown).to_be_visible() + + # 1. Default should be Artifact + expect(dropdown).to_have_value("Artifact") + # Verify Artifact ID input is visible + expect(page.locator("input[placeholder='workspace/dataset-alias']")).to_be_visible() + + # 2. Switch to BioImage Archive + dropdown.select_option("BioImage Archive") + # Verify Archive URL input is visible and Artifact ID input is hidden + archive_input = page.locator("input[placeholder*='ebi.ac.uk']") + expect(archive_input).to_be_visible() + expect( + page.locator("input[placeholder='workspace/dataset-alias']") + ).not_to_be_visible() + + # 3. Switch to Upload + dropdown.select_option("Upload") + # Verify Upload button is visible and Archive URL input is hidden + upload_btn = page.locator("button:has-text('Upload Folder')") + expect(upload_btn).to_be_visible() + expect(archive_input).not_to_be_visible() + + # Check Upload Modal opens + upload_btn.click() + expect(page.locator("h3:has-text('Upload Dataset')")).to_be_visible() + page.click("button:has-text('Cancel')") + expect(page.locator("h3:has-text('Upload Dataset')")).not_to_be_visible() + + +def test_inference_panel_not_blank_live(page: Page): + """Inference tab should render configuration controls, not an empty panel.""" + if not RUN_LIVE_REGRESSION_TESTS: + pytest.skip( + "Set RUN_LIVE_REGRESSION_TESTS=1 to run live inference-panel regression" + ) + + inference_logs = capture_console_logs(page, "[InferenceTab]") + open_and_connect(page) + page.click("text=Inference") + + expect(page.locator("h2:has-text('Live Inference')")).to_be_visible(timeout=20000) + + has_config = False + has_init_state = False + for _ in range(8): + page.wait_for_timeout(1500) + has_config = page.locator("h3:has-text('Configuration')").count() > 0 + has_init_state = ( + page.locator("h3:has-text('Initializing inference service')").count() > 0 + ) + if has_config or has_init_state: + break + + if not (has_config or has_init_state): + body_preview = page.locator("body").inner_text()[:1200] + raise AssertionError( + "Inference tab rendered blank (neither configuration nor initializing state visible). " + f"Captured logs: {inference_logs[-12:]}. Body preview: {body_preview}" + ) + + if has_config: + expect(page.locator("button:has-text('Run Segmentation')")).to_be_visible( + timeout=15000 + ) + if has_init_state: + expect(page.locator("button:has-text('Retry Connection')")).to_be_visible( + timeout=15000 + ) + + page.wait_for_timeout(1000) + assert ( + inference_logs + ), "No [InferenceTab] diagnostic logs captured after opening Inference tab" + + +def test_workspace_autocomplete_suggestions_live(page: Page): + """Workspace autocomplete should show options while typing workspace prefix.""" + if not RUN_LIVE_REGRESSION_TESTS: + pytest.skip( + "Set RUN_LIVE_REGRESSION_TESTS=1 to run live workspace-autocomplete regression" + ) + + artifact_logs = capture_console_logs(page, "[ArtifactAutocomplete]") + open_and_connect(page) + page.click("text=New Training") + + artifact_input = page.locator("input[placeholder='workspace/dataset-alias']") + expect(artifact_input).to_be_visible(timeout=15000) + artifact_input.fill(APP_WORKSPACE[:2]) + + suggestions = page.locator( + "input[placeholder='workspace/dataset-alias'] + button + div > div" + ) + try: + expect(suggestions.first).to_be_visible(timeout=30000) + except AssertionError as exc: + raise AssertionError( + "Workspace suggestions were not visible. " + f"Captured diagnostic logs: {artifact_logs[-12:]}" + ) from exc + + first_text = suggestions.first.inner_text().strip() + assert ( + "/" not in first_text + ), f"Expected workspace-only suggestion, got: {first_text}" + + +def test_artifact_autocomplete_suggestions_live(page: Page): + """Artifact autocomplete should show suggestions after entering workspace/ prefix.""" + if not RUN_LIVE_REGRESSION_TESTS: + pytest.skip( + "Set RUN_LIVE_REGRESSION_TESTS=1 to run live artifact-autocomplete regression" + ) + + artifact_logs = capture_console_logs(page, "[ArtifactAutocomplete]") + open_and_connect(page) + page.click("text=New Training") + artifact_input = page.locator("input[placeholder='workspace/dataset-alias']") + expect(artifact_input).to_be_visible(timeout=15000) + + artifact_input.fill(f"{APP_WORKSPACE}/c") + + suggestions = page.locator( + "input[placeholder='workspace/dataset-alias'] + button + div > div" + ) + try: + expect(suggestions.first).to_be_visible(timeout=30000) + except AssertionError as exc: + raise AssertionError( + "Autocomplete suggestions were not visible. " + f"Captured diagnostic logs: {artifact_logs[-12:]}" + ) from exc + + first_text = suggestions.first.inner_text() + assert first_text.startswith( + f"{APP_WORKSPACE}/" + ), f"Unexpected suggestion: {first_text}" + assert ( + artifact_logs + ), "No [ArtifactAutocomplete] diagnostic logs captured during autocomplete flow"