"""Long-lived user API for announcing and publishing experiment data."""
from __future__ import annotations
import threading
from typing import TYPE_CHECKING, Any, Literal, Self
from ._service import artifacts_with_manifests
from .authentication import AuthenticationStatus, LoginPolicy, TokenCredentials
from .errors import (
AuthenticationFailedError,
AuthenticationRequiredError,
AuthenticationUnsupportedError,
ConfigurationError,
MeasurementBackendRequiredError,
PublishError,
QHarborToolsError,
)
from .models import (
AnnouncementRequest,
Artifact,
BackendContext,
ExperimentConfig,
RemoteDataset,
RunReference,
SaveRequest,
)
from .receipts import MeasurementReceipt, Publication, PublicationReceipt
if TYPE_CHECKING:
from collections.abc import Callable, Mapping, Sequence
from typing import TypeVar
from .authentication import AuthenticationBackend
from .backend import MeasurementBackend, RemoteBackend
from .schemas import MeasurementSchema
T = TypeVar("T")
def _default_remote_backends() -> tuple[RemoteBackend, AuthenticationBackend]:
from ._adapters.qdrive import QDriveAuthenticationBackend, QDriveRemoteBackend # ruff: ignore[import-outside-top-level]
return QDriveRemoteBackend(), QDriveAuthenticationBackend()
def _resolve_remote_backends(
remote_backend: RemoteBackend | None,
authentication_backend: AuthenticationBackend | None,
) -> tuple[RemoteBackend, AuthenticationBackend | None]:
if remote_backend is not None:
return remote_backend, authentication_backend
default_remote, default_authentication = _default_remote_backends()
return default_remote, authentication_backend or default_authentication
[docs]
class Experiment:
"""Stable facade over remote publication and optional local persistence."""
def __init__(
self,
config: ExperimentConfig,
*,
measurement_backend: MeasurementBackend | None = None,
remote_backend: RemoteBackend | None = None,
authentication_backend: AuthenticationBackend | None = None,
login_policy: LoginPolicy | str = LoginPolicy.ON_DEMAND,
) -> None:
self.config: ExperimentConfig | None = config
self.measurement_backend = measurement_backend
self.remote_backend, self.authentication_backend = _resolve_remote_backends(
remote_backend,
authentication_backend,
)
self.login_policy = self._coerce_login_policy(login_policy)
self._authentication_lock = threading.Lock()
self._dataset: RemoteDataset | None = None
self._backend_prepared = False
self._closed = False
[docs]
@classmethod
def open(
cls,
uuid: str,
*,
measurement_backend: MeasurementBackend | None = None,
remote_backend: RemoteBackend | None = None,
authentication_backend: AuthenticationBackend | None = None,
login_policy: LoginPolicy | str = LoginPolicy.ON_DEMAND,
) -> Self:
"""Open an existing remote experiment, optionally attaching storage."""
instance = cls.__new__(cls)
instance.config = None
instance.measurement_backend = measurement_backend
instance.remote_backend, instance.authentication_backend = _resolve_remote_backends(
remote_backend,
authentication_backend,
)
instance.login_policy = instance._coerce_login_policy(login_policy)
instance._authentication_lock = threading.Lock()
instance._backend_prepared = False
instance._closed = False
instance._dataset = instance._run_remote("open", lambda: instance.remote_backend.open(uuid))
return instance
@property
def dataset(self) -> RemoteDataset:
"""Return the announced or opened remote dataset."""
return self._require_dataset()
@property
def metadata(self) -> Mapping[str, Any]:
"""Return the current remote dataset attributes."""
return self._require_dataset().attributes
@property
def files(self) -> tuple[str, ...]:
"""Return filenames known when the remote dataset was last refreshed."""
return self._require_dataset().files
[docs]
def authentication_status(self) -> AuthenticationStatus:
"""Return login state without opening an interactive prompt."""
self._ensure_active()
if self.authentication_backend is None:
return AuthenticationStatus.unsupported(self.remote_backend.name)
return self.authentication_backend.status()
[docs]
def login(self) -> AuthenticationStatus:
"""Authenticate interactively, serializing concurrent login attempts."""
self._ensure_active()
backend = self._require_authentication_backend()
with self._authentication_lock:
status = backend.status()
if status.authenticated:
return status
return backend.login()
[docs]
def login_with_token(self, credentials: TokenCredentials) -> AuthenticationStatus:
"""Authenticate with secret-safe token credentials."""
self._ensure_active()
backend = self._require_authentication_backend()
with self._authentication_lock:
status = backend.status()
if status.authenticated:
return status
return backend.login_with_token(credentials)
[docs]
def logout(self) -> AuthenticationStatus:
"""End the shared remote session; repeated calls are safe."""
self._ensure_active()
backend = self._require_authentication_backend()
with self._authentication_lock:
status = backend.status()
if not status.authenticated:
return status
backend.logout()
return backend.status()
[docs]
def announce(self) -> RemoteDataset:
"""Create the configured remote dataset once and return its reference."""
self._ensure_active()
if self._dataset is not None:
return self._dataset
if self.config is None:
raise ConfigurationError("An opened experiment cannot be announced again.", operation="announce")
dataset = self.config.dataset
request = AnnouncementRequest(
name=dataset.name,
description=dataset.description,
scope=dataset.scope,
attributes=dataset.attributes(self.config.user),
)
self._dataset = self._run_remote("announce", lambda: self.remote_backend.announce(request))
return self._dataset
[docs]
def refresh(self) -> RemoteDataset:
"""Reload metadata and filenames from the remote service."""
self._ensure_active()
dataset = self._require_dataset()
self._dataset = self._run_remote("open", lambda: self.remote_backend.open(dataset.uuid))
return self._dataset
[docs]
def publish(
self,
*,
name: str,
data: Any,
metadata: Mapping[str, Any] | None = None,
media_type: str = "application/octet-stream",
) -> PublicationReceipt:
"""Publish existing data directly, without a measurement backend."""
self._ensure_active()
artifact = Artifact(name=name, data=data, media_type=media_type, metadata=metadata or {})
publications = self._publish_artifacts(artifacts_with_manifests((artifact,)))
return PublicationReceipt(
dataset_uuid=self._require_dataset().uuid,
publications=publications,
)
[docs]
def save(
self,
*,
schema: MeasurementSchema,
results: Mapping[str, Any],
metadata: Mapping[str, Any] | None = None,
publish: Sequence[str] | None = None,
) -> MeasurementReceipt:
"""Persist one backend run, export it, and publish its artifacts."""
self._ensure_active()
dataset = self._require_dataset()
backend = self._require_measurement_backend()
schema.validate_results(results)
selected = tuple(publish) if publish is not None else None
if selected is not None:
unknown = set(selected) - set(schema.parameter_names)
if unknown:
raise ConfigurationError(f"Unknown publish variables: {sorted(unknown)}.", operation="validate")
self._prepare_backend()
subject_id = str(dataset.attributes.get("SubjectID", "unknown-subject"))
request = SaveRequest(
experiment_name=dataset.name,
subject_id=subject_id,
schema=schema,
results=results,
metadata=metadata or {},
)
stored_run = backend.save(request)
artifacts = backend.export(stored_run, variables=selected)
publications = self._publish_artifacts(
artifacts_with_manifests(artifacts),
run_reference=stored_run.reference,
)
return MeasurementReceipt(
dataset_uuid=dataset.uuid,
run_reference=stored_run.reference,
publications=publications,
metadata=metadata or {},
)
[docs]
def retry_publish(
self,
recovery: RunReference | PublishError,
*,
publish: Sequence[str] | None = None,
) -> MeasurementReceipt:
"""Republish a preserved backend run without creating a new local run."""
self._ensure_active()
dataset = self._require_dataset()
backend = self._require_measurement_backend()
self._prepare_backend()
if isinstance(recovery, PublishError) and recovery.pending_artifacts:
run_reference = recovery.run_reference
artifacts = recovery.pending_artifacts
completed_publications = recovery.completed_publications
else:
run_reference = recovery.run_reference if isinstance(recovery, PublishError) else recovery
if run_reference is None:
raise ConfigurationError("A run reference is required to retry publication.", operation="retry_publish")
stored_run = backend.load(run_reference)
artifacts = artifacts_with_manifests(backend.export(stored_run, variables=publish))
completed_publications = ()
if run_reference is None:
raise ConfigurationError("A run reference is required to retry publication.", operation="retry_publish")
publications = self._publish_artifacts(
artifacts,
run_reference=run_reference,
completed_publications=completed_publications,
)
return MeasurementReceipt(
dataset_uuid=dataset.uuid,
run_reference=run_reference,
publications=publications,
)
[docs]
def attach_backend(self, backend: MeasurementBackend) -> Self:
"""Attach local measurement storage to a remote-only experiment."""
self._ensure_active()
if self.measurement_backend is backend:
return self
if self.measurement_backend is not None:
self.measurement_backend.close()
self.measurement_backend = backend
self._backend_prepared = False
return self
[docs]
def close(self) -> None:
"""Close attached backends; repeated calls are safe."""
if self._closed:
return
errors: list[QHarborToolsError] = []
if self.measurement_backend is not None:
try:
self.measurement_backend.close()
except QHarborToolsError as error:
errors.append(error)
try:
self.remote_backend.close()
except QHarborToolsError as error:
errors.append(error)
self._closed = True
if errors:
primary = errors[0]
for secondary in errors[1:]:
primary.add_note(str(secondary))
raise primary
def __enter__(self) -> Self:
self._ensure_active()
return self
def __exit__(self, exception_type, exception, traceback) -> Literal[False]:
try:
self.close()
except QHarborToolsError as close_error:
if exception is None:
raise
exception.add_note(f"Cleanup also failed: {close_error}")
return False
def _prepare_backend(self) -> None:
if self._backend_prepared:
return
backend = self._require_measurement_backend()
dataset = self._require_dataset()
backend.prepare(
BackendContext(
experiment_name=dataset.name,
subject_id=str(dataset.attributes.get("SubjectID", "unknown-subject")),
dataset_uuid=dataset.uuid,
)
)
self._backend_prepared = True
def _publish_artifacts(
self,
artifacts: Sequence[Artifact],
*,
run_reference: RunReference | None = None,
completed_publications: Sequence[Publication] = (),
) -> tuple[Publication, ...]:
dataset = self._require_dataset()
completed = list(completed_publications)
for index, artifact in enumerate(artifacts):
try:
completed.append(
self._run_remote(
"publish",
lambda: self.remote_backend.publish(dataset, artifact),
)
)
except QHarborToolsError as error:
raise PublishError(
f"Could not publish artifact '{artifact.name}'.",
operation="publish",
dataset_uuid=dataset.uuid,
run_reference=run_reference,
artifact_name=artifact.name,
retryable=error.retryable,
completed_publications=completed,
pending_artifacts=artifacts[index:],
) from error
return tuple(completed)
def _require_dataset(self) -> RemoteDataset:
if self._dataset is None:
raise ConfigurationError("Call announce() or Experiment.open() before this operation.", operation="state")
return self._dataset
def _require_measurement_backend(self) -> MeasurementBackend:
if self.measurement_backend is None:
raise MeasurementBackendRequiredError(
"This operation requires a measurement backend. Use publish() for existing data or attach_backend() first.",
operation="state",
)
return self.measurement_backend
def _require_authentication_backend(self) -> AuthenticationBackend:
if self.authentication_backend is None:
raise AuthenticationUnsupportedError(
f"Remote backend '{self.remote_backend.name}' has no authentication provider.",
operation="login",
backend_name=self.remote_backend.name,
)
return self.authentication_backend
def _run_remote(self, operation: str, call: Callable[[], T]) -> T:
try:
return call()
except AuthenticationRequiredError:
if self.login_policy is LoginPolicy.EXPLICIT:
raise
self.login()
try:
return call()
except AuthenticationRequiredError as error:
raise AuthenticationFailedError(
f"Authentication did not allow remote operation '{operation}'.",
operation=operation,
backend_name=self.remote_backend.name,
retryable=True,
) from error
@staticmethod
def _coerce_login_policy(policy: LoginPolicy | str) -> LoginPolicy:
try:
return LoginPolicy(policy)
except (TypeError, ValueError) as error:
choices = ", ".join(item.value for item in LoginPolicy)
raise ConfigurationError(
f"login_policy must be one of: {choices}.",
operation="validate",
) from error
def _ensure_active(self) -> None:
if self._closed:
raise ConfigurationError(
"The experiment is closed. Create a new Experiment before rerunning the workflow.",
operation="state",
)