Skip to content

Weights & Biases

framework3.utils.wandb

WandbAgent

A class to create and run Weights & Biases agents for sweeps.

This class provides a callable interface to create and run Weights & Biases agents, which are used to execute sweep runs with specified configurations.

Key Features
  • Create and run Weights & Biases agents for sweeps
  • Execute custom functions with sweep configurations
  • Automatically handle initialization and teardown of Weights & Biases runs
Usage 
from framework3.utils.wandb import WandbAgent

def custom_function(config):
    # Your sweep run logic here
    result = ...
    return result

agent = WandbAgent()
agent("your_sweep_id", "your_project_name", custom_function)

Methods:

Name Description
__call__

str, project: str, function: Callable) -> None: Create and run a Weights & Biases agent for a specified sweep.
Source code in framework3/utils/wandb.py 
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
class WandbAgent:
    """
    A class to create and run Weights & Biases agents for sweeps.

    This class provides a callable interface to create and run Weights & Biases agents,
    which are used to execute sweep runs with specified configurations.

    Key Features:
        - Create and run Weights & Biases agents for sweeps
        - Execute custom functions with sweep configurations
        - Automatically handle initialization and teardown of Weights & Biases runs

    Usage:
        ```python
        from framework3.utils.wandb import WandbAgent

        def custom_function(config):
            # Your sweep run logic here
            result = ...
            return result

        agent = WandbAgent()
        agent("your_sweep_id", "your_project_name", custom_function)
        ```

    Methods:
        __call__(sweep_id: str, project: str, function: Callable) -> None:
            Create and run a Weights & Biases agent for a specified sweep.
    """

    @staticmethod
    def __call__(sweep_id: str, project: str, function: Callable) -> None:
        """
        Create and run a Weights & Biases agent for a specified sweep.

        This method initializes a Weights & Biases agent, executes the provided function
        with the sweep configuration, and handles the teardown of the Weights & Biases run.

        Args:
            sweep_id (str): The ID of the sweep to run.
            project (str): The name of the Weights & Biases project.
            function (Callable): A function that takes a configuration dictionary and returns a result to be logged.

        Returns:
            (None)
        """
        wandb.agent(  # type: ignore
            sweep_id,
            function=lambda: {
                wandb.init(reinit="finish_previous"),  # type: ignore
                wandb.log(function(dict(wandb.config))),  # type: ignore
            },
            project=project,
        )  # type: ignore
        wandb.teardown()  # type: ignore
__call__(sweep_id, project, function) staticmethod

Create and run a Weights & Biases agent for a specified sweep.

This method initializes a Weights & Biases agent, executes the provided function with the sweep configuration, and handles the teardown of the Weights & Biases run.

Parameters:

Name Type Description Default
sweep_id str

The ID of the sweep to run.

 required
project str

The name of the Weights & Biases project.

 required
function Callable

A function that takes a configuration dictionary and returns a result to be logged.

 required

Returns:

Type Description
None

(None)
Source code in framework3/utils/wandb.py 
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
@staticmethod
def __call__(sweep_id: str, project: str, function: Callable) -> None:
    """
    Create and run a Weights & Biases agent for a specified sweep.

    This method initializes a Weights & Biases agent, executes the provided function
    with the sweep configuration, and handles the teardown of the Weights & Biases run.

    Args:
        sweep_id (str): The ID of the sweep to run.
        project (str): The name of the Weights & Biases project.
        function (Callable): A function that takes a configuration dictionary and returns a result to be logged.

    Returns:
        (None)
    """
    wandb.agent(  # type: ignore
        sweep_id,
        function=lambda: {
            wandb.init(reinit="finish_previous"),  # type: ignore
            wandb.log(function(dict(wandb.config))),  # type: ignore
        },
        project=project,
    )  # type: ignore
    wandb.teardown()  # type: ignore

WandbRunLogger

Source code in framework3/utils/wandb.py 
269
class WandbRunLogger: ...

WandbSweepManager

A manager class for creating and handling Weights & Biases sweeps.

This class provides methods to generate sweep configurations, create sweeps, and manage sweep runs for hyperparameter optimization using Weights & Biases.

Key Features
  • Generate sweep configurations from BaseFilter pipelines
  • Create sweeps with customizable parameters and metrics
  • Retrieve sweep information and best configurations
  • Restart sweeps by deleting specific runs
