Tinkex.TrainingClient.Operations (Tinkex v0.3.4)

Request building and sending operations for TrainingClient.

This module handles:

Building and sending forward/backward requests
Building and sending optimizer step requests
Building and sending save/load weights requests
Creating sampling clients from saved weights
Handling response parsing and future extraction

Summary

Functions

build_linear_loss_data_safe(data, gradients)

Build linear loss data safely with validation.

ensure_model(opts, session_id, model_seq_id, config, service_api, telemetry_metadata)

Ensure a model exists, creating one if necessary.

handle_load_state_response(future, state, opts)

Handle load state response, polling if necessary.

handle_save_state_response(future, state, opts)

Handle save state response, polling if necessary.

handle_save_weights_response(future, state, opts)

Handle save weights for sampler response, polling if necessary.

merge_custom_metrics(output, metrics)

Merge custom loss metrics into ForwardBackwardOutput.

normalize_save_weights_opts(opts, state)

Normalize save weights options, generating sampling_session_seq_id if needed.

poll_forward_custom_loss(futures, opts, state)

Poll forward futures for custom loss computation.

send_backward_for_custom_loss(linear_data, seq_ids, opts, state)

Send backward pass for custom loss computation.

send_forward_backward_request(chunk, loss_fn, seq_id, opts, state)

Send a forward-backward request and return a future.

send_forward_request(chunk, loss_fn, seq_id, opts, state)

Send a forward-only request and return a future.

send_load_state_request(path, optimizer, seq_id, opts, state)

Send a load state (weights) request.

send_optim_step_request(adam_params, seq_id, opts, state)

Send an optimizer step request and return a future.

send_save_state_request(name, seq_id, opts, state)

Send a save state (weights) request.

send_save_weights_for_sampler_request(seq_id, opts, state)

Send a save weights for sampler request.

start_sampling_client_from_save(save_response, sampling_client_id, opts, state)

Start a sampling client from save response data.

Functions

build_linear_loss_data_safe(data, gradients)

@spec build_linear_loss_data_safe(list(), [Nx.Tensor.t()]) ::
  {:ok, list()} | {:error, Tinkex.Error.t()}

Build linear loss data safely with validation.

ensure_model(opts, session_id, model_seq_id, config, service_api, telemetry_metadata)

@spec ensure_model(keyword(), String.t(), integer(), map(), module(), map()) ::
  {:ok, String.t()} | {:error, Tinkex.Error.t()}

Ensure a model exists, creating one if necessary.

Returns {:ok, model_id} on success.

handle_load_state_response(future, state, opts)

@spec handle_load_state_response(
  map() | Tinkex.Types.LoadWeightsResponse.t(),
  map(),
  keyword()
) ::
  {:ok, Tinkex.Types.LoadWeightsResponse.t() | map()}
  | {:error, Tinkex.Error.t()}

Handle load state response, polling if necessary.

handle_save_state_response(future, state, opts)

@spec handle_save_state_response(
  map() | Tinkex.Types.SaveWeightsResponse.t(),
  map(),
  keyword()
) ::
  {:ok, Tinkex.Types.SaveWeightsResponse.t() | map()}
  | {:error, Tinkex.Error.t()}

Handle save state response, polling if necessary.

handle_save_weights_response(future, state, opts)

@spec handle_save_weights_response(
  map() | Tinkex.Types.SaveWeightsForSamplerResponse.t(),
  map(),
  keyword()
) ::
  {:ok, Tinkex.Types.SaveWeightsForSamplerResponse.t() | map()}
  | {:error, Tinkex.Error.t()}

Handle save weights for sampler response, polling if necessary.

merge_custom_metrics(output, metrics)

@spec merge_custom_metrics(Tinkex.Types.ForwardBackwardOutput.t(), map()) ::
  Tinkex.Types.ForwardBackwardOutput.t()

Merge custom loss metrics into ForwardBackwardOutput.

normalize_save_weights_opts(opts, state)

@spec normalize_save_weights_opts(
  keyword(),
  map()
) :: {keyword(), integer()}

Normalize save weights options, generating sampling_session_seq_id if needed.

Returns {normalized_opts, new_counter}.

poll_forward_custom_loss(futures, opts, state)

@spec poll_forward_custom_loss([map()], keyword(), map()) ::
  {:ok, [Tinkex.Types.ForwardBackwardOutput.t()]} | {:error, Tinkex.Error.t()}

Poll forward futures for custom loss computation.

send_backward_for_custom_loss(linear_data, seq_ids, opts, state)

@spec send_backward_for_custom_loss(list(), [integer()], keyword(), map()) ::
  {:ok, [Tinkex.Types.ForwardBackwardOutput.t()]} | {:error, Tinkex.Error.t()}

Send backward pass for custom loss computation.

send_forward_backward_request(chunk, loss_fn, seq_id, opts, state)

@spec send_forward_backward_request(
  list(),
  atom() | String.t(),
  integer(),
  keyword(),
  map()
) ::
  {:ok, map()} | {:error, Tinkex.Error.t()}

Send a forward-backward request and return a future.

send_forward_request(chunk, loss_fn, seq_id, opts, state)

@spec send_forward_request(list(), atom() | String.t(), integer(), keyword(), map()) ::
  {:ok, map()} | {:error, Tinkex.Error.t()}

Send a forward-only request and return a future.

send_load_state_request(path, optimizer, seq_id, opts, state)

@spec send_load_state_request(String.t(), boolean(), integer(), keyword(), map()) ::
  {:ok, map() | Tinkex.Types.LoadWeightsResponse.t()}
  | {:error, Tinkex.Error.t()}

Send a load state (weights) request.

send_optim_step_request(adam_params, seq_id, opts, state)

@spec send_optim_step_request(map(), integer(), keyword(), map()) ::
  {:ok, map()} | {:error, Tinkex.Error.t()}

Send an optimizer step request and return a future.

send_save_state_request(name, seq_id, opts, state)

@spec send_save_state_request(String.t(), integer(), keyword(), map()) ::
  {:ok, map() | Tinkex.Types.SaveWeightsResponse.t()}
  | {:error, Tinkex.Error.t()}

Send a save state (weights) request.

send_save_weights_for_sampler_request(seq_id, opts, state)

@spec send_save_weights_for_sampler_request(integer(), keyword(), map()) ::
  {:ok, map() | Tinkex.Types.SaveWeightsForSamplerResponse.t()}
  | {:error, Tinkex.Error.t()}

Send a save weights for sampler request.

start_sampling_client_from_save(save_response, sampling_client_id, opts, state)

@spec start_sampling_client_from_save(
  Tinkex.Types.SaveWeightsForSamplerResponse.t() | map(),
  integer(),
  keyword(),
  map()
) :: {:ok, pid()} | {:error, Tinkex.Error.t() | any()}

Start a sampling client from save response data.

Handles both path-based and sampling_session_id-based responses.