API: Experiment

Overview

An “Experiment” is the central concept in RapidFire AI to organize and streamline your multi-config AI experimentation.

Every experiment must have a unique user-given name that is used to collate plots on the metrics dashboard and for saving its artifacts. If you (mistakenly) reuse a previous experiment name, RapidFire AI will append a suffix akin to what filesystems do.

The Experiment class has the functions and semantics detailed below. As of this writing, an experiment can operate in either "evals" mode or "fit" mode but not both. We plan to allow both modes in the same experiment very soon.

Note

Only one function in this class can be run at a time. If you interrupt a long-running function, say, run_evals() or run_fit(), wait for up to 2min for the Python threads to be cleaned up before running another function.

Experiment Constructor

Constructor to instantiate a new experiment.

__init__(self, experiment_name: str, mode: str = 'evals', experiment_path: str = '$RF_HOME/rapidfire_experiments', num_cpus: int = None, num_gpus: int = None) None
Parameters:
  • experiment_name (str) – Unique name for this experiment

  • mode (str) – Mode of this experiment, either "fit" or "evals"; default is "evals". The default matches the rapidfireai init install default; "fit" is the --train opt-in.

  • experiment_path (str, optional) – Path to a folder to store this experiment’s artifacts. Default is "$RF_HOME/rapidfire_experiments"

  • num_cpus (int, optional) – Cap on the number of logical CPUs handed to the execution backend (Ray). Must be greater than 0 and no greater than the number available. If unset, defaults to 90% of available logical CPUs. Used in "evals" mode only; validated but otherwise ignored in "fit" mode.

  • num_gpus (int, optional) – Cap on the number of GPUs handed to the execution backend (Ray). Must be non-negative and no greater than the number available. If unset, defaults to all available GPUs. Used in "evals" mode only; validated but otherwise ignored in "fit" mode.

Returns:

None

Return type:

None

Example:

# Based on FiQA RAG chatbot tutorial notebook (evals is the default mode)
>>> experiment = Experiment(experiment_name="exp1-fiqa")
Allocating 8 CPUs and 1 GPUs to the experiment
Using 1 actors, 1.0 GPUs per actor, 8.0 CPUs per actor
Experiment exp1-fiqa created ...

# Based on SFT chatbot tutorial notebook
>>> experiment = Experiment(experiment_name="exp1-chatqa", mode="fit")
Experiment exp1-chatqa created ...

Notes:

You can instantiate as many experiment objects as you want. We recommend explicitly ending a previous experiment (see end() below) before starting a new one so that you are cognizant of your code and/or config changes across them.

If you are using Jupyter and if its kernel restarts or gets interrupted for whatever reason, you can just reconnect the kernel and pick up that experiment from where you left off by just rerunning its constructor cell as is (unless you ended the experiment explicitly). Using that object you can continue that experiment as before.

Resource Allocation Defaults and Override Semantics

In "evals" mode, the constructor auto-detects the available hardware at creation time and computes a default actor allocation. It then prints the allocation (CPUs/GPUs handed to the backend, number of actors, and per-actor budgets). These computed defaults are then used by run_evals(), and you can override them per call.

The defaults are derived as follows:

  • GPU machines: One actor per GPU, with 1 GPU per actor.

  • CPU-only machines: By default 90% of logical CPUs are handed to the backend and divided evenly across actors. The actor count is tiered by the number of CPUs allocated to the backend: 2 for up to 2 CPUs, 4 for up to 4, 8 for up to 32, 16 for up to 128, and 32 beyond that.

The num_cpus and num_gpus constructor arguments cap what is handed to the backend; num_cpus must be greater than 0; num_gpus must be non-negative; neither may exceed what is available.

run_evals() exposes num_actors, gpus_per_actor, and cpus_per_actor; any value you pass overrides the auto-detected default for that call, while any left as None falls back to the constructor’s computed value. The override is validated: num_actors * gpus_per_actor may not exceed the available GPUs and num_actors * cpus_per_actor may not exceed the available CPUs, otherwise the call errors out. If gpus_per_actor resolves to 0 (a CPU-only allocation), a warning reminds you to use external model APIs for inference.

In "fit" mode, this auto-allocation does not apply; control GPU usage per config via the num_gpus argument of run_fit() (and RFModelConfig) instead.

Run Fit

The main function to launch training (including LLM fine-tuning and post-training) and evaluation for a given config group in one go. See the Multi-Config Specification page for more details on how to construct a config group.

