src.utils

GPU

The GPU utilities.

src.utils.gpu.set_gpu_power_limit_if_needed

Helper function, that sets GPU power limit if specified GPU is used.

Parameters:

Name	Type	Description	Default
gpu_name	str	The name of the GPU to assign the new power limit to. Default: "NVIDIA GeForce RTX 3090".	'NVIDIA GeForce RTX 3090'
pw	int	The new power limit to set. Defaults to 250W.	250

Source code in src/utils/gpu.py

def set_gpu_power_limit_if_needed(gpu_name: str = "NVIDIA GeForce RTX 3090", pw: int = 250) -> None:
    """Helper function, that sets GPU power limit if specified GPU is used.

    Args:
        gpu_name: The name of the GPU to assign the new power limit to. Default: "NVIDIA GeForce RTX 3090".
        pw: The new power limit to set. Defaults to 250W.

    """
    gpu_list = os.popen("/usr/bin/nvidia-smi --query-gpu=gpu_name --format=csv").read()
    if gpu_name in gpu_list:
        os.system("/usr/bin/sudo /usr/bin/nvidia-smi -pm 1")
        os.system(f"/usr/bin/sudo /usr/bin/nvidia-smi -pl {pw}")

Logging

The logging utilities.

src.utils.logging.get_logger

Builds a Logger instance with provided name and log level.

Parameters:

Name	Type	Description	Default
name	str	The name for the logger.	required
log_level	int | str	The default log level.	INFO

Returns:

Type	Description
Logger	The logger.

Source code in src/utils/logging.py

def get_logger(name: str, log_level: int | str = logging.INFO) -> logging.Logger:
    """Builds a `Logger` instance with provided name and log level.

    Args:
        name: The name for the logger.
        log_level: The default log level.

    Returns:
        The logger.

    """
    logger = logging.getLogger(name=name)
    logger.setLevel(log_level)

    stream_handler = logging.StreamHandler()
    formatter = logging.Formatter(fmt=consts.logging.FORMAT)
    stream_handler.setFormatter(fmt=formatter)
    logger.addHandler(stream_handler)

    return logger

src.utils.logging.timed

This decorator prints the execution time for the decorated function.

Parameters:

Name	Type	Description	Default
func	Callable[P, T]	The function to wrap.	required

Returns:

Type	Description
Callable[P, T]	Wrapper around the function.

Source code in src/utils/logging.py

def timed(func: Callable[P, T]) -> Callable[P, T]:
    """This decorator prints the execution time for the decorated function.

    Args:
        func: The function to wrap.

    Returns:
        Wrapper around the function.

    """

    @wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
        _timed_logger.info("%(func_name)s is running...", {"func_name": func.__qualname__})
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        _timed_logger.info(
            "%(func_name)s ran in %(execution_time)s",
            {
                "func_name": func.__qualname__,
                "execution_time": f"{(end - start):.4f}",
            },
        )
        return result

    return wrapper

MLFlow

The MLFlow utilities.

src.utils.mlflow.resolve_experiment_name

Resolves MLFlow experiment name.

If environment variable "MLFLOW_EXPERIMENT_NAME" is set, then the experiment name will be resolved based on this environment variable. Otherwise, the default experiment name, passed as and argument will be used.

Useful when working with Azure ML.

Notes

This function will set the environment variable "MLFLOW_EXPERIMENT_NAME" if it is not set.

Parameters:

Name	Type	Description	Default
default_experiment_name	str	The default experiment name to use if environment variable is not set.	required

Examples:

When env variable is unset, the default exp name passed as argument will be used.

>>> del os.environ["MLFLOW_EXPERIMENT_NAME"]
>>> resolve_experiment_name("custom-mlflow-experiment-name")
'custom-mlflow-experiment-name'

Otherwise, the default exp name indicated by the env var will be used.

>>> os.environ["MLFLOW_EXPERIMENT_NAME"] = "env-defined-mlflow-experiment-name"
>>> resolve_experiment_name("different-mlflow-experiment-name")
'env-defined-mlflow-experiment-name'

Source code in src/utils/mlflow.py

def resolve_experiment_name(default_experiment_name: str) -> str:
    """Resolves MLFlow experiment name.

    If environment variable `"MLFLOW_EXPERIMENT_NAME"` is set, then the experiment name
    will be resolved based on this environment variable. Otherwise, the default experiment name, passed as
    and argument will be used.

    Useful when working with Azure ML.

    Notes:
        This function will set the environment variable `"MLFLOW_EXPERIMENT_NAME"` if it is not set.

    Args:
        default_experiment_name: The default experiment name to use if environment variable is not set.

    Returns: The resolved experiment name.

    Examples:
        When env variable is unset, the default exp name passed as argument will be used.

        >>> del os.environ["MLFLOW_EXPERIMENT_NAME"]
        >>> resolve_experiment_name("custom-mlflow-experiment-name")
        'custom-mlflow-experiment-name'

        Otherwise, the default exp name indicated by the env var will be used.

        >>> os.environ["MLFLOW_EXPERIMENT_NAME"] = "env-defined-mlflow-experiment-name"
        >>> resolve_experiment_name("different-mlflow-experiment-name")
        'env-defined-mlflow-experiment-name'

    """
    if os.environ.get("MLFLOW_EXPERIMENT_NAME", None) is None:
        os.environ["MLFLOW_EXPERIMENT_NAME"] = default_experiment_name
    return os.environ["MLFLOW_EXPERIMENT_NAME"]

src.utils.mlflow.run_id_from_context

Resolves the MLFlow Run ID from the context.

Returns:

Type	Description
str | None	The MLFlow Run ID based on "MLFLOW_RUN_ID" environment variable. If not set, returns None.

Source code in src/utils/mlflow.py

def run_id_from_context() -> str | None:
    """Resolves the MLFlow Run ID from the context.

    Returns:
        The MLFlow Run ID based on `"MLFLOW_RUN_ID"` environment variable. If not set, returns `None`.

    """
    return os.environ.get("MLFLOW_RUN_ID", None)

Serialization

The serialization utils.

src.utils.serialization.JsonEncoder

Bases: JSONEncoder

Custom JSON encoder that handles datatypes that are not out-of-the-box supported by the json package.

Source code in src/utils/serialization.py

class JsonEncoder(JSONEncoder):
    """Custom JSON encoder that handles datatypes that are not out-of-the-box supported by the `json` package."""

    def default(self, o: Any) -> str:
        if isinstance(o, (date, datetime)):
            return o.isoformat()

        if isinstance(o, Path):
            return o.as_posix()

        return super().default(o)  # type: ignore[no-any-return]