aitune.torch.tune_strategy.profiling_tune_strategy

Shared base for profiling-based tune strategies.

A profiling strategy profiles a TorchEager baseline, then builds, validates, and profiles every user-provided backend, selecting the one whose profiled metric is best. Subclasses define the metric (throughput, latency, …) by setting a few class attributes and implementing a small set of hooks; MaxThroughputStrategy, MinLatencyStrategy, and LatencyBudgetStrategy are the concrete implementations.

Module Contents

Classes

Name	Description
`BackendPerfResult`	Profiled metric and speedup result for a single backend.
`BackendProfilingResult`	Profiled backend result returned by profiling strategies.
`ProfilingTuneStrategy`	Base class for strategies that select a backend by a profiled metric.
`_TuneCandidate`	-

API

class aitune.torch.tune_strategy.profiling_tune_strategy.BackendPerfResult(
    backend_description: str,
    metric: float,
    baseline_metric: float,
    speedup: float,
    passed: bool
)

Dataclass

Profiled metric and speedup result for a single backend.

backend_description

str

baseline_metric

float

metric

float

passed

bool

speedup

float

class aitune.torch.tune_strategy.profiling_tune_strategy.BackendProfilingResult(
    selected_batch_size: int
)

DataclassAbstract

Profiled backend result returned by profiling strategies.

metric

float

Returns the scalar metric used to compare candidates.

selected_batch_size

int

aitune.torch.tune_strategy.profiling_tune_strategy.BackendProfilingResult.to_json_dict(
    metric_label: str
) -> dict[str, int | float]

Returns fields stored in strategy_results for this backend.

class aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy(
    backends: list[aitune.torch.backend.Backend] | None = None,
    profiling_config: aitune.torch.task.profiling.ProfilingConfig | None = None,
    kwargs: typing.Any = {}
)

Bases: FindMaxBatchSizeMixin

Base class for strategies that select a backend by a profiled metric.

Subclasses set _title, _description, _metric_label (e.g. “throughput”), _metric_unit (e.g. “samples/s”) and _value_fmt (a format spec such as “.2f”), and implement :meth:_measure, :meth:_is_better, and :meth:_speedup.

TorchEager is profiled in _pre_tune as a baseline (not injected into the backends list). When performance validation is enabled (default), the strategy falls back to TorchEager when no user-provided backend beats it. When disabled, the best user-provided backend wins regardless of speed, and the strategy raises if all user backends fail.

_backends

_baseline_backend

Backend | None = None

_baseline_result

BackendProfilingResult | None = None

_description

str = ''

_metric_label

str = ''

_metric_unit

str = ''

_performance_validation_enabled

bool = True

_title

str = ''

_value_fmt

str = '.2f'

perf_validation_results

list[BackendPerfResult] = []

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._default_backends() -> list[aitune.torch.backend.Backend]

Returns default backends.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._describe_parts() -> list[str]

Returns the parts of the description.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._fmt(
    value: float
) -> str

Formats a metric value with its unit, e.g. 12.34 samples/s.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._get_profiling_config(
    batching: bool,
    max_batch_size: int
) -> aitune.torch.task.profiling.ProfilingConfig

Gets profiling configuration.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._is_better(
    result: aitune.torch.tune_strategy.profiling_tune_strategy.BackendProfilingResult,
    other: aitune.torch.tune_strategy.profiling_tune_strategy.BackendProfilingResult
) -> bool

abstract

Returns True when result is better than other.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._log_baseline_selected(
    backend: aitune.torch.backend.Backend
) -> None

Emit an explicit message when the TorchEager baseline is the selected backend.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._measure(
    backend: aitune.torch.backend.Backend,
    name: str,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    profiling_cfg: aitune.torch.task.profiling.ProfilingConfig
) -> aitune.torch.tune_strategy.profiling_tune_strategy.BackendProfilingResult

abstract

Profiles the backend and returns its result.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._post_tune(
    backend: aitune.torch.backend.Backend | None,
    name: str,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample]
)

Emits a speedup line after tuning completes.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._pre_tune(
    module: torch.nn.Module,
    name: str,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    device: torch.device,
    cache_dir: pathlib.Path
)

Calls super()._pre_tune() (finds max batch size) then profiles TorchEager as baseline.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._record_perf_result(
    backend: aitune.torch.backend.Backend,
    result: aitune.torch.tune_strategy.profiling_tune_strategy.BackendProfilingResult
) -> None

Appends a BackendPerfResult for the given backend if a baseline is available.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._resolve_winner(
    best: aitune.torch.tune_strategy.profiling_tune_strategy._TuneCandidate | None
) -> aitune.torch.tune_strategy.profiling_tune_strategy._TuneCandidate

Returns the winning candidate, falling back to the TorchEager baseline when appropriate.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._run_backends(
    module: torch.nn.Module,
    name: str,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    device: torch.device,
    cache_dir: pathlib.Path,
    batching: bool,
    max_batch_size: int
) -> aitune.torch.tune_strategy.profiling_tune_strategy._TuneCandidate | None

Builds, validates, and profiles each backend; returns the best candidate.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._speedup(
    result: aitune.torch.tune_strategy.profiling_tune_strategy.BackendProfilingResult,
    baseline_result: aitune.torch.tune_strategy.profiling_tune_strategy.BackendProfilingResult
) -> float

abstract

Returns the speedup of result relative to baseline_result (>1 is faster).

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy._tune(
    module: torch.nn.Module,
    name: str,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    device: torch.device,
    cache_dir: pathlib.Path
) -> aitune.torch.backend.Backend

Tunes given torch module with provided graph_spec and data.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy.enable_performance_validation(
    enable: bool = True
) -> aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy

Enables or disables baseline validation.

aitune.torch.tune_strategy.profiling_tune_strategy.ProfilingTuneStrategy.to_json_dict() -> dict[str, typing.Any]

Returns config dict for the strategy.

class aitune.torch.tune_strategy.profiling_tune_strategy._TuneCandidate(
    backend: aitune.torch.backend.Backend,
    result: aitune.torch.tune_strategy.profiling_tune_strategy.BackendProfilingResult
)

Dataclass

backend

Backend

result

BackendProfilingResult