"""Turn a function into a tracked run, with lineage carried on its result."""
from __future__ import annotations
import functools
import inspect
from collections.abc import Callable
from dataclasses import dataclass
from typing import Concatenate, Generic, Protocol, TypeVar, cast, overload
from embedl_hub._internal.tracking.client import Client, _resolve_run_type
from embedl_hub._internal.tracking.rest_api import Run, RunType
_T = TypeVar("_T")
_V_co = TypeVar("_V_co", covariant=True)
_A0 = TypeVar("_A0")
_A1 = TypeVar("_A1")
_A2 = TypeVar("_A2")
_A3 = TypeVar("_A3")
_A4 = TypeVar("_A4")
_A5 = TypeVar("_A5")
_PLACEHOLDER = getattr(functools, "Placeholder", object())
[docs]
@dataclass(frozen=True)
class Tracked(Generic[_V_co]):
"""A value together with the run that produced it.
Returned by every :func:`hub_run` function. Passing it to another
:func:`hub_run` function records that run as the parent of the new
one, and hands the function :attr:`value` rather than the wrapper.
:param value: What the function returned.
:param run: The run the function ran in, or ``None`` when it ran
with ``safe=True`` and the run could not be created.
"""
value: _V_co
run: Run | None
class _HubRunDecorator(Protocol):
@overload
def __call__(
self, fn: Callable[[Client], _T], /
) -> Callable[..., Tracked[_T]]: ...
@overload
def __call__(
self, fn: Callable[[Client, _A0], _T], /
) -> Callable[
Concatenate[Client, _A0 | Tracked[_A0], ...], Tracked[_T]
]: ...
@overload
def __call__(
self, fn: Callable[[Client, _A0, _A1], _T], /
) -> Callable[
Concatenate[Client, _A0 | Tracked[_A0], _A1 | Tracked[_A1], ...],
Tracked[_T],
]: ...
@overload
def __call__(
self, fn: Callable[[Client, _A0, _A1, _A2], _T], /
) -> Callable[
Concatenate[
Client,
_A0 | Tracked[_A0],
_A1 | Tracked[_A1],
_A2 | Tracked[_A2],
...,
],
Tracked[_T],
]: ...
@overload
def __call__(
self, fn: Callable[[Client, _A0, _A1, _A2, _A3], _T], /
) -> Callable[
Concatenate[
Client,
_A0 | Tracked[_A0],
_A1 | Tracked[_A1],
_A2 | Tracked[_A2],
_A3 | Tracked[_A3],
...,
],
Tracked[_T],
]: ...
@overload
def __call__(
self, fn: Callable[[Client, _A0, _A1, _A2, _A3, _A4], _T], /
) -> Callable[
Concatenate[
Client,
_A0 | Tracked[_A0],
_A1 | Tracked[_A1],
_A2 | Tracked[_A2],
_A3 | Tracked[_A3],
_A4 | Tracked[_A4],
...,
],
Tracked[_T],
]: ...
@overload
def __call__(
self, fn: Callable[[Client, _A0, _A1, _A2, _A3, _A4, _A5], _T], /
) -> Callable[
Concatenate[
Client,
_A0 | Tracked[_A0],
_A1 | Tracked[_A1],
_A2 | Tracked[_A2],
_A3 | Tracked[_A3],
_A4 | Tracked[_A4],
_A5 | Tracked[_A5],
...,
],
Tracked[_T],
]: ...
@overload
def __call__(
self, fn: Callable[..., _T], /
) -> Callable[..., Tracked[_T]]: ...
[docs]
def hub_run(
run_type: RunType | str,
*,
name: str | None = None,
safe: bool = False,
) -> _HubRunDecorator:
"""Run each call of the decorated function as a Hub run.
The function takes the :class:`~embedl_hub.tracking.Client` as its
first parameter, after ``self`` on a method. Each call starts a run
with :meth:`~embedl_hub.tracking.Client.start_run` and returns the
result wrapped in a :class:`Tracked`. Every :class:`Tracked`
argument, positional or keyword, is unwrapped to its value before
the call, and the first one that has a run becomes the parent of the
new run. :class:`Tracked` values inside containers are passed as they
are.
.. code-block:: python
@hub_run("train")
def train(client: Client, data: Path) -> Path: ...
@hub_run("evaluate")
def evaluate(client: Client, model: Path) -> float: ...
model = train(client, data)
score = evaluate(client, model) # a child run of train
Type checkers check the required arguments after the client, up to
six, each as its declared type or a :class:`Tracked` of it. Pass the
client and these arguments positionally; given by keyword, they are
reported as missing. Arguments for parameters with defaults are not
checked. Methods, and functions with more than six required
arguments, are checked only for their return type.
:param run_type: A :class:`~embedl_hub.tracking.RunType` or a custom
type label, as :meth:`~embedl_hub.tracking.Client.start_run`
takes.
:param name: The run name. Defaults to the function's name.
:param safe: Passed to :meth:`~embedl_hub.tracking.Client.start_run`:
when the run cannot be created or finished, warn instead of
raising. A function whose run could not be created still runs,
and returns a :class:`Tracked` whose ``run`` is ``None``.
:return: A decorator for a synchronous function.
:raises TypeError: If the decorated function is async or a generator,
whose run would finish before its body runs; has no parameter
for the client; or is a ``classmethod`` or ``staticmethod``
object, which must be applied above :func:`hub_run` instead.
At call time, if the client argument is not a
:class:`~embedl_hub.tracking.Client`.
:raises ValueError: If ``run_type`` is ``RunType.CUSTOM``, which
needs a label: pass the label itself instead.
:raises TrackingUsageError: At call time, if a run is already active
on the client: runs do not nest.
"""
try:
_resolve_run_type(run_type)
except ValueError as exc:
raise ValueError(
"hub_run takes a custom run type as its label, e.g. "
"@hub_run('analyse'), not RunType.CUSTOM."
) from exc
def decorate(fn: Callable[..., _T]) -> Callable[..., Tracked[_T]]:
fn_name = getattr(fn, "__name__", type(fn).__name__)
if isinstance(fn, (classmethod, staticmethod)):
raise TypeError(
f"'{fn_name}' is a {type(fn).__name__} object. Put "
f"@{type(fn).__name__} above @hub_run, so hub_run wraps "
"the function itself."
)
target, bound_args, bound_kwargs = _unwrap_partial(fn)
body = (
target
if inspect.isfunction(target) or inspect.ismethod(target)
else getattr(type(target), "__call__", target)
)
if (
inspect.iscoroutinefunction(body)
or inspect.isgeneratorfunction(body)
or inspect.isasyncgenfunction(body)
):
raise TypeError(
f"'{fn_name}' is an async or generator function, which "
"hub_run cannot track: the run would finish before its "
"body runs. Decorate a function that returns its result, "
"or call Client.start_run inside the body instead."
)
run_name = name or fn_name
signature = inspect.signature(target)
client_param = _client_parameter(signature, fn_name)
@functools.wraps(fn)
def wrapper(*args: object, **kwargs: object) -> Tracked[_T]:
all_args = _fill_placeholders(bound_args, args)
missing = sum(
value is _PLACEHOLDER for value in all_args[: len(bound_args)]
)
if missing:
raise TypeError(
"missing positional arguments in 'partial' call for "
f"{missing} placeholder(s)"
)
all_kwargs = {**bound_kwargs, **kwargs}
bound = signature.bind(*all_args, **all_kwargs)
bound.apply_defaults()
client = bound.arguments.get(client_param.name)
if client_param.kind is inspect.Parameter.VAR_POSITIONAL:
client = client[0] if client else None
if not isinstance(client, Client):
raise TypeError(
f"'{fn_name}' must be called with a Client as its "
f"'{client_param.name}' argument, got "
f"{type(client).__name__}. Pass the client the run "
"should be tracked on."
)
parent = _first_parent(all_args, all_kwargs)
call_args = tuple(_unwrap(a) for a in all_args)
call_kwargs = {k: _unwrap(v) for k, v in all_kwargs.items()}
with client.start_run(
run_type,
name=run_name,
parent_run_id=parent.id if parent is not None else None,
safe=safe,
) as run:
value = target(*call_args, **call_kwargs)
return Tracked(cast(_T, value), run)
return wrapper
return cast(_HubRunDecorator, decorate)
def _unwrap_partial(
fn: Callable[..., object],
) -> tuple[Callable[..., object], tuple[object, ...], dict[str, object]]:
args: tuple[object, ...] = ()
kwargs: dict[str, object] = {}
while isinstance(fn, functools.partial):
args = _fill_placeholders(fn.args, args)
kwargs = {**fn.keywords, **kwargs}
fn = fn.func
return fn, args, kwargs
def _fill_placeholders(
bound: tuple[object, ...], args: tuple[object, ...]
) -> tuple[object, ...]:
rest = iter(args)
filled = tuple(
next(rest, _PLACEHOLDER) if value is _PLACEHOLDER else value
for value in bound
)
return (*filled, *rest)
def _client_parameter(
signature: inspect.Signature, fn_name: str
) -> inspect.Parameter:
params = list(signature.parameters.values())
if params and params[0].name in ("self", "cls"):
params = params[1:]
if not params or params[0].kind is inspect.Parameter.VAR_KEYWORD:
raise TypeError(
f"'{fn_name}' needs a parameter for the Client, first or "
"right after self or cls. Add one, e.g. "
f"'def {fn_name}(client, ...)'."
)
return params[0]
def _first_parent(
args: tuple[object, ...], kwargs: dict[str, object]
) -> Run | None:
for value in (*args, *kwargs.values()):
if isinstance(value, Tracked) and value.run is not None:
return value.run
return None
def _unwrap(value: object) -> object:
return value.value if isinstance(value, Tracked) else value
