track module

Utilities to support experiment tracking within Armory.

class track.Trackable

Bases: object

A mixin to make a class a trackable, meaning that it will automatically receive the parameters recorded during a tracking context in which the trackable instance is created.

If the subclass has an __init__ function, it must call super().__init__() in order for the mixin to be activated.

This has to be first mixin to work correctly with multiple inheritance.

Example:

from armory.track import Trackable, trackable_context, track_param

class MyClass(Trackable):
    pass

with trackable_context():
    track_param("key", "value")
    my_obj = MyClass()

assert my_obj.tracked_params == {"key": "value"}

track.get_current_params() → Dict[str, Any]

Get the parameters from the current tracking context

Returns:: Current parameters
Return type:: Dict[str, Any]

track.get_current_trackables() → List[Trackable]: Get the trackables from the current context

track.init_tracking_uri(armory_home: Path) → str

Initialize the MLFlow tracking URI.

If the MLFLOW_TRACKING_URI environment variable is set, no change is made to the default tracking URI. Otherwise, the mlruns directory under the given armory home path will be set as the tracking URI.

Parameters:: armory_home (Path) – Path to armory home directory
Returns:: Current MLFlow tracking URI
Return type:: str

track.reset_params(): Clear all parameters in the current tracking context

track.server(): Start the MLFlow server

track.track_call(func: Callable[[P], T], *args: P, **kwargs: P) → T

Invoke the given function or class’s initializer with the given arguments, recording keyword arguments as parameters in the current tracking context to be logged with MLFlow.

The effect is functionally equivalent to using the track_params or track_init_params decorators in an inline fashion with no prefix or ignored arguments. In order to use a custom prefix or to ignore arguments, the track_params or track_init_params functions must be used directly.

Example:

from armory.track import track_call

class MyDataset:
    def __init__(self, batch_size: int):
        pass


def load_model(name: str):
    pass

ds = track_call(MyDataset, batch_size=32)
model = track_call(load_model, name="model")

Parameters:

func (Callable[P, T]) – Function or class to be invoked
*args –
Positional arguments to be passed to the function or class
**kwargs –
Keyword arguments to be passed to the function or class, which will be recorded as parameters

Returns:

Result of the function or class invocation

Return type:

T

track.track_init_params(*, prefix: str | None = None, ignore: Sequence[str] | None = None) → Callable[[T], T]

track.track_init_params(_cls: T, *, prefix: str | None = None, ignore: Sequence[str] | None = None) → T

Create a decorator to record class dunder-init keyword arguments as parameters in the current tracking context to be logged with MLFlow.

Example:

from charmory.track import track_init_params

@track_init_params()
class MyDataset:
    def __init__(self, batch_size: int):
        pass

# Or for a third-party class that cannot have the decorator
# already applied, you can apply it inline
obj = track_init_params(ThirdPartyClass)(arg=42)

Parameters:

_cls (object, optional) – Optional class to be decorated, defaults to None
prefix (str, optional) – Optional prefix for all keyword argument names (default is inferred from decorated class name)
ignore (Sequence[str], optional) – Optional list of keyword arguments to be ignored, defaults to None

Returns:

Decorated class if _cls was provided, else a class decorator

track.track_metrics(metrics: Mapping[str, float | Sequence[float] | Tensor])

Log the given metrics with MLFlow.

Parameters:: metrics (Mapping[str, Union[float, Sequence[float], torch.Tensor]]) – _Mapping of metrics names to values

track.track_param(key: str, value: Any)

Record a parameter in the current tracking context to be logged with MLFlow.

Example:

from charmory.track import track_param

track_param("key", "value")

Parameters:

key (str) – Parameter name (should be unique or will overwrite previous values)
value (Any) – Parameter value

track.track_params(*, prefix: str | None = None, ignore: Sequence[str] | None = None) → Callable[[Callable[[P], T]], Callable[[P], T]]

track.track_params(_func: Callable[[P], T], *, prefix: str | None = None, ignore: Sequence[str] | None = None) → Callable[[P], T]

Create a decorator to record function keyword arguments as parameters in the current tracking context to be logged with MLFlow.

Example:

from charmory.track import track_params

@track_params()
def load_model(name: str, batch_size: int):
    pass

# Or for a third-party function that cannot have the decorator
# already applied, you can apply it inline
track_params(third_party_func)(arg=42)

:param _func:Optional function to be decorated, defaults to None :type _func: Callable, optional :param prefix: Optional prefix for all keyword argument names (default is

inferred from decorated function name), defaults to None

Parameters:: ignore (Sequence[str], optional) – Optional list of keyword arguments to be ignored, defaults to None
Returns:: Decorated function if _func was provided, else a function decorator
Return type:: Callable

track.track_system_metrics(run_id: str | None)

Create a context in which to track system metrics and log them to the given MLflow experiment run. System metrics include CPU, disk, and network utilization metrics. If the pynvml package is installed, then GPU utilization metrics will also be collected.

If the run ID is empty or undefined (which would be the case if running distributed on anything other than the rank zero node), then no tracking will be enabled.

Example:

from charmory.track import track_system_metrics
import mlflow

with mlflow.start_run() as active_run:
    with track_system_metrics(active_run.info.run_id)
        # Do something

Parameters:: run_id (str, optional) – MLflow experiment run ID of the run to which to record the system metrics. If empty or undefined, no tracking will be enabled
Returns:: Context manager
Return type:: _type_

track.trackable_context(nested: bool = False)

Create a new trackable context. Parameters recorded while the context is active will be automatically associated with any Trackable objects creating while the context is active. When the context begins, a new tracking context is created to isolate parameters from other contexts. Upon completion of the context, the tracked parameters are associated with the trackables and all parameters will be cleared.

Example:

from armory.track import Trackable, trackable_context, track_param

class MyClass(Trackable):
    pass

with trackable_context():
    track_param("key", "value")
    my_obj = MyClass()
    assert get_current_params() == {"key": "value"}

assert get_current_params() == {}
assert my_obj.tracked_params == {"key": "value"}

Parameters:: nested (bool, optional) – Copy parameters from the current context into the new tracking context, defaults to False
Returns:: Context manager

track.tracking_context(nested: bool = False)

Create a new tracking context. Parameters recorded while the context is active will be isolated from other contexts. Upon completion of the context, all parameters will be cleared.

Example:

from charmory.track import tracking_context, track_param

with tracking_context():
    track_param("key", "value1")

with tracking_context():
    track_param("key", "value2")

Parameters:: nested (bool, optional) – Copy parameters from the current context into the new tracking context, defaults to False
Returns:: Context manager
Return type:: _type_