aidevelopementtoolkit.logging_utils.mlflow_utils

start_mlflow_run(experiment_name: str, mlflow_kwargs: Dict[str, Any]) -> None

This function starts an MLflow run with the specified experiment name and additional keyword arguments.

Notes

The function checks for the presence of the MLFLOW_ENDPOINT_URL environment variable. If the provided MLflow tracking URI is a remote server (i.e., starts with "http"), it also checks for the presence of the following environment variables: - REQUESTS_CA_BUNDLE - MLFLOW_TRACKING_USERNAME - MLFLOW_TRACKING_PASSWORD - MLFLOW_TRACKING_SERVER_CERT_PATH - MLFLOW_ENDPOINT_URL - MLFLOW_S3_ENDPOINT_URL

If you are using a local MLflow server (i.e., the tracking URI starts with "sqlite:///"), the function will create an artifacts directory in the same location as the SQLite database.

Parameters:
  • experiment_name (str) –

    Name of the MLflow experiment. If the experiment does not exist, it will be created.

  • mlflow_kwargs (Dict[str, Any]) –

    Additional keyword arguments to pass to mlflow.start_run().

Examples:

>>> os.environ["MLFLOW_ENDPOINT_URL"] = https://myremote.com
>>> start_mlflow_run(
...         experiment_name="MNIST",
...         mlflow_kwargs={
...             "run_name": "ResNet50",
...             "tags": {
...                 "model": "resnet50",
...                 "pre-trained": "false",
...             },
...             "log_system_metrics": True,
...             "description": "Simple showcase."
...         }
...     )
See Also

mlflow.start_run : https://mlflow.org/docs/latest/python_api/mlflow.html#mlflow.start_run

Source code in aidevelopementtoolkit/logging_utils/mlflow_utils.py
 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
def start_mlflow_run(experiment_name: str, mlflow_kwargs: Dict[str, Any]) -> None:
    """This function starts an MLflow run with the specified 
    experiment name and additional keyword arguments.

    Notes
    -----
    The function checks for the presence of the `MLFLOW_ENDPOINT_URL` environment variable.
    If the provided MLflow tracking URI is a remote server (i.e., starts with "http"), it also checks 
    for the presence of the following environment variables:
    - `REQUESTS_CA_BUNDLE`
    - `MLFLOW_TRACKING_USERNAME`
    - `MLFLOW_TRACKING_PASSWORD`
    - `MLFLOW_TRACKING_SERVER_CERT_PATH`
    - `MLFLOW_ENDPOINT_URL`
    - `MLFLOW_S3_ENDPOINT_URL`

    If you are using a local MLflow server (i.e., the tracking URI starts with "sqlite:///"), the function
      will create an artifacts directory in the same location as the SQLite database.

    Parameters
    ----------
    experiment_name : str
        Name of the MLflow experiment. If the experiment does not exist, it will be created.

    mlflow_kwargs : Dict[str, Any]
        Additional keyword arguments to pass to `mlflow.start_run()`.

    Examples
    --------
    >>> os.environ["MLFLOW_ENDPOINT_URL"] = https://myremote.com
    >>> start_mlflow_run(
    ...         experiment_name="MNIST",
    ...         mlflow_kwargs={
    ...             "run_name": "ResNet50",
    ...             "tags": {
    ...                 "model": "resnet50",
    ...                 "pre-trained": "false",
    ...             },
    ...             "log_system_metrics": True,
    ...             "description": "Simple showcase."
    ...         }
    ...     )

    See Also
    --------
    mlflow.start_run : https://mlflow.org/docs/latest/python_api/mlflow.html#mlflow.start_run
    """

    logger = get_formatted_logger(name="mlflow", level="ERROR")

    if "MLFLOW_ENDPOINT_URL" not in environ:
        logger.error("`MLFLOW_ENDPOINT_URL` environment variable is not set.")
        raise ValueError()

    mlflow_tracking_uri = environ["MLFLOW_ENDPOINT_URL"]

    # Remote server
    artifacts_location = None
    if mlflow_tracking_uri.startswith("http"):

        needed_env_vars = [
            "REQUESTS_CA_BUNDLE",
            "MLFLOW_TRACKING_USERNAME",
            "MLFLOW_TRACKING_PASSWORD",
            "MLFLOW_TRACKING_SERVER_CERT_PATH",
            "MLFLOW_ENDPOINT_URL",
            "MLFLOW_S3_ENDPOINT_URL"
        ]

        for env_var in needed_env_vars:
            if env_var not in environ:
                logger.error(f"`{env_var}` environment variable is not set.")
                raise ValueError()

    # Local server
    elif mlflow_tracking_uri.startswith("sqlite:///"):

        if not mlflow_tracking_uri.endswith("mlflow.db"):
            logger.error(
                f"The provided `mlflow_tracking_uri`='{mlflow_tracking_uri}' is "
                "a local server but it's not ending with 'mlflow.db'."
            )
            raise ValueError()

        local_path = mlflow_tracking_uri.replace("sqlite:///", "").replace("mlflow.db", "")
        artifacts_location = os.path.join(local_path, "artifacts", experiment_name)

    else:
        logger.error(f"Invalid `MLFLOW_ENDPOINT_URL`: {mlflow_tracking_uri}")
        raise ValueError()

    # Set tracking URI
    mlflow.set_tracking_uri(mlflow_tracking_uri)

    # Set workspace if provided
    if "MLFLOW_WORKSPACE" in environ:
        mlflow.set_workspace(environ["MLFLOW_WORKSPACE"])

    # Create the experiment if it doesn't exist
    experiment = mlflow.get_experiment_by_name(experiment_name)
    if experiment is None:
        mlflow.create_experiment(
            name=experiment_name,
            artifact_location=artifacts_location,
        )

    # Set the experiment and start the run
    mlflow.set_experiment(experiment_name)
    mlflow.start_run(**mlflow_kwargs)