run_fit(self, param_config: Any, create_model_fn: Callable, train_dataset: Dataset, eval_dataset: Dataset, num_chunks: int, seed: int = 42, num_gpus: int = 1) None:
Parameters:
  • param_config (Train config-group or list as described in the Multi-Config Specification page) – A train config knob dictionary, a generated config group, or a list of configs or config groups

  • create_model_fn (Callable) – User-given function to create a model instance; a single cfg is passed as input by the system

  • train_dataset (Dataset) – Training dataset

  • eval_dataset (Dataset) – Evaluation dataset to measure eval metrics

  • num_chunks (int) – Number of logical splits of data to control degree of concurrency for multi-config execution (recommended: at least 4)

  • seed (int, optional) – Seed for any randomness used in your code (default: 42)

  • num_gpus (int, optional) – Number of GPUs to use per run/config for each config represented in param_config (default: 1); overriden by any num_gpus given in RFModelConfig for those associated configs.

Returns:

None

Return type:

None

Example:

# Based on SFT chatbot tutorial notebook
>>> experiment.run_fit(config_group, sample_create_model, train_dataset, eval_dataset, num_chunks=4, seed=42)
Started 4 worker processes successfully ...

Notes:

This method auto-generates the metrics files as per user specification and auto-plots them on the dashboard. Within an experiment, you can rerun run_fit() as many times as you want. All of them will be overlaid on the same plots on the metrics dashboard. Note that run_fit() must be actively running for you to be able to use Interactive Control (IC) ops on the dashboard.

The param_config argument is very versatile in allowing you to construct various knob combinations and launch them in one go. It can be a single config dictionary, a list of config dictionaries, a config group generator output (RFGridSearch() or RFRandomSearch() for now), or even a list with mix of configs or config group generator outputs as its elements. Please see the the Multi-Config Specification page for more details.

Each individual config is passed as input to your create_model_fn(). Inside it you can use whatever knob you set in the config group, e.g., model type or name to instantiate a model accordingly. You can import models from libraries such as HuggingFace transformers or load your own PyTorch checkpoints.

The num_chunks argument is a critical one that enables you to balance a higher degree of concurrency you desire for cross-config comparisons against the (relatively minor) extra swapping overhead incurred. We recommend at least 4, which means you will see results for all runs on 1/4th of the data at a time.

Run Evals

The main function to launch LLM evaluation (evals), including with optional RAG, for a given config group in one go. See the Multi-Config Specification page for more details on how to construct a config group.

run_evals(self, config_group: Any, dataset: Dataset, num_shards: int = 4, seed: int = 42, num_actors: int = None, gpus_per_actor: float = None, cpus_per_actor: float = None) dict[int, tuple[dict, dict]]:
Parameters:
  • config_group (Evals config-group or list as described in the Multi-Config Specification page) – Single evals config knob dictionary, a generated config group, or a list of configs or config groups

  • dataset (Dataset) – Evaluation dataset to measure eval metrics

  • num_shards (int) – Number of logical splits of data to control degree of concurrency for multi-config execution (recommended: at least 4)

  • seed (int, optional) – Seed to control randomness for online aggregation (default: 42)

  • num_actors (int, optional) – Number of parallel worker actors to control degree of concurrency. Defaults to None, in which case the value auto-detected by the constructor is used (one actor per GPU on GPU machines; a CPU-count-based tier otherwise). See the Resource Allocation section above.

  • gpus_per_actor (float, optional) – Number of GPUs assigned to each actor. Defaults to None (auto-detected: 1.0 on GPU machines, 0.0 on CPU-only machines).

  • cpus_per_actor (float, optional) – Number of CPUs assigned to each actor. Defaults to None (auto-detected from the constructor’s resource allocation).

Returns:

Dictionary keyed by run/config ID; see the description of the returned structure in the Notes below.

Return type:

dict[int, tuple[dict, dict]]

Example:

# Based on FiQA RAG chatbot tutorial notebook
>>> experiment.run_evals(config_group=config_group, dataset=fiqa_dataset, num_shards=4, num_actors=8, seed=42)
Started 8 actor processes ...

Notes:

This method auto-generates the metrics as per user specification and lists them in an auto-updated table shown on the notebook itself (and soon, on the metrics dashboard also). Alongside the metrics table, the Interactive Control (IC) Ops panel will also appear on the notebook itself. Note that run_evals() must be actively running for you to be able to use IC Ops.

Within an experiment, you can rerun run_evals() as many times as you want. All of them will be overlaid on the same plots on the metrics dashboard.

The config_group argument allows you to construct various knob combinations for inference pipelines and launch them in one go. These pipelines can involve LLMs running on your GPUs, or OpenAI API calls, or both.

Just like with run_fit() above, you can provide a single config dictionary, a list of config dictionaries, a config group generator output (RFGridSearch() or RFRandomSearch() for now), or even a list with mix of configs or config group generator outputs as its elements. Please see the the Multi-Config Specification page for more details.

