.. _edge-client-sdk-label:

=================
Edge Client SDK
=================

``EdgeClient`` is the object injected into your ``startup()`` function and is the primary interface between your ML code and the Scaleout Edge runtime. It handles callback registration, metric and telemetry logging, task-abort signalling, inference helpers, and outbound message prioritisation.

ScaleoutModel
=============

``ScaleoutModel`` is the immutable container passed to your training and validation callbacks. It holds serialized model weights, an optional full model representation (e.g. ONNX), and a metadata dict.

**Reading parameters** — use ``get_training_model(helper)`` to deserialize the weights in your callback:

.. code-block:: python

   from scaleoututil.helpers.helpers import get_helper

   helper = get_helper("numpyhelper")

   def train(self, scaleout_model: ScaleoutModel, settings: dict):
       params = scaleout_model.get_training_model(helper)
       # ... train ...

**Creating a model** — wrap your updated parameters with ``ScaleoutModel.from_training_model()`` before returning from the train callback:

.. code-block:: python

   return ScaleoutModel.from_training_model(updated_params, helper), metadata

**Useful properties:**

- ``model.model_id`` — unique identifier assigned by the platform
- ``model.metadata`` — shallow copy of the metadata dict (``model_id``, ``helper_type``, any user keys)
- ``model.has_inference_model`` — ``True`` if a full model artifact (e.g. ONNX) is bundled
- ``model.get_inference_model_format()`` — format string (``"onnx"``, ``"pt"``, …) or ``None``

To bundle an inference artifact alongside the weights, or to attach custom metadata, use the builder:

.. code-block:: python

   from scaleoututil.utils.model import ScaleoutModelBuilder

   model = (
       ScaleoutModelBuilder.from_training_model(params, helper)
       .set_metadata("accuracy", 0.94)
       .set_inference_model(onnx_bytes, "onnx")
       .build()
   )


Callback Registration
=====================

Register your ML callbacks inside ``startup()``. The client calls these when a task arrives from the server. See :doc:`project instructions <projects>` for more info on ``startup()``. 

``set_train_callback(callback)``
---------------------------------

Register the function that performs a training round.

The callback receives:

- ``scaleout_model`` (``ScaleoutModel``) — the current global model
- ``settings`` (``dict``) — hyperparameters sent by the server (epochs, learning rate, etc.)

It must return a tuple ``(updated_model, metadata_dict)``.

.. code-block:: python

   def train(self, scaleout_model: ScaleoutModel, settings: dict):
       model = load_parameters(scaleout_model)
       # ... training loop ...
       self.client.check_task_abort()
       self.client.log_metric({"loss": float(loss)})
       return save_parameters(model), {"num_examples": len(x_train)}

   client.set_train_callback(self.train)

``set_validate_callback(callback)``
-------------------------------------

Register the function that evaluates the global model after aggregation.

The callback receives:

- ``scaleout_model`` (``ScaleoutModel``) — the model to evaluate

It must return a ``dict`` of validation metrics. Scalar values are captured and visualised in the UI.

``set_inference_callback(callback)``
--------------------------------------

Register the function called during live inference.

The callback receives:

- ``model`` (``ScaleoutModel``) — the staged model
- ``params`` (``dict``) — runtime parameters

``set_stage_model_callback(callback)``
----------------------------------------

Register a function that is called immediately after a model is staged for inference via ``stage_model()``. Use this to load model weights into memory or prepare hardware.

``set_custom_callback(name, callback)`` / ``remove_custom_callback(name)``
---------------------------------------------------------------------------

Register (or remove) a named callback that the server can invoke with arbitrary parameters. The ``name`` is automatically prefixed with ``Custom_`` if not already present.

The callback receives a ``TaskRequest`` protobuf and must return a ``dict``.

.. code-block:: python

   def my_command(self, request):
       params = dict(request.data)
       return {"status": "ok", "received": params}

   client.set_custom_callback("my_command", self.my_command)


Logging & Observability
========================

All ``log_*`` methods enqueue a message to the priority queue for delivery to the server. They all accept an optional ``priority`` keyword argument — see `Priority Messages`_ for details on how to control delivery order.