Usage 
from framework3.utils.wandb import WandbSweepManager
from framework3.base import BaseFilter, BaseMetric, XYData

# Create a WandbSweepManager instance
sweep_manager = WandbSweepManager()

# Create a sweep
pipeline = YourPipeline()
scorer = YourScorer()
x_data = XYData(...)
y_data = XYData(...)
sweep_id = sweep_manager.create_sweep(pipeline, "your_project", scorer, x_data, y_data)

# Get the best configuration from a sweep
best_config = sweep_manager.get_best_config("your_project", sweep_id, order="ascending")

# Restart a sweep
sweep = sweep_manager.get_sweep("your_project", sweep_id)
sweep_manager.restart_sweep(sweep, states=["failed", "crashed"])

Methods:

Name Description
get_grid

Dict[str, Any], config: Dict[str, Any]) -> None: Recursively extract grid search parameters from a pipeline configuration.
generate_config_for_pipeline

BaseFilter) -> Dict[str, Any]: Generate a Weights & Biases sweep configuration from a BaseFilter pipeline.
create_sweep

BaseFilter, project_name: str, scorer: BaseMetric, x: XYData, y: XYData | None = None) -> str: Create a new sweep in Weights & Biases.
get_sweep

str, sweep_id: str) -> Any: Retrieve a sweep object from Weights & Biases.
get_best_config

str, sweep_id: str, order: str) -> Dict[str, Any]: Get the best configuration from a completed sweep.
restart_sweep

Any, states: List[str] | Literal["all"] = "all") -> None: Restart a sweep by deleting runs with specified states.
init