The num_shards argument is identical to the num_chunks argument of run_fit() above. That is, it let you balance the degree of concurrency for cross-config comparisons against the (minor) extra swapping overhead incurred. Again, we recommend at least 4, which means you will see results being updated for all runs on 1/4th of the data at a time.

For resource control via num_actors, gpus_per_actor, and cpus_per_actor, see the Resource Allocation Defaults and Override Semantics section above.

Return value:

Unlike run_fit(), this function returns a value: a dictionary keyed by run/config ID. Each value is a 2-tuple (results, metrics):

  • results — A dictionary of the raw per-query outputs accumulated across all processed shards (e.g., predictions, references, and any keys your postprocess/compute functions emit).

  • metrics — An ordered dictionary that leads with the run’s identity and configuration and then lists the eval metrics, in this order:

    • run_id

    • model_name — the model identifier

    • the configuration knobs set on the config (listed below)

    • Samples Processed

    • each of your eval metrics, carrying confidence-interval fields (value, lower_bound, upper_bound, margin_of_error) per Online Aggregation for Evals

The configuration knobs surfaced in the metrics dictionary — each present only if set on the config, and each wrapped as {"value": ...} — are:

  • text_splitter_cfg

  • embedding_cfg

  • vector_store_cfg

  • search_cfg

  • reranker_cfg

  • sampling_params

  • prompt_manager_k

  • model_config

The model identifier appears under model_name (the model key is removed from model_config to avoid redundancy). Because every entry is wrapped, read a knob as, e.g., metrics["search_cfg"]["value"].

This gives you a config-to-metrics mapping straight off the return value (join the results / metrics dictionaries by run/config ID). Note that this surfaced set is a curated subset of knobs, not the full config dictionary.

The returned dictionary includes runs that COMPLETED, were STOPPED (via IC Ops or Optuna pruning), or are ONGOING with partial results. It excludes runs that were DELETED or FAILED, as well as any run that produced no results (for example, a clone that was created but never processed a shard).

End

End the current experiment to clear out relevant system state and allow you to move on to a new experiment. Please do not run this when a run_fit() is still running.

end(self) None
Returns:

None

Return type:

None

Cancel Current

Cancel the currently running task. Works in both "evals" and "fit" modes.

cancel_current(self) None
Returns:

None

Return type:

None

Get Runs Information

Returns metadata about all the runs from across all run_fit() invocations in the current experiment.

Note

This function is available in "fit" mode only; calling it in "evals" mode raises a ValueError. In "evals" mode, per-run configuration and metrics are returned directly by run_evals() (see its Return value above).

get_runs_info(self) pd.DataFrame:
Returns:

A DataFrame with one row per run and the following columns:

  • run_id

  • status

  • mlflow_run_id

  • completed_steps

  • total_steps

  • start_chunk_id

  • num_chunks_visited_curr_epoch

  • num_epochs_completed

  • error

  • source

  • ended_by

  • warm_started_from

  • config — the full config dictionary

Return type:

pandas.DataFrame

Examples:

# Get metadata of all runs from this experiments so far; based on SFT notebook
all_runs_info = experiment.get_runs_info()
all_runs_info # Screenshot of output below
Outputs of get runs info

Notes:

This function is useful for programmatic post-processing and/or pre-processing of runs and their config knobs. For instance, you can use it as part of new custom AutoML procedure to launch a new run_fit() with new config knob values based on get_results() from past run_fit() invocations.

We plan to expand this API in the future to return other details about runs such as total runtime, GPU utilization, etc. based on feedback.

Get Results

Returns all metrics (including loss, eval loss, and any eval metrics defined) for all steps for all runs from across all run_fit() in the current experiment.

Note

This function is available in "fit" mode only; calling it in "evals" mode raises a ValueError. In "evals" mode, results are returned directly by run_evals() (see its Return value above).

get_results(self) pd.DataFrame
Returns:

A DataFrame with the following columns: run ID, step number, loss, and one column per metric plot displayed on the dashboard

Return type:

pandas.DataFrame

Examples:

# Get results of all runs from this experiments so far; based on SFT notebook
all_results = experiment.get_results()
print(all_results.columns) # Screenshot of output below
all_results # Screenshot of output below
Columns in DataFrame returned by get resultsOutputs of get results

Notes:

This function can be useful for programmatic post-processing of the results of your experiments. For instance, you can use it as part of new custom AutoML procedure if you’d like to adjust your config for a new run_fit() based on the results of your last run_fit().

Get Log File Path

Returns the path to a log file for the current experiment.

get_log_file_path(self, log_type: str = None) Path
Parameters:

log_type (str, optional) – Which log to return. None, "main", or "experiment" returns the main experiment log; "training" returns the training log. Any other value raises a ValueError.

Returns:

Path to the requested log file

Return type:

pathlib.Path