API ReferenceAhead-of-time Tuning API

aitune.torch.backend.backend

View as Markdown

Backend interface.

Module Contents

Classes

NameDescription
BackendBackend interface for tuning a module.
BackendBuildStepBase class for backend build step enums.
BackendConfigConfiguration for a backend.
BackendStateEnum representing the state of a backend.
DummyBackendDummy backend for testing purposes.

API

class aitune.torch.backend.backend.Backend()
Abstract

Backend interface for tuning a module.

_build_results
= []
_config
BackendConfig | None = None
_devices
list[str] = ['cpu', 'cuda']
_logger
Logger

Get a logger specific to this backend implementation.

device
device

Get the device of the backend.

is_active
bool

Returns True if the backend is active.

is_jit
bool

Returns True if the backend is a JIT backend.

This method ensures that the backend has this property defined.

name
str

Name of a backend.

state
= BackendState.INIT
aitune.torch.backend.backend.Backend._activate()
abstract

Activates backend.

After activating, the backend should be ready to do inference.

aitune.torch.backend.backend.Backend._assert_device(
    device: torch.device
)

Assert the device of the backend.

Parameters:

device
torch.device

The device to assert.

abstract

Build the model with the given arguments.

After building, the backend should be activated.

Parameters:

module
nn.Module

The module to build the backend on.

name

The name of the backend.

graph_spec
GraphSpec

The graph specification of the backend.

data
list[Sample]

The data to build the backend on.

cache_dir
Path

The cache directory to store the backend artifacts.

aitune.torch.backend.backend.Backend._clean_memory()

Clean up memory.

aitune.torch.backend.backend.Backend._deactivate()
abstract

Deactivates backend.

After deactivating, the backend cannot be used to do inference.

aitune.torch.backend.backend.Backend._deploy()
abstract

Deploys the backend.

After deploying, the backend is ready to do inference. Backend cannot be deactivated anymore.

aitune.torch.backend.backend.Backend._infer(
    args = (),
    kwargs = {}
)
abstract

Run inference with the given arguments.

Parameters:

*args
Defaults to ()

Variable length argument list.

**kwargs
Defaults to {}

Arbitrary keyword arguments.

Returns:

The result of the inference.

aitune.torch.backend.backend.Backend._set_device(
    device: torch.device | None
)

Set the device of the backend.

Parameters:

device
torch.device | None

The device to set the backend on.

aitune.torch.backend.backend.Backend._track_build_step(
    step: aitune.torch.backend.backend.BackendBuildStep
) -> collections.abc.Generator[dict[str, typing.Any], None, None]

Record a build step result.

aitune.torch.backend.backend.Backend.activate()

Activates backend.

After activating, the backend should be ready to do inference.

aitune.torch.backend.backend.Backend.build(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    device: torch.device,
    cache_dir: pathlib.Path,
    log_file: pathlib.Path | None = None
) -> aitune.torch.backend.backend.Backend

Build the model with the given arguments.

Building a backend should be idempotent i.e. do not cause side effects. A model is not necessarily pure functional and can have an internal state (like kv cache for LLMs). That is why build can call a sample of inputs at most once so that subsequent calls have exact same state as the first call for the given sample.

After building, the backend should be activated.

aitune.torch.backend.backend.Backend.deactivate()

Deactivates backend.

After deactivating, the backend cannot be used to do inference.

aitune.torch.backend.backend.Backend.deploy(
    device: torch.device | None
)

Deploys the backend.

After deploying, the backend is ready to do inference. Backend cannot be deactivated anymore.

Parameters:

device
torch.device | None

The device to deploy the backend on.

aitune.torch.backend.backend.Backend.describe() -> str
abstract

Returns the description of the backend.

aitune.torch.backend.backend.Backend.from_dict(
    module: torch.nn.Module | None,
    state_dict: dict
)
classmethodabstract

Creates a backend from a module and state_dict.