str, name: str, reinit: bool = True) -> Any: Initialize a new Weights & Biases run.
Source code in framework3/utils/wandb.py 
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
class WandbSweepManager:
    """
    A manager class for creating and handling Weights & Biases sweeps.

    This class provides methods to generate sweep configurations, create sweeps,
    and manage sweep runs for hyperparameter optimization using Weights & Biases.

    Key Features:
        - Generate sweep configurations from BaseFilter pipelines
        - Create sweeps with customizable parameters and metrics
        - Retrieve sweep information and best configurations
        - Restart sweeps by deleting specific runs

    Usage:
        ```python
        from framework3.utils.wandb import WandbSweepManager
        from framework3.base import BaseFilter, BaseMetric, XYData

        # Create a WandbSweepManager instance
        sweep_manager = WandbSweepManager()

        # Create a sweep
        pipeline = YourPipeline()
        scorer = YourScorer()
        x_data = XYData(...)
        y_data = XYData(...)
        sweep_id = sweep_manager.create_sweep(pipeline, "your_project", scorer, x_data, y_data)

        # Get the best configuration from a sweep
        best_config = sweep_manager.get_best_config("your_project", sweep_id, order="ascending")

        # Restart a sweep
        sweep = sweep_manager.get_sweep("your_project", sweep_id)
        sweep_manager.restart_sweep(sweep, states=["failed", "crashed"])
        ```

    Methods:
        get_grid(aux: Dict[str, Any], config: Dict[str, Any]) -> None:
            Recursively extract grid search parameters from a pipeline configuration.
        generate_config_for_pipeline(pipeline: BaseFilter) -> Dict[str, Any]:
            Generate a Weights & Biases sweep configuration from a BaseFilter pipeline.
        create_sweep(pipeline: BaseFilter, project_name: str, scorer: BaseMetric, x: XYData, y: XYData | None = None) -> str:
            Create a new sweep in Weights & Biases.
        get_sweep(project_name: str, sweep_id: str) -> Any:
            Retrieve a sweep object from Weights & Biases.
        get_best_config(project_name: str, sweep_id: str, order: str) -> Dict[str, Any]:
            Get the best configuration from a completed sweep.
        restart_sweep(sweep: Any, states: List[str] | Literal["all"] = "all") -> None:
            Restart a sweep by deleting runs with specified states.
        init(group: str, name: str, reinit: bool = True) -> Any:
            Initialize a new Weights & Biases run.
    """

    @staticmethod
    def get_grid(aux: Dict[str, Any], config: Dict[str, Any]):
        """
        Recursively extract grid search parameters from a pipeline configuration.

        Args:
            aux (Dict[str, Any]): The input configuration dictionary.
            config (Dict[str, Any]): The output configuration dictionary to be updated.
        """
        match aux["params"]:
            case {"filters": filters, **r}:
                for filter_config in filters:
                    WandbSweepManager.get_grid(filter_config, config)
            case {"pipeline": pipeline, **r}:  # noqa: F841
                WandbSweepManager.get_grid(pipeline, config)
            case p_params:
                if "_grid" in aux:
                    f_config = {}
                    for param, value in aux["_grid"].items():
                        print(f"categorical param: {param}: {value}")
                        p_params.update({param: value})
                        match type(value):
                            case list():
                                f_config[param] = {"values": value}
                            case dict():
                                f_config[param] = value
                            case _:
                                f_config[param] = {"values": value}
                    if len(f_config) > 0:
                        config["parameters"]["filters"]["parameters"][
                            str(aux["clazz"])
                        ] = {"parameters": f_config}

    @staticmethod
    def generate_config_for_pipeline(pipeline: BaseFilter) -> Dict[str, Any]:
        """
        Generate a Weights & Biases sweep configuration from a BaseFilter pipeline.

        Args:
            pipeline (BaseFilter): The pipeline to generate the configuration for.

        Returns:
            Dict[str, Any]: A Weights & Biases sweep configuration.
        """
        sweep_config: Dict[str, Dict[str, Dict[str, Any]]] = {
            "parameters": {"filters": {"parameters": {}}, "pipeline": {"value": {}}}
        }

        dumped_pipeline = pipeline.item_dump(include=["_grid"])

        WandbSweepManager.get_grid(dumped_pipeline, sweep_config)

        sweep_config["parameters"]["pipeline"]["value"] = dumped_pipeline

        return sweep_config

    def create_sweep(
        self,
        pipeline: BaseFilter,
        project_name: str,
        scorer: BaseMetric,
        x: XYData,
        y: XYData | None = None,
    ) -> str:
        """
        Create a new sweep in Weights & Biases.

        Args:
            pipeline (BaseFilter): The pipeline to be optimized.
            project_name (str): The name of the Weights & Biases project.
            scorer (BaseMetric): The metric used to evaluate the pipeline.
            x (XYData): The input data.
            y (XYData | None): The target data (optional).

        Returns:
            str: The ID of the created sweep.
        """
        sweep_config = WandbSweepManager.generate_config_for_pipeline(pipeline)
        sweep_config["method"] = "grid"
        sweep_config["parameters"]["x_dataset"] = {"value": x._hash}
        sweep_config["parameters"]["y_dataset"] = (
            {"value": y._hash} if y is not None else {"value": "None"}
        )
        sweep_config["metric"] = {
            "name": scorer.__class__.__name__,
            "goal": "maximize" if scorer.higher_better else "minimize",
        }
        print("______________________SWEE CONFIG_____________________")
        print(sweep_config)
        print("_____________________________________________________")
        return wandb.sweep(sweep_config, project=project_name)  # type: ignore

    def get_sweep(self, project_name, sweep_id) -> Any:
        """
        Retrieve a sweep object from Weights & Biases.

        Args:
            project_name (str): The name of the Weights & Biases project.
            sweep_id (str): The ID of the sweep to retrieve.

        Returns:
            (Any): The sweep object.
        """
        sweep = wandb.Api().sweep(f"citius-irlab/{project_name}/sweeps/{sweep_id}")  # type: ignore
        return sweep

    def get_best_config(self, project_name, sweep_id, order) -> Dict[str, Any]:
        """
        Get the best configuration from a completed sweep.

        Args:
            project_name (str): The name of the Weights & Biases project.
            sweep_id (str): The ID of the sweep.
            order (str): The order to use when selecting the best run ("ascending" or "descending").

        Returns:
            (Dict[str, Any]): The configuration of the best run.
        """
        sweep = self.get_sweep(project_name, sweep_id)
        winner_run = sweep.best_run(order=order)
        return dict(winner_run.config)

    def restart_sweep(self, sweep, states: List[str] | Literal["all"] = "all"):
        """
        Restart a sweep by deleting runs with specified states.

        Args:
            sweep (Any): The sweep object to restart.
            states (List[str] | Literal["all"]): The states of runs to delete, or "all" to delete all runs.
        """
        # Eliminar todas las ejecuciones fallidas
        for run in sweep.runs:
            if run.state in states or states == "all":
                run.delete()
                print("Deleting run:", run.id)

    def init(self, group: str, name: str, reinit=True) -> Any:
        """
        Initialize a new Weights & Biases run.

        Args:
            group (str): The group name for the run.
            name (str): The name of the run.
            reinit (bool): Whether to reinitialize if a run is already in progress.

        Returns:
            (Any): The initialized run object.
        """
        run = wandb.init(group=group, name=name, reinit=reinit)  # type: ignore
        return run