``log_metric(metrics, step, commit, check_task_abort, context)``
-----------------------------------------------------------------

Log scalar metrics tied to the current training round. These appear in the Scaleout Edge UI.

Requires an active ``LoggingContext`` (set automatically during training callbacks). Pass ``context=`` to override the thread-local context explicitly.

.. code-block:: python

   self.client.log_metric({"training_loss": 0.42, "accuracy": 0.91})

Parameters:

- ``metrics`` (``dict``) — key/value pairs of scalar metrics
- ``step`` (``int``, optional) — explicitly set the step counter; uses context step if omitted
- ``commit`` (``bool``, default ``True``) — increment the step counter after logging
- ``check_task_abort`` (``bool``, default ``True``) — check for server-side abort after logging
- ``context`` (``LoggingContext``, optional) — override the thread-local context

``log_attributes(attributes, priority, check_task_abort)``
-----------------------------------------------------------

Log structured key/value data as an artifact. Delivered with ``"artifact"`` priority by default.

.. code-block:: python

   self.client.log_attributes({"model_version": "v3", "framework": "pytorch"})

``log_telemetry(key, payload, priority, check_task_abort)``
-----------------------------------------------------------

Log a time-series observation. Delivered with ``"telemetry"`` priority by default (LIFO, demoted to ``backlog`` after 300 s of queuing).

.. code-block:: python

   self.client.log_telemetry("cpu_temp", {"value": 74.2, "unit": "celsius"})

Parameters:

- ``key`` (``str``) — observation name
- ``payload`` (``dict``, optional) — JSON-serialisable data; use ``{"value": <float>}`` for plottable scalars
- ``priority`` (``str``, optional) — override the default ``"telemetry"`` class

.. note::
   Passing a bare ``dict`` as the first positional argument (``log_telemetry({"loss": 0.5})``) is still accepted for backwards compatibility but emits a ``DeprecationWarning``. Prefer the ``key`` + ``payload`` form.

``log_inference_result(key, payload, context, priority)``
----------------------------------------------------------

Log a labelled inference observation (sighting, classification hit, etc.). Delivered with ``"inference"`` priority by default.

Requires an active ``InferenceContext`` which already exists for commands initiated from the control plane. Returns ``False`` and logs a warning if none is set.

.. code-block:: python

   with self.client.inference_context(InferenceContext(inference_id="run-1", model_id=model.id)):
       result = model.predict(frame)
       self.client.log_inference_result("detection", {"bbox": result.bbox, "score": result.score})


Task Control
============

``check_task_abort()``
-----------------------

Raises ``StoppedException`` if the server has aborted the current task. Call this periodically inside long training loops so the client can stop cleanly when requested.

.. code-block:: python

   for batch in dataloader:
       self.client.check_task_abort()
       # ... process batch ...


Inference
=========

``stage_model(model)``
-----------------------

Download and cache a model locally for inference. Accepts either a ``ScaleoutModel`` or a model-ID string. Returns the staged ``ScaleoutModel``. Triggers ``stage_model_callback`` if one is registered.

``run_inference(model, params)``
---------------------------------

Stage the model (if needed) and invoke the registered ``inference_callback``. Raises ``ValueError`` if no callback is registered or the model cannot be found.


Logging and Inference Contexts
================================

Contexts carry the session, model, and round identifiers that ``log_metric`` and ``log_inference_result`` stamp on every outgoing record. They are set automatically during training and validation callbacks. Use the context managers to set them explicitly in other code paths.

Logging Context
---------------

``LoggingContext`` carries ``model_id``, ``round_id``, ``session_id``, and a ``step`` counter.

.. code-block:: python

   from scaleout.client.logging_context import LoggingContext

   ctx = LoggingContext(model_id="m-1", round_id="r-3", session_id="s-7")
   with client.logging_context(ctx):
       client.log_metric({"eval_loss": 0.31})

The current context for the running thread is also accessible as ``client.current_logging_context``.

Inference Context
-----------------

``InferenceContext`` carries ``inference_id`` and ``model_id``.

.. code-block:: python

   from scaleout.client.inference_context import InferenceContext

   ctx = InferenceContext(inference_id="infer-42", model_id="m-1")
   with client.inference_context(ctx):
       client.log_inference_result("classification", {"label": "cat", "confidence": 0.98})