aitune.torch.backend.backend.Backend.infer(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Run inference with the given arguments.

Parameters:

args
AnyDefaults to ()

Variable length argument list.

kwargs
AnyDefaults to {}

Arbitrary keyword arguments.

Returns: Any

The result of the inference.

aitune.torch.backend.backend.Backend.key() -> str
abstract

Returns the key of the backend.

aitune.torch.backend.backend.Backend.to_dict()
abstract

Returns the state_dict of the backend.

Note: if there any binary artifacts (files) which should be stored by a backend, they must be passed as Python Path object. Such objects will be bundled with a checkpoint.

class aitune.torch.backend.backend.BackendBuildStep

Bases: enum.Enum

Base class for backend build step enums.

annotation
= 
class aitune.torch.backend.backend.BackendConfig()
Dataclass

Configuration for a backend.

aitune.torch.backend.backend.BackendConfig._default_describe_fields() -> list[str]

Returns the default fields to describe.

aitune.torch.backend.backend.BackendConfig._get_changed_fields(
    current,
    other,
    exclude: list[str] | None = None,
    include: list[str] | None = None
) -> list[str]

Returns the changed fields of the backend configuration.

aitune.torch.backend.backend.BackendConfig.describe() -> str

Describe the backend configuration. Display only changed fields.

classmethod

Initialise config from a plain dict (e.g. parsed from YAML).

The default implementation passes all keys as keyword arguments. Override in subclasses that need type conversion (e.g. nested dicts, enum values, or pickle-serialised objects).

aitune.torch.backend.backend.BackendConfig.key() -> str

Returns the keys of the backend configuration.

aitune.torch.backend.backend.BackendConfig.to_dict()

Returns the state_dict of the backend.

aitune.torch.backend.backend.BackendConfig.to_json(
    path: pathlib.Path
)

Saves the backend configuration to a file.

class aitune.torch.backend.backend.BackendState

Bases: enum.Enum

Enum representing the state of a backend.

State transitions:

stateDiagram
    [*] --> init
    init --> active: When backend is built successfully
    active --> inactive: When backend is deactivated
    inactive --> active: When backend is activated again
    checkpoint_loaded --> deployed: When backend is deployed from checkpoint
    checkpoint_loaded --> active: When backend is activated from checkpoint
    active --> deployed: When backend is deployed after being active
    deployed --> [*]

Note:

  • Changing state to deployed should be done in ModuleWrapper as it does a proxy of forward method and reverses that in order to deploy/recompile jit backends.
  • After loading a checkpoint you can still do activation when grokking with the code but bare in mind that jit backends do recompilation.
ACTIVE
= 'active'
CHECKPOINT_LOADED
= 'checkpoint_loaded'
DEPLOYED
= 'deployed'
INACTIVE
= 'inactive'
INIT
= 'init'
class aitune.torch.backend.backend.DummyBackend()

Bases: Backend

Dummy backend for testing purposes.

_devices
= ['cpu', 'cuda']
aitune.torch.backend.backend.DummyBackend._activate()

Activate the backend.

aitune.torch.backend.backend.DummyBackend._build(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    cache_dir: pathlib.Path
) -> aitune.torch.backend.backend.Backend

Build the model with the given arguments.

aitune.torch.backend.backend.DummyBackend._deactivate()

Deactivate the backend.

aitune.torch.backend.backend.DummyBackend._deploy()

Deploy the backend.

aitune.torch.backend.backend.DummyBackend._infer(
    args = (),
    kwargs = {}
)

Run inference with the given arguments.

aitune.torch.backend.backend.DummyBackend.describe() -> str

Returns the description of the backend.

aitune.torch.backend.backend.DummyBackend.from_dict(
    module: torch.nn.Module | None,
    state_dict: dict
)
classmethod

Creates a backend from a module and state_dict.

aitune.torch.backend.backend.DummyBackend.infer(
    args = (),
    kwargs = {}
)

Run inference with the given arguments.

aitune.torch.backend.backend.DummyBackend.key() -> str

Returns the key of the backend.

aitune.torch.backend.backend.DummyBackend.to_dict()

Returns the state_dict of the backend.