create_sweep(pipeline, project_name, scorer, x, y=None)

Create a new sweep in Weights & Biases.

Parameters:

Name Type Description Default
pipeline BaseFilter

The pipeline to be optimized.

 required
project_name str

The name of the Weights & Biases project.

 required
scorer BaseMetric

The metric used to evaluate the pipeline.

 required
x XYData

The input data.

 required
y XYData | None

The target data (optional).

 None

Returns:

Name Type Description
str str

The ID of the created sweep.
Source code in framework3/utils/wandb.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def create_sweep(
    self,
    pipeline: BaseFilter,
    project_name: str,
    scorer: BaseMetric,
    x: XYData,
    y: XYData | None = None,
) -> str:
    """
    Create a new sweep in Weights & Biases.

    Args:
        pipeline (BaseFilter): The pipeline to be optimized.
        project_name (str): The name of the Weights & Biases project.
        scorer (BaseMetric): The metric used to evaluate the pipeline.
        x (XYData): The input data.
        y (XYData | None): The target data (optional).

    Returns:
        str: The ID of the created sweep.
    """
    sweep_config = WandbSweepManager.generate_config_for_pipeline(pipeline)
    sweep_config["method"] = "grid"
    sweep_config["parameters"]["x_dataset"] = {"value": x._hash}
    sweep_config["parameters"]["y_dataset"] = (
        {"value": y._hash} if y is not None else {"value": "None"}
    )
    sweep_config["metric"] = {
        "name": scorer.__class__.__name__,
        "goal": "maximize" if scorer.higher_better else "minimize",
    }
    print("______________________SWEE CONFIG_____________________")
    print(sweep_config)
    print("_____________________________________________________")
    return wandb.sweep(sweep_config, project=project_name)  # type: ignore
generate_config_for_pipeline(pipeline) staticmethod

Generate a Weights & Biases sweep configuration from a BaseFilter pipeline.

Parameters:

Name Type Description Default
pipeline BaseFilter

The pipeline to generate the configuration for.

 required

Returns:

Type Description
Dict[str, Any]

Dict[str, Any]: A Weights & Biases sweep configuration.
Source code in framework3/utils/wandb.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
@staticmethod
def generate_config_for_pipeline(pipeline: BaseFilter) -> Dict[str, Any]:
    """
    Generate a Weights & Biases sweep configuration from a BaseFilter pipeline.

    Args:
        pipeline (BaseFilter): The pipeline to generate the configuration for.

    Returns:
        Dict[str, Any]: A Weights & Biases sweep configuration.
    """
    sweep_config: Dict[str, Dict[str, Dict[str, Any]]] = {
        "parameters": {"filters": {"parameters": {}}, "pipeline": {"value": {}}}
    }

    dumped_pipeline = pipeline.item_dump(include=["_grid"])

    WandbSweepManager.get_grid(dumped_pipeline, sweep_config)

    sweep_config["parameters"]["pipeline"]["value"] = dumped_pipeline

    return sweep_config
get_best_config(project_name, sweep_id, order)

Get the best configuration from a completed sweep.

Parameters:

Name Type Description Default
project_name str

The name of the Weights & Biases project.

 required
sweep_id str

The ID of the sweep.

 required
order str

The order to use when selecting the best run ("ascending" or "descending").

 required

Returns:

Type Description
Dict[str, Any]