The current context for the running thread is accessible as ``client.current_inference_context``.


Priority Messages
=================

Every ``log_*`` call enqueues a message. The runtime drains the queue using weighted priority scheduling, so high-urgency records (alerts, model updates) are sent before lower-priority bulk data (telemetry, inference logs).

Default Priority Classes
------------------------

Six classes are registered by default:

.. list-table::
   :header-rows: 1
   :widths: 20 15 15 50

   * - Class
     - Priority
     - Order
     - Decay
   * - ``alert``
     - 80
     - FIFO
     - —
   * - ``model_update``
     - 70
     - FIFO
     - —
   * - ``telemetry``
     - 40
     - LIFO
     - demote to ``backlog`` after 300 s
   * - ``inference``
     - 30
     - LIFO
     - demote to ``backlog`` after 300 s
   * - ``artifact``
     - 20
     - FIFO
     - —
   * - ``backlog``
     - 10
     - FIFO
     - —

LIFO ordering means the **newest** record in that class is sent first. This keeps the server receiving fresh observations even under load; records that age past the decay threshold are demoted to ``backlog`` rather than dropped.

PriorityClass
-------------

Custom classes are defined with ``PriorityClass`` from ``scaleoututil.queue``:

.. code-block:: python

   from scaleoututil.queue import PriorityClass, ClassOrder

   PriorityClass(
       name="urgent-alert",   # identifier used in priority= kwargs
       priority=90,           # scheduling weight; higher = more urgent
       order=ClassOrder.FIFO,
       capacity=100,          # optional: drop oldest when class exceeds this depth
       decay=None,            # optional: DecayPolicy
   )

``ClassOrder`` values are ``ClassOrder.FIFO`` and ``ClassOrder.LIFO``.

DecayPolicy
-----------

``DecayPolicy`` applies age-based rules to records inside a class:

.. code-block:: python

   from scaleoututil.queue import DecayPolicy

   DecayPolicy(
       drop_after=60.0,                       # discard records older than 60 s
       demote_after=(300.0, "backlog"),        # or move to another class after 300 s
   )

Only one of ``drop_after`` or ``demote_after`` is typically needed.

Managing Priority Classes
--------------------------

Use the following ``EdgeClient`` methods to inspect or modify the active priority classes at runtime.

``priority_classes`` (property)
   Return the list of currently registered ``PriorityClass`` objects (excludes the internal catch-all default).

``add_priority_class(cls)``
   Register a new class. Raises ``ValueError`` if a class with the same name already exists.

``remove_priority_class(name, force=False)``
   Remove a class by name. If ``force=True``, queued records are migrated to the default class; otherwise raises ``ValueError`` when the class is non-empty. Returns the number of migrated records.

``replace_priority_class(cls)``
   Update an existing class definition in-place. Queued records are not moved — they pick up the new parameters on the next scheduling pass. Raises ``KeyError`` if the name does not exist.

Example: add a high-urgency class and use it for a critical alert:

.. code-block:: python

   from scaleoututil.queue import PriorityClass, ClassOrder

   client.add_priority_class(
       PriorityClass(name="urgent-alert", priority=90, order=ClassOrder.FIFO)
   )

   # Later, in a callback:
   client.log_telemetry("intrusion_detected", {"score": 0.97}, priority="urgent-alert")

Example: cap the telemetry class to avoid accumulating stale samples during a slow link:

.. code-block:: python

   from scaleoututil.queue import PriorityClass, ClassOrder, DecayPolicy

   client.replace_priority_class(
       PriorityClass(
           name="telemetry",
           priority=40,
           order=ClassOrder.LIFO,
           capacity=50,
           decay=DecayPolicy(demote_after=(300.0, "backlog")),
       )
   )


.. meta::
   :description lang=en:
      Reference documentation for the EdgeClient SDK: callback registration, logging, task control, inference, and the priority message system.
   :keywords: Federated Learning, Scaleout, Edge Client, SDK, Priority Queue, Telemetry, Inference, Federated Learning Framework, Scaleout Edge