Experiment Tracking¶

Molfun integrates with popular experiment tracking platforms through a unified BaseTracker interface. Multiple trackers can be composed to log to several platforms simultaneously.

Quick Start¶

from molfun import MolfunStructureModel

model = MolfunStructureModel.from_pretrained("openfold_v2")

# Pass tracker name to fit()
model.fit(train_dataset=ds, strategy="lora", tracker="wandb", epochs=10)

# Or instantiate a tracker directly
from molfun.tracking import WandbTracker

tracker = WandbTracker(project="molfun-experiments", name="lora-r8")
model.fit(train_dataset=ds, strategy="lora", tracker=tracker, epochs=10)

# Compose multiple trackers
from molfun.tracking import CompositeTracker, WandbTracker, ConsoleTracker

tracker = CompositeTracker([
    WandbTracker(project="molfun"),
    ConsoleTracker(),
])

BaseTracker¶

BaseTracker ¶

Bases: ABC

Interface for experiment tracking.

Lifecycle

tracker.start_run("my-experiment", tags=["lora", "pairformer"]) tracker.log_config({"lr": 1e-4, "block": "pairformer"}) for epoch in range(epochs): tracker.log_metrics({"train_loss": 0.5, "val_loss": 0.3}, step=epoch) tracker.log_artifact("checkpoints/best.pt") tracker.log_text("Training complete. Best val: 0.28", tag="summary") tracker.end_run(status="completed")

start_run `abstractmethod` ¶

start_run(name: str | None = None, tags: list[str] | None = None, config: dict | None = None) -> None

Start a new tracked run/experiment.

log_metrics `abstractmethod` ¶

log_metrics(metrics: dict, step: int | None = None) -> None

Log scalar metrics (loss, accuracy, etc.).

log_config `abstractmethod` ¶

log_config(config: dict) -> None

Log hyperparameters / experiment configuration.

log_artifact `abstractmethod` ¶

log_artifact(path: str, name: str | None = None) -> None

Log a file artifact (checkpoint, plot, etc.).

log_text `abstractmethod` ¶

log_text(text: str, tag: str = 'log') -> None

Log free-form text (agent reasoning, summaries, etc.).

end_run `abstractmethod` ¶

end_run(status: str = 'completed') -> None

End the current run.

Abstract base class for all experiment trackers. Implement this to add support for a new tracking platform.

Required Methods¶

Method	Description
`start_run(**kwargs)`	Initialize a tracking run
`log_metrics(metrics, step)`	Log a dict of scalar metrics
`log_config(config)`	Log experiment configuration/hyperparameters
`log_artifact(path, name)`	Log a file artifact (model checkpoint, etc.)
`log_text(text, name)`	Log a text blob (predictions, summaries)
`end_run()`	Finalize and close the run

from molfun.tracking.base import BaseTracker

class MyTracker(BaseTracker):
    def start_run(self, **kwargs):
        ...
    def log_metrics(self, metrics: dict, step: int | None = None):
        ...
    def log_config(self, config: dict):
        ...
    def log_artifact(self, path: str, name: str | None = None):
        ...
    def log_text(self, text: str, name: str | None = None):
        ...
    def end_run(self):
        ...

WandbTracker¶

WandbTracker ¶

Bases: BaseTracker

Log experiments to Weights & Biases.

Usage::

tracker = WandbTracker(project="molfun-search", entity="my-team")
tracker.start_run("baseline-pairformer", tags=["lora"])
tracker.log_metrics({"val_loss": 0.3}, step=1)
tracker.end_run()

Weights & Biases experiment tracker.

from molfun.tracking import WandbTracker

tracker = WandbTracker(
    project="molfun-structure",
    name="lora-rank8-run1",
    entity="my-team",
    tags=["lora", "openfold"],
)

tracker.start_run()
tracker.log_config({"rank": 8, "lr": 1e-4})
tracker.log_metrics({"loss": 0.42, "rmsd": 1.5}, step=100)
tracker.log_artifact("./checkpoint.pt", name="best_model")
tracker.end_run()

Parameter	Type	Default	Description
`project`	`str`	`"molfun"`	W&B project name
`name`	`str \\| None`	`None`	Run name
`entity`	`str \\| None`	`None`	W&B team/entity
`tags`	`list[str]`	`[]`	Run tags

CometTracker¶

CometTracker ¶

Bases: BaseTracker

Log experiments to Comet ML.

Usage::

tracker = CometTracker(
    project_name="molfun-search",
    api_key="your-key",  # or set COMET_API_KEY env var
)
tracker.start_run("baseline-pairformer", tags=["lora"])
tracker.log_metrics({"val_loss": 0.3}, step=1)
tracker.end_run()

Comet ML experiment tracker.

from molfun.tracking import CometTracker

tracker = CometTracker(
    project_name="molfun-experiments",
    api_key="...",
)

Parameter	Type	Default	Description
`project_name`	`str`	required	Comet project name
`api_key`	`str \\| None`	`None`	API key (reads `COMET_API_KEY` env var if `None`)
`workspace`	`str \\| None`	`None`	Comet workspace

MLflowTracker¶

MLflowTracker ¶

Bases: BaseTracker

Log experiments to MLflow.

Usage::

tracker = MLflowTracker(
    experiment_name="molfun-search",
    tracking_uri="http://localhost:5000",  # or local ./mlruns
)
tracker.start_run("baseline-pairformer", tags=["lora"])
tracker.log_metrics({"val_loss": 0.3}, step=1)
tracker.end_run()

MLflow experiment tracker.

from molfun.tracking import MLflowTracker