The configuration of the best run.
Source code in framework3/utils/wandb.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
def get_best_config(self, project_name, sweep_id, order) -> Dict[str, Any]:
    """
    Get the best configuration from a completed sweep.

    Args:
        project_name (str): The name of the Weights & Biases project.
        sweep_id (str): The ID of the sweep.
        order (str): The order to use when selecting the best run ("ascending" or "descending").

    Returns:
        (Dict[str, Any]): The configuration of the best run.
    """
    sweep = self.get_sweep(project_name, sweep_id)
    winner_run = sweep.best_run(order=order)
    return dict(winner_run.config)
get_grid(aux, config) staticmethod

Recursively extract grid search parameters from a pipeline configuration.

Parameters:

Name Type Description Default
aux Dict[str, Any]

The input configuration dictionary.

 required
config Dict[str, Any]

The output configuration dictionary to be updated.

 required
Source code in framework3/utils/wandb.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
@staticmethod
def get_grid(aux: Dict[str, Any], config: Dict[str, Any]):
    """
    Recursively extract grid search parameters from a pipeline configuration.

    Args:
        aux (Dict[str, Any]): The input configuration dictionary.
        config (Dict[str, Any]): The output configuration dictionary to be updated.
    """
    match aux["params"]:
        case {"filters": filters, **r}:
            for filter_config in filters:
                WandbSweepManager.get_grid(filter_config, config)
        case {"pipeline": pipeline, **r}:  # noqa: F841
            WandbSweepManager.get_grid(pipeline, config)
        case p_params:
            if "_grid" in aux:
                f_config = {}
                for param, value in aux["_grid"].items():
                    print(f"categorical param: {param}: {value}")
                    p_params.update({param: value})
                    match type(value):
                        case list():
                            f_config[param] = {"values": value}
                        case dict():
                            f_config[param] = value
                        case _:
                            f_config[param] = {"values": value}
                if len(f_config) > 0:
                    config["parameters"]["filters"]["parameters"][
                        str(aux["clazz"])
                    ] = {"parameters": f_config}
get_sweep(project_name, sweep_id)

Retrieve a sweep object from Weights & Biases.

Parameters:

Name Type Description Default
project_name str

The name of the Weights & Biases project.

 required
sweep_id str

The ID of the sweep to retrieve.

 required

Returns:

Type Description
Any

The sweep object.
Source code in framework3/utils/wandb.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
def get_sweep(self, project_name, sweep_id) -> Any:
    """
    Retrieve a sweep object from Weights & Biases.

    Args:
        project_name (str): The name of the Weights & Biases project.
        sweep_id (str): The ID of the sweep to retrieve.

    Returns:
        (Any): The sweep object.
    """
    sweep = wandb.Api().sweep(f"citius-irlab/{project_name}/sweeps/{sweep_id}")  # type: ignore
    return sweep
init(group, name, reinit=True)

Initialize a new Weights & Biases run.

Parameters:

Name Type Description Default
group str

The group name for the run.

 required
name str

The name of the run.

 required
reinit bool

Whether to reinitialize if a run is already in progress.

 True

Returns:

Type Description
Any

The initialized run object.
Source code in framework3/utils/wandb.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
def init(self, group: str, name: str, reinit=True) -> Any:
    """
    Initialize a new Weights & Biases run.

    Args:
        group (str): The group name for the run.
        name (str): The name of the run.
        reinit (bool): Whether to reinitialize if a run is already in progress.

    Returns:
        (Any): The initialized run object.
    """
    run = wandb.init(group=group, name=name, reinit=reinit)  # type: ignore
    return run
restart_sweep(sweep, states='all')

Restart a sweep by deleting runs with specified states.

Parameters:

Name Type Description Default
sweep Any

The sweep object to restart.

 required
states List[str] | Literal['all']

The states of runs to delete, or "all" to delete all runs.

 'all'
Source code in framework3/utils/wandb.py 
182
183
184
185
186
187
188
189
190
191
192
193
194
def restart_sweep(self, sweep, states: List[str] | Literal["all"] = "all"):
    """
    Restart a sweep by deleting runs with specified states.

    Args:
        sweep (Any): The sweep object to restart.
        states (List[str] | Literal["all"]): The states of runs to delete, or "all" to delete all runs.
    """
    # Eliminar todas las ejecuciones fallidas
    for run in sweep.runs:
        if run.state in states or states == "all":
            run.delete()
            print("Deleting run:", run.id)