log_run_parameters(parameters: Dict[str, Any], prefix: str = '', recursive_logging: bool = False) -> None

Logs a dictionary (including nested dictionaries) to MLFlow as parameters.

By default, all key-value pairs are logged directly. If a value is a dictionary, it is logged recursively using nested key prefixes.

Parameters:
  • parameters (Dict[str, Any]) –

    Dictionary to be logged.

  • prefix (str, default: "" ) –

    Prefix for nested keys.

  • recursive_logging (bool, default: False ) –

    If True, when a value points to a YAML or JSON file, the file is loaded and its contents are recursively logged instead of logging the file path.

Examples:

>>> log_run_parameters({"lr": 0.001, "batch_size": 32})
>>> log_run_parameters({"train": {"lr": 0.001}}, prefix="model")
Source code in aidevelopementtoolkit/logging_utils/mlflow_utils.py
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
def log_run_parameters(
        parameters: Dict[str, Any], 
        prefix: str = "", 
        recursive_logging: bool = False
    ) -> None:
    """
    Logs a dictionary (including nested dictionaries) to MLFlow as parameters.

    By default, all key-value pairs are logged directly. If a value is a dictionary, it is logged recursively using 
    nested key prefixes.

    Parameters
    ----------
    parameters : Dict[str, Any]
        Dictionary to be logged.

    prefix : str, default=""
        Prefix for nested keys.

    recursive_logging : bool, default=False
        If True, when a value points to a YAML or JSON file, the file is loaded 
        and its contents are recursively logged instead of
        logging the file path.

    Examples
    --------
    >>> log_run_parameters({"lr": 0.001, "batch_size": 32})
    >>> log_run_parameters({"train": {"lr": 0.001}}, prefix="model")
    """

    for key, value in parameters.items():
        full_key = f"{prefix}/{key}" if prefix else key

        # Case 1: Nested dictionary, then recurse
        if isinstance(value, dict):
            log_run_parameters(value, full_key, recursive_logging=recursive_logging)
            continue

        # Case 2: Recursive logging enabled, then check if value is a YAML or JSON file
        if recursive_logging and isinstance(value, str) and os.path.isfile(value):
            if value.endswith((".yaml", ".yml")):
                nested_dict = load_file(value)
                if isinstance(nested_dict, dict):
                    log_run_parameters(nested_dict, full_key, recursive_logging=recursive_logging)
                    continue
            elif value.endswith(".json"):
                nested_dict = load_file(value)
                if isinstance(nested_dict, dict):
                    log_run_parameters(nested_dict, full_key, recursive_logging=recursive_logging)
                    continue

        # Case 3: Normal logging (base case)
        mlflow.log_param(full_key, value)

save_model_checkpoint(model: nn.Module, config: Dict[str, Any], checkpoint_name: str) -> None

This function saves the model checkpoints to a temporary local directory and logs that directory as an artifact to MLflow.

Parameters:
  • model (Module) –

    The PyTorch model to be saved.

  • config (Dict[str, Any]) –

    Configuration dictionary used to build the model.

  • checkpoint_name (str) –

    Name of the checkpoint (e.g., "last", "best") to be used in the artifact path in MLflow.

Examples:

>>> save_model_checkpoint(model, {"hidden_dim": 128}, "best")
Source code in aidevelopementtoolkit/logging_utils/mlflow_utils.py
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
210
211
212
213
def save_model_checkpoint(
        model: nn.Module, 
        config: Dict[str, Any],
        checkpoint_name: str,
    ) -> None:
    """This function saves the model checkpoints to a temporary local directory 
    and logs that directory as an artifact to MLflow.

    Parameters
    ----------
    model : nn.Module
        The PyTorch model to be saved.

    config : Dict[str, Any]
        Configuration dictionary used to build the model.

    checkpoint_name : str
        Name of the checkpoint (e.g., "last", "best") to be used in the artifact path in MLflow.

    Examples
    --------
    >>> save_model_checkpoint(model, {"hidden_dim": 128}, "best")
    """

    # Store in a temporary local directory
    local_dir = "temp_model_checkpoints"
    os.makedirs(local_dir, exist_ok=True)
    save_model(model, config, local_dir)

    # Log the directory as an artifact to MLflow
    mlflow.log_artifacts(local_dir, artifact_path=f"model_checkpoints/{checkpoint_name}")

    # Remove directory
    shutil.rmtree(local_dir)

is_numeric(val: Any) -> bool

This function checks whether a given variable is numeric.

Parameters:
  • val (Any) –

    Variable to be checked.

Returns:
  • bool

    True if the variable is numeric, False otherwise.

Source code in aidevelopementtoolkit/logging_utils/mlflow_utils.py
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
def is_numeric(val: Any) -> bool:
    """
    This function checks whether a given variable is numeric.

    Parameters
    ----------
    val : Any
        Variable to be checked.

    Returns
    -------
    bool
        `True` if the variable is numeric, `False` otherwise.
    """
    return isinstance(val, (int, float))