tracker = MLflowTracker(
    experiment_name="molfun-structure",
    tracking_uri="http://localhost:5000",
)

Parameter	Type	Default	Description
`experiment_name`	`str`	`"molfun"`	MLflow experiment name
`tracking_uri`	`str \\| None`	`None`	MLflow tracking server URI

LangfuseTracker¶

LangfuseTracker ¶

Bases: BaseTracker

Log experiments and agent traces to Langfuse.

Works with Langfuse Cloud or self-hosted instances. API keys can be passed directly or via environment variables LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY, LANGFUSE_HOST.

Langfuse tracker for monitoring AI/ML pipelines.

from molfun.tracking import LangfuseTracker

tracker = LangfuseTracker(
    public_key="...",
    secret_key="...",
)

Parameter	Type	Default	Description
`public_key`	`str \\| None`	`None`	Langfuse public key (or `LANGFUSE_PUBLIC_KEY` env var)
`secret_key`	`str \\| None`	`None`	Langfuse secret key (or `LANGFUSE_SECRET_KEY` env var)
`host`	`str`	`"https://cloud.langfuse.com"`	Langfuse host URL

HuggingFaceTracker¶

HuggingFaceTracker ¶

Bases: BaseTracker

Track experiments by pushing to Hugging Face Hub.

Metrics and config are accumulated during the run and written to the repo on end_run() (or when log_artifact() is called).

Usage::

from molfun.tracking import HuggingFaceTracker

tracker = HuggingFaceTracker(repo_id="user/my-model")
tracker.start_run(name="lora-pairformer", config={...})
tracker.log_metrics({"val_loss": 0.28}, step=10)
tracker.log_artifact("runs/checkpoint/")
tracker.end_run()

Combine with other trackers::

from molfun.tracking import CompositeTracker, WandbTracker

tracker = CompositeTracker([
    WandbTracker(project="molfun"),
    HuggingFaceTracker(repo_id="user/my-model"),
])

init ¶

__init__(repo_id: str, token: str | None = None, private: bool = False, repo_type: str = 'model')

Parameters:

Name	Type	Description	Default
`repo_id`	`str`	HF Hub repo (e.g. "username/model-name").	required
`token`	`str \| None`	HF API token (or set HF_TOKEN env var).	`None`
`private`	`bool`	Whether the repo should be private.	`False`
`repo_type`	`str`	"model" or "dataset".	`'model'`

log_artifact ¶

log_artifact(path: str, name: str | None = None) -> None

Upload a file or directory to the HF Hub repo.

end_run ¶

end_run(status: str = 'completed') -> None

Push accumulated metadata (config, metrics) to the repo.

upload_model_card ¶

upload_model_card(card_text: str) -> None

Push a README.md (model card) to the repo.

download_repo ¶

download_repo(local_dir: str, revision: str | None = None) -> str

Download the full repo to a local directory.

HuggingFace Hub tracker that logs metrics and model cards.

from molfun.tracking import HuggingFaceTracker

tracker = HuggingFaceTracker(repo_id="myorg/molfun-model")

Parameter	Type	Default	Description
`repo_id`	`str`	required	HuggingFace Hub repository ID
`token`	`str \\| None`	`None`	HF API token

ConsoleTracker¶

ConsoleTracker ¶

Bases: BaseTracker

Logs everything to stdout.

Usage::

tracker = ConsoleTracker(prefix="[molfun]")
tracker.start_run("baseline-pairformer")
tracker.log_metrics({"train_loss": 0.5}, step=1)

Simple tracker that prints metrics to the console. Useful for debugging and local development.

from molfun.tracking import ConsoleTracker

tracker = ConsoleTracker(print_every=10)

Parameter	Type	Default	Description
`print_every`	`int`	`1`	Print metrics every N steps

CompositeTracker¶

CompositeTracker ¶

Bases: BaseTracker

Fans out every call to multiple trackers.

If one tracker fails, the others still receive the call. Errors are logged but not raised to avoid interrupting training.

Compose multiple trackers to log to several platforms simultaneously.

from molfun.tracking import CompositeTracker, WandbTracker, ConsoleTracker, MLflowTracker

tracker = CompositeTracker([
    WandbTracker(project="molfun"),
    MLflowTracker(experiment_name="molfun"),
    ConsoleTracker(),
])

# All calls are forwarded to every child tracker
tracker.start_run()
tracker.log_metrics({"loss": 0.42}, step=1)
tracker.end_run()

Parameter	Type	Description
`trackers`	`list[BaseTracker]`	List of tracker instances to compose

Experiment Tracking¶

Quick Start¶

BaseTracker¶

BaseTracker ¶

start_run abstractmethod ¶

log_metrics abstractmethod ¶

log_config abstractmethod ¶

log_artifact abstractmethod ¶

log_text abstractmethod ¶

end_run abstractmethod ¶

Required Methods¶

WandbTracker¶

WandbTracker ¶

CometTracker¶

CometTracker ¶

MLflowTracker¶

MLflowTracker ¶

LangfuseTracker¶

LangfuseTracker ¶

HuggingFaceTracker¶

HuggingFaceTracker ¶

__init__ ¶

log_artifact ¶

end_run ¶

upload_model_card ¶

download_repo ¶

ConsoleTracker¶

ConsoleTracker ¶

CompositeTracker¶

CompositeTracker ¶

start_run `abstractmethod` ¶

log_metrics `abstractmethod` ¶

log_config `abstractmethod` ¶

log_artifact `abstractmethod` ¶

log_text `abstractmethod` ¶

end_run `abstractmethod` ¶

init ¶