biolm.plugins.mlflow package

Submodules

biolm.plugins.mlflow.datasets module

MLflow operations for dataset management.

This module provides functionality to manage datasets using MLflow runs. Each dataset is represented as an MLflow run with specific tags.

biolm.plugins.mlflow.datasets.download_dataset(dataset_id: str, output_path: str | Path, experiment_name: str | None = None, artifact_path: str | None = None, mlflow_uri: str | None = None) → Dict[str, Any]

Download artifacts from a dataset.

Args:

dataset_id: Dataset identifier (tag value or run_id) output_path: Local directory to save artifacts experiment_name: MLflow experiment name (default: “{username}/datasets”) artifact_path: Optional specific artifact path to download (default: all artifacts) mlflow_uri: Optional MLflow tracking URI

Returns:

Dictionary with download status and path information

biolm.plugins.mlflow.datasets.get_dataset(dataset_id: str, experiment_name: str | None = None, mlflow_uri: str | None = None) → Dict[str, Any] | None

Get a specific dataset by dataset_id tag or run_id.

Args:

dataset_id: Dataset identifier (tag value or run_id) experiment_name: MLflow experiment name (default: “{username}/datasets”) mlflow_uri: Optional MLflow tracking URI

Returns:

Dataset dictionary with run information, or None if not found

biolm.plugins.mlflow.datasets.list_datasets(experiment_name: str | None = None, mlflow_uri: str | None = None, all_runs: bool = False) → List[Dict[str, Any]]

List all datasets (MLflow runs with type: dataset tag).

Args:

experiment_name: MLflow experiment name (default: “{username}/datasets”) mlflow_uri: Optional MLflow tracking URI all_runs: If True, list all runs in experiment, not just those with type: dataset tag

Returns:

List of dataset dictionaries with run information

biolm.plugins.mlflow.datasets.upload_dataset(dataset_id: str, file_path: str | Path, experiment_name: str | None = None, name: str | None = None, mlflow_uri: str | None = None, recursive: bool = False) → Dict[str, Any]

Upload files as artifacts to a dataset (MLflow run).

Args:

dataset_id: Dataset identifier file_path: Path to file or directory to upload experiment_name: MLflow experiment name (default: “{username}/datasets”) name: Optional dataset name/description (stored as run name) mlflow_uri: Optional MLflow tracking URI recursive: If True, upload directory recursively

Returns:

Dictionary with run_id and status information

biolm.plugins.mlflow.protocols module

MLflow logging for protocol execution results.

This module provides functionality to log protocol results to MLflow based on the protocol’s outputs configuration. It handles result selection, template expression evaluation, and MLflow run creation.

exception biolm.plugins.mlflow.protocols.MLflowNotAvailableError

Bases: Exception

Raised when MLflow is not installed.

biolm.plugins.mlflow.protocols.combine_output_rules(results: List[Dict], output_rules: List[Dict]) → Dict[str, Any]

Combine results from multiple output rules.

Args:

results: Full list of result dictionaries output_rules: List of output rule configurations

Returns:

Dictionary mapping result index to combined logging data

biolm.plugins.mlflow.protocols.compute_aggregates(selected_results: List[Dict], all_results: List[Dict], aggregates: List[Dict], aggregate_over: str = 'selected') → Dict[str, Any]

Compute aggregate statistics.

Args:

selected_results: Results selected by output rules all_results: All results from protocol execution aggregates: List of aggregate specifications aggregate_over: “selected” or “all”

Returns:

Dictionary of aggregate metric names to values

biolm.plugins.mlflow.protocols.generate_seqparse(sequence: str, sequence_id: str | None = None, metadata: Dict | None = None) → str

Generate seqparse format JSON for a sequence.

According to https://github.com/Lattice-Automation/seqparse, the format is:

default Copy

{
  "name": string,
  "type": "dna" | "rna" | "aa" | "unknown",
  "seq": string,
  "annotations": Annotation[]
}

Args:

sequence: The sequence string sequence_id: Optional sequence identifier (used as name) metadata: Optional metadata dictionary (may contain annotations)

Returns:

JSON string in seqparse format

biolm.plugins.mlflow.protocols.load_outputs_config(outputs_config: List[Dict] | str | Dict) → List[Dict]

Load outputs configuration from various formats.

Args:

outputs_config: Can be:

• List of output rule dicts

• Path to YAML file containing outputs config

• Path to protocol YAML file (will extract outputs section)

• Protocol dict (will extract outputs section)

Returns:

List of output rule dictionaries

biolm.plugins.mlflow.protocols.load_results(results: List[Dict] | str) → List[Dict]

Load results from a list or JSONL file.

Supports plain JSONL and compressed formats: * .jsonl - plain text * .jsonl.gz - gzip compressed * .zip - zip archive containing a .jsonl file (e.g. results_{id}.jsonl.zip)

Args:

results: Either a list of dicts or a path to a JSONL file (optionally compressed)

Returns:

List of result dictionaries

biolm.plugins.mlflow.protocols.log_protocol_results(results: List[Dict] | str, outputs_config: List[Dict] | str | Dict, account_name: str, workspace_name: str, protocol_name: str, protocol_metadata: Dict | None = None, mlflow_uri: str | None = None, dry_run: bool = False, aggregate_over: str = 'selected') → Dict[str, Any]

Main entry point for logging protocol results to MLflow.

The MLflow experiment name is built as “{account_name}/{workspace_name}/{protocol_name}”.

Args:

results: List of result dicts or path to JSONL file outputs_config: Outputs config (list, dict, or file path) account_name: Account name for the experiment path workspace_name: Workspace name for the experiment path protocol_name: Protocol name (slug) for the experiment path protocol_metadata: Optional protocol metadata (name, version, inputs, etc.) mlflow_uri: Optional MLflow tracking URI (default: https://mlflow.biolm.ai/) dry_run: If True, prepare data but don’t log to MLflow aggregate_over: “selected” or “all” for aggregate computation

Returns:

Dictionary with logging results (includes experiment_name)

biolm.plugins.mlflow.protocols.log_to_mlflow(prepared_data: Dict[str, Any], experiment_name: str, mlflow_uri: str | None = None, dry_run: bool = False) → Dict[str, Any]

Log prepared data to MLflow (Stage 2: MLflow interaction only).

Args:

prepared_data: Prepared logging data from prepare_logging_data() experiment_name: MLflow experiment name mlflow_uri: Optional MLflow tracking URI dry_run: If True, don’t actually log to MLflow

Returns:

Dictionary with run IDs and status

biolm.plugins.mlflow.protocols.prepare_artifact(artifact_spec: Dict, result: Dict) → tuple[str, Union[str, bytes]]

Prepare an artifact for logging.

Args:

artifact_spec: Artifact specification from output rule result: Result dictionary

Returns:

Tuple of (artifact_name, artifact_content)

biolm.plugins.mlflow.protocols.prepare_logging_data(results: List[Dict], outputs_config: List[Dict], protocol_metadata: Dict | None = None, aggregate_over: str = 'selected') → Dict[str, Any]

Prepare all logging data (Stage 1: No MLflow interaction).

Args:

results: List of result dictionaries outputs_config: List of output rule configurations protocol_metadata: Optional protocol metadata (name, version, inputs, etc.) aggregate_over: “selected” or “all” for aggregate computation

Returns:

Dictionary with prepared logging data

biolm.plugins.mlflow.protocols.select_results(results: List[Dict], output_rule: Dict) → List[Dict]

Select results based on an output rule.

Args:

results: List of result dictionaries output_rule: Output rule configuration

Returns:

List of selected result dictionaries

Module contents

MLflow plugin: protocol result logging and dataset storage.

exception biolm.plugins.mlflow.MLflowNotAvailableError

Bases: Exception

Raised when MLflow is not installed.